From 5d0c01b9e37f9a0fb49934c3be6d6cec102ca3a7 Mon Sep 17 00:00:00 2001 From: Kubuxu Date: Tue, 11 Jul 2023 22:04:43 +0100 Subject: [PATCH] Use dtype_unet as specified, propagete types in gaussian --- modules/sd_hijack.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/modules/sd_hijack.py b/modules/sd_hijack.py index 1ee40d033..070a71b2e 100644 --- a/modules/sd_hijack.py +++ b/modules/sd_hijack.py @@ -2,6 +2,7 @@ from types import MethodType import torch from torch.nn.functional import silu import ldm.modules.attention +import ldm.modules.distributions.distributions import ldm.modules.diffusionmodules.model import ldm.modules.diffusionmodules.openaimodel import ldm.models.diffusion.ddim @@ -305,3 +306,6 @@ def register_buffer(self, name, attr): ldm.models.diffusion.ddim.DDIMSampler.register_buffer = register_buffer ldm.models.diffusion.plms.PLMSSampler.register_buffer = register_buffer + +# Ensure samping from Guassian for DDPM follows types +ldm.modules.distributions.distributions.DiagonalGaussianDistribution.sample = lambda self: self.mean.to(self.parameters.dtype) + self.std.to(self.parameters.dtype) * torch.randn(self.mean.shape, dtype=self.parameters.dtype).to(device=self.parameters.device)