Flux: pass joint_attention_kwargs when using gradient_checkpointing (huggingface#11814)

piercus · tolgacangoz · commit 94240876e631 · 2025-07-18T10:08:43.000+03:00
Flux: pass joint_attention_kwargs when gradient_checkpointing
diff --git a/src/diffusers/models/transformers/transformer_flux.py b/src/diffusers/models/transformers/transformer_flux.py
@@ -490,6 +490,7 @@ def forward(
                     encoder_hidden_states,
                     temb,
                     image_rotary_emb,
+                    joint_attention_kwargs,
                 )
 
             else:
@@ -521,6 +522,7 @@ def forward(
                     encoder_hidden_states,
                     temb,
                     image_rotary_emb,
+                    joint_attention_kwargs,
                 )
 
             else:
