update dep

diffusion.py

@@ -16,11 +16,27 @@ def make_sigmas(beta_start=0.00085, beta_end=0.012, num_train_timesteps=default_
     return sigmas
 
 
+_with_tqdm = False
+
+
+def set_with_tqdm(it):
+    global _with_tqdm
+
+    _with_tqdm = it
+
+
 @torch.no_grad()
 def rk_ode_solver_diffusion_loop(eps_theta, timesteps, sigmas, x_T, rk_steps_weights):
     x_t = x_T
 
-    for i in range(len(timesteps) - 1, -1, -1):
+    iter_over = range(len(timesteps) - 1, -1, -1)
+
+    if _with_tqdm:
+        from tqdm import tqdm
+
+        iter_over = tqdm(iter_over)
+
+    for i in iter_over:
         t = timesteps[i].unsqueeze(0)
         sigma = sigmas[t]
 

sdxl.py

@@ -9,7 +9,6 @@ import torch
 import torch.nn.functional as F
 import torchvision.transforms
 import torchvision.transforms.functional as TF
-import webdataset as wds
 from PIL import Image
 from torch.nn.parallel import DistributedDataParallel as DDP
 from torch.utils.data import default_collate
@@ -348,6 +347,8 @@ class SDXLTraining:
 
 
 def get_sdxl_dataset(train_shards: str, shuffle_buffer_size: int, batch_size: int, proportion_empty_prompts: float, get_sdxl_conditioning_images=None):
+    import webdataset as wds
+
     dataset = (
         wds.WebDataset(
             train_shards,
@@ -443,61 +444,85 @@ def get_random_crop_params(input_size: Tuple[int, int], output_size: Tuple[int,
     return i, j, th, tw
 
 
-def get_sdxl_conditioning_images(image, adapter_type=None, controlnet_type=None, controlnet_variant=None, open_pose=None, conditioning_image_mask=None):
+def get_adapter_openpose_conditioning_image(image, open_pose):
     resolution = image.width
 
-    if adapter_type == "openpose":
-        conditioning_image = open_pose(image, detect_resolution=resolution, image_resolution=resolution, return_pil=False)
+    conditioning_image = open_pose(image, detect_resolution=resolution, image_resolution=resolution, return_pil=False)
 
-        if (conditioning_image == 0).all():
-            return None, None
+    if (conditioning_image == 0).all():
+        return None, None
 
-        conditioning_image_as_pil = Image.fromarray(conditioning_image)
+    conditioning_image_as_pil = Image.fromarray(conditioning_image)
 
-        conditioning_image = TF.to_tensor(conditioning_image)
+    conditioning_image = TF.to_tensor(conditioning_image)
 
-    if controlnet_type == "canny":
-        import cv2
+    return dict(conditioning_image=conditioning_image, conditioning_image_as_pil=conditioning_image_as_pil)
 
-        conditioning_image = np.array(image)
-        conditioning_image = cv2.Canny(conditioning_image, 100, 200)
-        conditioning_image = conditioning_image[:, :, None]
-        conditioning_image = np.concatenate([conditioning_image, conditioning_image, conditioning_image], axis=2)
 
-        conditioning_image_as_pil = Image.fromarray(conditioning_image)
+def get_controlnet_canny_conditioning_image(image):
+    import cv2
 
-        conditioning_image = TF.to_tensor(conditioning_image)
+    conditioning_image = np.array(image)
+    conditioning_image = cv2.Canny(conditioning_image, 100, 200)
+    conditioning_image = conditioning_image[:, :, None]
+    conditioning_image = np.concatenate([conditioning_image, conditioning_image, conditioning_image], axis=2)
 
-    if controlnet_type == "inpainting":
-        if conditioning_image_mask is None:
-            if random.random() <= 0.25:
-                conditioning_image_mask = np.ones((resolution, resolution), np.float32)
-            else:
-                conditioning_image_mask = random.choice([make_random_rectangle_mask, make_random_irregular_mask, make_outpainting_mask])(resolution, resolution)
+    conditioning_image_as_pil = Image.fromarray(conditioning_image)
+
+    conditioning_image = TF.to_tensor(conditioning_image)
+
+    return dict(conditioning_image=conditioning_image, conditioning_image_as_pil=conditioning_image_as_pil)
+
+
+def get_controlnet_pre_encoded_controlnet_inpainting_conditioning_image(image, conditioning_image_mask):
+    resolution = image.width
+
+    if conditioning_image_mask is None:
+        if random.random() <= 0.25:
+            conditioning_image_mask = np.ones((resolution, resolution), np.float32)
+        else:
+            conditioning_image_mask = random.choice([make_random_rectangle_mask, make_random_irregular_mask, make_outpainting_mask])(resolution, resolution)
 
-            conditioning_image_mask = torch.from_numpy(conditioning_image_mask)
+        conditioning_image_mask = torch.from_numpy(conditioning_image_mask)
 
-            conditioning_image_mask = conditioning_image_mask[None, :, :]
+        conditioning_image_mask = conditioning_image_mask[None, :, :]
 
-        conditioning_image = TF.to_tensor(image)
+    conditioning_image = TF.to_tensor(image)
 
-        if controlnet_variant == "pre_encoded_controlnet_cond":
-            # where mask is 1, zero out the pixels. Note that this requires mask to be concattenated
-            # with the mask so that the network knows the zeroed out pixels are from the mask and
-            # are not just zero in the original image
-            conditioning_image = conditioning_image * (conditioning_image_mask < 0.5)
+    # where mask is 1, zero out the pixels. Note that this requires mask to be concattenated
+    # with the mask so that the network knows the zeroed out pixels are from the mask and
+    # are not just zero in the original image
+    conditioning_image = conditioning_image * (conditioning_image_mask < 0.5)
 
-            conditioning_image_as_pil = TF.to_pil_image(conditioning_image)
+    conditioning_image_as_pil = TF.to_pil_image(conditioning_image)
 
-            conditioning_image = TF.normalize(conditioning_image, [0.5], [0.5])
+    conditioning_image = TF.normalize(conditioning_image, [0.5], [0.5])
+
+    return dict(conditioning_image=conditioning_image, conditioning_image_mask=conditioning_image_mask, conditioning_image_as_pil=conditioning_image_as_pil)
+
+
+def get_controlnet_inpainting_conditioning_image(image, conditioning_image_mask):
+    resolution = image.width
+
+    if conditioning_image_mask is None:
+        if random.random() <= 0.25:
+            conditioning_image_mask = np.ones((resolution, resolution), np.float32)
         else:
-            # Just zero out the pixels which will be masked
-            conditioning_image_as_pil = TF.to_pil_image(conditioning_image * (conditioning_image_mask < 0.5))
+            conditioning_image_mask = random.choice([make_random_rectangle_mask, make_random_irregular_mask, make_outpainting_mask])(resolution, resolution)
+
+        conditioning_image_mask = torch.from_numpy(conditioning_image_mask)
+
+        conditioning_image_mask = conditioning_image_mask[None, :, :]
+
+    conditioning_image = TF.to_tensor(image)
 
-            # where mask is set to 1, set to -1 "special" masked image pixel.
-            # -1 is outside of the 0-1 range that the controlnet normalized
-            # input is in.
-            conditioning_image = conditioning_image * (conditioning_image_mask < 0.5) + -1.0 * (conditioning_image_mask >= 0.5)
+    # Just zero out the pixels which will be masked
+    conditioning_image_as_pil = TF.to_pil_image(conditioning_image * (conditioning_image_mask < 0.5))
+
+    # where mask is set to 1, set to -1 "special" masked image pixel.
+    # -1 is outside of the 0-1 range that the controlnet normalized
+    # input is in.
+    conditioning_image = conditioning_image * (conditioning_image_mask < 0.5) + -1.0 * (conditioning_image_mask >= 0.5)
 
     return dict(conditioning_image=conditioning_image, conditioning_image_mask=conditioning_image_mask, conditioning_image_as_pil=conditioning_image_as_pil)
 
@@ -830,102 +855,112 @@ def sdxl_eps_theta(
 
 known_negative_prompt = "text, watermark, low-quality, signature, moiré pattern, downsampling, aliasing, distorted, blurry, glossy, blur, jpeg artifacts, compression artifacts, poorly drawn, low-resolution, bad, distortion, twisted, excessive, exaggerated pose, exaggerated limbs, grainy, symmetrical, duplicate, error, pattern, beginner, pixelated, fake, hyper, glitch, overexposed, high-contrast, bad-contrast"
 
+if __name__ == "__main__":
+    from argparse import ArgumentParser
 
-# TODO probably just combine with sdxl_diffusion_loop
-def gen_sdxl_simplified_interface(
-    prompts: Union[str, List[str]],
-    negative_prompts: Optional[Union[str, List[str]]] = None,
-    controlnet_checkpoint: Optional[str] = None,
-    controlnet: Optional[Literal["SDXLControlNet", "SDXLContolNetFull", "SDXLControlNetPreEncodedControlnetCond"]] = None,
-    adapter_checkpoint: Optional[str] = None,
-    num_inference_steps=50,
-    images=None,
-    masks=None,
-    apply_conditioning: Optional[Literal["canny"]] = None,
-    num_images: int = 1,
-    guidance_scale=5.0,
-    device: Optional[str] = None,
-    text_encoder_one=None,
-    text_encoder_two=None,
-    unet=None,
-    vae=None,
-):
-    if device is None:
+    args = ArgumentParser()
+    args.add_argument("--prompts", required=True, type=str, nargs="+")
+    args.add_argument("--negative_prompts", required=False, type=str, nargs="+")
+    args.add_argument("--use_known_negative_prompt", action="store_true")
+    args.add_argument("--num_images_per_prompt", required=True, type=int, default=1)
+    args.add_argument("--num_inference_steps", required=False, type=int, default=50)
+    args.add_argument("--images", required=False, type=str, default=None, nargs="+")
+    args.add_argument("--masks", required=False, type=str, default=None, nargs="+")
+    args.add_argument("--controlnet_checkpoint", required=False, type=str, default=None)
+    args.add_argument("--controlnet", required=False, choices=["SDXLControlNet", "SDXLControlNetFull", "SDXLControNetPreEncodedControlnetCond"], default=None)
+    args.add_argument("--adapter_checkpoint", required=False, type=str, default=None)
+    args.add_argument("--device", required=False, default=None)
+    args.add_argument("--dtype", required=False, default="fp16", choices=["fp16", "fp32"])
+    args.add_argument("--guidance_scale", required=False, default=5.0, type=float)
+    args.add_argument("--seed", required=False, type=int)
+    args = args.parse_args()
+
+    if args.device is None:
         if torch.cuda.is_available():
             device = "cuda"
         elif torch.backends.mps.is_available():
             device = "mps"
 
-    if text_encoder_one is None:
+    if args.dtype == "fp16":
+        dtype = torch.float16
+
         text_encoder_one = CLIPTextModel.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", subfolder="text_encoder", variant="fp16", torch_dtype=torch.float16)
         text_encoder_one.to(device=device)
 
-    if text_encoder_two is None:
         text_encoder_two = CLIPTextModelWithProjection.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", subfolder="text_encoder_2", variant="fp16", torch_dtype=torch.float16)
         text_encoder_two.to(device=device)
 
-    if vae is None:
         vae = SDXLVae.load_fp16_fix(device=device)
+        vae.to(torch.float16)
 
-    if unet is None:
         unet = SDXLUNet.load_fp16(device=device)
+    elif args.dtype == "fp32":
+        dtype = torch.float32
 
-    if isinstance(controlnet, str) and controlnet_checkpoint is not None:
-        if controlnet == "SDXLControlNet":
-            controlnet = SDXLControlNet.load(controlnet_checkpoint, device=device, dtype=torch.float16)
-        elif controlnet == "SDXLControlNetFull":
-            controlnet = SDXLControlNetFull.load(controlnet_checkpoint, device=device, dtype=torch.float16)
-        elif controlnet == "SDXLControlNetPreEncodedControlnetCond":
-            controlnet = SDXLControlNetPreEncodedControlnetCond.load(controlnet_checkpoint, device=device, dtype=torch.float16)
-        else:
-            assert False
+        text_encoder_one = CLIPTextModel.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", subfolder="text_encoder")
+        text_encoder_one.to(device=device)
 
-    if adapter_checkpoint is not None:
-        adapter = SDXLAdapter.load(adapter_checkpoint, device=device, dtype=torch.float16)
+        text_encoder_two = CLIPTextModelWithProjection.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", subfolder="text_encoder_2")
+        text_encoder_two.to(device=device)
+
+        vae = SDXLVae.load_fp16_fix(device=device)
+
+        unet = SDXLUNet.load_fp32(device=device)
     else:
-        adapter = None
+        assert False
 
-    sigmas = make_sigmas()
+    if args.controlnet == "SDXLControlNet":
+        controlnet = SDXLControlNet.load(args.controlnet_checkpoint, device=device)
+        controlnet.to(dtype)
+    elif args.controlnet == "SDXLControlNetFull":
+        controlnet = SDXLControlNetFull.load(args.controlnet_checkpoint, device=device)
+        controlnet.to(dtype)
+    elif args.controlnet == "SDXLControlNetPreEncodedControlnetCond":
+        controlnet = SDXLControlNetPreEncodedControlnetCond.load(args.controlnet_checkpoint, device=device)
+        controlnet.to(dtype)
+    else:
+        controlnet = None
 
-    timesteps = torch.linspace(0, sigmas.numel() - 1, num_inference_steps, dtype=torch.long, device=unet.device)
+    if args.adapter_checkpoint is not None:
+        adapter = SDXLAdapter.load(args.adapter_checkpoint, device=device)
+        adapter.to(dtype)
+    else:
+        adapter = None
 
-    if images is not None:
-        if not isinstance(images, list):
-            images = [images]
+    sigmas = make_sigmas(device=device).to(unet.dtype)
 
-        if masks is not None and not isinstance(masks, list):
-            masks = [masks]
+    timesteps = torch.linspace(0, sigmas.numel() - 1, args.num_inference_steps, dtype=torch.long, device=unet.device)
 
-        images_ = []
+    prompts = []
+    for prompt in args.prompts:
+        prompts += [prompt] * args.num_images_per_prompt
 
-        for image_idx, image in enumerate(images):
-            if isinstance(image, str):
-                image = Image.open(image)
-                image = image.convert("RGB")
-                image = image.resize((1024, 1024))
-            elif isinstance(image, Image.Image):
-                ...
-            else:
-                assert False
+    if args.use_known_negative_prompt:
+        args.negative_prompts = [known_negative_prompt]
 
-            if apply_conditioning == "canny":
-                import cv2
+    if args.negative_prompts is None:
+        negative_prompts = None
+    elif len(args.negative_prompts) == 1:
+        negative_prompts = args.negative_prompts * len(prompts)
+    elif len(args.negative_prompts) == len(args.prompts):
+        negative_prompts = []
+        for negative_prompt in args.negative_prompts:
+            negative_prompts += [negative_prompt] * args.num_images_per_prompt
+    else:
+        assert False
 
-                image = np.array(image)
-                image = cv2.Canny(image, 100, 200)
-                image = image[:, :, None]
-                controlnet_image = np.concatenate([controlnet_image, controlnet_image, controlnet_image], axis=2)
+    if args.images is not None:
+        images = []
 
+        for image_idx, image in enumerate(args.images):
+            image = Image.open(image)
+            image = image.convert("RGB")
+            image = image.resize((1024, 1024))
             image = TF.to_tensor(image)
 
-            if masks is not None:
-                mask = masks[image_idx]
-                if isinstance(mask, str):
-                    mask = Image.open(mask)
-                elif isinstance(mask, Image.Image):
-                    ...
-                else:
-                    assert False
+            if args.masks is not None:
+                mask = args.masks[image_idx]
+                mask = Image.open(mask)
                 mask = mask.convert("L")
                 mask = mask.resize((1024, 1024))
                 mask = TF.to_tensor(mask)
@@ -933,83 +968,40 @@ def gen_sdxl_simplified_interface(
                 if isinstance(controlnet, SDXLControlNetPreEncodedControlnetCond):
                     image = image * (mask < 0.5)
                     image = TF.normalize(image, [0.5], [0.5])
-                    image = vae.encode(image[None, :, :, :].to(dtype=vae.dtype, device=vae.device)).to(dtype=unet.dtype, device=unet.device)
+                    image = vae.encode(image[None, :, :, :].to(dtype=vae.dtype, device=vae.device)).to(dtype=controlnet.dtype, device=controlnet.device)
                     mask = TF.resize(mask, (1024 // 8, 1024 // 8))[None, :, :, :].to(dtype=image.dtype, device=image.device)
                     image = torch.concat((image, mask), dim=1)
                 else:
-                    image = (image * (mask < 0.5) + -1.0 * (mask >= 0.5)).to(dtype=unet.dtype, device=unet.device)
+                    image = (image * (mask < 0.5) + -1.0 * (mask >= 0.5)).to(dtype=dtype, device=device)
                     image = image[None, :, :, :]
 
-            images_.append(image)
+            images += [image] * args.num_images_per_prompt
 
-        images_ = torch.concat(images_)
+        images = torch.concat(images)
     else:
-        images_ = None
+        images = None
 
-    if isinstance(prompts, str):
-        prompts = [prompts]
-    prompts_ = []
-    for prompt in prompts:
-        prompts_ += [prompt] * num_images
-
-    if negative_prompts is not None:
-        if isinstance(negative_prompts, str):
-            negative_prompts = [negative_prompts]
-        negative_prompts_ = []
-        for negative_prompt in negative_prompts:
-            negative_prompts_ += [negative_prompt] * num_images
+    if args.seed is None:
+        generator = None
     else:
-        negative_prompts_ = None
+        generator = torch.Generator(device).manual_seed(args.seed)
 
-    x_0 = sdxl_diffusion_loop(
-        prompts=prompts_,
-        negative_prompts=negative_prompts_,
+    images = sdxl_diffusion_loop(
+        prompts=prompts,
         unet=unet,
         text_encoder_one=text_encoder_one,
         text_encoder_two=text_encoder_two,
-        sigmas=sigmas,
-        timesteps=timesteps,
+        images=images,
         controlnet=controlnet,
         adapter=adapter,
-        images=images_,
-        guidance_scale=guidance_scale,
+        sigmas=sigmas,
+        timesteps=timesteps,
+        guidance_scale=args.guidance_scale,
+        negative_prompts=negative_prompts,
+        generator=generator,
     )
 
-    x_0 = vae.decode(x_0.to(vae.dtype))
-    x_0 = vae.output_tensor_to_pil(x_0)
-
-    return x_0
-
-
-if __name__ == "__main__":
-    from argparse import ArgumentParser
-
-    args = ArgumentParser()
-    args.add_argument("--prompt", required=True, type=str)
-    args.add_argument("--num_images", required=True, type=int, default=1)
-    args.add_argument("--num_inference_steps", required=False, type=int, default=50)
-    args.add_argument("--image", required=False, type=str, default=None)
-    args.add_argument("--mask", required=False, type=str, default=None)
-    args.add_argument("--controlnet_checkpoint", required=False, type=str, default=None)
-    args.add_argument("--controlnet", required=False, choices=["SDXLControlNet", "SDXLControlNetFull", "SDXLControNetPreEncodedControlnetCond"], default=None)
-    args.add_argument("--adapter_checkpoint", required=False, type=str, default=None)
-    args.add_argument("--apply_conditioning", choices=["canny"], required=False, default=None)
-    args.add_argument("--device", required=False, default=None)
-    args = args.parse_args()
-
-    images = gen_sdxl_simplified_interface(
-        prompt=args.prompt,
-        num_images=args.num_images,
-        num_inference_steps=args.num_inference_steps,
-        images=[args.image],
-        masks=[args.mask],
-        controlnet_checkpoint=args.controlnet_checkpoint,
-        controlnet=args.controlnet,
-        adapter_checkpoint=args.adapter_checkpoint,
-        apply_conditioning=args.apply_conditioning,
-        device=args.device,
-        negative_prompt=known_negative_prompt,
-    )
+    images = vae.output_tensor_to_pil(vae.decode(images))
 
     for i, image in enumerate(images):
         image.save(f"out_{i}.png")

sdxl_models.py

@@ -1,12 +1,11 @@
 import math
 import os
-from typing import List, Literal, Optional
+from typing import List, Literal, Optional, Union
 
 import safetensors.torch
 import torch
 import torch.nn.functional as F
 import torchvision.transforms.functional as TF
-import xformers.ops
 from PIL import Image
 from torch import nn
 
@@ -21,12 +20,14 @@ class ModelUtils:
         return next(self.parameters()).device
 
     @classmethod
-    def load(cls, load_from: str, device, overrides: Optional[List[str]] = None):
+    def load(cls, load_from: str, device, overrides: Optional[Union[str, List[str]]] = None):
         import load_state_dict_patch
 
         load_from = [load_from]
 
         if overrides is not None:
+            if isinstance(overrides, str):
+                overrides = [overrides]
             load_from += overrides
 
         state_dict = {}
@@ -1323,51 +1324,57 @@ class TransformerDecoderBlock(nn.Module):
         return hidden_states
 
 
-class AttentionMixin:
-    attention_implementation: Literal["xformers", "torch_2.0_scaled_dot_product"] = "xformers"
+_attention_implementation: Literal["xformers", "torch_2.0_scaled_dot_product"] = "torch_2.0_scaled_dot_product"
 
-    @classmethod
-    def attention(cls, to_q, to_k, to_v, to_out, head_dim, hidden_states, encoder_hidden_states=None):
-        batch_size, q_seq_len, channels = hidden_states.shape
 
-        if encoder_hidden_states is not None:
-            kv = encoder_hidden_states
-        else:
-            kv = hidden_states
+def set_attention_implementation(impl: Literal["xformers", "torch_2.0_scaled_dot_product"]):
+    global _attention_implementation
+    _attention_implementation = impl
 
-        kv_seq_len = kv.shape[1]
 
-        query = to_q(hidden_states)
-        key = to_k(kv)
-        value = to_v(kv)
+def attention(to_q, to_k, to_v, to_out, head_dim, hidden_states, encoder_hidden_states=None):
+    batch_size, q_seq_len, channels = hidden_states.shape
 
-        if AttentionMixin.attention_implementation == "xformers":
-            query = query.reshape(batch_size, q_seq_len, channels // head_dim, head_dim).contiguous()
-            key = key.reshape(batch_size, kv_seq_len, channels // head_dim, head_dim).contiguous()
-            value = value.reshape(batch_size, kv_seq_len, channels // head_dim, head_dim).contiguous()
+    if encoder_hidden_states is not None:
+        kv = encoder_hidden_states
+    else:
+        kv = hidden_states
 
-            hidden_states = xformers.ops.memory_efficient_attention(query, key, value)
+    kv_seq_len = kv.shape[1]
 
-            hidden_states = hidden_states.to(query.dtype)
-            hidden_states = hidden_states.reshape(batch_size, q_seq_len, channels).contiguous()
-        elif AttentionMixin.attention_implementation == "torch_2.0_scaled_dot_product":
-            query = query.reshape(batch_size, q_seq_len, channels // head_dim, head_dim).transpose(1, 2).contiguous()
-            key = key.reshape(batch_size, kv_seq_len, channels // head_dim, head_dim).transpose(1, 2).contiguous()
-            value = value.reshape(batch_size, kv_seq_len, channels // head_dim, head_dim).transpose(1, 2).contiguous()
+    query = to_q(hidden_states)
+    key = to_k(kv)
+    value = to_v(kv)
 
-            hidden_states = F.scaled_dot_product_attention(query, key, value)
+    if _attention_implementation == "xformers":
+        import xformers.ops
 
-            hidden_states = hidden_states.to(query.dtype)
-            hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, q_seq_len, channels).contiguous()
-        else:
-            assert False
+        query = query.reshape(batch_size, q_seq_len, channels // head_dim, head_dim).contiguous()
+        key = key.reshape(batch_size, kv_seq_len, channels // head_dim, head_dim).contiguous()
+        value = value.reshape(batch_size, kv_seq_len, channels // head_dim, head_dim).contiguous()
 
-        hidden_states = to_out(hidden_states)
+        hidden_states = xformers.ops.memory_efficient_attention(query, key, value)
 
-        return hidden_states
+        hidden_states = hidden_states.to(query.dtype)
+        hidden_states = hidden_states.reshape(batch_size, q_seq_len, channels).contiguous()
+    elif _attention_implementation == "torch_2.0_scaled_dot_product":
+        query = query.reshape(batch_size, q_seq_len, channels // head_dim, head_dim).transpose(1, 2).contiguous()
+        key = key.reshape(batch_size, kv_seq_len, channels // head_dim, head_dim).transpose(1, 2).contiguous()
+        value = value.reshape(batch_size, kv_seq_len, channels // head_dim, head_dim).transpose(1, 2).contiguous()
+
+        hidden_states = F.scaled_dot_product_attention(query, key, value)
+
+        hidden_states = hidden_states.to(query.dtype)
+        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, q_seq_len, channels).contiguous()
+    else:
+        assert False
+
+    hidden_states = to_out(hidden_states)
+
+    return hidden_states
 
 
-class Attention(nn.Module, AttentionMixin):
+class Attention(nn.Module):
     def __init__(self, channels, encoder_hidden_states_dim):
         super().__init__()
         self.to_q = nn.Linear(channels, channels, bias=False)
@@ -1376,10 +1383,10 @@ class Attention(nn.Module, AttentionMixin):
         self.to_out = nn.Sequential(nn.Linear(channels, channels), nn.Dropout(0.0))
 
     def forward(self, hidden_states, encoder_hidden_states=None):
-        return self.attention(self.to_q, self.to_k, self.to_v, self.to_out, 64, hidden_states, encoder_hidden_states)
+        return attention(self.to_q, self.to_k, self.to_v, self.to_out, 64, hidden_states, encoder_hidden_states)
 
 
-class VaeMidBlockAttention(nn.Module, AttentionMixin):
+class VaeMidBlockAttention(nn.Module):
     def __init__(self, channels):
         super().__init__()
         self.group_norm = nn.GroupNorm(32, channels, eps=1e-06)
@@ -1397,7 +1404,7 @@ class VaeMidBlockAttention(nn.Module, AttentionMixin):
 
         hidden_states = self.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
 
-        hidden_states = self.attention(self.to_q, self.to_k, self.to_v, self.to_out, self.head_dim, hidden_states)
+        hidden_states = attention(self.to_q, self.to_k, self.to_v, self.to_out, self.head_dim, hidden_states)
 
         hidden_states = hidden_states.transpose(1, 2).view(batch_size, channels, height, width)