Add support of Xlabs Controlnets #9638

Anghellia · 2024-10-10T17:45:54Z

What does this PR do?

Hi!
This PR brings support of Xlabs Controlnets, so it can be used with Diffusers.
We converted checkpoints to the Diffusers format and it can be downloaded here:

The request: #9378

Who can review?

Anyone in the community is free to review the PR once the tests have passed.
@sayakpaul

How to use

Here is the example of code to launch Canny Controlnet.

import torch
from diffusers.utils import load_image
from diffusers import FluxControlNetModel
from diffusers.pipelines import FluxControlNetPipeline
from PIL import Image
import numpy as np

generator = torch.Generator(device="cuda").manual_seed(87544357)

controlnet = FluxControlNetModel.from_pretrained(
  "Xlabs-AI/flux-controlnet-canny-diffusers",
  torch_dtype=torch.bfloat16,
  use_safetensors=True,
)
pipe = FluxControlNetPipeline.from_pretrained(
  "black-forest-labs/FLUX.1-dev",
  controlnet=controlnet,
  torch_dtype=torch.bfloat16
)
pipe.to("cuda")

control_image = load_image("https://huggingface.co/Xlabs-AI/flux-controlnet-canny-diffusers/resolve/main/canny_example.png")
prompt = "handsome girl with rainbow hair, anime"

image = pipe(
    prompt,
    control_image=control_image,
    controlnet_conditioning_scale=0.7,
    num_inference_steps=25,
    guidance_scale=3.5,
    height=1024,
    width=768,
    generator=generator,
    num_images_per_prompt=1,
).images[0]

image.save("output_test_controlnet.png")

Examples

"photo of village in the winter"

"it programmer sitting in the office"

"couple of man and woman in the water, dancing"

"photo of woman in the beach"

"futuristic bulding in the spain"

"2d art, girl in the magic city, sparkles, fantasy"

a-r-r-o-w

Just some minor comments that you can address/ignore based on what Sayak has to say

src/diffusers/models/controlnet_flux.py

a-r-r-o-w · 2024-10-10T21:56:17Z

src/diffusers/pipelines/flux/pipeline_flux_controlnet.py

