Improve docstrings and type hints in scheduling_amused.py

- Add complete type hints for helper functions (gumbel_noise, mask_by_random_topk)
- Enhance AmusedSchedulerOutput with proper Optional typing
- Add comprehensive docstrings for AmusedScheduler class
- Improve __init__, set_timesteps, step, and add_noise methods
- Fix type hints to match documentation conventions
- All changes follow project standards from issue #9567

Author: delmalih

src/diffusers/schedulers/scheduling_amused.py

@@ -9,13 +9,48 @@
 from .scheduling_utils import SchedulerMixin
 
 
-def gumbel_noise(t, generator=None):
+def gumbel_noise(t: torch.Tensor, generator: Optional[torch.Generator] = None) -> torch.Tensor:
+    """
+    Generate Gumbel noise for sampling.
+
+    Args:
+        t (`torch.Tensor`):
+            Input tensor to match the shape and dtype of the output noise.
+        generator (`torch.Generator`, *optional*):
+            A random number generator for reproducible sampling.
+
+    Returns:
+        `torch.Tensor`:
+            Gumbel-distributed noise with the same shape as the input tensor.
+    """
     device = generator.device if generator is not None else t.device
     noise = torch.zeros_like(t, device=device).uniform_(0, 1, generator=generator).to(t.device)
     return -torch.log((-torch.log(noise.clamp(1e-20))).clamp(1e-20))
 
 
-def mask_by_random_topk(mask_len, probs, temperature=1.0, generator=None):
+def mask_by_random_topk(
+    mask_len: torch.Tensor,
+    probs: torch.Tensor,
+    temperature: float = 1.0,
+    generator: Optional[torch.Generator] = None,
+) -> torch.Tensor:
+    """
+    Mask tokens by selecting the top-k lowest confidence scores with temperature-based randomness.
+
+    Args:
+        mask_len (`torch.Tensor`):
+            Number of tokens to mask per sample in the batch.
+        probs (`torch.Tensor`):
+            Probability scores for each token.
+        temperature (`float`, *optional*, defaults to 1.0):
+            Temperature parameter for controlling randomness in the masking process.
+        generator (`torch.Generator`, *optional*):
+            A random number generator for reproducible sampling.
+
+    Returns:
+        `torch.Tensor`:
+            Boolean mask indicating which tokens should be masked.
+    """
     confidence = torch.log(probs.clamp(1e-20)) + temperature * gumbel_noise(probs, generator=generator)
     sorted_confidence = torch.sort(confidence, dim=-1).values
     cut_off = torch.gather(sorted_confidence, 1, mask_len.long())
@@ -38,20 +73,31 @@ class AmusedSchedulerOutput(BaseOutput):
     """
 
     prev_sample: torch.Tensor
-    pred_original_sample: torch.Tensor = None
+    pred_original_sample: Optional[torch.Tensor] = None
 
 
 class AmusedScheduler(SchedulerMixin, ConfigMixin):
     order = 1
 
-    temperatures: torch.Tensor
+    temperatures: Optional[torch.Tensor]
+    timesteps: Optional[torch.Tensor]
 
     @register_to_config
     def __init__(
         self,
         mask_token_id: int,
         masking_schedule: str = "cosine",
     ):
+        """
+        Create a new AmusedScheduler instance.
+
+        Args:
+            mask_token_id (`int`):
+                The token ID used to represent masked tokens in the sequence.
+            masking_schedule (`str`, *optional*, defaults to `"cosine"`):
+                The schedule type for determining the mask ratio at each timestep. Can be either `"cosine"` or
+                `"linear"`.
+        """
         self.temperatures = None
         self.timesteps = None
 
@@ -60,7 +106,21 @@ def set_timesteps(
         num_inference_steps: int,
         temperature: Union[int, Tuple[int, int], List[int]] = (2, 0),
         device: Union[str, torch.device] = None,
-    ):
+    ) -> None:
+        """
+        Set the discrete timesteps used for the diffusion chain.
+
+        Args:
+            num_inference_steps (`int`):
+                The number of diffusion steps used when generating samples with a pre-trained model.
+            temperature (`Union[int, Tuple[int, int], List[int]]`, *optional*, defaults to `(2, 0)`):
+                Temperature parameter(s) for controlling the randomness of sampling. If a tuple or list is provided,
+                temperatures will be linearly interpolated between the first and second values across all timesteps. If
+                a single value is provided, temperatures will be linearly interpolated from that value to 0.01.
+            device (`Union[str, torch.device]`, *optional*):
+                The device to which the timesteps and temperatures should be moved. If not specified, uses the default
+                device.
+        """
         self.timesteps = torch.arange(num_inference_steps, device=device).flip(0)
 
         if isinstance(temperature, (tuple, list)):
@@ -71,12 +131,38 @@ def set_timesteps(
     def step(
         self,
         model_output: torch.Tensor,
-        timestep: torch.long,
+        timestep: int,
         sample: torch.LongTensor,
-        starting_mask_ratio: int = 1,
+        starting_mask_ratio: float = 1.0,
         generator: Optional[torch.Generator] = None,
         return_dict: bool = True,
     ) -> Union[AmusedSchedulerOutput, Tuple]:
+        """
+        Predict the sample at the previous timestep by masking tokens based on confidence scores.
+
+        Args:
+            model_output (`torch.Tensor`):
+                The direct output from the learned diffusion model. Typically of shape `(batch_size, num_tokens,
+                codebook_size)` or `(batch_size, codebook_size, height, width)` for 2D inputs.
+            timestep (`int`):
+                The current discrete timestep in the diffusion chain.
+            sample (`torch.LongTensor`):
+                A current instance of a sample created by the diffusion process. Contains token IDs, with masked
+                positions indicated by `mask_token_id`.
+            starting_mask_ratio (`float`, *optional*, defaults to 1.0):
+                A multiplier applied to the mask ratio schedule. Values less than 1.0 will result in fewer tokens being
+                masked at each step.
+            generator (`torch.Generator`, *optional*):
+                A random number generator for reproducible sampling.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether to return an [`~schedulers.scheduling_amused.AmusedSchedulerOutput`] or a plain tuple.
+
+        Returns:
+            [`~schedulers.scheduling_amused.AmusedSchedulerOutput`] or `tuple`:
+                If `return_dict` is `True`, returns [`~schedulers.scheduling_amused.AmusedSchedulerOutput`], otherwise
+                returns a tuple where the first element is the sample tensor and the second element is the predicted
+                original sample tensor.
+        """
         two_dim_input = sample.ndim == 3 and model_output.ndim == 4
 
         if two_dim_input:
@@ -137,7 +223,27 @@ def step(
 
         return AmusedSchedulerOutput(prev_sample, pred_original_sample)
 
-    def add_noise(self, sample, timesteps, generator=None):
+    def add_noise(
+        self,
+        sample: torch.LongTensor,
+        timesteps: int,
+        generator: Optional[torch.Generator] = None,
+    ) -> torch.LongTensor:
+        """
+        Add noise to a sample by randomly masking tokens according to the masking schedule.
+
+        Args:
+            sample (`torch.LongTensor`):
+                The input sample containing token IDs to be partially masked.
+            timesteps (`int`):
+                The timestep that determines how much masking to apply. Higher timesteps result in more masking.
+            generator (`torch.Generator`, *optional*):
+                A random number generator for reproducible masking.
+
+        Returns:
+            `torch.LongTensor`:
+                The sample with some tokens replaced by `mask_token_id` according to the masking schedule.
+        """
         step_idx = (self.timesteps == timesteps).nonzero()
         ratio = (step_idx + 1) / len(self.timesteps)