Support pass kwargs to cogvideox custom attention processor (#10456)

* Support pass kwargs to cogvideox custom attention processor * remove args in cogvideox attn processor * remove unused kwargs
huggingface · Jan 9, 2025 · 7116fd2 · 7116fd2
1 parent 553b138
commit 7116fd2
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/src/diffusers/models/transformers/cogvideox_transformer_3d.py b/src/diffusers/models/transformers/cogvideox_transformer_3d.py
@@ -120,8 +120,10 @@ def forward(
         encoder_hidden_states: torch.Tensor,
         temb: torch.Tensor,
         image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        attention_kwargs: Optional[Dict[str, Any]] = None,
     ) -> torch.Tensor:
         text_seq_length = encoder_hidden_states.size(1)
+        attention_kwargs = attention_kwargs or {}
 
         # norm & modulate
         norm_hidden_states, norm_encoder_hidden_states, gate_msa, enc_gate_msa = self.norm1(
@@ -133,6 +135,7 @@ def forward(
             hidden_states=norm_hidden_states,
             encoder_hidden_states=norm_encoder_hidden_states,
             image_rotary_emb=image_rotary_emb,
+            **attention_kwargs,
         )
 
         hidden_states = hidden_states + gate_msa * attn_hidden_states
@@ -498,6 +501,7 @@ def custom_forward(*inputs):
                     encoder_hidden_states,
                     emb,
                     image_rotary_emb,
+                    attention_kwargs,
                     **ckpt_kwargs,
                 )
             else:
@@ -506,6 +510,7 @@ def custom_forward(*inputs):
                     encoder_hidden_states=encoder_hidden_states,
                     temb=emb,
                     image_rotary_emb=image_rotary_emb,
+                    attention_kwargs=attention_kwargs,
                 )
 
         if not self.config.use_rotary_positional_embeddings: