[Core] add: controlnet support for SDXL

[sayakpaul]

This PR adds support for ControlNets with SDXL. The two primary components being added to this PR:

• Training script train_controlnet_sdxl.py.
• Pipeline StableDiffusionXLControlNetPipeline (with changes to ControlNetModel to accommodate the pipeline-level changes).

However, these seems to be something weird going on here.

I first started training on a small subset of dataset (the circles dataset) with the following command:

export MODEL_DIR="stabilityai/stable-diffusion-xl-base-0.9"
export OUTPUT_DIR="controlnet-sdxl-circles"

wget https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/controlnet_training/conditioning_image_1.png

wget https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/controlnet_training/conditioning_image_2.png

accelerate launch train_controlnet_sdxl.py \
 --pretrained_model_name_or_path=$MODEL_DIR \
 --output_dir=$OUTPUT_DIR \
 --dataset_name=fusing/fill50k \
 --mixed_precision="fp16" \
 --resolution=1024 \
 --learning_rate=5e-5 \
 --max_train_samples=500 \
 --max_train_steps=1000 \
 --validation_image "./conditioning_image_1.png" "./conditioning_image_2.png" \
 --validation_prompt "red circle with blue background" "cyan circle with brown floral background" \
 --validation_steps=25 \
 --train_batch_size=1 \
 --gradient_accumulation_steps=4 \
 --report_to="wandb" \
 --seed=42 \
 --push_to_hub

The trained checkpoints seem to only generate black images: https://huggingface.co/fusing/controlnet-sdxl-circles-fixed (only visible to the diffusers team members).

To further debug this, I tried:

from diffusers import StableDiffusionXLControlNetPipeline, ControlNetModel
from diffusers.utils import load_image
import torch 

base_ckpt_id = "stabilityai/stable-diffusion-xl-base-0.9"
controlnet_ckpt_id = "controlnet-sdxl-circles-fixed"

controlnet = ControlNetModel.from_pretrained(
	controlnet_ckpt_id, subfolder="checkpoint-500/controlnet", torch_dtype=torch.float16
).to("cuda")

pipeline = StableDiffusionXLControlNetPipeline.from_pretrained(
	base_ckpt_id, controlnet=controlnet, torch_dtype=torch.float16
).to("cuda")


cond_image = load_image(
    "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/controlnet_training/conditioning_image_1.png"
)
prompt = "red circle with blue background"

image = pipeline(prompt, image=cond_image).images[0]
image.save("controlnet@ckpt-500.png")

This doesn't generate the expected results (which is expected since the number of training steps is quite low) but doesn't generate all black images either.

@patrickvonplaten @williamberman could you take a deeper look here?

TODOs

• tests
• docs
• misc changes

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[gkorepanov]

The trained checkpoints seem to only generate black images:

Hi! Do you mean that in validation during training all images are black, but if you manually load trained checkpoint using external script, the images are fine?

[sayakpaul]

Hi! Do you mean that in validation during training all images are black, but if you manually load trained checkpoint using external script, the images are fine?

Exactly.

[gkorepanov]

Exactly.

I think that might relate to SDXL VAE producing NANs in some cases with fp16 mode.

From https://github.com/kohya-ss/sd-scripts/tree/sdxl:

The image generation during training is now available. However, the VAE for SDXL seems to produce NaNs in some cases when using fp16. The images will be black. Currently, the NaNs cannot be avoided even with --no_half_vae option. It works with bf16 or without mixed precision.

Also:
https://huggingface.co/stabilityai/sdxl-vae/discussions/6
https://huggingface.co/madebyollin/sdxl-vae-fp16-fix

[sayakpaul]

Thanks for being willing to help.

I think the issue with VAE is handled. See: https://github.com/huggingface/diffusers/blob/db78a4cb4e3f105cbc7534890f606e25e906e23a/src/diffusers/pipelines/controlnet/pipeline_controlnet_sd_xl.py#L1118C1-L1133C38.

Also, when I run the manual validation, it's in FP16 only.

[gkorepanov]

I think the issue with VAE is handled

Ah, really, seems so, thanks. BTW, in the code you mentioned there might be a small bug with unnecessary not which was recently fixed in the main branch: #4019

Also, to run your code, I had to put extra StableDiffusionXLControlNetPipeline imports in few places, I think you have forgotten to include few __init__.py into PR

[sayakpaul]

@gkorepanov thanks so much for your catches. I incorporated the fixes. Let me run the dummy experiment one more time to check quickly.

[laksjdjf]

Is it because autocast is used to generate the validation image?

diffusers/examples/controlnet/train_controlnet_sdxl.py

Lines 123 to 126 in 68f2c38

	with torch.autocast("cuda"):
	image = pipeline(
	validation_prompt, validation_image, num_inference_steps=20, generator=generator
	).images[0]

kohya-ss's problems also seem to have been caused by autocast.
kohya-ss/sd-scripts@814996b#diff-5f48c8e976d43e587007dc13a34100a96621cdc5fbe083ee772e920855648722R3877

[patrickvonplaten]

Cool! Let's make sure we have a working controlnet training run before merging this though :-)

[sayakpaul]

Cool! Let's make sure we have a working controlnet training run before merging this though :-)

There's a working script in the PR. The issues described in the original post is why I am seeking reviews for.

[gkorepanov]

There's a working script in the PR. The issues described in the original post is why I am seeking reviews for.

After disabling autocast in validation and using torch.float32 when loading the pipeline the validation looks better (at least images are not black anymore):

But images seem to be awkward.

[gkorepanov]

But images seem to be awkward.

The difference was caused by different resolution in inference. By default, controlnet pipeline takes height/width from control image

[sayakpaul]

@gkorepanov thanks again for your inputs! Very much appreciated.

Let me run a couple of experiments now.

[sayakpaul]

@gkorepanov may I know which GPU model did you use for your tests? I am currently using a 40GB A100 and when I try log_validation() in FP32, it OOMs.

[sayakpaul]

@patrickvonplaten thanks for all the reviews. A final review and I think we're good to go. Let me know.

[adhikjoshi]

Will existing controlnet 1.1 checkpoints work here?

[sayakpaul]

Will existing controlnet 1.1 checkpoints work here?

No.

[patrickvonplaten]

Cool!

[sayakpaul]

@gkorepanov we start a PR for adding switching support and MultiControlNet too since switching likely impacts that more.

Let me know :)

[sayakpaul]

@williamberman #4188.

[zdxpan]

tranning met loss is nan, and the pred_noise cntain nan, which will case the tranning fail， the traing target is predict noise with given noise, which mse alwas nearby 1 (~= 1),
possible reason
1、 lr too large
2、 data had some nan value
3、is there anty other reason?

and met the log_validate validate image alwas balck

[sayakpaul]

Try passing the following as your VAE: madebyollin/sdxl-vae-fp16-fix.

Additionally, you can ask questions on the repositories like the following, which leveraged our training scripts to obtain nice results: https://huggingface.co/thibaud/controlnet-openpose-sdxl-1.0/discussions.

[patrickvonplaten]

tranning met loss is nan, and the pred_noise cntain nan, which will case the tranning fail， the traing target is predict noise with given noise, which mse alwas nearby 1 (~= 1), possible reason 1、 lr too large 2、 data had some nan value 3、is there anty other reason?

and met the log_validate validate image alwas balck

@zdxpan please make sure to open a new issue instead of commenting on the PR here