huggingface / huggingface/diffusers

hunyuandit model/pipeline review

Open
#13,641 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
34.5k
Forks
7.3k
Avg merge
3d 3h
Merged PRs (30d)
91

Description

# `hunyuandit` model/pipeline review

Commit tested: `0f1abc4ae8b0eb2a3b40e82a310507281144c423`

Review performed against the repository review rules.

Reviewed: target model/pipeline files, lazy exports, config/loading paths, dtype/device/offload behavior, attention processors, docs, examples, and tests. Fast and slow tests exist for both base and ControlNet pipelines, but ControlNet has several no-op fast tests noted below. Local `.venv` snippets were run; full pytest collection was blocked because this `.venv` torch build is missing `torch._C._distributed_c10d`.

Duplicate search: searched GitHub Issues and PRs for `hunyuandit`, target class/file names, `from_transformer transformer_num_layers`, `prompt_embeds num_images_per_prompt attention_mask`, `MultiControlNet controlnet_conditioning_scale`, `_no_split_modules`, and `HunyuanAttnProcessor2_0 set_attention_backend`. No exact duplicates found. Related closed issue: https://github.com/huggingface/diffusers/issues/9142 covered older multi-ControlNet loading, not the runtime scale/validation failures below.

## Issue 1: `from_transformer()` cannot build a matching HunyuanDiT ControlNet

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/controlnets/controlnet_hunyuan.py#L177-L209

Problem:
`HunyuanDiT2DControlNetModel.from_transformer()` reads `config.transformer_num_layers`, but `HunyuanDiT2DModel` stores `num_layers`. Even when `transformer_num_layers` is passed manually, the method does not copy `pooled_projection_dim` or `use_style_cond_and_image_meta_size`, so weight loading can hit size mismatches.

Impact:
The public constructor for deriving a ControlNet from a transformer is broken for normal HunyuanDiT transformers and can also create a ControlNet with incompatible conditioning embeddings.

Reproduction:
```python
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])
```

Relevant precedent:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/controlnets/controlnet_flux.py#L135-L142
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/controlnets/controlnet_sd3.py#L257-L260

Suggested fix:
```python
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,
)
```

## Issue 2: Precomputed prompt masks are not repeated for `num_images_per_prompt`

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/hunyuandit/pipeline_hunyuandit.py#L349-L362
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/hunyuandit/pipeline_hunyuandit.py#L394-L409
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/controlnet_hunyuandit/pipeline_hunyuandit_controlnet.py#L377-L390
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/controlnet_hunyuandit/pipeline_hunyuandit_controlnet.py#L422-L437

Problem:
When users pass `prompt_embeds` and attention masks directly, `encode_prompt()` repeats embeddings for `num_images_per_prompt` but leaves the provided masks at the original batch size. Generated masks are repeated, but caller-provided masks are not.

Impact:
Batched precomputed embeddings fail when `num_images_per_prompt > 1`; with batch size 1 the mask broadcasts, hiding the bug.

Reproduction:
```python
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,
)
```

Relevant precedent:
The same method is copied into the ControlNet pipeline, so the fix should be applied at the source and propagated.

Suggested fix:
```python
# 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)
...
```

## Issue 3: Multi-ControlNet inputs are neither normalized nor length-checked

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/controlnet_hunyuandit/pipeline_hunyuandit_controlnet.py#L642
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/controlnet_hunyuandit/pipeline_hunyuandit_controlnet.py#L867-L888
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/controlnets/controlnet_hunyuan.py#L374-L400

Problem:
For `HunyuanDiT2DMultiControlNetModel`, `conditioning_scale` is iterated directly. The default scalar `1.0` raises `TypeError`, and mismatched lengths between `control_image`, `conditioning_scale`, and `self.nets` are silently truncated by `zip()`.

Impact:
Multiple ControlNets are fragile: the documented scalar default fails, and missing images/scales can silently skip ControlNets.

Reproduction:
```python
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)
```

Relevant precedent:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/controlnet_sd3/pipeline_stable_diffusion_3_controlnet.py#L981-L990
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/controlnet_sd3/pipeline_stable_diffusion_3_controlnet.py#L1172-L1178

Suggested fix:
```python
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.")
```

## Issue 4: HunyuanDiT models do not support `device_map="auto"`

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/hunyuan_transformer_2d.py#L201-L246
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/controlnets/controlnet_hunyuan.py#L40-L42

Problem:
Neither `HunyuanDiT2DModel` nor `HunyuanDiT2DControlNetModel` defines `_no_split_modules`. Diffusers rejects `device_map="auto"` for model classes without this attribute.

Impact:
Large HunyuanDiT checkpoints cannot use automatic model-level device placement, unlike related transformer/controlnet families.

Reproduction:
```python
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])
```

Relevant precedent:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_flux.py#L566-L568
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/controlnets/controlnet_sana.py#L43-L44

Suggested fix:
```python
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"]
```

## Issue 5: Attention backend selection is a no-op

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/hunyuan_transformer_2d.py#L20-L21
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/hunyuan_transformer_2d.py#L111-L120
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/attention_processor.py#L3124-L3201

Problem:
HunyuanDiT uses legacy `HunyuanAttnProcessor2_0` from `attention_processor.py`. That processor calls `F.scaled_dot_product_attention` directly and has no `_attention_backend`, so `model.set_attention_backend("native")` silently leaves all processors unchanged.

Impact:
Users cannot select modern attention backends for HunyuanDiT despite the model exposing the generic attention backend API.

Reproduction:
```python
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()])
```

Relevant precedent:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_flux.py#L75-L124

Suggested fix:
Move Hunyuan attention processors to the model-file pattern used by newer transformers, add `_attention_backend = None` and `_parallel_config = None`, and route attention through `dispatch_attention_fn(...)` instead of calling `F.scaled_dot_product_attention` directly.

## Issue 6: ControlNet fast tests contain passing TODO stubs

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/pipelines/controlnet_hunyuandit/test_controlnet_hunyuandit.py#L177-L187

Problem:
`test_sequential_cpu_offload_forward_pass`, `test_sequential_offload_forward_pass_twice`, and `test_save_load_optional_components` contain only TODO comments and `pass`, so pytest reports them as successful without checking anything.

Impact:
The target’s ControlNet offload and serialization coverage is weaker than it appears, and regressions in optional component save/load behavior are not exercised.

Reproduction:
```python
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)
```

Relevant precedent:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/pipelines/hunyuandit/test_hunyuan_dit.py#L129-L145
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/pipelines/hunyuandit/test_hunyuan_dit.py#L222

Suggested fix:
Implement the save/load optional component test for ControlNet, and either implement offload tests or explicitly skip them with the same unsupported-offload reason used by the base HunyuanDiT pipeline.

Contributor guide

Open the contributing guide

Research direction

Start with the affected HunyuanDiT model and pipeline files named in the review, then run the existing fast and slow tests for the base and ControlNet pipelines. The review lists separate failures in transformer-to-ControlNet construction, prompt-mask batching, multi-ControlNet validation, device mapping, and attention backend selection; completion would require resolving the agreed scope and covering each behavior with passing tests.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning, testing-qa
Issue type
Bug
Difficulty
5/5
Estimated time
Over a week
Activity status
Quiet
Clarity
Needs clarification
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.