lavinal712 / control-lora-v3

ControlLoRA Version 3: A Lightweight Neural Network To Control Stable Diffusion Spatial Information Version 3

ControlLoRA Version 3 is a neural network structure extended from ControlNet to control diffusion models by adding extra conditions.

Inspired by ControlLoRA, control-lora-v2 and script train_controlnet.py from diffusers, control-lora-v3 does not add new features, but provides a PEFT implement of ControlLoRA.

News

• Jun. 08, 2024. Norm layer is trainable.
• May. 19, 2024. Add DoRA.

Data

To train ControlLoRA, you should have image-conditioning_image-text datasets. Of course you can hardly train on LAION-5B dataset in direct like Stable Diffusion. Here are some:

• fusing/fill50k. I do not suggest you to train ControlLoRA seriously as it is simple and lack of diversity.
• HighCWu/diffusiondb_2m_first_5k_canny. A small canny dataset. Here is poloclub/diffusiondb dataset. For canny condition, you can easily generate your own dataset.
• Nahrawy/VIDIT-Depth-ControlNet. Depth map? Heat map? But it is good!
• SaffalPoosh/scribble_controlnet_dataset. Many duplicate images. I suggest you synthesize your dataset.

Model

Stable Diffusion v1-5 is the base model.

Stable Diffusion v1-4, Stable Diffusion v2-1 need to be vertified.

Stable Diffusion XL needs to be vertified, but probably does not work.

Train

You can train either ControlNet or ControlLoRA using script train_control_lora.py.

Train ControlNet

By observation, training 50000 steps with batch size of 4 is the balance between image quality, control ability and time.

accelerate launch train_control_lora.py \
 --pretrained_model_name_or_path="runwayml/stable-diffusion-v1-5" \
 --output_dir="controlnet-model" \
 --dataset_name="fusing/fill50k" \
 --resolution=512 \
 --learning_rate=1e-5 \
 --train_batch_size=4 \
 --max_train_steps=100000 \
 --tracker_project_name="controlnet" \
 --checkpointing_steps=5000 \
 --validation_steps=5000 \
 --report_to wandb

Train ControlLoRA

To train ControlLoRA, add --use_lora in start command to activate it.

accelerate launch train_control_lora.py \
 --pretrained_model_name_or_path="runwayml/stable-diffusion-v1-5" \
 --output_dir="control_lora-model" \
 --dataset_name="fusing/fill50k" \
 --resolution=512 \
 --learning_rate=1e-4 \
 --train_batch_size=4 \
 --max_train_steps=100000 \
 --tracker_project_name="control_lora" \
 --checkpointing_steps=5000 \
 --validation_steps=5000 \
 --report_to wandb \
 --use_lora \
 --lora_r=32

You can also train ControlLoRA / ControlNet with your own dataset.

accelerate launch train_control_lora.py \
 --pretrained_model_name_or_path="stable-diffusion-v1-5" \
 --output_dir="control_lora-model" \
 --conditioning_image_column="hint" \
 --image_column="jpg" \
 --caption_column="txt" \
 --resolution=512 \
 --learning_rate=1e-4 \
 --train_batch_size=4 \
 --num_train_epochs=3 \
 --max_train_steps=100000 \
 --tracker_project_name="control_lora" \
 --checkpointing_steps=5000 \
 --validation_steps=5000 \
 --report_to wandb \
 --use_lora \
 --lora_r=32 \
 --custom_dataset="fill50k"

Merge

If you want to train ControlNet, you have already got it. If you got a lora, merge it!

python merge_lora.py

Test

Original image:

Output:

Citation

@software{lavinal7122024controllorav3,
    author = {lavinal712},
    month = {5},
    title = {{ControlLoRA Version 3: A Lightweight Neural Network To Control Stable Diffusion Spatial Information Version 3}},
    url = {https://github.com/lavinal712/control-lora-v3},
    version = {1.0.0},
    year = {2024}
}

control-lora-v3's Issues

run with distilled model

Is it possible to run distilled model (Block-removed Knowledge-distilled Stable Diffusion, https://github.com/Nota-NetsPresso/BK-SDM).

When I tried to some distilled model. I got error message similar to "unknown mid_block_type" , probably because of unet structure changed.