From c71a7e6b203cb159c77d9396a8889849389abd04 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Mon, 26 Jun 2023 00:48:48 -0400 Subject: [PATCH] Fix ddim + inpainting not working. --- comfy/ldm/models/diffusion/ddim.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/comfy/ldm/models/diffusion/ddim.py b/comfy/ldm/models/diffusion/ddim.py index d5649089..108fce1c 100644 --- a/comfy/ldm/models/diffusion/ddim.py +++ b/comfy/ldm/models/diffusion/ddim.py @@ -180,6 +180,12 @@ class DDIMSampler(object): ) return samples, intermediates + def q_sample(self, x_start, t, noise=None): + if noise is None: + noise = torch.randn_like(x_start) + return (extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start + + extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise) + @torch.no_grad() def ddim_sampling(self, cond, shape, x_T=None, ddim_use_original_steps=False, @@ -214,7 +220,7 @@ class DDIMSampler(object): if mask is not None: assert x0 is not None - img_orig = self.model.q_sample(x0, ts) # TODO: deterministic forward pass? + img_orig = self.q_sample(x0, ts) # TODO: deterministic forward pass? img = img_orig * mask + (1. - mask) * img if ucg_schedule is not None: