Fix wrong param types, docs, and handles noise=None in scale_noise of FlowMatching schedulers (huggingface#11669)

Promisery · yiyixuxu · web-flow · commit 4fb44bdf910e · 2026-01-09T11:42:33.000-10:00
* Bug: Fix wrong params, docs, and handles noise=None

* make noise a required arg

---------

Co-authored-by: YiYi Xu &lt;yixu310@gmail.com&gt;
diff --git a/src/diffusers/schedulers/scheduling_flow_match_euler_discrete.py b/src/diffusers/schedulers/scheduling_flow_match_euler_discrete.py
@@ -171,17 +171,19 @@ def set_shift(self, shift: float):
     def scale_noise(
         self,
         sample: torch.FloatTensor,
-        timestep: Union[float, torch.FloatTensor],
-        noise: Optional[torch.FloatTensor] = None,
+        timestep: torch.FloatTensor,
+        noise: torch.FloatTensor,
     ) -> torch.FloatTensor:
         """
         Forward process in flow-matching
 
         Args:
             sample (`torch.FloatTensor`):
                 The input sample.
-            timestep (`int`, *optional*):
+            timestep (`torch.FloatTensor`):
                 The current timestep in the diffusion chain.
+            noise (`torch.FloatTensor`):
+                The noise tensor.
 
         Returns:
             `torch.FloatTensor`:
diff --git a/src/diffusers/schedulers/scheduling_flow_match_heun_discrete.py b/src/diffusers/schedulers/scheduling_flow_match_heun_discrete.py
@@ -110,17 +110,19 @@ def set_begin_index(self, begin_index: int = 0):
     def scale_noise(
         self,
         sample: torch.FloatTensor,
-        timestep: Union[float, torch.FloatTensor],
-        noise: Optional[torch.FloatTensor] = None,
+        timestep: torch.FloatTensor,
+        noise: torch.FloatTensor,
     ) -> torch.FloatTensor:
         """
         Forward process in flow-matching
 
         Args:
             sample (`torch.FloatTensor`):
                 The input sample.
-            timestep (`int`, *optional*):
+            timestep (`torch.FloatTensor`):
                 The current timestep in the diffusion chain.
+            noise (`torch.FloatTensor`):
+                The noise tensor.
 
         Returns:
             `torch.FloatTensor`:
@@ -130,6 +132,7 @@ def scale_noise(
             self._init_step_index(timestep)
 
         sigma = self.sigmas[self.step_index]
+
         sample = sigma * noise + (1.0 - sigma) * sample
 
         return sample
diff --git a/src/diffusers/schedulers/scheduling_flow_match_lcm.py b/src/diffusers/schedulers/scheduling_flow_match_lcm.py
@@ -192,17 +192,19 @@ def set_scale_factors(self, scale_factors: list, upscale_mode):
     def scale_noise(
         self,
         sample: torch.FloatTensor,
-        timestep: Union[float, torch.FloatTensor],
-        noise: Optional[torch.FloatTensor] = None,
+        timestep: torch.FloatTensor,
+        noise: torch.FloatTensor,
     ) -> torch.FloatTensor:
         """
         Forward process in flow-matching
 
         Args:
             sample (`torch.FloatTensor`):
                 The input sample.
-            timestep (`int`, *optional*):
+            timestep (`torch.FloatTensor`):
                 The current timestep in the diffusion chain.
+            noise (`torch.FloatTensor`):
+                The noise tensor.
 
         Returns:
             `torch.FloatTensor`:
