fix many type hint errors

src/diffusers/models/attention.py

@@ -674,7 +674,7 @@ def forward(
  encoder_hidden_states: torch.FloatTensor,
  temb: torch.FloatTensor,
  joint_attention_kwargs: Optional[Dict[str, Any]] = None,
- ):
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
  joint_attention_kwargs = joint_attention_kwargs or {}
  if self.use_dual_attention:
  norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp, norm_hidden_states2, gate_msa2 = self.norm1(

src/diffusers/models/transformers/auraflow_transformer_2d.py

@@ -173,7 +173,7 @@ def forward(
  hidden_states: torch.FloatTensor,
  temb: torch.FloatTensor,
  attention_kwargs: Optional[Dict[str, Any]] = None,
- ):
+ ) -> torch.Tensor:
  residual = hidden_states
  attention_kwargs = attention_kwargs or {}
 

src/diffusers/models/transformers/cogvideox_transformer_3d.py

@@ -122,7 +122,7 @@ def forward(
  temb: torch.Tensor,
  image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
  attention_kwargs: Optional[Dict[str, Any]] = None,
- ) -> torch.Tensor:
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
  text_seq_length = encoder_hidden_states.size(1)
  attention_kwargs = attention_kwargs or {}
 

src/diffusers/models/transformers/consisid_transformer_3d.py

@@ -315,7 +315,7 @@ def forward(
  encoder_hidden_states: torch.Tensor,
  temb: torch.Tensor,
  image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
- ) -> torch.Tensor:
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
  text_seq_length = encoder_hidden_states.size(1)
 
  # norm & modulate

src/diffusers/models/transformers/lumina_nextdit2d.py

@@ -124,7 +124,7 @@ def forward(
  encoder_mask: torch.Tensor,
  temb: torch.Tensor,
  cross_attention_kwargs: Optional[Dict[str, Any]] = None,
- ):
+ ) -> torch.Tensor:
  """
  Perform a forward pass through the LuminaNextDiTBlock.
 

src/diffusers/models/transformers/transformer_bria.py

@@ -472,7 +472,7 @@ def forward(
  temb: torch.Tensor,
  image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
  attention_kwargs: Optional[Dict[str, Any]] = None,
- ) -> torch.Tensor:
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
  text_seq_len = encoder_hidden_states.shape[1]
  hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
 

src/diffusers/models/transformers/transformer_cogview3plus.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 
-from typing import Dict, Union
+from typing import Dict, Union, Tuple
 
 import torch
 import torch.nn as nn
@@ -79,7 +79,7 @@ def forward(
  hidden_states: torch.Tensor,
  encoder_hidden_states: torch.Tensor,
  emb: torch.Tensor,
- ) -> torch.Tensor:
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
  text_seq_length = encoder_hidden_states.size(1)
 
  # norm & modulate

src/diffusers/models/transformers/transformer_cogview4.py

@@ -494,7 +494,7 @@ def forward(
  ] = None,
  attention_mask: Optional[Dict[str, torch.Tensor]] = None,
  attention_kwargs: Optional[Dict[str, Any]] = None,
- ) -> torch.Tensor:
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
  # 1. Timestep conditioning
  (
  norm_hidden_states,

src/diffusers/models/transformers/transformer_hidream_image.py

@@ -534,7 +534,7 @@ def forward(
  encoder_hidden_states: Optional[torch.Tensor] = None,
  temb: Optional[torch.Tensor] = None,
  image_rotary_emb: torch.Tensor = None,
- ) -> torch.Tensor:
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
  wtype = hidden_states.dtype
  (
  shift_msa_i,
@@ -592,7 +592,7 @@ def forward(
  encoder_hidden_states: Optional[torch.Tensor] = None,
  temb: Optional[torch.Tensor] = None,
  image_rotary_emb: torch.Tensor = None,
- ) -> torch.Tensor:
+ ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
  return self.block(
  hidden_states=hidden_states,
  hidden_states_masks=hidden_states_masks,

src/diffusers/models/transformers/transformer_hunyuan_video.py

@@ -684,7 +684,7 @@ def forward(
  image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
  token_replace_emb: torch.Tensor = None,
  num_tokens: int = None,
- ) -> torch.Tensor:
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
  text_seq_length = encoder_hidden_states.shape[1]
  hidden_states = torch.cat([hidden_states, encoder_hidden_states], dim=1)