[WIP] [Core] Add support for ControlNet LoRA

[sayakpaul]

This PR introduces support for ControlNet LoRA.

Why the design?

You might wonder why not unify everything in load_lora_weights(). The reason why I decided not to do that is because there might be use cases where people load UNet LoRA checkpoints as well as ControlNet LoRA checkpoints. Plus ControlNet LoRA will always depend on a ControlNet that's initialized from a UNet.

The current design

Load things:

fromdiffusersimport (
StableDiffusionXLControlNetPipeline, ControlNetModel, UNet2DConditionModel
)
importtorchpipe_id="stabilityai/stable-diffusion-xl-base-1.0"lora_id="stabilityai/control-lora"lora_filename="control-LoRAs-rank128/control-lora-canny-rank128.safetensors"unet=UNet2DConditionModel.from_pretrained(pipe_id, subfolder="unet", torch_dtype=torch.float16).to("cuda")
controlnet=ControlNetModel.from_unet(unet).to(device="cuda", dtype=torch.float16)
controlnet.load_lora_weights(lora_id, weight_name=lora_filename, controlnet_config=controlnet.config)

Infer:

fromdiffusersimportAutoencoderKLfromdiffusers.utilsimportload_image, make_image_gridfromPILimportImageimportnumpyasnpimportcv2prompt="aerial view, a futuristic research complex in a bright foggy jungle, hard lighting"negative_prompt="low quality, bad quality, sketches"image=load_image("https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/hf-logo.png")
controlnet_conditioning_scale=0.5# recommended for good generalizationvae=AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
pipe=StableDiffusionXLControlNetPipeline.from_pretrained(
pipe_id,
unet=unet,
controlnet=controlnet,
vae=vae,
torch_dtype=torch.float16,
).to("cuda")
image=np.array(image)
image=cv2.Canny(image, 100, 200)
image=image[:, :, None]
image=np.concatenate([image, image, image], axis=2)
image=Image.fromarray(image)
images=pipe(
prompt, negative_prompt=negative_prompt, image=image,
controlnet_conditioning_scale=controlnet_conditioning_scale,
num_images_per_prompt=4
).imagesfinal_image= [image] +imagesgrid=make_image_grid(final_image, 1, 5)
grid

Results don't seem to be following the conditioning image at all. I will play with the hyperparameters a bit to see how the results get affected.

---

@patrickvonplaten could you give this a look and comment on what I might be missing out on?

[sayakpaul]

@patrickvonplaten here's what I have found out so far. Note that this is entirely my understanding and I am happy to stand corrected.

Control-LoRA workflow as per ComfyUI

For debugging purposes, I developed a simplified Colab Notebook out of the official Colab Notebook provided in the ComfyUI repository.

From there one, I narrowed it down to the following LoCs that are responsible for initializing a ControlNet model and populating it with the Control-LoRA parameters (such as the ones found here and here).

(For debugging, I am using this checkpoint: https://huggingface.co/stabilityai/control-lora/blob/main/control-LoRAs-rank128/control-lora-canny-rank128.safetensors (Canny).)

• https://github.com/comfyanonymous/ComfyUI/blob/1c8ae9dbb249ed5326d61d16b4e6b5807c09c0e1/comfy/controlnet.py#L291 -- delegates the provided checkpoint to ControlNet, one that is compatible with the LoRA checkpoints.
• https://github.com/comfyanonymous/ComfyUI/blob/1c8ae9dbb249ed5326d61d16b4e6b5807c09c0e1/comfy/controlnet.py#L260C1-L269C141 -- loads the LoRA checkpoint for the compatible layers. You would also notice that some checkpoints from the underlying UNet (UNet of SDXL base in this case; which you can find out from the Colab as well) are also used during initialization.
• LoRA-specific computation bits:
  • https://github.com/comfyanonymous/ComfyUI/blob/1c8ae9dbb249ed5326d61d16b4e6b5807c09c0e1/comfy/controlnet.py#L190
  • https://github.com/comfyanonymous/ComfyUI/blob/1c8ae9dbb249ed5326d61d16b4e6b5807c09c0e1/comfy/controlnet.py#L229

What we're doing

Now, based on this observation, I tried to do a similar thing in this PR:

• Point of entry is this class: ControlLoRALoaderMixin:
  

  diffusers/src/diffusers/loaders.py

  Line 2517 in 86f5980

  classControlLoRAMixin(LoraLoaderMixin):

• Then the original checkpoint is prepared, leveraging a modified version convert_ldm_unet_checkpoint():

  diffusers/src/diffusers/loaders.py

  Line 2532 in 86f5980

  converted_state_dict=convert_ldm_unet_checkpoint(

• Then load the converted checkpoint into the underlying ControlNet for whatever keys that are getting matched:

  diffusers/src/diffusers/loaders.py

  Line 2537 in 86f5980

  load_state_dict_results=self.load_state_dict(converted_state_dict, strict=False)

  (following the similar approach as ComfyUI).
• Then I populate the LoRA-specific parameters. Follow from here:

  diffusers/src/diffusers/loaders.py

  Line 2544 in 86f5980

  remaining_state_dict= {

  .

This along with the code snippet I provided in my PR description, does follow what ComfyUI is doing. But as mentioned above, I am happy to stand corrected.

[patrickvonplaten]

@patrickvonplaten here's what I have found out so far. Note that this is entirely my understanding and I am happy to stand corrected.

Control-LoRA workflow as per ComfyUI

For debugging purposes, I developed a simplified Colab Notebook out of the official Colab Notebook provided in the ComfyUI repository.

From there one, I narrowed it down to the following LoCs that are responsible for initializing a ControlNet model and populating it with the Control-LoRA parameters (such as the ones found here and here).

(For debugging, I am using this checkpoint: https://huggingface.co/stabilityai/control-lora/blob/main/control-LoRAs-rank128/control-lora-canny-rank128.safetensors (Canny).)

• https://github.com/comfyanonymous/ComfyUI/blob/1c8ae9dbb249ed5326d61d16b4e6b5807c09c0e1/comfy/controlnet.py#L291 -- delegates the provided checkpoint to ControlNet, one that is compatible with the LoRA checkpoints.

• https://github.com/comfyanonymous/ComfyUI/blob/1c8ae9dbb249ed5326d61d16b4e6b5807c09c0e1/comfy/controlnet.py#L260C1-L269C141 -- loads the LoRA checkpoint for the compatible layers. You would also notice that some checkpoints from the underlying UNet (UNet of SDXL base in this case; which you can find out from the Colab as well) are also used during initialization.

• LoRA-specific computation bits:

  • https://github.com/comfyanonymous/ComfyUI/blob/1c8ae9dbb249ed5326d61d16b4e6b5807c09c0e1/comfy/controlnet.py#L190
  • https://github.com/comfyanonymous/ComfyUI/blob/1c8ae9dbb249ed5326d61d16b4e6b5807c09c0e1/comfy/controlnet.py#L229

What we're doing

Now, based on this observation, I tried to do a similar thing in this PR:

• Point of entry is this class: ControlLoRALoaderMixin:
  

  diffusers/src/diffusers/loaders.py

  Line 2517 in 86f5980

  classControlLoRAMixin(LoraLoaderMixin):

• Then the original checkpoint is prepared, leveraging a modified version convert_ldm_unet_checkpoint():

  diffusers/src/diffusers/loaders.py

  Line 2532 in 86f5980

  converted_state_dict=convert_ldm_unet_checkpoint(

• Then load the converted checkpoint into the underlying ControlNet for whatever keys that are getting matched:

  diffusers/src/diffusers/loaders.py

  Line 2537 in 86f5980

  load_state_dict_results=self.load_state_dict(converted_state_dict, strict=False)

  
  (following the similar approach as ComfyUI).
• Then I populate the LoRA-specific parameters. Follow from here:

  diffusers/src/diffusers/loaders.py

  Line 2544 in 86f5980

  remaining_state_dict= {

  
  .

This along with the code snippet I provided in my PR description, does follow what ComfyUI is doing. But as mentioned above, I am happy to stand corrected.

Where does a divergence appear? If you get ComfyUI working where is the point that the official implementation starts to differ from ours. E.g. where do numerical values diverge?

[sayakpaul]

Where does a divergence appear? If you get ComfyUI working where is the point that the official implementation starts to differ from ours. E.g. where do numerical values diverge?

For now, I am only considering the final results. I couldn't identify any implementation-related differences in our implementation which is why I tried to detail as much as possible in my previous comment.

[patrickvonplaten]

only

Can you set intermediate print statements to find out at which point the results start to diverge?

[clarencechen]

@sayakpaul This gist (https://gist.github.com/madebyollin/e6e217a77db94e2a960ab6ccbd627db9) by @madebyollin helped him debug and fix the SDXL VAE when running in FP16. It may also be helpful with comparing and debugging the Control LoRA activations.

[sayakpaul]

@clarencechen you mean we could reuse the gist to see if there's a problem FP16 inference here?

[clarencechen]

@clarencechen you mean we could reuse the gist to see if there's a problem FP16 inference here?

I think you can reuse the gist to help track and diagnose most numerical discrepancies in Pytorch models.

[adhikjoshi]

this is not a priority anymore?

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

[sayakpaul]

Closing for #5938.

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

[lavinal712]

I use diffusers to load control-lora canny from https://huggingface.co/stabilityai/control-lora. Here is the result.

[lavinal712]

@sayakpaul Does this meet your expected outcome?