From 40cf818f7234eef443bbd9491f831d4113b14f94 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Mon, 8 Jun 2026 00:09:58 +0100 Subject: [PATCH] feat(samplers): apply sigma transforms to ER-SDE flow sigmas In flow mode ER-SDE only ran karras/beta/exponential on the VP path and silently dropped them, unlike FlowMatchEuler and DPM FlowMatch which redistribute the shifted flow sigmas. Apply the same transform to the flow sigmas in _setup_flow so the sigma method works in flow mode and the ER-SDE FlowMatch variants gain karras/beta/exponential. Default flow schedule is unchanged. --- modules/schedulers/scheduler_ersde.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/modules/schedulers/scheduler_ersde.py b/modules/schedulers/scheduler_ersde.py index 8866988a1..682d0a79e 100644 --- a/modules/schedulers/scheduler_ersde.py +++ b/modules/schedulers/scheduler_ersde.py @@ -219,6 +219,17 @@ class ERSDEScheduler(SchedulerMixin, ConfigMixin): sigmas = self._time_shift(mu, 1.0, sigmas) else: sigmas = self.config.shift * sigmas / (1 + (self.config.shift - 1) * sigmas) + # redistribute the shifted flow sigmas via the configured schedule (parity with FlowMatchEuler) + if self.config.use_karras_sigmas or self.config.use_exponential_sigmas or self.config.use_beta_sigmas: + arr = sigmas.detach().cpu().numpy() + n = len(arr) + if self.config.use_karras_sigmas: + arr = self._convert_to_karras(arr, n) + elif self.config.use_exponential_sigmas: + arr = self._convert_to_exponential(arr, n) + elif self.config.use_beta_sigmas: + arr = self._convert_to_beta(arr, n) + sigmas = torch.from_numpy(np.asarray(arr, dtype=np.float64)).to(device=sigmas.device, dtype=sigmas.dtype) flow_sigmas = sigmas.clamp(min=1e-8, max=1.0 - 1e-8) flow_alphas = torch.ones_like(flow_sigmas) flow_lambdas = flow_sigmas.clone()