Fixed sampler not being restored if interrupted

2023-06-30 15:22:19 -04:00
parent f18f401051
commit b4e7e042e0
1 changed files with 5 additions and 3 deletions
@@ -58,7 +58,9 @@ def restart_sampling(model, seed, steps, cfg, sampler_name, scheduler, positive,
        return x

    setattr(k_diffusion_sampling, sample_func_name, restart_wrapper)
-    samples = common_ksampler(model, seed, steps, cfg, sampler_name, scheduler,
-                              positive, negative, latent_image, denoise=denoise)
-    setattr(k_diffusion_sampling, sample_func_name, sampler)
+    try:
+        samples = common_ksampler(model, seed, steps, cfg, sampler_name, scheduler,
+                                  positive, negative, latent_image, denoise=denoise)
+    finally:
+        setattr(k_diffusion_sampling, sample_func_name, sampler)
    return samples