From 5d0c01b9e37f9a0fb49934c3be6d6cec102ca3a7 Mon Sep 17 00:00:00 2001 From: Kubuxu Date: Tue, 11 Jul 2023 22:04:43 +0100 Subject: [PATCH 1/3] 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) From a966a0d5adb1375e78c8be2f65e2584a30de15e9 Mon Sep 17 00:00:00 2001 From: Kubuxu Date: Tue, 11 Jul 2023 22:06:14 +0100 Subject: [PATCH 2/3] Use float16 for image processing, force dtype_vae for encoding --- modules/processing.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/modules/processing.py b/modules/processing.py index 5ef1dd226..4b2fe1905 100644 --- a/modules/processing.py +++ b/modules/processing.py @@ -216,7 +216,7 @@ class StableDiffusionProcessing: conditioning_mask = torch.nn.functional.interpolate(conditioning_mask, size=latent_image.shape[-2:]) conditioning_mask = conditioning_mask.expand(conditioning_image.shape[0], -1, -1, -1) image_conditioning = torch.cat([conditioning_mask, conditioning_image], dim=1) - image_conditioning = image_conditioning.to(shared.device).type(self.sd_model.dtype) + image_conditioning = image_conditioning.to(device = shared.device, dtype = source_image.dtype) return image_conditioning def img2img_image_conditioning(self, source_image, latent_image, image_mask=None): @@ -1020,7 +1020,7 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing): image = np.moveaxis(image, 2, 0) batch_images.append(image) decoded_samples = torch.from_numpy(np.array(batch_images)) - decoded_samples = decoded_samples.to(shared.device) + decoded_samples = decoded_samples.to(device=shared.device, dtype=devices.dtype_vae) decoded_samples = 2. * decoded_samples - 1. if shared.opts.sd_vae_sliced_encode and len(decoded_samples) > 1: samples = torch.stack([ @@ -1149,7 +1149,7 @@ class StableDiffusionProcessingImg2Img(StableDiffusionProcessing): raise RuntimeError(f"bad number of images passed: {len(imgs)}; expecting {self.batch_size} or less") image = torch.from_numpy(batch_images) image = 2. * image - 1. - image = image.to(shared.device) + image = image.to(device=shared.device, dtype=devices.dtype_vae) if shared.backend == Backend.ORIGINAL: self.init_latent = self.sd_model.get_first_stage_encoding(self.sd_model.encode_first_stage(image)) @@ -1166,8 +1166,8 @@ class StableDiffusionProcessingImg2Img(StableDiffusionProcessing): latmask = latmask[0] latmask = np.around(latmask) latmask = np.tile(latmask[None], (4, 1, 1)) - self.mask = torch.asarray(1.0 - latmask).to(shared.device).type(self.sd_model.dtype) - self.nmask = torch.asarray(latmask).to(shared.device).type(self.sd_model.dtype) + self.mask = torch.asarray(1.0 - latmask).to(device=shared.device, dtype=self.sd_model.dtype) + self.nmask = torch.asarray(latmask).to(device=shared.device, dtype=self.sd_model.dtype) # this needs to be fixed to be done in sample() using actual seeds for batches if self.inpainting_fill == 2: self.init_latent = self.init_latent * self.mask + create_random_tensors(self.init_latent.shape[1:], all_seeds[0:self.init_latent.shape[0]]) * self.nmask From 2eb705df15449110c7883c9ac5225647eb82dd5b Mon Sep 17 00:00:00 2001 From: Kubuxu Date: Wed, 12 Jul 2023 01:16:37 +0100 Subject: [PATCH 3/3] Intoduce attention heads dimension into sdp_attnblock_forward This enables flash-attention and memory-efficient attention optimizations. --- modules/sd_hijack_optimizations.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/modules/sd_hijack_optimizations.py b/modules/sd_hijack_optimizations.py index a757a3755..730f0f09c 100644 --- a/modules/sd_hijack_optimizations.py +++ b/modules/sd_hijack_optimizations.py @@ -505,7 +505,11 @@ def sdp_attnblock_forward(self, x): k = self.k(h_) v = self.v(h_) b, c, h, w = q.shape # pylint: disable=unused-variable - q, k, v = (rearrange(t, 'b c h w -> b (h w) c') for t in (q, k, v)) + + # SDP optimization kenels are built for operations with multiple attention heads. + # Four dimensional tensors are required for mem_efficient and flash attention to work. + # We add an attention head dimension `a` to allow these kernels to be used. + q, k, v = (rearrange(t, '(b a) c h w -> b a (h w) c', a=1) for t in (q, k, v)) dtype = q.dtype if shared.opts.upcast_attn: q, k, v = q.float(), k.float(), v.float() @@ -514,7 +518,7 @@ def sdp_attnblock_forward(self, x): v = v.contiguous() out = torch.nn.functional.scaled_dot_product_attention(q, k, v, dropout_p=0.0, is_causal=False) out = out.to(dtype) - out = rearrange(out, 'b (h w) c -> b c h w', h=h) + out = rearrange(out, 'b a (h w) c -> (b a) c h w', h=h) # remove the one attention head dimension `a` out = self.proj_out(out) return x + out