@@ -773,6 +773,17 @@ def __call__(
                control_mode = torch.tensor(control_mode).to(device, dtype=torch.long)
                control_mode = control_mode.view(-1, 1).expand(control_image.shape[0], 1)

+        elif isinstance(self.controlnet, FluxControlNetModel) and self.controlnet.is_xlabs_controlnet:


I think these changes will have to be propagated to other pipeline files as well? Maybe you could rewrite this as:

if isinstance(self.controlnet): control_image = self.prepare_image(...) if self.controlnet.is_xlabs_controlnet: # remaining mismatching logic

Please check it now.

sayakpaul · 2024-10-11T08:43:06Z

@Anghellia thanks so much for this <3

Could you supplement this PR with an example code snippet and some resultant images? Ccing @asomoza for doing a test drive, too.

sayakpaul

Beautiful PR.

src/diffusers/models/transformers/transformer_flux.py

sayakpaul · 2024-10-11T08:55:44Z

src/diffusers/models/controlnet_flux.py

@@ -55,6 +56,7 @@ def __init__(
        guidance_embeds: bool = False,
        axes_dims_rope: List[int] = [16, 56, 56],
        num_mode: int = None,
+        is_xlabs_controlnet: bool = False,


Is it possible to determine if a ControlNet is of type xlabs? If not, then it's fine!

Hmm, I am not sure. Actually, xlabs ControlNets are not so different from others. I see two main changes:

We use num_layers=2 (the depth of ControlNet is 2). However, I don't think it's correct to rely solely on this, as one could also train a ControlNet with num_layers=2 using the diffusers script.

We use input_hint_block, but we can specify this only if we check for a specific keyword in the model's state_dict. I think this may not apply in our case.

Maybe both could be combined to create a condition to determine if it's an Xlabs ControlNet? @yiyixuxu would love to know your thoughts here.

Nevermind I think #9638 (comment) should cut the deal for us unless there's some differences in the forward method.

Anghellia · 2024-10-11T13:53:23Z

@Anghellia thanks so much for this <3

Could you supplement this PR with an example code snippet and some resultant images? Ccing @asomoza for doing a test drive, too.

Thank you! Updated the PR with examples 🤗

yiyixuxu

thanks a lot for the PR! I left some suggestions, let me know if they would work!

yiyixuxu · 2024-10-11T16:54:38Z

src/diffusers/models/controlnet_flux.py

@@ -55,6 +55,7 @@ def __init__(
        guidance_embeds: bool = False,
        axes_dims_rope: List[int] = [16, 56, 56],
        num_mode: int = None,
+        is_xlabs_controlnet: bool = False,


Suggested change

is_xlabs_controlnet: bool = False,

conditioning_embedding_channels: int = None,

we can add a new config conditioning_embedding_channels to the flux controlnet that defaults to None

Thanks for your suggestion, it works!

src/diffusers/models/controlnet_flux.py

src/diffusers/pipelines/flux/pipeline_flux_controlnet.py

RimoChan · 2024-10-14T03:13:19Z

I encountered an issue where the FluxMultiControlNetModel is not compatible with this branch. I used the following command to install diffusers:

pip install git+https://github.com/XLabs-AI/diffusers.git@xlabs_controlnet_support

After installation, I tried executing the following code:

import torch
from diffusers.utils import load_image
from diffusers import FluxControlNetModel, FluxMultiControlNetModel
from diffusers.pipelines import FluxControlNetPipeline



generator = torch.Generator(device="cuda").manual_seed(87544357)

controlnet = FluxMultiControlNetModel([
    FluxControlNetModel.from_pretrained(
        "Xlabs-AI/flux-controlnet-canny-diffusers",
        torch_dtype=torch.bfloat16,
        use_safetensors=True,
    ),
    FluxControlNetModel.from_pretrained(
        "Xlabs-AI/flux-controlnet-canny-diffusers",
        torch_dtype=torch.bfloat16,
        use_safetensors=True,
    ),
])
pipe = FluxControlNetPipeline.from_pretrained(
  '/mypath/to/FLUX.1-dev_official',
  controlnet=controlnet,
  torch_dtype=torch.bfloat16
)
pipe.to("cuda")

control_image = load_image("https://huggingface.co/Xlabs-AI/flux-controlnet-canny-diffusers/resolve/main/canny_example.png")

image = pipe(
    "handsome girl with rainbow hair, anime",
    control_image=[control_image, control_image],
    controlnet_conditioning_scale=[0.7, 0.7],
    num_inference_steps=25,
    guidance_scale=3.5,
    height=1024,
    width=768,
    generator=generator,
    num_images_per_prompt=1,
).images[0]

image.save("output_test_controlnet.png")

However, I encountered this error:

Traceback (most recent call last):
  File "/opt/tiger/test_1/t2s.py", line 33, in <module>
    image = pipe(
  File "/opt/tiger/miniconda3/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/opt/tiger/miniconda3/lib/python3.10/site-packages/diffusers/pipelines/flux/pipeline_flux_controlnet.py", line 897, in __call__
    controlnet_block_samples, controlnet_single_block_samples = self.controlnet(
  File "/opt/tiger/miniconda3/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/opt/tiger/miniconda3/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
  File "/opt/tiger/miniconda3/lib/python3.10/site-packages/diffusers/models/controlnet_flux.py", line 503, in forward
    block_samples, single_block_samples = controlnet(
  File "/opt/tiger/miniconda3/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/opt/tiger/miniconda3/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
  File "/opt/tiger/miniconda3/lib/python3.10/site-packages/diffusers/models/controlnet_flux.py", line 281, in forward
    controlnet_cond = self.input_hint_block(controlnet_cond)
  File "/opt/tiger/miniconda3/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/opt/tiger/miniconda3/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
  File "/opt/tiger/miniconda3/lib/python3.10/site-packages/diffusers/models/controlnet.py", line 99, in forward
    embedding = self.conv_in(conditioning)
  File "/opt/tiger/miniconda3/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/opt/tiger/miniconda3/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
  File "/opt/tiger/miniconda3/lib/python3.10/site-packages/torch/nn/modules/conv.py", line 460, in forward
    return self._conv_forward(input, self.weight, self.bias)
  File "/opt/tiger/miniconda3/lib/python3.10/site-packages/torch/nn/modules/conv.py", line 456, in _conv_forward
    return F.conv2d(input, weight, bias, self.stride,
RuntimeError: Given groups=1, weight of size [16, 3, 3, 3], expected input[1, 1, 3072, 64] to have 3 channels, but got 1 channels instead

Could you please investigate this incompatibility?

sayakpaul · 2024-10-14T03:19:50Z

MultiControlNet compatibility can be incorporated after this initial PR is merged.

… config

yiyixuxu

thanks, I left one more feedbacks
let's merge this soon!

yiyixuxu · 2024-10-14T21:04:18Z

src/diffusers/models/transformers/transformer_flux.py

@@ -508,7 +508,11 @@ def custom_forward(*inputs):
            if controlnet_block_samples is not None:
                interval_control = len(self.transformer_blocks) / len(controlnet_block_samples)
                interval_control = int(np.ceil(interval_control))
-                hidden_states = hidden_states + controlnet_block_samples[index_block // interval_control]
+                # For Xlabs ControlNet.
+                if len(controlnet_block_samples) == 2:


I would suggest passing a flag down here, controlnet_repeat_interleave = False maybe?
This would break if someone trained a controlnet with 2 blocks but want to use the other indexing method

Agree
Updated with controlnet_blocks_repeat flag

Anghellia · 2024-10-15T09:24:53Z

I encountered an issue where the FluxMultiControlNetModel is not compatible with this branch. I used the following command to install diffusers:

pip install git+https://github.com/XLabs-AI/diffusers.git@xlabs_controlnet_support

After installation, I tried executing the following code:

import torch
from diffusers.utils import load_image
from diffusers import FluxControlNetModel, FluxMultiControlNetModel
from diffusers.pipelines import FluxControlNetPipeline



generator = torch.Generator(device="cuda").manual_seed(87544357)

controlnet = FluxMultiControlNetModel([
    FluxControlNetModel.from_pretrained(
        "Xlabs-AI/flux-controlnet-canny-diffusers",
        torch_dtype=torch.bfloat16,
        use_safetensors=True,
    ),
    FluxControlNetModel.from_pretrained(
        "Xlabs-AI/flux-controlnet-canny-diffusers",
        torch_dtype=torch.bfloat16,
        use_safetensors=True,
    ),
])
pipe = FluxControlNetPipeline.from_pretrained(
  '/mypath/to/FLUX.1-dev_official',
  controlnet=controlnet,
  torch_dtype=torch.bfloat16
)
pipe.to("cuda")

control_image = load_image("https://huggingface.co/Xlabs-AI/flux-controlnet-canny-diffusers/resolve/main/canny_example.png")

image = pipe(
    "handsome girl with rainbow hair, anime",
    control_image=[control_image, control_image],
    controlnet_conditioning_scale=[0.7, 0.7],
    num_inference_steps=25,
    guidance_scale=3.5,
    height=1024,
    width=768,
    generator=generator,
    num_images_per_prompt=1,
).images[0]

image.save("output_test_controlnet.png")

However, I encountered this error:

Traceback (most recent call last):
  File "/opt/tiger/test_1/t2s.py", line 33, in <module>
    image = pipe(
  File "/opt/tiger/miniconda3/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/opt/tiger/miniconda3/lib/python3.10/site-packages/diffusers/pipelines/flux/pipeline_flux_controlnet.py", line 897, in __call__
    controlnet_block_samples, controlnet_single_block_samples = self.controlnet(
  File "/opt/tiger/miniconda3/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/opt/tiger/miniconda3/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
  File "/opt/tiger/miniconda3/lib/python3.10/site-packages/diffusers/models/controlnet_flux.py", line 503, in forward
    block_samples, single_block_samples = controlnet(
  File "/opt/tiger/miniconda3/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/opt/tiger/miniconda3/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
  File "/opt/tiger/miniconda3/lib/python3.10/site-packages/diffusers/models/controlnet_flux.py", line 281, in forward
    controlnet_cond = self.input_hint_block(controlnet_cond)
  File "/opt/tiger/miniconda3/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/opt/tiger/miniconda3/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
  File "/opt/tiger/miniconda3/lib/python3.10/site-packages/diffusers/models/controlnet.py", line 99, in forward
    embedding = self.conv_in(conditioning)
  File "/opt/tiger/miniconda3/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/opt/tiger/miniconda3/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
  File "/opt/tiger/miniconda3/lib/python3.10/site-packages/torch/nn/modules/conv.py", line 460, in forward
    return self._conv_forward(input, self.weight, self.bias)
  File "/opt/tiger/miniconda3/lib/python3.10/site-packages/torch/nn/modules/conv.py", line 456, in _conv_forward
    return F.conv2d(input, weight, bias, self.stride,
RuntimeError: Given groups=1, weight of size [16, 3, 3, 3], expected input[1, 1, 3072, 64] to have 3 channels, but got 1 channels instead

Could you please investigate this incompatibility?

Fixed

a-r-r-o-w

I think this is look good to merge now after @yiyixuxu gives a final review!

For the failing style tests, could you run make style and push? Thanks

sayakpaul · 2024-10-15T14:03:36Z

We could add some tests in a follow-up PR.

HuggingFaceDocBuilderDev · 2024-10-15T14:04:59Z

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

Anghellia · 2024-10-15T15:14:40Z

@a-r-r-o-w please launch tests

* Add support of Xlabs Controlnets --------- Co-authored-by: Anzhella Pankratova <[email protected]>

yiyixuxu · 2024-10-15T22:13:08Z

hey thanks for the PR!
I merged it in here #9687 since I cannot push into your PR
It is branched off your PR so all your commits are there and you're an author there:)

yiyixuxu · 2024-10-19T12:36:16Z

closing the PR now since we already merged it!

* Add support of Xlabs Controlnets --------- Co-authored-by: Anzhella Pankratova <[email protected]>

Add support of Xlabs Controlnets

cdca0bf

a-r-r-o-w reviewed Oct 10, 2024

View reviewed changes

a-r-r-o-w requested a review from sayakpaul October 10, 2024 21:58

sayakpaul reviewed Oct 11, 2024

View reviewed changes

use torch reshape instead of einops, fix pipeline_flux_controlnet.py

f38671e

Use ControlNetConditioningEmbedding for input_hint_block

b8c9496

yiyixuxu reviewed Oct 11, 2024

View reviewed changes

sayakpaul mentioned this pull request Oct 13, 2024

Draft PR: [Flux ControlNet] Support Xlabs ControlNet in diffusers #9385

Closed

6 tasks

Use conditioning_embedding_channels instead of is_xlabs_controlnet in…

7be937e

… config

Anghellia requested a review from yiyixuxu October 14, 2024 09:07

yiyixuxu reviewed Oct 14, 2024

View reviewed changes

Anghellia added 2 commits October 15, 2024 11:16

Add controlnet_blocks_repeat to Flux forward

b07b48d

Fix for FluxMultiControlNetModel

bff8644

a-r-r-o-w approved these changes Oct 15, 2024

View reviewed changes

Fix import order

8cf0105

yiyixuxu mentioned this pull request Oct 15, 2024

[authored by @Anghellia) Add support of Xlabs Controlnets #9638 #9687

Merged

yiyixuxu added a commit that referenced this pull request Oct 15, 2024

[authored by @Anghellia) Add support of Xlabs Controlnets #9638 (#9687)

3e9a28a

* Add support of Xlabs Controlnets --------- Co-authored-by: Anzhella Pankratova <[email protected]>

yiyixuxu closed this Oct 19, 2024

sayakpaul pushed a commit that referenced this pull request Dec 23, 2024

[authored by @Anghellia) Add support of Xlabs Controlnets #9638 (#9687)

dd83a81

* Add support of Xlabs Controlnets --------- Co-authored-by: Anzhella Pankratova <[email protected]>

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add support of Xlabs Controlnets #9638

Add support of Xlabs Controlnets #9638

Anghellia commented Oct 10, 2024 •

edited

Loading

a-r-r-o-w left a comment

a-r-r-o-w Oct 10, 2024

Anghellia Oct 11, 2024

sayakpaul commented Oct 11, 2024

sayakpaul left a comment

sayakpaul Oct 11, 2024

Anghellia Oct 11, 2024

sayakpaul Oct 12, 2024

sayakpaul Oct 12, 2024

Anghellia commented Oct 11, 2024

yiyixuxu left a comment

yiyixuxu Oct 11, 2024

yiyixuxu Oct 11, 2024

Anghellia Oct 14, 2024

RimoChan commented Oct 14, 2024 •

edited

Loading

sayakpaul commented Oct 14, 2024

yiyixuxu left a comment

yiyixuxu Oct 14, 2024

Anghellia Oct 15, 2024

Anghellia commented Oct 15, 2024

a-r-r-o-w left a comment

sayakpaul commented Oct 15, 2024

HuggingFaceDocBuilderDev commented Oct 15, 2024

Anghellia commented Oct 15, 2024

yiyixuxu commented Oct 15, 2024

yiyixuxu commented Oct 19, 2024

	is_xlabs_controlnet: bool = False,
	conditioning_embedding_channels: int = None,

Add support of Xlabs Controlnets #9638

Add support of Xlabs Controlnets #9638

Conversation

Anghellia commented Oct 10, 2024 • edited Loading

What does this PR do?

Who can review?

How to use

Examples

a-r-r-o-w left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

sayakpaul commented Oct 11, 2024

sayakpaul left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Anghellia commented Oct 11, 2024

yiyixuxu left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

RimoChan commented Oct 14, 2024 • edited Loading

sayakpaul commented Oct 14, 2024

yiyixuxu left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Anghellia commented Oct 15, 2024

a-r-r-o-w left a comment

Choose a reason for hiding this comment

sayakpaul commented Oct 15, 2024

HuggingFaceDocBuilderDev commented Oct 15, 2024

Anghellia commented Oct 15, 2024

yiyixuxu commented Oct 15, 2024

yiyixuxu commented Oct 19, 2024

Anghellia commented Oct 10, 2024 •

edited

Loading

RimoChan commented Oct 14, 2024 •

edited

Loading