ControlnetUnion not working correctly

[asomoza]

Describe the bug

After #10723 the way controlnet union works with the control images is not correct.

The more clear and concise way of showing this is using a depth map with the tile condition, the tile condition should return almost the same image but in this case it return an image as if it was using a depth map.

I noticed that with this issue, the inpainting also suffers a lot and mostly generates subpar inpainted images compared to the original implementation or how it was before that PR.

Reproduction

Using this code before and after the relevant PR.

import torch

from diffusers import ControlNetUnionModel, StableDiffusionXLControlNetUnionPipeline
from diffusers.utils import load_image


torch_dtype = torch.float16
prompt = "A cat"

control_image = load_image(
    "https://huggingface.co/datasets/OzzyGT/testing-resources/resolve/main/resources/cat_depthmap.png"
)

controlnet_model = ControlNetUnionModel.from_pretrained(
    "OzzyGT/controlnet-union-promax-sdxl-1.0", variant="fp16", torch_dtype=torch_dtype
)

pipe = StableDiffusionXLControlNetUnionPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    variant="fp16",
    controlnet=controlnet_model,
    torch_dtype=torch_dtype,
)

pipe.to("cuda")

image = pipe(
    prompt=prompt,
    guidance_scale=5.0,
    num_inference_steps=20,
    control_image=control_image,
    controlnet_conditioning_scale=0.5,
    control_mode=6,
).images[0]

image.save("controlnet_union_test.png")

Before	After

Who can help?

I'm opening this issue to the community since I don't have time right now, if no one wants to tackle it, I'll do it when I have time.

[hlky]

diff --git a/src/diffusers/models/controlnets/controlnet_union.py b/src/diffusers/models/controlnets/controlnet_union.py
index d731ef137..19dfdc414 100644
--- a/src/diffusers/models/controlnets/controlnet_union.py
+++ b/src/diffusers/models/controlnets/controlnet_union.py
@@ -752,7 +752,7 @@ class ControlNetUnionModel(ModelMixin, ConfigMixin, FromOriginalModelMixin):
             condition = self.controlnet_cond_embedding(cond)
             feat_seq = torch.mean(condition, dim=(2, 3))
             feat_seq = feat_seq + self.task_embedding[control_idx]
-            if from_multi:
+            if from_multi or len(control_type_idx) == 1:
                 inputs.append(feat_seq.unsqueeze(1))
                 condition_list.append(condition)
             else:
@@ -772,7 +772,7 @@ class ControlNetUnionModel(ModelMixin, ConfigMixin, FromOriginalModelMixin):
         for (idx, condition), scale in zip(enumerate(condition_list[:-1]), conditioning_scale):
             alpha = self.spatial_ch_projs(x[:, idx])
             alpha = alpha.unsqueeze(-1).unsqueeze(-1)
-            if from_multi:
+            if from_multi or len(control_type_idx) == 1:
                 controlnet_cond_fuser += condition + alpha
             else:
                 controlnet_cond_fuser += condition + alpha * scale
@@ -823,7 +823,7 @@ class ControlNetUnionModel(ModelMixin, ConfigMixin, FromOriginalModelMixin):
                 scales = scales * conditioning_scale[0]
             down_block_res_samples = [sample * scale for sample, scale in zip(down_block_res_samples, scales)]
             mid_block_res_sample = mid_block_res_sample * scales[-1]  # last one
-        elif from_multi:
+        elif from_multi or len(control_type_idx) == 1:
             down_block_res_samples = [sample * conditioning_scale[0] for sample in down_block_res_samples]
             mid_block_res_sample = mid_block_res_sample * conditioning_scale[0]
 

[asomoza]

thanks a lot @hlky I'll do a PR with this and the other one, is ok if I add you as an author?