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

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:

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: