docs: improve docstring scheduling_flow_match_heun_discrete.py (#13130)

Improve docstring scheduling flow match heun discrete

Author: delmalih

src/diffusers/schedulers/scheduling_flow_match_heun_discrete.py

@@ -51,9 +51,6 @@ class FlowMatchHeunDiscreteScheduler(SchedulerMixin, ConfigMixin):
     Args:
         num_train_timesteps (`int`, defaults to 1000):
             The number of diffusion steps to train the model.
-        timestep_spacing (`str`, defaults to `"linspace"`):
-            The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and
-            Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information.
         shift (`float`, defaults to 1.0):
             The shift value for the timestep schedule.
     """
@@ -110,7 +107,7 @@ def set_begin_index(self, begin_index: int = 0):
     def scale_noise(
         self,
         sample: torch.FloatTensor,
-        timestep: torch.FloatTensor,
+        timestep: Union[float, torch.FloatTensor],
         noise: torch.FloatTensor,
     ) -> torch.FloatTensor:
         """
@@ -119,7 +116,7 @@ def scale_noise(
         Args:
             sample (`torch.FloatTensor`):
                 The input sample.
-            timestep (`torch.FloatTensor`):
+            timestep (`float` or `torch.FloatTensor`):
                 The current timestep in the diffusion chain.
             noise (`torch.FloatTensor`):
                 The noise tensor.
@@ -137,10 +134,14 @@ def scale_noise(
 
         return sample
 
-    def _sigma_to_t(self, sigma):
+    def _sigma_to_t(self, sigma: float) -> float:
         return sigma * self.config.num_train_timesteps
 
-    def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.device] = None):
+    def set_timesteps(
+        self,
+        num_inference_steps: int,
+        device: Union[str, torch.device] = None,
+    ) -> None:
         """
         Sets the discrete timesteps used for the diffusion chain (to be run before inference).
 
@@ -153,7 +154,9 @@ def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.devic
         self.num_inference_steps = num_inference_steps
 
         timesteps = np.linspace(
-            self._sigma_to_t(self.sigma_max), self._sigma_to_t(self.sigma_min), num_inference_steps
+            self._sigma_to_t(self.sigma_max),
+            self._sigma_to_t(self.sigma_min),
+            num_inference_steps,
         )
 
         sigmas = timesteps / self.config.num_train_timesteps
@@ -174,7 +177,24 @@ def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.devic
         self._step_index = None
         self._begin_index = None
 
-    def index_for_timestep(self, timestep, schedule_timesteps=None):
+    def index_for_timestep(
+        self,
+        timestep: Union[float, torch.FloatTensor],
+        schedule_timesteps: Optional[torch.FloatTensor] = None,
+    ) -> int:
+        """
+        Find the index of a given timestep in the timestep schedule.
+
+        Args:
+            timestep (`float` or `torch.FloatTensor`):
+                The timestep value to find in the schedule.
+            schedule_timesteps (`torch.FloatTensor`, *optional*):
+                The timestep schedule to search in. If `None`, uses `self.timesteps`.
+
+        Returns:
+            `int`:
+                The index of the timestep in the schedule.
+        """
         if schedule_timesteps is None:
             schedule_timesteps = self.timesteps
 
@@ -188,7 +208,7 @@ def index_for_timestep(self, timestep, schedule_timesteps=None):
 
         return indices[pos].item()
 
-    def _init_step_index(self, timestep):
+    def _init_step_index(self, timestep: Union[float, torch.FloatTensor]) -> None:
         if self.begin_index is None:
             if isinstance(timestep, torch.Tensor):
                 timestep = timestep.to(self.timesteps.device)
@@ -197,7 +217,10 @@ def _init_step_index(self, timestep):
             self._step_index = self._begin_index
 
     @property
-    def state_in_first_order(self):
+    def state_in_first_order(self) -> bool:
+        """
+        Returns whether the scheduler is in the first-order state.
+        """
         return self.dt is None
 
     def step(
@@ -219,13 +242,19 @@ def step(
         Args:
             model_output (`torch.FloatTensor`):
                 The direct output from learned diffusion model.
-            timestep (`float`):
+            timestep (`float` or `torch.FloatTensor`):
                 The current discrete timestep in the diffusion chain.
             sample (`torch.FloatTensor`):
                 A current instance of a sample created by the diffusion process.
             s_churn (`float`):
-            s_tmin  (`float`):
-            s_tmax  (`float`):
+                Stochasticity parameter that controls the amount of noise added during sampling. Higher values increase
+                randomness.
+            s_tmin (`float`):
+                Minimum timestep threshold for applying stochasticity. Only timesteps above this value will have noise
+                added.
+            s_tmax (`float`):
+                Maximum timestep threshold for applying stochasticity. Only timesteps below this value will have noise
+                added.
             s_noise (`float`, defaults to 1.0):
                 Scaling factor for noise added to the sample.
             generator (`torch.Generator`, *optional*):
@@ -274,7 +303,10 @@ def step(
 
         if gamma > 0:
             noise = randn_tensor(
-                model_output.shape, dtype=model_output.dtype, device=model_output.device, generator=generator
+                model_output.shape,
+                dtype=model_output.dtype,
+                device=model_output.device,
+                generator=generator,
             )
             eps = noise * s_noise
             sample = sample + eps * (sigma_hat**2 - sigma**2) ** 0.5
@@ -320,5 +352,5 @@ def step(
 
         return FlowMatchHeunDiscreteSchedulerOutput(prev_sample=prev_sample)
 
-    def __len__(self):
+    def __len__(self) -> int:
         return self.config.num_train_timesteps