From ee1c09a3a6de138eecd6fac6a870d547ec8d1cb4 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Fri, 23 Oct 2026 06:05:03 +0000 Subject: [PATCH] fix docstring inconsistencies. --- .../condition_embedder_ace_step.py | 42 +++++++++++++++++++ .../condition_embedder_audioldm2.py | 17 +++++++- .../condition_embedder_clip_image.py | 11 ++++- .../condition_embedder_flux_redux.py | 9 ++++ .../condition_embedder_ltx2.py | 4 ++ .../condition_embedder_stable_audio.py | 17 +++++++- .../latent_upscaler/latent_upscaler_ltx.py | 10 ++++- .../latent_upscaler/latent_upscaler_ltx2.py | 10 ++++- 8 files changed, 114 insertions(+), 6 deletions(-) diff --git a/src/diffusers/models/condition_embedders/condition_embedder_ace_step.py b/src/diffusers/models/condition_embedders/condition_embedder_ace_step.py index d918cc6d02..8c753850a6 100644 --- a/src/diffusers/models/condition_embedders/condition_embedder_ace_step.py +++ b/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( diff --git a/src/diffusers/models/condition_embedders/condition_embedder_audioldm2.py b/src/diffusers/models/condition_embedders/condition_embedder_audioldm2.py index 4fb7f62867..616b885ba8 100644 --- a/src/diffusers/models/condition_embedders/condition_embedder_audioldm2.py +++ b/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 diff --git a/src/diffusers/models/condition_embedders/condition_embedder_clip_image.py b/src/diffusers/models/condition_embedders/condition_embedder_clip_image.py index a373f69fa3..f717ff5f51 100644 --- a/src/diffusers/models/condition_embedders/condition_embedder_clip_image.py +++ b/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) diff --git a/src/diffusers/models/condition_embedders/condition_embedder_flux_redux.py b/src/diffusers/models/condition_embedders/condition_embedder_flux_redux.py index 526a8ddba9..e8bae87fda 100644 --- a/src/diffusers/models/condition_embedders/condition_embedder_flux_redux.py +++ b/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) diff --git a/src/diffusers/models/condition_embedders/condition_embedder_ltx2.py b/src/diffusers/models/condition_embedders/condition_embedder_ltx2.py index c119b6a0f6..3e5f2d269b 100644 --- a/src/diffusers/models/condition_embedders/condition_embedder_ltx2.py +++ b/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] diff --git a/src/diffusers/models/condition_embedders/condition_embedder_stable_audio.py b/src/diffusers/models/condition_embedders/condition_embedder_stable_audio.py index 28bf563159..0dac6ca016 100644 --- a/src/diffusers/models/condition_embedders/condition_embedder_stable_audio.py +++ b/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) ) diff --git a/src/diffusers/models/latent_upscaler/latent_upscaler_ltx.py b/src/diffusers/models/latent_upscaler/latent_upscaler_ltx.py index 76bb35ead8..f93ddddf13 100644 --- a/src/diffusers/models/latent_upscaler/latent_upscaler_ltx.py +++ b/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: diff --git a/src/diffusers/models/latent_upscaler/latent_upscaler_ltx2.py b/src/diffusers/models/latent_upscaler/latent_upscaler_ltx2.py index d43a779201..da3087a5b7 100644 --- a/src/diffusers/models/latent_upscaler/latent_upscaler_ltx2.py +++ b/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: