[Pipelines] Add a ControlNet pipeline

[apolinario]

Model/Pipeline/Scheduler description

ControlNet by @lllyasviel is a neural network structure to control diffusion models by adding extra conditions.

It has integration with Stable Diffusion and 8 pre-trained models that conditions the models on different attributes (such as edge detection, scribbles, depth maps, semantic segmentations and more)

Would be great if this was added to diffusers as a pipeline (and probably with a method to load the different models).

Scribble maps condition model	Human pose condition model

Open source status

• The model implementation is available
• The model weights are available (Only relevant if addition is not a scheduler).

Provide useful links for the implementation

Original code: https://github.com/lllyasviel/ControlNet
Pre-trained models: https://huggingface.co/lllyasviel/ControlNet/tree/main/models
Hugging Face Spaces demo: https://huggingface.co/spaces/RamAnanth1/ControlNet

[takuma104]

Hi @apolinario,

ControlNet seems to offer great output quality and seems to be very versatile. I am very interested in getting involved.

I'm thinking of first trying to proceed with the PoC level this next week or so as follows. This is just my idea, so please let me know if you have any suggestions on how this should be done.

1. Load and convert ckpt: needed for E2E comparison with reference implementation
   • add function to convert_from_ckpt.py
2. ControlNet class equivalent: inherit from UNet2DConditionModel class Mixins
   • Porting input_hint_block and zero_conv
3. ControlledUnetModel class equivalent: inherit from UNet2DConditionModel class
   • Override forward()
4. Minimal pipeline.
   • Do not implement the control embedding (c_concat) calculation in the pipeline, but make it an argument of pipe.__call__() to avoid pipeline abuse.
5. Testing
   • Ensure that all current unit tests pass
   • E2E compare with reference implementation. Output is using the same sampler (DDIM) to check pixel match.
   • Additional unit tests for new ControlNet and ControlledUnetModel equivalent classes and new pipeline.

[takuma104]

This is a WIP code for 1.
main...takuma104:diffusers:controlnet

It can now convert to Diffusers format without error messages. However, the ControlNetModel implementation is not yet complete, so the data is not correct.
After writing all this, I wonder if I should write it as a Community Pipeline? I am concerned about the impact as this will be a fairly large modification.

[apolinario]

Very cool @takuma104, thanks a lot! cc @patrickvonplaten, @patil-suraj, @williamberman on best practices and trade-off of impact/modifications.

I think this technique is showing how adding extra conditions to diffusion models (in general) has huge potential and how the pre-trained models are impactful

So I think an official pipeline makes sense here, but the trade-off of the modifications it requires may be worth debating a bit.

[takuma104]

@apolinario,
Thanks! Sounds nice, I too think this model is worth making official.

@patrickvonplaten, @patil-suraj, @williamberman
Here is an overview of my current code diff. My basic policy is to try to avoid interfering with the original code as much as possible, but I'm not sure if this is a good idea in terms of future maintainability, so I'd like your feedback.

• scripts/convert_controlnet_to_diffusers.py

  • This is mostly a copy of convert_original_stable_diffusion_to_diffusers.py. This file is pretty simple and has a good outlook, so I think I'll leave it as is. There may be some adjustments to some of the command line options.

• src/diffusers/models/unet_2d_condition_controlnet.py

  • This is a copy of unet_2d_condition.py, the file where I am trying to create the ControlNet model. It's a work in progress so I can't see what's coming, but I can assume that a lot of the code will be copied from UNet2DConditionModel.

• src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_controlnet.py

  • I will write this by copying from pipeline_stable_diffusion.py. I will modify it based on a copy of __call__().

• src/diffusers/pipelines/stable_diffusion/convert_from_ckpt.py

  • Add load_pipeline_from_control_net_ckpt() as an API. This is used by convert_controlnet_to_diffusers.py to convert to Diffusers format.
  • Most of the additional code is copied from the original and is quite verbose. I wonder if I should modify the original code by making it as splitting functions, or if this way of copying is better. Or should I split the files themselves? (e.g. by partially calling the function in convert_from_ckpt.py)

• I've also modified the __init__.py in each folder hierarchy to propagate the class names.

• No test files have been added yet, one file each would be added to tests/models and test/pipelines.

[takuma104]

About file placement:
I think the file layout depends on how we view ControlNet: is it an extension of StableDiffusion like LoRA or Hypernetwork, or is it a different model since it uses two Unet equivalents and is quite special?

If the latter, we could put all the necessary dependency files in the src/diffusers/pipelines/control_net folder, including the model file. Would this be preferable?

[xvjiarui]

Wow, nicely done! Looking forward to the completed version!

Since one important application of ControlNet is to initialize from Stable Diffusion and fine-tune on user's customized dataset, would it make sense to have some instructions and methods for this?

[takuma104]

Hi @xvjiarui , thanks! I agree with you. I was thinking of minimizing the PoC, so I'm thinking of just supporting inference. Once the PoC is complete, could you try writing some code for fine-tune? :)

[xvjiarui]

Sure. Looking forward to it.

[williamberman]

Super cool model!

I've only skimmed the paper so these are just preliminary thoughts. I think we can add this model cleanly but there are some nuanced considerations. We definitely don't want to add a separate unet implementation if possible. I think we have to mainly consider inference here as for training and weight loading, there is a whole independent unet that has to be trained/loaded (i.e. there are no lora-esque considerations for efficient weight training/loading maybe modulo loading the original unet).

My basic understanding:
controlnet uses two separate unets.

unet 1: Same architecture and weights as original stable diffusion. The only difference here is that residual connections are added together with outputs from unet2. This means that our existing conditional unet has to take these "additional residuals" as input¹.

unet 2: zero conv modified architecture stable diffusion unet. This is the trained unet which takes in the additional conditioning. The only (?) architecture modification from the original unet is the zero conv blocks. It looks like the output of the decoder blocks are fed through the zero convs before being passed to the residual connections in unet1.

We will have to return the outputs of the unet decoder for passing to the residuals of unet1².

If the zero conv blocks are only used "between" unets i.e. unet2 doesn't use the zero conv output before passing to the next layer in the decoder, we should be able to avoid any changes to existing unet blocks. The zero conv layers can all be stored as a separate model independent of unet 2³.

If the zero conv blocks are used w/in unet2, we will either have to duplicate the unet blocks to separate unet blocks that include the zero conv layer, or we will have to add a flag to the relevant unet blocks for them to include an optional zero conv layer⁴.

Separately,
There is an additional model maps the additional conditioning into the same dimensionality as the latent images. It sounds straight forward to add.

Footnotes

1. This is a small and ok change. ↩

2. This is also a small and ok change. ↩

3. This is just a helper model local to the pipeline, also ok. ↩

4. The flag is a small change but it is network specific. The separate blocks are isolated but will be more copy-paste. I'm leaning towards the flag but one of the two should be ok ↩

[takuma104]

Hi @williamberman ,

Thanks for pointing me in the right direction! I was aware that I was building a bunch of copy-paste and was concerned about maintainability. As you said, it would not be difficult to implement unet1 and unet2 by changing the behavior of UNet2DConditionModel's __init__() and forward() arguments for both unet1 and unet2. I'll make a plan for implementation.

I think I have implemented a PoC version of the forward path equivalent to unet1 and unet2, and the difference from the original code is as follows. FYI.

• unet1 as ControlledUNet2DConditionModel
• unet2 as ControlNetModel

[williamberman]

Thanks! Feel free to keep me updated :)