bmaltais/kohya_ss
Kohya Ss
kohya_ss is a graphical user interface and workbench for fine-tuning diffusion models, specifically designed for Stable Diffusion. It provides a suite of tools for training generative AI models, including specialized interfaces for creating Low-Rank Adaptation weights and training ControlNet spatial control networks.
The project distinguishes itself through integrated VRAM usage optimization and hardware acceleration, featuring specific support for Intel GPUs via XPU-accelerated libraries. It implements parameter-efficient training methods and memory-saving techniques like gradient checkpointing to enable the training of large models on consumer hardware.
The platform covers the entire training lifecycle, from dataset preparation with image bucket organization and caption control to the execution of fine-tuning scripts. It includes capabilities for real-time progress monitoring through in-training sample generation, state recovery via model checkpointing, and the application of advanced training techniques such as masked loss and custom learning schedules.
The software includes automation for environment bootstrapping, dependency management, and containerized deployment options.
Features
- Generative Model Fine-Tuning - Provides a comprehensive graphical workbench for fine-tuning generative diffusion models.
- Model Fine-Tuning - Provides a comprehensive suite for fine-tuning Stable Diffusion models with custom images and captions.
- Custom Diffusion Model Training - Ships a full suite of tools for configuring optimizers and datasets to train diffusion-based neural networks.
- Dataset Preparation Tools - Includes utilities for organizing images into buckets and managing caption tokens for generative AI datasets.
- Stable Diffusion Web Interfaces - Provides a graphical web interface for fine-tuning Stable Diffusion models using custom hyperparameters.
- GPU-Accelerated Training Workbenches - Provides a managed workbench for allocating VRAM and optimizing compute resources during large-scale model training.
- LoRA Training - Implements specialized tools for creating Low-Rank Adaptation weights to inject styles or characters into models.
- Low-Rank Adaptation - Enables parameter-efficient fine-tuning using low-rank adaptation to customize models with minimal hardware requirements.
- Memory Optimization - Minimizes VRAM consumption using techniques like gradient checkpointing and caching to prevent out-of-memory errors.
- Model Training Interfaces - Provides a browser-based graphical interface for configuring and executing diffusion model training scripts.
- Spatial Control Interfaces - Provides a specialized interface for training lightweight ControlNet spatial control networks.
- Lightweight Control Networks - Supports the training of lightweight ControlNet-LLLite networks using image pairs for spatial control.
- Spatial Control Networks - Provides specialized tooling for training ControlNet spatial control networks using image pairs.
- Training Optimizations - Implements training optimizations to manage VRAM and batch sizes for training on consumer hardware.
- VRAM & Compute Optimization - Deno-based balancing of memory usage and processing load via optimizer selection and batch size adjustments.
- Web-Based Configuration Dashboards - Provides a browser-based dashboard for managing training settings and configuring model hyperparameters.
- Aspect Ratio Bucketing - Implements image bucketing to maintain spatial proportions and minimize padding during the dataset preparation process.
- Caption Augmentation Strategies - Provides controls for shuffling word order and dropping captions to influence how the model learns from text descriptions.
- Image Augmentation - Increases dataset variety by applying random image flipping and color hue alterations during the training process.
- Hardware Resource Allocation - Deno-based definition of specific GPU hardware identifiers to be used during the training process.
- Gradient Checkpointing - Implements memory-saving techniques that recompute intermediate activations during backpropagation to reduce VRAM usage.
- Learning Rate Schedulers - Provides control over learning rates through various schedulers and optimizer algorithms.
- Training Utilities - Ships built-in tools for preparing training data and organizing directories for optimal training.
- Training Checkpointers - Provides mechanisms for saving and resuming model training states to ensure session continuity.
- Training Progress Monitors - Features a GPU-accelerated visualization dashboard for real-time monitoring of training metrics and logs.
- In-Training Sample Generation - Produces sample images using active weights and specific prompts during training to assess progress.
- XPU Accelerators - Provides a configuration layer that enables training on Intel GPUs via XPU-accelerated libraries.
- High Resolution Training - Enables fine-tuning of high-resolution generative models using specific resolution settings and network constraints.
- Caption Dataset Utilities - Manages image datasets by organizing them into size-based buckets and handling caption token lengths.
- GPU Device Assignment - Ships an interface for assigning specific GPU hardware identifiers to the training process.
- Intel Hardware Acceleration - Includes specific environment configurations and support for Intel GPUs via XPU-accelerated libraries.
- Resource Allocation - Allows for the allocation of GPU resources and precision levels to maximize compute performance.
- Compute Environment Toggles - Provides toggles to optimize performance across different hardware architectures and compute environments.
