from diffusers import ControlNetUnionModel, MultiControlNetModel
print(ControlNetUnionModel, MultiControlNetModel)

from diffusers import MultiControlNetUnionModel

# src/diffusers/__init__.py
# Add "MultiControlNetUnionModel" next to "MultiControlNetModel"
# in both the lazy _import_structure["models"] list and TYPE_CHECKING imports.

from diffusers import ControlNetUnionModel

ControlNetUnionModel(
    in_channels=4,
    conditioning_channels=3,
    down_block_types=("DownBlock2D",),
    block_out_channels=(8,),
    layers_per_block=1,
    norm_num_groups=4,
    cross_attention_dim=16,
    attention_head_dim=1,
    num_trans_channel=8,
    num_trans_head=1,
    num_proj_channel=8,
    conditioning_embedding_out_channels=(4, 8),
)

if addition_time_embed_dim is None:
    raise ValueError("`addition_time_embed_dim` must be set for `ControlNetUnionModel`.")

from types import MethodType
from diffusers import ControlNetUnionModel, StableDiffusionXLControlNetUnionPipeline

pipe = object.__new__(StableDiffusionXLControlNetUnionPipeline)
pipe._callback_tensor_inputs = []
pipe.check_image = MethodType(lambda self, image, prompt, prompt_embeds: None, pipe)
pipe.controlnet = ControlNetUnionModel(
    in_channels=4,
    conditioning_channels=3,
    down_block_types=("DownBlock2D",),
    block_out_channels=(8,),
    layers_per_block=1,
    norm_num_groups=4,
    cross_attention_dim=16,
    attention_head_dim=1,
    addition_time_embed_dim=8,
    num_trans_channel=8,
    num_trans_head=1,
    num_proj_channel=8,
    conditioning_embedding_out_channels=(4, 8),
)

control_mode = None
if not isinstance(control_mode, list):
    control_mode = [control_mode]

pipe.check_inputs(
    prompt="a prompt",
    prompt_2=None,
    image=[object()],
    control_guidance_start=[0.0],
    control_guidance_end=[1.0],
    control_mode=control_mode,
    callback_on_step_end_tensor_inputs=[],
)

if control_mode is None:
    control_mode = 0
if not isinstance(control_mode, list):
    control_mode = [control_mode]

from types import MethodType
import torch
from diffusers import MultiControlNetModel, StableDiffusionControlNetPipeline

class DummyControlNet(torch.nn.Module):
    pass

pipe = object.__new__(StableDiffusionControlNetPipeline)
pipe._callback_tensor_inputs = []
pipe.controlnet = MultiControlNetModel([DummyControlNet(), DummyControlNet()])
pipe.check_image = MethodType(lambda self, image, prompt, prompt_embeds: None, pipe)

pipe.check_inputs(
    prompt="a prompt",
    image=[object(), object()],
    callback_steps=None,
    callback_on_step_end_tensor_inputs=[],
    controlnet_conditioning_scale=[0.5],
    control_guidance_start=[0.0, 0.0],
    control_guidance_end=[1.0, 1.0],
)
print("No error, but one scale for two ControlNets should be rejected.")

if isinstance(controlnet_conditioning_scale, list):
    if any(isinstance(i, list) for i in controlnet_conditioning_scale):
        raise ValueError("Batched varying conditioning scales are not supported.")
    if len(controlnet_conditioning_scale) != len(self.controlnet.nets):
        raise ValueError("Scale list length must match the number of ControlNets.")

import torch
from diffusers.models import MultiControlNetUnionModel

class DummyUnion(torch.nn.Module):
    config = type("Config", (), {"num_control_type": 6})()

    def forward(self, *args, **kwargs):
        return [torch.ones(1, 1, 1, 1)], torch.ones(1, 1, 1, 1)

multi = MultiControlNetUnionModel([DummyUnion()])
down, mid = multi(
    sample=torch.zeros(1, 4, 8, 8),
    timestep=0,
    encoder_hidden_states=torch.zeros(1, 1, 4),
    controlnet_cond=[torch.zeros(1, 3, 8, 8)],
    control_type=[torch.zeros(1, 6)],
    control_type_idx=[[0]],
    conditioning_scale=[0.0],
    return_dict=False,
)

[torch.cat([torch.zeros_like(d), d]) for d in down]

# Do not skip zero scales. Let the child ControlNet return correctly shaped zero residuals.
# Remove:
if scale == 0.0:
    continue

import torch
from diffusers import ControlNetUnionModel

model = ControlNetUnionModel(
    in_channels=4,
    conditioning_channels=3,
    down_block_types=("DownBlock2D",),
    block_out_channels=(8,),
    layers_per_block=1,
    norm_num_groups=4,
    cross_attention_dim=16,
    attention_head_dim=1,
    addition_time_embed_dim=8,
    num_trans_channel=8,
    num_trans_head=1,
    num_proj_channel=8,
    conditioning_embedding_out_channels=(4, 8),
    controlnet_conditioning_channel_order="bgr",
)

model(
    sample=torch.randn(1, 4, 8, 8),
    timestep=0,
    encoder_hidden_states=torch.randn(1, 2, 16),
    controlnet_cond=[torch.randn(1, 3, 8, 8)],
    control_type=torch.zeros(1, 6),
    control_type_idx=[0],
    return_dict=False,
)

if channel_order == "rgb":
    pass
elif channel_order == "bgr":
    controlnet_cond = [torch.flip(cond, dims=[1]) for cond in controlnet_cond]
else:
    raise ValueError(f"unknown `controlnet_conditioning_channel_order`: {channel_order}")

from pathlib import Path
from tempfile import TemporaryDirectory
import torch
from diffusers import MultiControlNetModel
from diffusers.models import MultiControlNetUnionModel

class DummyControlNet(torch.nn.Module):
    def save_pretrained(self, save_directory, **kwargs):
        print(save_directory)

for cls in (MultiControlNetModel, MultiControlNetUnionModel):
    with TemporaryDirectory() as tmp:
        try:
            cls([DummyControlNet()]).save_pretrained(Path(tmp) / "controlnet")
        except Exception as e:
            print(cls.__name__, type(e).__name__, e)

save_directory = os.fspath(save_directory)
...
model_path_to_load = os.fspath(pretrained_model_path)
...
if not return_dict:
    return down_block_res_samples, mid_block_res_sample
return ControlNetOutput(
    down_block_res_samples=down_block_res_samples,
    mid_block_res_sample=mid_block_res_sample,
)

from pathlib import Path

roots = [Path("tests/pipelines/controlnet"), Path("tests/models/controlnets"), Path("tests/single_file")]
files = [p for root in roots for p in root.rglob("*.py")]

for label, terms in {
    "Union": ["ControlNetUnionModel", "MultiControlNetUnionModel", "StableDiffusionXLControlNetUnion"],
    "BLIP": ["BlipDiffusionControlNetPipeline"],
    "Flax": ["FlaxControlNetModel", "FlaxStableDiffusionControlNetPipeline"],
}.items():
    hits = [str(p) for p in files if any(term in p.read_text(encoding="utf-8") for term in terms)]
    print(label, hits or "NO TEST REFERENCES")

for p in sorted(Path("tests/pipelines/controlnet").glob("test_controlnet*.py")):
    print(p, "@slow" in p.read_text(encoding="utf-8"))