fix docstring inconsistencies.

src/diffusers/models/condition_embedders/condition_embedder_ace_step.py

@@ -183,6 +183,16 @@ class AceStepLyricEncoder(ModelMixin, ConfigMixin):
         inputs_embeds: torch.FloatTensor,
         attention_mask: torch.Tensor,
     ) -> torch.Tensor:
+        r"""
+        Args:
+            inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, text_hidden_dim)`):
+                Lyric embeddings from the text encoder.
+            attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`):
+                Mask with 1 for valid tokens and 0 for padding.
+
+        Returns:
+            `torch.Tensor` of shape `(batch_size, sequence_length, hidden_size)`: The encoded lyric hidden states.
+        """
         inputs_embeds = self.embed_tokens(inputs_embeds)
 
         seq_len = inputs_embeds.shape[1]
@@ -322,6 +332,18 @@ class AceStepTimbreEncoder(ModelMixin, ConfigMixin):
         refer_audio_acoustic_hidden_states_packed: torch.FloatTensor,
         refer_audio_order_mask: torch.LongTensor,
     ) -> Tuple[torch.Tensor, torch.Tensor]:
+        r"""
+        Args:
+            refer_audio_acoustic_hidden_states_packed (`torch.FloatTensor` of shape `(num_references, sequence_length, timbre_hidden_dim)`):
+                VAE latents of the reference audio clips, packed across the batch.
+            refer_audio_order_mask (`torch.LongTensor` of shape `(num_references,)`):
+                Index of the batch item each reference clip belongs to.
+
+        Returns:
+            `Tuple[torch.Tensor, torch.Tensor]`: The timbre embeddings of shape `(batch_size, max_references,
+            hidden_size)` and a mask of shape `(batch_size, max_references)` with 1 for valid embeddings and 0 for
+            padding.
+        """
         inputs_embeds = self.embed_tokens(refer_audio_acoustic_hidden_states_packed)
 
         seq_len = inputs_embeds.shape[1]
@@ -457,6 +479,26 @@ class AceStepConditionEncoder(ModelMixin, ConfigMixin):
         refer_audio_acoustic_hidden_states_packed: torch.FloatTensor,
         refer_audio_order_mask: torch.LongTensor,
     ) -> Tuple[torch.Tensor, torch.Tensor]:
+        r"""
+        Args:
+            text_hidden_states (`torch.FloatTensor` of shape `(batch_size, text_length, text_hidden_dim)`):
+                Prompt embeddings from the text encoder.
+            text_attention_mask (`torch.Tensor` of shape `(batch_size, text_length)`):
+                Mask with 1 for valid prompt tokens and 0 for padding.
+            lyric_hidden_states (`torch.FloatTensor` of shape `(batch_size, lyric_length, text_hidden_dim)`):
+                Lyric embeddings from the text encoder.
+            lyric_attention_mask (`torch.Tensor` of shape `(batch_size, lyric_length)`):
+                Mask with 1 for valid lyric tokens and 0 for padding.
+            refer_audio_acoustic_hidden_states_packed (`torch.FloatTensor` of shape `(num_references, sequence_length, timbre_hidden_dim)`):
+                VAE latents of the reference audio clips, packed across the batch.
+            refer_audio_order_mask (`torch.LongTensor` of shape `(num_references,)`):
+                Index of the batch item each reference clip belongs to.
+
+        Returns:
+            `Tuple[torch.Tensor, torch.Tensor]`: The packed lyric, timbre, and prompt hidden states of shape
+            `(batch_size, sequence_length, hidden_size)`, with valid tokens first, and their attention mask of shape
+            `(batch_size, sequence_length)`.
+        """
         text_hidden_states = self.text_projector(text_hidden_states)
 
         lyric_hidden_states = self.lyric_encoder(

src/diffusers/models/condition_embedders/condition_embedder_audioldm2.py

@@ -107,7 +107,22 @@ class AudioLDM2ProjectionModel(ModelMixin, ConfigMixin):
         hidden_states_1: torch.Tensor | None = None,
         attention_mask: torch.LongTensor | None = None,
         attention_mask_1: torch.LongTensor | None = None,
-    ):
+    ) -> AudioLDM2ProjectionModelOutput:
+        r"""
+        Args:
+            hidden_states (`torch.Tensor` of shape `(batch_size, sequence_length, text_encoder_dim)`, *optional*):
+                Text embeddings from the first text encoder (CLAP).
+            hidden_states_1 (`torch.Tensor` of shape `(batch_size, sequence_length_1, text_encoder_1_dim)`, *optional*):
+                Text embeddings from the second text encoder (T5 or VITS).
+            attention_mask (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+                Attention mask for `hidden_states`.
+            attention_mask_1 (`torch.LongTensor` of shape `(batch_size, sequence_length_1)`, *optional*):
+                Attention mask for `hidden_states_1`.
+
+        Returns:
+            [`AudioLDM2ProjectionModelOutput`]: The projected text embeddings of both text encoders, concatenated along
+            the sequence dimension, and their combined attention mask.
+        """
         hidden_states = self.projection(hidden_states)
         hidden_states, attention_mask = add_special_tokens(
             hidden_states, attention_mask, sos_token=self.sos_embed, eos_token=self.eos_embed

src/diffusers/models/condition_embedders/condition_embedder_clip_image.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import torch
 from torch import nn
 
 from ...configuration_utils import ConfigMixin, register_to_config
@@ -25,5 +26,13 @@ class CLIPImageProjection(ModelMixin, ConfigMixin):
         self.hidden_size = hidden_size
         self.project = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        r"""
+        Args:
+            x (`torch.Tensor` of shape `(..., hidden_size)`):
+                CLIP image embeddings to project.
+
+        Returns:
+            `torch.Tensor` of shape `(..., hidden_size)`: The projected embeddings.
+        """
         return self.project(x)

src/diffusers/models/condition_embedders/condition_embedder_flux_redux.py

@@ -41,6 +41,15 @@ class ReduxImageEncoder(ModelMixin, ConfigMixin):
         self.redux_down = nn.Linear(txt_in_features * 3, txt_in_features)
 
     def forward(self, x: torch.Tensor) -> ReduxImageEncoderOutput:
+        r"""
+        Args:
+            x (`torch.Tensor` of shape `(batch_size, sequence_length, redux_dim)`):
+                Image features from the image encoder.
+
+        Returns:
+            [`ReduxImageEncoderOutput`]: The image embeddings of shape `(batch_size, sequence_length,
+            txt_in_features)`.
+        """
         projected_x = self.redux_down(nn.functional.silu(self.redux_up(x)))
 
         return ReduxImageEncoderOutput(image_embeds=projected_x)

src/diffusers/models/condition_embedders/condition_embedder_ltx2.py

@@ -421,6 +421,10 @@ class LTX2TextConnectors(ModelMixin, PeftAdapterMixin, ConfigMixin):
             scale_factor (`int`, *optional*, defaults to `8`):
                 Scale factor for masked mean/range normalization. Only used if `per_modality_projections` is `False`
                 (LTX-2.0 models).
+
+        Returns:
+            `tuple[torch.Tensor, torch.Tensor, torch.Tensor]`: The video text embeddings, the audio text embeddings,
+            and a binary attention mask of shape `(batch_size, seq_len)` for the video text embeddings.
         """
         if text_encoder_hidden_states.ndim == 3:
             # Ensure shape is [batch_size, seq_len, caption_channels, text_proj_in_factor]

src/diffusers/models/condition_embedders/condition_embedder_stable_audio.py

@@ -96,7 +96,7 @@ class StableAudioNumberConditioner(nn.Module):
     def forward(
         self,
         floats: torch.Tensor,
-    ):
+    ) -> torch.Tensor:
         floats = floats.clamp(self.min_value, self.max_value)
 
         normalized_floats = (floats - self.min_value) / (self.max_value - self.min_value)
@@ -140,7 +140,20 @@ class StableAudioProjectionModel(ModelMixin, ConfigMixin):
         text_hidden_states: torch.Tensor | None = None,
         start_seconds: torch.Tensor | None = None,
         end_seconds: torch.Tensor | None = None,
-    ):
+    ) -> StableAudioProjectionModelOutput:
+        r"""
+        Args:
+            text_hidden_states (`torch.Tensor` of shape `(batch_size, sequence_length, text_encoder_dim)`, *optional*):
+                Text embeddings from the text encoder.
+            start_seconds (`torch.Tensor` of shape `(batch_size,)`, *optional*):
+                Audio start times in seconds.
+            end_seconds (`torch.Tensor` of shape `(batch_size,)`, *optional*):
+                Audio end times in seconds.
+
+        Returns:
+            [`StableAudioProjectionModelOutput`]: The projected text, start-time, and end-time hidden states. Inputs
+            passed as `None` are returned as `None`.
+        """
         text_hidden_states = (
             text_hidden_states if text_hidden_states is None else self.text_projection(text_hidden_states)
         )

src/diffusers/models/latent_upscaler/latent_upscaler_ltx.py

@@ -53,7 +53,7 @@ class PixelShuffleND(torch.nn.Module):
         if dims not in [1, 2, 3]:
             raise ValueError("dims must be 1, 2, or 3")
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         if self.dims == 3:
             # spatiotemporal: b (c p1 p2 p3) d h w -> b c (d p1) (h p2) (w p3)
             return (
@@ -144,6 +144,14 @@ class LTXLatentUpsamplerModel(ModelMixin, ConfigMixin):
         self.final_conv = ConvNd(mid_channels, in_channels, kernel_size=3, padding=1)
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        r"""
+        Args:
+            hidden_states (`torch.Tensor` of shape `(batch_size, in_channels, num_frames, height, width)`):
+                VAE latents to upsample.
+
+        Returns:
+            `torch.Tensor`: The upsampled latents with `in_channels` channels.
+        """
         batch_size, num_channels, num_frames, height, width = hidden_states.shape
 
         if self.dims == 2:

src/diffusers/models/latent_upscaler/latent_upscaler_ltx2.py

@@ -66,7 +66,7 @@ class PixelShuffleND(torch.nn.Module):
         if dims not in [1, 2, 3]:
             raise ValueError("dims must be 1, 2, or 3")
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         if self.dims == 3:
             # spatiotemporal: b (c p1 p2 p3) d h w -> b c (d p1) (h p2) (w p3)
             return (
@@ -243,6 +243,14 @@ class LTX2LatentUpsamplerModel(ModelMixin, ConfigMixin):
         self.final_conv = ConvNd(mid_channels, in_channels, kernel_size=3, padding=1)
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        r"""
+        Args:
+            hidden_states (`torch.Tensor` of shape `(batch_size, in_channels, num_frames, height, width)`):
+                VAE latents to upsample.
+
+        Returns:
+            `torch.Tensor`: The upsampled latents with `in_channels` channels.
+        """
         batch_size, num_channels, num_frames, height, width = hidden_states.shape
 
         if self.dims == 2: