From ee8f8ee07fb141e5a5ce3abf602ed0fa2e50cf7b Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Mon, 17 Jul 2023 09:35:19 -0400 Subject: [PATCH] Fix regression with ddim and uni_pc when batch size > 1. --- comfy/k_diffusion/external.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/comfy/k_diffusion/external.py b/comfy/k_diffusion/external.py index 680a356..7335d56 100644 --- a/comfy/k_diffusion/external.py +++ b/comfy/k_diffusion/external.py @@ -92,8 +92,8 @@ class DiscreteSchedule(nn.Module): def predict_eps_discrete_timestep(self, input, t, **kwargs): sigma = self.t_to_sigma(t.round()) - input = input * ((sigma ** 2 + 1.0) ** 0.5) - return (input - self(input, sigma, **kwargs)) / sigma + input = input * ((utils.append_dims(sigma, input.ndim) ** 2 + 1.0) ** 0.5) + return (input - self(input, sigma, **kwargs)) / utils.append_dims(sigma, input.ndim) class DiscreteEpsDDPMDenoiser(DiscreteSchedule): """A wrapper for discrete schedule DDPM models that output eps (the predicted