schedulers fix zero-sigma final-step

Signed-off-by: Vladimir Mandic <[email protected]>

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
 

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)

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:

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]

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)

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)

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,

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,