docs: improve docstring scheduling folder - last batch

src/diffusers/schedulers/scheduling_ltx_euler_ancestral_rf.py

@@ -81,30 +81,37 @@ def __init__(
         num_train_timesteps: int = 1000,
         eta: float = 1.0,
         s_noise: float = 1.0,
-    ):
+    ) -> None:
         # Note: num_train_timesteps is kept only for config compatibility.
-        self.num_inference_steps: int = None
+        self.num_inference_steps: int | None = None
         self.sigmas: torch.Tensor | None = None
         self.timesteps: torch.Tensor | None = None
-        self._step_index: int = None
-        self._begin_index: int = None
+        self._step_index: int | None = None
+        self._begin_index: int | None = None
 
     @property
-    def step_index(self) -> int:
+    def step_index(self) -> int | None:
+        """
+        The index counter for the current timestep. It increases by one after each scheduler step.
+        """
         return self._step_index
 
     @property
-    def begin_index(self) -> int:
+    def begin_index(self) -> int | None:
         """
         The index for the first timestep. It can be set from a pipeline with `set_begin_index` to support
         image-to-image like workflows that start denoising part-way through the schedule.
         """
         return self._begin_index
 
-    def set_begin_index(self, begin_index: int = 0):
+    def set_begin_index(self, begin_index: int = 0) -> None:
         """
         Included for API compatibility; not strictly needed here but kept to allow pipelines that call
         `set_begin_index`.
+
+        Args:
+            begin_index (`int`, defaults to `0`):
+                The index of the first timestep in the denoising schedule.
         """
         self._begin_index = begin_index
 
@@ -117,6 +124,16 @@ def index_for_timestep(
         This follows the convention used in other discrete schedulers: if the same timestep value appears multiple
         times in the schedule (which can happen when starting in the middle of the schedule), the *second* occurrence
         is used for the first `step` call so that no sigma is accidentally skipped.
+
+        Args:
+            timestep (`float` or `torch.Tensor`):
+                The timestep value to find in the schedule.
+            schedule_timesteps (`torch.Tensor`, *optional*):
+                The timestep schedule to search. If `None`, `self.timesteps` is used.
+
+        Returns:
+            `int`:
+                The index of the timestep in the schedule.
         """
         if schedule_timesteps is None:
             if self.timesteps is None:
@@ -141,9 +158,13 @@ def index_for_timestep(
 
         return indices[pos].item()
 
-    def _init_step_index(self, timestep: float | torch.Tensor):
+    def _init_step_index(self, timestep: float | torch.Tensor) -> None:
         """
         Initialize the internal step index based on a given timestep.
+
+        Args:
+            timestep (`float` or `torch.Tensor`):
+                The current timestep used to initialize the step index.
         """
         if self.timesteps is None:
             raise ValueError("Timesteps have not been set. Call `set_timesteps` first.")
@@ -162,8 +183,8 @@ def set_timesteps(
         sigmas: list[float] | torch.Tensor | None = None,
         timesteps: list[float] | torch.Tensor | None = None,
         mu: float | None = None,
-        **kwargs,
-    ):
+        **kwargs: object,
+    ) -> None:
         """
         Set the sigma / timestep schedule for sampling.
 
@@ -186,6 +207,8 @@ def set_timesteps(
             mu (`float`, *optional*):
                 Optional shift parameter used when delegating to [`FlowMatchEulerDiscreteScheduler.set_timesteps`] and
                 `config.use_dynamic_shifting` is `True`.
+            kwargs (`object`):
+                Additional keyword arguments accepted for API compatibility.
         """
         # 1. Auto-generate schedule (FlowMatch-style) when no explicit sigmas/timesteps are given
         if sigmas is None and timesteps is None:
@@ -269,6 +292,16 @@ def set_timesteps(
     def _sigma_broadcast(self, sigma: torch.Tensor, sample: torch.Tensor) -> torch.Tensor:
         """
         Helper to broadcast a scalar sigma to the shape of `sample`.
+
+        Args:
+            sigma (`torch.Tensor`):
+                The scalar sigma tensor to broadcast.
+            sample (`torch.Tensor`):
+                The sample tensor whose number of dimensions determines the broadcast shape.
+
+        Returns:
+            `torch.Tensor`:
+                The sigma tensor reshaped for broadcasting with `sample`.
         """
         while sigma.ndim < sample.ndim:
             sigma = sigma.view(*sigma.shape, 1)
@@ -295,9 +328,15 @@ def step(
                 Current latent sample `x_t`.
             generator (`torch.Generator`, *optional*):
                 Optional generator for reproducible noise.
-            return_dict (`bool`):
+            return_dict (`bool`, defaults to `True`):
                 If `True`, return a `LTXEulerAncestralRFSchedulerOutput`; otherwise return a tuple where the first
                 element is the updated sample.
+
+        Returns:
+            [`~schedulers.scheduling_ltx_euler_ancestral_rf.LTXEulerAncestralRFSchedulerOutput`] or `tuple`:
+                If `return_dict` is `True`,
+                [`~schedulers.scheduling_ltx_euler_ancestral_rf.LTXEulerAncestralRFSchedulerOutput`] is returned.
+                Otherwise, a tuple is returned where the first element is the updated sample.
         """
 
         if isinstance(timestep, (int, torch.IntTensor, torch.LongTensor)):
@@ -380,6 +419,11 @@ def step(
         return LTXEulerAncestralRFSchedulerOutput(prev_sample=prev_sample)
 
     def __len__(self) -> int:
-        # For compatibility with other schedulers; used e.g. in some training
-        # utilities to infer the maximum number of training timesteps.
+        """
+        Return the number of training timesteps for compatibility with scheduler utilities.
+
+        Returns:
+            `int`:
+                The configured number of training timesteps.
+        """
         return int(getattr(self.config, "num_train_timesteps", 1000))

src/diffusers/schedulers/scheduling_utils.py

@@ -15,6 +15,7 @@
 import os
 from dataclasses import dataclass
 from enum import Enum
+from typing import Any
 
 import torch
 from huggingface_hub.utils import validate_hf_hub_args
@@ -31,6 +32,10 @@
 # When it's used as a type in pipelines, it really is a Union because the actual
 # scheduler instance is passed in.
 class KarrasDiffusionSchedulers(Enum):
+    """
+    Enumeration of schedulers compatible with Karras diffusion pipelines.
+    """
+
     DDIMScheduler = 1
     DDPMScheduler = 2
     PNDMScheduler = 3
@@ -97,13 +102,13 @@ def from_pretrained(
         cls,
         pretrained_model_name_or_path: str | os.PathLike | None = None,
         subfolder: str | None = None,
-        return_unused_kwargs=False,
-        **kwargs,
-    ) -> Self:
+        return_unused_kwargs: bool = False,
+        **kwargs: Any,
+    ) -> Self | tuple[Self, dict[str, Any]]:
         r"""
         Instantiate a scheduler from a pre-defined JSON configuration file in a local directory or Hub repository.
 
-        Parameters:
+        Args:
             pretrained_model_name_or_path (`str` or `os.PathLike`, *optional*):
                 Can be either:
 
@@ -115,28 +120,32 @@ def from_pretrained(
                 The subfolder location of a model file within a larger model repository on the Hub or locally.
             return_unused_kwargs (`bool`, *optional*, defaults to `False`):
                 Whether kwargs that are not consumed by the Python class should be returned or not.
-            cache_dir (`str | os.PathLike`, *optional*):
+            cache_dir (`str` or `os.PathLike`, *optional*):
                 Path to a directory where a downloaded pretrained model configuration is cached if the standard cache
                 is not used.
             force_download (`bool`, *optional*, defaults to `False`):
                 Whether or not to force the (re-)download of the model weights and configuration files, overriding the
                 cached versions if they exist.
-
             proxies (`dict[str, str]`, *optional*):
                 A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128',
                 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.
-            output_loading_info(`bool`, *optional*, defaults to `False`):
+            output_loading_info (`bool`, *optional*, defaults to `False`):
                 Whether or not to also return a dictionary containing missing keys, unexpected keys and error messages.
-            local_files_only(`bool`, *optional*, defaults to `False`):
+            local_files_only (`bool`, *optional*, defaults to `False`):
                 Whether to only load local model weights and configuration files or not. If set to `True`, the model
                 won't be downloaded from the Hub.
-            token (`str` or *bool*, *optional*):
+            token (`str` or `bool`, *optional*):
                 The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from
                 `diffusers-cli login` (stored in `~/.huggingface`) is used.
             revision (`str`, *optional*, defaults to `"main"`):
                 The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier
                 allowed by Git.
 
+        Returns:
+            [`~schedulers.scheduling_utils.SchedulerMixin`] or `tuple[SchedulerMixin, dict[str, Any]]`:
+                The scheduler instantiated from the configuration. If `return_unused_kwargs` is `True`, a tuple is
+                returned where the second element is a dictionary of unused keyword arguments.
+
         > [!TIP] > To use private or [gated models](https://huggingface.co/docs/hub/models-gated#gated-models), log-in
         with `hf > auth login`. You can also activate the special >
         ["offline-mode"](https://huggingface.co/diffusers/installation.html#offline-mode) to use this method in a >
@@ -152,7 +161,7 @@ def from_pretrained(
         )
         return cls.from_config(config, return_unused_kwargs=return_unused_kwargs, **kwargs)
 
-    def save_pretrained(self, save_directory: str | os.PathLike, push_to_hub: bool = False, **kwargs):
+    def save_pretrained(self, save_directory: str | os.PathLike, push_to_hub: bool = False, **kwargs: Any) -> None:
         """
         Save a scheduler configuration object to a directory so that it can be reloaded using the
         [`~SchedulerMixin.from_pretrained`] class method.
@@ -170,17 +179,18 @@ def save_pretrained(self, save_directory: str | os.PathLike, push_to_hub: bool =
         self.save_config(save_directory=save_directory, push_to_hub=push_to_hub, **kwargs)
 
     @property
-    def compatibles(self):
+    def compatibles(self) -> list[type[Self]]:
         """
-        Returns all schedulers that are compatible with this scheduler
+        Return all scheduler classes that are compatible with this scheduler.
 
         Returns:
-            `list[SchedulerMixin]`: list of compatible schedulers
+            `list[type[SchedulerMixin]]`:
+                A list of compatible scheduler classes.
         """
         return self._get_compatibles()
 
     @classmethod
-    def _get_compatibles(cls):
+    def _get_compatibles(cls) -> list[type[Self]]:
         compatible_classes_str = list(set([cls.__name__] + cls._compatibles))
         diffusers_library = importlib.import_module(__name__.split(".")[0])
         compatible_classes = [