minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[diffusion] fix: fix unipc device placement + flowunipc sigma_min crash (#23294)

Browse Source
This commit is contained in:
jy-song-hub
2026-05-18 09:01:07 +08:00
committed by GitHub
parent c67b287056
commit b3803164cb
2 changed files with 8 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_unipc_multistep.py
+1 -1
View File
@@ -208,7 +208,7 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
sigmas = shift * sigmas / (1 + (shift - 1) * sigmas) # pyright: ignore sigmas = shift * sigmas / (1 + (shift - 1) * sigmas) # pyright: ignore
if self.config.final_sigmas_type == "sigma_min": 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": elif self.config.final_sigmas_type == "zero":
sigma_last = 0 sigma_last = 0
else: else:
python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_unipc_multistep.py
+7 -1
View File
@@ -492,7 +492,11 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
) )
sigmas = np.concatenate([sigmas, [sigma_last]]).astype(np.float32) 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) 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( self.timesteps = torch.from_numpy(timesteps).to(
device=device, dtype=torch.int64 device=device, dtype=torch.int64
) )
@@ -510,7 +514,9 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
# add an index counter for schedulers that allow duplicated timesteps # add an index counter for schedulers that allow duplicated timesteps
self._step_index = None self._step_index = None
self._begin_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 # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler._threshold_sample
def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor: def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor: