Navigate the Universe of Diffusion models with Unified workflow.

UniDiffusion is a toolbox that provides state-of-the-art training and inference algorithms, based on diffusers. UniDiffusion is aimed at researchers and users who wish to deeply customize the training of stable diffusion. We hope that this code repository can provide excellent support for future research and application extensions.

If you also want to implement the following things, have fun with UniDiffusion

• Train only cross attention (or convolution / feedforward / ...) layer.
• Set different lr / weight decay / ... for different layers.
• Using or supporting PEFT/PETL methods for different layers and easily merging them, e.g., finetune the convolution layer and update attention layer with lora.
• Train all parameter in stable diffusion, including unet, vae, text_encoder, and automatically save and load.

Note: UniDiffusion is still under development. Some modules are borrowed from other code repositories and have not been tested yet, especially the components that are not enabled by default in the configuration system. We are working hard to improve this project.

• Modular Design. UniDiffusion is designed with a modular architecture. The modular design enables easy implementation of new methods.
• Config System. LazyConfig System for more flexible syntax and cleaner config files.
• Easy to Use.
  • Distributed Training: Using accelerate to support all distributed training environment.
  • Experiment Tracker: Using wandb to log all training information.
  • Distributed Evaluation: Evaluate ✅FID, ✅IS, CLIP Score during training

In UniDiffusion, all training methods are decomposed into three dimensions

• Learnable parameters: which layer or which module will be updated.
• PEFT/PETL method: how to update them. E.g., finetune, low-rank adaption, adapter, etc.
• Training process: default to diffuion-denoising, which can be extended like XTI.

It allows we conduct a unified training pipeline with strong config system.

python train_dreambooth.py --arg ......
python train_finetune.py --arg ......

# text-to-image finetune
unet.training_args = {'': {'mode': 'finetune'}}
# text-to-image finetune with lora
unet.training_args = {'': {'mode': 'lora'}}
# update cross attention with lora
unet.training_args = {'attn2': {'mode': 'lora'}}

# dreambooth
unet.training_args = {'': {'mode': 'finetune'}}
text_encoder.training_args = {'text_embedding': {'initial': True}}
# dreambooth with small lr for text-encoder
unet.training_args = {'': {'mode': 'finetune'}}
text_encoder.training_args = {'text_embedding': {'initial': True, 'optim_kwargs': {'lr': 1e-6}}}

accelerate launch scripts/train.py --config-file /path/to/your/config

In UniDiffusion, we provide a regular matching system for module selection. It allows us to select modules by regular matching. See Regular Matching for Module Selection for more details.

We provide a powerful support for PEFT/PETL methods. See PEFT/PETL Methods for more details.

1. Install prerequisites

• Python 3.10
• Pytorch 2.0 + CUDA11.8
• CUDNN

2. Install requirements

pip install -e requirements.txt

3. Configuring accelerate and wandb

accelerate config
wandb login

See Train textual inversion / Dreambooth / LoRA / text-to-image Finetune for details.

accelerate launch scrits/common.py --config-file configs/train/text_to_image_finetune.py

1. Train textual inversion / Dreambooth / LoRA / text-to-image Finetune.
2. Customize your training process.

1. [TODO] Supporting new dataset.
2. [TODO] Supporting new PETL method.
3. [TODO] Supporting new training pipeline.

We are going to add the following features in the future. We also welcome contributions from the community. Feel free to pull requests or open an issue to discuss ideas for new features.

• Methods:
  • preservation of class semantic priors (dreambooth).
  • XTI & Custom Diffusion.
  • RepAdapter and LyCORIS.
• Features:
  • Merge PEFT to original model.
  • Convert model to diffusers and webui format.
  • Webui extension.

We welcome contributions from the open-source community!

• Diffusion Trainer is built based on diffusers.
• A lot of module design is borrowed from detectron2 and detrex.
• Some implementations of methods is borrowed from diffusers and LyCORIS.

If you use this toolbox in your research or wish to refer to the baseline results published here, please use the following BibTeX entries:

• Citing UniDiffusion:

@misc{pu2022diffusion,
  author =       {Pu Cao, Tianrui Huang, Lu Yang, Qing Song},
  title =        {UniDiffusion},
  howpublished = {\url{https://github.com/PRIV-Creation/UniDiffusion}},
  year =         {2023}
}