From 2d6e1efdc9cd4a9042a7f4be5109f94417717aa9 Mon Sep 17 00:00:00 2001
From: sayakpaul <spsayakpaul@gmail.com>
Date: Tue, 20 Oct 2026 21:27:32 +0530
Subject: [PATCH] update return type annotations for models.

---
 .../models/autoencoders/autoencoder_dc.py        |  6 +++---
 .../autoencoders/autoencoder_kandinsky6_sr.py    |  4 ++--
 .../models/autoencoders/autoencoder_kl.py        |  8 ++++----
 .../autoencoders/autoencoder_kl_allegro.py       |  4 ++--
 .../autoencoders/autoencoder_kl_cogvideox.py     | 12 +++++++-----
 .../models/autoencoders/autoencoder_kl_flux2.py  |  8 ++++----
 .../autoencoders/autoencoder_kl_hunyuan_video.py | 10 +++++-----
 .../autoencoders/autoencoder_kl_hunyuanimage.py  |  6 +++---
 .../autoencoder_kl_hunyuanimage_refiner.py       |  4 ++--
 .../autoencoder_kl_hunyuanvideo15.py             |  4 ++--
 .../models/autoencoders/autoencoder_kl_kvae.py   |  8 ++++----
 .../autoencoders/autoencoder_kl_kvae_audio.py    |  4 ++--
 .../autoencoders/autoencoder_kl_kvae_video.py    |  4 ++--
 .../models/autoencoders/autoencoder_kl_ltx.py    | 10 +++++-----
 .../models/autoencoders/autoencoder_kl_ltx2.py   | 10 +++++-----
 .../autoencoders/autoencoder_kl_ltx2_audio.py    |  4 ++--
 .../models/autoencoders/autoencoder_kl_magvit.py |  8 ++++----
 .../models/autoencoders/autoencoder_kl_mochi.py  |  8 ++++----
 .../autoencoders/autoencoder_kl_qwenimage.py     |  6 +++---
 .../autoencoders/autoencoder_kl_qwenimage21.py   |  6 +++---
 .../autoencoder_kl_temporal_decoder.py           |  4 ++--
 .../models/autoencoders/autoencoder_kl_wan.py    |  6 +++---
 .../models/autoencoders/autoencoder_mmaudio.py   |  4 ++--
 .../models/autoencoders/autoencoder_oobleck.py   |  6 +++---
 .../models/autoencoders/autoencoder_same.py      |  6 +++---
 .../models/autoencoders/autoencoder_vidtok.py    |  2 +-
 .../autoencoders/ltx2_diffusion_decoder.py       |  2 +-
 .../models/autoencoders/mmaudio_vocoder.py       |  2 +-
 src/diffusers/models/autoencoders/vq_model.py    |  4 ++--
 .../models/transformers/ace_step_transformer.py  |  2 +-
 .../transformers/stable_audio_transformer.py     |  2 +-
 .../models/transformers/transformer_anyflow.py   |  2 +-
 .../transformers/transformer_anyflow_far.py      | 13 +++++++++----
 .../models/transformers/transformer_bria.py      |  8 ++++----
 .../models/transformers/transformer_bria_fibo.py |  8 ++++----
 .../models/transformers/transformer_chroma.py    |  2 +-
 .../transformers/transformer_chronoedit.py       |  4 ++--
 .../models/transformers/transformer_cogview4.py  |  4 +++-
 .../models/transformers/transformer_cosmos.py    |  4 ++--
 .../transformers/transformer_easyanimate.py      |  2 +-
 .../models/transformers/transformer_flux.py      | 10 +++++-----
 .../models/transformers/transformer_flux2.py     | 10 +++++-----
 .../models/transformers/transformer_flux3.py     |  2 +-
 .../models/transformers/transformer_helios.py    |  2 +-
 .../transformers/transformer_hidream_image.py    |  4 ++--
 .../transformers/transformer_hunyuan_video.py    |  4 ++--
 .../transformers/transformer_hunyuan_video15.py  |  4 ++--
 .../transformers/transformer_hunyuanimage.py     |  8 ++++----
 .../transformer_joyimage_edit_plus.py            |  2 +-
 .../models/transformers/transformer_krea2.py     |  2 +-
 .../transformers/transformer_longcat_image.py    |  8 ++++----
 .../models/transformers/transformer_ltx.py       |  2 +-
 .../models/transformers/transformer_ltx2.py      |  4 ++--
 .../models/transformers/transformer_lumina2.py   |  2 +-
 .../models/transformers/transformer_mochi.py     |  4 ++--
 .../transformers/transformer_motif_video.py      | 10 +++++-----
 .../transformers/transformer_nucleusmoe_image.py |  2 +-
 .../transformers/transformer_ovis_image.py       |  8 ++++----
 .../models/transformers/transformer_qwenimage.py |  4 ++--
 .../transformers/transformer_qwenimage21.py      |  2 +-
 .../transformers/transformer_sana_video.py       |  2 +-
 .../models/transformers/transformer_sd3.py       |  2 +-
 .../transformers/transformer_skyreels_v2.py      |  4 ++--
 .../transformers/transformer_stable_audio3.py    |  2 +-
 .../models/transformers/transformer_temporal.py  |  2 +-
 .../models/transformers/transformer_wan.py       |  4 ++--
 .../transformers/transformer_wan_animate.py      |  4 ++--
 .../models/transformers/transformer_wan_vace.py  |  4 ++--
 src/diffusers/models/unets/unet_1d.py            |  2 +-
 src/diffusers/models/unets/unet_1d_blocks.py     | 16 ++++++++++++----
 src/diffusers/models/unets/unet_2d.py            |  2 +-
 src/diffusers/models/unets/unet_2d_condition.py  |  2 +-
 .../unets/unet_spatio_temporal_condition.py      |  2 +-
 73 files changed, 190 insertions(+), 173 deletions(-)

diff --git a/src/diffusers/models/autoencoders/autoencoder_dc.py b/src/diffusers/models/autoencoders/autoencoder_dc.py
index 859a4a6850b2..0c9f6cc36ff2 100644
--- a/src/diffusers/models/autoencoders/autoencoder_dc.py
+++ b/src/diffusers/models/autoencoders/autoencoder_dc.py
@@ -616,7 +616,7 @@ def blend_h(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.
             b[:, :, :, x] = a[:, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, x] * (x / blend_extent)
         return b
 
-    def tiled_encode(self, x: torch.Tensor, return_dict: bool = True) -> torch.Tensor:
+    def tiled_encode(self, x: torch.Tensor, return_dict: bool = True) -> EncoderOutput | tuple[torch.Tensor]:
         batch_size, num_channels, height, width = x.shape
         latent_height = height // self.spatial_compression_ratio
         latent_width = width // self.spatial_compression_ratio
@@ -664,7 +664,7 @@ def tiled_encode(self, x: torch.Tensor, return_dict: bool = True) -> torch.Tenso
             return (encoded,)
         return EncoderOutput(latent=encoded)
 
-    def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor:
+    def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple[torch.Tensor]:
         batch_size, num_channels, height, width = z.shape
 
         tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
@@ -705,7 +705,7 @@ def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutp
             return (decoded,)
         return DecoderOutput(sample=decoded)
 
-    def forward(self, sample: torch.Tensor, return_dict: bool = True) -> torch.Tensor:
+    def forward(self, sample: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple[torch.Tensor]:
         r"""
         Args:
             sample (`torch.Tensor`): Input sample.
diff --git a/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py b/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py
index 25ed70941f82..5b20f648c14f 100644
--- a/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py
+++ b/src/diffusers/models/autoencoders/autoencoder_kandinsky6_sr.py
@@ -609,7 +609,7 @@ def encode(self, x: torch.Tensor, return_dict: bool = True) -> AutoencoderKLOutp
         return AutoencoderKLOutput(latent_dist=posterior)
 
     @apply_forward_hook
-    def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple:
+    def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple[torch.Tensor]:
         r"""
         Decode latents into a video.
 
@@ -645,7 +645,7 @@ def forward(
         sample_posterior: bool = False,
         return_dict: bool = True,
         generator: torch.Generator | None = None,
-    ) -> DecoderOutput | tuple:
+    ) -> DecoderOutput | tuple[torch.Tensor]:
         r"""
         Args:
             sample (`torch.Tensor` of shape `(batch_size, channels, num_frames, height, width)`):
diff --git a/src/diffusers/models/autoencoders/autoencoder_kl.py b/src/diffusers/models/autoencoders/autoencoder_kl.py
index 4434b735d949..fc2e1fa18a92 100644
--- a/src/diffusers/models/autoencoders/autoencoder_kl.py
+++ b/src/diffusers/models/autoencoders/autoencoder_kl.py
@@ -196,7 +196,7 @@ def encode(
 
         return AutoencoderKLOutput(latent_dist=posterior)
 
-    def _decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor:
+    def _decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple[torch.Tensor]:
         if self.use_tiling and (z.shape[-1] > self.tile_latent_min_size or z.shape[-2] > self.tile_latent_min_size):
             return self.tiled_decode(z, return_dict=return_dict)
 
@@ -213,7 +213,7 @@ def _decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput |
     @apply_forward_hook
     def decode(
         self, z: torch.FloatTensor, return_dict: bool = True, generator=None
-    ) -> DecoderOutput | torch.FloatTensor:
+    ) -> DecoderOutput | tuple[torch.Tensor]:
         """
         Decode a batch of images.
 
@@ -361,7 +361,7 @@ def tiled_encode(self, x: torch.Tensor, return_dict: bool = True) -> Autoencoder
 
         return AutoencoderKLOutput(latent_dist=posterior)
 
-    def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor:
+    def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple[torch.Tensor]:
         r"""
         Decode a batch of images using a tiled decoder.
 
@@ -416,7 +416,7 @@ def forward(
         sample_posterior: bool = False,
         return_dict: bool = True,
         generator: torch.Generator | None = None,
-    ) -> DecoderOutput | torch.Tensor:
+    ) -> DecoderOutput | tuple[torch.Tensor]:
         r"""
         Args:
             sample (`torch.Tensor`): Input sample.
diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_allegro.py b/src/diffusers/models/autoencoders/autoencoder_kl_allegro.py
index 5983c08a6f86..02ff3288da57 100644
--- a/src/diffusers/models/autoencoders/autoencoder_kl_allegro.py
+++ b/src/diffusers/models/autoencoders/autoencoder_kl_allegro.py
@@ -841,7 +841,7 @@ def _decode(self, z: torch.Tensor) -> torch.Tensor:
         raise NotImplementedError("Decoding without tiling has not been implemented yet.")
 
     @apply_forward_hook
-    def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor:
+    def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple[torch.Tensor]:
         """
         Decode a batch of videos.
 
@@ -1044,7 +1044,7 @@ def forward(
         sample_posterior: bool = False,
         return_dict: bool = True,
         generator: torch.Generator | None = None,
-    ) -> DecoderOutput | torch.Tensor:
+    ) -> DecoderOutput | tuple[torch.Tensor]:
         r"""
         Args:
             sample (`torch.Tensor`): Input sample.
diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_cogvideox.py b/src/diffusers/models/autoencoders/autoencoder_kl_cogvideox.py
index ed624dc9e62e..a5dfad8c79af 100644
--- a/src/diffusers/models/autoencoders/autoencoder_kl_cogvideox.py
+++ b/src/diffusers/models/autoencoders/autoencoder_kl_cogvideox.py
@@ -132,7 +132,9 @@ def fake_context_parallel_forward(
                 inputs = torch.cat(cached_inputs + [inputs], dim=2)
         return inputs
 
-    def forward(self, inputs: torch.Tensor, conv_cache: torch.Tensor | None = None) -> torch.Tensor:
+    def forward(
+        self, inputs: torch.Tensor, conv_cache: torch.Tensor | None = None
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
         inputs = self.fake_context_parallel_forward(inputs, conv_cache)
 
         if self.pad_mode == "replicate":
@@ -1176,7 +1178,7 @@ def encode(
             return (posterior,)
         return AutoencoderKLOutput(latent_dist=posterior)
 
-    def _decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor:
+    def _decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple[torch.Tensor]:
         batch_size, num_channels, num_frames, height, width = z.shape
 
         if self.use_tiling and (width > self.tile_latent_min_width or height > self.tile_latent_min_height):
@@ -1205,7 +1207,7 @@ def _decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput |
         return DecoderOutput(sample=dec)
 
     @apply_forward_hook
-    def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor:
+    def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple[torch.Tensor]:
         """
         Decode a batch of images.
 
@@ -1319,7 +1321,7 @@ def tiled_encode(self, x: torch.Tensor) -> torch.Tensor:
         enc = torch.cat(result_rows, dim=3)
         return enc
 
-    def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor:
+    def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple[torch.Tensor]:
         r"""
         Decode a batch of images using a tiled decoder.
 
@@ -1408,7 +1410,7 @@ def forward(
         sample_posterior: bool = False,
         return_dict: bool = True,
         generator: torch.Generator | None = None,
-    ) -> torch.Tensor | torch.Tensor:
+    ) -> DecoderOutput | tuple[torch.Tensor]:
         r"""
         Args:
             sample (`torch.Tensor`): Input sample.
diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_flux2.py b/src/diffusers/models/autoencoders/autoencoder_kl_flux2.py
index 37d69c8a1dc2..3a501768e7a8 100644
--- a/src/diffusers/models/autoencoders/autoencoder_kl_flux2.py
+++ b/src/diffusers/models/autoencoders/autoencoder_kl_flux2.py
@@ -208,7 +208,7 @@ def encode(
 
         return AutoencoderKLOutput(latent_dist=posterior)
 
-    def _decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor:
+    def _decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple[torch.Tensor]:
         if self.use_tiling and (z.shape[-1] > self.tile_latent_min_size or z.shape[-2] > self.tile_latent_min_size):
             return self.tiled_decode(z, return_dict=return_dict)
 
@@ -225,7 +225,7 @@ def _decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput |
     @apply_forward_hook
     def decode(
         self, z: torch.FloatTensor, return_dict: bool = True, generator=None
-    ) -> DecoderOutput | torch.FloatTensor:
+    ) -> DecoderOutput | tuple[torch.Tensor]:
         """
         Decode a batch of images.
 
@@ -373,7 +373,7 @@ def tiled_encode(self, x: torch.Tensor, return_dict: bool = True) -> Autoencoder
 
         return AutoencoderKLOutput(latent_dist=posterior)
 
-    def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor:
+    def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple[torch.Tensor]:
         r"""
         Decode a batch of images using a tiled decoder.
 
@@ -428,7 +428,7 @@ def forward(
         sample_posterior: bool = False,
         return_dict: bool = True,
         generator: torch.Generator | None = None,
-    ) -> DecoderOutput | torch.Tensor:
+    ) -> DecoderOutput | tuple[torch.Tensor]:
         r"""
         Args:
             sample (`torch.Tensor`): Input sample.
diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_hunyuan_video.py b/src/diffusers/models/autoencoders/autoencoder_kl_hunyuan_video.py
index fece756ebec6..48b5bcf375aa 100644
--- a/src/diffusers/models/autoencoders/autoencoder_kl_hunyuan_video.py
+++ b/src/diffusers/models/autoencoders/autoencoder_kl_hunyuan_video.py
@@ -802,7 +802,7 @@ def encode(
             return (posterior,)
         return AutoencoderKLOutput(latent_dist=posterior)
 
-    def _decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor:
+    def _decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple[torch.Tensor]:
         batch_size, num_channels, num_frames, height, width = z.shape
         tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
         tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio
@@ -823,7 +823,7 @@ def _decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput |
         return DecoderOutput(sample=dec)
 
     @apply_forward_hook
-    def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor:
+    def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple[torch.Tensor]:
         r"""
         Decode a batch of images.
 
@@ -922,7 +922,7 @@ def tiled_encode(self, x: torch.Tensor) -> AutoencoderKLOutput:
         enc = torch.cat(result_rows, dim=3)[:, :, :, :latent_height, :latent_width]
         return enc
 
-    def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor:
+    def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple[torch.Tensor]:
         r"""
         Decode a batch of images using a tiled decoder.
 
@@ -1011,7 +1011,7 @@ def _temporal_tiled_encode(self, x: torch.Tensor) -> AutoencoderKLOutput:
         enc = torch.cat(result_row, dim=2)[:, :, :latent_num_frames]
         return enc
 
-    def _temporal_tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor:
+    def _temporal_tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple[torch.Tensor]:
         batch_size, num_channels, num_frames, height, width = z.shape
         num_sample_frames = (num_frames - 1) * self.temporal_compression_ratio + 1
 
@@ -1053,7 +1053,7 @@ def forward(
         sample_posterior: bool = False,
         return_dict: bool = True,
         generator: torch.Generator | None = None,
-    ) -> DecoderOutput | torch.Tensor:
+    ) -> DecoderOutput | tuple[torch.Tensor]:
         r"""
         Args:
             sample (`torch.Tensor`): Input sample.
diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_hunyuanimage.py b/src/diffusers/models/autoencoders/autoencoder_kl_hunyuanimage.py
index 4d46b7c32ec4..e8ea3d57eefc 100644
--- a/src/diffusers/models/autoencoders/autoencoder_kl_hunyuanimage.py
+++ b/src/diffusers/models/autoencoders/autoencoder_kl_hunyuanimage.py
@@ -538,7 +538,7 @@ def _decode(self, z: torch.Tensor, return_dict: bool = True):
         return DecoderOutput(sample=dec)
 
     @apply_forward_hook
-    def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor:
+    def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple[torch.Tensor]:
         r"""
         Decode a batch of images.
 
@@ -619,7 +619,7 @@ def tiled_encode(self, x: torch.Tensor) -> torch.Tensor:
 
         return moments
 
-    def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor:
+    def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple[torch.Tensor]:
         """
         Decode latent using spatial tiling strategy.
 
@@ -670,7 +670,7 @@ def forward(
         sample_posterior: bool = False,
         return_dict: bool = True,
         generator: torch.Generator | None = None,
-    ) -> DecoderOutput | torch.Tensor:
+    ) -> DecoderOutput | tuple[torch.Tensor]:
         """
         Args:
             sample (`torch.Tensor`): Input sample.
diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_hunyuanimage_refiner.py b/src/diffusers/models/autoencoders/autoencoder_kl_hunyuanimage_refiner.py
index 091304660ba8..f924670810e6 100644
--- a/src/diffusers/models/autoencoders/autoencoder_kl_hunyuanimage_refiner.py
+++ b/src/diffusers/models/autoencoders/autoencoder_kl_hunyuanimage_refiner.py
@@ -741,7 +741,7 @@ def _decode(self, z: torch.Tensor) -> torch.Tensor:
         return dec
 
     @apply_forward_hook
-    def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor:
+    def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple[torch.Tensor]:
         r"""
         Decode a batch of images.
 
@@ -896,7 +896,7 @@ def forward(
         sample_posterior: bool = False,
         return_dict: bool = True,
         generator: torch.Generator | None = None,
-    ) -> DecoderOutput | torch.Tensor:
+    ) -> DecoderOutput | tuple[torch.Tensor]:
         r"""
         Args:
             sample (`torch.Tensor`): Input sample.
diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_hunyuanvideo15.py b/src/diffusers/models/autoencoders/autoencoder_kl_hunyuanvideo15.py
index 00e28ee620b6..4c0382b86dcf 100644
--- a/src/diffusers/models/autoencoders/autoencoder_kl_hunyuanvideo15.py
+++ b/src/diffusers/models/autoencoders/autoencoder_kl_hunyuanvideo15.py
@@ -777,7 +777,7 @@ def _decode(self, z: torch.Tensor) -> torch.Tensor:
         return dec
 
     @apply_forward_hook
-    def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor:
+    def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple[torch.Tensor]:
         r"""
         Decode a batch of images.
 
@@ -929,7 +929,7 @@ def forward(
         sample_posterior: bool = False,
         return_dict: bool = True,
         generator: torch.Generator | None = None,
-    ) -> DecoderOutput | torch.Tensor:
+    ) -> DecoderOutput | tuple[torch.Tensor]:
         r"""
         Args:
             sample (`torch.Tensor`): Input sample.
diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_kvae.py b/src/diffusers/models/autoencoders/autoencoder_kl_kvae.py
index dc8b9e4c36e7..7e7fd7c05114 100644
--- a/src/diffusers/models/autoencoders/autoencoder_kl_kvae.py
+++ b/src/diffusers/models/autoencoders/autoencoder_kl_kvae.py
@@ -628,7 +628,7 @@ def encode(
 
         return AutoencoderKLOutput(latent_dist=posterior)
 
-    def _decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]:
+    def _decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, tuple[torch.Tensor]]:
         if self.use_tiling and (z.shape[-1] > self.tile_latent_min_size or z.shape[-2] > self.tile_latent_min_size):
             return self.tiled_decode(z, return_dict=return_dict)
 
@@ -642,7 +642,7 @@ def _decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOut
     @apply_forward_hook
     def decode(
         self, z: torch.FloatTensor, return_dict: bool = True, generator=None
-    ) -> Union[DecoderOutput, torch.FloatTensor]:
+    ) -> Union[DecoderOutput, tuple[torch.Tensor]]:
         """
         Decode a batch of images.
 
@@ -726,7 +726,7 @@ def _tiled_encode(self, x: torch.Tensor) -> torch.Tensor:
         enc = torch.cat(result_rows, dim=2)
         return enc
 
-    def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]:
+    def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, tuple[torch.Tensor]]:
         r"""
         Decode a batch of images using a tiled decoder.
 
@@ -779,7 +779,7 @@ def forward(
         sample_posterior: bool = False,
         return_dict: bool = True,
         generator: Optional[torch.Generator] = None,
-    ) -> Union[DecoderOutput, torch.Tensor]:
+    ) -> Union[DecoderOutput, tuple[torch.Tensor]]:
         r"""
         Args:
             sample (`torch.Tensor`): Input sample.
diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_kvae_audio.py b/src/diffusers/models/autoencoders/autoencoder_kl_kvae_audio.py
index ad99f19e2a82..dac72f7ba298 100644
--- a/src/diffusers/models/autoencoders/autoencoder_kl_kvae_audio.py
+++ b/src/diffusers/models/autoencoders/autoencoder_kl_kvae_audio.py
@@ -352,7 +352,7 @@ def _decode(self, z: torch.Tensor) -> torch.Tensor:
         return self.decoder(z)
 
     @apply_forward_hook
-    def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor:
+    def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple[torch.Tensor]:
         """
         Decode a batch of latents into audio waveforms.
 
@@ -382,7 +382,7 @@ def forward(
         sample_posterior: bool = False,
         return_dict: bool = True,
         generator: Optional[torch.Generator] = None,
-    ) -> DecoderOutput | torch.Tensor:
+    ) -> DecoderOutput | tuple[torch.Tensor]:
         r"""
         Args:
             sample (`torch.Tensor` of shape `(batch_size, num_channels, num_samples)`):
diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_kvae_video.py b/src/diffusers/models/autoencoders/autoencoder_kl_kvae_video.py
index 26a7d5b2ef1c..d81a8253dc57 100644
--- a/src/diffusers/models/autoencoders/autoencoder_kl_kvae_video.py
+++ b/src/diffusers/models/autoencoders/autoencoder_kl_kvae_video.py
@@ -913,7 +913,7 @@ def _decode(self, z: torch.Tensor, seg_len: int = 16) -> torch.Tensor:
         return torch.cat(recs, dim=2)
 
     @apply_forward_hook
-    def decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]:
+    def decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, tuple[torch.Tensor]]:
         """
         Decode a batch of videos.
 
@@ -941,7 +941,7 @@ def forward(
         sample_posterior: bool = False,
         return_dict: bool = True,
         generator: Optional[torch.Generator] = None,
-    ) -> Union[DecoderOutput, torch.Tensor]:
+    ) -> Union[DecoderOutput, tuple[torch.Tensor]]:
         r"""
         Args:
             sample (`torch.Tensor`): Input sample.
diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_ltx.py b/src/diffusers/models/autoencoders/autoencoder_kl_ltx.py
index 8cb646e8b5db..6713ac81dbd6 100644
--- a/src/diffusers/models/autoencoders/autoencoder_kl_ltx.py
+++ b/src/diffusers/models/autoencoders/autoencoder_kl_ltx.py
@@ -1259,7 +1259,7 @@ def encode(
 
     def _decode(
         self, z: torch.Tensor, temb: torch.Tensor | None = None, return_dict: bool = True
-    ) -> DecoderOutput | torch.Tensor:
+    ) -> DecoderOutput | tuple[torch.Tensor]:
         batch_size, num_channels, num_frames, height, width = z.shape
         tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
         tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio
@@ -1281,7 +1281,7 @@ def _decode(
     @apply_forward_hook
     def decode(
         self, z: torch.Tensor, temb: torch.Tensor | None = None, return_dict: bool = True
-    ) -> DecoderOutput | torch.Tensor:
+    ) -> DecoderOutput | tuple[torch.Tensor]:
         """
         Decode a batch of images.
 
@@ -1388,7 +1388,7 @@ def tiled_encode(self, x: torch.Tensor) -> torch.Tensor:
 
     def tiled_decode(
         self, z: torch.Tensor, temb: torch.Tensor | None, return_dict: bool = True
-    ) -> DecoderOutput | torch.Tensor:
+    ) -> DecoderOutput | tuple[torch.Tensor]:
         r"""
         Decode a batch of images using a tiled decoder.
 
@@ -1478,7 +1478,7 @@ def _temporal_tiled_encode(self, x: torch.Tensor) -> AutoencoderKLOutput:
 
     def _temporal_tiled_decode(
         self, z: torch.Tensor, temb: torch.Tensor | None, return_dict: bool = True
-    ) -> DecoderOutput | torch.Tensor:
+    ) -> DecoderOutput | tuple[torch.Tensor]:
         batch_size, num_channels, num_frames, height, width = z.shape
         num_sample_frames = (num_frames - 1) * self.temporal_compression_ratio + 1
 
@@ -1521,7 +1521,7 @@ def forward(
         sample_posterior: bool = False,
         return_dict: bool = True,
         generator: torch.Generator | None = None,
-    ) -> torch.Tensor | torch.Tensor:
+    ) -> DecoderOutput | tuple[torch.Tensor]:
         r"""
         Args:
             sample (`torch.Tensor`): Input sample.
diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_ltx2.py b/src/diffusers/models/autoencoders/autoencoder_kl_ltx2.py
index e64a9347f63d..8237a5df0696 100644
--- a/src/diffusers/models/autoencoders/autoencoder_kl_ltx2.py
+++ b/src/diffusers/models/autoencoders/autoencoder_kl_ltx2.py
@@ -1269,7 +1269,7 @@ def _decode(
         temb: torch.Tensor | None = None,
         causal: bool | None = None,
         return_dict: bool = True,
-    ) -> DecoderOutput | torch.Tensor:
+    ) -> DecoderOutput | tuple[torch.Tensor]:
         batch_size, num_channels, num_frames, height, width = z.shape
         tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
         tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio
@@ -1295,7 +1295,7 @@ def decode(
         temb: torch.Tensor | None = None,
         causal: bool | None = None,
         return_dict: bool = True,
-    ) -> DecoderOutput | torch.Tensor:
+    ) -> DecoderOutput | tuple[torch.Tensor]:
         """
         Decode a batch of images.
 
@@ -1404,7 +1404,7 @@ def tiled_encode(self, x: torch.Tensor, causal: bool | None = None) -> torch.Ten
 
     def tiled_decode(
         self, z: torch.Tensor, temb: torch.Tensor | None, causal: bool | None = None, return_dict: bool = True
-    ) -> DecoderOutput | torch.Tensor:
+    ) -> DecoderOutput | tuple[torch.Tensor]:
         r"""
         Decode a batch of images using a tiled decoder.
 
@@ -1496,7 +1496,7 @@ def _temporal_tiled_encode(self, x: torch.Tensor, causal: bool | None = None) ->
 
     def _temporal_tiled_decode(
         self, z: torch.Tensor, temb: torch.Tensor | None, causal: bool | None = None, return_dict: bool = True
-    ) -> DecoderOutput | torch.Tensor:
+    ) -> DecoderOutput | tuple[torch.Tensor]:
         batch_size, num_channels, num_frames, height, width = z.shape
         num_sample_frames = (num_frames - 1) * self.temporal_compression_ratio + 1
 
@@ -1541,7 +1541,7 @@ def forward(
         decoder_causal: bool | None = None,
         return_dict: bool = True,
         generator: torch.Generator | None = None,
-    ) -> torch.Tensor | torch.Tensor:
+    ) -> DecoderOutput | tuple[torch.Tensor]:
         r"""
         Args:
             sample (`torch.Tensor`): Input sample.
diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_ltx2_audio.py b/src/diffusers/models/autoencoders/autoencoder_kl_ltx2_audio.py
index fb773dbdc01e..2b1c4a3b8ddf 100644
--- a/src/diffusers/models/autoencoders/autoencoder_kl_ltx2_audio.py
+++ b/src/diffusers/models/autoencoders/autoencoder_kl_ltx2_audio.py
@@ -773,7 +773,7 @@ def _decode(self, z: torch.Tensor) -> torch.Tensor:
         return self.decoder(z)
 
     @apply_forward_hook
-    def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor:
+    def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple[torch.Tensor]:
         if self.use_slicing and z.shape[0] > 1:
             decoded_slices = [self._decode(z_slice) for z_slice in z.split(1)]
             decoded = torch.cat(decoded_slices)
@@ -791,7 +791,7 @@ def forward(
         sample_posterior: bool = False,
         return_dict: bool = True,
         generator: torch.Generator | None = None,
-    ) -> DecoderOutput | torch.Tensor:
+    ) -> DecoderOutput | tuple[torch.Tensor]:
         r"""
         Args:
             sample (`torch.Tensor`): Input sample.
diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_magvit.py b/src/diffusers/models/autoencoders/autoencoder_kl_magvit.py
index 9f9718e13584..9326efa96639 100644
--- a/src/diffusers/models/autoencoders/autoencoder_kl_magvit.py
+++ b/src/diffusers/models/autoencoders/autoencoder_kl_magvit.py
@@ -862,7 +862,7 @@ def encode(
             return (posterior,)
         return AutoencoderKLOutput(latent_dist=posterior)
 
-    def _decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor:
+    def _decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple[torch.Tensor]:
         batch_size, num_channels, num_frames, height, width = z.shape
         tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
         tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio
@@ -889,7 +889,7 @@ def _decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput |
         return DecoderOutput(sample=dec)
 
     @apply_forward_hook
-    def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor:
+    def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple[torch.Tensor]:
         """
         Decode a batch of images.
 
@@ -982,7 +982,7 @@ def tiled_encode(self, x: torch.Tensor, return_dict: bool = True) -> Autoencoder
         moments = torch.cat(result_rows, dim=3)[:, :, :, :latent_height, :latent_width]
         return moments
 
-    def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor:
+    def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple[torch.Tensor]:
         batch_size, num_channels, num_frames, height, width = z.shape
         sample_height = height * self.spatial_compression_ratio
         sample_width = width * self.spatial_compression_ratio
@@ -1049,7 +1049,7 @@ def forward(
         sample_posterior: bool = False,
         return_dict: bool = True,
         generator: torch.Generator | None = None,
-    ) -> DecoderOutput | torch.Tensor:
+    ) -> DecoderOutput | tuple[torch.Tensor]:
         r"""
         Args:
             sample (`torch.Tensor`): Input sample.
diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_mochi.py b/src/diffusers/models/autoencoders/autoencoder_kl_mochi.py
index bb447015c54d..5383298542f7 100644
--- a/src/diffusers/models/autoencoders/autoencoder_kl_mochi.py
+++ b/src/diffusers/models/autoencoders/autoencoder_kl_mochi.py
@@ -880,7 +880,7 @@ def encode(
             return (posterior,)
         return AutoencoderKLOutput(latent_dist=posterior)
 
-    def _decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor:
+    def _decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple[torch.Tensor]:
         batch_size, num_channels, num_frames, height, width = z.shape
         tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
         tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio
@@ -910,7 +910,7 @@ def _decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput |
         return DecoderOutput(sample=dec)
 
     @apply_forward_hook
-    def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor:
+    def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple[torch.Tensor]:
         """
         Decode a batch of images.
 
@@ -1008,7 +1008,7 @@ def tiled_encode(self, x: torch.Tensor) -> torch.Tensor:
         enc = torch.cat(result_rows, dim=3)[:, :, :, :latent_height, :latent_width]
         return enc
 
-    def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor:
+    def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple[torch.Tensor]:
         r"""
         Decode a batch of images using a tiled decoder.
 
@@ -1092,7 +1092,7 @@ def forward(
         sample_posterior: bool = False,
         return_dict: bool = True,
         generator: torch.Generator | None = None,
-    ) -> DecoderOutput | torch.Tensor:
+    ) -> DecoderOutput | tuple[torch.Tensor]:
         r"""
         Args:
             sample (`torch.Tensor`): Input sample.
diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_qwenimage.py b/src/diffusers/models/autoencoders/autoencoder_kl_qwenimage.py
index fa9f31ac8512..0ff82463a59b 100644
--- a/src/diffusers/models/autoencoders/autoencoder_kl_qwenimage.py
+++ b/src/diffusers/models/autoencoders/autoencoder_kl_qwenimage.py
@@ -864,7 +864,7 @@ def _decode(self, z: torch.Tensor, return_dict: bool = True):
         return DecoderOutput(sample=out)
 
     @apply_forward_hook
-    def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor:
+    def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple[torch.Tensor]:
         r"""
         Decode a batch of images.
 
@@ -970,7 +970,7 @@ def tiled_encode(self, x: torch.Tensor) -> AutoencoderKLOutput:
         enc = torch.cat(result_rows, dim=3)[:, :, :, :latent_height, :latent_width]
         return enc
 
-    def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor:
+    def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple[torch.Tensor]:
         r"""
         Decode a batch of images using a tiled decoder.
 
@@ -1039,7 +1039,7 @@ def forward(
         sample_posterior: bool = False,
         return_dict: bool = True,
         generator: torch.Generator | None = None,
-    ) -> DecoderOutput | torch.Tensor:
+    ) -> DecoderOutput | tuple[torch.Tensor]:
         """
         Args:
             sample (`torch.Tensor`): Input sample.
diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_qwenimage21.py b/src/diffusers/models/autoencoders/autoencoder_kl_qwenimage21.py
index ac6af47ab493..aaba75706d54 100644
--- a/src/diffusers/models/autoencoders/autoencoder_kl_qwenimage21.py
+++ b/src/diffusers/models/autoencoders/autoencoder_kl_qwenimage21.py
@@ -1335,7 +1335,7 @@ def _decode(self, z: torch.Tensor, return_dict: bool = True):
 
     @apply_forward_hook
     # Copied from diffusers.models.autoencoders.autoencoder_kl_wan.AutoencoderKLWan.decode
-    def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor:
+    def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple[torch.Tensor]:
         r"""
         Decode a batch of images.
 
@@ -1451,7 +1451,7 @@ def tiled_encode(self, x: torch.Tensor) -> AutoencoderKLOutput:
         return enc
 
     # Copied from diffusers.models.autoencoders.autoencoder_kl_wan.AutoencoderKLWan.tiled_decode with unpatchify->_unpatchify
-    def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor:
+    def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple[torch.Tensor]:
         r"""
         Decode a batch of images using a tiled decoder.
 
@@ -1535,7 +1535,7 @@ def forward(
         sample_posterior: bool = False,
         return_dict: bool = True,
         generator: torch.Generator | None = None,
-    ) -> DecoderOutput | torch.Tensor:
+    ) -> DecoderOutput | tuple[torch.Tensor]:
         """
         Args:
             sample (`torch.Tensor`): Input sample.
diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_temporal_decoder.py b/src/diffusers/models/autoencoders/autoencoder_kl_temporal_decoder.py
index 8b0e5806d8ef..651c0dba74b6 100644
--- a/src/diffusers/models/autoencoders/autoencoder_kl_temporal_decoder.py
+++ b/src/diffusers/models/autoencoders/autoencoder_kl_temporal_decoder.py
@@ -248,7 +248,7 @@ def decode(
         z: torch.Tensor,
         num_frames: int,
         return_dict: bool = True,
-    ) -> DecoderOutput | torch.Tensor:
+    ) -> DecoderOutput | tuple[torch.Tensor]:
         """
         Decode a batch of images.
 
