hidream_image model/pipeline review

hidream_image model/pipeline review

Commit tested: 0f1abc4ae8b0eb2a3b40e82a310507281144c423

Review performed against the repository review rules.

Duplicate search performed with gh for HiDream, class/file names, and specific failure modes. Existing duplicate found only for the torch.compile item: https://github.com/huggingface/diffusers/pull/11477 and https://github.com/huggingface/diffusers/issues/11430.

Issue 1: height and width are silently replaced with default-area dimensions

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/hidream_image/pipeline_hidream_image.py#L873-L880

Problem:
The pipeline rescales any requested height/width to the default sample area. A user asking for 64x64 gets default-area latents/images instead of 64x64.

Impact:
User-visible output dimensions do not match the public API or docstring, and precomputed latents for the requested size are rejected.

Reproduction:

import torch
from diffusers import FlowMatchEulerDiscreteScheduler, HiDreamImagePipeline, HiDreamImageTransformer2DModel

m = HiDreamImageTransformer2DModel(patch_size=2, in_channels=4, out_channels=4, num_layers=0, num_single_layers=0,
    attention_head_dim=8, num_attention_heads=1, caption_channels=[32, 16], text_emb_dim=64,
    num_routed_experts=0, axes_dims_rope=(4, 2, 2), max_resolution=(32, 32), llama_layers=())
p = HiDreamImagePipeline(FlowMatchEulerDiscreteScheduler(), None, None, None, None, None, None, None, None, None, m)
out = p(height=64, width=64, num_inference_steps=0, guidance_scale=1.0, output_type="latent",
        pooled_prompt_embeds=torch.randn(1, 64), prompt_embeds_t5=torch.randn(1, 1, 32),
        prompt_embeds_llama3=torch.randn(0, 1, 1, 16)).images
print(out.shape)  # torch.Size([1, 4, 128, 128])

Relevant precedent:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/flux/pipeline_flux.py#L854-L865

Suggested fix:

height = height or self.default_sample_size * self.vae_scale_factor
width = width or self.default_sample_size * self.vae_scale_factor

division = self.vae_scale_factor * self.transformer.config.patch_size
height = int(height) // division * division
width = int(width) // division * division
image_seq_len = (height // self.vae_scale_factor // self.transformer.config.patch_size) * (
    width // self.vae_scale_factor // self.transformer.config.patch_size
)
if image_seq_len > self.transformer.max_seq:
    raise ValueError(f"Requested image has {image_seq_len} latent tokens, but this model supports {self.transformer.max_seq}.")

Issue 2: attention_kwargs is accepted but never forwarded to the transformer

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/hidream_image/pipeline_hidream_image.py#L996-L1003

Problem:
__call__ stores attention_kwargs, and the transformer forward is decorated with @apply_lora_scale("attention_kwargs"), but the denoising call never passes the kwargs.

Impact:
Runtime LoRA scaling via attention_kwargs={"scale": ...} is ignored for HiDream.

Reproduction:

import torch
from diffusers import FlowMatchEulerDiscreteScheduler, HiDreamImagePipeline, HiDreamImageTransformer2DModel

m = HiDreamImageTransformer2DModel(patch_size=2, in_channels=4, out_channels=4, num_layers=1, num_single_layers=0,
    attention_head_dim=8, num_attention_heads=1, caption_channels=[32, 16], text_emb_dim=64,
    num_routed_experts=0, axes_dims_rope=(4, 2, 2), max_resolution=(32, 32), llama_layers=(0,)).eval()
p = HiDreamImagePipeline(FlowMatchEulerDiscreteScheduler(), None, None, None, None, None, None, None, None, None, m)
seen = []
h = p.transformer.register_forward_pre_hook(lambda module, args, kwargs: seen.append(kwargs.get("attention_kwargs")), with_kwargs=True)
p(height=128, width=128, num_inference_steps=1, guidance_scale=1.0, output_type="latent",
  attention_kwargs={"scale": 0.25}, pooled_prompt_embeds=torch.randn(1, 64),
  prompt_embeds_t5=torch.randn(1, 4, 32), prompt_embeds_llama3=torch.randn(1, 1, 4, 16))
h.remove()
print(seen)  # [None]

Relevant precedent:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/qwenimage/pipeline_qwenimage.py#L695-L703

Suggested fix:

noise_pred = self.transformer(
    hidden_states=latent_model_input,
    timesteps=timestep,
    encoder_hidden_states_t5=prompt_embeds_t5,
    encoder_hidden_states_llama3=prompt_embeds_llama3,
    pooled_embeds=pooled_prompt_embeds,
    attention_kwargs=self.attention_kwargs,
    return_dict=False,
)[0]

Issue 3: Padded tokens are not actually masked in attention

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_hidream_image.py#L227-L260

Problem:
For non-square/padded latent sequences, hidden_states_masks only multiplies image keys by 0. It does not mask attention logits and does not mask values, so masked tokens still change valid outputs.

Impact:
Non-square generation can be contaminated by padded latent tokens, and externally supplied padded hidden states are not semantically masked.

Reproduction:

import torch
from diffusers import HiDreamImageTransformer2DModel

torch.manual_seed(0)
m = HiDreamImageTransformer2DModel(patch_size=2, in_channels=1, out_channels=1, num_layers=1, num_single_layers=0,
    attention_head_dim=8, num_attention_heads=1, caption_channels=[3, 2], text_emb_dim=5,
    num_routed_experts=0, axes_dims_rope=(4, 2, 2), max_resolution=(4, 4), llama_layers=(0,)).eval()
hidden = torch.randn(1, 4, 4)
kwargs = dict(timesteps=torch.tensor([1]), encoder_hidden_states_t5=torch.randn(1, 1, 3),
    encoder_hidden_states_llama3=torch.randn(1, 1, 1, 2), pooled_embeds=torch.randn(1, 5),
    img_ids=torch.zeros(1, 4, 3), img_sizes=torch.tensor([[1, 2]]),
    hidden_states_masks=torch.tensor([[1., 1., 0., 0.]]), return_dict=False)
with torch.no_grad():
    out1 = m(hidden_states=hidden.clone(), **kwargs)[0]
    hidden[:, 2:] += 1000
    out2 = m(hidden_states=hidden, **kwargs)[0]
print((out1 - out2).abs().max().item())  # non-zero

Relevant precedent:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_qwenimage.py#L562-L570

Suggested fix:
Route HiDream attention through dispatch_attention_fn with a real boolean attention mask. For double-stream blocks, concatenate the image mask with all-true text masks before dispatch; for single-stream blocks, pass the already-expanded image/text mask.

Issue 4: HiDream attention bypasses diffusers attention backend dispatch

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_hidream_image.py#L201-L260

Problem:
HiDreamAttnProcessor calls F.scaled_dot_product_attention directly and has no _attention_backend / _parallel_config. model.set_attention_backend(...) therefore has no effect.

Impact:
HiDream cannot use diffusers attention backends such as Flash/Sage/Flex through the standard API, and context-parallel backend plumbing is bypassed.

Reproduction:

from diffusers import HiDreamImageTransformer2DModel

m = HiDreamImageTransformer2DModel(patch_size=2, in_channels=4, out_channels=4, num_layers=1, num_single_layers=0,
    attention_head_dim=8, num_attention_heads=1, caption_channels=[32, 16], text_emb_dim=64,
    num_routed_experts=0, axes_dims_rope=(4, 2, 2), max_resolution=(32, 32), llama_layers=(0,))
processors = [mod.processor for mod in m.modules() if mod.__class__.__name__ == "HiDreamAttention"]
print(hasattr(processors[0], "_attention_backend"))  # False
m.set_attention_backend("native")
print(getattr(processors[0], "_attention_backend", None))  # None

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

Suggested fix:
Make HiDreamAttention follow the new custom-attention pattern: inherit AttentionModuleMixin, declare _default_processor_cls / _available_processors, add _attention_backend and _parallel_config to the processor, and replace the direct SDPA call with dispatch_attention_fn(...).

Issue 5: Transformer config defaults are invalid

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_hidream_image.py#L612-L625

Problem:
Several registered config defaults are unusable: patch_size=None, caption_channels=None, llama_layers=None, and axes_dims_rope=(32, 32) for three-axis ids.

Impact:
Default construction and configs missing these fields fail with low-level errors instead of a clear config error, which weakens serialization/backwards-compatibility behavior.

Reproduction:

from diffusers import HiDreamImageTransformer2DModel

try:
    HiDreamImageTransformer2DModel()
except Exception as e:
    print(type(e).__name__, e)
# TypeError unsupported operand type(s) for *: 'int' and 'NoneType'

Relevant precedent:
Flux and Qwen transformer constructors use internally consistent config defaults.

Suggested fix:

if patch_size is None:
    raise ValueError("`patch_size` must be set for HiDreamImageTransformer2DModel.")
if caption_channels is None or len(caption_channels) != 2:
    raise ValueError("`caption_channels` must contain [t5_dim, llama_dim].")
if llama_layers is None:
    raise ValueError("`llama_layers` must be set.")
if len(axes_dims_rope) != 3:
    raise ValueError("`axes_dims_rope` must contain three axes for HiDream image ids.")

Issue 6: MoE inference path is not torch.compile-friendly

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_hidream_image.py#L388-L392

Problem:
moe_infer performs bincount().cpu().numpy().cumsum(0) inside model forward. This creates host synchronization and graph breaks.

Impact:
torch.compile(fullgraph=True) fails for the default MoE inference path.

Reproduction:

import torch
from diffusers import HiDreamImageTransformer2DModel

m = HiDreamImageTransformer2DModel(patch_size=2, in_channels=1, out_channels=1, num_layers=1, num_single_layers=0,
    attention_head_dim=8, num_attention_heads=1, caption_channels=[3, 2], text_emb_dim=5,
    num_routed_experts=2, num_activated_experts=1, axes_dims_rope=(4, 2, 2), max_resolution=(4, 4),
    llama_layers=(0,)).eval()
compiled = torch.compile(m, fullgraph=True)
try:
    compiled(hidden_states=torch.randn(1, 1, 4, 4), timesteps=torch.tensor([1]),
        encoder_hidden_states_t5=torch.randn(1, 1, 3), encoder_hidden_states_llama3=torch.randn(1, 1, 1, 2),
        pooled_embeds=torch.randn(1, 5), return_dict=False)
except Exception as e:
    print(type(e).__name__, str(e).splitlines()[0])

Relevant precedent:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_nucleusmoe_image.py#L487-L498

Suggested fix:
Existing duplicate: https://github.com/huggingface/diffusers/pull/11477. Continue that PR or replace this routing with a torch-only grouped implementation that avoids NumPy and host-side dynamic routing in the compiled path.

Issue 7: No slow HiDream pipeline/model coverage exists

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/pipelines/hidream_image/test_pipeline_hidream.py#L45-L135

Problem:
HiDream has fast pipeline and model tests, plus GGUF coverage, but no @slow test for a published HiDream pipeline/model checkpoint.

Impact:
Large-checkpoint loading, real component configs, offload behavior, real output slices, and docs example compatibility are not covered.

Reproduction:

from pathlib import Path

paths = [Path("tests/pipelines/hidream_image"), Path("tests/models/transformers/test_models_transformer_hidream.py")]
hits = []
for path in paths:
    files = [path] if path.is_file() else list(path.rglob("*.py"))
    hits += [(str(f), i) for f in files for i, line in enumerate(f.read_text().splitlines(), 1) if "@slow" in line or "slow(" in line]
print(hits)  # []

Relevant precedent:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/pipelines/stable_diffusion_3/test_pipeline_stable_diffusion_3.py#L227-L245

Suggested fix:
Add a @slow HiDream test class that loads a real or hf-internal-testing published HiDream pipeline checkpoint through from_pretrained, runs a deterministic minimal inference, and asserts shape plus an output slice.

Local verification note: targeted snippets were run with .venv. Direct pytest collection for the HiDream model and pipeline test files failed in this environment because the installed torch build lacks torch._C._distributed_c10d, imported by shared test mixins.

Commencez par les points d’entrée concernés dans src/diffusers/pipelines/hidream_image/pipeline_hidream_image.py et src/diffusers/models/transformers/transformer_hidream_image.py, puis examinez tests/pipelines/hidream_image/test_pipeline_hidream.py. Reproduisez les cas répertoriés concernant la taille, l’attention, le masquage, le backend, la configuration et torch.compile avant de choisir un périmètre ciblé. Le travail est considéré comme terminé lorsque le comportement sélectionné est corrigé avec une couverture de régression, y compris la couverture manquante de la pipeline/du modèle slow le cas échéant.