Add ControlNet-XS support

[UmerHA]

What does this PR do?

Adds ControlNet-XS support (and therefore fixes #5168).
Project page: https://vislearn.github.io/ControlNet-XS/

See here for a full working example

This PR is work in progress. Still to do:

• Add other version of ControlNet-XS: SD canny ✅, SD depth ✅, SDXL depth ✅
• Add documentation ✅
• A few other (iiuc) minor things ✅

Still, I would love your feedback!

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline?
• Did you read our philosophy doc (important for complex PRs)?
• Was this discussed/approved via a GitHub issue or the forum? Please add a link to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of who to tag.
Please tag fewer than 3 people.

Core library:

• Pipelines: @patrickvonplaten and @sayakpaul

[sayakpaul]

Cc: @DN6

[patrickvonplaten]

Amazing job @UmerHA

[universewill]

@UmerHA Great work! Can you provide controlnet-xs traning example?

[UmerHA]

@universewill Sure - see https://github.com/UmerHA/diffusers/tree/cnxs-training/examples/controlnet_xs. I've tested that they they run, but haven't fully tested full training runs. When I have more time, I'll do that and open a PR.

In the meantime, let me know if you encounter any issues!

[patrickvonplaten]

Sorry, we had to move the implementation to the research folder for now as the design was not in line with the usual diffusers design (e.g. the unet is forwarded into the controlnet-xs function etc...). We should have caught that when reviewing the PR, but sadly failed to do so. We still very much want to add ControlNet-XS to diffusers, but we'll need to apply the changes as suggested by the following new PR review.

Very sorry @UmerHA that we missed these things in the initial review 🙏

[patrickvonplaten]

The control_model should not be a UNet2DConditionModel if we have to apply a lot of surgery afterwards. Let's make sure we directly instantiate the correct torch.nn.Modules right away

[patrickvonplaten]

We can't do surgery here, let's make sure to instead instantiate the correct classes right away

[patrickvonplaten]

They should instead never have been instantiated

[patrickvonplaten]

Let's not provide a set_attention_slice() operation anymore since with FlashAttention it's pretty useless to slice the attention

[patrickvonplaten]

We can't pass the base_model here into the forward method

[patrickvonplaten]

Do we need the class_labels parameter?

[patrickvonplaten]

Do we really need this? I think the class labels are only required for Stable Diffusion Upsampling

[patrickvonplaten]

We should not call the base model here - instead self.add_embedding(...) should be called

[patrickvonplaten]

We should not have to call such a method in the forward pass of the controlnet

[patrickvonplaten]

Let's try to avoid map here as this breaks torch.compile

[patrickvonplaten]

Let's not use such a try-except function, this breaks torch.compile

[patrickvonplaten]

We should not have to use such surgery methods

[UmerHA]

Hi @patrickvonplaten, no hurt feelings & fully understand. I'll start to change it to better fit diffusers, and open a new PR.

Happy New Year btw :)

[UmerHA]

Hi @patrickvonplaten, could you answer two questions before I start the new implementation:

---

1. Does the following design work for you?

# UNet as it currently is. Will not be changed. Shown for completeness
class UNet2DConditionModel(...):
    # ...

# Class containing the controlnet-relevant weights
# Having a separate class ensures users can save/load the controlnet-relevant weights separately,
# and don't always have to save/load the large base model
class ControlNetXSAddonModel(...):
    def __init__(self):
        # Initiate what currently is called the control model.
        # Instead of starting with a UNet2DConditionModel and doing model surgery, build it correctly from scratch
        
    def forward(self, x):
        raise ValueError("A ControlNetXSAddonModel cannot be run by itself. Pass it into a ControlNetXSModel model instead.")

# Full ControlNetXSModel that will be compilable
class ControlNetXSModel(...):
    def __init__(self, base_model: UNet2DConditionModel, ctrl_model: ControlNetXSAddonModel):
        self.ctrl_model = ctrl_model
        self.base_model = base_model

    def forward(self, x):
        # other inputs (time, conditioning, ...) omitted for brevity 

        # decompose both models into the parts that need to be run interwovenly.
        # for brevity, let's use blocks. in actual implementation we need to use "subblocks".
        base_down_blocks = self.base_model.down_blocks
        ctrl_down_blocks = self.ctrl_model.down_blocks

        x_base = x_ctrl = x

        for b, c in zip(base_down_blocks, ctrl_down_blocks):
            x_ctrl = torch.cat([x_base, x_ctrl])  # concat information from base to ctrl
            x_ctrl = c(x_ctrl)                    # execute ctrl block
            x_base = b(x_base)                    # execute base block
            x_base += x_ctrl                      # add information from ctrl to base

        return x_base

class StableDiffusionControlNetXSPipeline(...):
    def __init__(
        ...
        unet: UNet2DConditionModel,
        controlnet_addon: ControlNetXSAddonModel,
        ...
    ):
        self.controlnet = ControlNetXSModel(
            base_model=unet,
            ctrl_model=controlnet_addon
        )
        ...
    
    def __call__(self, prompt)
        # other inputs omitted for brevity 
        ...
        latent = ...
        timesteps = ..
        ...
        for t in timesteps:
            noise_prediction = self.controlnet(sample=latent, timestep=t, ...)
        ...
        return ..

controlnet_addon = ControlNetXSAddonModel.from_pretrained(...)
pipe = StableDiffusionControlNetXSPipeline.from_pretrained(..., controlnet_addon=controlnet_addon)
pipe("A dog")

2. Do you want ControlNet-XS to be part of diffusers "core" or community examples? I'm asking as they have different requirements re documentation, testing, ...

Thanks!

[patrickvonplaten]

Hey @UmerHA,

Thanks for the write-up! Yes this design makes a lot of sense to me :-)

2. Do you want ControlNet-XS to be part of diffusers "core" or community examples? I'm asking as they have different requirements re documentation, testing, ...

I think ControlNet-XS should go into "core"