@@ -279,7 +279,7 @@ def forward(
         return_dict: bool = True,
         generator: torch.Generator | None = None,
         num_frames: int = 1,
-    ) -> DecoderOutput | torch.Tensor:
+    ) -> DecoderOutput | tuple[torch.Tensor]:
         r"""
         Args:
             sample (`torch.Tensor`): Input sample.
diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_wan.py b/src/diffusers/models/autoencoders/autoencoder_kl_wan.py
index dcb280ff21a2..dd6dfa8c7fe8 100644
--- a/src/diffusers/models/autoencoders/autoencoder_kl_wan.py
+++ b/src/diffusers/models/autoencoders/autoencoder_kl_wan.py
@@ -1220,7 +1220,7 @@ def _decode(self, z: torch.Tensor, return_dict: bool = True):
         return DecoderOutput(sample=out)
 
     @apply_forward_hook
-    def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor:
+    def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple[torch.Tensor]:
         r"""
         Decode a batch of images.
 
@@ -1332,7 +1332,7 @@ def tiled_encode(self, x: torch.Tensor) -> AutoencoderKLOutput:
         enc = torch.cat(result_rows, dim=3)[:, :, :, :latent_height, :latent_width]
         return enc
 
-    def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor:
+    def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple[torch.Tensor]:
         r"""
         Decode a batch of images using a tiled decoder.
 
@@ -1416,7 +1416,7 @@ def forward(
         sample_posterior: bool = False,
         return_dict: bool = True,
         generator: torch.Generator | None = None,
-    ) -> DecoderOutput | torch.Tensor:
+    ) -> DecoderOutput | tuple[torch.Tensor]:
         """
         Args:
             sample (`torch.Tensor`): Input sample.
diff --git a/src/diffusers/models/autoencoders/autoencoder_mmaudio.py b/src/diffusers/models/autoencoders/autoencoder_mmaudio.py
index c19ad4fd3619..d33b82cfa987 100644
--- a/src/diffusers/models/autoencoders/autoencoder_mmaudio.py
+++ b/src/diffusers/models/autoencoders/autoencoder_mmaudio.py
@@ -418,7 +418,7 @@ def encode(self, audio: torch.Tensor, return_dict: bool = True) -> AutoencoderKL
         return AutoencoderKLOutput(latent_dist=posterior)
 
     @apply_forward_hook
-    def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple:
+    def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple[torch.Tensor]:
         r"""
         Decode latents into a mel spectrogram.
 
@@ -442,7 +442,7 @@ def forward(
         sample_posterior: bool = False,
         return_dict: bool = True,
         generator: torch.Generator | None = None,
-    ) -> DecoderOutput | tuple:
+    ) -> DecoderOutput | tuple[torch.Tensor]:
         r"""
         Args:
             sample (`torch.Tensor` of shape `(batch_size, num_samples)`):
diff --git a/src/diffusers/models/autoencoders/autoencoder_oobleck.py b/src/diffusers/models/autoencoders/autoencoder_oobleck.py
index 639199a8dd36..fa0069c86efd 100644
--- a/src/diffusers/models/autoencoders/autoencoder_oobleck.py
+++ b/src/diffusers/models/autoencoders/autoencoder_oobleck.py
@@ -440,7 +440,7 @@ def _tiled_encode(self, x: torch.Tensor) -> torch.Tensor:
 
         return torch.cat(tiles, dim=-1)
 
-    def _decode(self, z: torch.Tensor, return_dict: bool = True) -> OobleckDecoderOutput | torch.Tensor:
+    def _decode(self, z: torch.Tensor, return_dict: bool = True) -> OobleckDecoderOutput | tuple[torch.Tensor]:
         if self.use_tiling and z.shape[-1] > self.tile_latent_min_length:
             dec = self._tiled_decode(z)
         else:
@@ -488,7 +488,7 @@ def _tiled_decode(self, z: torch.Tensor) -> torch.Tensor:
     @apply_forward_hook
     def decode(
         self, z: torch.FloatTensor, return_dict: bool = True, generator=None
-    ) -> OobleckDecoderOutput | torch.FloatTensor:
+    ) -> OobleckDecoderOutput | tuple[torch.Tensor]:
         """
         Decode a batch of images.
 
@@ -520,7 +520,7 @@ def forward(
         sample_posterior: bool = False,
         return_dict: bool = True,
         generator: torch.Generator | None = None,
-    ) -> OobleckDecoderOutput | torch.Tensor:
+    ) -> OobleckDecoderOutput | tuple[torch.Tensor]:
         r"""
         Args:
             sample (`torch.Tensor`): Input sample.
diff --git a/src/diffusers/models/autoencoders/autoencoder_same.py b/src/diffusers/models/autoencoders/autoencoder_same.py
index 23dccf7732d9..33e397ae11e7 100644
--- a/src/diffusers/models/autoencoders/autoencoder_same.py
+++ b/src/diffusers/models/autoencoders/autoencoder_same.py
@@ -735,7 +735,7 @@ def encode(
         self,
         sample: torch.Tensor,
         return_dict: bool = True,
-    ) -> Union[AutoencoderSAMEOutput, torch.Tensor]:
+    ) -> Union[AutoencoderSAMEOutput, tuple[torch.Tensor]]:
         """
         Encode stereo audio to latents.
 
@@ -760,7 +760,7 @@ def decode(
         self,
         latents: torch.Tensor,
         return_dict: bool = True,
-    ) -> Union[AutoencoderSAMEDecoderOutput, torch.Tensor]:
+    ) -> Union[AutoencoderSAMEDecoderOutput, tuple[torch.Tensor]]:
         """
         Decode latents back to stereo audio.
 
@@ -784,7 +784,7 @@ def forward(
         self,
         sample: torch.Tensor,
         return_dict: bool = True,
-    ) -> Union[AutoencoderSAMEDecoderOutput, torch.Tensor]:
+    ) -> Union[AutoencoderSAMEDecoderOutput, tuple[torch.Tensor]]:
         """
         Encode and immediately decode *sample* (reconstruction).
 
diff --git a/src/diffusers/models/autoencoders/autoencoder_vidtok.py b/src/diffusers/models/autoencoders/autoencoder_vidtok.py
index 296c7bd8d85a..13f15c639d7f 100644
--- a/src/diffusers/models/autoencoders/autoencoder_vidtok.py
+++ b/src/diffusers/models/autoencoders/autoencoder_vidtok.py
@@ -1439,7 +1439,7 @@ def forward(
         encoder_mode: bool = False,
         return_dict: bool = True,
         generator: Optional[torch.Generator] = None,
-    ) -> Union[torch.Tensor, DecoderOutput]:
+    ) -> Union[torch.Tensor, DecoderOutput, tuple[torch.Tensor], tuple[torch.Tensor, torch.Tensor]]:
         r"""
         Args:
             sample (`torch.Tensor`): Input sample.
diff --git a/src/diffusers/models/autoencoders/ltx2_diffusion_decoder.py b/src/diffusers/models/autoencoders/ltx2_diffusion_decoder.py
index 41388991e0b4..06ccc8d5865d 100644
--- a/src/diffusers/models/autoencoders/ltx2_diffusion_decoder.py
+++ b/src/diffusers/models/autoencoders/ltx2_diffusion_decoder.py
@@ -992,7 +992,7 @@ def decode(
         generator: torch.Generator | None = None,
         num_inference_steps: int | None = None,
         return_dict: bool = True,
-    ) -> DecoderOutput | torch.Tensor:
+    ) -> DecoderOutput | tuple[torch.Tensor]:
         """Decode a batch of latents.
 
         `z` is expected to be denormalized already (the pipeline applies `latents_mean` / `latents_std`), matching
diff --git a/src/diffusers/models/autoencoders/mmaudio_vocoder.py b/src/diffusers/models/autoencoders/mmaudio_vocoder.py
index d1f12fb5ef62..e5482c7524da 100644
--- a/src/diffusers/models/autoencoders/mmaudio_vocoder.py
+++ b/src/diffusers/models/autoencoders/mmaudio_vocoder.py
@@ -210,7 +210,7 @@ def __init__(
         self.activation_post = MMAudioActivation1d(channels)
         self.conv_post = nn.Conv1d(channels, 1, 7, padding=3, bias=False)
 
-    def forward(self, mel: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple:
+    def forward(self, mel: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple[torch.Tensor]:
         r"""
         Args:
             mel (`torch.Tensor` of shape `(batch_size, num_mels, num_mel_frames)`):
diff --git a/src/diffusers/models/autoencoders/vq_model.py b/src/diffusers/models/autoencoders/vq_model.py
index 619327dde417..af97946297cf 100644
--- a/src/diffusers/models/autoencoders/vq_model.py
+++ b/src/diffusers/models/autoencoders/vq_model.py
@@ -130,7 +130,7 @@ def __init__(
         )
 
     @apply_forward_hook
-    def encode(self, x: torch.Tensor, return_dict: bool = True) -> VQEncoderOutput:
+    def encode(self, x: torch.Tensor, return_dict: bool = True) -> VQEncoderOutput | tuple[torch.Tensor]:
         h = self.encoder(x)
         h = self.quant_conv(h)
 
@@ -142,7 +142,7 @@ def encode(self, x: torch.Tensor, return_dict: bool = True) -> VQEncoderOutput:
     @apply_forward_hook
     def decode(
         self, h: torch.Tensor, force_not_quantize: bool = False, return_dict: bool = True, shape=None
-    ) -> DecoderOutput | torch.Tensor:
+    ) -> DecoderOutput | tuple[torch.Tensor, torch.Tensor]:
         # also go through quantization layer
         if not force_not_quantize:
             quant, commit_loss, _ = self.quantize(h)
diff --git a/src/diffusers/models/transformers/ace_step_transformer.py b/src/diffusers/models/transformers/ace_step_transformer.py
index 821c7ad1491a..31b2fb5fd783 100644
--- a/src/diffusers/models/transformers/ace_step_transformer.py
+++ b/src/diffusers/models/transformers/ace_step_transformer.py
@@ -539,7 +539,7 @@ def forward(
         context_latents: torch.Tensor,
         attention_kwargs: Optional[dict] = None,
         return_dict: bool = True,
-    ) -> Union[torch.Tensor, Transformer2DModelOutput]:
+    ) -> Union[Transformer2DModelOutput, tuple[torch.Tensor]]:
         """The [`AceStepTransformer1DModel`] forward method.
 
         Args:
diff --git a/src/diffusers/models/transformers/stable_audio_transformer.py b/src/diffusers/models/transformers/stable_audio_transformer.py
index 72c42048db71..2849bec52fa5 100644
--- a/src/diffusers/models/transformers/stable_audio_transformer.py
+++ b/src/diffusers/models/transformers/stable_audio_transformer.py
@@ -289,7 +289,7 @@ def forward(
         return_dict: bool = True,
         attention_mask: torch.LongTensor | None = None,
         encoder_attention_mask: torch.LongTensor | None = None,
-    ) -> torch.FloatTensor | Transformer2DModelOutput:
+    ) -> Transformer2DModelOutput | tuple[torch.Tensor]:
         """
         The [`StableAudioDiTModel`] forward method.
 
diff --git a/src/diffusers/models/transformers/transformer_anyflow.py b/src/diffusers/models/transformers/transformer_anyflow.py
index 1388c63096b0..b0b8e03c5e5d 100644
--- a/src/diffusers/models/transformers/transformer_anyflow.py
+++ b/src/diffusers/models/transformers/transformer_anyflow.py
@@ -633,7 +633,7 @@ def forward(
         encoder_hidden_states_image: Optional[torch.Tensor] = None,
         attention_kwargs: Optional[Dict[str, Any]] = None,
         return_dict: bool = True,
-    ) -> Union[Transformer2DModelOutput, Tuple]:
+    ) -> Union[Transformer2DModelOutput, Tuple[torch.Tensor]]:
         """
         Bidirectional flow-map forward pass. ``hidden_states`` is laid out as ``(B, F, C, H, W)`` (per-frame latents).
         The input is patchified with the standard ``patch_embedding`` (kernel = stride = ``patch_size``) and denoised
diff --git a/src/diffusers/models/transformers/transformer_anyflow_far.py b/src/diffusers/models/transformers/transformer_anyflow_far.py
index a5fdf84ac829..bb8856b1f422 100644
--- a/src/diffusers/models/transformers/transformer_anyflow_far.py
+++ b/src/diffusers/models/transformers/transformer_anyflow_far.py
@@ -1111,7 +1111,12 @@ def forward(
         attention_mask: Optional[BlockMask] = None,
         attention_kwargs: Optional[Dict[str, Any]] = None,
         return_dict: bool = True,
-    ) -> Union[Transformer2DModelOutput, AnyFlowFARTransformerOutput, Tuple]:
+    ) -> Union[
+        Transformer2DModelOutput,
+        AnyFlowFARTransformerOutput,
+        Tuple[torch.Tensor],
+        Tuple[Optional[torch.Tensor], List[Dict[str, torch.Tensor]]],
+    ]:
         """
         FAR causal forward pass. Dispatches to one of three internal paths:
 
