From 08c79cbf450d13cd05f4c7d6699849c18baf25ce Mon Sep 17 00:00:00 2001 From: conanna Date: Mon, 13 Apr 2026 02:55:17 +0800 Subject: [PATCH] fix: align sigmas device in UniPCMultistepScheduler.set_timesteps --- src/diffusers/schedulers/scheduling_unipc_multistep.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/diffusers/schedulers/scheduling_unipc_multistep.py b/src/diffusers/schedulers/scheduling_unipc_multistep.py index 21f81bc381b1..4e65c7ff13df 100644 --- a/src/diffusers/schedulers/scheduling_unipc_multistep.py +++ b/src/diffusers/schedulers/scheduling_unipc_multistep.py @@ -478,7 +478,8 @@ def set_timesteps( # add an index counter for schedulers that allow duplicated timesteps self._step_index = None self._begin_index = None - self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication + # self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication + self.sigmas = self.sigmas.to(device=device) # Copied from diffusers.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler.time_shift def time_shift(self, mu: float, sigma: float, t: torch.Tensor):