Scalable Stable Diffusion Trainer
- Can run with <8GB VRAM
- Aspect Ratio Bucketing
- CLIP skip
- WandB logging
Customizable training objective, including:
Linux is recommended, on Windows you have to install bitsandbytes
manually for int8 optimizers.
conda env create -f environment.yml
conda activate ssdt
Python 3.10 is required. CUDA toolkit and torch should be installed manually.
pip install -r requirements.txt
Documentation: configs/README.md
.
(Link)
In configs
directory, native.yaml
(for so-called native training), dreambooth.yaml
, lora.yaml
provided as
examples.
python train.py --config configs/your_config.yaml
But if you are running DreamBooth, run this first to generate regularization images:
python gen_class_imgs.py --config configs/your_config.yaml
WebUI may not directly load the checkpoints due to the safe loading mechanism. To solve this issue and reduce checkpoint size:
python ckpt_tool.py prune INPUT OUTPUT --unet-dtype fp16
INPUT
is the path to the trained SCAL-SDT checkpoint.
If OUTPUT
has suffix .safetensors
then safetensors format will be used.
A ~1.6GB file will be created, which can be loaded by WebUI, containing fp16 UNet states.
If you are not using WebUI and having issues, specify both --text-encoder
and --vae
and remove --unet-dtype fp16
to get a full checkpoint.