@@ -1285,7 +1290,7 @@ def _forward_inference(
         return_dict: bool = True,
         kv_cache=None,
         kv_cache_flag=None,
-    ) -> Union[torch.Tensor, Dict[str, torch.Tensor]]:
+    ) -> Union[AnyFlowFARTransformerOutput, Tuple[torch.Tensor, List[Dict[str, torch.Tensor]]]]:
         hidden_states = hidden_states.permute(0, 2, 1, 3, 4)
 
         batch_size, num_channels, num_frames, height, width = hidden_states.shape
@@ -1405,7 +1410,7 @@ def _forward_cache(
         clean_timestep=None,
         kv_cache=None,
         kv_cache_flag=None,
-    ) -> Union[torch.Tensor, Dict[str, torch.Tensor]]:
+    ) -> Union[AnyFlowFARTransformerOutput, Tuple[None, List[Dict[str, torch.Tensor]]]]:
         hidden_states = hidden_states.permute(0, 2, 1, 3, 4)
         if clean_hidden_states is not None:
             clean_hidden_states = clean_hidden_states.permute(0, 2, 1, 3, 4)
@@ -1510,7 +1515,7 @@ def _forward_train(
         return_dict: bool = True,
         clean_hidden_states=None,
         clean_timestep=None,
-    ) -> Union[torch.Tensor, Dict[str, torch.Tensor]]:
+    ) -> Union[Transformer2DModelOutput, Tuple[torch.Tensor]]:
         hidden_states = hidden_states.permute(0, 2, 1, 3, 4)
         if clean_hidden_states is not None:
             clean_hidden_states = clean_hidden_states.permute(0, 2, 1, 3, 4)
diff --git a/src/diffusers/models/transformers/transformer_bria.py b/src/diffusers/models/transformers/transformer_bria.py
index a97b5b772552..7b3652f30f9a 100644
--- a/src/diffusers/models/transformers/transformer_bria.py
+++ b/src/diffusers/models/transformers/transformer_bria.py
@@ -133,7 +133,7 @@ def __call__(
         encoder_hidden_states: torch.Tensor = None,
         attention_mask: torch.Tensor | None = None,
         image_rotary_emb: torch.Tensor | None = None,
-    ) -> torch.Tensor:
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
         query, key, value, encoder_query, encoder_key, encoder_value = _get_qkv_projections(
             attn, hidden_states, encoder_hidden_states
         )
@@ -252,7 +252,7 @@ def forward(
         attention_mask: torch.Tensor | None = None,
         image_rotary_emb: torch.Tensor | None = None,
         **kwargs,
-    ) -> torch.Tensor:
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
         attn_parameters = set(inspect.signature(self.processor.__call__).parameters.keys())
         quiet_attn_parameters = {"ip_adapter_masks", "ip_hidden_states"}
         unused_kwargs = [k for k, _ in kwargs.items() if k not in attn_parameters and k not in quiet_attn_parameters]
@@ -271,7 +271,7 @@ def __init__(self, theta: int, axes_dim: list[int]):
         self.theta = theta
         self.axes_dim = axes_dim
 
