[WIP] chore: add utilities to check if call/forward methods are documented.

.github/workflows/pr_modular_tests.yml

@@ -73,6 +73,7 @@ jobs:
           python utils/check_copies.py
           python utils/check_dummies.py
           python utils/check_support_list.py
+          python utils/check_forward_call_docstrings.py
           make deps_table_check_updated
       - name: Check if failure
         if: ${{ failure() }}

.github/workflows/pr_tests.yml

@@ -68,6 +68,7 @@ jobs:
           python utils/check_copies.py
           python utils/check_dummies.py
           python utils/check_support_list.py
+          python utils/check_forward_call_docstrings.py
           make deps_table_check_updated
       - name: Check if failure
         if: ${{ failure() }}

.github/workflows/pr_tests_gpu.yml

@@ -69,6 +69,7 @@ jobs:
           python utils/check_copies.py
           python utils/check_dummies.py
           python utils/check_support_list.py
+          python utils/check_forward_call_docstrings.py
           make deps_table_check_updated
       - name: Check if failure
         if: ${{ failure() }}

Makefile

@@ -36,6 +36,7 @@ repo-consistency:
 	python utils/check_dummies.py
 	python utils/check_repo.py
 	python utils/check_inits.py
+	python utils/check_forward_call_docstrings.py
 
 # this target runs checks on all files
 
@@ -74,6 +75,10 @@ fix-copies:
 modular-autodoctrings:
 	python utils/modular_auto_docstring.py
 
+# Verify forward() / __call__() arguments are documented in their docstrings
+check-forward-call-docstrings:
+	python utils/check_forward_call_docstrings.py
+
 # Run tests for the library
 
 test:

src/diffusers/models/adapter.py

@@ -269,6 +269,10 @@ def forward(self, x: torch.Tensor) -> list[torch.Tensor]:
         each representing information extracted at a different scale from the input. The length of the list is
         determined by the number of downsample blocks in the Adapter, as specified by the `channels` and
         `num_res_blocks` parameters during initialization.
+
+        Args:
+            x (`torch.Tensor`):
+                The input tensor to process through the adapter model.
         """
         return self.adapter(x)
 

src/diffusers/models/autoencoders/autoencoder_asym_kl.py

@@ -166,6 +166,9 @@ def forward(
                 Whether to sample from the posterior.
             return_dict (`bool`, *optional*, defaults to `True`):
                 Whether or not to return a [`DecoderOutput`] instead of a plain tuple.
+            generator (`torch.Generator`, *optional*):
+                A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make sampling
+                deterministic.
         """
         x = sample
         posterior = self.encode(x).latent_dist

src/diffusers/models/autoencoders/autoencoder_dc.py

@@ -706,6 +706,12 @@ def tiled_decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutp
         return DecoderOutput(sample=decoded)
 
     def forward(self, sample: torch.Tensor, return_dict: bool = True) -> torch.Tensor:
+        r"""
+        Args:
+            sample (`torch.Tensor`): Input sample.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`DecoderOutput`] instead of a plain tuple.
+        """
         encoded = self.encode(sample, return_dict=False)[0]
         decoded = self.decode(encoded, return_dict=False)[0]
         if not return_dict:

src/diffusers/models/autoencoders/autoencoder_kl.py

@@ -424,6 +424,9 @@ def forward(
                 Whether to sample from the posterior.
             return_dict (`bool`, *optional*, defaults to `True`):
                 Whether or not to return a [`DecoderOutput`] instead of a plain tuple.
+            generator (`torch.Generator`, *optional*):
+                A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make sampling
+                deterministic.
         """
         x = sample
         posterior = self.encode(x).latent_dist

src/diffusers/models/autoencoders/autoencoder_kl_cogvideox.py

@@ -1409,6 +1409,17 @@ def forward(
         return_dict: bool = True,
         generator: torch.Generator | None = None,
     ) -> torch.Tensor | torch.Tensor:
+        r"""
+        Args:
+            sample (`torch.Tensor`): Input sample.
+            sample_posterior (`bool`, *optional*, defaults to `False`):
+                Whether to sample from the posterior.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`DecoderOutput`] instead of a plain tuple.
+            generator (`torch.Generator`, *optional*):
+                A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make sampling
+                deterministic.
+        """
         x = sample
         posterior = self.encode(x).latent_dist
         if sample_posterior:

src/diffusers/models/autoencoders/autoencoder_kl_cosmos.py

@@ -1078,6 +1078,17 @@ def forward(
         return_dict: bool = True,
         generator: torch.Generator | None = None,
     ) -> tuple[torch.Tensor] | DecoderOutput:
+        r"""
+        Args:
+            sample (`torch.Tensor`): Input sample.
+            sample_posterior (`bool`, *optional*, defaults to `False`):
+                Whether to sample from the posterior.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`DecoderOutput`] instead of a plain tuple.
+            generator (`torch.Generator`, *optional*):
+                A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make sampling
+                deterministic.
+        """
         x = sample
         posterior = self.encode(x).latent_dist
         if sample_posterior:

src/diffusers/models/autoencoders/autoencoder_kl_flux2.py

@@ -441,6 +441,9 @@ def forward(
                 Whether to sample from the posterior.
             return_dict (`bool`, *optional*, defaults to `True`):
                 Whether or not to return a [`DecoderOutput`] instead of a plain tuple.
+            generator (`torch.Generator`, *optional*):
+                A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make sampling
+                deterministic.
         """
         x = sample
         posterior = self.encode(x).latent_dist

src/diffusers/models/autoencoders/autoencoder_kl_hunyuan_video.py

