From b3803164cb422bbcbd2bf8db03149c0207120bc5 Mon Sep 17 00:00:00 2001 From: jy-song-hub Date: Sun, 17 May 2026 18:01:07 -0700 Subject: [PATCH] [diffusion] fix: fix unipc device placement + flowunipc sigma_min crash (#23294) --- .../models/schedulers/scheduling_flow_unipc_multistep.py | 2 +- .../models/schedulers/scheduling_unipc_multistep.py | 8 +++++++- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_unipc_multistep.py b/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_unipc_multistep.py index 4874f6966..5d9ea035d 100644 --- a/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_unipc_multistep.py +++ b/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_unipc_multistep.py @@ -208,7 +208,7 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): sigmas = shift * sigmas / (1 + (shift - 1) * sigmas) # pyright: ignore if self.config.final_sigmas_type == "sigma_min": - sigma_last = ((1 - self.alphas_cumprod[0]) / self.alphas_cumprod[0]) ** 0.5 + sigma_last = sigmas[-1] elif self.config.final_sigmas_type == "zero": sigma_last = 0 else: diff --git a/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_unipc_multistep.py b/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_unipc_multistep.py index df5e9b834..cca8b81b0 100644 --- a/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_unipc_multistep.py +++ b/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_unipc_multistep.py @@ -492,7 +492,11 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): ) sigmas = np.concatenate([sigmas, [sigma_last]]).astype(np.float32) + # Keep sigmas on the same device as the sampling tensors. + # If `device` is None, keep them on CPU. self.sigmas = torch.from_numpy(sigmas) + if device is not None: + self.sigmas = self.sigmas.to(device=device) self.timesteps = torch.from_numpy(timesteps).to( device=device, dtype=torch.int64 ) @@ -510,7 +514,9 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): # 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 + # Keep CPU sigmas only for CPU inference; for GPU/MPS this would cause device mismatch. + if device is None or torch.device(device).type == "cpu": + self.sigmas = self.sigmas.to("cpu") # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler._threshold_sample def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor: