Update sd-webui-smea/scripts/sd-webui-smea.py

sd-webui-smea/scripts/sd-webui-smea.py

@@ -1297,43 +1297,102 @@ def sample_euler_h_m_b_c(model, x, sigmas, extra_args=None, callback=None, disab
     return x
 
 @torch.no_grad()
-def sample_euler_h_m_b_c_pp(model, x, sigmas, extra_args=None, callback=None, disable=None, s_churn=0., s_tmin=0., s_tmax=float('inf'), s_noise=1., noise_sampler=None):
+def sample_euler_h_m_b_c_pp(
+    model,
+    x,
+    sigmas,
+    extra_args=None,
+    callback=None,
+    disable=None,
+    s_churn=0.,
+    s_tmin=0.,
+    s_tmax=float('inf'),
+    s_noise=1.,
+    noise_sampler=None,
+):
     extra_args = {} if extra_args is None else extra_args
     s_in = x.new_ones([x.shape[0]])
-    for i in trange(len(sigmas) - 1, disable=disable):
-        wave = math.cos(math.pi * 0.5 * i)/(0.5 * i + 1.5) + 1
-        sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
-        s_tmin, s_tmax = sigma_min if s_tmin == 0. else s_tmin, sigma_max if s_tmax == float('inf') else s_tmax
-        gamma = min(wave * ((2 ** 0.5 - 1) + s_churn) / (len(sigmas) - 1), 2 ** 0.5 - 1) if s_tmin <= sigmas[i] <= s_tmax else 0.
-        eps = k_diffusion.sampling.BrownianTreeNoiseSampler(x, s_tmin, s_tmax, 0) if noise_sampler is None else noise_sampler
-        gammaup = gamma + 1
-        sigma_hat = sigmas[i] * gammaup
-        if gamma > 0:	
-            x = x + eps(sigmas[i], sigmas[i + 1]) * s_noise * (sigma_hat ** 2 - sigmas[i] ** 2) ** 0.5
-        denoised = model(x, sigma_hat * s_in, **extra_args)
-        last_noise_uncond = model.last_noise_uncond
-        d = to_d(x, sigma_hat, denoised)
-        dt = sigmas[i + 1] - sigma_hat
-        if callback is not None:
-            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigma_hat, 'denoised': denoised}) 
-        if i == 0:
-            x = x + d * dt
-        elif i <= len(sigmas) - 4:
-            x_2 = x + d * dt
-            d_2 = to_d(x_2, sigmas[i + 1] * gammaup, denoised)
-            x_3 = x_2 + d_2 * dt
-            d_3 = to_d(x_3, sigmas[i + 2] * gammaup, last_noise_uncond)
-            d_prime = d * 0.5 + d_2 * 0.375 + d_3 * 0.125		
-            x = x + d_prime * dt
-        elif sigmas[i + 1] > 0:
-            x_2 = x + d * dt
-            d_2 = to_d(x_2, sigmas[i + 1] * gammaup, denoised)
-            d_prime = d * 0.5 + d_2 * 0.5		
-            x = x + d_prime * dt
-        else:
-            # Euler method
-            x = x + d * dt
-    return x
+
+    old_need_last = getattr(model, 'need_last_noise_uncond', False)
+    model.need_last_noise_uncond = True
+
+    try:
+        for i in trange(len(sigmas) - 1, disable=disable):
+            wave = math.cos(math.pi * 0.5 * i) / (0.5 * i + 1.5) + 1
+            sigma_min = sigmas[sigmas > 0].min()
+            sigma_max = sigmas.max()
+            s_tmin_eff = sigma_min if s_tmin == 0. else s_tmin
+            s_tmax_eff = sigma_max if s_tmax == float('inf') else s_tmax
+
+            gamma = (
+                min(
+                    wave * ((2 ** 0.5 - 1) + s_churn) / (len(sigmas) - 1),
+                    2 ** 0.5 - 1,
+                )
+                if s_tmin_eff <= sigmas[i] <= s_tmax_eff
+                else 0.
+            )
+
+            eps = (
+                k_diffusion.sampling.BrownianTreeNoiseSampler(x, s_tmin_eff, s_tmax_eff, 0)
+                if noise_sampler is None
+                else noise_sampler
+            )
+
+            gammaup = gamma + 1
+            sigma_hat = sigmas[i] * gammaup
+
+            if gamma > 0:
+                x = x + eps(sigmas[i], sigmas[i + 1]) * s_noise * (sigma_hat ** 2 - sigmas[i] ** 2) ** 0.5
+
+            denoised = model(x, sigma_hat * s_in, **extra_args)
+            d = to_d(x, sigma_hat, denoised)
+            dt = sigmas[i + 1] - sigma_hat
+
+            if callback is not None:
+                callback({
+                    'x': x,
+                    'i': i,
+                    'sigma': sigmas[i],
+                    'sigma_hat': sigma_hat,
+                    'denoised': denoised,
+                })
+
+            if i == 0:
+                x = x + d * dt
+
+            elif i <= len(sigmas) - 4:
+                x_2 = x + d * dt
+                d_2 = to_d(x_2, sigmas[i + 1] * gammaup, denoised)
+
+                x_3 = x_2 + d_2 * dt
+                sigma_3 = sigmas[i + 2] * gammaup
+
+                # Fresh evaluation at the actual third-point location.
+                # With the A1111 CFG++ core hotfix, model.last_noise_uncond is stored
+                # as derivative/noise estimate, which is exactly what this branch needs.
+                _ = model(x_3, sigma_3 * s_in, **extra_args)
+                d_3 = model.last_noise_uncond
+
+                if d_3 is None:
+                    d_3 = d_2
+
+                d_prime = d * 0.5 + d_2 * 0.375 + d_3 * 0.125
+                x = x + d_prime * dt
+
+            elif sigmas[i + 1] > 0:
+                x_2 = x + d * dt
+                d_2 = to_d(x_2, sigmas[i + 1] * gammaup, denoised)
+                d_prime = d * 0.5 + d_2 * 0.5
+                x = x + d_prime * dt
+
+            else:
+                x = x + d * dt
+
+        return x
+
+    finally:
+        model.need_last_noise_uncond = old_need_last
 
 @torch.no_grad()
 def sample_euler_smea_max(model, x, sigmas, extra_args=None, callback=None, disable=None, s_churn=0., s_tmin=0., s_tmax=float('inf'), s_noise=1., smooth=False):