Merge pull request #1648 from Kubuxu/feat/unet-dtype

Use float16 for image processing, use dummy attention head to enable spd flash-attention and memory-efficient attention

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

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)

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