-    def forward(self, ids: torch.Tensor) -> torch.Tensor:
+    def forward(self, ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
         n_axes = ids.shape[-1]
         cos_out = []
         sin_out = []
@@ -338,7 +338,7 @@ def __init__(self, theta: int, axes_dim: list[int]):
         self.theta = theta
         self.axes_dim = axes_dim
 
-    def forward(self, ids: torch.Tensor) -> torch.Tensor:
+    def forward(self, ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
         n_axes = ids.shape[-1]
         cos_out = []
         sin_out = []
diff --git a/src/diffusers/models/transformers/transformer_bria_fibo.py b/src/diffusers/models/transformers/transformer_bria_fibo.py
index 3dff8bcf55c7..893b4e8ca525 100644
--- a/src/diffusers/models/transformers/transformer_bria_fibo.py
+++ b/src/diffusers/models/transformers/transformer_bria_fibo.py
@@ -79,7 +79,7 @@ def __call__(
         encoder_hidden_states: torch.Tensor = None,
         attention_mask: torch.Tensor | None = None,
         image_rotary_emb: torch.Tensor | None = None,
-    ) -> torch.Tensor:
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
         query, key, value, encoder_query, encoder_key, encoder_value = _get_qkv_projections(
             attn, hidden_states, encoder_hidden_states
         )
@@ -197,7 +197,7 @@ def forward(
         attention_mask: torch.Tensor | None = None,
         image_rotary_emb: torch.Tensor | None = None,
         **kwargs,
-    ) -> torch.Tensor:
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
         attn_parameters = set(inspect.signature(self.processor.__call__).parameters.keys())
         quiet_attn_parameters = {"ip_adapter_masks", "ip_hidden_states"}
         unused_kwargs = [k for k, _ in kwargs.items() if k not in attn_parameters and k not in quiet_attn_parameters]
@@ -216,7 +216,7 @@ def __init__(self, theta: int, axes_dim: list[int]):
         self.theta = theta
         self.axes_dim = axes_dim
 
-    def forward(self, ids: torch.Tensor) -> torch.Tensor:
+    def forward(self, ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
         n_axes = ids.shape[-1]
         cos_out = []
         sin_out = []
@@ -520,7 +520,7 @@ def forward(
         guidance: torch.Tensor = None,
         joint_attention_kwargs: dict[str, Any] | None = None,
         return_dict: bool = True,
-    ) -> torch.FloatTensor | Transformer2DModelOutput:
+    ) -> Transformer2DModelOutput | tuple[torch.Tensor]:
         """
 
         Args:
diff --git a/src/diffusers/models/transformers/transformer_chroma.py b/src/diffusers/models/transformers/transformer_chroma.py
index 81a3dadc0998..103febb7b7d9 100644
--- a/src/diffusers/models/transformers/transformer_chroma.py
+++ b/src/diffusers/models/transformers/transformer_chroma.py
@@ -476,7 +476,7 @@ def forward(
         controlnet_single_block_samples=None,
         return_dict: bool = True,
         controlnet_blocks_repeat: bool = False,
-    ) -> torch.Tensor | Transformer2DModelOutput:
+    ) -> Transformer2DModelOutput | tuple[torch.Tensor]:
         """
         The [`FluxTransformer2DModel`] forward method.
 
diff --git a/src/diffusers/models/transformers/transformer_chronoedit.py b/src/diffusers/models/transformers/transformer_chronoedit.py
index c676a437cd16..326eb0048bbb 100644
--- a/src/diffusers/models/transformers/transformer_chronoedit.py
+++ b/src/diffusers/models/transformers/transformer_chronoedit.py
@@ -396,7 +396,7 @@ def __init__(
         self.register_buffer("freqs_cos", torch.cat(freqs_cos, dim=1), persistent=False)
         self.register_buffer("freqs_sin", torch.cat(freqs_sin, dim=1), persistent=False)
 
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+    def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
         batch_size, num_channels, num_frames, height, width = hidden_states.shape
         p_t, p_h, p_w = self.patch_size
         ppf, pph, ppw = num_frames // p_t, height // p_h, width // p_w
@@ -650,7 +650,7 @@ def forward(
         encoder_hidden_states_image: torch.Tensor | None = None,
         return_dict: bool = True,
         attention_kwargs: dict[str, Any] | None = None,
-    ) -> torch.Tensor | dict[str, torch.Tensor]:
+    ) -> Transformer2DModelOutput | tuple[torch.Tensor]:
         """
         The [`ChronoEditTransformer3DModel`] forward method.
 
diff --git a/src/diffusers/models/transformers/transformer_cogview4.py b/src/diffusers/models/transformers/transformer_cogview4.py
index 2856fffd2a63..6de561ebf47e 100644
--- a/src/diffusers/models/transformers/transformer_cogview4.py
+++ b/src/diffusers/models/transformers/transformer_cogview4.py
@@ -48,7 +48,9 @@ def __init__(
         self.proj = nn.Linear(in_channels * patch_size**2, hidden_size)
         self.text_proj = nn.Linear(text_hidden_size, hidden_size)
 
-    def forward(self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor) -> torch.Tensor:
+    def forward(
+        self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         batch_size, channel, height, width = hidden_states.shape
         post_patch_height = height // self.patch_size
         post_patch_width = width // self.patch_size
diff --git a/src/diffusers/models/transformers/transformer_cosmos.py b/src/diffusers/models/transformers/transformer_cosmos.py
index d901bb5809de..04cefac49f9a 100644
--- a/src/diffusers/models/transformers/transformer_cosmos.py
+++ b/src/diffusers/models/transformers/transformer_cosmos.py
@@ -74,7 +74,7 @@ def __init__(self, embedding_dim: int, condition_dim: int) -> None:
         self.t_embedder = CosmosTimestepEmbedding(embedding_dim, condition_dim)
         self.norm = RMSNorm(embedding_dim, eps=1e-6, elementwise_affine=True)
 
-    def forward(self, hidden_states: torch.Tensor, timestep: torch.LongTensor) -> torch.Tensor:
+    def forward(self, hidden_states: torch.Tensor, timestep: torch.LongTensor) -> tuple[torch.Tensor, torch.Tensor]:
         timesteps_proj = self.time_proj(timestep).type_as(hidden_states)
         temb = self.t_embedder(timesteps_proj)
         embedded_timestep = self.norm(timesteps_proj)
@@ -130,7 +130,7 @@ def forward(
         hidden_states: torch.Tensor,
         embedded_timestep: torch.Tensor,
         temb: torch.Tensor | None = None,
-    ) -> torch.Tensor:
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         embedded_timestep = self.activation(embedded_timestep)
         embedded_timestep = self.linear_1(embedded_timestep)
         embedded_timestep = self.linear_2(embedded_timestep)
diff --git a/src/diffusers/models/transformers/transformer_easyanimate.py b/src/diffusers/models/transformers/transformer_easyanimate.py
index 24c874ad40ef..d1983ae1543a 100755
--- a/src/diffusers/models/transformers/transformer_easyanimate.py
+++ b/src/diffusers/models/transformers/transformer_easyanimate.py
@@ -128,7 +128,7 @@ def __call__(
         encoder_hidden_states: torch.Tensor,
         attention_mask: torch.Tensor | None = None,
         image_rotary_emb: torch.Tensor | None = None,
-    ) -> torch.Tensor:
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         if attn.add_q_proj is None and encoder_hidden_states is not None:
             hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
 
diff --git a/src/diffusers/models/transformers/transformer_flux.py b/src/diffusers/models/transformers/transformer_flux.py
index c8ba2d517988..a0fe19bcd517 100644
--- a/src/diffusers/models/transformers/transformer_flux.py
+++ b/src/diffusers/models/transformers/transformer_flux.py
@@ -88,7 +88,7 @@ def __call__(
         encoder_hidden_states: torch.Tensor = None,
         attention_mask: torch.Tensor | None = None,
         image_rotary_emb: torch.Tensor | None = None,
-    ) -> torch.Tensor:
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
         query, key, value, encoder_query, encoder_key, encoder_value = _get_qkv_projections(
             attn, hidden_states, encoder_hidden_states
         )
@@ -192,7 +192,7 @@ def __call__(
         image_rotary_emb: torch.Tensor | None = None,
         ip_hidden_states: list[torch.Tensor] | None = None,
         ip_adapter_masks: torch.Tensor | None = None,
-    ) -> torch.Tensor:
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         batch_size = hidden_states.shape[0]
 
         query, key, value, encoder_query, encoder_key, encoder_value = _get_qkv_projections(
@@ -343,7 +343,7 @@ def forward(
         attention_mask: torch.Tensor | None = None,
         image_rotary_emb: torch.Tensor | None = None,
         **kwargs,
-    ) -> torch.Tensor:
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor] | tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         attn_parameters = set(inspect.signature(self.processor.__call__).parameters.keys())
         quiet_attn_parameters = {"ip_adapter_masks", "ip_hidden_states"}
         unused_kwargs = [k for k, _ in kwargs.items() if k not in attn_parameters and k not in quiet_attn_parameters]
@@ -504,7 +504,7 @@ def __init__(self, theta: int, axes_dim: list[int]):
         self.theta = theta
         self.axes_dim = axes_dim
 
-    def forward(self, ids: torch.Tensor) -> torch.Tensor:
+    def forward(self, ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
         n_axes = ids.shape[-1]
         cos_out = []
         sin_out = []
@@ -682,7 +682,7 @@ def forward(
         controlnet_single_block_samples=None,
         return_dict: bool = True,
         controlnet_blocks_repeat: bool = False,
-    ) -> torch.Tensor | Transformer2DModelOutput:
+    ) -> Transformer2DModelOutput | tuple[torch.Tensor]:
         """
         The [`FluxTransformer2DModel`] forward method.
 
diff --git a/src/diffusers/models/transformers/transformer_flux2.py b/src/diffusers/models/transformers/transformer_flux2.py
index be56ecb935a1..7bbed714df3d 100644
--- a/src/diffusers/models/transformers/transformer_flux2.py
+++ b/src/diffusers/models/transformers/transformer_flux2.py
@@ -339,7 +339,7 @@ def __call__(
         encoder_hidden_states: torch.Tensor = None,
         attention_mask: torch.Tensor | None = None,
         image_rotary_emb: torch.Tensor | None = None,
-    ) -> torch.Tensor:
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
         query, key, value, encoder_query, encoder_key, encoder_value = _get_qkv_projections(
             attn, hidden_states, encoder_hidden_states
         )
@@ -421,7 +421,7 @@ def __call__(
         kv_cache: Flux2KVLayerCache | None = None,
         kv_cache_mode: str | None = None,
         num_ref_tokens: int = 0,
-    ) -> torch.Tensor:
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
         query, key, value, encoder_query, encoder_key, encoder_value = _get_qkv_projections(
             attn, hidden_states, encoder_hidden_states
         )
@@ -558,7 +558,7 @@ def forward(
         attention_mask: torch.Tensor | None = None,
         image_rotary_emb: torch.Tensor | None = None,
         **kwargs,
-    ) -> torch.Tensor:
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
         attn_parameters = set(inspect.signature(self.processor.__call__).parameters.keys())
         unused_kwargs = [k for k, _ in kwargs.items() if k not in attn_parameters]
         if len(unused_kwargs) > 0:
@@ -975,7 +975,7 @@ def __init__(self, theta: int, axes_dim: list[int]):
         self.theta = theta
         self.axes_dim = axes_dim
 
-    def forward(self, ids: torch.Tensor) -> torch.Tensor:
+    def forward(self, ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
         # Expected ids shape: [S, len(self.axes_dim)]
         cos_out = []
         sin_out = []
@@ -1237,7 +1237,7 @@ def forward(
         kv_cache_mode: str | None = None,
         num_ref_tokens: int = 0,
         ref_fixed_timestep: float = 0.0,
-    ) -> torch.Tensor | Flux2Transformer2DModelOutput:
+    ) -> Flux2Transformer2DModelOutput | tuple[torch.Tensor] | tuple[torch.Tensor, Flux2KVCache]:
         """
         The [`Flux2Transformer2DModel`] forward method.
 
diff --git a/src/diffusers/models/transformers/transformer_flux3.py b/src/diffusers/models/transformers/transformer_flux3.py
index 2c5fe7a40b5d..6a6021fe06c4 100644
--- a/src/diffusers/models/transformers/transformer_flux3.py
+++ b/src/diffusers/models/transformers/transformer_flux3.py
@@ -300,7 +300,7 @@ def __init__(self, theta: int, axes_dim: list[int]):
         self.theta = theta
         self.axes_dim = axes_dim
 
-    def forward(self, ids: torch.Tensor) -> torch.Tensor:
+    def forward(self, ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
         # Expected ids shape: [S, len(self.axes_dim)]
         cos_out = []
         sin_out = []
diff --git a/src/diffusers/models/transformers/transformer_helios.py b/src/diffusers/models/transformers/transformer_helios.py
index 6733004bf1e5..d72b72dae5ff 100644
--- a/src/diffusers/models/transformers/transformer_helios.py
+++ b/src/diffusers/models/transformers/transformer_helios.py
@@ -674,7 +674,7 @@ def forward(
         latents_history_long=None,
         return_dict: bool = True,
         attention_kwargs: dict[str, Any] | None = None,
-    ) -> torch.Tensor | dict[str, torch.Tensor]:
+    ) -> Transformer2DModelOutput | tuple[torch.Tensor]:
         """
         The [`HeliosTransformer3DModel`] forward method.
 
diff --git a/src/diffusers/models/transformers/transformer_hidream_image.py b/src/diffusers/models/transformers/transformer_hidream_image.py
index afb1f8b8edb9..9c6297da45e7 100644
--- a/src/diffusers/models/transformers/transformer_hidream_image.py
+++ b/src/diffusers/models/transformers/transformer_hidream_image.py
@@ -185,7 +185,7 @@ def forward(
         hidden_states_masks: torch.Tensor = None,
         norm_encoder_hidden_states: torch.Tensor = None,
         image_rotary_emb: torch.Tensor = None,
-    ) -> torch.Tensor:
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
         return self.processor(
             self,
             hidden_states=norm_hidden_states,
@@ -207,7 +207,7 @@ def __call__(
         image_rotary_emb: torch.Tensor = None,
         *args,
         **kwargs,
-    ) -> torch.Tensor:
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
         dtype = hidden_states.dtype
         batch_size = hidden_states.shape[0]
 
diff --git a/src/diffusers/models/transformers/transformer_hunyuan_video.py b/src/diffusers/models/transformers/transformer_hunyuan_video.py
index 3730cc8ffa56..e25624c2e711 100644
--- a/src/diffusers/models/transformers/transformer_hunyuan_video.py
+++ b/src/diffusers/models/transformers/transformer_hunyuan_video.py
@@ -59,7 +59,7 @@ def __call__(
         encoder_hidden_states: torch.Tensor | None = None,
         attention_mask: torch.Tensor | None = None,
         image_rotary_emb: torch.Tensor | None = None,
-    ) -> torch.Tensor:
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         if attn.add_q_proj is None and encoder_hidden_states is not None:
             hidden_states = torch.cat([hidden_states, encoder_hidden_states], dim=1)
 
@@ -484,7 +484,7 @@ def __init__(self, patch_size: int, patch_size_t: int, rope_dim: list[int], thet
         self.rope_dim = rope_dim
         self.theta = theta
 
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+    def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
         batch_size, num_channels, num_frames, height, width = hidden_states.shape
         rope_sizes = [num_frames // self.patch_size_t, height // self.patch_size, width // self.patch_size]
 
diff --git a/src/diffusers/models/transformers/transformer_hunyuan_video15.py b/src/diffusers/models/transformers/transformer_hunyuan_video15.py
index 64c18e541d7c..5a5d67296728 100644
--- a/src/diffusers/models/transformers/transformer_hunyuan_video15.py
+++ b/src/diffusers/models/transformers/transformer_hunyuan_video15.py
@@ -58,7 +58,7 @@ def __call__(
         encoder_hidden_states: torch.Tensor | None = None,
         attention_mask: torch.Tensor | None = None,
         image_rotary_emb: torch.Tensor | None = None,
-    ) -> torch.Tensor:
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         # 1. QKV projections
         query = attn.to_q(hidden_states)
         key = attn.to_k(hidden_states)
@@ -367,7 +367,7 @@ def __init__(self, patch_size: int, patch_size_t: int, rope_dim: list[int], thet
         self.rope_dim = rope_dim
         self.theta = theta
 
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+    def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
         batch_size, num_channels, num_frames, height, width = hidden_states.shape
         rope_sizes = [num_frames // self.patch_size_t, height // self.patch_size, width // self.patch_size]
 
diff --git a/src/diffusers/models/transformers/transformer_hunyuanimage.py b/src/diffusers/models/transformers/transformer_hunyuanimage.py
index dd2176a4096f..17833bf7a72e 100644
--- a/src/diffusers/models/transformers/transformer_hunyuanimage.py
+++ b/src/diffusers/models/transformers/transformer_hunyuanimage.py
@@ -60,7 +60,7 @@ def __call__(
         encoder_hidden_states: torch.Tensor | None = None,
         attention_mask: torch.Tensor | None = None,
         image_rotary_emb: torch.Tensor | None = None,
-    ) -> torch.Tensor:
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         if attn.add_q_proj is None and encoder_hidden_states is not None:
             hidden_states = torch.cat([hidden_states, encoder_hidden_states], dim=1)
 
@@ -428,7 +428,7 @@ def __init__(self, patch_size: tuple | list[int], rope_dim: tuple | list[int], t
         self.rope_dim = rope_dim
         self.theta = theta
 
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+    def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
         if hidden_states.ndim == 5:
             _, _, frame, height, width = hidden_states.shape
             patch_size_frame, patch_size_height, patch_size_width = self.patch_size
@@ -498,7 +498,7 @@ def forward(
         image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
         *args,
         **kwargs,
-    ) -> 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)
 
@@ -753,7 +753,7 @@ def forward(
         guidance: torch.Tensor | None = None,
         attention_kwargs: dict[str, Any] | None = None,
         return_dict: bool = True,
-    ) -> torch.Tensor | dict[str, torch.Tensor]:
+    ) -> Transformer2DModelOutput | tuple[torch.Tensor]:
         """
         The [`HunyuanImageTransformer2DModel`] forward method.
 
diff --git a/src/diffusers/models/transformers/transformer_joyimage_edit_plus.py b/src/diffusers/models/transformers/transformer_joyimage_edit_plus.py
index a81027b5f40b..135e8e759ce0 100644
--- a/src/diffusers/models/transformers/transformer_joyimage_edit_plus.py
+++ b/src/diffusers/models/transformers/transformer_joyimage_edit_plus.py
@@ -449,7 +449,7 @@ def forward(
         encoder_hidden_states_mask: torch.Tensor | None = None,
         shape_list: list[list[tuple[int, int, int]]] = None,
         return_dict: bool = True,
-    ) -> torch.Tensor | tuple:
+    ) -> Transformer2DModelOutput | tuple[torch.Tensor]:
         """
         Args:
             hidden_states: [B, max_patches, C, pt, ph, pw] - patchified latent input.
diff --git a/src/diffusers/models/transformers/transformer_krea2.py b/src/diffusers/models/transformers/transformer_krea2.py
index 980d942831e6..575008b283ed 100644
--- a/src/diffusers/models/transformers/transformer_krea2.py
+++ b/src/diffusers/models/transformers/transformer_krea2.py
@@ -314,7 +314,7 @@ def __init__(self, theta: int, axes_dim: list[int]):
         self.theta = theta
         self.axes_dim = axes_dim
 
-    def forward(self, ids: torch.Tensor) -> torch.Tensor:
+    def forward(self, ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
         n_axes = ids.shape[-1]
         cos_out = []
         sin_out = []
diff --git a/src/diffusers/models/transformers/transformer_longcat_image.py b/src/diffusers/models/transformers/transformer_longcat_image.py
index 55d4771090e9..d63fdee8d92b 100644
--- a/src/diffusers/models/transformers/transformer_longcat_image.py
+++ b/src/diffusers/models/transformers/transformer_longcat_image.py
@@ -80,7 +80,7 @@ def __call__(
         encoder_hidden_states: torch.Tensor = None,
         attention_mask: torch.Tensor | None = None,
         image_rotary_emb: torch.Tensor | None = None,
-    ) -> torch.Tensor:
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
         query, key, value, encoder_query, encoder_key, encoder_value = _get_qkv_projections(
             attn, hidden_states, encoder_hidden_states
         )
@@ -199,7 +199,7 @@ def forward(
         attention_mask: torch.Tensor | None = None,
         image_rotary_emb: torch.Tensor | None = None,
         **kwargs,
-    ) -> torch.Tensor:
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
         attn_parameters = set(inspect.signature(self.processor.__call__).parameters.keys())
         quiet_attn_parameters = {"ip_adapter_masks", "ip_hidden_states"}
         unused_kwargs = [k for k, _ in kwargs.items() if k not in attn_parameters and k not in quiet_attn_parameters]
@@ -356,7 +356,7 @@ def __init__(self, theta: int, axes_dim: list[int]):
         self.theta = theta
         self.axes_dim = axes_dim
 
-    def forward(self, ids: torch.Tensor) -> torch.Tensor:
+    def forward(self, ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
         n_axes = ids.shape[-1]
         cos_out = []
         sin_out = []
@@ -470,7 +470,7 @@ def forward(
         txt_ids: torch.Tensor = None,
         guidance: torch.Tensor = None,
         return_dict: bool = True,
-    ) -> torch.FloatTensor | Transformer2DModelOutput:
+    ) -> Transformer2DModelOutput | tuple[torch.Tensor]:
         """
         The forward method.
 
diff --git a/src/diffusers/models/transformers/transformer_ltx.py b/src/diffusers/models/transformers/transformer_ltx.py
index c33e0f6141fc..b6bb692a1a86 100644
--- a/src/diffusers/models/transformers/transformer_ltx.py
+++ b/src/diffusers/models/transformers/transformer_ltx.py
@@ -505,7 +505,7 @@ def forward(
         video_coords: torch.Tensor | None = None,
         attention_kwargs: dict[str, Any] | None = None,
         return_dict: bool = True,
-    ) -> torch.Tensor:
+    ) -> Transformer2DModelOutput | tuple[torch.Tensor]:
         """
         The [`LTXVideoTransformer3DModel`] forward method.
 
diff --git a/src/diffusers/models/transformers/transformer_ltx2.py b/src/diffusers/models/transformers/transformer_ltx2.py
index e5462e4b047a..d85db610f0f7 100644
--- a/src/diffusers/models/transformers/transformer_ltx2.py
+++ b/src/diffusers/models/transformers/transformer_ltx2.py
@@ -622,7 +622,7 @@ def forward(
         use_v2a_cross_attention: bool = True,
         perturbation_mask: torch.Tensor | None = None,
         all_perturbed: bool | None = None,
-    ) -> torch.Tensor:
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         batch_size = hidden_states.size(0)
 
         # 1. Video and Audio Self-Attention
@@ -1390,7 +1390,7 @@ def forward(
         video_self_attention_mask: torch.Tensor | None = None,
         video_keyframes_mask: torch.Tensor | None = None,
         return_dict: bool = True,
-    ) -> torch.Tensor:
+    ) -> AudioVisualModelOutput | tuple[torch.Tensor, torch.Tensor]:
         """
         Forward pass for LTX-2.0 audiovisual video transformer.
 
diff --git a/src/diffusers/models/transformers/transformer_lumina2.py b/src/diffusers/models/transformers/transformer_lumina2.py
index 6d51b2925ef0..874b3349a4c7 100644
--- a/src/diffusers/models/transformers/transformer_lumina2.py
+++ b/src/diffusers/models/transformers/transformer_lumina2.py
@@ -466,7 +466,7 @@ def forward(
         encoder_attention_mask: torch.Tensor,
         attention_kwargs: dict[str, Any] | None = None,
         return_dict: bool = True,
-    ) -> torch.Tensor | Transformer2DModelOutput:
+    ) -> Transformer2DModelOutput | tuple[torch.Tensor]:
         """
         The [`Lumina2Transformer2DModel`] forward method.
 
diff --git a/src/diffusers/models/transformers/transformer_mochi.py b/src/diffusers/models/transformers/transformer_mochi.py
index 1e544890a4b8..2bfde45db74e 100644
--- a/src/diffusers/models/transformers/transformer_mochi.py
+++ b/src/diffusers/models/transformers/transformer_mochi.py
@@ -282,7 +282,7 @@ def _get_positions(
         positions = torch.stack([grid_t, grid_h, grid_w], dim=-1).view(-1, 3)
         return positions
 
-    def _create_rope(self, freqs: torch.Tensor, pos: torch.Tensor) -> torch.Tensor:
+    def _create_rope(self, freqs: torch.Tensor, pos: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
         with torch.autocast(freqs.device.type, torch.float32):
             # Always run ROPE freqs computation in FP32
             freqs = torch.einsum("nd,dhf->nhf", pos.to(torch.float32), freqs.to(torch.float32))
@@ -413,7 +413,7 @@ def forward(
         encoder_attention_mask: torch.Tensor,
         attention_kwargs: dict[str, Any] | None = None,
         return_dict: bool = True,
-    ) -> torch.Tensor:
+    ) -> Transformer2DModelOutput | tuple[torch.Tensor]:
         """
         The [`MochiTransformer3DModel`] forward method.
 
diff --git a/src/diffusers/models/transformers/transformer_motif_video.py b/src/diffusers/models/transformers/transformer_motif_video.py
index fb3ff0666f95..4493432e7ff4 100644
--- a/src/diffusers/models/transformers/transformer_motif_video.py
+++ b/src/diffusers/models/transformers/transformer_motif_video.py
@@ -123,7 +123,7 @@ def __call__(
         encoder_hidden_states: Optional[torch.Tensor] = None,
         attention_mask: Optional[torch.Tensor] = None,
         image_rotary_emb: Optional[torch.Tensor] = None,
-    ) -> torch.Tensor:
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
         # Concatenate hidden states with encoder hidden states for joint attention if needed
         if attn.add_q_proj is None and encoder_hidden_states is not None:
             hidden_states = torch.cat([hidden_states, encoder_hidden_states], dim=1)
@@ -375,7 +375,7 @@ def forward(
         attention_mask: torch.Tensor | None = None,
         image_rotary_emb: torch.Tensor | None = None,
         **kwargs,
-    ) -> torch.Tensor:
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
         attn_parameters = set(inspect.signature(self.processor.__call__).parameters.keys())
         unused_kwargs = [k for k, _ in kwargs.items() if k not in attn_parameters]
         if len(unused_kwargs) > 0:
@@ -468,7 +468,7 @@ def __init__(
         self.rope_dim = rope_dim
         self.theta = theta
 
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+    def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
         batch_size, num_channels, num_frames, height, width = hidden_states.shape
         rope_sizes = [
             num_frames // self.patch_size_t,
@@ -573,7 +573,7 @@ def forward(
         attention_mask: Optional[torch.Tensor] = None,
         image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
         image_embed_seq_len: int = 0,
-    ) -> torch.Tensor:
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         encoder_seq_length = encoder_hidden_states.shape[1]
         hidden_states = torch.cat([hidden_states, encoder_hidden_states], dim=1)
 
@@ -894,7 +894,7 @@ def forward(
         image_embeds: torch.Tensor | None = None,
         attention_kwargs: Optional[Dict[str, Any]] = None,
         return_dict: bool = True,
-    ) -> Union[torch.Tensor, Dict[str, torch.Tensor]]:
+    ) -> Union[Transformer2DModelOutput, tuple[torch.Tensor]]:
         """
         Forward pass of the MotifVideoTransformer3DModel.
 
diff --git a/src/diffusers/models/transformers/transformer_nucleusmoe_image.py b/src/diffusers/models/transformers/transformer_nucleusmoe_image.py
index 5488c02c579f..25b9617f7ed0 100644
--- a/src/diffusers/models/transformers/transformer_nucleusmoe_image.py
+++ b/src/diffusers/models/transformers/transformer_nucleusmoe_image.py
@@ -841,7 +841,7 @@ def forward(
         timestep: torch.LongTensor = None,
         attention_kwargs: dict[str, Any] | None = None,
         return_dict: bool = True,
-    ) -> torch.Tensor | Transformer2DModelOutput:
+    ) -> Transformer2DModelOutput | tuple[torch.Tensor]:
         """
         The [`NucleusMoEImageTransformer2DModel`] forward method.
 
diff --git a/src/diffusers/models/transformers/transformer_ovis_image.py b/src/diffusers/models/transformers/transformer_ovis_image.py
index 44723bc44fd0..f7e5d94948cc 100644
--- a/src/diffusers/models/transformers/transformer_ovis_image.py
+++ b/src/diffusers/models/transformers/transformer_ovis_image.py
@@ -80,7 +80,7 @@ def __call__(
         encoder_hidden_states: torch.Tensor = None,
         attention_mask: torch.Tensor | None = None,
         image_rotary_emb: torch.Tensor | None = None,
-    ) -> torch.Tensor:
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
         query, key, value, encoder_query, encoder_key, encoder_value = _get_qkv_projections(
             attn, hidden_states, encoder_hidden_states
         )
@@ -199,7 +199,7 @@ def forward(
         attention_mask: torch.Tensor | None = None,
         image_rotary_emb: torch.Tensor | None = None,
         **kwargs,
-    ) -> torch.Tensor:
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
         attn_parameters = set(inspect.signature(self.processor.__call__).parameters.keys())
         quiet_attn_parameters = {"ip_adapter_masks", "ip_hidden_states"}
         unused_kwargs = [k for k, _ in kwargs.items() if k not in attn_parameters and k not in quiet_attn_parameters]
@@ -359,7 +359,7 @@ def __init__(self, theta: int, axes_dim: list[int]):
         self.theta = theta
         self.axes_dim = axes_dim
 
-    def forward(self, ids: torch.Tensor) -> torch.Tensor:
+    def forward(self, ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
         n_axes = ids.shape[-1]
         cos_out = []
         sin_out = []
@@ -483,7 +483,7 @@ def forward(
         txt_ids: torch.Tensor = None,
         joint_attention_kwargs: dict[str, Any] | None = None,
         return_dict: bool = True,
-    ) -> torch.Tensor | Transformer2DModelOutput:
+    ) -> Transformer2DModelOutput | tuple[torch.Tensor]:
         """
         The [`OvisImageTransformer2DModel`] forward method.
 
diff --git a/src/diffusers/models/transformers/transformer_qwenimage.py b/src/diffusers/models/transformers/transformer_qwenimage.py
index 29a98c0e9e57..590ad7aca395 100644
--- a/src/diffusers/models/transformers/transformer_qwenimage.py
+++ b/src/diffusers/models/transformers/transformer_qwenimage.py
@@ -493,7 +493,7 @@ def __call__(
         encoder_hidden_states_mask: torch.FloatTensor = None,
         attention_mask: torch.FloatTensor | None = None,
         image_rotary_emb: torch.Tensor | None = None,
-    ) -> torch.FloatTensor:
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         if encoder_hidden_states is None:
             raise ValueError("QwenDoubleStreamAttnProcessor2_0 requires encoder_hidden_states (text stream)")
 
@@ -878,7 +878,7 @@ def forward(
         controlnet_block_samples=None,
         additional_t_cond=None,
         return_dict: bool = True,
-    ) -> torch.Tensor | Transformer2DModelOutput:
+    ) -> Transformer2DModelOutput | tuple[torch.Tensor]:
         """
         The [`QwenTransformer2DModel`] forward method.
 
diff --git a/src/diffusers/models/transformers/transformer_qwenimage21.py b/src/diffusers/models/transformers/transformer_qwenimage21.py
index 1a9ef250c519..db871ae242f0 100644
--- a/src/diffusers/models/transformers/transformer_qwenimage21.py
+++ b/src/diffusers/models/transformers/transformer_qwenimage21.py
@@ -872,7 +872,7 @@ def forward(
         kv_cache: QwenImage21KVCache | None = None,
         kv_cache_mode: str | None = None,
         return_dict: bool = True,
-    ) -> torch.Tensor | Transformer2DModelOutput:
+    ) -> Transformer2DModelOutput | tuple[torch.Tensor]:
         r"""
         Args:
             hidden_states (`torch.Tensor` of shape `(batch_size, image_sequence_length, in_channels)`):
diff --git a/src/diffusers/models/transformers/transformer_sana_video.py b/src/diffusers/models/transformers/transformer_sana_video.py
index 84c1fa18a9bb..9499922d17a3 100644
--- a/src/diffusers/models/transformers/transformer_sana_video.py
+++ b/src/diffusers/models/transformers/transformer_sana_video.py
@@ -213,7 +213,7 @@ def __init__(
         self.register_buffer("freqs_cos", torch.cat(freqs_cos, dim=1), persistent=False)
         self.register_buffer("freqs_sin", torch.cat(freqs_sin, dim=1), persistent=False)
 
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+    def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
         batch_size, num_channels, num_frames, height, width = hidden_states.shape
         p_t, p_h, p_w = self.patch_size
         ppf, pph, ppw = num_frames // p_t, height // p_h, width // p_w
diff --git a/src/diffusers/models/transformers/transformer_sd3.py b/src/diffusers/models/transformers/transformer_sd3.py
index eaafe057fd82..e88fbb305855 100644
--- a/src/diffusers/models/transformers/transformer_sd3.py
+++ b/src/diffusers/models/transformers/transformer_sd3.py
@@ -256,7 +256,7 @@ def forward(
         joint_attention_kwargs: dict[str, Any] | None = None,
         return_dict: bool = True,
         skip_layers: list[int] | None = None,
-    ) -> torch.Tensor | Transformer2DModelOutput:
+    ) -> Transformer2DModelOutput | tuple[torch.Tensor]:
         """
         The [`SD3Transformer2DModel`] forward method.
 
diff --git a/src/diffusers/models/transformers/transformer_skyreels_v2.py b/src/diffusers/models/transformers/transformer_skyreels_v2.py
index 5687cfcf2964..7abce77aa8b9 100644
--- a/src/diffusers/models/transformers/transformer_skyreels_v2.py
+++ b/src/diffusers/models/transformers/transformer_skyreels_v2.py
@@ -411,7 +411,7 @@ def __init__(
         self.register_buffer("freqs_cos", torch.cat(freqs_cos, dim=1), persistent=False)
         self.register_buffer("freqs_sin", torch.cat(freqs_sin, dim=1), persistent=False)
 
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+    def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
         batch_size, num_channels, num_frames, height, width = hidden_states.shape
         p_t, p_h, p_w = self.patch_size
         ppf, pph, ppw = num_frames // p_t, height // p_h, width // p_w
@@ -641,7 +641,7 @@ def forward(
         fps: torch.Tensor | None = None,
         return_dict: bool = True,
         attention_kwargs: dict[str, Any] | None = None,
-    ) -> torch.Tensor | dict[str, torch.Tensor]:
+    ) -> Transformer2DModelOutput | tuple[torch.Tensor]:
         """
         The [`SkyReelsV2Transformer3DModel`] forward method.
 
diff --git a/src/diffusers/models/transformers/transformer_stable_audio3.py b/src/diffusers/models/transformers/transformer_stable_audio3.py
index aab80b744d8f..42b10220cb66 100644
--- a/src/diffusers/models/transformers/transformer_stable_audio3.py
+++ b/src/diffusers/models/transformers/transformer_stable_audio3.py
@@ -527,7 +527,7 @@ def forward(
         encoder_attention_mask: Optional[torch.Tensor] = None,
         local_add_cond: Optional[torch.Tensor] = None,
         return_dict: bool = True,
-    ) -> Union[StableAudio3DiTModelOutput, tuple]:
+    ) -> Union[StableAudio3DiTModelOutput, tuple[torch.Tensor]]:
         """
         Args:
             hidden_states (`torch.Tensor`): Noisy latent audio `(batch, io_channels, T)`.
diff --git a/src/diffusers/models/transformers/transformer_temporal.py b/src/diffusers/models/transformers/transformer_temporal.py
index 5141f07b8235..fff0031b5a7a 100644
--- a/src/diffusers/models/transformers/transformer_temporal.py
+++ b/src/diffusers/models/transformers/transformer_temporal.py
@@ -129,7 +129,7 @@ def forward(
         num_frames: int = 1,
         cross_attention_kwargs: dict[str, Any] | None = None,
         return_dict: bool = True,
-    ) -> TransformerTemporalModelOutput:
+    ) -> TransformerTemporalModelOutput | tuple[torch.Tensor]:
         """
         The [`TransformerTemporal`] forward method.
 
diff --git a/src/diffusers/models/transformers/transformer_wan.py b/src/diffusers/models/transformers/transformer_wan.py
index b9e0b75f46e8..7a1cc0b583e8 100644
--- a/src/diffusers/models/transformers/transformer_wan.py
+++ b/src/diffusers/models/transformers/transformer_wan.py
@@ -392,7 +392,7 @@ def __init__(
         self.register_buffer("freqs_cos", torch.cat(freqs_cos, dim=1), persistent=False)
         self.register_buffer("freqs_sin", torch.cat(freqs_sin, dim=1), persistent=False)
 
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+    def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
         batch_size, num_channels, num_frames, height, width = hidden_states.shape
         p_t, p_h, p_w = self.patch_size
         ppf, pph, ppw = num_frames // p_t, height // p_h, width // p_w
@@ -634,7 +634,7 @@ def forward(
         encoder_hidden_states_image: torch.Tensor | None = None,
         return_dict: bool = True,
         attention_kwargs: dict[str, Any] | None = None,
-    ) -> torch.Tensor | dict[str, torch.Tensor]:
+    ) -> Transformer2DModelOutput | tuple[torch.Tensor]:
         """
         The [`WanTransformer3DModel`] forward method.
 
diff --git a/src/diffusers/models/transformers/transformer_wan_animate.py b/src/diffusers/models/transformers/transformer_wan_animate.py
index 3e188ce1ec51..609a71a3fe70 100644
--- a/src/diffusers/models/transformers/transformer_wan_animate.py
+++ b/src/diffusers/models/transformers/transformer_wan_animate.py
@@ -871,7 +871,7 @@ def __init__(
         self.register_buffer("freqs_cos", torch.cat(freqs_cos, dim=1), persistent=False)
         self.register_buffer("freqs_sin", torch.cat(freqs_sin, dim=1), persistent=False)
 
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+    def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
         batch_size, num_channels, num_frames, height, width = hidden_states.shape
         p_t, p_h, p_w = self.patch_size
         ppf, pph, ppw = num_frames // p_t, height // p_h, width // p_w
@@ -1163,7 +1163,7 @@ def forward(
         motion_encode_batch_size: int | None = None,
         return_dict: bool = True,
         attention_kwargs: dict[str, Any] | None = None,
-    ) -> torch.Tensor | dict[str, torch.Tensor]:
+    ) -> Transformer2DModelOutput | tuple[torch.Tensor]:
         """
         Forward pass of Wan2.2-Animate transformer model.
 
diff --git a/src/diffusers/models/transformers/transformer_wan_vace.py b/src/diffusers/models/transformers/transformer_wan_vace.py
index 440417a80845..0bb2f32c9378 100644
--- a/src/diffusers/models/transformers/transformer_wan_vace.py
+++ b/src/diffusers/models/transformers/transformer_wan_vace.py
@@ -98,7 +98,7 @@ def forward(
         control_hidden_states: torch.Tensor,
         temb: torch.Tensor,
         rotary_emb: torch.Tensor,
-    ) -> torch.Tensor:
+    ) -> tuple[torch.Tensor | None, torch.Tensor]:
         if self.proj_in is not None:
             control_hidden_states = self.proj_in(control_hidden_states)
             control_hidden_states = control_hidden_states + hidden_states
@@ -274,7 +274,7 @@ def forward(
         control_hidden_states_scale: torch.Tensor = None,
         return_dict: bool = True,
         attention_kwargs: dict[str, Any] | None = None,
-    ) -> torch.Tensor | dict[str, torch.Tensor]:
+    ) -> Transformer2DModelOutput | tuple[torch.Tensor]:
         """
         The [`WanVACETransformer3DModel`] forward method.
 
diff --git a/src/diffusers/models/unets/unet_1d.py b/src/diffusers/models/unets/unet_1d.py
index 959e82e9d7cd..c3e2d8312f26 100644
--- a/src/diffusers/models/unets/unet_1d.py
+++ b/src/diffusers/models/unets/unet_1d.py
@@ -207,7 +207,7 @@ def forward(
         sample: torch.Tensor,
         timestep: torch.Tensor | float | int,
         return_dict: bool = True,
-    ) -> UNet1DOutput | tuple:
+    ) -> UNet1DOutput | tuple[torch.Tensor]:
         r"""
         The [`UNet1DModel`] forward method.
 
diff --git a/src/diffusers/models/unets/unet_1d_blocks.py b/src/diffusers/models/unets/unet_1d_blocks.py
index f4d5c7d93a12..e0d8c0e779e8 100644
--- a/src/diffusers/models/unets/unet_1d_blocks.py
+++ b/src/diffusers/models/unets/unet_1d_blocks.py
@@ -65,7 +65,9 @@ def __init__(
         if add_downsample:
             self.downsample = Downsample1D(out_channels, use_conv=True, padding=1)
 
-    def forward(self, hidden_states: torch.Tensor, temb: torch.Tensor | None = None) -> torch.Tensor:
+    def forward(
+        self, hidden_states: torch.Tensor, temb: torch.Tensor | None = None
+    ) -> tuple[torch.Tensor, tuple[torch.Tensor]]:
         output_states = ()
 
         hidden_states = self.resnets[0](hidden_states, temb)
@@ -465,7 +467,9 @@ def __init__(self, out_channels: int, in_channels: int, mid_channels: int | None
         self.attentions = nn.ModuleList(attentions)
         self.resnets = nn.ModuleList(resnets)
 
-    def forward(self, hidden_states: torch.Tensor, temb: torch.Tensor | None = None) -> torch.Tensor:
+    def forward(
+        self, hidden_states: torch.Tensor, temb: torch.Tensor | None = None
+    ) -> tuple[torch.Tensor, tuple[torch.Tensor]]:
         hidden_states = self.down(hidden_states)
 
         for resnet, attn in zip(self.resnets, self.attentions):
@@ -489,7 +493,9 @@ def __init__(self, out_channels: int, in_channels: int, mid_channels: int | None
 
         self.resnets = nn.ModuleList(resnets)
 
-    def forward(self, hidden_states: torch.Tensor, temb: torch.Tensor | None = None) -> torch.Tensor:
+    def forward(
+        self, hidden_states: torch.Tensor, temb: torch.Tensor | None = None
+    ) -> tuple[torch.Tensor, tuple[torch.Tensor]]:
         hidden_states = self.down(hidden_states)
 
         for resnet in self.resnets:
@@ -511,7 +517,9 @@ def __init__(self, out_channels: int, in_channels: int, mid_channels: int | None
 
         self.resnets = nn.ModuleList(resnets)
 
-    def forward(self, hidden_states: torch.Tensor, temb: torch.Tensor | None = None) -> torch.Tensor:
+    def forward(
+        self, hidden_states: torch.Tensor, temb: torch.Tensor | None = None
+    ) -> tuple[torch.Tensor, tuple[torch.Tensor]]:
         hidden_states = torch.cat([hidden_states, temb], dim=1)
         for resnet in self.resnets:
             hidden_states = resnet(hidden_states)
diff --git a/src/diffusers/models/unets/unet_2d.py b/src/diffusers/models/unets/unet_2d.py
index 4bbe0535e94a..c9c026d349ae 100644
--- a/src/diffusers/models/unets/unet_2d.py
+++ b/src/diffusers/models/unets/unet_2d.py
@@ -252,7 +252,7 @@ def forward(
         timestep: torch.Tensor | float | int,
         class_labels: torch.Tensor | None = None,
         return_dict: bool = True,
-    ) -> UNet2DOutput | tuple:
+    ) -> UNet2DOutput | tuple[torch.Tensor]:
         r"""
         The [`UNet2DModel`] forward method.
 
diff --git a/src/diffusers/models/unets/unet_2d_condition.py b/src/diffusers/models/unets/unet_2d_condition.py
index af44f0e9d2cb..a5e2ed726dc3 100644
--- a/src/diffusers/models/unets/unet_2d_condition.py
+++ b/src/diffusers/models/unets/unet_2d_condition.py
@@ -991,7 +991,7 @@ def forward(
         down_intrablock_additional_residuals: tuple[torch.Tensor] | None = None,
         encoder_attention_mask: torch.Tensor | None = None,
         return_dict: bool = True,
-    ) -> UNet2DConditionOutput | tuple:
+    ) -> UNet2DConditionOutput | tuple[torch.Tensor]:
         r"""
         The [`UNet2DConditionModel`] forward method.
 
diff --git a/src/diffusers/models/unets/unet_spatio_temporal_condition.py b/src/diffusers/models/unets/unet_spatio_temporal_condition.py
index 7d6b861df82b..572b89e24e05 100644
--- a/src/diffusers/models/unets/unet_spatio_temporal_condition.py
+++ b/src/diffusers/models/unets/unet_spatio_temporal_condition.py
@@ -296,7 +296,7 @@ def forward(
         encoder_hidden_states: torch.Tensor,
         added_time_ids: torch.Tensor,
         return_dict: bool = True,
-    ) -> UNetSpatioTemporalConditionOutput | tuple:
+    ) -> UNetSpatioTemporalConditionOutput | tuple[torch.Tensor]:
         r"""
         The [`UNetSpatioTemporalConditionModel`] forward method.