From 512051bafafd28efb4916bfc03a4501af8543fc5 Mon Sep 17 00:00:00 2001 From: resonantsky Date: Tue, 14 Apr 2026 14:28:09 +0200 Subject: [PATCH] samplers: fix UniPCMultistepScheduler sigmas device mismatch at step >= 2 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit upstream UniPCMultistepScheduler.set_timesteps unconditionally moves self.sigmas to CPU after building them. multistep_uni_p_bh_update / multistep_uni_c_bh_update then constructs a torch.ones(..., device=sample.device) tensor and calls torch.stack([..., self.sigmas[...]]) — crashing at inference step >= 2 whenever the model runs on a non-CPU device (CUDA, ROCm, MPS). Monkey-patch set_timesteps so that, after the upstream call, self.sigmas is moved back to the requested device. Applied once at import time inside the existing sampler-load try/except block so failures are silent-logged and never break the rest of the sampler registry. --- modules/sd_samplers_diffusers.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/modules/sd_samplers_diffusers.py b/modules/sd_samplers_diffusers.py index 632d584ec..5676950e8 100644 --- a/modules/sd_samplers_diffusers.py +++ b/modules/sd_samplers_diffusers.py @@ -98,6 +98,21 @@ except Exception as e: if os.environ.get('SD_SAMPLER_DEBUG', None) is not None: errors.display(e, 'Samplers') +# Patch UniPCMultistepScheduler.set_timesteps: upstream forces self.sigmas to CPU after building them, +# but multistep_uni_p/c_bh_update mixes those CPU sigmas with torch.ones(..., device=sample.device), +# crashing torch.stack at step >= 2. Keep sigmas on the compute device instead. +try: + _orig_unipc_set_timesteps = UniPCMultistepScheduler.set_timesteps + + def _unipc_set_timesteps_device_fix(self, num_inference_steps=None, device=None, **kwargs): + _orig_unipc_set_timesteps(self, num_inference_steps=num_inference_steps, device=device, **kwargs) + if device is not None: + self.sigmas = self.sigmas.to(device) + + UniPCMultistepScheduler.set_timesteps = _unipc_set_timesteps_device_fix +except Exception as e: + log.error(f'Sampler patch: UniPCMultistepScheduler.set_timesteps error: {e}') + config = { # beta_start, beta_end are typically per-scheduler, but we don't want them as they should be taken from the model itself as those are values model was trained on # prediction_type is ideally set in model as well, but it maybe needed that we do auto-detect of model type in the future