some new files

app.py

@@ -0,0 +1,305 @@
+import os
+import torch
+import gradio as gr
+import pandas as pd
+import random
+import copy
+import numpy as np
+from PIL import Image
+from torchvision import transforms
+from huggingface_hub import snapshot_download
+from diffusers import DDPMScheduler
+
+# ----------------------------
+# Load config + custom modules
+# ----------------------------
+from defaults import pose_transfer_C as cfg
+from pose_transfer_train import build_model
+from models import UNet, VariationalAutoencoder
+from pose_utils import (cords_to_map, draw_pose_from_cords, load_pose_cords_from_strings)
+
+# ----------------------------
+# Globals
+# ----------------------------
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model = None
+unet = None
+vae = None
+noise_scheduler = None
+annotation_file = None
+test_pairs = None
+model_dir = None  # Will store the downloaded model directory path
+
+# ----------------------------
+# Helper: Build pose image
+# ----------------------------
+def build_pose_img(annotation_file, img_path):
+    """Build pose image from annotation file and image path"""
+    string = annotation_file.loc[os.path.basename(img_path)]
+    array = load_pose_cords_from_strings(string['keypoints_y'], string['keypoints_x'])
+    pose_map = torch.tensor(
+        cords_to_map(array, (256, 256), (256, 176)).transpose(2, 0, 1),
+        dtype=torch.float32
+    )
+    pose_img = torch.tensor(
+        draw_pose_from_cords(array, (256, 256), (256, 176)).transpose(2, 0, 1) / 255.,
+        dtype=torch.float32
+    )
+    pose_img = torch.cat([pose_img, pose_map], dim=0)
+    return pose_img
+
+# ----------------------------
+# Model loader (runs ONCE)
+# ----------------------------
+def load_models():
+    global model, unet, vae, noise_scheduler, annotation_file, test_pairs, model_dir
+
+    if model is not None:  # already loaded
+        return
+
+    print("⏳ Downloading models & data from repository...")
+
+    # Download everything from model repository (models + fashion data)
+    repo_id = "recky101/new_l_cfld_model"
+    model_dir = snapshot_download(repo_id=repo_id)
+    
+    print(f"📁 Downloaded to: {model_dir}")
+
+    # Load schedulers & models
+    print("🔧 Loading scheduler...")
+    noise_scheduler = DDPMScheduler.from_pretrained(
+        os.path.join(model_dir, "pretrained_models/scheduler")
+    )
+    
+    print("🔧 Loading VAE...")
+    vae = VariationalAutoencoder(
+        pretrained_path=os.path.join(model_dir, "pretrained_models/vae")
+    ).eval().requires_grad_(False).to(device)
+
+    print("🔧 Loading main model...")
+    model = build_model(cfg).eval().requires_grad_(False).to(device)
+    
+    print("🔧 Loading UNet...")
+    unet = UNet(cfg).eval().requires_grad_(False).to(device)
+
+    print("📦 Loading model weights...")
+    model.load_state_dict(
+        torch.load(os.path.join(model_dir, "checkpoints/pytorch_model.bin"), map_location="cpu"),
+        strict=False
+    )
+    unet.load_state_dict(
+        torch.load(os.path.join(model_dir, "checkpoints/pytorch_model_1.bin"), map_location="cpu"),
+        strict=False
+    )
+
+    print("📊 Loading fashion dataset...")
+    # Load fashion dataset from model repository
+    test_pairs = pd.read_csv(os.path.join(model_dir, "fashion", "fasion-resize-pairs-test.csv"))
+    annotation_file = pd.read_csv(os.path.join(model_dir, "fashion", "fasion-resize-annotation-test.csv"), sep=":")
+    annotation_file = annotation_file.set_index("name")
+
+    print("✅ Everything loaded successfully!")
+    print(f"📈 Loaded {len(test_pairs)} test pairs")
+
+# ----------------------------
+# Inference function
+# ----------------------------
+def infer(img_from: Image.Image, random_index: int = None):
+    """Perform pose transfer inference"""
+    try:
+        # Load models if not already loaded
+        load_models()
+
+        if img_from is None:
+            return None
+
+        # Handle random index
+        if random_index is None:
+            random_index = random.choice(range(len(test_pairs)))
+        else:
+            # Ensure random_index is within bounds
+            random_index = max(0, min(int(random_index), len(test_pairs) - 1))
+
+        img_to_path = test_pairs.iloc[random_index]["to"]
+        print(f"🎯 Using pose from image: {img_to_path} (index: {random_index})")
+
+        # Preprocess source image
+        trans = transforms.Compose([
+            transforms.Resize([256, 256], interpolation=transforms.InterpolationMode.BICUBIC, antialias=True),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
+        ])
+        img_from_tensor = trans(img_from).unsqueeze(0).to(device)
+
+        # Build pose image tensor
+        pose_img_tensor = build_pose_img(annotation_file, img_to_path).unsqueeze(0).to(device)
+
+        print("🚀 Running inference...")
+        # Inference
+        with torch.no_grad():
+            c_new, down_block_additional_residuals, up_block_additional_residuals = model({
+                "img_cond": img_from_tensor, "pose_img": pose_img_tensor
+            })
+            noisy_latents = torch.randn((1, 4, 64, 64)).to(device)
+            weight_dtype = torch.float32
+            bsz = 1
+
+            c_new = torch.cat([c_new[:bsz], c_new[:bsz], c_new[bsz:]])
+            down_block_additional_residuals = [
+                torch.cat([torch.zeros_like(sample), sample, sample]).to(dtype=weight_dtype)
+                for sample in down_block_additional_residuals
+            ]
+            up_block_additional_residuals = {
+                k: torch.cat([torch.zeros_like(v), torch.zeros_like(v), v]).to(dtype=weight_dtype)
+                for k, v in up_block_additional_residuals.items()
+            }
+
+            noise_scheduler.set_timesteps(cfg.TEST.NUM_INFERENCE_STEPS)
+            for t in noise_scheduler.timesteps:
+                inputs = torch.cat([noisy_latents, noisy_latents, noisy_latents], dim=0)
+                inputs = noise_scheduler.scale_model_input(inputs, timestep=t)
+                noise_pred = unet(
+                    sample=inputs, 
+                    timestep=t, 
+                    encoder_hidden_states=c_new,
+                    down_block_additional_residuals=copy.deepcopy(down_block_additional_residuals),
+                    up_block_additional_residuals=copy.deepcopy(up_block_additional_residuals)
+                )
+
+                noise_pred_uc, noise_pred_down, noise_pred_full = noise_pred.chunk(3)
+                noise_pred = noise_pred_uc + \
+                            cfg.TEST.DOWN_BLOCK_GUIDANCE_SCALE * (noise_pred_down - noise_pred_uc) + \
+                            cfg.TEST.FULL_GUIDANCE_SCALE * (noise_pred_full - noise_pred_down)
+
+                noisy_latents = noise_scheduler.step(noise_pred, t, noisy_latents)[0]
+
+            sampling_imgs = vae.decode(noisy_latents) * 0.5 + 0.5
+            sampling_imgs = sampling_imgs.clamp(0, 1)
+
+        # Convert to PIL image
+        output_img = Image.fromarray(
+            (sampling_imgs[0] * 255.).permute((1, 2, 0)).long().cpu().numpy().astype(np.uint8)
+        ).resize((256, 256))
+
+        print("✅ Inference completed successfully!")
+        return output_img
+
+    except Exception as e:
+        print(f"❌ Error in inference: {e}")
+        import traceback
+        traceback.print_exc()
+        return None
+
+# ----------------------------
+# Gradio Interface
+# ----------------------------
+with gr.Blocks(
+    title="CFLD Pose Transfer Demo",
+    theme=gr.themes.Soft(),
+    css="""
+    .gradio-container {
+        max-width: 1200px !important;
+    }
+    .gr-button-primary {
+        background: linear-gradient(90deg, #ff6b6b, #4ecdc4) !important;
+        border: none !important;
+    }
+    """
+) as demo:
+    gr.Markdown(
+        """
+        # 👗 CFLD Pose Transfer Demo
+        
+        Upload a person image and transfer their pose using our CFLD (Controllable Fashion Layout Diffusion) model!
+        
+        🎯 **How it works:** Upload an image → Select a target pose (or use random) → Get pose-transferred result!
+        """
+    )
+    
+    # Status indicator
+    with gr.Row():
+        status_text = gr.Markdown("🔄 **Status:** Loading models... Please wait.")
+    
+    with gr.Row(equal_height=True):
+        with gr.Column(scale=1):
+            gr.Markdown("### 📤 Input")
+            inp = gr.Image(
+                label="Upload Source Image", 
+                type="pil",
+                height=400
+            )
+            
+            idx = gr.Number(
+                label="Target Pose Index (0-4499, leave blank for random)", 
+                value=None,
+                precision=0,
+                minimum=0,
+                maximum=4499
+            )
+            
+            btn = gr.Button(
+                "🚀 Generate Pose Transfer", 
+                variant="primary",
+                size="lg"
+            )
+            
+            gr.Markdown(
+                """
+                ### 💡 Tips:
+                - Upload clear images of people
+                - Works best with full-body or upper-body shots
+                - Try different pose indices for variety
+                - Leave index blank for random poses
+                """
+            )
+        
+        with gr.Column(scale=1):
+            gr.Markdown("### 📥 Result")
+            out = gr.Image(
+                label="Pose Transferred Result",
+                height=400
+            )
+            
+            # Progress and info
+            info_text = gr.Markdown("Upload an image and click generate to start!")
+
+    # Event handlers
+    def update_status():
+        if model is not None:
+            return "✅ **Status:** Models loaded and ready!"
+        else:
+            return "🔄 **Status:** Loading models... Please wait."
+    
+    def infer_with_status(img, idx):
+        if img is None:
+            return None, "❌ Please upload an image first!"
+        
+        if model is None:
+            return None, "⏳ Models are still loading, please wait..."
+        
+        result = infer(img, idx)
+        
+        if result is None:
+            return None, "❌ Failed to generate result. Please try again."
+        
+        used_idx = idx if idx is not None else "random"
+        return result, f"✅ Success! Used pose index: {used_idx}"
+    
+    # Connect events
+    btn.click(
+        fn=infer_with_status,
+        inputs=[inp, idx],
+        outputs=[out, info_text]
+    )
+    
+    # Auto-update status every few seconds
+    demo.load(
+        fn=update_status,
+        outputs=[status_text],
+        every=3
+    )
+
+# ----------------------------
+# Launch the demo
+# ----------------------------
+demo.launch()

configs/ablation_study/ape.yaml

@@ -0,0 +1,17 @@
+ACCELERATE:
+  RUN_NAME: "ape"
+  EVAL_PERIOD: 10
+
+MODEL:
+  APPEARANCE_GUIDANCE_CONFIG:
+    CONVIN_KERNEL_SIZE: [8, 8, 8, 4, 4, 4, 2, 2, 2]
+    CONVIN_STRIDE: [8, 8, 8, 4, 4, 4, 2, 2, 2]
+    CONVIN_PADDING: [0, 0, 0, 0, 0, 0, 0, 0, 0]
+    CTX_DIMS: [768, 768, 768, 768, 768, 768, 768, 768, 768]
+    TO_QUERIES: False
+    TO_KEYS: True
+    TO_VALUES: True
+
+  DECODER_CONFIG:
+    N_CTX: 64
+    DEPTH: -1

configs/ablation_study/clip_model.yaml

@@ -0,0 +1,25 @@
+ACCELERATE:
+  RUN_NAME: "clip_model"
+  EVAL_PERIOD: 10
+
+MODEL:
+  COND_STAGE_CONFIG:
+    DEPTHS: []
+
+  APPEARANCE_GUIDANCE_CONFIG:
+    ATTN_RESIDUAL_BLOCK_IDX: []
+    INNER_DIMS: []
+    CTX_DIMS: []
+    EMBED_DIMS: []
+    HEADS: []
+    CONVIN_KERNEL_SIZE: []
+    CONVIN_STRIDE: []
+    CONVIN_PADDING: []
+
+  DECODER_CONFIG:
+    N_CTX: 1
+    DEPTH: 0
+
+INPUT:
+  COND:
+    IMG_SIZE: [224, 224]

configs/ablation_study/no_app.yaml

@@ -0,0 +1,14 @@
+ACCELERATE:
+  RUN_NAME: "no_app"
+  EVAL_PERIOD: 10
+
+MODEL:
+  APPEARANCE_GUIDANCE_CONFIG:
+    ATTN_RESIDUAL_BLOCK_IDX: []
+    INNER_DIMS: []
+    CTX_DIMS: []
+    EMBED_DIMS: []
+    HEADS: []
+    CONVIN_KERNEL_SIZE: []
+    CONVIN_STRIDE: []
+    CONVIN_PADDING: []

configs/ablation_study/no_app_trainq.yaml

@@ -0,0 +1,17 @@
+ACCELERATE:
+  RUN_NAME: "no_app_trainq"
+  EVAL_PERIOD: 10
+
+MODEL:
+  UNET_CONFIG:
+    TRAIN_CROSS_ATTN_Q: True
+
+  APPEARANCE_GUIDANCE_CONFIG:
+    ATTN_RESIDUAL_BLOCK_IDX: []
+    INNER_DIMS: []
+    CTX_DIMS: []
+    EMBED_DIMS: []
+    HEADS: []
+    CONVIN_KERNEL_SIZE: []
+    CONVIN_STRIDE: []
+    CONVIN_PADDING: []

configs/ablation_study/swin.yaml

@@ -0,0 +1,8 @@
+ACCELERATE:
+  RUN_NAME: "swin"
+  EVAL_PERIOD: 10
+
+MODEL:
+  DECODER_CONFIG:
+    N_CTX: 64
+    DEPTH: -2

configs/fashion_256.yaml

@@ -0,0 +1,2 @@
+TEST:
+  IMG_SIZE: [256, 176]

configs/fashion_512.yaml

@@ -0,0 +1,2 @@
+TEST:
+  IMG_SIZE: [512, 352]

datasets/__init__.py

@@ -0,0 +1,2 @@
+from .deepfashion import (FidRealDeepFashion, PisTestDeepFashion,
+                          PisTrainDeepFashion)

datasets/deepfashion.py

@@ -0,0 +1,258 @@
+"""
+@author: Yanzuo Lu
+@author: oliveryanzuolu@gmail.com
+"""
+
+import glob
+import logging
+import math
+import os
+import random
+
+import numpy as np
+import pandas as pd
+import torch
+import torchvision.transforms as transforms
+from PIL import Image
+from torch.utils.data import Dataset
+
+from pose_utils import (cords_to_map, draw_pose_from_cords,
+                        load_pose_cords_from_strings)
+
+logger = logging.getLogger()
+
+
+class PisTrainDeepFashion(Dataset):
+    def __init__(self, root_dir, gt_img_size, pose_img_size, cond_img_size, min_scale,
+                 log_aspect_ratio, pred_ratio, pred_ratio_var, psz):
+        super().__init__()
+        self.pose_img_size = pose_img_size
+        self.cond_img_size = cond_img_size
+        self.log_aspect_ratio = log_aspect_ratio
+        self.pred_ratio = pred_ratio
+        self.pred_ratio_var = pred_ratio_var
+        self.psz = psz
+
+        # root_dir = os.path.join(root_dir, "DeepFashion")
+        train_dir = os.path.join(root_dir, "train_highres")
+        train_pairs = os.path.join(root_dir, "fasion-resize-pairs-train.csv")
+        train_pairs = pd.read_csv(train_pairs)
+        self.img_items = self.process_dir(train_dir, train_pairs)
+
+        self.annotation_file = pd.read_csv(os.path.join(root_dir, "fasion-resize-annotation-train.csv"), sep=':')
+        self.annotation_file = self.annotation_file.set_index('name')
+
+        self.transform_gt = transforms.Compose([
+            transforms.Resize(gt_img_size, interpolation=transforms.InterpolationMode.BICUBIC, antialias=True),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
+        ])
+        self.transform_cond = transforms.Compose([
+            transforms.Resize(cond_img_size, interpolation=transforms.InterpolationMode.BICUBIC, antialias=True),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
+        ])
+
+        aspect_ratio = cond_img_size[1] / cond_img_size[0]
+        self.transform = transforms.Compose([
+            transforms.RandomResizedCrop(cond_img_size, scale=(min_scale, 1.), ratio=(aspect_ratio*3./4., aspect_ratio*4./3.),
+                                         interpolation=transforms.InterpolationMode.BICUBIC, antialias=True),
+            transforms.RandomHorizontalFlip(p=0.5),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
+        ]) if min_scale < 1.0 else transforms.Compose([
+            transforms.Resize(cond_img_size, interpolation=transforms.InterpolationMode.BICUBIC, antialias=True),
+            transforms.RandomHorizontalFlip(p=0.5),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
+        ])
+
+    def process_dir(self, root_dir, csv_file):
+        data = []
+        for i in range(len(csv_file)):
+            data.append((os.path.join(root_dir, csv_file.iloc[i]["from"]),
+                         os.path.join(root_dir, csv_file.iloc[i]["to"])))
+        return data
+
+    def get_pred_ratio(self):
+        pred_ratio = []
+        for prm, prv in zip(self.pred_ratio, self.pred_ratio_var):
+            assert prm >= prv
+            pr = random.uniform(prm - prv, prm + prv) if prv > 0 else prm
+            pred_ratio.append(pr)
+        pred_ratio = random.choice(pred_ratio)
+        return pred_ratio
+
+    def __len__(self):
+        return len(self.img_items)
+
+    def __getitem__(self, index):
+        img_path_from, img_path_to = self.img_items[index]
+        with open(img_path_from, 'rb') as f:
+            img_from = Image.open(f).convert('RGB')
+        with open(img_path_to, 'rb') as f:
+            img_to = Image.open(f).convert('RGB')
+
+        img_src = self.transform_gt(img_from)
+        img_tgt = self.transform_gt(img_to)
+        img_cond = self.transform(img_from)
+        pose_img_src = self.build_pose_img(img_path_from)
+        pose_img_tgt = self.build_pose_img(img_path_to)
+
+        mask = None
+        if len(self.pred_ratio) > 0:
+            H, W = self.cond_img_size[0] // self.psz, self.cond_img_size[1] // self.psz
+            high = self.get_pred_ratio() * H * W
+
+            # following BEiT (https://arxiv.org/abs/2106.08254), see at
+            # https://github.com/microsoft/unilm/blob/b94ec76c36f02fb2b0bf0dcb0b8554a2185173cd/beit/masking_generator.py#L55
+            mask = np.zeros((H, W), dtype=bool)
+            mask_count = 0
+            while mask_count < high:
+                max_mask_patches = high - mask_count
+
+                delta = 0
+                for attempt in range(10):
+                    low = (min(H, W) // 3) ** 2
+                    target_area = random.uniform(low, max_mask_patches)
+                    aspect_ratio = math.exp(random.uniform(*self.log_aspect_ratio))
+                    h = int(round(math.sqrt(target_area * aspect_ratio)))
+                    w = int(round(math.sqrt(target_area / aspect_ratio)))
+                    if w < W and h < H:
+                        top = random.randint(0, H - h)
+                        left = random.randint(0, W - w)
+
+                        num_masked = mask[top: top + h, left: left + w].sum()
+                        if 0 < h * w - num_masked <= max_mask_patches:
+                            for i in range(top, top + h):
+                                for j in range(left, left + w):
+                                    if mask[i, j] == 0:
+                                        mask[i, j] = 1
+                                        delta += 1
+
+                    if delta > 0:
+                        break
+
+                if delta == 0:
+                    break
+                else:
+                    mask_count += delta
+
+        return_dict = {
+            "img_src": img_src,
+            "img_tgt": img_tgt,
+            "img_cond": img_cond,
+            "pose_img_src": pose_img_src,
+            "pose_img_tgt": pose_img_tgt
+        }
+        if mask is not None:
+            return_dict["mask"] = mask
+        return return_dict
+
+    def build_pose_img(self, img_path):
+        string = self.annotation_file.loc[os.path.basename(img_path)]
+        array = load_pose_cords_from_strings(string['keypoints_y'], string['keypoints_x'])
+        pose_map = torch.tensor(cords_to_map(array, tuple(self.pose_img_size), (256, 176)).transpose(2, 0, 1), dtype=torch.float32)
+        pose_img = torch.tensor(draw_pose_from_cords(array, tuple(self.pose_img_size), (256, 176)).transpose(2, 0, 1) / 255., dtype=torch.float32)
+        pose_img = torch.cat([pose_img, pose_map], dim=0)
+        return pose_img
+
+
+class PisTestDeepFashion(Dataset):
+    def __init__(self, root_dir, gt_img_size, pose_img_size, cond_img_size, test_img_size):
+        super().__init__()
+        self.pose_img_size = pose_img_size
+
+        # root_dir = os.path.join(root_dir, "DeepFashion")
+        test_pairs = os.path.join(root_dir, "fasion-resize-pairs-test.csv")
+        test_pairs = pd.read_csv(test_pairs)
+        self.img_items = self.process_dir(root_dir, test_pairs)
+
+        self.annotation_file = pd.read_csv(os.path.join(root_dir, "fasion-resize-annotation-test.csv"), sep=':')
+        self.annotation_file = self.annotation_file.set_index('name')
+
+        self.transform_gt = transforms.Compose([
+            transforms.Resize(gt_img_size, interpolation=transforms.InterpolationMode.BICUBIC, antialias=True),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
+        ])
+        self.transform_cond = transforms.Compose([
+            transforms.Resize(cond_img_size, interpolation=transforms.InterpolationMode.BICUBIC, antialias=True),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
+        ])
+        self.transform_test = transforms.Compose([
+            transforms.Resize(test_img_size, interpolation=transforms.InterpolationMode.BICUBIC, antialias=True),
+            transforms.ToTensor()
+        ])
+
+    def process_dir(self, root_dir, csv_file):
+        data = []
+        for i in range(len(csv_file)):
+            data.append((os.path.join(root_dir, "test_highres", csv_file.iloc[i]["from"]),
+                         os.path.join(root_dir, "test_highres", csv_file.iloc[i]["to"])))
+        return data
+
+    def __len__(self):
+        return len(self.img_items)
+
+    def __getitem__(self, index):
+        img_path_from, img_path_to = self.img_items[index]
+        with open(img_path_from, 'rb') as f:
+            img_from = Image.open(f).convert('RGB')
+        with open(img_path_to, 'rb') as f:
+            img_to = Image.open(f).convert('RGB')
+
+        img_src = self.transform_gt(img_from) # for visualization
+        img_tgt = self.transform_gt(img_to) # for visualization
+        img_gt = self.transform_test(img_to) # for metrics, 3x256x176
+        img_cond_from = self.transform_cond(img_from)
+
+        pose_img_from = self.build_pose_img(img_path_from)
+        pose_img_to = self.build_pose_img(img_path_to)
+
+        return {
+            "img_src": img_src,
+            "img_tgt": img_tgt,
+            "img_gt": img_gt,
+            "img_cond_from": img_cond_from,
+            "pose_img_from": pose_img_from,
+            "pose_img_to": pose_img_to
+        }
+
+    def build_pose_img(self, img_path):
+        string = self.annotation_file.loc[os.path.basename(img_path)]
+        array = load_pose_cords_from_strings(string['keypoints_y'], string['keypoints_x'])
+        pose_map = torch.tensor(cords_to_map(array, tuple(self.pose_img_size), (256, 176)).transpose(2, 0, 1), dtype=torch.float32)
+        pose_img = torch.tensor(draw_pose_from_cords(array, tuple(self.pose_img_size), (256, 176)).transpose(2, 0, 1) / 255., dtype=torch.float32)
+        pose_img = torch.cat([pose_img, pose_map], dim=0)
+        return pose_img
+
+
+class FidRealDeepFashion(Dataset):
+    def __init__(self, root_dir, test_img_size):
+        super().__init__()
+        # root_dir = os.path.join(root_dir, "DeepFashion")
+        train_dir = os.path.join(root_dir, "train_highres")
+        self.img_items = self.process_dir(train_dir)
+
+        self.transform_test = transforms.Compose([
+            transforms.Resize(test_img_size, interpolation=transforms.InterpolationMode.BICUBIC, antialias=True),
+            transforms.ToTensor()
+        ])
+
+    def process_dir(self, root_dir):
+        data = []
+        img_paths = glob.glob(os.path.join(root_dir, '*.jpg'))
+        for img_path in img_paths:
+            data.append(img_path)
+        return data
+
+    def __len__(self):
+        return len(self.img_items)
+
+    def __getitem__(self, index):
+        img_path = self.img_items[index]
+        with open(img_path, 'rb') as f:
+            img = Image.open(f).convert('RGB')
+        return self.transform_test(img)

defaults/__init__.py

@@ -0,0 +1 @@
+from .deepfashion import _C as pose_transfer_C

defaults/deepfashion.py

@@ -0,0 +1,125 @@
+from yacs.config import CfgNode as CN
+
+_C = CN()
+
+_C.ACCELERATE = CN()
+_C.ACCELERATE.PROJECT_NAME = "CFLD"
+_C.ACCELERATE.RUN_NAME = "debug"
+_C.ACCELERATE.MIXED_PRECISION = "fp16"
+_C.ACCELERATE.ALLOW_TF32 = True
+_C.ACCELERATE.SEED = 3407
+_C.ACCELERATE.GRADIENT_ACCUMULATION_STEPS = 1
+_C.ACCELERATE.LOG_PERIOD = 10
+_C.ACCELERATE.EVAL_PERIOD = 5
+
+_C.MODEL = CN()
+_C.MODEL.PRETRAINED_PATH = ""
+_C.MODEL.LAST_EPOCH = 0
+_C.MODEL.U_COND_PERCENT = 0.2
+_C.MODEL.U_COND_DOWN_BLOCK_GUIDANCE = False
+_C.MODEL.U_COND_UP_BLOCK_GUIDANCE = False
+
+_C.MODEL.FIRST_STAGE_CONFIG = CN()
+_C.MODEL.FIRST_STAGE_CONFIG.PRETRAINED_PATH = "pretrained_models/vae"
+
+_C.MODEL.UNET_CONFIG = CN()
+_C.MODEL.UNET_CONFIG.PRETRAINED_PATH = "pretrained_models/unet"
+_C.MODEL.UNET_CONFIG.TRAINABLE_BLOCK_IDX = [11, 10, 9, 8, 7, 6, 5, 4, 3]
+_C.MODEL.UNET_CONFIG.TRAIN_SELF_ATTN_Q = False
+_C.MODEL.UNET_CONFIG.TRAIN_SELF_ATTN_K = False
+_C.MODEL.UNET_CONFIG.TRAIN_SELF_ATTN_V = False
+_C.MODEL.UNET_CONFIG.TRAIN_CROSS_ATTN_Q = False
+_C.MODEL.UNET_CONFIG.TRAIN_CROSS_ATTN_K = True
+_C.MODEL.UNET_CONFIG.TRAIN_CROSS_ATTN_V = True
+
+_C.MODEL.SCHEDULER_CONFIG = CN()
+_C.MODEL.SCHEDULER_CONFIG.NAME = "ddpm"
+_C.MODEL.SCHEDULER_CONFIG.PRETRAINED_PATH = "pretrained_models/scheduler"
+_C.MODEL.SCHEDULER_CONFIG.CUBIC_SAMPLING = True
+
+_C.MODEL.COND_STAGE_CONFIG = CN()
+_C.MODEL.COND_STAGE_CONFIG.PRETRAINED_PATH = "pretrained_models/swin/swin_base_patch4_window12_384_22kto1k.pth"
+_C.MODEL.COND_STAGE_CONFIG.EMBED_DIM = 128
+_C.MODEL.COND_STAGE_CONFIG.DEPTHS = [2, 2, 18, 2]
+_C.MODEL.COND_STAGE_CONFIG.NUM_HEADS = [4, 8, 16, 32]
+_C.MODEL.COND_STAGE_CONFIG.WINDOW_SIZE = 16
+_C.MODEL.COND_STAGE_CONFIG.DROP_PATH_RATE = 0.2
+_C.MODEL.COND_STAGE_CONFIG.LAST_NORM = False
+
+_C.MODEL.APPEARANCE_GUIDANCE_CONFIG = CN()
+_C.MODEL.APPEARANCE_GUIDANCE_CONFIG.CONVIN_KERNEL_SIZE = [1, 1, 1, 1, 1, 1, 1, 1, 1]
+_C.MODEL.APPEARANCE_GUIDANCE_CONFIG.CONVIN_STRIDE = [1, 1, 1, 1, 1, 1, 1, 1, 1]
+_C.MODEL.APPEARANCE_GUIDANCE_CONFIG.CONVIN_PADDING = [0, 0, 0, 0, 0, 0, 0, 0, 0]
+_C.MODEL.APPEARANCE_GUIDANCE_CONFIG.ATTN_RESIDUAL_BLOCK_IDX = [11, 10, 9, 8, 7, 6, 5, 4, 3]
+_C.MODEL.APPEARANCE_GUIDANCE_CONFIG.INNER_DIMS = [128, 128, 128, 256, 256, 256, 512, 512, 512]
+_C.MODEL.APPEARANCE_GUIDANCE_CONFIG.CTX_DIMS = [320, 320, 320, 640, 640, 640, 1280, 1280, 1280]
+_C.MODEL.APPEARANCE_GUIDANCE_CONFIG.EMBED_DIMS = [64, 64, 64, 128, 128, 128, 256, 256, 256]
+_C.MODEL.APPEARANCE_GUIDANCE_CONFIG.HEADS = [2, 2, 2, 4, 4, 4, 8, 8, 8]
+_C.MODEL.APPEARANCE_GUIDANCE_CONFIG.DEPTH = 4
+_C.MODEL.APPEARANCE_GUIDANCE_CONFIG.TO_SELF_ATTN = False
+_C.MODEL.APPEARANCE_GUIDANCE_CONFIG.TO_QUERIES = True
+_C.MODEL.APPEARANCE_GUIDANCE_CONFIG.TO_KEYS = False
+_C.MODEL.APPEARANCE_GUIDANCE_CONFIG.TO_VALUES = False
+_C.MODEL.APPEARANCE_GUIDANCE_CONFIG.DETACH_INPUT = False
+
+_C.MODEL.POSE_GUIDANCE_CONFIG = CN()
+_C.MODEL.POSE_GUIDANCE_CONFIG.DOWNSCALE_FACTOR = 4
+_C.MODEL.POSE_GUIDANCE_CONFIG.POSE_CHANNELS = 21
+_C.MODEL.POSE_GUIDANCE_CONFIG.IN_CHANNELS = 320
+_C.MODEL.POSE_GUIDANCE_CONFIG.CHANNELS = [320, 640, 1280]
+
+_C.MODEL.DECODER_CONFIG = CN()
+_C.MODEL.DECODER_CONFIG.N_CTX = 16
+_C.MODEL.DECODER_CONFIG.CTX_DIM = 768
+_C.MODEL.DECODER_CONFIG.DEPTH = 8
+_C.MODEL.DECODER_CONFIG.HEADS = 24
+_C.MODEL.DECODER_CONFIG.POSE_QUERY = False
+
+_C.OPTIMIZER = CN()
+_C.OPTIMIZER.NAME = "adam"
+_C.OPTIMIZER.EPOCHS = 100
+_C.OPTIMIZER.WARMUP_STEPS = 1000
+_C.OPTIMIZER.DECAY_EPOCHS = [50]
+_C.OPTIMIZER.LR = 1.0e-4
+_C.OPTIMIZER.SCALE_LR = False
+_C.OPTIMIZER.WARMUP_RATE = 0.1
+_C.OPTIMIZER.DECAY_RATE = 0.1
+_C.OPTIMIZER.OVERRIDE_LR = 0.
+
+_C.INPUT = CN()
+_C.INPUT.ROOT_DIR = "fashion"
+_C.INPUT.BATCH_SIZE = 224
+_C.INPUT.NUM_WORKERS = 8
+
+_C.INPUT.GT = CN()
+_C.INPUT.GT.IMG_SIZE = [512, 512]
+
+_C.INPUT.COND = CN()
+_C.INPUT.COND.IMG_SIZE = [256, 256]
+_C.INPUT.COND.PRED_ASPECT_RATIO = [0.3, 1/0.3]
+_C.INPUT.COND.PRED_RATIO = []
+_C.INPUT.COND.PRED_RATIO_VAR = []
+_C.INPUT.COND.MASK_PATCH_SIZE = 8
+_C.INPUT.COND.MIN_SCALE = 1.0
+
+_C.INPUT.POSE = CN()
+_C.INPUT.POSE.IMG_SIZE = [256, 256]
+
+_C.TEST = CN()
+_C.TEST.NUM_INFERENCE_STEPS = 50
+_C.TEST.MICRO_BATCH_SIZE = 16
+_C.TEST.NUM_WORKERS = 8
+_C.TEST.IMG_SIZE = [256, 176]
+
+_C.TEST.DDIM_INVERSION_STEPS = 0
+_C.TEST.DDIM_INVERSION_DOWN_BLOCK_GUIDANCE = False
+_C.TEST.DDIM_INVERSION_UP_BLOCK_GUIDANCE = False
+_C.TEST.DDIM_INVERSION_UNCONDITIONAL = True
+
+# "uc_full", "updown_full", "down_full", "uc_down_full", "uc_down_updown_cdown", "uc_down_updown_full"
+_C.TEST.GUIDANCE_TYPE = "uc_down_full"
+_C.TEST.GUIDANCE_SCALE = 2.0
+_C.TEST.DOWN_BLOCK_GUIDANCE_SCALE = 2.0
+_C.TEST.UP_BLOCK_GUIDANCE_SCALE = 2.0
+_C.TEST.ALL_BLOCK_GUIDANCE_SCALE = 2.0
+_C.TEST.FULL_GUIDANCE_SCALE = 2.0

generate_fashion_datasets.py

@@ -0,0 +1,59 @@
+import os
+import shutil
+
+IMG_EXTENSIONS = [
+    '.jpg', '.JPG', '.jpeg', '.JPEG',
+    '.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP',
+]
+
+def is_image_file(filename):
+	return any(filename.endswith(extension) for extension in IMG_EXTENSIONS)
+
+
+def make_dataset(dir):
+	assert os.path.isdir(dir), '%s is not a valid directory' % dir
+
+	train_root = './fashion/train_highres'
+	if not os.path.exists(train_root):
+		os.mkdir(train_root)
+
+	test_root = './fashion/test_highres'
+	if not os.path.exists(test_root):
+		os.mkdir(test_root)
+
+	train_images = []
+	train_f = open('./fashion/train.lst', 'r')
+	for lines in train_f:
+		lines = lines.strip()
+		if lines.endswith('.jpg'):
+			train_images.append(lines)
+
+	test_images = []
+	test_f = open('./fashion/test.lst', 'r')
+	for lines in test_f:
+		lines = lines.strip()
+		if lines.endswith('.jpg'):
+			test_images.append(lines)
+
+
+	print("Walaking Direstory: ", dir)
+	
+	for root, _, fnames in sorted(os.walk(dir)):
+		print('root:', root)	
+		print('fnames:', fnames)
+		for fname in fnames:
+
+			if is_image_file(fname):
+				path = os.path.join(root, fname)
+				path_names = path.split('\\')
+			
+				print("pathNames = ", path_names)
+				path_names[3] = path_names[3].replace('_', '')
+				path_names[4] = path_names[4].split('_')[0] + "_" + "".join(path_names[4].split('_')[1:])
+				path_names = "".join(path_names)
+				if path_names in train_images:
+					shutil.copy(path, os.path.join(train_root, path_names))
+				if path_names in test_images:
+					shutil.copy(path, os.path.join(test_root, path_names))
+
+make_dataset('fashion')

lr_scheduler.py

@@ -0,0 +1,32 @@
+"""
+@author: Yanzuo Lu
+@author: oliveryanzuolu@gmail.com
+"""
+
+from bisect import bisect_right
+
+import torch.optim.lr_scheduler
+
+
+class LinearWarmupMultiStepDecayLRScheduler(torch.optim.lr_scheduler._LRScheduler):
+    def __init__(self, optimizer, warmup_steps, warmup_rate, decay_rate,
+                 num_epochs, decay_epochs, iters_per_epoch, override_lr=0.,
+                 last_epoch=-1, verbose=False):
+        self.warmup_steps = warmup_steps
+        self.warmup_rate = warmup_rate
+        self.decay_rate = decay_rate
+        self.decay_epochs = [decay_epoch * iters_per_epoch for decay_epoch in decay_epochs]
+        self.num_epochs = num_epochs * iters_per_epoch
+        self.override_lr = override_lr
+        super(LinearWarmupMultiStepDecayLRScheduler, self).__init__(optimizer, last_epoch, verbose)
+
+    def get_lr(self):
+        if self.last_epoch < self.warmup_steps:
+            alpha = (self.last_epoch + 1) / self.warmup_steps
+            return [base_lr * (self.warmup_rate + (1. - self.warmup_rate) * alpha) \
+                    for base_lr in self.base_lrs]
+        else:
+            if self.override_lr > 0.:
+                return [self.override_lr for _ in self.base_lrs]
+            e = bisect_right(self.decay_epochs, self.last_epoch)
+            return [base_lr * (self.decay_rate ** e) for base_lr in self.base_lrs]

models/__init__.py

@@ -0,0 +1,7 @@
+from .appearance_encoder import AppearanceEncoder
+from .decoder import Decoder
+from .metrics import build_metric
+from .pose_encoder import PoseEncoder
+from .swin_transformer import build_backbone
+from .unet import UNet
+from .vae import VariationalAutoencoder

models/appearance_encoder.py

@@ -0,0 +1,103 @@
+"""
+@author: Yanzuo Lu
+@author: oliveryanzuolu@gmail.com
+"""
+
+import torch
+import torch.nn as nn
+from diffusers.models.attention import BasicTransformerBlock
+
+
+class AppearanceEncoder(nn.Module):
+    def __init__(self, attn_residual_block_idx, inner_dims, ctx_dims, embed_dims, heads, depth,
+                 to_self_attn, to_queries, to_keys, to_values, aspect_ratio, detach_input,
+                 convin_kernel_size, convin_stride, convin_padding):
+        super().__init__()
+        self.attn_residual_block_idx = attn_residual_block_idx
+        self.inner_dims = inner_dims
+        self.ctx_dims = ctx_dims
+        self.embed_dims = embed_dims
+        self.to_self_attn = to_self_attn
+        self.to_queries = to_queries
+        self.to_keys = to_keys
+        self.to_values = to_values
+        self.aspect_ratio = aspect_ratio
+        self.detach_input = detach_input
+
+        self.zero_conv_ins = []
+        self.zero_conv_outs = []
+        self.blocks = []
+        for inner_dim, embed_dim, ctx_dim, num_head, kernel_size, stride, padding in \
+            zip(inner_dims, self.embed_dims, self.ctx_dims, heads, convin_kernel_size, convin_stride, convin_padding):
+            self.zero_conv_ins.append(nn.Conv2d(inner_dim, embed_dim, kernel_size=kernel_size,
+                                                stride=stride, padding=padding))
+            self.zero_conv_outs.append(nn.Conv2d(embed_dim, ctx_dim, kernel_size=1, stride=1, padding=0))
+            self.blocks.append(nn.Sequential(*[BasicTransformerBlock(
+                dim=embed_dim,
+                num_attention_heads=num_head,
+                attention_head_dim=embed_dim//num_head,
+                double_self_attention=True
+            ) for _ in range(depth)]))
+
+        self.blocks = nn.ModuleList(self.blocks)
+        self.zero_conv_ins = nn.ModuleList(self.zero_conv_ins)
+        self.zero_conv_outs = nn.ModuleList(self.zero_conv_outs)
+
+        for n in self.zero_conv_ins.parameters():
+            nn.init.zeros_(n)
+        for n in self.zero_conv_outs.parameters():
+            nn.init.zeros_(n)
+
+        # enable xformers
+        def fn_recursive_set_mem_eff(module: torch.nn.Module):
+            if hasattr(module, "set_use_memory_efficient_attention_xformers"):
+                module.set_use_memory_efficient_attention_xformers(True, attention_op=None)
+
+            for child in module.children():
+                fn_recursive_set_mem_eff(child)
+
+        for module in self.children():
+            if isinstance(module, torch.nn.Module):
+                fn_recursive_set_mem_eff(module)
+
+    def forward(self, features):
+        additional_residuals = {}
+
+        for i, block in enumerate(self.blocks):
+            hidden_states = features[0]
+            if self.detach_input:
+                hidden_states = hidden_states.detach()
+
+            in_H = in_W = int(features[0].shape[1] ** 0.5)
+            hidden_states = features[0].permute(0, 2, 1).reshape(-1, self.inner_dims[i], in_H, in_W)
+            hidden_states = self.zero_conv_ins[i](hidden_states)
+            H = W = hidden_states.shape[2]
+            hidden_states = hidden_states.reshape(-1, self.embed_dims[i], H * W).permute(0, 2, 1)
+
+            hidden_states = block(hidden_states)
+
+            hidden_states = hidden_states.permute(0, 2, 1).reshape(-1, self.embed_dims[i], H, W)
+            hidden_states = self.zero_conv_outs[i](hidden_states)
+            hidden_states = hidden_states.reshape(-1, self.ctx_dims[i], H * W).permute(0, 2, 1)
+
+            if self.to_self_attn:
+                if self.to_queries:
+                    additional_residuals[f"block_{self.attn_residual_block_idx[i]}_self_attn_q"] = hidden_states
+                elif self.to_keys:
+                    additional_residuals[f"block_{self.attn_residual_block_idx[i]}_self_attn_k"] = hidden_states
+                elif self.to_values:
+                    additional_residuals[f"block_{self.attn_residual_block_idx[i]}_self_attn_v"] = hidden_states
+            else:
+                if self.to_keys and self.to_values:
+                    additional_residuals[f"block_{self.attn_residual_block_idx[i]}_cross_attn_c"] = hidden_states
+                elif self.to_queries:
+                    additional_residuals[f"block_{self.attn_residual_block_idx[i]}_cross_attn_q"] = hidden_states
+                elif self.to_keys:
+                    additional_residuals[f"block_{self.attn_residual_block_idx[i]}_cross_attn_k"] = hidden_states
+                elif self.to_values:
+                    additional_residuals[f"block_{self.attn_residual_block_idx[i]}_cross_attn_v"] = hidden_states
+
+            if i != len(self.blocks) - 1 and self.inner_dims[i] != self.inner_dims[i + 1]:
+                features.pop(0)
+
+        return additional_residuals

models/decoder.py

@@ -0,0 +1,193 @@
+"""
+@author: Yanzuo Lu
+@email:  luyz5@mail2.sysu.edu.cn
+"""
+
+from typing import Any, Dict, Optional
+
+import torch
+import torch.nn as nn
+from diffusers.models.attention import BasicTransformerBlock
+
+from .xf import FrozenCLIPImageEmbedder
+
+
+class CrossAttnFirstTransformerBlock(BasicTransformerBlock):
+    def forward(
+        self,
+        hidden_states: torch.FloatTensor,
+        query_pos: torch.FloatTensor,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        encoder_hidden_states: Optional[torch.FloatTensor] = None,
+        encoder_attention_mask: Optional[torch.FloatTensor] = None,
+        timestep: Optional[torch.LongTensor] = None,
+        cross_attention_kwargs: Dict[str, Any] = None,
+        class_labels: Optional[torch.LongTensor] = None,
+    ):
+        # Notice that normalization is always applied before the real computation in the following blocks.
+        cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {}
+
+        # 1. Cross-Attention
+        if self.attn2 is not None:
+            hidden_states = hidden_states + query_pos
+            norm_hidden_states = (
+                self.norm2(hidden_states, timestep) if self.use_ada_layer_norm else self.norm2(hidden_states)
+            )
+
+            attn_output = self.attn2(
+                norm_hidden_states,
+                encoder_hidden_states=encoder_hidden_states,
+                attention_mask=encoder_attention_mask,
+                **cross_attention_kwargs,
+            )
+            hidden_states = attn_output + hidden_states
+
+        # 2. Self-Attention
+        hidden_states = hidden_states + query_pos
+        if self.use_ada_layer_norm:
+            norm_hidden_states = self.norm1(hidden_states, timestep)
+        elif self.use_ada_layer_norm_zero:
+            norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(
+                hidden_states, timestep, class_labels, hidden_dtype=hidden_states.dtype
+            )
+        else:
+            norm_hidden_states = self.norm1(hidden_states)
+
+        attn_output = self.attn1(
+            norm_hidden_states,
+            encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None,
+            attention_mask=attention_mask,
+            **cross_attention_kwargs,
+        )
+        if self.use_ada_layer_norm_zero:
+            attn_output = gate_msa.unsqueeze(1) * attn_output
+        hidden_states = attn_output + hidden_states
+
+        # 3. Feed-forward
+        norm_hidden_states = self.norm3(hidden_states)
+
+        if self.use_ada_layer_norm_zero:
+            norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
+
+        if self._chunk_size is not None:
+            # "feed_forward_chunk_size" can be used to save memory
+            if norm_hidden_states.shape[self._chunk_dim] % self._chunk_size != 0:
+                raise ValueError(
+                    f"`hidden_states` dimension to be chunked: {norm_hidden_states.shape[self._chunk_dim]} has to be divisible by chunk size: {self._chunk_size}. Make sure to set an appropriate `chunk_size` when calling `unet.enable_forward_chunking`."
+                )
+
+            num_chunks = norm_hidden_states.shape[self._chunk_dim] // self._chunk_size
+            ff_output = torch.cat(
+                [self.ff(hid_slice) for hid_slice in norm_hidden_states.chunk(num_chunks, dim=self._chunk_dim)],
+                dim=self._chunk_dim,
+            )
+        else:
+            ff_output = self.ff(norm_hidden_states)
+
+        if self.use_ada_layer_norm_zero:
+            ff_output = gate_mlp.unsqueeze(1) * ff_output
+
+        hidden_states = ff_output + hidden_states
+
+        return hidden_states
+
+
+class Decoder(nn.Module):
+    def __init__(self, n_ctx, ctx_dim, heads, depth, last_norm, img_size,
+                 embed_dim, depths, pose_query, pose_channel):
+        super().__init__()
+        self.last_norm = last_norm
+        self.pose_query = pose_query
+        self.pose_channel = pose_channel
+        self.ctx_dim = ctx_dim
+        self.depth = depth
+
+        if self.depth > 0:
+            n_layers = len(depths)
+            embed_dim = embed_dim * 2 ** (n_layers - 1)
+
+            if not self.pose_query:
+                self.query_feat = nn.Parameter(torch.zeros(n_ctx, ctx_dim))
+                nn.init.normal_(self.query_feat, std=0.02)
+            else:
+                self.decoder_fc = nn.Linear(pose_channel, ctx_dim, bias=False)
+
+            self.pos_embed = nn.Parameter(torch.zeros(n_ctx, ctx_dim))
+            nn.init.normal_(self.pos_embed, std=0.02)
+
+            self.blocks = []
+            for _ in range(depth):
+                self.blocks.append(CrossAttnFirstTransformerBlock(
+                    dim=ctx_dim,
+                    num_attention_heads=heads,
+                    attention_head_dim=ctx_dim//heads,
+                    cross_attention_dim=embed_dim
+                ))
+            self.blocks = nn.ModuleList(self.blocks)
+
+            if not self.last_norm:
+                H, W = img_size[0] // 32, img_size[1] // 32
+                self.kv_pos_embed = nn.Parameter(torch.zeros(1, H*W, embed_dim))
+                nn.init.normal_(self.kv_pos_embed, std=0.02)
+
+            # enable xformers
+            def fn_recursive_set_mem_eff(module: torch.nn.Module):
+                if hasattr(module, "set_use_memory_efficient_attention_xformers"):
+                    module.set_use_memory_efficient_attention_xformers(True, attention_op=None)
+
+                for child in module.children():
+                    fn_recursive_set_mem_eff(child)
+
+            for module in self.children():
+                if isinstance(module, torch.nn.Module):
+                    fn_recursive_set_mem_eff(module)
+        elif self.depth == 0:
+            self.clip_model = FrozenCLIPImageEmbedder()
+        elif self.depth == -2:
+            n_layers = len(depths)
+            embed_dim = embed_dim * 2 ** (n_layers - 1)
+            self.decoder_fc = nn.Linear(embed_dim, ctx_dim, bias=False)
+
+    def forward(self, x, features, pose_features):
+        if self.depth > 0:
+            if self.last_norm:
+                B, C = x.shape
+                encoder_hidden_states = x.unsqueeze(1)
+            else:
+                B, L, C = features[-1].shape
+                encoder_hidden_states = features.pop()
+                kv_pos_embed = self.kv_pos_embed.expand(B, -1, -1)
+                encoder_hidden_states = encoder_hidden_states + kv_pos_embed
+
+            if self.pose_query:
+                hidden_states = pose_features.pop()
+                if self.training:
+                    hidden_states = hidden_states.reshape(B*2, self.pose_channel, -1).permute(0, 2, 1)
+                    pos_embed = self.pos_embed.expand(B*2, -1, -1)
+                    encoder_hidden_states = torch.cat([encoder_hidden_states, encoder_hidden_states])
+                else:
+                    hidden_states = hidden_states.reshape(B, self.pose_channel, -1).permute(0, 2, 1)
+                    pos_embed = self.pos_embed.expand(B, -1, -1)
+
+                hidden_states = self.decoder_fc(hidden_states)
+            else:
+                hidden_states = self.query_feat.expand(B, -1, -1)
+                pos_embed = self.pos_embed.expand(B, -1, -1)
+
+            for blk in self.blocks:
+                hidden_states = blk(hidden_states, pos_embed, encoder_hidden_states=encoder_hidden_states)
+            return hidden_states
+        elif self.depth == 0:
+            x = x * 0.5 + 0.5
+            x = x - torch.tensor([0.48145466, 0.4578275, 0.40821073]).view(1, 3, 1, 1).to(dtype=x.dtype, device=x.device)
+            x = x / torch.tensor([0.26862954, 0.26130258, 0.27577711]).view(1, 3, 1, 1).to(dtype=x.dtype, device=x.device)
+            return self.clip_model(x)
+        elif self.depth == -1:
+            encoder_hidden_states = features.pop()
+            encoder_hidden_states = encoder_hidden_states * 0.
+            encoder_hidden_states = encoder_hidden_states.mean(dim=2, keepdim=True).expand(-1, -1, self.ctx_dim)
+            return encoder_hidden_states
+        elif self.depth == -2:
+            encoder_hidden_states = features.pop()
+            encoder_hidden_states = self.decoder_fc(encoder_hidden_states)
+            return encoder_hidden_states

models/metrics.py

@@ -0,0 +1,95 @@
+"""
+@author: Yanzuo Lu
+@author: oliveryanzuolu@gmail.com
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from lpips import LPIPS
+from skimage.metrics import peak_signal_noise_ratio as compare_psnr
+from skimage.metrics import structural_similarity as compare_ssim
+from torchvision import models
+
+
+class build_metric(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+        # FID
+        inception = models.inception_v3(weights=models.Inception_V3_Weights.IMAGENET1K_V1)
+        self.inception_blocks = nn.Sequential(
+            inception.Conv2d_1a_3x3,
+            inception.Conv2d_2a_3x3,
+            inception.Conv2d_2b_3x3,
+            nn.MaxPool2d(kernel_size=3, stride=2),
+            inception.Conv2d_3b_1x1,
+            inception.Conv2d_4a_3x3,
+            nn.MaxPool2d(kernel_size=3, stride=2),
+            inception.Mixed_5b,
+            inception.Mixed_5c,
+            inception.Mixed_5d,
+            inception.Mixed_6a,
+            inception.Mixed_6b,
+            inception.Mixed_6c,
+            inception.Mixed_6d,
+            inception.Mixed_6e,
+            inception.Mixed_7a,
+            inception.Mixed_7b,
+            inception.Mixed_7c,
+            nn.AdaptiveAvgPool2d(output_size=(1, 1))
+        )
+
+        # LPIPS
+        self.lpips_model = LPIPS(net="alex", verbose=False)
+
+        # freeze
+        self.eval()
+        self.requires_grad_(False)
+
+    def forward(self, gt, pred=None):
+        if pred is None:
+            return self.forward_inception(gt).reshape(gt.shape[0], -1) # fid real
+
+        # inputs should be [0,1] here
+        assert gt.shape[0] == pred.shape[0]
+        bsz = gt.shape[0]
+
+        # FID
+        out = self.forward_inception(pred).reshape(bsz, -1)
+
+        # LPIPS
+        lpips = self.lpips_model(pred, gt, normalize=True).reshape(bsz, -1)
+
+        # PSNR & SSIM
+        img_gts = gt.cpu().numpy()
+        img_preds = pred.cpu().numpy()
+        psnr = []
+        ssim = []
+        ssim_256 = []
+
+        for i in range(bsz):
+            img_gt = img_gts[i]
+            img_pred = img_preds[i]
+
+            psnr.append(compare_psnr(img_gt, img_pred, data_range=1))
+            ssim.append(compare_ssim(img_gt, img_pred, data_range=1, win_size=51, channel_axis=0))
+
+            img_gt_256 = img_gt * 255.0
+            img_pred_256 = img_pred * 255.0
+            ssim_256.append(compare_ssim(img_gt_256, img_pred_256, gaussian_weights=True, sigma=1.5,
+                                         use_sample_covariance=False, channel_axis=0,
+                                         data_range=img_pred_256.max() - img_pred_256.min()))
+
+        psnr = torch.tensor(psnr).to(gt.device).reshape(bsz, -1)
+        ssim = torch.tensor(ssim).to(gt.device).reshape(bsz, -1)
+        ssim_256 = torch.tensor(ssim_256).to(gt.device).reshape(bsz, -1)
+        return out, lpips, psnr, ssim, ssim_256
+
+    def forward_inception(self, x):
+        x = F.interpolate(x, size=(299, 299), mode='bilinear')
+        x[:, 0] = x[:, 0] * (0.229 / 0.5) + (0.485 - 0.5) / 0.5
+        x[:, 1] = x[:, 1] * (0.224 / 0.5) + (0.456 - 0.5) / 0.5
+        x[:, 2] = x[:, 2] * (0.225 / 0.5) + (0.406 - 0.5) / 0.5
+        out = self.inception_blocks(x)
+        return out

models/pose_encoder.py

@@ -0,0 +1,46 @@
+"""
+@author: Yanzuo Lu
+@author: oliveryanzuolu@gmail.com
+"""
+
+import torch.nn as nn
+from diffusers.models.resnet import ResnetBlock2D, Downsample2D
+
+
+class PoseEncoder(nn.Module):
+    def __init__(self, downscale_factor, pose_channels, in_channels, channels):
+        super().__init__()
+        self.unshuffle = nn.PixelUnshuffle(downscale_factor)
+        self.conv_in = nn.Conv2d(int(pose_channels * (downscale_factor ** 2)), in_channels, kernel_size=1)
+
+        resnets = []
+        downsamplers = []
+        for i in range(len(channels)):
+            in_channels = in_channels if i == 0 else channels[i - 1]
+            out_channels = channels[i]
+
+            resnets.append(ResnetBlock2D(
+                in_channels=in_channels,
+                out_channels=out_channels,
+                temb_channels=None, # no time embed
+            ))
+            downsamplers.append(Downsample2D(
+                out_channels,
+                use_conv=False,
+                out_channels=out_channels,
+                padding=1,
+                name="op"
+            ) if i != len(channels) - 1 else nn.Identity())
+
+        self.resnets = nn.ModuleList(resnets)
+        self.downsamplers = nn.ModuleList(downsamplers)
+
+    def forward(self, hidden_states):
+        features = []
+        hidden_states = self.unshuffle(hidden_states)
+        hidden_states = self.conv_in(hidden_states)
+        for resnet, downsampler in zip(self.resnets, self.downsamplers):
+            hidden_states = resnet(hidden_states, temb=None)
+            features.append(hidden_states)
+            hidden_states = downsampler(hidden_states)
+        return features

models/swin_transformer.py

@@ -0,0 +1,724 @@
+"""
+@author: Yanzuo Lu
+@author: oliveryanzuolu@gmail.com
+"""
+
+import logging
+import numpy as np
+
+import torch
+import torch.nn as nn
+import torch.utils.checkpoint as checkpoint
+from timm.models.layers import DropPath, to_2tuple, trunc_normal_
+
+# we find the kernel to cause nan, simply omit it
+WindowProcess = None
+WindowProcessReverse = None
+
+logger = logging.getLogger()
+
+
+class Mlp(nn.Module):
+    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+def window_partition(x, window_size):
+    """
+    Args:
+        x: (B, H, W, C)
+        window_size (int): window size
+
+    Returns:
+        windows: (num_windows*B, window_size, window_size, C)
+    """
+    B, H, W, C = x.shape
+    x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)
+    windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)
+    return windows
+
+
+def window_reverse(windows, window_size, H, W):
+    """
+    Args:
+        windows: (num_windows*B, window_size, window_size, C)
+        window_size (int): Window size
+        H (int): Height of image
+        W (int): Width of image
+
+    Returns:
+        x: (B, H, W, C)
+    """
+    B = int(windows.shape[0] / (H * W / window_size / window_size))
+    x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1)
+    x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)
+    return x
+
+
+class WindowAttention(nn.Module):
+    r""" Window based multi-head self attention (W-MSA) module with relative position bias.
+    It supports both of shifted and non-shifted window.
+
+    Args:
+        dim (int): Number of input channels.
+        window_size (tuple[int]): The height and width of the window.
+        num_heads (int): Number of attention heads.
+        qkv_bias (bool, optional):  If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set
+        attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0
+        proj_drop (float, optional): Dropout ratio of output. Default: 0.0
+    """
+
+    def __init__(self, dim, window_size, num_heads, qkv_bias=True, qk_scale=None, attn_drop=0., proj_drop=0.):
+
+        super().__init__()
+        self.dim = dim
+        self.window_size = window_size  # Wh, Ww
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = qk_scale or head_dim ** -0.5
+
+        # define a parameter table of relative position bias
+        self.relative_position_bias_table = nn.Parameter(
+            torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads))  # 2*Wh-1 * 2*Ww-1, nH
+
+        # get pair-wise relative position index for each token inside the window
+        coords_h = torch.arange(self.window_size[0])
+        coords_w = torch.arange(self.window_size[1])
+        coords = torch.stack(torch.meshgrid([coords_h, coords_w]))  # 2, Wh, Ww
+        coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww
+        relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]  # 2, Wh*Ww, Wh*Ww
+        relative_coords = relative_coords.permute(1, 2, 0).contiguous()  # Wh*Ww, Wh*Ww, 2
+        relative_coords[:, :, 0] += self.window_size[0] - 1  # shift to start from 0
+        relative_coords[:, :, 1] += self.window_size[1] - 1
+        relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1
+        relative_position_index = relative_coords.sum(-1)  # Wh*Ww, Wh*Ww
+        self.register_buffer("relative_position_index", relative_position_index)
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+        trunc_normal_(self.relative_position_bias_table, std=.02)
+        self.softmax = nn.Softmax(dim=-1)
+
+    def forward(self, x, mask=None):
+        """
+        Args:
+            x: input features with shape of (num_windows*B, N, C)
+            mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None
+        """
+        B_, N, C = x.shape
+        qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]  # make torchscript happy (cannot use tensor as tuple)
+
+        q = q * self.scale
+        attn = (q @ k.transpose(-2, -1))
+
+        relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view(
+            self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1)  # Wh*Ww,Wh*Ww,nH
+        relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()  # nH, Wh*Ww, Wh*Ww
+        attn = attn + relative_position_bias.unsqueeze(0)
+
+        if mask is not None:
+            nW = mask.shape[0]
+            attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)
+            attn = attn.view(-1, self.num_heads, N, N)
+            attn = self.softmax(attn)
+        else:
+            attn = self.softmax(attn)
+
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B_, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+    def extra_repr(self) -> str:
+        return f'dim={self.dim}, window_size={self.window_size}, num_heads={self.num_heads}'
+
+    def flops(self, N):
+        # calculate flops for 1 window with token length of N
+        flops = 0
+        # qkv = self.qkv(x)
+        flops += N * self.dim * 3 * self.dim
+        # attn = (q @ k.transpose(-2, -1))
+        flops += self.num_heads * N * (self.dim // self.num_heads) * N
+        #  x = (attn @ v)
+        flops += self.num_heads * N * N * (self.dim // self.num_heads)
+        # x = self.proj(x)
+        flops += N * self.dim * self.dim
+        return flops
+
+
+class SwinTransformerBlock(nn.Module):
+    r""" Swin Transformer Block.
+
+    Args:
+        dim (int): Number of input channels.
+        input_resolution (tuple[int]): Input resulotion.
+        num_heads (int): Number of attention heads.
+        window_size (int): Window size.
+        shift_size (int): Shift size for SW-MSA.
+        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
+        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
+        drop (float, optional): Dropout rate. Default: 0.0
+        attn_drop (float, optional): Attention dropout rate. Default: 0.0
+        drop_path (float, optional): Stochastic depth rate. Default: 0.0
+        act_layer (nn.Module, optional): Activation layer. Default: nn.GELU
+        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm
+        fused_window_process (bool, optional): If True, use one kernel to fused window shift & window partition for acceleration, similar for the reversed part. Default: False
+    """
+
+    def __init__(self, dim, input_resolution, num_heads, window_size=7, shift_size=0,
+                 mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0., drop_path=0.,
+                 act_layer=nn.GELU, norm_layer=nn.LayerNorm,
+                 fused_window_process=False):
+        super().__init__()
+        self.dim = dim
+        self.input_resolution = input_resolution
+        self.num_heads = num_heads
+        self.window_size = window_size
+        self.shift_size = shift_size
+        self.mlp_ratio = mlp_ratio
+        if min(self.input_resolution) <= self.window_size:
+            # if window size is larger than input resolution, we don't partition windows
+            self.shift_size = 0
+            self.window_size = min(self.input_resolution)
+        assert 0 <= self.shift_size < self.window_size, "shift_size must in 0-window_size"
+
+        self.norm1 = norm_layer(dim)
+        self.attn = WindowAttention(
+            dim, window_size=to_2tuple(self.window_size), num_heads=num_heads,
+            qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
+
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
+
+        if self.shift_size > 0:
+            # calculate attention mask for SW-MSA
+            H, W = self.input_resolution
+            img_mask = torch.zeros((1, H, W, 1))  # 1 H W 1
+            h_slices = (slice(0, -self.window_size),
+                        slice(-self.window_size, -self.shift_size),
+                        slice(-self.shift_size, None))
+            w_slices = (slice(0, -self.window_size),
+                        slice(-self.window_size, -self.shift_size),
+                        slice(-self.shift_size, None))
+            cnt = 0
+            for h in h_slices:
+                for w in w_slices:
+                    img_mask[:, h, w, :] = cnt
+                    cnt += 1
+
+            mask_windows = window_partition(img_mask, self.window_size)  # nW, window_size, window_size, 1
+            mask_windows = mask_windows.view(-1, self.window_size * self.window_size)
+            attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
+            attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))
+        else:
+            attn_mask = None
+
+        self.register_buffer("attn_mask", attn_mask)
+        self.fused_window_process = fused_window_process
+
+    def forward(self, x):
+        H, W = self.input_resolution
+        B, L, C = x.shape
+        assert L == H * W, "input feature has wrong size"
+
+        shortcut = x
+        x = self.norm1(x)
+        x = x.view(B, H, W, C)
+
+        # cyclic shift
+        if self.shift_size > 0:
+            if not self.fused_window_process:
+                shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))
+                # partition windows
+                x_windows = window_partition(shifted_x, self.window_size)  # nW*B, window_size, window_size, C
+            else:
+                x_windows = WindowProcess.apply(x, B, H, W, C, -self.shift_size, self.window_size)
+        else:
+            shifted_x = x
+            # partition windows
+            x_windows = window_partition(shifted_x, self.window_size)  # nW*B, window_size, window_size, C
+
+        x_windows = x_windows.view(-1, self.window_size * self.window_size, C)  # nW*B, window_size*window_size, C
+
+        # W-MSA/SW-MSA
+        attn_windows = self.attn(x_windows, mask=self.attn_mask)  # nW*B, window_size*window_size, C
+
+        # merge windows
+        attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C)
+
+        # reverse cyclic shift
+        if self.shift_size > 0:
+            if not self.fused_window_process:
+                shifted_x = window_reverse(attn_windows, self.window_size, H, W)  # B H' W' C
+                x = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2))
+            else:
+                x = WindowProcessReverse.apply(attn_windows, B, H, W, C, self.shift_size, self.window_size)
+        else:
+            shifted_x = window_reverse(attn_windows, self.window_size, H, W)  # B H' W' C
+            x = shifted_x
+        x = x.view(B, H * W, C)
+        x = shortcut + self.drop_path(x)
+
+        # FFN
+        x = x + self.drop_path(self.mlp(self.norm2(x)))
+
+        return x
+
+    def extra_repr(self) -> str:
+        return f"dim={self.dim}, input_resolution={self.input_resolution}, num_heads={self.num_heads}, " \
+               f"window_size={self.window_size}, shift_size={self.shift_size}, mlp_ratio={self.mlp_ratio}"
+
+    def flops(self):
+        flops = 0
+        H, W = self.input_resolution
+        # norm1
+        flops += self.dim * H * W
+        # W-MSA/SW-MSA
+        nW = H * W / self.window_size / self.window_size
+        flops += nW * self.attn.flops(self.window_size * self.window_size)
+        # mlp
+        flops += 2 * H * W * self.dim * self.dim * self.mlp_ratio
+        # norm2
+        flops += self.dim * H * W
+        return flops
+
+
+class PatchMerging(nn.Module):
+    r""" Patch Merging Layer.
+
+    Args:
+        input_resolution (tuple[int]): Resolution of input feature.
+        dim (int): Number of input channels.
+        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm
+    """
+
+    def __init__(self, input_resolution, dim, norm_layer=nn.LayerNorm):
+        super().__init__()
+        self.input_resolution = input_resolution
+        self.dim = dim
+        self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False)
+        self.norm = norm_layer(4 * dim)
+
+    def forward(self, x):
+        """
+        x: B, H*W, C
+        """
+        H, W = self.input_resolution
+        B, L, C = x.shape
+        assert L == H * W, "input feature has wrong size"
+        assert H % 2 == 0 and W % 2 == 0, f"x size ({H}*{W}) are not even."
+
+        x = x.view(B, H, W, C)
+
+        x0 = x[:, 0::2, 0::2, :]  # B H/2 W/2 C
+        x1 = x[:, 1::2, 0::2, :]  # B H/2 W/2 C
+        x2 = x[:, 0::2, 1::2, :]  # B H/2 W/2 C
+        x3 = x[:, 1::2, 1::2, :]  # B H/2 W/2 C
+        x = torch.cat([x0, x1, x2, x3], -1)  # B H/2 W/2 4*C
+        x = x.view(B, -1, 4 * C)  # B H/2*W/2 4*C
+
+        x = self.norm(x)
+        x = self.reduction(x)
+
+        return x
+
+    def extra_repr(self) -> str:
+        return f"input_resolution={self.input_resolution}, dim={self.dim}"
+
+    def flops(self):
+        H, W = self.input_resolution
+        flops = H * W * self.dim
+        flops += (H // 2) * (W // 2) * 4 * self.dim * 2 * self.dim
+        return flops
+
+
+class BasicLayer(nn.Module):
+    """ A basic Swin Transformer layer for one stage.
+
+    Args:
+        dim (int): Number of input channels.
+        input_resolution (tuple[int]): Input resolution.
+        depth (int): Number of blocks.
+        num_heads (int): Number of attention heads.
+        window_size (int): Local window size.
+        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
+        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
+        drop (float, optional): Dropout rate. Default: 0.0
+        attn_drop (float, optional): Attention dropout rate. Default: 0.0
+        drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0
+        norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm
+        downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None
+        use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False.
+        fused_window_process (bool, optional): If True, use one kernel to fused window shift & window partition for acceleration, similar for the reversed part. Default: False
+    """
+
+    def __init__(self, dim, input_resolution, depth, num_heads, window_size,
+                 mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0.,
+                 drop_path=0., norm_layer=nn.LayerNorm, downsample=None, use_checkpoint=False,
+                 fused_window_process=False):
+
+        super().__init__()
+        self.dim = dim
+        self.input_resolution = input_resolution
+        self.depth = depth
+        self.use_checkpoint = use_checkpoint
+
+        # build blocks
+        self.blocks = nn.ModuleList([
+            SwinTransformerBlock(dim=dim, input_resolution=input_resolution,
+                                 num_heads=num_heads, window_size=window_size,
+                                 shift_size=0 if (i % 2 == 0) else window_size // 2,
+                                 mlp_ratio=mlp_ratio,
+                                 qkv_bias=qkv_bias, qk_scale=qk_scale,
+                                 drop=drop, attn_drop=attn_drop,
+                                 drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path,
+                                 norm_layer=norm_layer,
+                                 fused_window_process=fused_window_process)
+            for i in range(depth)])
+
+        # patch merging layer
+        if downsample is not None:
+            self.downsample = downsample(input_resolution, dim=dim, norm_layer=norm_layer)
+        else:
+            self.downsample = None
+
+    def forward(self, x):
+        for blk in self.blocks:
+            if self.use_checkpoint:
+                x = checkpoint.checkpoint(blk, x)
+            else:
+                x = blk(x)
+        feature = x
+        if self.downsample is not None:
+            x = self.downsample(x)
+        return x, feature
+
+    def extra_repr(self) -> str:
+        return f"dim={self.dim}, input_resolution={self.input_resolution}, depth={self.depth}"
+
+    def flops(self):
+        flops = 0
+        for blk in self.blocks:
+            flops += blk.flops()
+        if self.downsample is not None:
+            flops += self.downsample.flops()
+        return flops
+
+
+class PatchEmbed(nn.Module):
+    r""" Image to Patch Embedding
+
+    Args:
+        img_size (int): Image size.  Default: 224.
+        patch_size (int): Patch token size. Default: 4.
+        in_chans (int): Number of input image channels. Default: 3.
+        embed_dim (int): Number of linear projection output channels. Default: 96.
+        norm_layer (nn.Module, optional): Normalization layer. Default: None
+    """
+
+    def __init__(self, img_size=224, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None):
+        super().__init__()
+        img_size = to_2tuple(img_size)
+        patch_size = to_2tuple(patch_size)
+        patches_resolution = [img_size[0] // patch_size[0], img_size[1] // patch_size[1]]
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.patches_resolution = patches_resolution
+        self.num_patches = patches_resolution[0] * patches_resolution[1]
+
+        self.in_chans = in_chans
+        self.embed_dim = embed_dim
+
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
+        if norm_layer is not None:
+            self.norm = norm_layer(embed_dim)
+        else:
+            self.norm = None
+
+    def forward(self, x):
+        # B, C, H, W = x.shape
+        # FIXME look at relaxing size constraints
+        # assert H == self.img_size[0] and W == self.img_size[1], \
+        #     f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
+        # x = self.proj(x).flatten(2).transpose(1, 2)  # B Ph*Pw C
+        x = self.proj(x)
+        B, C, H, W = x.shape
+        x = x.flatten(2).transpose(1, 2)
+        if self.norm is not None:
+            x = self.norm(x)
+        return x.transpose(1, 2).reshape(B, C, H, W)
+
+    def flops(self):
+        Ho, Wo = self.patches_resolution
+        flops = Ho * Wo * self.embed_dim * self.in_chans * (self.patch_size[0] * self.patch_size[1])
+        if self.norm is not None:
+            flops += Ho * Wo * self.embed_dim
+        return flops
+
+
+class SwinTransformer(nn.Module):
+    r""" Swin Transformer
+        A PyTorch impl of : `Swin Transformer: Hierarchical Vision Transformer using Shifted Windows`  -
+          https://arxiv.org/pdf/2103.14030
+
+    Args:
+        img_size (int | tuple(int)): Input image size. Default 224
+        patch_size (int | tuple(int)): Patch size. Default: 4
+        in_chans (int): Number of input image channels. Default: 3
+        num_classes (int): Number of classes for classification head. Default: 1000
+        embed_dim (int): Patch embedding dimension. Default: 96
+        depths (tuple(int)): Depth of each Swin Transformer layer.
+        num_heads (tuple(int)): Number of attention heads in different layers.
+        window_size (int): Window size. Default: 7
+        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4
+        qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float): Override default qk scale of head_dim ** -0.5 if set. Default: None
+        drop_rate (float): Dropout rate. Default: 0
+        attn_drop_rate (float): Attention dropout rate. Default: 0
+        drop_path_rate (float): Stochastic depth rate. Default: 0.1
+        norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm.
+        ape (bool): If True, add absolute position embedding to the patch embedding. Default: False
+        patch_norm (bool): If True, add normalization after patch embedding. Default: True
+        use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False
+        fused_window_process (bool, optional): If True, use one kernel to fused window shift & window partition for acceleration, similar for the reversed part. Default: False
+    """
+
+    def __init__(self, img_size=224, patch_size=4, in_chans=3, num_classes=0,
+                 embed_dim=96, depths=[2, 2, 6, 2], num_heads=[3, 6, 12, 24],
+                 window_size=7, mlp_ratio=4., qkv_bias=True, qk_scale=None,
+                 drop_rate=0., attn_drop_rate=0., drop_path_rate=0.1,
+                 norm_layer=nn.LayerNorm, ape=False, patch_norm=True,
+                 use_checkpoint=False, fused_window_process=False,
+                 mask=True, last_norm=True, **kwargs):
+        super().__init__()
+
+        self.num_classes = num_classes
+        self.num_layers = len(depths)
+        self.embed_dim = embed_dim
+        self.ape = ape
+        self.patch_norm = patch_norm
+        self.num_features = int(embed_dim * 2 ** (self.num_layers - 1))
+        self.mlp_ratio = mlp_ratio
+        self.mask = mask
+
+        # split image into non-overlapping patches
+        self.patch_embed = PatchEmbed(
+            img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim,
+            norm_layer=norm_layer if self.patch_norm else None)
+        num_patches = self.patch_embed.num_patches
+        patches_resolution = self.patch_embed.patches_resolution
+        self.patches_resolution = patches_resolution
+
+        # absolute position embedding
+        if self.ape:
+            self.absolute_pos_embed = nn.Parameter(torch.zeros(1, num_patches, embed_dim))
+            trunc_normal_(self.absolute_pos_embed, std=.02)
+
+        self.pos_drop = nn.Dropout(p=drop_rate)
+
+        # stochastic depth
+        self.dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]  # stochastic depth decay rule
+
+        # build layers
+        self.layers = nn.ModuleList()
+        for i_layer in range(self.num_layers):
+            layer = BasicLayer(dim=int(embed_dim * 2 ** i_layer),
+                               input_resolution=(patches_resolution[0] // (2 ** i_layer),
+                                                 patches_resolution[1] // (2 ** i_layer)),
+                               depth=depths[i_layer],
+                               num_heads=num_heads[i_layer],
+                               window_size=window_size,
+                               mlp_ratio=self.mlp_ratio,
+                               qkv_bias=qkv_bias, qk_scale=qk_scale,
+                               drop=drop_rate, attn_drop=attn_drop_rate,
+                               drop_path=self.dpr[sum(depths[:i_layer]):sum(depths[:i_layer + 1])],
+                               norm_layer=norm_layer,
+                               downsample=PatchMerging if (i_layer < self.num_layers - 1) else None,
+                               use_checkpoint=use_checkpoint,
+                               fused_window_process=fused_window_process)
+            self.layers.append(layer)
+
+        self.norm = norm_layer(self.num_features) if last_norm else nn.Identity()
+        self.avgpool = nn.AdaptiveAvgPool1d(1)
+        self.head = nn.Linear(self.num_features, num_classes) if num_classes > 0 else nn.Identity()
+
+        self.apply(self._init_weights)
+        if self.mask:
+            self.masked_embed = nn.Parameter(torch.zeros(1, embed_dim))
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        return {'absolute_pos_embed'}
+
+    @torch.jit.ignore
+    def no_weight_decay_keywords(self):
+        return {'relative_position_bias_table'}
+
+    def forward_features(self, x, mask=None):
+        x = self.patch_embed(x)
+        if self.mask and mask is not None:
+            x = self.mask_model(x, mask)
+        x = x.flatten(2).transpose(1, 2)
+
+        if self.ape:
+            x = x + self.absolute_pos_embed
+        x = self.pos_drop(x)
+
+        features = []
+        for layer in self.layers:
+            x, feature = layer(x)
+            features.append(feature)
+
+        x = self.norm(x)  # B L C
+        x = self.avgpool(x.transpose(1, 2))  # B C 1
+        x = torch.flatten(x, 1)
+        return x, features
+
+    def mask_model(self, x, mask):
+        if x.shape[-2:] != mask.shape[-2:]:
+            htimes, wtimes = np.array(x.shape[-2:]) // np.array(mask.shape[-2:])
+            mask = mask.repeat_interleave(htimes, -2).repeat_interleave(wtimes, -1)
+
+        # mask embed
+        x.permute(0, 2, 3, 1)[mask, :] = self.masked_embed.to(x.dtype)
+
+        return x
+
+    def forward(self, x, mask=None):
+        x, features = self.forward_features(x, mask)
+        x = self.head(x)
+        return x, features
+
+    def flops(self):
+        flops = 0
+        flops += self.patch_embed.flops()
+        for i, layer in enumerate(self.layers):
+            flops += layer.flops()
+        flops += self.num_features * self.patches_resolution[0] * self.patches_resolution[1] // (2 ** self.num_layers)
+        flops += self.num_features * self.num_classes
+        return flops
+
+
+def load_pretrained(model, pretrained_path, checkpoint_key="model", checkpoint_prefix=None):
+    # logger.info(f"==============> Loading weight {pretrained_path} for fine-tuning......")
+    state_dict = torch.load(pretrained_path, map_location='cpu')
+    if checkpoint_key:
+        state_dict = state_dict[checkpoint_key]
+    if checkpoint_prefix:
+        state_dict = {k[len(checkpoint_prefix):]: v for k, v in state_dict.items() \
+                      if k.startswith(checkpoint_prefix)}
+
+    # delete relative_position_index since we always re-init it
+    relative_position_index_keys = [k for k in state_dict.keys() if "relative_position_index" in k]
+    for k in relative_position_index_keys:
+        del state_dict[k]
+
+    # delete relative_coords_table since we always re-init it
+    relative_position_index_keys = [k for k in state_dict.keys() if "relative_coords_table" in k]
+    for k in relative_position_index_keys:
+        del state_dict[k]
+
+    # delete attn_mask since we always re-init it
+    attn_mask_keys = [k for k in state_dict.keys() if "attn_mask" in k]
+    for k in attn_mask_keys:
+        del state_dict[k]
+
+    # bicubic interpolate relative_position_bias_table if not match
+    relative_position_bias_table_keys = [k for k in state_dict.keys() if "relative_position_bias_table" in k]
+    for k in relative_position_bias_table_keys:
+        relative_position_bias_table_pretrained = state_dict[k]
+        relative_position_bias_table_current = model.state_dict()[k]
+        L1, nH1 = relative_position_bias_table_pretrained.size()
+        L2, nH2 = relative_position_bias_table_current.size()
+        if nH1 != nH2:
+            logger.info(f"Error in loading {k}, passing......")
+        else:
+            if L1 != L2:
+                # bicubic interpolate relative_position_bias_table if not match
+                S1 = int(L1 ** 0.5)
+                S2 = int(L2 ** 0.5)
+                relative_position_bias_table_pretrained_resized = torch.nn.functional.interpolate(
+                    relative_position_bias_table_pretrained.permute(1, 0).view(1, nH1, S1, S1), size=(S2, S2),
+                    mode='bicubic')
+                state_dict[k] = relative_position_bias_table_pretrained_resized.view(nH2, L2).permute(1, 0)
+
+    # bicubic interpolate absolute_pos_embed if not match
+    absolute_pos_embed_keys = [k for k in state_dict.keys() if "absolute_pos_embed" in k]
+    for k in absolute_pos_embed_keys:
+        # dpe
+        absolute_pos_embed_pretrained = state_dict[k]
+        absolute_pos_embed_current = model.state_dict()[k]
+        _, L1, C1 = absolute_pos_embed_pretrained.size()
+        _, L2, C2 = absolute_pos_embed_current.size()
+        if C1 != C1:
+            logger.info(f"Error in loading {k}, passing......")
+        else:
+            if L1 != L2:
+                S1 = int(L1 ** 0.5)
+                S2 = int(L2 ** 0.5)
+                absolute_pos_embed_pretrained = absolute_pos_embed_pretrained.reshape(-1, S1, S1, C1)
+                absolute_pos_embed_pretrained = absolute_pos_embed_pretrained.permute(0, 3, 1, 2)
+                absolute_pos_embed_pretrained_resized = torch.nn.functional.interpolate(
+                    absolute_pos_embed_pretrained, size=(S2, S2), mode='bicubic')
+                absolute_pos_embed_pretrained_resized = absolute_pos_embed_pretrained_resized.permute(0, 2, 3, 1)
+                absolute_pos_embed_pretrained_resized = absolute_pos_embed_pretrained_resized.flatten(1, 2)
+                state_dict[k] = absolute_pos_embed_pretrained_resized
+
+    msg = model.load_state_dict(state_dict, strict=False)
+    logger.info(msg)
+
+    # logger.info(f"=> loaded successfully '{pretrained_path}'")
+
+    # del checkpoint
+    torch.cuda.empty_cache()
+
+
+def build_backbone(img_size, embed_dim, depths, num_heads, window_size, drop_path_rate, mask, last_norm, pretrained_path):
+    if len(depths) > 0:
+        model = SwinTransformer(img_size=img_size, embed_dim=embed_dim, depths=depths, num_heads=num_heads,
+                                window_size=window_size, drop_path_rate=drop_path_rate, mask=mask, last_norm=last_norm)
+        if pretrained_path:
+            load_pretrained(model, pretrained_path)
+    else:
+        class Identity(nn.Module):
+            def __init__(self):
+                super().__init__()
+
+            def forward(self, x, mask=None):
+                return x, []
+        model = Identity()
+    return model

models/unet.py

@@ -0,0 +1,1946 @@
+"""
+@author: Yanzuo Lu
+@author: oliveryanzuolu@gmail.com
+"""
+
+import torch
+import torch.nn as nn
+from diffusers.models.attention import *
+from diffusers.models.attention_processor import *
+from diffusers.models.resnet import *
+from diffusers.models.transformer_2d import *
+from diffusers.models.unet_2d_blocks import *
+from diffusers.models.unet_2d_condition import *
+
+
+class ResidualXFormersAttnProcessor(XFormersAttnProcessor):
+    def __call__(
+        self,
+        attn: Attention,
+        hidden_states: torch.FloatTensor,
+        encoder_hidden_states: Optional[torch.FloatTensor] = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        temb: Optional[torch.FloatTensor] = None,
+        block_idx: Optional[int] = None,
+        additional_residuals: Optional[Dict[str, torch.FloatTensor]] = None,
+        is_self_attn: Optional[bool] = None
+    ):
+        residual = hidden_states
+
+        if attn.spatial_norm is not None:
+            hidden_states = attn.spatial_norm(hidden_states, temb)
+
+        input_ndim = hidden_states.ndim
+
+        if input_ndim == 4:
+            batch_size, channel, height, width = hidden_states.shape
+            hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
+
+        batch_size, key_tokens, _ = (
+            hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+        )
+
+        attention_mask = attn.prepare_attention_mask(attention_mask, key_tokens, batch_size)
+        if attention_mask is not None:
+            # expand our mask's singleton query_tokens dimension:
+            #   [batch*heads,            1, key_tokens] ->
+            #   [batch*heads, query_tokens, key_tokens]
+            # so that it can be added as a bias onto the attention scores that xformers computes:
+            #   [batch*heads, query_tokens, key_tokens]
+            # we do this explicitly because xformers doesn't broadcast the singleton dimension for us.
+            _, query_tokens, _ = hidden_states.shape
+            attention_mask = attention_mask.expand(-1, query_tokens, -1)
+
+        if attn.group_norm is not None:
+            hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
+
+        query = attn.to_q(hidden_states)
+
+        # newly added
+        if is_self_attn and additional_residuals and f"block_{block_idx}_self_attn_q" in additional_residuals:
+            query = query + additional_residuals[f"block_{block_idx}_self_attn_q"]
+        elif not is_self_attn and additional_residuals and f"block_{block_idx}_cross_attn_q" in additional_residuals:
+            query = query + additional_residuals[f"block_{block_idx}_cross_attn_q"]
+
+        if encoder_hidden_states is None:
+            encoder_hidden_states = hidden_states
+        elif attn.norm_cross:
+            encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
+
+        if not is_self_attn and additional_residuals and f"block_{block_idx}_cross_attn_c" in additional_residuals:
+            not_uc = torch.abs(encoder_hidden_states - torch.zeros_like(encoder_hidden_states)).mean(dim=[1, 2], keepdim=True) < 1e-4
+            encoder_hidden_states = encoder_hidden_states + additional_residuals[f"block_{block_idx}_cross_attn_c"] * not_uc
+            # encoder_hidden_states[not_uc] = encoder_hidden_states[not_uc] + \
+            #     additional_residuals[f"block_{block_idx}_cross_attn_c"][not_uc]
+            # encoder_hidden_states[~not_uc] = encoder_hidden_states[~not_uc] + \
+            #     additional_residuals[f"block_{block_idx}_cross_attn_c"][~not_uc] * 0.
+
+        key = attn.to_k(encoder_hidden_states)
+        value = attn.to_v(encoder_hidden_states)
+
+        # newly added
+        if is_self_attn and additional_residuals and f"block_{block_idx}_self_attn_k" in additional_residuals:
+            key = key + additional_residuals[f"block_{block_idx}_self_attn_k"]
+        elif not is_self_attn and additional_residuals and f"block_{block_idx}_cross_attn_k" in additional_residuals:
+            key = key + additional_residuals[f"block_{block_idx}_cross_attn_k"]
+
+        if is_self_attn and additional_residuals and f"block_{block_idx}_self_attn_v" in additional_residuals:
+            value = value + additional_residuals[f"block_{block_idx}_self_attn_v"]
+        elif not is_self_attn and additional_residuals and f"block_{block_idx}_cross_attn_v" in additional_residuals:
+            value = value + additional_residuals[f"block_{block_idx}_cross_attn_v"]
+
+        query = attn.head_to_batch_dim(query).contiguous()
+        key = attn.head_to_batch_dim(key).contiguous()
+        value = attn.head_to_batch_dim(value).contiguous()
+
+        hidden_states = xformers.ops.memory_efficient_attention(
+            query, key, value, attn_bias=attention_mask, op=self.attention_op, scale=attn.scale
+        )
+        hidden_states = hidden_states.to(query.dtype)
+        hidden_states = attn.batch_to_head_dim(hidden_states)
+
+        # linear proj
+        hidden_states = attn.to_out[0](hidden_states)
+        # dropout
+        hidden_states = attn.to_out[1](hidden_states)
+
+        if input_ndim == 4:
+            hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
+
+        if attn.residual_connection:
+            hidden_states = hidden_states + residual
+
+        hidden_states = hidden_states / attn.rescale_output_factor
+
+        return hidden_states
+
+
+class ResidualAttention(Attention):
+    def set_use_memory_efficient_attention_xformers(
+        self, use_memory_efficient_attention_xformers: bool, attention_op: Optional[Callable] = None
+    ):
+        is_lora = hasattr(self, "processor") and isinstance(
+            self.processor,
+            LORA_ATTENTION_PROCESSORS,
+        )
+        is_custom_diffusion = hasattr(self, "processor") and isinstance(
+            self.processor, (CustomDiffusionAttnProcessor, CustomDiffusionXFormersAttnProcessor)
+        )
+        is_added_kv_processor = hasattr(self, "processor") and isinstance(
+            self.processor,
+            (
+                AttnAddedKVProcessor,
+                AttnAddedKVProcessor2_0,
+                SlicedAttnAddedKVProcessor,
+                XFormersAttnAddedKVProcessor,
+                LoRAAttnAddedKVProcessor,
+            ),
+        )
+
+        if use_memory_efficient_attention_xformers:
+            if is_added_kv_processor and (is_lora or is_custom_diffusion):
+                raise NotImplementedError(
+                    f"Memory efficient attention is currently not supported for LoRA or custom diffuson for attention processor type {self.processor}"
+                )
+            if not is_xformers_available():
+                raise ModuleNotFoundError(
+                    (
+                        "Refer to https://github.com/facebookresearch/xformers for more information on how to install"
+                        " xformers"
+                    ),
+                    name="xformers",
+                )
+            elif not torch.cuda.is_available():
+                raise ValueError(
+                    "torch.cuda.is_available() should be True but is False. xformers' memory efficient attention is"
+                    " only available for GPU "
+                )
+            else:
+                try:
+                    # Make sure we can run the memory efficient attention
+                    _ = xformers.ops.memory_efficient_attention(
+                        torch.randn((1, 2, 40), device="cuda"),
+                        torch.randn((1, 2, 40), device="cuda"),
+                        torch.randn((1, 2, 40), device="cuda"),
+                    )
+                except Exception as e:
+                    raise e
+
+            if is_lora:
+                # TODO (sayakpaul): should we throw a warning if someone wants to use the xformers
+                # variant when using PT 2.0 now that we have LoRAAttnProcessor2_0?
+                processor = LoRAXFormersAttnProcessor(
+                    hidden_size=self.processor.hidden_size,
+                    cross_attention_dim=self.processor.cross_attention_dim,
+                    rank=self.processor.rank,
+                    attention_op=attention_op,
+                )
+                processor.load_state_dict(self.processor.state_dict())
+                processor.to(self.processor.to_q_lora.up.weight.device)
+            elif is_custom_diffusion:
+                processor = CustomDiffusionXFormersAttnProcessor(
+                    train_kv=self.processor.train_kv,
+                    train_q_out=self.processor.train_q_out,
+                    hidden_size=self.processor.hidden_size,
+                    cross_attention_dim=self.processor.cross_attention_dim,
+                    attention_op=attention_op,
+                )
+                processor.load_state_dict(self.processor.state_dict())
+                if hasattr(self.processor, "to_k_custom_diffusion"):
+                    processor.to(self.processor.to_k_custom_diffusion.weight.device)
+            elif is_added_kv_processor:
+                # TODO(Patrick, Suraj, William) - currently xformers doesn't work for UnCLIP
+                # which uses this type of cross attention ONLY because the attention mask of format
+                # [0, ..., -10.000, ..., 0, ...,] is not supported
+                # throw warning
+                logger.info(
+                    "Memory efficient attention with `xformers` might currently not work correctly if an attention mask is required for the attention operation."
+                )
+                processor = XFormersAttnAddedKVProcessor(attention_op=attention_op)
+            else:
+                processor = ResidualXFormersAttnProcessor(attention_op=attention_op)
+        else:
+            if is_lora:
+                attn_processor_class = (
+                    LoRAAttnProcessor2_0 if hasattr(F, "scaled_dot_product_attention") else LoRAAttnProcessor
+                )
+                processor = attn_processor_class(
+                    hidden_size=self.processor.hidden_size,
+                    cross_attention_dim=self.processor.cross_attention_dim,
+                    rank=self.processor.rank,
+                )
+                processor.load_state_dict(self.processor.state_dict())
+                processor.to(self.processor.to_q_lora.up.weight.device)
+            elif is_custom_diffusion:
+                processor = CustomDiffusionAttnProcessor(
+                    train_kv=self.processor.train_kv,
+                    train_q_out=self.processor.train_q_out,
+                    hidden_size=self.processor.hidden_size,
+                    cross_attention_dim=self.processor.cross_attention_dim,
+                )
+                processor.load_state_dict(self.processor.state_dict())
+                if hasattr(self.processor, "to_k_custom_diffusion"):
+                    processor.to(self.processor.to_k_custom_diffusion.weight.device)
+            else:
+                # set attention processor
+                # We use the AttnProcessor2_0 by default when torch 2.x is used which uses
+                # torch.nn.functional.scaled_dot_product_attention for native Flash/memory_efficient_attention
+                # but only if it has the default `scale` argument. TODO remove scale_qk check when we move to torch 2.1
+                processor = (
+                    AttnProcessor2_0()
+                    if hasattr(F, "scaled_dot_product_attention") and self.scale_qk
+                    else AttnProcessor()
+                )
+
+        self.set_processor(processor)
+
+    def forward(self, hidden_states, encoder_hidden_states=None, attention_mask=None,
+                block_idx: Optional[int] = None, additional_residuals: Optional[Dict[str, torch.FloatTensor]] = None,
+                is_self_attn: Optional[bool] = None, **cross_attention_kwargs):
+        # The `Attention` class can call different attention processors / attention functions
+        # here we simply pass along all tensors to the selected processor class
+        # For standard processors that are defined here, `**cross_attention_kwargs` is empty
+        return self.processor(
+            self,
+            hidden_states,
+            encoder_hidden_states=encoder_hidden_states,
+            attention_mask=attention_mask,
+            block_idx=block_idx,
+            additional_residuals=additional_residuals,
+            is_self_attn=is_self_attn,
+            **cross_attention_kwargs,
+        )
+
+
+class ResidualTransformerBlock(BasicTransformerBlock):
+    def __init__(
+        self,
+        dim: int,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        dropout=0.0,
+        cross_attention_dim: Optional[int] = None,
+        activation_fn: str = "geglu",
+        num_embeds_ada_norm: Optional[int] = None,
+        attention_bias: bool = False,
+        only_cross_attention: bool = False,
+        double_self_attention: bool = False,
+        upcast_attention: bool = False,
+        norm_elementwise_affine: bool = True,
+        norm_type: str = "layer_norm",
+        final_dropout: bool = False,
+    ):
+        super(BasicTransformerBlock, self).__init__()
+        self.only_cross_attention = only_cross_attention
+
+        self.use_ada_layer_norm_zero = (num_embeds_ada_norm is not None) and norm_type == "ada_norm_zero"
+        self.use_ada_layer_norm = (num_embeds_ada_norm is not None) and norm_type == "ada_norm"
+
+        if norm_type in ("ada_norm", "ada_norm_zero") and num_embeds_ada_norm is None:
+            raise ValueError(
+                f"`norm_type` is set to {norm_type}, but `num_embeds_ada_norm` is not defined. Please make sure to"
+                f" define `num_embeds_ada_norm` if setting `norm_type` to {norm_type}."
+            )
+
+        # Define 3 blocks. Each block has its own normalization layer.
+        # 1. Self-Attn
+        if self.use_ada_layer_norm:
+            self.norm1 = AdaLayerNorm(dim, num_embeds_ada_norm)
+        elif self.use_ada_layer_norm_zero:
+            self.norm1 = AdaLayerNormZero(dim, num_embeds_ada_norm)
+        else:
+            self.norm1 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine)
+        self.attn1 = ResidualAttention(
+            query_dim=dim,
+            heads=num_attention_heads,
+            dim_head=attention_head_dim,
+            dropout=dropout,
+            bias=attention_bias,
+            cross_attention_dim=cross_attention_dim if only_cross_attention else None,
+            upcast_attention=upcast_attention,
+        )
+
+        # 2. Cross-Attn
+        if cross_attention_dim is not None or double_self_attention:
+            # We currently only use AdaLayerNormZero for self attention where there will only be one attention block.
+            # I.e. the number of returned modulation chunks from AdaLayerZero would not make sense if returned during
+            # the second cross attention block.
+            self.norm2 = (
+                AdaLayerNorm(dim, num_embeds_ada_norm)
+                if self.use_ada_layer_norm
+                else nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine)
+            )
+            self.attn2 = ResidualAttention(
+                query_dim=dim,
+                cross_attention_dim=cross_attention_dim if not double_self_attention else None,
+                heads=num_attention_heads,
+                dim_head=attention_head_dim,
+                dropout=dropout,
+                bias=attention_bias,
+                upcast_attention=upcast_attention,
+            )  # is self-attn if encoder_hidden_states is none
+        else:
+            self.norm2 = None
+            self.attn2 = None
+
+        # 3. Feed-forward
+        self.norm3 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine)
+        self.ff = FeedForward(dim, dropout=dropout, activation_fn=activation_fn, final_dropout=final_dropout)
+
+        # let chunk size default to None
+        self._chunk_size = None
+        self._chunk_dim = 0
+
+    def forward(
+        self,
+        hidden_states: torch.FloatTensor,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        encoder_hidden_states: Optional[torch.FloatTensor] = None,
+        encoder_attention_mask: Optional[torch.FloatTensor] = None,
+        timestep: Optional[torch.LongTensor] = None,
+        cross_attention_kwargs: Dict[str, Any] = None,
+        class_labels: Optional[torch.LongTensor] = None,
+        block_idx: Optional[int] = None,
+        additional_residuals: Optional[Dict[str, torch.FloatTensor]] = None
+    ):
+        # Notice that normalization is always applied before the real computation in the following blocks.
+        # 1. Self-Attention
+        if self.use_ada_layer_norm:
+            norm_hidden_states = self.norm1(hidden_states, timestep)
+        elif self.use_ada_layer_norm_zero:
+            norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(
+                hidden_states, timestep, class_labels, hidden_dtype=hidden_states.dtype
+            )
+        else:
+            norm_hidden_states = self.norm1(hidden_states)
+
+        cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {}
+
+        attn_output = self.attn1(
+            norm_hidden_states,
+            encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None,
+            attention_mask=attention_mask,
+            block_idx=block_idx,
+            additional_residuals=additional_residuals,
+            is_self_attn=True,
+            **cross_attention_kwargs,
+        )
+        if self.use_ada_layer_norm_zero:
+            attn_output = gate_msa.unsqueeze(1) * attn_output
+        hidden_states = attn_output + hidden_states
+
+        # 2. Cross-Attention
+        if self.attn2 is not None:
+            norm_hidden_states = (
+                self.norm2(hidden_states, timestep) if self.use_ada_layer_norm else self.norm2(hidden_states)
+            )
+
+            attn_output = self.attn2(
+                norm_hidden_states,
+                encoder_hidden_states=encoder_hidden_states,
+                attention_mask=encoder_attention_mask,
+                block_idx=block_idx,
+                additional_residuals=additional_residuals,
+                is_self_attn=False,
+                **cross_attention_kwargs,
+            )
+            hidden_states = attn_output + hidden_states
+
+        # 3. Feed-forward
+        norm_hidden_states = self.norm3(hidden_states)
+
+        if self.use_ada_layer_norm_zero:
+            norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
+
+        if self._chunk_size is not None:
+            # "feed_forward_chunk_size" can be used to save memory
+            if norm_hidden_states.shape[self._chunk_dim] % self._chunk_size != 0:
+                raise ValueError(
+                    f"`hidden_states` dimension to be chunked: {norm_hidden_states.shape[self._chunk_dim]} has to be divisible by chunk size: {self._chunk_size}. Make sure to set an appropriate `chunk_size` when calling `unet.enable_forward_chunking`."
+                )
+
+            num_chunks = norm_hidden_states.shape[self._chunk_dim] // self._chunk_size
+            ff_output = torch.cat(
+                [self.ff(hid_slice) for hid_slice in norm_hidden_states.chunk(num_chunks, dim=self._chunk_dim)],
+                dim=self._chunk_dim,
+            )
+        else:
+            ff_output = self.ff(norm_hidden_states)
+
+        if self.use_ada_layer_norm_zero:
+            ff_output = gate_mlp.unsqueeze(1) * ff_output
+
+        hidden_states = ff_output + hidden_states
+
+        return hidden_states
+
+
+class ResidualResnetBlock2D(ResnetBlock2D):
+    def forward(self, input_tensor, temb, block_idx: Optional[int] = None,
+                additional_residuals: Optional[Dict[str, torch.FloatTensor]] = None):
+        hidden_states = input_tensor
+
+        if self.time_embedding_norm == "ada_group" or self.time_embedding_norm == "spatial":
+            hidden_states = self.norm1(hidden_states, temb)
+        else:
+            hidden_states = self.norm1(hidden_states)
+
+        hidden_states = self.nonlinearity(hidden_states)
+
+        if self.upsample is not None:
+            # upsample_nearest_nhwc fails with large batch sizes. see https://github.com/huggingface/diffusers/issues/984
+            if hidden_states.shape[0] >= 64:
+                input_tensor = input_tensor.contiguous()
+                hidden_states = hidden_states.contiguous()
+            input_tensor = self.upsample(input_tensor)
+            hidden_states = self.upsample(hidden_states)
+        elif self.downsample is not None:
+            input_tensor = self.downsample(input_tensor)
+            hidden_states = self.downsample(hidden_states)
+
+        hidden_states = self.conv1(hidden_states)
+
+        if self.time_emb_proj is not None:
+            if not self.skip_time_act:
+                temb = self.nonlinearity(temb)
+            temb = self.time_emb_proj(temb)[:, :, None, None]
+
+        if temb is not None and self.time_embedding_norm == "default":
+            hidden_states = hidden_states + temb
+
+        if self.time_embedding_norm == "ada_group" or self.time_embedding_norm == "spatial":
+            hidden_states = self.norm2(hidden_states, temb)
+        else:
+            hidden_states = self.norm2(hidden_states)
+
+        if temb is not None and self.time_embedding_norm == "scale_shift":
+            scale, shift = torch.chunk(temb, 2, dim=1)
+            hidden_states = hidden_states * (1 + scale) + shift
+
+        hidden_states = self.nonlinearity(hidden_states)
+
+        hidden_states = self.dropout(hidden_states)
+        hidden_states = self.conv2(hidden_states)
+
+        if self.conv_shortcut is not None:
+            input_tensor = self.conv_shortcut(input_tensor)
+
+        if additional_residuals and f"block_{block_idx}_resnet_feat" in additional_residuals:
+            hidden_states = hidden_states + additional_residuals[f"block_{block_idx}_resnet_feat"]
+
+        output_tensor = (input_tensor + hidden_states) / self.output_scale_factor
+
+        return output_tensor
+
+
+class ResidualTransformer2DModel(Transformer2DModel):
+    @register_to_config
+    def __init__(
+        self,
+        num_attention_heads: int = 16,
+        attention_head_dim: int = 88,
+        in_channels: Optional[int] = None,
+        out_channels: Optional[int] = None,
+        num_layers: int = 1,
+        dropout: float = 0.0,
+        norm_num_groups: int = 32,
+        cross_attention_dim: Optional[int] = None,
+        attention_bias: bool = False,
+        sample_size: Optional[int] = None,
+        num_vector_embeds: Optional[int] = None,
+        patch_size: Optional[int] = None,
+        activation_fn: str = "geglu",
+        num_embeds_ada_norm: Optional[int] = None,
+        use_linear_projection: bool = False,
+        only_cross_attention: bool = False,
+        upcast_attention: bool = False,
+        norm_type: str = "layer_norm",
+        norm_elementwise_affine: bool = True,
+    ):
+        super(Transformer2DModel, self).__init__()
+        self.use_linear_projection = use_linear_projection
+        self.num_attention_heads = num_attention_heads
+        self.attention_head_dim = attention_head_dim
+        inner_dim = num_attention_heads * attention_head_dim
+
+        # 1. Transformer2DModel can process both standard continuous images of shape `(batch_size, num_channels, width, height)` as well as quantized image embeddings of shape `(batch_size, num_image_vectors)`
+        # Define whether input is continuous or discrete depending on configuration
+        self.is_input_continuous = (in_channels is not None) and (patch_size is None)
+        self.is_input_vectorized = num_vector_embeds is not None
+        self.is_input_patches = in_channels is not None and patch_size is not None
+
+        if norm_type == "layer_norm" and num_embeds_ada_norm is not None:
+            deprecation_message = (
+                f"The configuration file of this model: {self.__class__} is outdated. `norm_type` is either not set or"
+                " incorrectly set to `'layer_norm'`.Make sure to set `norm_type` to `'ada_norm'` in the config."
+                " Please make sure to update the config accordingly as leaving `norm_type` might led to incorrect"
+                " results in future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it"
+                " would be very nice if you could open a Pull request for the `transformer/config.json` file"
+            )
+            deprecate("norm_type!=num_embeds_ada_norm", "1.0.0", deprecation_message, standard_warn=False)
+            norm_type = "ada_norm"
+
+        if self.is_input_continuous and self.is_input_vectorized:
+            raise ValueError(
+                f"Cannot define both `in_channels`: {in_channels} and `num_vector_embeds`: {num_vector_embeds}. Make"
+                " sure that either `in_channels` or `num_vector_embeds` is None."
+            )
+        elif self.is_input_vectorized and self.is_input_patches:
+            raise ValueError(
+                f"Cannot define both `num_vector_embeds`: {num_vector_embeds} and `patch_size`: {patch_size}. Make"
+                " sure that either `num_vector_embeds` or `num_patches` is None."
+            )
+        elif not self.is_input_continuous and not self.is_input_vectorized and not self.is_input_patches:
+            raise ValueError(
+                f"Has to define `in_channels`: {in_channels}, `num_vector_embeds`: {num_vector_embeds}, or patch_size:"
+                f" {patch_size}. Make sure that `in_channels`, `num_vector_embeds` or `num_patches` is not None."
+            )
+
+        # 2. Define input layers
+        if self.is_input_continuous:
+            self.in_channels = in_channels
+
+            self.norm = torch.nn.GroupNorm(num_groups=norm_num_groups, num_channels=in_channels, eps=1e-6, affine=True)
+            if use_linear_projection:
+                self.proj_in = LoRACompatibleLinear(in_channels, inner_dim)
+            else:
+                self.proj_in = LoRACompatibleConv(in_channels, inner_dim, kernel_size=1, stride=1, padding=0)
+        elif self.is_input_vectorized:
+            assert sample_size is not None, "Transformer2DModel over discrete input must provide sample_size"
+            assert num_vector_embeds is not None, "Transformer2DModel over discrete input must provide num_embed"
+
+            self.height = sample_size
+            self.width = sample_size
+            self.num_vector_embeds = num_vector_embeds
+            self.num_latent_pixels = self.height * self.width
+
+            self.latent_image_embedding = ImagePositionalEmbeddings(
+                num_embed=num_vector_embeds, embed_dim=inner_dim, height=self.height, width=self.width
+            )
+        elif self.is_input_patches:
+            assert sample_size is not None, "Transformer2DModel over patched input must provide sample_size"
+
+            self.height = sample_size
+            self.width = sample_size
+
+            self.patch_size = patch_size
+            self.pos_embed = PatchEmbed(
+                height=sample_size,
+                width=sample_size,
+                patch_size=patch_size,
+                in_channels=in_channels,
+                embed_dim=inner_dim,
+            )
+
+        # 3. Define transformers blocks
+        self.transformer_blocks = nn.ModuleList(
+            [
+                ResidualTransformerBlock(
+                    inner_dim,
+                    num_attention_heads,
+                    attention_head_dim,
+                    dropout=dropout,
+                    cross_attention_dim=cross_attention_dim,
+                    activation_fn=activation_fn,
+                    num_embeds_ada_norm=num_embeds_ada_norm,
+                    attention_bias=attention_bias,
+                    only_cross_attention=only_cross_attention,
+                    upcast_attention=upcast_attention,
+                    norm_type=norm_type,
+                    norm_elementwise_affine=norm_elementwise_affine,
+                )
+                for d in range(num_layers)
+            ]
+        )
+
+        # 4. Define output layers
+        self.out_channels = in_channels if out_channels is None else out_channels
+        if self.is_input_continuous:
+            # TODO: should use out_channels for continuous projections
+            if use_linear_projection:
+                self.proj_out = LoRACompatibleLinear(inner_dim, in_channels)
+            else:
+                self.proj_out = LoRACompatibleConv(inner_dim, in_channels, kernel_size=1, stride=1, padding=0)
+        elif self.is_input_vectorized:
+            self.norm_out = nn.LayerNorm(inner_dim)
+            self.out = nn.Linear(inner_dim, self.num_vector_embeds - 1)
+        elif self.is_input_patches:
+            self.norm_out = nn.LayerNorm(inner_dim, elementwise_affine=False, eps=1e-6)
+            self.proj_out_1 = nn.Linear(inner_dim, 2 * inner_dim)
+            self.proj_out_2 = nn.Linear(inner_dim, patch_size * patch_size * self.out_channels)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: Optional[torch.Tensor] = None,
+        timestep: Optional[torch.LongTensor] = None,
+        class_labels: Optional[torch.LongTensor] = None,
+        cross_attention_kwargs: Dict[str, Any] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        encoder_attention_mask: Optional[torch.Tensor] = None,
+        block_idx: Optional[int] = None,
+        additional_residuals: Optional[Dict[str, torch.FloatTensor]] = None,
+        return_dict: bool = True,
+    ):
+        """
+        The [`Transformer2DModel`] forward method.
+
+        Args:
+            hidden_states (`torch.LongTensor` of shape `(batch size, num latent pixels)` if discrete, `torch.FloatTensor` of shape `(batch size, channel, height, width)` if continuous):
+                Input `hidden_states`.
+            encoder_hidden_states ( `torch.FloatTensor` of shape `(batch size, sequence len, embed dims)`, *optional*):
+                Conditional embeddings for cross attention layer. If not given, cross-attention defaults to
+                self-attention.
+            timestep ( `torch.LongTensor`, *optional*):
+                Used to indicate denoising step. Optional timestep to be applied as an embedding in `AdaLayerNorm`.
+            class_labels ( `torch.LongTensor` of shape `(batch size, num classes)`, *optional*):
+                Used to indicate class labels conditioning. Optional class labels to be applied as an embedding in
+                `AdaLayerZeroNorm`.
+            encoder_attention_mask ( `torch.Tensor`, *optional*):
+                Cross-attention mask applied to `encoder_hidden_states`. Two formats supported:
+
+                    * Mask `(batch, sequence_length)` True = keep, False = discard.
+                    * Bias `(batch, 1, sequence_length)` 0 = keep, -10000 = discard.
+
+                If `ndim == 2`: will be interpreted as a mask, then converted into a bias consistent with the format
+                above. This bias will be added to the cross-attention scores.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~models.unet_2d_condition.UNet2DConditionOutput`] instead of a plain
+                tuple.
+
+        Returns:
+            If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a
+            `tuple` where the first element is the sample tensor.
+        """
+        # ensure attention_mask is a bias, and give it a singleton query_tokens dimension.
+        #   we may have done this conversion already, e.g. if we came here via UNet2DConditionModel#forward.
+        #   we can tell by counting dims; if ndim == 2: it's a mask rather than a bias.
+        # expects mask of shape:
+        #   [batch, key_tokens]
+        # adds singleton query_tokens dimension:
+        #   [batch,                    1, key_tokens]
+        # this helps to broadcast it as a bias over attention scores, which will be in one of the following shapes:
+        #   [batch,  heads, query_tokens, key_tokens] (e.g. torch sdp attn)
+        #   [batch * heads, query_tokens, key_tokens] (e.g. xformers or classic attn)
+        if attention_mask is not None and attention_mask.ndim == 2:
+            # assume that mask is expressed as:
+            #   (1 = keep,      0 = discard)
+            # convert mask into a bias that can be added to attention scores:
+            #       (keep = +0,     discard = -10000.0)
+            attention_mask = (1 - attention_mask.to(hidden_states.dtype)) * -10000.0
+            attention_mask = attention_mask.unsqueeze(1)
+
+        # convert encoder_attention_mask to a bias the same way we do for attention_mask
+        if encoder_attention_mask is not None and encoder_attention_mask.ndim == 2:
+            encoder_attention_mask = (1 - encoder_attention_mask.to(hidden_states.dtype)) * -10000.0
+            encoder_attention_mask = encoder_attention_mask.unsqueeze(1)
+
+        # 1. Input
+        if self.is_input_continuous:
+            batch, _, height, width = hidden_states.shape
+            residual = hidden_states
+
+            hidden_states = self.norm(hidden_states)
+            if not self.use_linear_projection:
+                hidden_states = self.proj_in(hidden_states)
+                inner_dim = hidden_states.shape[1]
+                hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(batch, height * width, inner_dim)
+            else:
+                inner_dim = hidden_states.shape[1]
+                hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(batch, height * width, inner_dim)
+                hidden_states = self.proj_in(hidden_states)
+        elif self.is_input_vectorized:
+            hidden_states = self.latent_image_embedding(hidden_states)
+        elif self.is_input_patches:
+            hidden_states = self.pos_embed(hidden_states)
+
+        # 2. Blocks
+        for block in self.transformer_blocks:
+            hidden_states = block(
+                hidden_states,
+                attention_mask=attention_mask,
+                encoder_hidden_states=encoder_hidden_states,
+                encoder_attention_mask=encoder_attention_mask,
+                timestep=timestep,
+                cross_attention_kwargs=cross_attention_kwargs,
+                class_labels=class_labels,
+                block_idx=block_idx,
+                additional_residuals=additional_residuals
+            )
+
+        # 3. Output
+        if self.is_input_continuous:
+            if not self.use_linear_projection:
+                hidden_states = hidden_states.reshape(batch, height, width, inner_dim).permute(0, 3, 1, 2).contiguous()
+                hidden_states = self.proj_out(hidden_states)
+            else:
+                hidden_states = self.proj_out(hidden_states)
+                hidden_states = hidden_states.reshape(batch, height, width, inner_dim).permute(0, 3, 1, 2).contiguous()
+
+            output = hidden_states + residual
+        elif self.is_input_vectorized:
+            hidden_states = self.norm_out(hidden_states)
+            logits = self.out(hidden_states)
+            # (batch, self.num_vector_embeds - 1, self.num_latent_pixels)
+            logits = logits.permute(0, 2, 1)
+
+            # log(p(x_0))
+            output = F.log_softmax(logits.double(), dim=1).float()
+        elif self.is_input_patches:
+            # TODO: cleanup!
+            conditioning = self.transformer_blocks[0].norm1.emb(
+                timestep, class_labels, hidden_dtype=hidden_states.dtype
+            )
+            shift, scale = self.proj_out_1(F.silu(conditioning)).chunk(2, dim=1)
+            hidden_states = self.norm_out(hidden_states) * (1 + scale[:, None]) + shift[:, None]
+            hidden_states = self.proj_out_2(hidden_states)
+
+            # unpatchify
+            height = width = int(hidden_states.shape[1] ** 0.5)
+            hidden_states = hidden_states.reshape(
+                shape=(-1, height, width, self.patch_size, self.patch_size, self.out_channels)
+            )
+            hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states)
+            output = hidden_states.reshape(
+                shape=(-1, self.out_channels, height * self.patch_size, width * self.patch_size)
+            )
+
+        if not return_dict:
+            return (output,)
+
+        return Transformer2DModelOutput(sample=output)
+
+
+class ResidualUpBlock2D(UpBlock2D):
+    def __init__(
+        self,
+        in_channels: int,
+        prev_output_channel: int,
+        out_channels: int,
+        temb_channels: int,
+        dropout: float = 0.0,
+        num_layers: int = 1,
+        resnet_eps: float = 1e-6,
+        resnet_time_scale_shift: str = "default",
+        resnet_act_fn: str = "swish",
+        resnet_groups: int = 32,
+        resnet_pre_norm: bool = True,
+        output_scale_factor=1.0,
+        add_upsample=True,
+    ):
+        super(UpBlock2D, self).__init__()
+        resnets = []
+
+        for i in range(num_layers):
+            res_skip_channels = in_channels if (i == num_layers - 1) else out_channels
+            resnet_in_channels = prev_output_channel if i == 0 else out_channels
+
+            resnets.append(
+                ResidualResnetBlock2D(
+                    in_channels=resnet_in_channels + res_skip_channels,
+                    out_channels=out_channels,
+                    temb_channels=temb_channels,
+                    eps=resnet_eps,
+                    groups=resnet_groups,
+                    dropout=dropout,
+                    time_embedding_norm=resnet_time_scale_shift,
+                    non_linearity=resnet_act_fn,
+                    output_scale_factor=output_scale_factor,
+                    pre_norm=resnet_pre_norm,
+                )
+            )
+
+        self.resnets = nn.ModuleList(resnets)
+
+        if add_upsample:
+            self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)])
+        else:
+            self.upsamplers = None
+
+        self.gradient_checkpointing = False
+
+    def forward(self, hidden_states, res_hidden_states_tuple, temb=None, upsample_size=None,
+                additional_residuals: Optional[Dict[str, torch.FloatTensor]] = None):
+        for j, resnet in enumerate(self.resnets):
+            # pop res hidden states
+            res_hidden_states = res_hidden_states_tuple[-1]
+            res_hidden_states_tuple = res_hidden_states_tuple[:-1]
+            hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)
+
+            if self.training and self.gradient_checkpointing:
+
+                def create_custom_forward(module):
+                    def custom_forward(*inputs):
+                        return module(*inputs)
+
+                    return custom_forward
+
+                if is_torch_version(">=", "1.11.0"):
+                    hidden_states = torch.utils.checkpoint.checkpoint(
+                        create_custom_forward(resnet), hidden_states, temb, use_reentrant=False
+                    )
+                else:
+                    hidden_states = torch.utils.checkpoint.checkpoint(
+                        create_custom_forward(resnet), hidden_states, temb
+                    )
+            else:
+                hidden_states = resnet(hidden_states, temb, block_idx=j, additional_residuals=additional_residuals)
+
+        if self.upsamplers is not None:
+            for upsampler in self.upsamplers:
+                hidden_states = upsampler(hidden_states, upsample_size)
+
+        return hidden_states
+
+
+class ResidualCrossAttnUpBlock2D(CrossAttnUpBlock2D):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        prev_output_channel: int,
+        temb_channels: int,
+        dropout: float = 0.0,
+        num_layers: int = 1,
+        transformer_layers_per_block: int = 1,
+        resnet_eps: float = 1e-6,
+        resnet_time_scale_shift: str = "default",
+        resnet_act_fn: str = "swish",
+        resnet_groups: int = 32,
+        resnet_pre_norm: bool = True,
+        num_attention_heads=1,
+        cross_attention_dim=1280,
+        output_scale_factor=1.0,
+        add_upsample=True,
+        dual_cross_attention=False,
+        use_linear_projection=False,
+        only_cross_attention=False,
+        upcast_attention=False,
+    ):
+        super(CrossAttnUpBlock2D, self).__init__()
+        resnets = []
+        attentions = []
+
+        self.has_cross_attention = True
+        self.num_attention_heads = num_attention_heads
+
+        for i in range(num_layers):
+            res_skip_channels = in_channels if (i == num_layers - 1) else out_channels
+            resnet_in_channels = prev_output_channel if i == 0 else out_channels
+
+            resnets.append(
+                ResidualResnetBlock2D(
+                    in_channels=resnet_in_channels + res_skip_channels,
+                    out_channels=out_channels,
+                    temb_channels=temb_channels,
+                    eps=resnet_eps,
+                    groups=resnet_groups,
+                    dropout=dropout,
+                    time_embedding_norm=resnet_time_scale_shift,
+                    non_linearity=resnet_act_fn,
+                    output_scale_factor=output_scale_factor,
+                    pre_norm=resnet_pre_norm,
+                )
+            )
+            if not dual_cross_attention:
+                attentions.append(
+                    ResidualTransformer2DModel(
+                        num_attention_heads,
+                        out_channels // num_attention_heads,
+                        in_channels=out_channels,
+                        num_layers=transformer_layers_per_block,
+                        cross_attention_dim=cross_attention_dim,
+                        norm_num_groups=resnet_groups,
+                        use_linear_projection=use_linear_projection,
+                        only_cross_attention=only_cross_attention,
+                        upcast_attention=upcast_attention,
+                    )
+                )
+            else:
+                attentions.append(
+                    DualTransformer2DModel(
+                        num_attention_heads,
+                        out_channels // num_attention_heads,
+                        in_channels=out_channels,
+                        num_layers=1,
+                        cross_attention_dim=cross_attention_dim,
+                        norm_num_groups=resnet_groups,
+                    )
+                )
+        self.attentions = nn.ModuleList(attentions)
+        self.resnets = nn.ModuleList(resnets)
+
+        if add_upsample:
+            self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)])
+        else:
+            self.upsamplers = None
+
+        self.gradient_checkpointing = False
+
+    def forward(
+        self,
+        hidden_states: torch.FloatTensor,
+        res_hidden_states_tuple: Tuple[torch.FloatTensor, ...],
+        temb: Optional[torch.FloatTensor] = None,
+        encoder_hidden_states: Optional[torch.FloatTensor] = None,
+        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+        upsample_size: Optional[int] = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        encoder_attention_mask: Optional[torch.FloatTensor] = None,
+        block_idx: Optional[int] = None,
+        additional_residuals: Optional[Dict[str, torch.FloatTensor]] = None
+    ):
+        for j, (resnet, attn) in enumerate(zip(self.resnets, self.attentions)):
+            # pop res hidden states
+            res_hidden_states = res_hidden_states_tuple[-1]
+            res_hidden_states_tuple = res_hidden_states_tuple[:-1]
+            hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)
+
+            if self.training and self.gradient_checkpointing:
+
+                def create_custom_forward(module, return_dict=None):
+                    def custom_forward(*inputs):
+                        if return_dict is not None:
+                            return module(*inputs, return_dict=return_dict)
+                        else:
+                            return module(*inputs)
+
+                    return custom_forward
+
+                ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+                hidden_states = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(resnet),
+                    hidden_states,
+                    temb,
+                    block_idx * len(self.resnets) + j,
+                    additional_residuals,
+                    **ckpt_kwargs,
+                )
+                hidden_states = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(attn, return_dict=False),
+                    hidden_states,
+                    encoder_hidden_states,
+                    None,  # timestep
+                    None,  # class_labels
+                    cross_attention_kwargs,
+                    attention_mask,
+                    encoder_attention_mask,
+                    block_idx * len(self.resnets) + j,
+                    additional_residuals,
+                    **ckpt_kwargs,
+                )[0]
+            else:
+                hidden_states = resnet(hidden_states, temb,
+                                       block_idx * len(self.resnets) + j,
+                                       additional_residuals)
+                hidden_states = attn(
+                    hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    cross_attention_kwargs=cross_attention_kwargs,
+                    attention_mask=attention_mask,
+                    encoder_attention_mask=encoder_attention_mask,
+                    return_dict=False,
+                    block_idx=block_idx * len(self.resnets) + j,
+                    additional_residuals=additional_residuals
+                )[0]
+
+        if self.upsamplers is not None:
+            for upsampler in self.upsamplers:
+                hidden_states = upsampler(hidden_states, upsample_size)
+
+        return hidden_states
+
+
+def get_residual_up_block(
+    up_block_type,
+    num_layers,
+    in_channels,
+    out_channels,
+    prev_output_channel,
+    temb_channels,
+    add_upsample,
+    resnet_eps,
+    resnet_act_fn,
+    transformer_layers_per_block=1,
+    num_attention_heads=None,
+    resnet_groups=None,
+    cross_attention_dim=None,
+    dual_cross_attention=False,
+    use_linear_projection=False,
+    only_cross_attention=False,
+    upcast_attention=False,
+    resnet_time_scale_shift="default",
+    resnet_skip_time_act=False,
+    resnet_out_scale_factor=1.0,
+    cross_attention_norm=None,
+    attention_head_dim=None,
+    upsample_type=None,
+):
+    # If attn head dim is not defined, we default it to the number of heads
+    if attention_head_dim is None:
+        logger.warn(
+            f"It is recommended to provide `attention_head_dim` when calling `get_up_block`. Defaulting `attention_head_dim` to {num_attention_heads}."
+        )
+        attention_head_dim = num_attention_heads
+
+    up_block_type = up_block_type[7:] if up_block_type.startswith("UNetRes") else up_block_type
+    if up_block_type == "UpBlock2D":
+        return ResidualUpBlock2D(
+            num_layers=num_layers,
+            in_channels=in_channels,
+            out_channels=out_channels,
+            prev_output_channel=prev_output_channel,
+            temb_channels=temb_channels,
+            add_upsample=add_upsample,
+            resnet_eps=resnet_eps,
+            resnet_act_fn=resnet_act_fn,
+            resnet_groups=resnet_groups,
+            resnet_time_scale_shift=resnet_time_scale_shift,
+        )
+    elif up_block_type == "ResnetUpsampleBlock2D":
+        return ResnetUpsampleBlock2D(
+            num_layers=num_layers,
+            in_channels=in_channels,
+            out_channels=out_channels,
+            prev_output_channel=prev_output_channel,
+            temb_channels=temb_channels,
+            add_upsample=add_upsample,
+            resnet_eps=resnet_eps,
+            resnet_act_fn=resnet_act_fn,
+            resnet_groups=resnet_groups,
+            resnet_time_scale_shift=resnet_time_scale_shift,
+            skip_time_act=resnet_skip_time_act,
+            output_scale_factor=resnet_out_scale_factor,
+        )
+    elif up_block_type == "CrossAttnUpBlock2D":
+        if cross_attention_dim is None:
+            raise ValueError("cross_attention_dim must be specified for CrossAttnUpBlock2D")
+        return ResidualCrossAttnUpBlock2D(
+            num_layers=num_layers,
+            transformer_layers_per_block=transformer_layers_per_block,
+            in_channels=in_channels,
+            out_channels=out_channels,
+            prev_output_channel=prev_output_channel,
+            temb_channels=temb_channels,
+            add_upsample=add_upsample,
+            resnet_eps=resnet_eps,
+            resnet_act_fn=resnet_act_fn,
+            resnet_groups=resnet_groups,
+            cross_attention_dim=cross_attention_dim,
+            num_attention_heads=num_attention_heads,
+            dual_cross_attention=dual_cross_attention,
+            use_linear_projection=use_linear_projection,
+            only_cross_attention=only_cross_attention,
+            upcast_attention=upcast_attention,
+            resnet_time_scale_shift=resnet_time_scale_shift,
+        )
+    elif up_block_type == "SimpleCrossAttnUpBlock2D":
+        if cross_attention_dim is None:
+            raise ValueError("cross_attention_dim must be specified for SimpleCrossAttnUpBlock2D")
+        return SimpleCrossAttnUpBlock2D(
+            num_layers=num_layers,
+            in_channels=in_channels,
+            out_channels=out_channels,
+            prev_output_channel=prev_output_channel,
+            temb_channels=temb_channels,
+            add_upsample=add_upsample,
+            resnet_eps=resnet_eps,
+            resnet_act_fn=resnet_act_fn,
+            resnet_groups=resnet_groups,
+            cross_attention_dim=cross_attention_dim,
+            attention_head_dim=attention_head_dim,
+            resnet_time_scale_shift=resnet_time_scale_shift,
+            skip_time_act=resnet_skip_time_act,
+            output_scale_factor=resnet_out_scale_factor,
+            only_cross_attention=only_cross_attention,
+            cross_attention_norm=cross_attention_norm,
+        )
+    elif up_block_type == "AttnUpBlock2D":
+        if add_upsample is False:
+            upsample_type = None
+        else:
+            upsample_type = upsample_type or "conv"  # default to 'conv'
+
+        return AttnUpBlock2D(
+            num_layers=num_layers,
+            in_channels=in_channels,
+            out_channels=out_channels,
+            prev_output_channel=prev_output_channel,
+            temb_channels=temb_channels,
+            resnet_eps=resnet_eps,
+            resnet_act_fn=resnet_act_fn,
+            resnet_groups=resnet_groups,
+            attention_head_dim=attention_head_dim,
+            resnet_time_scale_shift=resnet_time_scale_shift,
+            upsample_type=upsample_type,
+        )
+    elif up_block_type == "SkipUpBlock2D":
+        return SkipUpBlock2D(
+            num_layers=num_layers,
+            in_channels=in_channels,
+            out_channels=out_channels,
+            prev_output_channel=prev_output_channel,
+            temb_channels=temb_channels,
+            add_upsample=add_upsample,
+            resnet_eps=resnet_eps,
+            resnet_act_fn=resnet_act_fn,
+            resnet_time_scale_shift=resnet_time_scale_shift,
+        )
+    elif up_block_type == "AttnSkipUpBlock2D":
+        return AttnSkipUpBlock2D(
+            num_layers=num_layers,
+            in_channels=in_channels,
+            out_channels=out_channels,
+            prev_output_channel=prev_output_channel,
+            temb_channels=temb_channels,
+            add_upsample=add_upsample,
+            resnet_eps=resnet_eps,
+            resnet_act_fn=resnet_act_fn,
+            attention_head_dim=attention_head_dim,
+            resnet_time_scale_shift=resnet_time_scale_shift,
+        )
+    elif up_block_type == "UpDecoderBlock2D":
+        return UpDecoderBlock2D(
+            num_layers=num_layers,
+            in_channels=in_channels,
+            out_channels=out_channels,
+            add_upsample=add_upsample,
+            resnet_eps=resnet_eps,
+            resnet_act_fn=resnet_act_fn,
+            resnet_groups=resnet_groups,
+            resnet_time_scale_shift=resnet_time_scale_shift,
+            temb_channels=temb_channels,
+        )
+    elif up_block_type == "AttnUpDecoderBlock2D":
+        return AttnUpDecoderBlock2D(
+            num_layers=num_layers,
+            in_channels=in_channels,
+            out_channels=out_channels,
+            add_upsample=add_upsample,
+            resnet_eps=resnet_eps,
+            resnet_act_fn=resnet_act_fn,
+            resnet_groups=resnet_groups,
+            attention_head_dim=attention_head_dim,
+            resnet_time_scale_shift=resnet_time_scale_shift,
+            temb_channels=temb_channels,
+        )
+    elif up_block_type == "KUpBlock2D":
+        return KUpBlock2D(
+            num_layers=num_layers,
+            in_channels=in_channels,
+            out_channels=out_channels,
+            temb_channels=temb_channels,
+            add_upsample=add_upsample,
+            resnet_eps=resnet_eps,
+            resnet_act_fn=resnet_act_fn,
+        )
+    elif up_block_type == "KCrossAttnUpBlock2D":
+        return KCrossAttnUpBlock2D(
+            num_layers=num_layers,
+            in_channels=in_channels,
+            out_channels=out_channels,
+            temb_channels=temb_channels,
+            add_upsample=add_upsample,
+            resnet_eps=resnet_eps,
+            resnet_act_fn=resnet_act_fn,
+            cross_attention_dim=cross_attention_dim,
+            attention_head_dim=attention_head_dim,
+        )
+
+    raise ValueError(f"{up_block_type} does not exist.")
+
+
+class ResidualUNet2DConditionModel(UNet2DConditionModel):
+    @register_to_config
+    def __init__(
+        self,
+        sample_size: Optional[int] = None,
+        in_channels: int = 4,
+        out_channels: int = 4,
+        center_input_sample: bool = False,
+        flip_sin_to_cos: bool = True,
+        freq_shift: int = 0,
+        down_block_types: Tuple[str] = (
+            "CrossAttnDownBlock2D",
+            "CrossAttnDownBlock2D",
+            "CrossAttnDownBlock2D",
+            "DownBlock2D",
+        ),
+        mid_block_type: Optional[str] = "UNetMidBlock2DCrossAttn",
+        up_block_types: Tuple[str] = ("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D"),
+        only_cross_attention: Union[bool, Tuple[bool]] = False,
+        block_out_channels: Tuple[int] = (320, 640, 1280, 1280),
+        layers_per_block: Union[int, Tuple[int]] = 2,
+        downsample_padding: int = 1,
+        mid_block_scale_factor: float = 1,
+        act_fn: str = "silu",
+        norm_num_groups: Optional[int] = 32,
+        norm_eps: float = 1e-5,
+        cross_attention_dim: Union[int, Tuple[int]] = 1280,
+        transformer_layers_per_block: Union[int, Tuple[int]] = 1,
+        encoder_hid_dim: Optional[int] = None,
+        encoder_hid_dim_type: Optional[str] = None,
+        attention_head_dim: Union[int, Tuple[int]] = 8,
+        num_attention_heads: Optional[Union[int, Tuple[int]]] = None,
+        dual_cross_attention: bool = False,
+        use_linear_projection: bool = False,
+        class_embed_type: Optional[str] = None,
+        addition_embed_type: Optional[str] = None,
+        addition_time_embed_dim: Optional[int] = None,
+        num_class_embeds: Optional[int] = None,
+        upcast_attention: bool = False,
+        resnet_time_scale_shift: str = "default",
+        resnet_skip_time_act: bool = False,
+        resnet_out_scale_factor: int = 1.0,
+        time_embedding_type: str = "positional",
+        time_embedding_dim: Optional[int] = None,
+        time_embedding_act_fn: Optional[str] = None,
+        timestep_post_act: Optional[str] = None,
+        time_cond_proj_dim: Optional[int] = None,
+        conv_in_kernel: int = 3,
+        conv_out_kernel: int = 3,
+        projection_class_embeddings_input_dim: Optional[int] = None,
+        class_embeddings_concat: bool = False,
+        mid_block_only_cross_attention: Optional[bool] = None,
+        cross_attention_norm: Optional[str] = None,
+        addition_embed_type_num_heads=64,
+    ):
+        super(UNet2DConditionModel, self).__init__()
+
+        self.sample_size = sample_size
+
+        if num_attention_heads is not None:
+            raise ValueError(
+                "At the moment it is not possible to define the number of attention heads via `num_attention_heads` because of a naming issue as described in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131. Passing `num_attention_heads` will only be supported in diffusers v0.19."
+            )
+
+        # If `num_attention_heads` is not defined (which is the case for most models)
+        # it will default to `attention_head_dim`. This looks weird upon first reading it and it is.
+        # The reason for this behavior is to correct for incorrectly named variables that were introduced
+        # when this library was created. The incorrect naming was only discovered much later in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131
+        # Changing `attention_head_dim` to `num_attention_heads` for 40,000+ configurations is too backwards breaking
+        # which is why we correct for the naming here.
+        num_attention_heads = num_attention_heads or attention_head_dim
+
+        # Check inputs
+        if len(down_block_types) != len(up_block_types):
+            raise ValueError(
+                f"Must provide the same number of `down_block_types` as `up_block_types`. `down_block_types`: {down_block_types}. `up_block_types`: {up_block_types}."
+            )
+
+        if len(block_out_channels) != len(down_block_types):
+            raise ValueError(
+                f"Must provide the same number of `block_out_channels` as `down_block_types`. `block_out_channels`: {block_out_channels}. `down_block_types`: {down_block_types}."
+            )
+
+        if not isinstance(only_cross_attention, bool) and len(only_cross_attention) != len(down_block_types):
+            raise ValueError(
+                f"Must provide the same number of `only_cross_attention` as `down_block_types`. `only_cross_attention`: {only_cross_attention}. `down_block_types`: {down_block_types}."
+            )
+
+        if not isinstance(num_attention_heads, int) and len(num_attention_heads) != len(down_block_types):
+            raise ValueError(
+                f"Must provide the same number of `num_attention_heads` as `down_block_types`. `num_attention_heads`: {num_attention_heads}. `down_block_types`: {down_block_types}."
+            )
+
+        if not isinstance(attention_head_dim, int) and len(attention_head_dim) != len(down_block_types):
+            raise ValueError(
+                f"Must provide the same number of `attention_head_dim` as `down_block_types`. `attention_head_dim`: {attention_head_dim}. `down_block_types`: {down_block_types}."
+            )
+
+        if isinstance(cross_attention_dim, list) and len(cross_attention_dim) != len(down_block_types):
+            raise ValueError(
+                f"Must provide the same number of `cross_attention_dim` as `down_block_types`. `cross_attention_dim`: {cross_attention_dim}. `down_block_types`: {down_block_types}."
+            )
+
+        if not isinstance(layers_per_block, int) and len(layers_per_block) != len(down_block_types):
+            raise ValueError(
+                f"Must provide the same number of `layers_per_block` as `down_block_types`. `layers_per_block`: {layers_per_block}. `down_block_types`: {down_block_types}."
+            )
+
+        # input
+        conv_in_padding = (conv_in_kernel - 1) // 2
+        self.conv_in = nn.Conv2d(
+            in_channels, block_out_channels[0], kernel_size=conv_in_kernel, padding=conv_in_padding
+        )
+
+        # time
+        if time_embedding_type == "fourier":
+            time_embed_dim = time_embedding_dim or block_out_channels[0] * 2
+            if time_embed_dim % 2 != 0:
+                raise ValueError(f"`time_embed_dim` should be divisible by 2, but is {time_embed_dim}.")
+            self.time_proj = GaussianFourierProjection(
+                time_embed_dim // 2, set_W_to_weight=False, log=False, flip_sin_to_cos=flip_sin_to_cos
+            )
+            timestep_input_dim = time_embed_dim
+        elif time_embedding_type == "positional":
+            time_embed_dim = time_embedding_dim or block_out_channels[0] * 4
+
+            self.time_proj = Timesteps(block_out_channels[0], flip_sin_to_cos, freq_shift)
+            timestep_input_dim = block_out_channels[0]
+        else:
+            raise ValueError(
+                f"{time_embedding_type} does not exist. Please make sure to use one of `fourier` or `positional`."
+            )
+
+        self.time_embedding = TimestepEmbedding(
+            timestep_input_dim,
+            time_embed_dim,
+            act_fn=act_fn,
+            post_act_fn=timestep_post_act,
+            cond_proj_dim=time_cond_proj_dim,
+        )
+
+        if encoder_hid_dim_type is None and encoder_hid_dim is not None:
+            encoder_hid_dim_type = "text_proj"
+            self.register_to_config(encoder_hid_dim_type=encoder_hid_dim_type)
+            logger.info("encoder_hid_dim_type defaults to 'text_proj' as `encoder_hid_dim` is defined.")
+
+        if encoder_hid_dim is None and encoder_hid_dim_type is not None:
+            raise ValueError(
+                f"`encoder_hid_dim` has to be defined when `encoder_hid_dim_type` is set to {encoder_hid_dim_type}."
+            )
+
+        if encoder_hid_dim_type == "text_proj":
+            self.encoder_hid_proj = nn.Linear(encoder_hid_dim, cross_attention_dim)
+        elif encoder_hid_dim_type == "text_image_proj":
+            # image_embed_dim DOESN'T have to be `cross_attention_dim`. To not clutter the __init__ too much
+            # they are set to `cross_attention_dim` here as this is exactly the required dimension for the currently only use
+            # case when `addition_embed_type == "text_image_proj"` (Kadinsky 2.1)`
+            self.encoder_hid_proj = TextImageProjection(
+                text_embed_dim=encoder_hid_dim,
+                image_embed_dim=cross_attention_dim,
+                cross_attention_dim=cross_attention_dim,
+            )
+        elif encoder_hid_dim_type == "image_proj":
+            # Kandinsky 2.2
+            self.encoder_hid_proj = ImageProjection(
+                image_embed_dim=encoder_hid_dim,
+                cross_attention_dim=cross_attention_dim,
+            )
+        elif encoder_hid_dim_type is not None:
+            raise ValueError(
+                f"encoder_hid_dim_type: {encoder_hid_dim_type} must be None, 'text_proj' or 'text_image_proj'."
+            )
+        else:
+            self.encoder_hid_proj = None
+
+        # class embedding
+        if class_embed_type is None and num_class_embeds is not None:
+            self.class_embedding = nn.Embedding(num_class_embeds, time_embed_dim)
+        elif class_embed_type == "timestep":
+            self.class_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim, act_fn=act_fn)
+        elif class_embed_type == "identity":
+            self.class_embedding = nn.Identity(time_embed_dim, time_embed_dim)
+        elif class_embed_type == "projection":
+            if projection_class_embeddings_input_dim is None:
+                raise ValueError(
+                    "`class_embed_type`: 'projection' requires `projection_class_embeddings_input_dim` be set"
+                )
+            # The projection `class_embed_type` is the same as the timestep `class_embed_type` except
+            # 1. the `class_labels` inputs are not first converted to sinusoidal embeddings
+            # 2. it projects from an arbitrary input dimension.
+            #
+            # Note that `TimestepEmbedding` is quite general, being mainly linear layers and activations.
+            # When used for embedding actual timesteps, the timesteps are first converted to sinusoidal embeddings.
+            # As a result, `TimestepEmbedding` can be passed arbitrary vectors.
+            self.class_embedding = TimestepEmbedding(projection_class_embeddings_input_dim, time_embed_dim)
+        elif class_embed_type == "simple_projection":
+            if projection_class_embeddings_input_dim is None:
+                raise ValueError(
+                    "`class_embed_type`: 'simple_projection' requires `projection_class_embeddings_input_dim` be set"
+                )
+            self.class_embedding = nn.Linear(projection_class_embeddings_input_dim, time_embed_dim)
+        else:
+            self.class_embedding = None
+
+        if addition_embed_type == "text":
+            if encoder_hid_dim is not None:
+                text_time_embedding_from_dim = encoder_hid_dim
+            else:
+                text_time_embedding_from_dim = cross_attention_dim
+
+            self.add_embedding = TextTimeEmbedding(
+                text_time_embedding_from_dim, time_embed_dim, num_heads=addition_embed_type_num_heads
+            )
+        elif addition_embed_type == "text_image":
+            # text_embed_dim and image_embed_dim DON'T have to be `cross_attention_dim`. To not clutter the __init__ too much
+            # they are set to `cross_attention_dim` here as this is exactly the required dimension for the currently only use
+            # case when `addition_embed_type == "text_image"` (Kadinsky 2.1)`
+            self.add_embedding = TextImageTimeEmbedding(
+                text_embed_dim=cross_attention_dim, image_embed_dim=cross_attention_dim, time_embed_dim=time_embed_dim
+            )
+        elif addition_embed_type == "text_time":
+            self.add_time_proj = Timesteps(addition_time_embed_dim, flip_sin_to_cos, freq_shift)
+            self.add_embedding = TimestepEmbedding(projection_class_embeddings_input_dim, time_embed_dim)
+        elif addition_embed_type == "image":
+            # Kandinsky 2.2
+            self.add_embedding = ImageTimeEmbedding(image_embed_dim=encoder_hid_dim, time_embed_dim=time_embed_dim)
+        elif addition_embed_type == "image_hint":
+            # Kandinsky 2.2 ControlNet
+            self.add_embedding = ImageHintTimeEmbedding(image_embed_dim=encoder_hid_dim, time_embed_dim=time_embed_dim)
+        elif addition_embed_type is not None:
+            raise ValueError(f"addition_embed_type: {addition_embed_type} must be None, 'text' or 'text_image'.")
+
+        if time_embedding_act_fn is None:
+            self.time_embed_act = None
+        else:
+            self.time_embed_act = get_activation(time_embedding_act_fn)
+
+        self.down_blocks = nn.ModuleList([])
+        self.up_blocks = nn.ModuleList([])
+
+        if isinstance(only_cross_attention, bool):
+            if mid_block_only_cross_attention is None:
+                mid_block_only_cross_attention = only_cross_attention
+
+            only_cross_attention = [only_cross_attention] * len(down_block_types)
+
+        if mid_block_only_cross_attention is None:
+            mid_block_only_cross_attention = False
+
+        if isinstance(num_attention_heads, int):
+            num_attention_heads = (num_attention_heads,) * len(down_block_types)
+
+        if isinstance(attention_head_dim, int):
+            attention_head_dim = (attention_head_dim,) * len(down_block_types)
+
+        if isinstance(cross_attention_dim, int):
+            cross_attention_dim = (cross_attention_dim,) * len(down_block_types)
+
+        if isinstance(layers_per_block, int):
+            layers_per_block = [layers_per_block] * len(down_block_types)
+
+        if isinstance(transformer_layers_per_block, int):
+            transformer_layers_per_block = [transformer_layers_per_block] * len(down_block_types)
+
+        if class_embeddings_concat:
+            # The time embeddings are concatenated with the class embeddings. The dimension of the
+            # time embeddings passed to the down, middle, and up blocks is twice the dimension of the
+            # regular time embeddings
+            blocks_time_embed_dim = time_embed_dim * 2
+        else:
+            blocks_time_embed_dim = time_embed_dim
+
+        # down
+        output_channel = block_out_channels[0]
+        for i, down_block_type in enumerate(down_block_types):
+            input_channel = output_channel
+            output_channel = block_out_channels[i]
+            is_final_block = i == len(block_out_channels) - 1
+
+            down_block = get_down_block(
+                down_block_type,
+                num_layers=layers_per_block[i],
+                transformer_layers_per_block=transformer_layers_per_block[i],
+                in_channels=input_channel,
+                out_channels=output_channel,
+                temb_channels=blocks_time_embed_dim,
+                add_downsample=not is_final_block,
+                resnet_eps=norm_eps,
+                resnet_act_fn=act_fn,
+                resnet_groups=norm_num_groups,
+                cross_attention_dim=cross_attention_dim[i],
+                num_attention_heads=num_attention_heads[i],
+                downsample_padding=downsample_padding,
+                dual_cross_attention=dual_cross_attention,
+                use_linear_projection=use_linear_projection,
+                only_cross_attention=only_cross_attention[i],
+                upcast_attention=upcast_attention,
+                resnet_time_scale_shift=resnet_time_scale_shift,
+                resnet_skip_time_act=resnet_skip_time_act,
+                resnet_out_scale_factor=resnet_out_scale_factor,
+                cross_attention_norm=cross_attention_norm,
+                attention_head_dim=attention_head_dim[i] if attention_head_dim[i] is not None else output_channel,
+            )
+            self.down_blocks.append(down_block)
+
+        # mid
+        if mid_block_type == "UNetMidBlock2DCrossAttn":
+            self.mid_block = UNetMidBlock2DCrossAttn(
+                transformer_layers_per_block=transformer_layers_per_block[-1],
+                in_channels=block_out_channels[-1],
+                temb_channels=blocks_time_embed_dim,
+                resnet_eps=norm_eps,
+                resnet_act_fn=act_fn,
+                output_scale_factor=mid_block_scale_factor,
+                resnet_time_scale_shift=resnet_time_scale_shift,
+                cross_attention_dim=cross_attention_dim[-1],
+                num_attention_heads=num_attention_heads[-1],
+                resnet_groups=norm_num_groups,
+                dual_cross_attention=dual_cross_attention,
+                use_linear_projection=use_linear_projection,
+                upcast_attention=upcast_attention,
+            )
+        elif mid_block_type == "UNetMidBlock2DSimpleCrossAttn":
+            self.mid_block = UNetMidBlock2DSimpleCrossAttn(
+                in_channels=block_out_channels[-1],
+                temb_channels=blocks_time_embed_dim,
+                resnet_eps=norm_eps,
+                resnet_act_fn=act_fn,
+                output_scale_factor=mid_block_scale_factor,
+                cross_attention_dim=cross_attention_dim[-1],
+                attention_head_dim=attention_head_dim[-1],
+                resnet_groups=norm_num_groups,
+                resnet_time_scale_shift=resnet_time_scale_shift,
+                skip_time_act=resnet_skip_time_act,
+                only_cross_attention=mid_block_only_cross_attention,
+                cross_attention_norm=cross_attention_norm,
+            )
+        elif mid_block_type is None:
+            self.mid_block = None
+        else:
+            raise ValueError(f"unknown mid_block_type : {mid_block_type}")
+
+        # count how many layers upsample the images
+        self.num_upsamplers = 0
+
+        # up
+        reversed_block_out_channels = list(reversed(block_out_channels))
+        reversed_num_attention_heads = list(reversed(num_attention_heads))
+        reversed_layers_per_block = list(reversed(layers_per_block))
+        reversed_cross_attention_dim = list(reversed(cross_attention_dim))
+        reversed_transformer_layers_per_block = list(reversed(transformer_layers_per_block))
+        only_cross_attention = list(reversed(only_cross_attention))
+
+        output_channel = reversed_block_out_channels[0]
+        for i, up_block_type in enumerate(up_block_types):
+            is_final_block = i == len(block_out_channels) - 1
+
+            prev_output_channel = output_channel
+            output_channel = reversed_block_out_channels[i]
+            input_channel = reversed_block_out_channels[min(i + 1, len(block_out_channels) - 1)]
+
+            # add upsample block for all BUT final layer
+            if not is_final_block:
+                add_upsample = True
+                self.num_upsamplers += 1
+            else:
+                add_upsample = False
+
+            up_block = get_residual_up_block(
+                up_block_type,
+                num_layers=reversed_layers_per_block[i] + 1,
+                transformer_layers_per_block=reversed_transformer_layers_per_block[i],
+                in_channels=input_channel,
+                out_channels=output_channel,
+                prev_output_channel=prev_output_channel,
+                temb_channels=blocks_time_embed_dim,
+                add_upsample=add_upsample,
+                resnet_eps=norm_eps,
+                resnet_act_fn=act_fn,
+                resnet_groups=norm_num_groups,
+                cross_attention_dim=reversed_cross_attention_dim[i],
+                num_attention_heads=reversed_num_attention_heads[i],
+                dual_cross_attention=dual_cross_attention,
+                use_linear_projection=use_linear_projection,
+                only_cross_attention=only_cross_attention[i],
+                upcast_attention=upcast_attention,
+                resnet_time_scale_shift=resnet_time_scale_shift,
+                resnet_skip_time_act=resnet_skip_time_act,
+                resnet_out_scale_factor=resnet_out_scale_factor,
+                cross_attention_norm=cross_attention_norm,
+                attention_head_dim=attention_head_dim[i] if attention_head_dim[i] is not None else output_channel,
+            )
+            self.up_blocks.append(up_block)
+            prev_output_channel = output_channel
+
+        # out
+        if norm_num_groups is not None:
+            self.conv_norm_out = nn.GroupNorm(
+                num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=norm_eps
+            )
+
+            self.conv_act = get_activation(act_fn)
+
+        else:
+            self.conv_norm_out = None
+            self.conv_act = None
+
+        conv_out_padding = (conv_out_kernel - 1) // 2
+        self.conv_out = nn.Conv2d(
+            block_out_channels[0], out_channels, kernel_size=conv_out_kernel, padding=conv_out_padding
+        )
+
+    def forward(
+        self,
+        sample: torch.FloatTensor,
+        timestep: Union[torch.Tensor, float, int],
+        encoder_hidden_states: torch.Tensor,
+        class_labels: Optional[torch.Tensor] = None,
+        timestep_cond: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+        added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None,
+        down_block_additional_residuals: Optional[Tuple[torch.Tensor]] = None,
+        mid_block_additional_residual: Optional[torch.Tensor] = None,
+        up_block_additional_residuals: Optional[Dict[str, torch.Tensor]] = None, # newly added
+        encoder_attention_mask: Optional[torch.Tensor] = None,
+        return_dict: bool = True,
+    ) -> Union[UNet2DConditionOutput, Tuple]:
+        r"""
+        The [`UNet2DConditionModel`] forward method.
+
+        Args:
+            sample (`torch.FloatTensor`):
+                The noisy input tensor with the following shape `(batch, channel, height, width)`.
+            timestep (`torch.FloatTensor` or `float` or `int`): The number of timesteps to denoise an input.
+            encoder_hidden_states (`torch.FloatTensor`):
+                The encoder hidden states with shape `(batch, sequence_length, feature_dim)`.
+            encoder_attention_mask (`torch.Tensor`):
+                A cross-attention mask of shape `(batch, sequence_length)` is applied to `encoder_hidden_states`. If
+                `True` the mask is kept, otherwise if `False` it is discarded. Mask will be converted into a bias,
+                which adds large negative values to the attention scores corresponding to "discard" tokens.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~models.unet_2d_condition.UNet2DConditionOutput`] instead of a plain
+                tuple.
+            cross_attention_kwargs (`dict`, *optional*):
+                A kwargs dictionary that if specified is passed along to the [`AttnProcessor`].
+            added_cond_kwargs: (`dict`, *optional*):
+                A kwargs dictionary containin additional embeddings that if specified are added to the embeddings that
+                are passed along to the UNet blocks.
+
+        Returns:
+            [`~models.unet_2d_condition.UNet2DConditionOutput`] or `tuple`:
+                If `return_dict` is True, an [`~models.unet_2d_condition.UNet2DConditionOutput`] is returned, otherwise
+                a `tuple` is returned where the first element is the sample tensor.
+        """
+        # By default samples have to be AT least a multiple of the overall upsampling factor.
+        # The overall upsampling factor is equal to 2 ** (# num of upsampling layers).
+        # However, the upsampling interpolation output size can be forced to fit any upsampling size
+        # on the fly if necessary.
+        default_overall_up_factor = 2**self.num_upsamplers
+
+        # upsample size should be forwarded when sample is not a multiple of `default_overall_up_factor`
+        forward_upsample_size = False
+        upsample_size = None
+
+        if any(s % default_overall_up_factor != 0 for s in sample.shape[-2:]):
+            logger.info("Forward upsample size to force interpolation output size.")
+            forward_upsample_size = True
+
+        # ensure attention_mask is a bias, and give it a singleton query_tokens dimension
+        # expects mask of shape:
+        #   [batch, key_tokens]
+        # adds singleton query_tokens dimension:
+        #   [batch,                    1, key_tokens]
+        # this helps to broadcast it as a bias over attention scores, which will be in one of the following shapes:
+        #   [batch,  heads, query_tokens, key_tokens] (e.g. torch sdp attn)
+        #   [batch * heads, query_tokens, key_tokens] (e.g. xformers or classic attn)
+        if attention_mask is not None:
+            # assume that mask is expressed as:
+            #   (1 = keep,      0 = discard)
+            # convert mask into a bias that can be added to attention scores:
+            #       (keep = +0,     discard = -10000.0)
+            attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0
+            attention_mask = attention_mask.unsqueeze(1)
+
+        # convert encoder_attention_mask to a bias the same way we do for attention_mask
+        if encoder_attention_mask is not None:
+            encoder_attention_mask = (1 - encoder_attention_mask.to(sample.dtype)) * -10000.0
+            encoder_attention_mask = encoder_attention_mask.unsqueeze(1)
+
+        # 0. center input if necessary
+        if self.config.center_input_sample:
+            sample = 2 * sample - 1.0
+
+        # 1. time
+        timesteps = timestep
+        if not torch.is_tensor(timesteps):
+            # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can
+            # This would be a good case for the `match` statement (Python 3.10+)
+            is_mps = sample.device.type == "mps"
+            if isinstance(timestep, float):
+                dtype = torch.float32 if is_mps else torch.float64
+            else:
+                dtype = torch.int32 if is_mps else torch.int64
+            timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device)
+        elif len(timesteps.shape) == 0:
+            timesteps = timesteps[None].to(sample.device)
+
+        # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+        timesteps = timesteps.expand(sample.shape[0])
+
+        t_emb = self.time_proj(timesteps)
+
+        # `Timesteps` does not contain any weights and will always return f32 tensors
+        # but time_embedding might actually be running in fp16. so we need to cast here.
+        # there might be better ways to encapsulate this.
+        t_emb = t_emb.to(dtype=sample.dtype)
+
+        emb = self.time_embedding(t_emb, timestep_cond)
+        aug_emb = None
+
+        if self.class_embedding is not None:
+            if class_labels is None:
+                raise ValueError("class_labels should be provided when num_class_embeds > 0")
+
+            if self.config.class_embed_type == "timestep":
+                class_labels = self.time_proj(class_labels)
+
+                # `Timesteps` does not contain any weights and will always return f32 tensors
+                # there might be better ways to encapsulate this.
+                class_labels = class_labels.to(dtype=sample.dtype)
+
+            class_emb = self.class_embedding(class_labels).to(dtype=sample.dtype)
+
+            if self.config.class_embeddings_concat:
+                emb = torch.cat([emb, class_emb], dim=-1)
+            else:
+                emb = emb + class_emb
+
+        if self.config.addition_embed_type == "text":
+            aug_emb = self.add_embedding(encoder_hidden_states)
+        elif self.config.addition_embed_type == "text_image":
+            # Kandinsky 2.1 - style
+            if "image_embeds" not in added_cond_kwargs:
+                raise ValueError(
+                    f"{self.__class__} has the config param `addition_embed_type` set to 'text_image' which requires the keyword argument `image_embeds` to be passed in `added_cond_kwargs`"
+                )
+
+            image_embs = added_cond_kwargs.get("image_embeds")
+            text_embs = added_cond_kwargs.get("text_embeds", encoder_hidden_states)
+            aug_emb = self.add_embedding(text_embs, image_embs)
+        elif self.config.addition_embed_type == "text_time":
+            # SDXL - style
+            if "text_embeds" not in added_cond_kwargs:
+                raise ValueError(
+                    f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' which requires the keyword argument `text_embeds` to be passed in `added_cond_kwargs`"
+                )
+            text_embeds = added_cond_kwargs.get("text_embeds")
+            if "time_ids" not in added_cond_kwargs:
+                raise ValueError(
+                    f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' which requires the keyword argument `time_ids` to be passed in `added_cond_kwargs`"
+                )
+            time_ids = added_cond_kwargs.get("time_ids")
+            time_embeds = self.add_time_proj(time_ids.flatten())
+            time_embeds = time_embeds.reshape((text_embeds.shape[0], -1))
+
+            add_embeds = torch.concat([text_embeds, time_embeds], dim=-1)
+            add_embeds = add_embeds.to(emb.dtype)
+            aug_emb = self.add_embedding(add_embeds)
+        elif self.config.addition_embed_type == "image":
+            # Kandinsky 2.2 - style
+            if "image_embeds" not in added_cond_kwargs:
+                raise ValueError(
+                    f"{self.__class__} has the config param `addition_embed_type` set to 'image' which requires the keyword argument `image_embeds` to be passed in `added_cond_kwargs`"
+                )
+            image_embs = added_cond_kwargs.get("image_embeds")
+            aug_emb = self.add_embedding(image_embs)
+        elif self.config.addition_embed_type == "image_hint":
+            # Kandinsky 2.2 - style
+            if "image_embeds" not in added_cond_kwargs or "hint" not in added_cond_kwargs:
+                raise ValueError(
+                    f"{self.__class__} has the config param `addition_embed_type` set to 'image_hint' which requires the keyword arguments `image_embeds` and `hint` to be passed in `added_cond_kwargs`"
+                )
+            image_embs = added_cond_kwargs.get("image_embeds")
+            hint = added_cond_kwargs.get("hint")
+            aug_emb, hint = self.add_embedding(image_embs, hint)
+            sample = torch.cat([sample, hint], dim=1)
+
+        emb = emb + aug_emb if aug_emb is not None else emb
+
+        if self.time_embed_act is not None:
+            emb = self.time_embed_act(emb)
+
+        if self.encoder_hid_proj is not None and self.config.encoder_hid_dim_type == "text_proj":
+            encoder_hidden_states = self.encoder_hid_proj(encoder_hidden_states)
+        elif self.encoder_hid_proj is not None and self.config.encoder_hid_dim_type == "text_image_proj":
+            # Kadinsky 2.1 - style
+            if "image_embeds" not in added_cond_kwargs:
+                raise ValueError(
+                    f"{self.__class__} has the config param `encoder_hid_dim_type` set to 'text_image_proj' which requires the keyword argument `image_embeds` to be passed in  `added_conditions`"
+                )
+
+            image_embeds = added_cond_kwargs.get("image_embeds")
+            encoder_hidden_states = self.encoder_hid_proj(encoder_hidden_states, image_embeds)
+        elif self.encoder_hid_proj is not None and self.config.encoder_hid_dim_type == "image_proj":
+            # Kandinsky 2.2 - style
+            if "image_embeds" not in added_cond_kwargs:
+                raise ValueError(
+                    f"{self.__class__} has the config param `encoder_hid_dim_type` set to 'image_proj' which requires the keyword argument `image_embeds` to be passed in  `added_conditions`"
+                )
+            image_embeds = added_cond_kwargs.get("image_embeds")
+            encoder_hidden_states = self.encoder_hid_proj(image_embeds)
+        # 2. pre-process
+        sample = self.conv_in(sample)
+
+        # 3. down
+
+        is_controlnet = mid_block_additional_residual is not None and down_block_additional_residuals is not None
+        is_adapter = mid_block_additional_residual is None and down_block_additional_residuals is not None
+
+        down_block_res_samples = (sample,)
+        for downsample_block in self.down_blocks:
+            if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention:
+                # For t2i-adapter CrossAttnDownBlock2D
+                additional_residuals = {}
+                if is_adapter and len(down_block_additional_residuals) > 0:
+                    additional_residuals["additional_residuals"] = down_block_additional_residuals.pop(0)
+
+                sample, res_samples = downsample_block(
+                    hidden_states=sample,
+                    temb=emb,
+                    encoder_hidden_states=encoder_hidden_states,
+                    attention_mask=attention_mask,
+                    cross_attention_kwargs=cross_attention_kwargs,
+                    encoder_attention_mask=encoder_attention_mask,
+                    **additional_residuals,
+                )
+            else:
+                sample, res_samples = downsample_block(hidden_states=sample, temb=emb)
+
+                if is_adapter and len(down_block_additional_residuals) > 0:
+                    sample += down_block_additional_residuals.pop(0)
+
+            down_block_res_samples += res_samples
+
+        if is_controlnet:
+            new_down_block_res_samples = ()
+
+            for down_block_res_sample, down_block_additional_residual in zip(
+                down_block_res_samples, down_block_additional_residuals
+            ):
+                down_block_res_sample = down_block_res_sample + down_block_additional_residual
+                new_down_block_res_samples = new_down_block_res_samples + (down_block_res_sample,)
+
+            down_block_res_samples = new_down_block_res_samples
+
+        # 4. mid
+        if self.mid_block is not None:
+            sample = self.mid_block(
+                sample,
+                emb,
+                encoder_hidden_states=encoder_hidden_states,
+                attention_mask=attention_mask,
+                cross_attention_kwargs=cross_attention_kwargs,
+                encoder_attention_mask=encoder_attention_mask,
+            )
+
+        if is_controlnet:
+            sample = sample + mid_block_additional_residual
+
+        # 5. up
+        for i, upsample_block in enumerate(self.up_blocks):
+            is_final_block = i == len(self.up_blocks) - 1
+
+            res_samples = down_block_res_samples[-len(upsample_block.resnets) :]
+            down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)]
+
+            # if we have not reached the final block and need to forward the
+            # upsample size, we do it here
+            if not is_final_block and forward_upsample_size:
+                upsample_size = down_block_res_samples[-1].shape[2:]
+
+            if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention:
+                sample = upsample_block(
+                    hidden_states=sample,
+                    temb=emb,
+                    res_hidden_states_tuple=res_samples,
+                    encoder_hidden_states=encoder_hidden_states,
+                    cross_attention_kwargs=cross_attention_kwargs,
+                    upsample_size=upsample_size,
+                    attention_mask=attention_mask,
+                    encoder_attention_mask=encoder_attention_mask,
+                    block_idx=i, # newly added
+                    additional_residuals=up_block_additional_residuals, # newly added
+                )
+            else:
+                sample = upsample_block(
+                    hidden_states=sample, temb=emb, res_hidden_states_tuple=res_samples, upsample_size=upsample_size,
+                    additional_residuals=up_block_additional_residuals # newly added
+                )
+
+        # 6. post-process
+        if self.conv_norm_out:
+            sample = self.conv_norm_out(sample)
+            sample = self.conv_act(sample)
+        sample = self.conv_out(sample)
+
+        if not return_dict:
+            return (sample,)
+
+        return UNet2DConditionOutput(sample=sample)
+
+
+class UNet(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+
+        self.model = ResidualUNet2DConditionModel.from_pretrained(
+            cfg.MODEL.UNET_CONFIG.PRETRAINED_PATH, use_safetensors = True)
+        self.model.requires_grad_(False)
+        self.model.enable_xformers_memory_efficient_attention()
+
+        self.model.enable_gradient_checkpointing()
+        for i, up_block in enumerate(self.model.up_blocks):
+            if isinstance(up_block, ResidualCrossAttnUpBlock2D):
+                for j, attn in enumerate(up_block.attentions):
+                    assert isinstance(attn, ResidualTransformer2DModel)
+                    block_idx = i * len(up_block.attentions) + j
+                    if block_idx not in cfg.MODEL.UNET_CONFIG.TRAINABLE_BLOCK_IDX:
+                        continue
+
+                    assert len(attn.transformer_blocks) == 1
+                    assert isinstance(attn.transformer_blocks[0], ResidualTransformerBlock)
+
+                    self_attn = attn.transformer_blocks[0].attn1
+                    assert isinstance(self_attn, ResidualAttention)
+                    if cfg.MODEL.UNET_CONFIG.TRAIN_SELF_ATTN_Q:
+                        self_attn.to_q.requires_grad_(True)
+                    if cfg.MODEL.UNET_CONFIG.TRAIN_SELF_ATTN_K:
+                        self_attn.to_k.requires_grad_(True)
+                    if cfg.MODEL.UNET_CONFIG.TRAIN_SELF_ATTN_V:
+                        self_attn.to_v.requires_grad_(True)
+
+                    cross_attn = attn.transformer_blocks[0].attn2
+                    assert isinstance(cross_attn, ResidualAttention)
+                    if cfg.MODEL.UNET_CONFIG.TRAIN_CROSS_ATTN_Q:
+                        cross_attn.to_q.requires_grad_(True)
+                    if cfg.MODEL.UNET_CONFIG.TRAIN_CROSS_ATTN_K:
+                        cross_attn.to_k.requires_grad_(True)
+                    if cfg.MODEL.UNET_CONFIG.TRAIN_CROSS_ATTN_V:
+                        cross_attn.to_v.requires_grad_(True)
+
+    def forward(self, sample, timestep, **kwargs):
+        return self.model(sample, timestep, **kwargs).sample

models/vae.py

@@ -0,0 +1,29 @@
+"""
+@author: Yanzuo Lu
+@author: oliveryanzuolu@gmail.com
+"""
+
+import diffusers
+import torch
+import torch.nn as nn
+
+
+class VariationalAutoencoder(nn.Module):
+    def __init__(self, pretrained_path):
+        super().__init__()
+        self.model = diffusers.AutoencoderKL.from_pretrained(pretrained_path, use_safetensors=True)
+        self.model.requires_grad_(False)
+        self.model.enable_slicing()
+
+    @torch.no_grad()
+    def encode(self, x):
+        z = self.model.encode(x).latent_dist
+        z = z.sample()
+        z = self.model.scaling_factor * z
+        return z
+
+    @torch.no_grad()
+    def decode(self, z):
+        z = 1. / self.model.scaling_factor * z
+        x = self.model.decode(z).sample
+        return x

models/xf.py

@@ -0,0 +1,155 @@
+"""
+@author: Yanzuo Lu
+@author: oliveryanzuolu@gmail.com
+"""
+
+import math
+
+import torch as th
+import torch.nn as nn
+from transformers import CLIPVisionModel
+
+
+class LayerNorm(nn.LayerNorm):
+    """
+    Implementation that supports fp16 inputs but fp32 gains/biases.
+    """
+
+    def forward(self, x: th.Tensor):
+        return super().forward(x.float()).to(x.dtype)
+
+
+class MultiheadAttention(nn.Module):
+    def __init__(self, n_ctx, width, heads):
+        super().__init__()
+        self.n_ctx = n_ctx
+        self.width = width
+        self.heads = heads
+        self.c_qkv = nn.Linear(width, width * 3)
+        self.c_proj = nn.Linear(width, width)
+        self.attention = QKVMultiheadAttention(heads, n_ctx)
+
+    def forward(self, x):
+        x = self.c_qkv(x)
+        x = self.attention(x)
+        x = self.c_proj(x)
+        return x
+
+
+class MLP(nn.Module):
+    def __init__(self, width):
+        super().__init__()
+        self.width = width
+        self.c_fc = nn.Linear(width, width * 4)
+        self.c_proj = nn.Linear(width * 4, width)
+        self.gelu = nn.GELU()
+
+    def forward(self, x):
+        return self.c_proj(self.gelu(self.c_fc(x)))
+
+
+class QKVMultiheadAttention(nn.Module):
+    def __init__(self, n_heads: int, n_ctx: int):
+        super().__init__()
+        self.n_heads = n_heads
+        self.n_ctx = n_ctx
+
+    def forward(self, qkv):
+        bs, n_ctx, width = qkv.shape
+        attn_ch = width // self.n_heads // 3
+        scale = 1 / math.sqrt(math.sqrt(attn_ch))
+        qkv = qkv.view(bs, n_ctx, self.n_heads, -1)
+        q, k, v = th.split(qkv, attn_ch, dim=-1)
+        weight = th.einsum(
+            "bthc,bshc->bhts", q * scale, k * scale
+        )  # More stable with f16 than dividing afterwards
+        wdtype = weight.dtype
+        weight = th.softmax(weight.float(), dim=-1).type(wdtype)
+        return th.einsum("bhts,bshc->bthc", weight, v).reshape(bs, n_ctx, -1)
+
+
+class ResidualAttentionBlock(nn.Module):
+    def __init__(
+        self,
+        n_ctx: int,
+        width: int,
+        heads: int,
+    ):
+        super().__init__()
+
+        self.attn = MultiheadAttention(
+            n_ctx,
+            width,
+            heads,
+        )
+        self.ln_1 = LayerNorm(width)
+        self.mlp = MLP(width)
+        self.ln_2 = LayerNorm(width)
+
+    def forward(self, x: th.Tensor):
+        x = x + self.attn(self.ln_1(x))
+        x = x + self.mlp(self.ln_2(x))
+        return x
+
+
+class Transformer(nn.Module):
+    def __init__(
+        self,
+        n_ctx: int,
+        width: int,
+        layers: int,
+        heads: int,
+    ):
+        super().__init__()
+        self.n_ctx = n_ctx
+        self.width = width
+        self.layers = layers
+        self.resblocks = nn.ModuleList(
+            [
+                ResidualAttentionBlock(
+                    n_ctx,
+                    width,
+                    heads,
+                )
+                for _ in range(layers)
+            ]
+        )
+
+    def forward(self, x: th.Tensor):
+        for block in self.resblocks:
+            x = block(x)
+        return x
+
+
+class FrozenCLIPImageEmbedder(nn.Module):
+    """Uses the CLIP transformer encoder for text (from Hugging Face)"""
+    def __init__(self, version="openai/clip-vit-large-patch14"):
+        super().__init__()
+        self.transformer = CLIPVisionModel.from_pretrained("pretrained_models/clip", use_safetensors=True)
+        self.final_ln = LayerNorm(768)
+        self.mapper = nn.Sequential(
+            nn.Linear(1024, 768, bias=False),
+            Transformer(1, 768, 5, 1)
+        )
+
+        self.freeze()
+
+    def freeze(self):
+        self.transformer = self.transformer.eval()
+        for param in self.parameters():
+            param.requires_grad = False
+        for param in self.mapper.parameters():
+            param.requires_grad = True
+        for param in self.final_ln.parameters():
+            param.requires_grad = True
+
+    def forward(self, image):
+        outputs = self.transformer(pixel_values=image)
+        z = outputs.pooler_output
+        z = z.unsqueeze(1)
+        z = self.mapper(z)
+        z = self.final_ln(z)
+        return z
+
+    def encode(self, image):
+        return self(image)

pose_transfer_test.py

@@ -0,0 +1,511 @@
+"""
+@author: Yanzuo Lu
+@author: oliveryanzuolu@gmail.com
+"""
+
+import argparse
+import copy
+import datetime
+import logging
+import os
+import sys
+import time
+import warnings
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from accelerate import Accelerator
+from accelerate.tracking import TensorBoardTracker, WandBTracker
+from accelerate.utils import set_seed
+from diffusers import (DDIMInverseScheduler, DDIMScheduler, DDPMScheduler,
+                       EulerDiscreteScheduler, PNDMScheduler)
+from einops import rearrange
+from PIL import Image
+from scipy.linalg import sqrtm
+from torch.utils.data import DataLoader
+from torchvision.utils import make_grid
+
+from datasets import FidRealDeepFashion, PisTestDeepFashion
+from defaults import pose_transfer_C as cfg
+from models import UNet, VariationalAutoencoder, build_metric
+from utils import AverageMeter
+
+warnings.filterwarnings("ignore")
+logger = logging.getLogger()
+
+
+def build_test_loader(cfg):
+    test_data = PisTestDeepFashion(
+        cfg.INPUT.ROOT_DIR, cfg.INPUT.GT.IMG_SIZE, cfg.INPUT.POSE.IMG_SIZE,
+        cfg.INPUT.COND.IMG_SIZE, cfg.TEST.IMG_SIZE)
+    test_loader = DataLoader(
+        test_data,
+        cfg.TEST.MICRO_BATCH_SIZE,
+        num_workers=cfg.TEST.NUM_WORKERS,
+        pin_memory=True
+    )
+
+    fid_real_data = FidRealDeepFashion(cfg.INPUT.ROOT_DIR, cfg.TEST.IMG_SIZE)
+    fid_real_loader = DataLoader(
+        fid_real_data,
+        cfg.TEST.MICRO_BATCH_SIZE,
+        num_workers=cfg.TEST.NUM_WORKERS,
+        pin_memory=True
+    )
+    return test_loader, fid_real_loader, test_data, fid_real_data
+
+
+def eval(cfg, model, test_loader, fid_real_loader, weight_dtype, save_dir,
+         test_data, fid_real_data, global_step, accelerator, metric,
+         noise_scheduler, inverse_noise_scheduler, vae, unet):
+    logger.info("start sampling...")
+    model.eval()
+    unet.eval()
+
+    gt_out_gathered = []
+    pred_out_gathered = []
+    lpips_gathered = []
+    psnr_gathered = []
+    ssim_gathered = []
+    ssim_256_gathered = []
+
+    with torch.no_grad():
+        end_time = time.time()
+        batch_time = AverageMeter()
+
+        for i, test_batch in enumerate(test_loader):
+            gt_imgs = test_batch["img_gt"]
+            img_size = test_batch["img_tgt"].shape[2:]
+            bsz = gt_imgs.shape[0]
+
+            if cfg.TEST.DDIM_INVERSION_STEPS > 0:
+                if cfg.TEST.DDIM_INVERSION_DOWN_BLOCK_GUIDANCE:
+                    c, down_block_additional_residuals, up_block_additional_residuals = model({
+                        "img_cond": test_batch["img_cond_from"], "pose_img": test_batch["pose_img_from"]})
+                else:
+                    c, down_block_additional_residuals, up_block_additional_residuals = model({
+                        "img_cond": test_batch["img_cond_from"], "pose_img": test_batch["pose_img_to"]})
+
+                noisy_latents = inverse_sample(
+                    cfg.TEST.DDIM_INVERSION_STEPS, accelerator, inverse_noise_scheduler, vae, unet,
+                    test_batch["img_src"], c[:bsz] if cfg.TEST.DDIM_INVERSION_UNCONDITIONAL else c[bsz:],
+                    [sample.to(dtype=weight_dtype) for sample in down_block_additional_residuals] if cfg.TEST.DDIM_INVERSION_DOWN_BLOCK_GUIDANCE else None,
+                    {k: v.to(dtype=weight_dtype) for k, v in up_block_additional_residuals.items()} if cfg.TEST.DDIM_INVERSION_UP_BLOCK_GUIDANCE else None)
+            else:
+                c, down_block_additional_residuals, up_block_additional_residuals = model({
+                    "img_cond": test_batch["img_cond_from"], "pose_img": test_batch["pose_img_to"]})
+                noisy_latents = torch.randn((bsz, 4, img_size[0]//8, img_size[1]//8)).to(accelerator.device)
+
+            if cfg.TEST.DDIM_INVERSION_STEPS > 0 and cfg.TEST.DDIM_INVERSION_DOWN_BLOCK_GUIDANCE:
+                c, down_block_additional_residuals, up_block_additional_residuals = model({
+                    "img_cond": test_batch["img_cond_from"], "pose_img": test_batch["pose_img_to"]})
+
+            sampling_imgs = sample(
+                cfg, weight_dtype, accelerator, noise_scheduler, vae, unet, noisy_latents,
+                c, down_block_additional_residuals, up_block_additional_residuals)
+
+            # log one-batch sampling results for visualization
+            if i == 0:
+                src_imgs = test_batch["img_src"] * 0.5 + 0.5
+                tgt_imgs = test_batch["img_tgt"] * 0.5 + 0.5
+                pose_imgs = F.interpolate(test_batch["pose_img_to"][:, :3, :, :],
+                                          tuple(test_batch["img_src"].shape[2:]),
+                                          mode="bicubic", antialias=True)
+                save_img = torch.stack([src_imgs, pose_imgs, tgt_imgs, sampling_imgs])
+                save_img = postprocess_image(save_img, nrow=save_img.shape[0]*2)
+                save_img.save(os.path.join(save_dir, f"inpainting_test_{accelerator.process_index}_{i}.jpg"))
+
+            sampling_imgs = F.interpolate(sampling_imgs, tuple(gt_imgs.shape[2:]), mode="bicubic", antialias=True)
+            sampling_imgs = sampling_imgs.float() * 255.0
+            sampling_imgs = sampling_imgs.clamp(0, 255).to(dtype=torch.uint8) # can save all images here!!!
+            sampling_imgs = sampling_imgs.to(torch.float32) / 255.
+
+            pred_out, lpips, psnr, ssim, ssim_256 = metric(gt_imgs, sampling_imgs)
+            pred_out_gathered.append(accelerator.gather_for_metrics(pred_out).cpu().numpy())
+            lpips_gathered.append(accelerator.gather_for_metrics(lpips).cpu().numpy())
+            psnr_gathered.append(accelerator.gather_for_metrics(psnr).cpu().numpy())
+            ssim_gathered.append(accelerator.gather_for_metrics(ssim).cpu().numpy())
+            ssim_256_gathered.append(accelerator.gather_for_metrics(ssim_256).cpu().numpy())
+
+            batch_time.update(time.time() - end_time)
+            end_time = time.time()
+
+            if (i + 1) % cfg.ACCELERATE.LOG_PERIOD == 0 or i == len(test_loader) - 1:
+                etas = batch_time.avg * (len(test_loader) - 1 - i)
+                logger.info(
+                    f"Sampling ({i+1}/{len(test_loader)})  "
+                    f"Time {batch_time.val:.4f}({batch_time.avg:.4f})  "
+                    f"Eta {datetime.timedelta(seconds=int(etas))}")
+                if os.environ.get("WANDB_MODE", None) == "offline":
+                    break
+
+        end_time = time.time()
+        batch_time = AverageMeter()
+        for i, fid_real_imgs in enumerate(fid_real_loader):
+            gt_out = metric(fid_real_imgs)
+            gt_out_gathered.append(accelerator.gather_for_metrics(gt_out).cpu().numpy())
+
+            batch_time.update(time.time() - end_time)
+            end_time = time.time()
+
+            if (i + 1) % cfg.ACCELERATE.LOG_PERIOD == 0 or i == len(fid_real_loader) - 1:
+                etas = batch_time.avg * (len(fid_real_loader) - 1 - i)
+                logger.info(
+                    f"FidReal ({i+1}/{len(fid_real_loader)})  "
+                    f"Time {batch_time.val:.4f}({batch_time.avg:.4f})  "
+                    f"Eta {datetime.timedelta(seconds=int(etas))}")
+
+        if accelerator.is_main_process:
+            gt_out_gathered = np.concatenate(gt_out_gathered, axis=0)
+            pred_out_gathered = np.concatenate(pred_out_gathered, axis=0)
+            lpips_gathered = np.concatenate(lpips_gathered, axis=0)
+            psnr_gathered = np.concatenate(psnr_gathered, axis=0)
+            ssim_gathered = np.concatenate(ssim_gathered, axis=0)
+            ssim_256_gathered = np.concatenate(ssim_256_gathered, axis=0)
+            if os.environ.get("WANDB_MODE", None) != "offline":
+                assert len(gt_out_gathered) == len(fid_real_data)
+                assert len(pred_out_gathered) == len(lpips_gathered) == len(psnr_gathered) == \
+                    len(ssim_gathered) == len(ssim_256_gathered) == len(test_data)
+
+            mu1 = np.mean(gt_out_gathered, axis=0)
+            sigma1 = np.cov(gt_out_gathered, rowvar=False)
+            mu2 = np.mean(pred_out_gathered, axis=0)
+            sigma2 = np.cov(pred_out_gathered, rowvar=False)
+
+            mu1 = np.atleast_1d(mu1)
+            mu2 = np.atleast_1d(mu2)
+            sigma1 = np.atleast_2d(sigma1)
+            sigma2 = np.atleast_2d(sigma2)
+
+            diff = mu1 - mu2
+
+            # Product might be almost singular
+            covmean, _ = sqrtm(sigma1.dot(sigma2), disp=False)
+            if not np.isfinite(covmean).all():
+                msg = ('fid calculation produces singular product; '
+                    'adding %s to diagonal of cov estimates') % 1e-6
+                logger.info(msg)
+                offset = np.eye(sigma1.shape[0]) * 1e-6
+                covmean = sqrtm((sigma1 + offset).dot(sigma2 + offset))
+
+            # Numerical error might give slight imaginary component
+            if np.iscomplexobj(covmean):
+                if not np.allclose(np.diagonal(covmean).imag, 0, atol=1e-3):
+                    m = np.max(np.abs(covmean.imag))
+                    raise ValueError('Imaginary component {}'.format(m))
+                covmean = covmean.real
+
+            tr_covmean = np.trace(covmean)
+
+            score_fid = diff.dot(diff) + np.trace(sigma1) + np.trace(sigma2) - 2 * tr_covmean
+            score_lpips = np.mean(lpips_gathered)
+            score_ssim = np.mean(ssim_gathered)
+            score_ssim_256 = np.mean(ssim_256_gathered)
+            score_psnr = np.mean(psnr_gathered)
+
+            logger.info("Evaluation Results:")
+            logger.info(f"FID: {score_fid:.3f}")
+            logger.info(f"LPIPS: {score_lpips:.4f}")
+            logger.info(f"SSIM: {score_ssim:.4f}")
+            logger.info(f"SSIM_256: {score_ssim_256:.4f}")
+            logger.info(f"PSNR: {score_psnr:.3f}")
+
+            accelerator.log({
+                "score_fid": score_fid,
+                "score_lpips": score_lpips,
+                "score_ssim": score_ssim,
+                "score_ssim_256": score_ssim_256,
+                "score_psnr": score_psnr
+            }, step=global_step)
+
+        accelerator.wait_for_everyone()
+        torch.cuda.empty_cache()
+
+
+def sample(cfg, weight_dtype, accelerator, noise_scheduler, vae, unet, noisy_latents,
+           c_new, down_block_additional_residuals, up_block_additional_residuals):
+    bsz = noisy_latents.shape[0]
+    noise_scheduler.set_timesteps(cfg.TEST.NUM_INFERENCE_STEPS)
+
+    if cfg.TEST.GUIDANCE_TYPE == "uc_full":
+        down_block_additional_residuals = [torch.cat([torch.zeros_like(sample), sample]).to(dtype=weight_dtype) \
+                                           for sample in down_block_additional_residuals]
+        up_block_additional_residuals = {k: torch.cat([torch.zeros_like(v), v]).to(dtype=weight_dtype) \
+                                         for k, v in up_block_additional_residuals.items()}
+
+        for t in noise_scheduler.timesteps:
+            inputs = torch.cat([noisy_latents, noisy_latents], dim=0)
+            inputs = noise_scheduler.scale_model_input(inputs, timestep=t)
+            with accelerator.autocast():
+                noise_pred = unet(sample=inputs, timestep=t, encoder_hidden_states=c_new,
+                    down_block_additional_residuals=copy.deepcopy(down_block_additional_residuals),
+                    up_block_additional_residuals=copy.deepcopy(up_block_additional_residuals))
+
+            noise_pred_uc, noise_pred_full = noise_pred.chunk(2)
+            noise_pred = noise_pred_uc + cfg.TEST.FULL_GUIDANCE_SCALE * (noise_pred_full - noise_pred_uc)
+            noisy_latents = noise_scheduler.step(noise_pred, t, noisy_latents)[0]
+
+    elif cfg.TEST.GUIDANCE_TYPE == "updown_full":
+        down_block_additional_residuals = [torch.cat([sample, sample]).to(dtype=weight_dtype) \
+                                           for sample in down_block_additional_residuals]
+        up_block_additional_residuals = {k: torch.cat([v, v]).to(dtype=weight_dtype) \
+                                         for k, v in up_block_additional_residuals.items()}
+
+        for t in noise_scheduler.timesteps:
+            inputs = torch.cat([noisy_latents, noisy_latents], dim=0)
+            inputs = noise_scheduler.scale_model_input(inputs, timestep=t)
+            with accelerator.autocast():
+                noise_pred = unet(sample=inputs, timestep=t, encoder_hidden_states=c_new,
+                    down_block_additional_residuals=copy.deepcopy(down_block_additional_residuals),
+                    up_block_additional_residuals=copy.deepcopy(up_block_additional_residuals))
+
+            noise_pred_updown, noise_pred_full = noise_pred.chunk(2)
+            noise_pred = noise_pred_updown + cfg.TEST.FULL_GUIDANCE_SCALE * (noise_pred_full - noise_pred_updown)
+            noisy_latents = noise_scheduler.step(noise_pred, t, noisy_latents)[0]
+
+    elif cfg.TEST.GUIDANCE_TYPE == "down_full":
+        down_block_additional_residuals = [torch.cat([sample, sample]).to(dtype=weight_dtype) \
+                                           for sample in down_block_additional_residuals]
+        up_block_additional_residuals = {k: torch.cat([torch.zeros_like(v), v]).to(dtype=weight_dtype) \
+                                         for k, v in up_block_additional_residuals.items()}
+
+        for t in noise_scheduler.timesteps:
+            inputs = torch.cat([noisy_latents, noisy_latents], dim=0)
+            inputs = noise_scheduler.scale_model_input(inputs, timestep=t)
+            with accelerator.autocast():
+                noise_pred = unet(sample=inputs, timestep=t, encoder_hidden_states=c_new,
+                    down_block_additional_residuals=copy.deepcopy(down_block_additional_residuals),
+                    up_block_additional_residuals=copy.deepcopy(up_block_additional_residuals))
+
+            noise_pred_down, noise_pred_full = noise_pred.chunk(2)
+            noise_pred = noise_pred_down + cfg.TEST.FULL_GUIDANCE_SCALE * (noise_pred_full - noise_pred_down)
+            noisy_latents = noise_scheduler.step(noise_pred, t, noisy_latents)[0]
+
+    elif cfg.TEST.GUIDANCE_TYPE == "uc_down_full":
+        c_new = torch.cat([c_new[:bsz], c_new[:bsz], c_new[bsz:]])
+        down_block_additional_residuals = [torch.cat([torch.zeros_like(sample), sample, sample]).to(dtype=weight_dtype) \
+                                           for sample in down_block_additional_residuals]
+        up_block_additional_residuals = {k: torch.cat([torch.zeros_like(v), torch.zeros_like(v), v]).to(dtype=weight_dtype) \
+                                         for k, v in up_block_additional_residuals.items()}
+
+        for t in noise_scheduler.timesteps:
+            inputs = torch.cat([noisy_latents, noisy_latents, noisy_latents], dim=0)
+            inputs = noise_scheduler.scale_model_input(inputs, timestep=t)
+            with accelerator.autocast():
+                noise_pred = unet(sample=inputs, timestep=t, encoder_hidden_states=c_new,
+                    down_block_additional_residuals=copy.deepcopy(down_block_additional_residuals),
+                    up_block_additional_residuals=copy.deepcopy(up_block_additional_residuals))
+
+            noise_pred_uc, noise_pred_down, noise_pred_full = noise_pred.chunk(3)
+            noise_pred = noise_pred_uc + \
+                         cfg.TEST.DOWN_BLOCK_GUIDANCE_SCALE * (noise_pred_down - noise_pred_uc) + \
+                         cfg.TEST.FULL_GUIDANCE_SCALE * (noise_pred_full - noise_pred_down)
+            noisy_latents = noise_scheduler.step(noise_pred, t, noisy_latents)[0]
+
+    elif cfg.TEST.GUIDANCE_TYPE == "uc_down_updown_cdown":
+        c_new = torch.cat([c_new[:bsz], c_new[:bsz], c_new[:bsz], c_new[bsz:]])
+        down_block_additional_residuals = [torch.cat([torch.zeros_like(sample), sample, sample, sample]).to(dtype=weight_dtype) \
+                                           for sample in down_block_additional_residuals]
+        up_block_additional_residuals = {k: torch.cat([torch.zeros_like(v), torch.zeros_like(v), v, torch.zeros_like(v)]).to(dtype=weight_dtype) \
+                                         for k, v in up_block_additional_residuals.items()}
+
+        for t in noise_scheduler.timesteps:
+            inputs = torch.cat([noisy_latents, noisy_latents, noisy_latents, noisy_latents], dim=0)
+            inputs = noise_scheduler.scale_model_input(inputs, timestep=t)
+            with accelerator.autocast():
+                noise_pred = unet(sample=inputs, timestep=t, encoder_hidden_states=c_new,
+                    down_block_additional_residuals=copy.deepcopy(down_block_additional_residuals),
+                    up_block_additional_residuals=copy.deepcopy(up_block_additional_residuals))
+
+            noise_pred_uc, noise_pred_down, noise_pred_updown, noise_pred_cdown = noise_pred.chunk(4)
+            noise_pred = noise_pred_uc + \
+                         cfg.TEST.DOWN_BLOCK_GUIDANCE_SCALE * (noise_pred_down - noise_pred_uc) + \
+                         cfg.TEST.ALL_BLOCK_GUIDANCE_SCALE * (noise_pred_updown - noise_pred_down) + \
+                         cfg.TEST.GUIDANCE_SCALE * (noise_pred_cdown - noise_pred_down)
+            noisy_latents = noise_scheduler.step(noise_pred, t, noisy_latents)[0]
+
+    elif cfg.TEST.GUIDANCE_TYPE == "uc_down_updown_full":
+        c_new = torch.cat([c_new[:bsz], c_new[:bsz], c_new[:bsz], c_new[bsz:]])
+        down_block_additional_residuals = [torch.cat([torch.zeros_like(sample), sample, sample, sample]).to(dtype=weight_dtype) \
+                                           for sample in down_block_additional_residuals]
+        up_block_additional_residuals = {k: torch.cat([torch.zeros_like(v), torch.zeros_like(v), v, v]).to(dtype=weight_dtype) \
+                                         for k, v in up_block_additional_residuals.items()}
+
+        for t in noise_scheduler.timesteps:
+            inputs = torch.cat([noisy_latents, noisy_latents, noisy_latents, noisy_latents], dim=0)
+            inputs = noise_scheduler.scale_model_input(inputs, timestep=t)
+            with accelerator.autocast():
+                noise_pred = unet(sample=inputs, timestep=t, encoder_hidden_states=c_new,
+                    down_block_additional_residuals=copy.deepcopy(down_block_additional_residuals),
+                    up_block_additional_residuals=copy.deepcopy(up_block_additional_residuals))
+
+            noise_pred_uc, noise_pred_down, noise_pred_updown, noise_pred_full = noise_pred.chunk(4)
+            noise_pred = noise_pred_uc + \
+                         cfg.TEST.DOWN_BLOCK_GUIDANCE_SCALE * (noise_pred_down - noise_pred_uc) + \
+                         cfg.TEST.ALL_BLOCK_GUIDANCE_SCALE * (noise_pred_updown - noise_pred_down) + \
+                         cfg.TEST.FULL_GUIDANCE_SCALE * (noise_pred_full - noise_pred_updown)
+            noisy_latents = noise_scheduler.step(noise_pred, t, noisy_latents)[0]
+
+    with accelerator.autocast():
+        sampling_imgs = vae.decode(noisy_latents) * 0.5 + 0.5 # denormalize
+        sampling_imgs = sampling_imgs.clamp(0, 1)
+    return sampling_imgs
+
+
+def inverse_sample(num_inference_steps, accelerator, inverse_noise_scheduler, vae, unet, img_src,
+                   c_new, down_block_additional_residuals=None, up_block_additional_residuals=None):
+    inverse_noise_scheduler.set_timesteps(num_inference_steps)
+    with accelerator.autocast():
+        noisy_latents = vae.encode(img_src)
+
+    for t in inverse_noise_scheduler.timesteps:
+        inputs = noisy_latents
+        with accelerator.autocast():
+            noise_pred = unet(sample=inputs, timestep=t, encoder_hidden_states=c_new,
+                down_block_additional_residuals=copy.deepcopy(down_block_additional_residuals) if down_block_additional_residuals else None,
+                up_block_additional_residuals=copy.deepcopy(up_block_additional_residuals) if up_block_additional_residuals else None)
+        noisy_latents = inverse_noise_scheduler.step(noise_pred, t, noisy_latents)[0]
+
+    return noisy_latents
+
+
+def postprocess_image(tensor, nrow):
+    tensor = tensor * 255.
+    tensor = torch.clamp(tensor, min=0., max=255.)
+    tensor = rearrange(tensor, 'n b c h w -> b n c h w')
+    tensor = rearrange(tensor, 'b n c h w -> (b n) c h w')
+    tensor = make_grid(tensor, nrow=nrow)
+    img = tensor.cpu().numpy().transpose(1, 2, 0).astype(np.uint8)
+    return Image.fromarray(img)
+
+
+def main(cfg):
+    project_dir = os.path.join("outputs", cfg.ACCELERATE.PROJECT_NAME)
+    run_dir = os.path.join(project_dir, cfg.ACCELERATE.RUN_NAME)
+    os.makedirs(run_dir, exist_ok=True)
+
+    accelerator = Accelerator(
+        log_with = ["wandb", "tensorboard"],
+        project_dir = project_dir,
+        mixed_precision = cfg.ACCELERATE.MIXED_PRECISION
+    )
+    torch.backends.cuda.matmul.allow_tf32 = cfg.ACCELERATE.ALLOW_TF32
+    set_seed(cfg.ACCELERATE.SEED)
+
+    if accelerator.is_main_process:
+        accelerator.trackers = []
+        accelerator.trackers.append(WandBTracker(
+            cfg.ACCELERATE.PROJECT_NAME, name=cfg.ACCELERATE.RUN_NAME, config=cfg, dir=project_dir))
+        accelerator.trackers.append(TensorBoardTracker(cfg.ACCELERATE.RUN_NAME, project_dir))
+
+        with open(os.path.join(run_dir, "config.yaml"), "w") as f:
+            f.write(cfg.dump())
+    accelerator.wait_for_everyone()
+
+    fmt = "[%(asctime)s %(filename)s:%(lineno)s] %(message)s"
+    datefmt = "%Y-%m-%d %H:%M:%S"
+    logging.basicConfig(
+        level = logging.INFO,
+        format = fmt,
+        datefmt = datefmt,
+        filename = f"{run_dir}/log_rank{accelerator.process_index}.txt",
+        filemode = "a"
+    )
+    if accelerator.is_main_process:
+        console_handler = logging.StreamHandler(sys.stdout)
+        console_handler.setLevel(logging.INFO)
+        console_handler.setFormatter(logging.Formatter(fmt, datefmt))
+        logger.addHandler(console_handler)
+
+    logger.info(f"running with config:\n{str(cfg)}")
+
+    logger.info("preparing datasets...")
+    test_loader, fid_real_loader, test_data, fid_real_data = build_test_loader(cfg)
+
+    logger.info("preparing model...")
+    weight_dtype = torch.float32
+    if accelerator.mixed_precision == "fp16":
+        weight_dtype = torch.float16
+    elif accelerator.mixed_precision == "bf16":
+        weight_dtype = torch.bfloat16
+
+    # not trained, move to 16-bit to save memory
+    vae = VariationalAutoencoder(
+        pretrained_path=cfg.MODEL.FIRST_STAGE_CONFIG.PRETRAINED_PATH
+    ).to(accelerator.device, dtype=weight_dtype)
+
+    if cfg.MODEL.SCHEDULER_CONFIG.NAME == "euler":
+        noise_scheduler = EulerDiscreteScheduler.from_pretrained(cfg.MODEL.SCHEDULER_CONFIG.PRETRAINED_PATH)
+    elif cfg.MODEL.SCHEDULER_CONFIG.NAME == "pndm":
+        noise_scheduler = PNDMScheduler.from_pretrained(cfg.MODEL.SCHEDULER_CONFIG.PRETRAINED_PATH)
+    elif cfg.MODEL.SCHEDULER_CONFIG.NAME == "ddim":
+        noise_scheduler = DDIMScheduler.from_pretrained(cfg.MODEL.SCHEDULER_CONFIG.PRETRAINED_PATH)
+    elif cfg.MODEL.SCHEDULER_CONFIG.NAME == "ddpm":
+        noise_scheduler = DDPMScheduler.from_pretrained(cfg.MODEL.SCHEDULER_CONFIG.PRETRAINED_PATH)
+
+    inverse_noise_scheduler = DDIMInverseScheduler(
+        num_train_timesteps=noise_scheduler.num_train_timesteps,
+        beta_start=noise_scheduler.beta_start,
+        beta_end=noise_scheduler.beta_end,
+        beta_schedule=noise_scheduler.beta_schedule,
+        trained_betas=noise_scheduler.trained_betas,
+        clip_sample=noise_scheduler.clip_sample,
+        set_alpha_to_one=noise_scheduler.set_alpha_to_one,
+        steps_offset=noise_scheduler.steps_offset,
+        prediction_type=noise_scheduler.prediction_type,
+        timestep_spacing=noise_scheduler.timestep_spacing
+    )
+
+    from pose_transfer_train import build_model
+    model = build_model(cfg)
+    unet = UNet(cfg)
+    metric = build_metric().to(accelerator.device)
+
+    logger.info(model.load_state_dict(torch.load(
+        os.path.join(cfg.MODEL.PRETRAINED_PATH, "pytorch_model.bin"), map_location="cpu"
+    ), strict=False))
+    logger.info(unet.load_state_dict(torch.load(
+        os.path.join(cfg.MODEL.PRETRAINED_PATH, "pytorch_model_1.bin"), map_location="cpu"
+    ), strict=False))
+
+    logger.info("preparing accelerator...")
+    model, unet, test_loader, fid_real_loader = accelerator.prepare(model, unet, test_loader, fid_real_loader)
+
+    save_dir = os.path.join(run_dir, "log_images")
+    os.makedirs(save_dir, exist_ok=True)
+
+    eval(
+        cfg=cfg,
+        model=model,
+        test_loader=test_loader,
+        fid_real_loader=fid_real_loader,
+        weight_dtype=weight_dtype,
+        save_dir=save_dir,
+        test_data=test_data,
+        fid_real_data=fid_real_data,
+        global_step=None,
+        accelerator=accelerator,
+        metric=metric,
+        noise_scheduler=noise_scheduler,
+        inverse_noise_scheduler=inverse_noise_scheduler,
+        vae=vae,
+        unet=unet
+    )
+
+    accelerator.end_training()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Pose Transfer Testing")
+    parser.add_argument("--config_file", type=str, default="", help="path to config file")
+    parser.add_argument("opts", default=None, nargs=argparse.REMAINDER, help=
+                        "modify config options using the command-line")
+    args = parser.parse_args()
+
+    if args.config_file:
+        cfg.merge_from_file(args.config_file)
+    cfg.merge_from_list(args.opts)
+    cfg.freeze()
+
+    main(cfg)

pose_transfer_train.py

@@ -0,0 +1,385 @@
+"""
+@author: Yanzuo Lu
+@author: oliveryanzuolu@gmail.com
+"""
+
+import argparse
+import datetime
+import logging
+import os
+import sys
+import time
+import warnings
+
+import torch
+import torch.nn as nn
+from accelerate import Accelerator
+from accelerate.tracking import TensorBoardTracker, WandBTracker
+from accelerate.utils import DistributedDataParallelKwargs, set_seed
+from diffusers import (DDIMInverseScheduler, DDIMScheduler, DDPMScheduler,
+                       EulerDiscreteScheduler, PNDMScheduler)
+from torch.utils.data import DataLoader
+
+from datasets import PisTrainDeepFashion
+from defaults import pose_transfer_C as cfg
+from lr_scheduler import LinearWarmupMultiStepDecayLRScheduler
+from models import (AppearanceEncoder, Decoder, PoseEncoder, UNet,
+                    VariationalAutoencoder, build_backbone, build_metric)
+from pose_transfer_test import build_test_loader, eval
+from utils import AverageMeter
+
+warnings.filterwarnings("ignore")
+logger = logging.getLogger()
+
+
+class build_model(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.pose_query = cfg.MODEL.DECODER_CONFIG.POSE_QUERY
+
+        self.backbone = build_backbone(
+            img_size=cfg.INPUT.COND.IMG_SIZE,
+            embed_dim=cfg.MODEL.COND_STAGE_CONFIG.EMBED_DIM,
+            depths=cfg.MODEL.COND_STAGE_CONFIG.DEPTHS,
+            num_heads=cfg.MODEL.COND_STAGE_CONFIG.NUM_HEADS,
+            window_size=cfg.MODEL.COND_STAGE_CONFIG.WINDOW_SIZE,
+            drop_path_rate=cfg.MODEL.COND_STAGE_CONFIG.DROP_PATH_RATE,
+            mask=len(cfg.INPUT.COND.PRED_RATIO) > 0,
+            last_norm=cfg.MODEL.COND_STAGE_CONFIG.LAST_NORM,
+            pretrained_path=cfg.MODEL.COND_STAGE_CONFIG.PRETRAINED_PATH
+        )
+
+        self.appearance_encoder = AppearanceEncoder(
+            attn_residual_block_idx=cfg.MODEL.APPEARANCE_GUIDANCE_CONFIG.ATTN_RESIDUAL_BLOCK_IDX,
+            inner_dims=cfg.MODEL.APPEARANCE_GUIDANCE_CONFIG.INNER_DIMS,
+            ctx_dims=cfg.MODEL.APPEARANCE_GUIDANCE_CONFIG.CTX_DIMS,
+            embed_dims=cfg.MODEL.APPEARANCE_GUIDANCE_CONFIG.EMBED_DIMS,
+            heads=cfg.MODEL.APPEARANCE_GUIDANCE_CONFIG.HEADS,
+            depth=cfg.MODEL.APPEARANCE_GUIDANCE_CONFIG.DEPTH,
+            to_self_attn=cfg.MODEL.APPEARANCE_GUIDANCE_CONFIG.TO_SELF_ATTN,
+            to_queries=cfg.MODEL.APPEARANCE_GUIDANCE_CONFIG.TO_QUERIES,
+            to_keys=cfg.MODEL.APPEARANCE_GUIDANCE_CONFIG.TO_KEYS,
+            to_values=cfg.MODEL.APPEARANCE_GUIDANCE_CONFIG.TO_VALUES,
+            aspect_ratio=cfg.INPUT.COND.IMG_SIZE[0] // cfg.INPUT.COND.IMG_SIZE[1],
+            detach_input=cfg.MODEL.APPEARANCE_GUIDANCE_CONFIG.DETACH_INPUT,
+            convin_kernel_size=cfg.MODEL.APPEARANCE_GUIDANCE_CONFIG.CONVIN_KERNEL_SIZE,
+            convin_stride=cfg.MODEL.APPEARANCE_GUIDANCE_CONFIG.CONVIN_STRIDE,
+            convin_padding=cfg.MODEL.APPEARANCE_GUIDANCE_CONFIG.CONVIN_PADDING
+        )
+
+        self.pose_encoder = PoseEncoder(
+            downscale_factor=cfg.MODEL.POSE_GUIDANCE_CONFIG.DOWNSCALE_FACTOR,
+            pose_channels=cfg.MODEL.POSE_GUIDANCE_CONFIG.POSE_CHANNELS,
+            in_channels=cfg.MODEL.POSE_GUIDANCE_CONFIG.IN_CHANNELS,
+            channels=cfg.MODEL.POSE_GUIDANCE_CONFIG.CHANNELS
+        )
+
+        self.decoder = Decoder(
+            n_ctx=cfg.MODEL.DECODER_CONFIG.N_CTX,
+            ctx_dim=cfg.MODEL.DECODER_CONFIG.CTX_DIM,
+            heads=cfg.MODEL.DECODER_CONFIG.HEADS,
+            depth=cfg.MODEL.DECODER_CONFIG.DEPTH,
+            last_norm=cfg.MODEL.COND_STAGE_CONFIG.LAST_NORM,
+            img_size=cfg.INPUT.COND.IMG_SIZE,
+            embed_dim=cfg.MODEL.COND_STAGE_CONFIG.EMBED_DIM,
+            depths=cfg.MODEL.COND_STAGE_CONFIG.DEPTHS,
+            pose_query=cfg.MODEL.DECODER_CONFIG.POSE_QUERY,
+            pose_channel=cfg.MODEL.POSE_GUIDANCE_CONFIG.CHANNELS[-1]
+        )
+
+        self.learnable_vector = nn.Parameter(torch.randn((1, cfg.MODEL.DECODER_CONFIG.N_CTX, cfg.MODEL.DECODER_CONFIG.CTX_DIM)))
+        self.u_cond_percent = cfg.MODEL.U_COND_PERCENT
+        self.u_cond_down_block_guidance = cfg.MODEL.U_COND_DOWN_BLOCK_GUIDANCE
+        self.u_cond_up_block_guidance = cfg.MODEL.U_COND_UP_BLOCK_GUIDANCE
+
+    def forward(self, batched_inputs):
+        mask = batched_inputs["mask"] if "mask" in batched_inputs else None
+        x, features = self.backbone(batched_inputs["img_cond"], mask=mask)
+        up_block_additional_residuals = self.appearance_encoder(features)
+
+        bsz = x.shape[0]
+        if self.training:
+            bsz = bsz * 2
+            down_block_additional_residuals = self.pose_encoder(torch.cat([batched_inputs["pose_img_src"], batched_inputs["pose_img_tgt"]]))
+            up_block_additional_residuals = {k: torch.cat([v, v]) for k, v in up_block_additional_residuals.items()}
+            c = self.decoder(x, features, down_block_additional_residuals)
+            if not self.pose_query:
+                c = torch.cat([c, c])
+
+            u_cond_prop = torch.rand(bsz, 1, 1)
+            u_cond_prop = (u_cond_prop < self.u_cond_percent).to(dtype=x.dtype, device=x.device)
+            c = self.learnable_vector.expand(bsz, -1, -1).to(dtype=x.dtype) * u_cond_prop + c * (1 - u_cond_prop)
+            if self.u_cond_down_block_guidance:
+                down_block_additional_residuals = [torch.zeros_like(sample) * u_cond_prop.unsqueeze(1) + \
+                                                   sample * (1 - u_cond_prop.unsqueeze(1)) \
+                                                   for sample in down_block_additional_residuals]
+            if self.u_cond_up_block_guidance:
+                up_block_additional_residuals = {k: torch.zeros_like(v) * u_cond_prop + v * (1 - u_cond_prop) \
+                                                 for k, v in up_block_additional_residuals.items()}
+        else:
+            down_block_additional_residuals = self.pose_encoder(batched_inputs["pose_img"])
+            c = self.decoder(x, features, down_block_additional_residuals)
+            c = torch.cat([self.learnable_vector.expand(bsz, -1, -1).to(dtype=x.dtype), c], dim=0)
+
+        return c, down_block_additional_residuals, up_block_additional_residuals
+
+
+def main(cfg):
+    project_dir = os.path.join("outputs", cfg.ACCELERATE.PROJECT_NAME)
+    run_dir = os.path.join(project_dir, cfg.ACCELERATE.RUN_NAME)
+    os.makedirs(run_dir, exist_ok=True)
+
+    accelerator = Accelerator(
+        log_with=["wandb", "tensorboard"],
+        project_dir=project_dir,
+        mixed_precision=cfg.ACCELERATE.MIXED_PRECISION,
+        gradient_accumulation_steps=cfg.ACCELERATE.GRADIENT_ACCUMULATION_STEPS,
+        kwargs_handlers=[DistributedDataParallelKwargs(bucket_cap_mb=200, gradient_as_bucket_view=True)]
+    )
+    torch.backends.cuda.matmul.allow_tf32 = cfg.ACCELERATE.ALLOW_TF32
+    set_seed(cfg.ACCELERATE.SEED)
+
+    if accelerator.is_main_process:
+        accelerator.trackers = []
+        accelerator.trackers.append(WandBTracker(
+            cfg.ACCELERATE.PROJECT_NAME, name=cfg.ACCELERATE.RUN_NAME, config=cfg, dir=project_dir))
+        accelerator.trackers.append(TensorBoardTracker(cfg.ACCELERATE.RUN_NAME, project_dir))
+
+        with open(os.path.join(run_dir, "config.yaml"), "w") as f:
+            f.write(cfg.dump())
+    accelerator.wait_for_everyone()
+
+    fmt = "[%(asctime)s %(filename)s:%(lineno)s] %(message)s"
+    datefmt = "%Y-%m-%d %H:%M:%S"
+    logging.basicConfig(
+        level=logging.INFO,
+        format=fmt,
+        datefmt=datefmt,
+        filename=f"{run_dir}/log_rank{accelerator.process_index}.txt",
+        filemode="a"
+    )
+    if accelerator.is_main_process:
+        console_handler = logging.StreamHandler(sys.stdout)
+        console_handler.setLevel(logging.INFO)
+        console_handler.setFormatter(logging.Formatter(fmt, datefmt))
+        logger.addHandler(console_handler)
+
+    logger.info(f"running with config:\n{str(cfg)}")
+
+    logger.info("preparing datasets...")
+    train_data = PisTrainDeepFashion(
+        root_dir=cfg.INPUT.ROOT_DIR,
+        gt_img_size=cfg.INPUT.GT.IMG_SIZE,
+        pose_img_size=cfg.INPUT.POSE.IMG_SIZE,
+        cond_img_size=cfg.INPUT.COND.IMG_SIZE,
+        min_scale=cfg.INPUT.COND.MIN_SCALE,
+        log_aspect_ratio=cfg.INPUT.COND.PRED_ASPECT_RATIO,
+        pred_ratio=cfg.INPUT.COND.PRED_RATIO,
+        pred_ratio_var=cfg.INPUT.COND.PRED_RATIO_VAR,
+        psz=cfg.INPUT.COND.MASK_PATCH_SIZE
+    )
+    train_loader = DataLoader(
+        train_data,
+        cfg.INPUT.BATCH_SIZE // accelerator.num_processes // cfg.ACCELERATE.GRADIENT_ACCUMULATION_STEPS,
+        shuffle = True,
+        drop_last = True,
+        num_workers = cfg.INPUT.NUM_WORKERS,
+        pin_memory = True
+    )
+
+    test_loader, fid_real_loader, test_data, fid_real_data = build_test_loader(cfg)
+
+    logger.info("preparing model...")
+    weight_dtype = torch.float32
+    if accelerator.mixed_precision == "fp16":
+        weight_dtype = torch.float16
+    elif accelerator.mixed_precision == "bf16":
+        weight_dtype = torch.bfloat16
+
+    # not trained, move to 16-bit to save memory
+    vae = VariationalAutoencoder(
+        pretrained_path=cfg.MODEL.FIRST_STAGE_CONFIG.PRETRAINED_PATH
+    ).to(accelerator.device, dtype=weight_dtype)
+
+    if cfg.MODEL.SCHEDULER_CONFIG.NAME == "euler":
+        noise_scheduler = EulerDiscreteScheduler.from_pretrained(cfg.MODEL.SCHEDULER_CONFIG.PRETRAINED_PATH)
+    elif cfg.MODEL.SCHEDULER_CONFIG.NAME == "pndm":
+        noise_scheduler = PNDMScheduler.from_pretrained(cfg.MODEL.SCHEDULER_CONFIG.PRETRAINED_PATH)
+    elif cfg.MODEL.SCHEDULER_CONFIG.NAME == "ddim":
+        noise_scheduler = DDIMScheduler.from_pretrained(cfg.MODEL.SCHEDULER_CONFIG.PRETRAINED_PATH)
+    elif cfg.MODEL.SCHEDULER_CONFIG.NAME == "ddpm":
+        noise_scheduler = DDPMScheduler.from_pretrained(cfg.MODEL.SCHEDULER_CONFIG.PRETRAINED_PATH)
+
+    inverse_noise_scheduler = DDIMInverseScheduler(
+        num_train_timesteps=noise_scheduler.num_train_timesteps,
+        beta_start=noise_scheduler.beta_start,
+        beta_end=noise_scheduler.beta_end,
+        beta_schedule=noise_scheduler.beta_schedule,
+        trained_betas=noise_scheduler.trained_betas,
+        clip_sample=noise_scheduler.clip_sample,
+        set_alpha_to_one=noise_scheduler.set_alpha_to_one,
+        steps_offset=noise_scheduler.steps_offset,
+        prediction_type=noise_scheduler.prediction_type,
+        timestep_spacing=noise_scheduler.timestep_spacing
+    )
+
+    model = build_model(cfg)
+    unet = UNet(cfg)
+    metric = build_metric().to(accelerator.device)
+    trainable_params = sum([p.numel() for p in model.parameters() if p.requires_grad] + \
+                           [p.numel() for p in unet.parameters() if p.requires_grad])
+    logger.info(f"number of trainable parameters: {trainable_params}")
+
+    logger.info("preparing optimizer...")
+    lr = cfg.OPTIMIZER.LR * cfg.INPUT.BATCH_SIZE if cfg.OPTIMIZER.SCALE_LR else cfg.OPTIMIZER.LR
+    params = [p for p in model.parameters() if p.requires_grad] + \
+             [p for p in unet.parameters() if p.requires_grad]
+    optimizer = torch.optim.Adam(params, lr=lr)
+
+    logger.info("preparing accelerator...")
+    model, unet, optimizer, train_loader, test_loader, fid_real_loader = accelerator.prepare(
+        model, unet, optimizer, train_loader, test_loader, fid_real_loader)
+
+    last_epoch = cfg.MODEL.LAST_EPOCH
+    if cfg.MODEL.PRETRAINED_PATH:
+        logger.info(f"loading states from {cfg.MODEL.PRETRAINED_PATH}")
+        accelerator.load_state(cfg.MODEL.PRETRAINED_PATH)
+    global_step = last_epoch * len(train_loader)
+
+    logger.info("preparing lr scheduler...")
+    lr_scheduler = LinearWarmupMultiStepDecayLRScheduler(
+        optimizer, cfg.OPTIMIZER.WARMUP_STEPS, cfg.OPTIMIZER.WARMUP_RATE, cfg.OPTIMIZER.DECAY_RATE,
+        cfg.OPTIMIZER.EPOCHS, cfg.OPTIMIZER.DECAY_EPOCHS, len(train_loader),
+        last_epoch=len(train_loader)*last_epoch-1, override_lr=cfg.OPTIMIZER.OVERRIDE_LR)
+
+    logger.info("start training...")
+    start_time = time.time()
+    end_time = time.time()
+
+    for epoch in range(last_epoch, cfg.OPTIMIZER.EPOCHS, 1):
+        model.train()
+        unet.train()
+
+        epoch_time = time.time()
+        logger.info(f"epoch {epoch + 1} start")
+        batch_time = AverageMeter()
+        total_loss = AverageMeter()
+
+        for i, batch in enumerate(train_loader):
+            with accelerator.accumulate(model, unet):
+                optimizer.zero_grad()
+
+                # Convert images to latent space
+                with accelerator.autocast():
+                    latents = vae.encode(torch.cat([batch["img_src"], batch["img_tgt"]]))
+
+                # Sample noise that we'll add to the latents
+                noise = torch.randn_like(latents)
+                bsz = latents.shape[0]
+
+                if cfg.MODEL.SCHEDULER_CONFIG.CUBIC_SAMPLING:
+                    # Cubic sampling to sample a random timestep for each image
+                    timesteps = torch.rand((bsz, ), device=accelerator.device)
+                    timesteps = (1 - timesteps**3) * noise_scheduler.config.num_train_timesteps
+                    timesteps = timesteps.long()
+                    timesteps = torch.clamp(timesteps, 0, noise_scheduler.config.num_train_timesteps - 1)
+                else:
+                    # Uniform sampling to sample a random timestep for each image
+                    timesteps = torch.randint(noise_scheduler.config.num_train_timesteps, (bsz, ), device=accelerator.device)
+
+                # Add noise to the latents according to the noise magnitude at each timestep
+                # (this is the forward diffusion process)
+                noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)
+
+                # get embedding
+                c, down_block_additional_residuals, up_block_additional_residuals = model(batch)
+                down_block_additional_residuals = [sample.to(dtype=weight_dtype) for sample in down_block_additional_residuals]
+                up_block_additional_residuals = {k: v.to(dtype=weight_dtype) for k, v in up_block_additional_residuals.items()}
+
+                # predict
+                with accelerator.autocast():
+                    encoder_hidden_states = c.to(dtype=weight_dtype)
+                    model_pred = unet(
+                        sample=noisy_latents, timestep=timesteps, encoder_hidden_states=encoder_hidden_states,
+                        down_block_additional_residuals=down_block_additional_residuals,
+                        up_block_additional_residuals=up_block_additional_residuals)
+
+                loss_simple = (noise - model_pred) ** 2
+                loss_simple = loss_simple.mean()
+                loss = loss_simple / cfg.ACCELERATE.GRADIENT_ACCUMULATION_STEPS
+                if torch.isnan(loss).any():
+                    accelerator.set_trigger()
+                if accelerator.check_trigger():
+                    logger.info("loss is nan, stop training")
+                    accelerator.end_training()
+                    time.sleep(86400) # waiting for...
+
+                accelerator.backward(loss)
+                if accelerator.sync_gradients:
+                    global_step += 1
+                optimizer.step()
+                lr_scheduler.step()
+
+            total_loss.update(loss_simple.item())
+            batch_time.update(time.time() - end_time)
+            end_time = time.time()
+
+            if (i + 1) % cfg.ACCELERATE.LOG_PERIOD == 0 or i == len(train_loader) - 1:
+                accelerator.log({
+                    "loss": loss_simple.item(),
+                    "loss_avg": total_loss.avg,
+                    "lr": optimizer.param_groups[-1]["lr"]
+                }, step=global_step)
+
+                etas = batch_time.avg * (len(train_loader) - 1 - i)
+                logger.info(
+                    f"Train [{epoch+1}/{cfg.OPTIMIZER.EPOCHS}]({i+1}/{len(train_loader)})  "
+                    f"Time {batch_time.val:.4f}({batch_time.avg:.4f})  "
+                    f"Loss {total_loss.val:.4f}({total_loss.avg:.4f})  "
+                    f"Lr {optimizer.param_groups[-1]['lr']:.8f}  "
+                    f"Eta {datetime.timedelta(seconds=int(etas))}")
+
+        logger.info(f"epoch {epoch + 1} finished, running time {datetime.timedelta(seconds=int(time.time() - epoch_time))}")
+        save_dir = os.path.join(run_dir, f"epochs_{(epoch+1):03d}")
+
+        if (epoch + 1) % cfg.ACCELERATE.EVAL_PERIOD == 0:
+            accelerator.save_state(os.path.join(save_dir, "checkpoints"))
+            save_dir = os.path.join(save_dir, "log_images")
+            os.makedirs(save_dir, exist_ok=True)
+
+            eval(
+                cfg=cfg,
+                model=model,
+                test_loader=test_loader,
+                fid_real_loader=fid_real_loader,
+                weight_dtype=weight_dtype,
+                save_dir=save_dir,
+                test_data=test_data,
+                fid_real_data=fid_real_data,
+                global_step=None,
+                accelerator=accelerator,
+                metric=metric,
+                noise_scheduler=noise_scheduler,
+                inverse_noise_scheduler=inverse_noise_scheduler,
+                vae=vae,
+                unet=unet
+            )
+
+    train_time = time.time() - start_time
+    logger.info(f'training completed, running time {datetime.timedelta(seconds=int(train_time))}')
+    accelerator.end_training()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Pose Transfer Training")
+    parser.add_argument("--config_file", type=str, default="", help="path to config file")
+    parser.add_argument("opts", default=None, nargs=argparse.REMAINDER, help=
+                        "modify config options using the command-line")
+    args = parser.parse_args()
+
+    if args.config_file:
+        cfg.merge_from_file(args.config_file)
+    cfg.merge_from_list(args.opts)
+    cfg.freeze()
+
+    main(cfg)

pose_utils.py

@@ -0,0 +1,72 @@
+"""
+@author: Yanzuo Lu
+@author: oliveryanzuolu@gmail.com
+"""
+
+import json
+import logging
+import cv2
+
+import numpy as np
+
+logger = logging.getLogger()
+
+BONES = [[1,2], [1,5], [2,3], [3,4], [5,6], [6,7], [1,8], [8,9],
+         [9,10], [1,11], [11,12], [12,13], [1,0], [0,14], [14,16],
+         [0,15], [15,17]]
+
+JOINT_COLORS = [[255, 0, 0], [255, 85, 0], [255, 170, 0], [255, 255, 0], [170, 255, 0], [85, 255, 0], [0, 255, 0],
+                [0, 255, 85], [0, 255, 170], [0, 255, 255], [0, 170, 255], [0, 85, 255], [0, 0, 255], [85, 0, 255],
+                [170, 0, 255], [255, 0, 255], [255, 0, 170], [255, 0, 85]]
+
+BONE_COLORS = [[153, 0, 0], [153, 51, 0], [153, 102, 0], [153, 153, 0], [102, 153, 0], [51, 153, 0], [0, 153, 0], [0, 153, 51],
+               [0, 153, 102], [0, 153, 153], [0, 102, 153], [0, 51, 153], [0, 0, 153], [51, 0, 153], [102, 0, 153],
+               [153, 0, 153], [153, 0, 102]]
+
+def load_pose_cords_from_strings(y_str, x_str):
+    y_cords = json.loads(y_str)
+    x_cords = json.loads(x_str)
+    return np.concatenate([np.expand_dims(y_cords, -1), np.expand_dims(x_cords, -1)], axis=1)
+
+
+def cords_to_map(cords, img_size, old_size=(128, 64), affine_matrix=None, sigma=6):
+    old_size = img_size if old_size is None else old_size
+    cords = cords.astype(float)
+    result = np.zeros(img_size + cords.shape[0:1], dtype='float32')
+    for i, point in enumerate(cords):
+        if point[0] == -1 or point[1] == -1:
+            continue
+        point[0] = point[0]/old_size[0] * img_size[0]
+        point[1] = point[1]/old_size[1] * img_size[1]
+        if affine_matrix is not None:
+            point_ =np.dot(affine_matrix, np.matrix([point[1], point[0], 1]).reshape(3,1))
+            point_0 = int(point_[1])
+            point_1 = int(point_[0])
+        else:
+            point_0 = int(point[0])
+            point_1 = int(point[1])
+        xx, yy = np.meshgrid(np.arange(img_size[1]), np.arange(img_size[0]))
+        result[..., i] = np.exp(-((yy - point_0) ** 2 + (xx - point_1) ** 2) / (2 * sigma ** 2))
+    return result
+
+
+def draw_pose_from_cords(array, img_size, old_size=(128, 64), radius=2, draw_bones=True):
+    colors = np.zeros(shape=img_size + (3, ), dtype=np.uint8)
+    scale_y = img_size[0] / old_size[0]
+    scale_x = img_size[1] / old_size[1]
+
+    if draw_bones:
+        for i, (f, t) in enumerate(BONES):
+            from_missing = array[f][0] == -1 or array[f][1] == -1
+            to_missing = array[t][0] == -1 or array[t][1] == -1
+            if from_missing or to_missing:
+                continue
+            cv2.line(colors, (int(array[f][1] * scale_x), int(array[f][0] * scale_y)),
+                     (int(array[t][1] * scale_x), int(array[t][0] * scale_y)), BONE_COLORS[i], radius, cv2.LINE_AA)
+
+    for i, joint in enumerate(array):
+        if array[i][0] == -1 or array[i][1] == -1:
+            continue
+        cv2.circle(colors, (int(joint[1] * scale_x), int(joint[0] * scale_y)), radius + 1, JOINT_COLORS[i], -1, cv2.LINE_AA)
+
+    return colors

requirements.txt

@@ -0,0 +1,29 @@
+torch==2.0.1 --index-url https://download.pytorch.org/whl/cu118
+torchvision==0.15.2 --index-url https://download.pytorch.org/whl/cu118
+torchaudio==2.0.2 --index-url https://download.pytorch.org/whl/cu118
+xformers==0.0.21
+diffusers==0.19.3
+transformers==4.35.0
+accelerate==0.23.0
+einops==0.3.0
+opencv-python==4.7.0.72
+timm==0.9.7
+safetensors==0.3.1
+scipy==1.10.1
+lpips==0.1.4
+tensorboard==2.13.0
+wandb==0.15.11
+numpy==1.23.1
+yacs==0.1.6
+pandas==2.0.3
+scikit-image==0.20.0
+huggingface-hub==0.17.3
+jax==0.4.13
+jaxlib==0.4.13
+flax==0.7.0
+fastapi==0.104.1
+uvicorn==0.24.0
+python-multipart==0.0.6
+gradio==4.7.1
+pillow==10.0.1
+matplotlib==3.7.2z

scripts/multi_gpu/pose_transfer_test.sh

@@ -0,0 +1,24 @@
+#!/bin/bash
+
+source $(dirname "${CONDA_PYTHON_EXE}")/activate CFLD
+export CUDA_VISIBLE_DEVICES=$1
+export NUM_GPUS=$(echo $CUDA_VISIBLE_DEVICES | awk -F "," '{print NF}')
+shift
+
+while true # find unused tcp port
+do
+    PORT=$(( ((RANDOM<<15)|RANDOM) % 49152 + 10000 ))
+    status="$(nc -z 127.0.0.1 $PORT < /dev/null &>/dev/null; echo $?)"
+    if [ "${status}" != "0" ]; then
+        break;
+    fi
+done
+
+accelerate launch \
+    --multi_gpu \
+    --num_processes $NUM_GPUS \
+    --num_machines 1 \
+    --dynamo_backend "no" \
+    --main_process_port $PORT \
+    pose_transfer_test.py $@ \
+    INPUT.ROOT_DIR ./fashion

scripts/multi_gpu/pose_transfer_train.sh

@@ -0,0 +1,24 @@
+#!/bin/bash
+
+source $(dirname "${CONDA_PYTHON_EXE}")/activate CFLD
+export CUDA_VISIBLE_DEVICES=$1
+export NUM_GPUS=$(echo $CUDA_VISIBLE_DEVICES | awk -F "," '{print NF}')
+shift
+
+while true # find unused tcp port
+do
+    PORT=$(( ((RANDOM<<15)|RANDOM) % 49152 + 10000 ))
+    status="$(nc -z 127.0.0.1 $PORT < /dev/null &>/dev/null; echo $?)"
+    if [ "${status}" != "0" ]; then
+        break;
+    fi
+done
+
+accelerate launch \
+    --multi_gpu \
+    --num_processes $NUM_GPUS \
+    --num_machines 1 \
+    --dynamo_backend "no" \
+    --main_process_port $PORT \
+    pose_transfer_train.py $@ \
+    INPUT.ROOT_DIR ./fashion

scripts/single_gpu/pose_transfer_test.sh

@@ -0,0 +1,8 @@
+#!/bin/bash
+
+source $(dirname "${CONDA_PYTHON_EXE}")/activate CFLD
+export CUDA_VISIBLE_DEVICES=$1
+shift
+
+python pose_transfer_test.py $@ \
+    INPUT.ROOT_DIR ./fashion

scripts/single_gpu/pose_transfer_train.sh

@@ -0,0 +1,8 @@
+#!/bin/bash
+
+source $(dirname "${CONDA_PYTHON_EXE}")/activate CFLD
+export CUDA_VISIBLE_DEVICES=$1
+shift
+
+python pose_transfer_train.py $@ \
+    INPUT.ROOT_DIR ./fashion

utils.py

@@ -0,0 +1,22 @@
+"""
+@author: Yanzuo Lu
+@author: oliveryanzuolu@gmail.com
+"""
+
+import logging
+
+logger = logging.getLogger()
+
+
+class AverageMeter:
+    def __init__(self):
+        self.val = 0
+        self.avg = 0
+        self.sum = 0
+        self.count = 0
+
+    def update(self, val, n=1):
+        self.val = val
+        self.sum += val * n
+        self.count += n
+        self.avg = self.sum / self.count