Flux: pass joint_attention_kwargs when gradient_checkpointing

piercus · web-flow · commit 1b61c34e6bb1 · 2025-06-26T18:51:33.000+02:00
diff --git a/src/diffusers/models/transformers/transformer_flux.py b/src/diffusers/models/transformers/transformer_flux.py
@@ -485,6 +485,7 @@ def forward(
                     encoder_hidden_states,
                     temb,
                     image_rotary_emb,
+                    joint_attention_kwargs,
                 )
 
             else:
@@ -516,6 +517,7 @@ def forward(
                     hidden_states,
                     temb,
                     image_rotary_emb,
+                    joint_attention_kwargs,
                 )
 
             else:
