diff --git a/src/diffusers/schedulers/scheduling_flow_match_euler_discrete.py b/src/diffusers/schedulers/scheduling_flow_match_euler_discrete.py index 575423ee80e7..93260fcc5ceb 100644 --- a/src/diffusers/schedulers/scheduling_flow_match_euler_discrete.py +++ b/src/diffusers/schedulers/scheduling_flow_match_euler_discrete.py @@ -171,8 +171,8 @@ 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 @@ -180,8 +180,10 @@ def scale_noise( 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 index 2addc5f3eeec..630e93c4de88 100644 --- a/src/diffusers/schedulers/scheduling_flow_match_heun_discrete.py +++ b/src/diffusers/schedulers/scheduling_flow_match_heun_discrete.py @@ -110,8 +110,8 @@ 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 @@ -119,8 +119,10 @@ def scale_noise( 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 index d79556ae8077..c6eb9116553e 100644 --- a/src/diffusers/schedulers/scheduling_flow_match_lcm.py +++ b/src/diffusers/schedulers/scheduling_flow_match_lcm.py @@ -192,8 +192,8 @@ 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 @@ -201,8 +201,10 @@ def scale_noise( 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`: