from diffusers import HunyuanDiT2DControlNetModel, HunyuanDiT2DModel

transformer = HunyuanDiT2DModel(
    sample_size=8, num_layers=4, patch_size=2,
    attention_head_dim=4, num_attention_heads=2, in_channels=4,
    cross_attention_dim=8, cross_attention_dim_t5=8,
    pooled_projection_dim=4, hidden_size=8, text_len=4, text_len_t5=4,
    use_style_cond_and_image_meta_size=False,
)

try:
    HunyuanDiT2DControlNetModel.from_transformer(transformer)
except Exception as e:
    print(type(e).__name__, str(e).splitlines()[0])

try:
    HunyuanDiT2DControlNetModel.from_transformer(transformer, transformer_num_layers=4)
except Exception as e:
    print(type(e).__name__, str(e).splitlines()[0])

config = transformer.config
transformer_num_layers = transformer_num_layers or config.num_layers

controlnet = cls(
    conditioning_channels=conditioning_channels,
    transformer_num_layers=transformer_num_layers,
    activation_fn=config.activation_fn,
    attention_head_dim=config.attention_head_dim,
    cross_attention_dim=config.cross_attention_dim,
    cross_attention_dim_t5=config.cross_attention_dim_t5,
    hidden_size=config.hidden_size,
    in_channels=config.in_channels,
    mlp_ratio=config.mlp_ratio,
    num_attention_heads=config.num_attention_heads,
    patch_size=config.patch_size,
    sample_size=config.sample_size,
    pooled_projection_dim=config.pooled_projection_dim,
    text_len=config.text_len,
    text_len_t5=config.text_len_t5,
    use_style_cond_and_image_meta_size=config.use_style_cond_and_image_meta_size,
)

import torch
from diffusers import AutoencoderKL, DDPMScheduler, HunyuanDiT2DModel, HunyuanDiTPipeline

transformer = HunyuanDiT2DModel(
    sample_size=16, num_layers=2, patch_size=2,
    attention_head_dim=8, num_attention_heads=3, in_channels=4,
    cross_attention_dim=32, cross_attention_dim_t5=32,
    pooled_projection_dim=16, hidden_size=24,
).eval()

pipe = HunyuanDiTPipeline(
    vae=AutoencoderKL(), text_encoder=None, tokenizer=None,
    transformer=transformer, scheduler=DDPMScheduler(),
    safety_checker=None, feature_extractor=None,
    text_encoder_2=None, tokenizer_2=None, requires_safety_checker=False,
)
pipe.set_progress_bar_config(disable=True)

pipe(
    prompt_embeds=torch.randn(2, 77, 32),
    prompt_attention_mask=torch.ones(2, 77, dtype=torch.long),
    prompt_embeds_2=torch.randn(2, 256, 32),
    prompt_attention_mask_2=torch.ones(2, 256, dtype=torch.long),
    num_images_per_prompt=2, guidance_scale=1.0,
    num_inference_steps=1, height=16, width=16,
    output_type="latent", use_resolution_binning=False,
)

# Keep generated masks unrepeated in the generation branch, then normalize both paths once.
if prompt_embeds is None:
    ...
    prompt_attention_mask = text_inputs.attention_mask.to(device)
    prompt_embeds = text_encoder(...)[0]

prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
prompt_attention_mask = prompt_attention_mask.to(device).repeat(num_images_per_prompt, 1)

...
if do_classifier_free_guidance and negative_prompt_embeds is None:
    ...
    negative_prompt_attention_mask = uncond_input.attention_mask.to(device)
    negative_prompt_embeds = text_encoder(...)[0]

if do_classifier_free_guidance:
    negative_prompt_embeds = negative_prompt_embeds.to(dtype=dtype, device=device)
    negative_prompt_attention_mask = negative_prompt_attention_mask.to(device).repeat(num_images_per_prompt, 1)
    ...

import torch
from diffusers import HunyuanDiT2DMultiControlNetModel

multi = HunyuanDiT2DMultiControlNetModel([torch.nn.Identity(), torch.nn.Identity()])
try:
    multi(torch.zeros(1), torch.zeros(1), [torch.zeros(1), torch.zeros(1)], conditioning_scale=1.0)
except Exception as e:
    print(type(e).__name__, str(e))

class DummyControlNet(torch.nn.Module):
    def __init__(self, value):
        super().__init__()
        self.value = value
        self.calls = 0
    def forward(self, *args, **kwargs):
        self.calls += 1
        return ([torch.tensor(float(self.value))],)

first, second = DummyControlNet(1), DummyControlNet(2)
multi = HunyuanDiT2DMultiControlNetModel([first, second])
out = multi(torch.zeros(1), torch.zeros(1), [torch.zeros(1)], conditioning_scale=[1.0, 1.0], return_dict=False)
print(out[0][0].item(), first.calls, second.calls)

if isinstance(self.controlnet, HunyuanDiT2DMultiControlNetModel):
    count = len(self.controlnet.nets)

    if not isinstance(controlnet_conditioning_scale, list):
        controlnet_conditioning_scale = [controlnet_conditioning_scale] * count

    if not isinstance(control_image, list) or len(control_image) != count:
        raise ValueError(f"`control_image` must be a list of {count} images for multiple ControlNets.")

    if len(controlnet_conditioning_scale) != count:
        raise ValueError(f"`controlnet_conditioning_scale` must contain {count} values.")

from diffusers import HunyuanDiT2DControlNetModel, HunyuanDiT2DModel

common = dict(
    sample_size=8, patch_size=2, attention_head_dim=4, num_attention_heads=2,
    in_channels=4, cross_attention_dim=8, cross_attention_dim_t5=8,
    pooled_projection_dim=4, hidden_size=8, text_len=4, text_len_t5=4,
)

for model in [
    HunyuanDiT2DModel(num_layers=1, **common),
    HunyuanDiT2DControlNetModel(transformer_num_layers=4, **common),
]:
    try:
        print(model.__class__.__name__, model._get_no_split_modules("auto"))
    except Exception as e:
        print(model.__class__.__name__, type(e).__name__, str(e).splitlines()[0])

class HunyuanDiT2DModel(...):
    _no_split_modules = ["HunyuanDiTBlock", "PatchEmbed"]
    _skip_layerwise_casting_patterns = ["pos_embed", "norm", "pooler"]

class HunyuanDiT2DControlNetModel(...):
    _no_split_modules = ["HunyuanDiTBlock", "PatchEmbed"]
    _skip_layerwise_casting_patterns = ["pos_embed", "norm", "pooler"]

from diffusers import HunyuanDiT2DModel

model = HunyuanDiT2DModel(
    sample_size=8, num_layers=1, patch_size=2,
    attention_head_dim=4, num_attention_heads=2, in_channels=4,
    cross_attention_dim=8, cross_attention_dim_t5=8,
    pooled_projection_dim=4, hidden_size=8, text_len=4, text_len_t5=4,
)

print([hasattr(p, "_attention_backend") for p in model.attn_processors.values()])
model.set_attention_backend("native")
print([getattr(p, "_attention_backend", None) for p in model.attn_processors.values()])

from pathlib import Path

text = Path("tests/pipelines/controlnet_hunyuandit/test_controlnet_hunyuandit.py").read_text()
for name in [
    "test_sequential_cpu_offload_forward_pass",
    "test_sequential_offload_forward_pass_twice",
    "test_save_load_optional_components",
]:
    start = text.index(f"def {name}")
    end = text.find("\n    def ", start + 1)
    block = text[start:end if end != -1 else len(text)]
    print(name, "TODO" in block and "pass" in block)