From 870cca30fa5f0a37c16a8c7e124a542d105c48e7 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Mon, 6 Jul 2026 13:21:19 +0200 Subject: [PATCH] schedulers fix zero-sigma final-step Signed-off-by: Vladimir Mandic --- CHANGELOG.md | 2 + modules/processing_callbacks.py | 8 ++++ modules/processing_vae.py | 45 +++++++++++++------ modules/schedulers/scheduler_dpm_flowmatch.py | 13 +++++- modules/schedulers/scheduler_ersde.py | 11 +++-- modules/schedulers/scheduler_flashflow.py | 27 ++++++----- modules/sdnq/quantizer.py | 4 +- test/test-schedulers.py | 20 +++++++++ 8 files changed, 101 insertions(+), 29 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 77239cf56..793cb56e5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -93,6 +93,7 @@ Also couple of *experimental* features: see below for details... - processors: init code and multiple fixes - pulid: import paths - python: experimental/ignore version checks + - scheduler: handle zero-sigma for i2i/inpaint flowmatch workflows - sdnq: warn instead of error for `triton` - startup: faster model storage checks - text encoder: load non-t5 single-file overrides as their actual class and quantize under sdnq @@ -102,6 +103,7 @@ Also couple of *experimental* features: see below for details... - ui: networks details scrollbars - vae: restore hijack on pipeline switch - vae: scale factor improved detection + - vae: better detection of invalid/nan values ## Update for 2026-06-16 diff --git a/modules/processing_callbacks.py b/modules/processing_callbacks.py index 9f5b88e47..be9d8727a 100644 --- a/modules/processing_callbacks.py +++ b/modules/processing_callbacks.py @@ -91,6 +91,14 @@ def diffusers_callback(pipe, step: int = 0, timestep: int = 0, kwargs: dict | No if latents is None or p is None: return kwargs + """ + if torch.isnan(latents).any().item(): + log.error(f'Callback: step={step} timestep={timestep} latents={latents.shape}:{latents.device}:{latents.dtype} error="contains NaN values"') + if (shared.state.current_latent is not None) and (shared.state.current_latent.shape == latents.shape): + log.error(f'Callback: step={step} timestep={timestep} latents={latents.shape}:{latents.device}:{latents.dtype} error="replacing with previous latent"') + latents = shared.state.current_latent + """ + if len(getattr(p, 'ip_adapter_names', [])) > 0 and p.ip_adapter_names[0] != 'None': ip_adapter_scales = list(p.ip_adapter_scales) ip_adapter_starts = list(p.ip_adapter_starts) diff --git a/modules/processing_vae.py b/modules/processing_vae.py index 84e4b79d1..638eb8f86 100644 --- a/modules/processing_vae.py +++ b/modules/processing_vae.py @@ -249,24 +249,28 @@ def vae_postprocess(tensor, model, output_type='np'): if tensor.ndim == 4 and tensor.shape[1] == 3: tensor = tensor.unsqueeze(2) try: - images = model.video_processor.postprocess_video(tensor, output_type='pil') - except Exception as e: - log.warning(f'VAE postprocess: type=video tensor={tensor.shape}:{tensor.device}:{tensor.dtype} error={e}') + with np.errstate(all='raise'): + images = model.video_processor.postprocess_video(tensor, output_type='pil') + except (Exception, FloatingPointError) as e: + amin, amax = tensor.min().item(), tensor.max().item() + log.warning(f'VAE postprocess: type=video tensor={tensor.shape}:{tensor.device}:{tensor.dtype} min={amin} max={amax} error="{e}"') images = tensor if debug: - errors.display(e, 'VAE postprocess video') + errors.display(e, 'VAE postprocess: type=video') if isinstance(images, list) and len(images) > 0 and isinstance(images[0], list): images = [frame for batch in images for frame in batch] elif hasattr(model, 'image_processor'): if tensor.ndim == 5 and tensor.shape[1] == 3: # Qwen Image tensor = tensor[:, :, 0] try: - images = model.image_processor.postprocess(tensor, output_type=output_type) - except Exception as e: - log.warning(f'VAE postprocess: type=image tensor={tensor.shape}:{tensor.device}:{tensor.dtype} error={e}') + with np.errstate(all='raise'): + images = model.image_processor.postprocess(tensor, output_type=output_type) + except (Exception, FloatingPointError) as e: + amin, amax = tensor.min().item(), tensor.max().item() + log.warning(f'VAE postprocess: type=image tensor={tensor.shape}:{tensor.device}:{tensor.dtype} min={amin} max={amax} error="{e}"') images = tensor if debug: - errors.display(e, 'VAE postprocess image') + errors.display(e, 'VAE postprocess: type=image') elif hasattr(model, "vqgan"): images = tensor.permute(0, 2, 3, 1).cpu().float().numpy() if output_type == "pil": @@ -277,12 +281,27 @@ def vae_postprocess(tensor, model, output_type='np'): if tensor.ndim == 5 and tensor.shape[1] == 3: # Qwen Image tensor = tensor[:, :, 0] images = model.image_processor.postprocess(tensor, output_type=output_type) + if torch.is_tensor(images): # failed to postprocess, do naive conversion - images = images.permute(0, 2, 3, 1).cpu().float().numpy() - if images.min() < 0 or images.max() > 1: - images = (images - images.min()) / (images.max() - images.min()) # naive normalization - if output_type == "pil": - images = model.numpy_to_pil(images) + try: + if torch.isnan(images).any().item(): + log.error(f'VAE postprocess: type=fallback tensor={images.shape}:{images.device}:{images.dtype} error="image contains invalid NaN values"') + images.nan_to_num_(nan=0.0) + while images.ndim > 4: + images = images.squeeze(0) + if images.shape[0] == 3: + images = images.permute(1, 2, 3, 0).cpu().float().numpy() + else: + images = images.permute(0, 2, 3, 1).cpu().float().numpy() + if images.min() < 0 or images.max() > 1: + images = (images - images.min()) / (images.max() - images.min()) # naive normalization + if output_type == "pil": + images = model.numpy_to_pil(images) + except (Exception, FloatingPointError) as e: + amin, amax = images.min().item(), images.max().item() + log.warning(f'VAE postprocess: type=fallback tensor={images.shape}:{images.device}:{images.dtype} min={amin} max={amax} error="{e}"') + if debug: + errors.display(e, 'VAE postprocess unknown') else: images = tensor if isinstance(tensor, list) or isinstance(tensor, np.ndarray) else [tensor] except Exception as e: diff --git a/modules/schedulers/scheduler_dpm_flowmatch.py b/modules/schedulers/scheduler_dpm_flowmatch.py index 8289aa2f5..e4158773a 100644 --- a/modules/schedulers/scheduler_dpm_flowmatch.py +++ b/modules/schedulers/scheduler_dpm_flowmatch.py @@ -490,11 +490,22 @@ class FlowMatchDPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin): if self.step_index is None: self._init_step_index(timestep) + sigma = self.sigmas[self.step_index] + if sigma == 0 or torch.isclose(sigma, torch.tensor(0.0, device=sigma.device, dtype=sigma.dtype)): + prev_sample = sample.to(model_output.dtype) + self._step_index += 1 + torch.cuda.empty_cache() + if not return_dict: + return (prev_sample,) + return FlowMatchDPMSolverMultistepSchedulerOutput(prev_sample=prev_sample) + + def _is_zero(value: torch.Tensor) -> bool: + return bool(value == 0 or torch.isclose(value, torch.tensor(0.0, device=value.device, dtype=value.dtype))) + if self.config.algorithm_type in ["dpmsolver2", "dpmsolver2A"]: pass else: # Flow Match needs to solve an integral of the data prediction model. - sigma = self.sigmas[self.step_index] model_output = sample - sigma * model_output for i in range(self.config.solver_order - 1): self.model_outputs[i] = self.model_outputs[i + 1] diff --git a/modules/schedulers/scheduler_ersde.py b/modules/schedulers/scheduler_ersde.py index 682d0a79e..fc10afc1b 100644 --- a/modules/schedulers/scheduler_ersde.py +++ b/modules/schedulers/scheduler_ersde.py @@ -331,9 +331,14 @@ class ERSDEScheduler(SchedulerMixin, ConfigMixin): alpha = sqrt(acp), sigma_vp = sqrt(1-acp). """ if self._is_flow: - alpha = self._flow_alphas[step_idx] - sigma = self._flow_sigmas[step_idx] - lam = self._flow_lambdas[step_idx] + if step_idx >= len(self._flow_alphas): + alpha = torch.tensor(1.0, dtype=torch.float64) + sigma = torch.tensor(0.0, dtype=torch.float64) + lam = torch.tensor(0.0, dtype=torch.float64) + else: + alpha = self._flow_alphas[step_idx] + sigma = self._flow_sigmas[step_idx] + lam = self._flow_lambdas[step_idx] elif step_idx >= len(self.sigmas): # Past the last entry: fully denoised alpha = torch.tensor(1.0, dtype=torch.float64) diff --git a/modules/schedulers/scheduler_flashflow.py b/modules/schedulers/scheduler_flashflow.py index 506eb0f51..7d0b1f7d1 100644 --- a/modules/schedulers/scheduler_flashflow.py +++ b/modules/schedulers/scheduler_flashflow.py @@ -244,7 +244,10 @@ class FlashFlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin): timesteps = sigmas * self.config.num_train_timesteps sigmas = torch.cat([sigmas, torch.ones(1, device=sigmas.device)]) else: - sigmas = torch.cat([sigmas, torch.zeros(1, device=sigmas.device)]) + if sigmas[-1].abs() < 1e-8: + sigmas = sigmas + else: + sigmas = torch.cat([sigmas, torch.zeros(1, device=sigmas.device)]) self.timesteps = timesteps.to(device=device) self.sigmas = sigmas @@ -357,15 +360,19 @@ class FlashFlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin): if self.step_index < self.num_inference_steps - 1: sigma_next = self.sigmas[self.step_index + 1] - noise = randn_tensor( - model_output.shape, - generator=generator, - device=model_output.device, - dtype=denoised.dtype, - ) - if noise_clip_std > 0.0: - noise = noise.clamp(-noise_clip_std, noise_clip_std) - sample = sigma_next * s_noise * noise + (1.0 - sigma_next) * denoised + at_final_sigma = sigma_next.abs() < 1e-8 + if at_final_sigma: + sample = denoised + else: + noise = randn_tensor( + model_output.shape, + generator=generator, + device=model_output.device, + dtype=denoised.dtype, + ) + if noise_clip_std > 0.0: + noise = noise.clamp(-noise_clip_std, noise_clip_std) + sample = sigma_next * s_noise * noise + (1.0 - sigma_next) * denoised self._step_index += 1 sample = sample.to(model_output.dtype) diff --git a/modules/sdnq/quantizer.py b/modules/sdnq/quantizer.py index 0b864c3c1..6f6adf3d9 100644 --- a/modules/sdnq/quantizer.py +++ b/modules/sdnq/quantizer.py @@ -616,7 +616,7 @@ class SDNQQuantizer(DiffusersQuantizer, HfQuantizer): return False @devices.inference_context() - def create_quantized_param( # pylint: disable=unused-argument + def create_quantized_param( # pylint: disable=unused-argument,arguments-differ self, model: torch.nn.Module, param_value: torch.FloatTensor, @@ -671,7 +671,7 @@ class SDNQQuantizer(DiffusersQuantizer, HfQuantizer): parent_module, tensor_name = get_module_from_name(model, param_name.removesuffix(tensor_name).removesuffix(".")) setattr(parent_module, tensor_name, layer) - def _process_model_before_weight_loading( # pylint: disable=unused-argument + def _process_model_before_weight_loading( # pylint: disable=unused-argument,arguments-differ self, model: torch.nn.Module, device_map, diff --git a/test/test-schedulers.py b/test/test-schedulers.py index ad3fbbc13..c49c9f600 100644 --- a/test/test-schedulers.py +++ b/test/test-schedulers.py @@ -245,6 +245,26 @@ def run_tests(): for cls in flow_schedulers: test_scheduler(cls.__name__, cls, {"prediction_type": "flow_prediction", "use_flow_sigmas": True}) + log.warning('type="flow-custom-sigmas"') + custom_sigmas = torch.linspace(0.8, 0.0, 10) + for cls, name in [ + (FlowMatchDPMSolverMultistepScheduler, "FlowMatchDPMSolverMultistepScheduler"), + (ERSDEScheduler, "ERSDEScheduler"), + (FlashFlowMatchEulerDiscreteScheduler, "FlashFlowMatchEulerDiscreteScheduler"), + ]: + try: + scheduler = cls() + scheduler.set_timesteps(sigmas=custom_sigmas, device='cpu') + sample = torch.randn((1, 4, 64, 64), dtype=torch.float32) + for t in scheduler.timesteps: + model_output = torch.randn_like(sample) + sample = scheduler.step(model_output, t, sample).prev_sample + if torch.isnan(sample).any() or torch.isinf(sample).any(): + log.error(f'scheduler="{name}" error="custom sigmas produced NaN/Inf"') + break + except Exception as e: + log.error(f'scheduler="{name}" error="custom sigmas test exception: {e}"') + log.warning('type="sdnext"') extended_schedulers = [ VDMScheduler,