From 51c65457c432612ca4fa043c1078a1bcb43085ec Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Wed, 31 Aug 2022 15:29:01 +0300
Subject: [PATCH] workaround for a weird bug in DDIM's make_schedule

---
 webui.py | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/webui.py b/webui.py
index 80952b792..7c7acefbe 100644
--- a/webui.py
+++ b/webui.py
@@ -847,7 +847,6 @@ class VanillaStableDiffusionSampler:
     def __init__(self, constructor):
         self.sampler = constructor(sd_model)
         self.orig_p_sample_ddim = self.sampler.p_sample_ddim
-        self.sampler.p_sample_ddim = lambda x_dec, cond, ts, *args, **kwargs: p_sample_ddim_hook(self, x_dec, cond, ts, *args, **kwargs)
         self.mask = None
         self.nmask = None
         self.init_latent = None
@@ -855,9 +854,15 @@ class VanillaStableDiffusionSampler:
     def sample_img2img(self, p, x, noise, conditioning, unconditional_conditioning):
         t_enc = int(min(p.denoising_strength, 0.999) * p.steps)
 
-        self.sampler.make_schedule(ddim_num_steps=p.steps, ddim_eta=0.0, verbose=False)
+        # existing code fail with cetin step counts, like 9
+        try:
+            self.sampler.make_schedule(ddim_num_steps=p.steps, verbose=False)
+        except Exception:
+            self.sampler.make_schedule(ddim_num_steps=p.steps+1, verbose=False)
+
         x1 = self.sampler.stochastic_encode(x, torch.tensor([t_enc] * int(x.shape[0])).to(device), noise=noise)
 
+        self.sampler.p_sample_ddim = lambda x_dec, cond, ts, *args, **kwargs: p_sample_ddim_hook(self, x_dec, cond, ts, *args, **kwargs)
         self.mask = p.mask
         self.nmask = p.nmask
         self.init_latent = p.init_latent