@@ -1061,6 +1061,9 @@ def forward(
                 Whether to sample from the posterior.
             return_dict (`bool`, *optional*, defaults to `True`):
                 Whether or not to return a [`DecoderOutput`] instead of a plain tuple.
+            generator (`torch.Generator`, *optional*):
+                A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make sampling
+                deterministic.
         """
         x = sample
         posterior = self.encode(x).latent_dist

src/diffusers/models/autoencoders/autoencoder_kl_hunyuanimage.py

@@ -674,8 +674,13 @@ def forward(
         """
         Args:
             sample (`torch.Tensor`): Input sample.
+            sample_posterior (`bool`, *optional*, defaults to `False`):
+                Whether to sample from the posterior.
             return_dict (`bool`, *optional*, defaults to `True`):
                 Whether or not to return a [`DecoderOutput`] instead of a plain tuple.
+            generator (`torch.Generator`, *optional*):
+                A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make sampling
+                deterministic.
         """
         posterior = self.encode(sample).latent_dist
         if sample_posterior:

src/diffusers/models/autoencoders/autoencoder_kl_hunyuanimage_refiner.py

@@ -908,6 +908,9 @@ def forward(
                 Whether to sample from the posterior.
             return_dict (`bool`, *optional*, defaults to `True`):
                 Whether or not to return a [`DecoderOutput`] instead of a plain tuple.
+            generator (`torch.Generator`, *optional*):
+                A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make sampling
+                deterministic.
         """
         x = sample
         posterior = self.encode(x).latent_dist

src/diffusers/models/autoencoders/autoencoder_kl_hunyuanvideo15.py

@@ -941,6 +941,9 @@ def forward(
                 Whether to sample from the posterior.
             return_dict (`bool`, *optional*, defaults to `True`):
                 Whether or not to return a [`DecoderOutput`] instead of a plain tuple.
+            generator (`torch.Generator`, *optional*):
+                A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make sampling
+                deterministic.
         """
         x = sample
         posterior = self.encode(x).latent_dist

src/diffusers/models/autoencoders/autoencoder_kl_kvae.py

@@ -787,6 +787,9 @@ def forward(
                 Whether to sample from the posterior.
             return_dict (`bool`, *optional*, defaults to `True`):
                 Whether or not to return a [`DecoderOutput`] instead of a plain tuple.
+            generator (`torch.Generator`, *optional*):
+                A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make sampling
+                deterministic.
         """
         x = sample
         posterior = self.encode(x).latent_dist

src/diffusers/models/autoencoders/autoencoder_kl_kvae_video.py

@@ -942,6 +942,17 @@ def forward(
         return_dict: bool = True,
         generator: Optional[torch.Generator] = None,
     ) -> Union[DecoderOutput, torch.Tensor]:
+        r"""
+        Args:
+            sample (`torch.Tensor`): Input sample.
+            sample_posterior (`bool`, *optional*, defaults to `False`):
+                Whether to sample from the posterior.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`DecoderOutput`] instead of a plain tuple.
+            generator (`torch.Generator`, *optional*):
+                A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make sampling
+                deterministic.
+        """
         x = sample
         posterior = self.encode(x).latent_dist
         if sample_posterior:

src/diffusers/models/autoencoders/autoencoder_kl_ltx.py

@@ -1522,6 +1522,19 @@ def forward(
         return_dict: bool = True,
         generator: torch.Generator | None = None,
     ) -> torch.Tensor | torch.Tensor:
+        r"""
+        Args:
+            sample (`torch.Tensor`): Input sample.
+            temb (`torch.Tensor`, *optional*):
+                Optional timestep embedding tensor used to condition the decoder.
+            sample_posterior (`bool`, *optional*, defaults to `False`):
+                Whether to sample from the posterior.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`DecoderOutput`] instead of a plain tuple.
+            generator (`torch.Generator`, *optional*):
+                A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make sampling
+                deterministic.
+        """
         x = sample
         posterior = self.encode(x).latent_dist
         if sample_posterior:

src/diffusers/models/autoencoders/autoencoder_kl_ltx2.py

@@ -1542,6 +1542,23 @@ def forward(
         return_dict: bool = True,
         generator: torch.Generator | None = None,
     ) -> torch.Tensor | torch.Tensor:
+        r"""
+        Args:
+            sample (`torch.Tensor`): Input sample.
+            temb (`torch.Tensor`, *optional*):
+                Optional timestep embedding tensor used to condition the decoder.
+            sample_posterior (`bool`, *optional*, defaults to `False`):
+                Whether to sample from the posterior.
+            encoder_causal (`bool`, *optional*):
+                Whether the encoder should use causal convolutions. If `None`, falls back to the model default.
+            decoder_causal (`bool`, *optional*):
+                Whether the decoder should use causal convolutions. If `None`, falls back to the model default.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`DecoderOutput`] instead of a plain tuple.
+            generator (`torch.Generator`, *optional*):
+                A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make sampling
+                deterministic.
+        """
         x = sample
         posterior = self.encode(x, causal=encoder_causal).latent_dist
         if sample_posterior:

src/diffusers/models/autoencoders/autoencoder_kl_ltx2_audio.py

@@ -792,6 +792,17 @@ def forward(
         return_dict: bool = True,
         generator: torch.Generator | None = None,
     ) -> DecoderOutput | torch.Tensor:
+        r"""
+        Args:
+            sample (`torch.Tensor`): Input sample.
+            sample_posterior (`bool`, *optional*, defaults to `False`):
+                Whether to sample from the posterior.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`DecoderOutput`] instead of a plain tuple.
+            generator (`torch.Generator`, *optional*):
+                A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make sampling
+                deterministic.
+        """
         posterior = self.encode(sample).latent_dist
         if sample_posterior:
             z = posterior.sample(generator=generator)

src/diffusers/models/autoencoders/autoencoder_kl_magvit.py

@@ -1057,6 +1057,9 @@ def forward(
                 Whether to sample from the posterior.
             return_dict (`bool`, *optional*, defaults to `True`):
                 Whether or not to return a [`DecoderOutput`] instead of a plain tuple.
+            generator (`torch.Generator`, *optional*):
+                A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make sampling
+                deterministic.
         """
         x = sample
         posterior = self.encode(x).latent_dist

src/diffusers/models/autoencoders/autoencoder_kl_mochi.py

@@ -1093,6 +1093,17 @@ def forward(
         return_dict: bool = True,
         generator: torch.Generator | None = None,
     ) -> torch.Tensor | torch.Tensor:
+        r"""
+        Args:
+            sample (`torch.Tensor`): Input sample.
+            sample_posterior (`bool`, *optional*, defaults to `False`):
+                Whether to sample from the posterior.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`DecoderOutput`] instead of a plain tuple.
+            generator (`torch.Generator`, *optional*):
+                A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make sampling
+                deterministic.
+        """
         x = sample
         posterior = self.encode(x).latent_dist
         if sample_posterior:

src/diffusers/models/autoencoders/autoencoder_kl_qwenimage.py

@@ -1043,8 +1043,13 @@ def forward(
         """
         Args:
             sample (`torch.Tensor`): Input sample.
+            sample_posterior (`bool`, *optional*, defaults to `False`):
+                Whether to sample from the posterior.
             return_dict (`bool`, *optional*, defaults to `True`):
                 Whether or not to return a [`DecoderOutput`] instead of a plain tuple.
+            generator (`torch.Generator`, *optional*):
+                A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make sampling
+                deterministic.
         """
         x = sample
         posterior = self.encode(x).latent_dist

src/diffusers/models/autoencoders/autoencoder_kl_temporal_decoder.py

@@ -287,6 +287,11 @@ def forward(
                 Whether to sample from the posterior.
             return_dict (`bool`, *optional*, defaults to `True`):
                 Whether or not to return a [`DecoderOutput`] instead of a plain tuple.
+            generator (`torch.Generator`, *optional*):
+                A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make sampling
+                deterministic.
+            num_frames (`int`, *optional*, defaults to 1):
+                The number of frames to decode per batch.
         """
         x = sample
         posterior = self.encode(x).latent_dist

src/diffusers/models/autoencoders/autoencoder_kl_wan.py

@@ -1416,8 +1416,13 @@ def forward(
         """
         Args:
             sample (`torch.Tensor`): Input sample.
+            sample_posterior (`bool`, *optional*, defaults to `False`):
+                Whether to sample from the posterior.
             return_dict (`bool`, *optional*, defaults to `True`):
                 Whether or not to return a [`DecoderOutput`] instead of a plain tuple.
+            generator (`torch.Generator`, *optional*):
+                A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make sampling
+                deterministic.
         """
         x = sample
         posterior = self.encode(x).latent_dist

src/diffusers/models/autoencoders/autoencoder_longcat_audio_dit.py

@@ -393,6 +393,17 @@ def forward(
         return_dict: bool = True,
         generator: torch.Generator | None = None,
     ) -> LongCatAudioDiTVaeDecoderOutput | tuple[torch.Tensor]:
+        r"""
+        Args:
+            sample (`torch.Tensor`): Input sample.
+            sample_posterior (`bool`, *optional*, defaults to `False`):
+                Whether to sample from the posterior.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`LongCatAudioDiTVaeDecoderOutput`] instead of a plain tuple.
+            generator (`torch.Generator`, *optional*):
+                A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make sampling
+                deterministic.
+        """
         latents = self.encode(sample, sample_posterior=sample_posterior, return_dict=True, generator=generator).latents
         decoded = self.decode(latents, return_dict=True).sample
         if not return_dict:

src/diffusers/models/autoencoders/autoencoder_oobleck.py

@@ -528,6 +528,9 @@ def forward(
                 Whether to sample from the posterior.
             return_dict (`bool`, *optional*, defaults to `True`):
                 Whether or not to return a [`OobleckDecoderOutput`] instead of a plain tuple.
+            generator (`torch.Generator`, *optional*):
+                A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make sampling
+                deterministic.
         """
         x = sample
         posterior = self.encode(x).latent_dist

src/diffusers/models/autoencoders/autoencoder_rae.py

@@ -682,6 +682,15 @@ def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | t
     def forward(
         self, sample: torch.Tensor, return_dict: bool = True, generator: torch.Generator | None = None
     ) -> DecoderOutput | tuple[torch.Tensor]:
+        r"""
+        Args:
+            sample (`torch.Tensor`): Input sample.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`DecoderOutput`] instead of a plain tuple.
+            generator (`torch.Generator`, *optional*):
+                A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make sampling
+                deterministic.
+        """
         latents = self.encode(sample, return_dict=False, generator=generator)[0]
         decoded = self.decode(latents, return_dict=False)[0]
         if not return_dict:

src/diffusers/models/autoencoders/autoencoder_vidtok.py

@@ -1440,6 +1440,19 @@ def forward(
         return_dict: bool = True,
         generator: Optional[torch.Generator] = None,
     ) -> Union[torch.Tensor, DecoderOutput]:
+        r"""
+        Args:
+            sample (`torch.Tensor`): Input sample.
+            sample_posterior (`bool`, *optional*, defaults to `True`):
+                Whether to sample from the posterior.
+            encoder_mode (`bool`, *optional*, defaults to `False`):
+                If `True`, only run the encoder and return the encoded latent without decoding.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`DecoderOutput`] instead of a plain tuple.
+            generator (`torch.Generator`, *optional*):
+                A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make sampling
+                deterministic.
+        """
         x = sample
         res = 1 if self.is_causal else 0
         if self.is_causal: