Skip noise addition in DPM++ 2M SDE if eta is 0

crowsonkb · Aug 12, 2023 · d911c4b · d911c4b
1 parent 7e5c201
commit d911c4b
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/k_diffusion/sampling.py b/k_diffusion/sampling.py
@@ -644,10 +644,11 @@ def sample_dpmpp_2m_sde(model, x, sigmas, extra_args=None, callback=None, disabl
                 elif solver_type == 'midpoint':
                     x = x + 0.5 * (-h - eta_h).expm1().neg() * (1 / r) * (denoised - old_denoised)
 
-            x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * sigmas[i + 1] * (-2 * eta_h).expm1().neg().sqrt() * s_noise
-            h_last = h
+            if eta:
+                x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * sigmas[i + 1] * (-2 * eta_h).expm1().neg().sqrt() * s_noise
 
         old_denoised = denoised
+        h_last = h
     return x