torch.compile fullgraph compatibility for Hunyuan Video

[a-r-r-o-w]

Fixes#11431 (review)

Will run the full model for testing in some time.

testing code

importtorchfromdiffusersimportHunyuanVideoTransformer3DModel@torch.no_grad()defmain():
device="cuda"dtype=torch.bfloat16batch_size=1num_channels=4num_frames=2height=4width=4sequence_length=8transformer=HunyuanVideoTransformer3DModel(
in_channels=4,
out_channels=4,
num_attention_heads=2,
attention_head_dim=10,
num_layers=2,
num_single_layers=2,
num_refiner_layers=1,
patch_size=1,
patch_size_t=1,
guidance_embeds=True,
text_embed_dim=16,
pooled_projection_dim=8,
rope_axes_dim=(2, 4, 4),
).to(device=device, dtype=dtype)
hidden_states=torch.randn(batch_size, num_channels, num_frames, height, width, dtype=dtype, device=device)
timestep=torch.randint(0, 1000, (batch_size,), dtype=torch.long, device=device)
encoder_hidden_states=torch.randn(batch_size, sequence_length, 16, dtype=dtype, device=device)
encoder_attention_mask=torch.zeros(batch_size, sequence_length, dtype=torch.bool, device=device)
encoder_attention_mask[:, :sequence_length-2] =Truepooled_projections=torch.randn(batch_size, 8, dtype=dtype, device=device)
guidance=torch.randint(0, 1000, (batch_size,), dtype=torch.long, device=device)
transformer=torch.compile(transformer, mode="max-autotune", fullgraph=True)
output=transformer(
hidden_states=hidden_states,
timestep=timestep,
encoder_hidden_states=encoder_hidden_states,
encoder_attention_mask=encoder_attention_mask,
pooled_projections=pooled_projections,
guidance=guidance,
return_dict=False,
)[0]
print(output.shape)
if__name__=="__main__":
main()

[Copilot]

Pull Request Overview

This PR improves fullgraph compatibility for Hunyuan Video by replacing a loop‑based attention mask construction with a vectorized approach.

• Replaces manual per‐batch loop with a vectorized masked_fill operation.
• Updates the attention mask initialization from zeros to ones and adds appropriate unsqueezing for broadcasting.

Comments suppressed due to low confidence (1)

src/diffusers/models/transformers/transformer_hunyuan_video.py:1071

• [nitpick] The refactored attention mask construction is more efficient; consider adding an inline comment that explains the logic behind initializing with ones and using masked_fill for clarity.

attention_mask = torch.ones(batch_size, sequence_length, device=hidden_states.device, dtype=torch.bool)

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[sayakpaul]

Looks solid, thanks for working on this!

To confirm, I first checked out the PR branch of #11431 and merged your PR branch there and then ran RUN_SLOW=1 RUN_COMPILE=1 pytest tests/models/transformers/test_models_transformer_hunyuan_video.py -k "test_torch_compile_recompilation_and_graph_break". Everything was green.