fix unbound local error crowsonkb#72

ai-tnc · Oct 18, 2023 · f79fe90 · f79fe90
1 parent 0455157
commit f79fe90
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/k_diffusion/sampling.py b/k_diffusion/sampling.py
@@ -621,6 +621,7 @@ def sample_dpmpp_2m_sde(model, x, sigmas, extra_args=None, callback=None, disabl
 
     old_denoised = None
     h_last = None
+    h = None
 
     for i in trange(len(sigmas) - 1, disable=disable):
         denoised = model(x, sigmas[i] * s_in, **extra_args)
@@ -647,8 +648,8 @@ def sample_dpmpp_2m_sde(model, x, sigmas, extra_args=None, callback=None, disabl
             if eta:
                 x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * sigmas[i + 1] * (-2 * eta_h).expm1().neg().sqrt() * s_noise
 
+            h_last = h
         old_denoised = denoised
-        h_last = h
     return x
 
 
@@ -663,6 +664,7 @@ def sample_dpmpp_3m_sde(model, x, sigmas, extra_args=None, callback=None, disabl
 
     denoised_1, denoised_2 = None, None
     h_1, h_2 = None, None
+    h = None
 
     for i in trange(len(sigmas) - 1, disable=disable):
         denoised = model(x, sigmas[i] * s_in, **extra_args)
@@ -697,6 +699,6 @@ def sample_dpmpp_3m_sde(model, x, sigmas, extra_args=None, callback=None, disabl
             if eta:
                 x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * sigmas[i + 1] * (-2 * h * eta).expm1().neg().sqrt() * s_noise
 
+            h_1, h_2 = h, h_1
         denoised_1, denoised_2 = denoised, denoised_1
-        h_1, h_2 = h, h_1
     return x