init

src/audio2vid.py

@@ -0,0 +1,255 @@
+import os
+import ffmpeg
+from datetime import datetime
+from pathlib import Path
+import numpy as np
+import cv2
+import torch
+# import spaces
+from scipy.spatial.transform import Rotation as R
+from scipy.interpolate import interp1d
+
+from diffusers import AutoencoderKL, DDIMScheduler
+from einops import repeat
+from omegaconf import OmegaConf
+from PIL import Image
+from torchvision import transforms
+from transformers import CLIPVisionModelWithProjection
+
+
+from src.models.pose_guider import PoseGuider
+from src.models.unet_2d_condition import UNet2DConditionModel
+from src.models.unet_3d import UNet3DConditionModel
+from src.pipelines.pipeline_pose2vid_long import Pose2VideoPipeline
+from src.utils.util import save_videos_grid
+
+from src.audio_models.model import Audio2MeshModel
+from src.utils.audio_util import prepare_audio_feature
+from src.utils.mp_utils  import LMKExtractor
+from src.utils.draw_util import FaceMeshVisualizer
+from src.utils.pose_util import project_points
+
+
+def matrix_to_euler_and_translation(matrix):
+    rotation_matrix = matrix[:3, :3]
+    translation_vector = matrix[:3, 3]
+    rotation = R.from_matrix(rotation_matrix)
+    euler_angles = rotation.as_euler('xyz', degrees=True)
+    return euler_angles, translation_vector
+
+
+def smooth_pose_seq(pose_seq, window_size=5):
+    smoothed_pose_seq = np.zeros_like(pose_seq)
+
+    for i in range(len(pose_seq)):
+        start = max(0, i - window_size // 2)
+        end = min(len(pose_seq), i + window_size // 2 + 1)
+        smoothed_pose_seq[i] = np.mean(pose_seq[start:end], axis=0)
+
+    return smoothed_pose_seq
+
+def get_headpose_temp(input_video):
+    lmk_extractor = LMKExtractor()
+    cap = cv2.VideoCapture(input_video)
+
+    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    fps = cap.get(cv2.CAP_PROP_FPS)  
+
+    trans_mat_list = []
+    while cap.isOpened():
+        ret, frame = cap.read()
+        if not ret:
+            break
+
+        result = lmk_extractor(frame)
+        trans_mat_list.append(result['trans_mat'].astype(np.float32))
+    cap.release()
+
+    trans_mat_arr = np.array(trans_mat_list)
+
+    # compute delta pose
+    trans_mat_inv_frame_0 = np.linalg.inv(trans_mat_arr[0])
+    pose_arr = np.zeros([trans_mat_arr.shape[0], 6])
+
+    for i in range(pose_arr.shape[0]):
+        pose_mat = trans_mat_inv_frame_0 @ trans_mat_arr[i]
+        euler_angles, translation_vector = matrix_to_euler_and_translation(pose_mat)
+        pose_arr[i, :3] =  euler_angles
+        pose_arr[i, 3:6] =  translation_vector
+
+    # interpolate to 30 fps
+    new_fps = 30
+    old_time = np.linspace(0, total_frames / fps, total_frames)
+    new_time = np.linspace(0, total_frames / fps, int(total_frames * new_fps / fps))
+
+    pose_arr_interp = np.zeros((len(new_time), 6))
+    for i in range(6):
+        interp_func = interp1d(old_time, pose_arr[:, i])
+        pose_arr_interp[:, i] = interp_func(new_time)
+
+    pose_arr_smooth = smooth_pose_seq(pose_arr_interp)
+    
+    return pose_arr_smooth
+
+# @spaces.GPU
+def audio2video(input_audio, ref_img, headpose_video=None, size=512, steps=25, length=150, seed=42):
+    fps = 30
+    cfg = 3.5
+
+    config = OmegaConf.load('./configs/prompts/animation_audio.yaml')
+
+    if config.weight_dtype == "fp16":
+        weight_dtype = torch.float16
+    else:
+        weight_dtype = torch.float32
+        
+    audio_infer_config = OmegaConf.load(config.audio_inference_config)
+    # prepare model
+    a2m_model = Audio2MeshModel(audio_infer_config['a2m_model'])
+    a2m_model.load_state_dict(torch.load(audio_infer_config['pretrained_model']['a2m_ckpt']), strict=False)
+    a2m_model.cuda().eval()
+
+    vae = AutoencoderKL.from_pretrained(
+        config.pretrained_vae_path,
+    ).to("cuda", dtype=weight_dtype)
+
+    reference_unet = UNet2DConditionModel.from_pretrained(
+        config.pretrained_base_model_path,
+        subfolder="unet",
+    ).to(dtype=weight_dtype, device="cuda")
+
+    inference_config_path = config.inference_config
+    infer_config = OmegaConf.load(inference_config_path)
+    denoising_unet = UNet3DConditionModel.from_pretrained_2d(
+        config.pretrained_base_model_path,
+        config.motion_module_path,
+        subfolder="unet",
+        unet_additional_kwargs=infer_config.unet_additional_kwargs,
+    ).to(dtype=weight_dtype, device="cuda")
+
+
+    pose_guider = PoseGuider(noise_latent_channels=320, use_ca=True).to(device="cuda", dtype=weight_dtype) # not use cross attention
+
+    image_enc = CLIPVisionModelWithProjection.from_pretrained(
+        config.image_encoder_path
+    ).to(dtype=weight_dtype, device="cuda")
+
+    sched_kwargs = OmegaConf.to_container(infer_config.noise_scheduler_kwargs)
+    scheduler = DDIMScheduler(**sched_kwargs)
+
+    generator = torch.manual_seed(seed)
+
+    width, height = size, size
+
+    # load pretrained weights
+    denoising_unet.load_state_dict(
+        torch.load(config.denoising_unet_path, map_location="cpu"),
+        strict=False,
+    )
+    reference_unet.load_state_dict(
+        torch.load(config.reference_unet_path, map_location="cpu"),
+    )
+    pose_guider.load_state_dict(
+        torch.load(config.pose_guider_path, map_location="cpu"),
+    )
+
+    pipe = Pose2VideoPipeline(
+        vae=vae,
+        image_encoder=image_enc,
+        reference_unet=reference_unet,
+        denoising_unet=denoising_unet,
+        pose_guider=pose_guider,
+        scheduler=scheduler,
+    )
+    pipe = pipe.to("cuda", dtype=weight_dtype)
+
+    date_str = datetime.now().strftime("%Y%m%d")
+    time_str = datetime.now().strftime("%H%M")
+    save_dir_name = f"{time_str}--seed_{seed}-{size}x{size}"
+
+    save_dir = Path(f"output/{date_str}/{save_dir_name}")
+    save_dir.mkdir(exist_ok=True, parents=True)
+
+    lmk_extractor = LMKExtractor()
+    vis = FaceMeshVisualizer(forehead_edge=False)
+
+    ref_image_np = cv2.cvtColor(ref_img, cv2.COLOR_RGB2BGR)
+    # TODO: 人脸检测+裁剪
+    ref_image_np = cv2.resize(ref_image_np, (size, size))
+    ref_image_pil = Image.fromarray(cv2.cvtColor(ref_image_np, cv2.COLOR_BGR2RGB))
+    
+    face_result = lmk_extractor(ref_image_np)
+    if face_result is None: 
+        return None
+    
+    lmks = face_result['lmks'].astype(np.float32)
+    ref_pose = vis.draw_landmarks((ref_image_np.shape[1], ref_image_np.shape[0]), lmks, normed=True)
+    
+    sample = prepare_audio_feature(input_audio, wav2vec_model_path=audio_infer_config['a2m_model']['model_path'])
+    sample['audio_feature'] = torch.from_numpy(sample['audio_feature']).float().cuda()
+    sample['audio_feature'] = sample['audio_feature'].unsqueeze(0)
+
+    # inference
+    pred = a2m_model.infer(sample['audio_feature'], sample['seq_len'])
+    pred = pred.squeeze().detach().cpu().numpy()
+    pred = pred.reshape(pred.shape[0], -1, 3)
+    pred = pred + face_result['lmks3d']
+    
+    if headpose_video is not None:
+        pose_seq = get_headpose_temp(headpose_video)
+    else:
+        pose_seq = np.load(config['pose_temp'])
+    mirrored_pose_seq = np.concatenate((pose_seq, pose_seq[-2:0:-1]), axis=0)
+    cycled_pose_seq = np.tile(mirrored_pose_seq, (sample['seq_len'] // len(mirrored_pose_seq) + 1, 1))[:sample['seq_len']]
+
+    # project 3D mesh to 2D landmark
+    projected_vertices = project_points(pred, face_result['trans_mat'], cycled_pose_seq, [height, width])
+
+    pose_images = []
+    for i, verts in enumerate(projected_vertices):
+        lmk_img = vis.draw_landmarks((width, height), verts, normed=False)
+        pose_images.append(lmk_img)
+
+    pose_list = []
+    pose_tensor_list = []
+
+    pose_transform = transforms.Compose(
+        [transforms.Resize((height, width)), transforms.ToTensor()]
+    )
+    args_L = len(pose_images) if length==0 or length > len(pose_images) else length
+    for pose_image_np in pose_images[: args_L]:
+        pose_image_pil = Image.fromarray(cv2.cvtColor(pose_image_np, cv2.COLOR_BGR2RGB))
+        pose_tensor_list.append(pose_transform(pose_image_pil))
+        pose_image_np = cv2.resize(pose_image_np,  (width, height))
+        pose_list.append(pose_image_np)
+    
+    pose_list = np.array(pose_list)
+    
+    video_length = len(pose_tensor_list)
+
+    video = pipe(
+        ref_image_pil,
+        pose_list,
+        ref_pose,
+        width,
+        height,
+        video_length,
+        steps,
+        cfg,
+        generator=generator,
+    ).videos
+
+    save_path = f"{save_dir}/{size}x{size}_{time_str}_noaudio.mp4"
+    save_videos_grid(
+        video,
+        save_path,
+        n_rows=1,
+        fps=fps,
+    )
+    
+    stream = ffmpeg.input(save_path)
+    audio = ffmpeg.input(input_audio)
+    ffmpeg.output(stream.video, audio.audio, save_path.replace('_noaudio.mp4', '.mp4'), vcodec='copy', acodec='aac').run()
+    os.remove(save_path)
+    
+    return save_path.replace('_noaudio.mp4', '.mp4')

src/audio_models/mish.py

@@ -0,0 +1,51 @@
+"""
+Applies the mish function element-wise:
+mish(x) = x * tanh(softplus(x)) = x * tanh(ln(1 + exp(x)))
+"""
+
+# import pytorch
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+@torch.jit.script
+def mish(input):
+    """
+    Applies the mish function element-wise:
+    mish(x) = x * tanh(softplus(x)) = x * tanh(ln(1 + exp(x)))
+    See additional documentation for mish class.
+    """
+    return input * torch.tanh(F.softplus(input))
+
+class Mish(nn.Module):
+    """
+    Applies the mish function element-wise:
+    mish(x) = x * tanh(softplus(x)) = x * tanh(ln(1 + exp(x)))
+
+    Shape:
+        - Input: (N, *) where * means, any number of additional
+          dimensions
+        - Output: (N, *), same shape as the input
+
+    Examples:
+        >>> m = Mish()
+        >>> input = torch.randn(2)
+        >>> output = m(input)
+
+    Reference: https://pytorch.org/docs/stable/generated/torch.nn.Mish.html
+    """
+
+    def __init__(self):
+        """
+        Init method.
+        """
+        super().__init__()
+
+    def forward(self, input):
+        """
+        Forward pass of the function.
+        """
+        if torch.__version__ >= "1.9":
+            return F.mish(input)
+        else:
+            return mish(input)

src/audio_models/model.py

@@ -0,0 +1,71 @@
+import os
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import Wav2Vec2Config
+
+from .torch_utils import get_mask_from_lengths
+from .wav2vec2 import Wav2Vec2Model
+
+
+class Audio2MeshModel(nn.Module):
+    def __init__(
+        self,
+        config
+    ):  
+        super().__init__()
+        out_dim = config['out_dim']
+        latent_dim = config['latent_dim']
+        model_path = config['model_path']
+        only_last_fetures = config['only_last_fetures']
+        from_pretrained = config['from_pretrained']
+
+        self._only_last_features = only_last_fetures
+
+        self.audio_encoder_config = Wav2Vec2Config.from_pretrained(model_path, local_files_only=True)
+        if from_pretrained:
+            self.audio_encoder = Wav2Vec2Model.from_pretrained(model_path, local_files_only=True)
+        else:
+            self.audio_encoder = Wav2Vec2Model(self.audio_encoder_config)
+        self.audio_encoder.feature_extractor._freeze_parameters()
+
+        hidden_size = self.audio_encoder_config.hidden_size
+
+        self.in_fn = nn.Linear(hidden_size, latent_dim)
+       
+        self.out_fn = nn.Linear(latent_dim, out_dim)
+        nn.init.constant_(self.out_fn.weight, 0)
+        nn.init.constant_(self.out_fn.bias, 0)
+
+    def forward(self, audio, label, audio_len=None):
+        attention_mask = ~get_mask_from_lengths(audio_len) if audio_len else None
+        
+        seq_len = label.shape[1]
+
+        embeddings = self.audio_encoder(audio, seq_len=seq_len, output_hidden_states=True,
+                                        attention_mask=attention_mask)
+
+        if self._only_last_features:
+            hidden_states = embeddings.last_hidden_state
+        else:
+            hidden_states = sum(embeddings.hidden_states) / len(embeddings.hidden_states)
+
+        layer_in = self.in_fn(hidden_states)
+        out = self.out_fn(layer_in)
+
+        return out, None
+
+    def infer(self, input_value, seq_len):
+        embeddings = self.audio_encoder(input_value, seq_len=seq_len, output_hidden_states=True)
+
+        if self._only_last_features:
+            hidden_states = embeddings.last_hidden_state
+        else:
+            hidden_states = sum(embeddings.hidden_states) / len(embeddings.hidden_states)
+
+        layer_in = self.in_fn(hidden_states)
+        out = self.out_fn(layer_in)
+
+        return out
+    
+

src/audio_models/torch_utils.py

@@ -0,0 +1,25 @@
+import torch
+import torch.nn.functional as F
+
+
+def get_mask_from_lengths(lengths, max_len=None):
+    lengths = lengths.to(torch.long)
+    if max_len is None:
+        max_len = torch.max(lengths).item()
+
+    ids = torch.arange(0, max_len).unsqueeze(0).expand(lengths.shape[0], -1).to(lengths.device)
+    mask = ids < lengths.unsqueeze(1).expand(-1, max_len)
+
+    return mask
+
+
+def linear_interpolation(features, seq_len):
+    features = features.transpose(1, 2)
+    output_features = F.interpolate(features, size=seq_len, align_corners=True, mode='linear')
+    return output_features.transpose(1, 2)
+
+
+if __name__ == "__main__":
+    import numpy as np
+    mask = ~get_mask_from_lengths(torch.from_numpy(np.array([4,6])))
+    import pdb; pdb.set_trace()

src/audio_models/wav2vec2.py

@@ -0,0 +1,125 @@
+from transformers import Wav2Vec2Config, Wav2Vec2Model
+from transformers.modeling_outputs import BaseModelOutput
+
+from .torch_utils import linear_interpolation
+
+# the implementation of Wav2Vec2Model is borrowed from
+# https://github.com/huggingface/transformers/blob/HEAD/src/transformers/models/wav2vec2/modeling_wav2vec2.py
+# initialize our encoder with the pre-trained wav2vec 2.0 weights.
+class Wav2Vec2Model(Wav2Vec2Model):
+    def __init__(self, config: Wav2Vec2Config):
+        super().__init__(config)
+
+    def forward(
+        self,
+        input_values,
+        seq_len,
+        attention_mask=None,
+        mask_time_indices=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+        self.config.output_attentions = True
+
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        extract_features = self.feature_extractor(input_values)
+        extract_features = extract_features.transpose(1, 2)
+        extract_features = linear_interpolation(extract_features, seq_len=seq_len)
+
+        if attention_mask is not None:
+            # compute reduced attention_mask corresponding to feature vectors
+            attention_mask = self._get_feature_vector_attention_mask(
+                extract_features.shape[1], attention_mask, add_adapter=False
+            )
+
+        hidden_states, extract_features = self.feature_projection(extract_features)
+        hidden_states = self._mask_hidden_states(
+            hidden_states, mask_time_indices=mask_time_indices, attention_mask=attention_mask
+        )
+
+        encoder_outputs = self.encoder(
+            hidden_states,
+            attention_mask=attention_mask,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        hidden_states = encoder_outputs[0]
+
+        if self.adapter is not None:
+            hidden_states = self.adapter(hidden_states)
+
+        if not return_dict:
+            return (hidden_states, ) + encoder_outputs[1:]
+        return BaseModelOutput(
+            last_hidden_state=hidden_states,
+            hidden_states=encoder_outputs.hidden_states,
+            attentions=encoder_outputs.attentions,
+        )
+
+
+    def feature_extract(
+        self,
+        input_values,
+        seq_len,
+    ):
+        extract_features = self.feature_extractor(input_values)
+        extract_features = extract_features.transpose(1, 2)
+        extract_features = linear_interpolation(extract_features, seq_len=seq_len)
+
+        return extract_features
+
+    def encode(
+        self,
+        extract_features,
+        attention_mask=None,
+        mask_time_indices=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+        self.config.output_attentions = True
+
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if attention_mask is not None:
+            # compute reduced attention_mask corresponding to feature vectors
+            attention_mask = self._get_feature_vector_attention_mask(
+                extract_features.shape[1], attention_mask, add_adapter=False
+            )
+            
+
+        hidden_states, extract_features = self.feature_projection(extract_features)
+        hidden_states = self._mask_hidden_states(
+            hidden_states, mask_time_indices=mask_time_indices, attention_mask=attention_mask
+        )
+
+        encoder_outputs = self.encoder(
+            hidden_states,
+            attention_mask=attention_mask,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        hidden_states = encoder_outputs[0]
+
+        if self.adapter is not None:
+            hidden_states = self.adapter(hidden_states)
+
+        if not return_dict:
+            return (hidden_states, ) + encoder_outputs[1:]
+        return BaseModelOutput(
+            last_hidden_state=hidden_states,
+            hidden_states=encoder_outputs.hidden_states,
+            attentions=encoder_outputs.attentions,
+        )

src/models/attention.py

@@ -0,0 +1,840 @@
+# Adapted from https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention.py
+
+from typing import Any, Dict, Optional
+
+import torch
+from diffusers.models.attention import AdaLayerNorm, Attention, FeedForward
+from diffusers.models.embeddings import SinusoidalPositionalEmbedding
+from einops import rearrange
+from torch import nn
+
+from diffusers.models.attention import *
+from diffusers.models.attention_processor import *  
+
+class BasicTransformerBlock(nn.Module):
+    r"""
+    A basic Transformer block.
+
+    Parameters:
+        dim (`int`): The number of channels in the input and output.
+        num_attention_heads (`int`): The number of heads to use for multi-head attention.
+        attention_head_dim (`int`): The number of channels in each head.
+        dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
+        cross_attention_dim (`int`, *optional*): The size of the encoder_hidden_states vector for cross attention.
+        activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward.
+        num_embeds_ada_norm (:
+            obj: `int`, *optional*): The number of diffusion steps used during training. See `Transformer2DModel`.
+        attention_bias (:
+            obj: `bool`, *optional*, defaults to `False`): Configure if the attentions should contain a bias parameter.
+        only_cross_attention (`bool`, *optional*):
+            Whether to use only cross-attention layers. In this case two cross attention layers are used.
+        double_self_attention (`bool`, *optional*):
+            Whether to use two self-attention layers. In this case no cross attention layers are used.
+        upcast_attention (`bool`, *optional*):
+            Whether to upcast the attention computation to float32. This is useful for mixed precision training.
+        norm_elementwise_affine (`bool`, *optional*, defaults to `True`):
+            Whether to use learnable elementwise affine parameters for normalization.
+        norm_type (`str`, *optional*, defaults to `"layer_norm"`):
+            The normalization layer to use. Can be `"layer_norm"`, `"ada_norm"` or `"ada_norm_zero"`.
+        final_dropout (`bool` *optional*, defaults to False):
+            Whether to apply a final dropout after the last feed-forward layer.
+        attention_type (`str`, *optional*, defaults to `"default"`):
+            The type of attention to use. Can be `"default"` or `"gated"` or `"gated-text-image"`.
+        positional_embeddings (`str`, *optional*, defaults to `None`):
+            The type of positional embeddings to apply to.
+        num_positional_embeddings (`int`, *optional*, defaults to `None`):
+            The maximum number of positional embeddings to apply.
+    """
+
+    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",  # 'layer_norm', 'ada_norm', 'ada_norm_zero', 'ada_norm_single'
+        norm_eps: float = 1e-5,
+        final_dropout: bool = False,
+        attention_type: str = "default",
+        positional_embeddings: Optional[str] = None,
+        num_positional_embeddings: Optional[int] = None,
+    ):
+        super().__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"
+        self.use_ada_layer_norm_single = norm_type == "ada_norm_single"
+        self.use_layer_norm = norm_type == "layer_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}."
+            )
+
+        if positional_embeddings and (num_positional_embeddings is None):
+            raise ValueError(
+                "If `positional_embedding` type is defined, `num_positition_embeddings` must also be defined."
+            )
+
+        if positional_embeddings == "sinusoidal":
+            self.pos_embed = SinusoidalPositionalEmbedding(
+                dim, max_seq_length=num_positional_embeddings
+            )
+        else:
+            self.pos_embed = None
+
+        # 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, eps=norm_eps
+            )
+
+        self.attn1 = Attention(
+            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, eps=norm_eps
+                )
+            )
+            self.attn2 = Attention(
+                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
+        if not self.use_ada_layer_norm_single:
+            self.norm3 = nn.LayerNorm(
+                dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps
+            )
+
+        self.ff = FeedForward(
+            dim,
+            dropout=dropout,
+            activation_fn=activation_fn,
+            final_dropout=final_dropout,
+        )
+
+        # 4. Fuser
+        if attention_type == "gated" or attention_type == "gated-text-image":
+            self.fuser = GatedSelfAttentionDense(
+                dim, cross_attention_dim, num_attention_heads, attention_head_dim
+            )
+
+        # 5. Scale-shift for PixArt-Alpha.
+        if self.use_ada_layer_norm_single:
+            self.scale_shift_table = nn.Parameter(torch.randn(6, dim) / dim**0.5)
+
+        # let chunk size default to None
+        self._chunk_size = None
+        self._chunk_dim = 0
+
+    def set_chunk_feed_forward(self, chunk_size: Optional[int], dim: int = 0):
+        # Sets chunk feed-forward
+        self._chunk_size = chunk_size
+        self._chunk_dim = dim
+
+    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,
+    ) -> torch.FloatTensor:
+        # Notice that normalization is always applied before the real computation in the following blocks.
+        # 0. Self-Attention
+        batch_size = hidden_states.shape[0]
+
+        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
+            )
+        elif self.use_layer_norm:
+            norm_hidden_states = self.norm1(hidden_states)
+        elif self.use_ada_layer_norm_single:
+            shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
+                self.scale_shift_table[None] + timestep.reshape(batch_size, 6, -1)
+            ).chunk(6, dim=1)
+            norm_hidden_states = self.norm1(hidden_states)
+            norm_hidden_states = norm_hidden_states * (1 + scale_msa) + shift_msa
+            norm_hidden_states = norm_hidden_states.squeeze(1)
+        else:
+            raise ValueError("Incorrect norm used")
+
+        if self.pos_embed is not None:
+            norm_hidden_states = self.pos_embed(norm_hidden_states)
+
+        # 1. Retrieve lora scale.
+        lora_scale = (
+            cross_attention_kwargs.get("scale", 1.0)
+            if cross_attention_kwargs is not None
+            else 1.0
+        )
+
+        # 2. Prepare GLIGEN inputs
+        cross_attention_kwargs = (
+            cross_attention_kwargs.copy() if cross_attention_kwargs is not None else {}
+        )
+        gligen_kwargs = cross_attention_kwargs.pop("gligen", None)
+
+        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
+        elif self.use_ada_layer_norm_single:
+            attn_output = gate_msa * attn_output
+
+        hidden_states = attn_output + hidden_states
+        if hidden_states.ndim == 4:
+            hidden_states = hidden_states.squeeze(1)
+
+        # 2.5 GLIGEN Control
+        if gligen_kwargs is not None:
+            hidden_states = self.fuser(hidden_states, gligen_kwargs["objs"])
+
+        # 3. Cross-Attention
+        if self.attn2 is not None:
+            if self.use_ada_layer_norm:
+                norm_hidden_states = self.norm2(hidden_states, timestep)
+            elif self.use_ada_layer_norm_zero or self.use_layer_norm:
+                norm_hidden_states = self.norm2(hidden_states)
+            elif self.use_ada_layer_norm_single:
+                # For PixArt norm2 isn't applied here:
+                # https://github.com/PixArt-alpha/PixArt-alpha/blob/0f55e922376d8b797edd44d25d0e7464b260dcab/diffusion/model/nets/PixArtMS.py#L70C1-L76C103
+                norm_hidden_states = hidden_states
+            else:
+                raise ValueError("Incorrect norm")
+
+            if self.pos_embed is not None and self.use_ada_layer_norm_single is False:
+                norm_hidden_states = self.pos_embed(norm_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
+
+        # 4. Feed-forward
+        if not self.use_ada_layer_norm_single:
+            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.use_ada_layer_norm_single:
+            norm_hidden_states = self.norm2(hidden_states)
+            norm_hidden_states = norm_hidden_states * (1 + scale_mlp) + shift_mlp
+
+        ff_output = self.ff(norm_hidden_states, scale=lora_scale)
+
+        if self.use_ada_layer_norm_zero:
+            ff_output = gate_mlp.unsqueeze(1) * ff_output
+        elif self.use_ada_layer_norm_single:
+            ff_output = gate_mlp * ff_output
+
+        hidden_states = ff_output + hidden_states
+        if hidden_states.ndim == 4:
+            hidden_states = hidden_states.squeeze(1)
+
+        return hidden_states
+
+
+class TemporalBasicTransformerBlock(nn.Module):
+    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,
+        upcast_attention: bool = False,
+        unet_use_cross_frame_attention=None,
+        unet_use_temporal_attention=None,
+    ):
+        super().__init__()
+        self.only_cross_attention = only_cross_attention
+        self.use_ada_layer_norm = num_embeds_ada_norm is not None
+        self.unet_use_cross_frame_attention = unet_use_cross_frame_attention
+        self.unet_use_temporal_attention = unet_use_temporal_attention
+
+        # SC-Attn
+        self.attn1 = Attention(
+            query_dim=dim,
+            heads=num_attention_heads,
+            dim_head=attention_head_dim,
+            dropout=dropout,
+            bias=attention_bias,
+            upcast_attention=upcast_attention,
+        )
+        self.norm1 = (
+            AdaLayerNorm(dim, num_embeds_ada_norm)
+            if self.use_ada_layer_norm
+            else nn.LayerNorm(dim)
+        )
+
+        # Cross-Attn
+        if cross_attention_dim is not None:
+            self.attn2 = Attention(
+                query_dim=dim,
+                cross_attention_dim=cross_attention_dim,
+                heads=num_attention_heads,
+                dim_head=attention_head_dim,
+                dropout=dropout,
+                bias=attention_bias,
+                upcast_attention=upcast_attention,
+            )
+        else:
+            self.attn2 = None
+
+        if cross_attention_dim is not None:
+            self.norm2 = (
+                AdaLayerNorm(dim, num_embeds_ada_norm)
+                if self.use_ada_layer_norm
+                else nn.LayerNorm(dim)
+            )
+        else:
+            self.norm2 = None
+
+        # Feed-forward
+        self.ff = FeedForward(dim, dropout=dropout, activation_fn=activation_fn)
+        self.norm3 = nn.LayerNorm(dim)
+        self.use_ada_layer_norm_zero = False
+
+        # Temp-Attn
+        assert unet_use_temporal_attention is not None
+        if unet_use_temporal_attention:
+            self.attn_temp = Attention(
+                query_dim=dim,
+                heads=num_attention_heads,
+                dim_head=attention_head_dim,
+                dropout=dropout,
+                bias=attention_bias,
+                upcast_attention=upcast_attention,
+            )
+            nn.init.zeros_(self.attn_temp.to_out[0].weight.data)
+            self.norm_temp = (
+                AdaLayerNorm(dim, num_embeds_ada_norm)
+                if self.use_ada_layer_norm
+                else nn.LayerNorm(dim)
+            )
+
+    def forward(
+        self,
+        hidden_states,
+        encoder_hidden_states=None,
+        timestep=None,
+        attention_mask=None,
+        video_length=None,
+    ):
+        norm_hidden_states = (
+            self.norm1(hidden_states, timestep)
+            if self.use_ada_layer_norm
+            else self.norm1(hidden_states)
+        )
+
+        if self.unet_use_cross_frame_attention:
+            hidden_states = (
+                self.attn1(
+                    norm_hidden_states,
+                    attention_mask=attention_mask,
+                    video_length=video_length,
+                )
+                + hidden_states
+            )
+        else:
+            hidden_states = (
+                self.attn1(norm_hidden_states, attention_mask=attention_mask)
+                + hidden_states
+            )
+
+        if self.attn2 is not None:
+            # Cross-Attention
+            norm_hidden_states = (
+                self.norm2(hidden_states, timestep)
+                if self.use_ada_layer_norm
+                else self.norm2(hidden_states)
+            )
+            hidden_states = (
+                self.attn2(
+                    norm_hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    attention_mask=attention_mask,
+                )
+                + hidden_states
+            )
+
+        # Feed-forward
+        hidden_states = self.ff(self.norm3(hidden_states)) + hidden_states
+
+        # Temporal-Attention
+        if self.unet_use_temporal_attention:
+            d = hidden_states.shape[1]
+            hidden_states = rearrange(
+                hidden_states, "(b f) d c -> (b d) f c", f=video_length
+            )
+            norm_hidden_states = (
+                self.norm_temp(hidden_states, timestep)
+                if self.use_ada_layer_norm
+                else self.norm_temp(hidden_states)
+            )
+            hidden_states = self.attn_temp(norm_hidden_states) + hidden_states
+            hidden_states = rearrange(hidden_states, "(b d) f c -> (b f) d c", d=d)
+
+        return hidden_states
+
+
+class ResidualTemporalBasicTransformerBlock(TemporalBasicTransformerBlock):
+    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,
+        upcast_attention: bool = False,
+        unet_use_cross_frame_attention=None,
+        unet_use_temporal_attention=None,
+    ):
+        super(TemporalBasicTransformerBlock, self).__init__()
+        self.only_cross_attention = only_cross_attention
+        self.use_ada_layer_norm = num_embeds_ada_norm is not None
+        self.unet_use_cross_frame_attention = unet_use_cross_frame_attention
+        self.unet_use_temporal_attention = unet_use_temporal_attention
+
+        # SC-Attn
+        self.attn1 = ResidualAttention(
+            query_dim=dim,
+            heads=num_attention_heads,
+            dim_head=attention_head_dim,
+            dropout=dropout,
+            bias=attention_bias,
+            upcast_attention=upcast_attention,
+        )
+        self.norm1 = (
+            AdaLayerNorm(dim, num_embeds_ada_norm)
+            if self.use_ada_layer_norm
+            else nn.LayerNorm(dim)
+        )
+
+        # Cross-Attn
+        if cross_attention_dim is not None:
+            self.attn2 = ResidualAttention(
+                query_dim=dim,
+                cross_attention_dim=cross_attention_dim,
+                heads=num_attention_heads,
+                dim_head=attention_head_dim,
+                dropout=dropout,
+                bias=attention_bias,
+                upcast_attention=upcast_attention,
+            )
+        else:
+            self.attn2 = None
+
+        if cross_attention_dim is not None:
+            self.norm2 = (
+                AdaLayerNorm(dim, num_embeds_ada_norm)
+                if self.use_ada_layer_norm
+                else nn.LayerNorm(dim)
+            )
+        else:
+            self.norm2 = None
+
+        # Feed-forward
+        self.ff = FeedForward(dim, dropout=dropout, activation_fn=activation_fn)
+        self.norm3 = nn.LayerNorm(dim)
+        self.use_ada_layer_norm_zero = False
+
+        # Temp-Attn
+        assert unet_use_temporal_attention is not None
+        if unet_use_temporal_attention:
+            self.attn_temp = Attention(
+                query_dim=dim,
+                heads=num_attention_heads,
+                dim_head=attention_head_dim,
+                dropout=dropout,
+                bias=attention_bias,
+                upcast_attention=upcast_attention,
+            )
+            nn.init.zeros_(self.attn_temp.to_out[0].weight.data)
+            self.norm_temp = (
+                AdaLayerNorm(dim, num_embeds_ada_norm)
+                if self.use_ada_layer_norm
+                else nn.LayerNorm(dim)
+            )
+
+    def forward(
+        self,
+        hidden_states,
+        encoder_hidden_states=None,
+        timestep=None,
+        attention_mask=None,
+        video_length=None,
+        block_idx: Optional[int] = None,
+        additional_residuals: Optional[Dict[str, torch.FloatTensor]] = None
+    ):
+        norm_hidden_states = (
+            self.norm1(hidden_states, timestep)
+            if self.use_ada_layer_norm
+            else self.norm1(hidden_states)
+        )
+
+        if self.unet_use_cross_frame_attention:
+            hidden_states = (
+                self.attn1(
+                    norm_hidden_states,
+                    attention_mask=attention_mask,
+                    video_length=video_length,
+                    block_idx=block_idx,
+                    additional_residuals=additional_residuals,
+                )
+                + hidden_states
+            )
+        else:
+            hidden_states = (
+                self.attn1(norm_hidden_states, attention_mask=attention_mask,
+                    block_idx=block_idx,
+                    additional_residuals=additional_residuals
+                )
+                + hidden_states
+            )
+
+        if self.attn2 is not None:
+            # Cross-Attention
+            norm_hidden_states = (
+                self.norm2(hidden_states, timestep)
+                if self.use_ada_layer_norm
+                else self.norm2(hidden_states)
+            )
+            hidden_states = (
+                self.attn2(
+                    norm_hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    attention_mask=attention_mask,
+                    block_idx=block_idx,
+                    additional_residuals=additional_residuals,
+                )
+                + hidden_states
+            )
+
+        # Feed-forward
+        hidden_states = self.ff(self.norm3(hidden_states)) + hidden_states
+
+        # Temporal-Attention
+        if self.unet_use_temporal_attention:
+            d = hidden_states.shape[1]
+            hidden_states = rearrange(
+                hidden_states, "(b f) d c -> (b d) f c", f=video_length
+            )
+            norm_hidden_states = (
+                self.norm_temp(hidden_states, timestep)
+                if self.use_ada_layer_norm
+                else self.norm_temp(hidden_states)
+            )
+            hidden_states = self.attn_temp(norm_hidden_states) + hidden_states
+            hidden_states = rearrange(hidden_states, "(b d) f c -> (b f) d c", d=d)
+
+        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 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

src/models/motion_module.py

@@ -0,0 +1,388 @@
+# Adapt from https://github.com/guoyww/AnimateDiff/blob/main/animatediff/models/motion_module.py
+import math
+from dataclasses import dataclass
+from typing import Callable, Optional
+
+import torch
+from diffusers.models.attention import FeedForward
+from diffusers.models.attention_processor import Attention, AttnProcessor
+from diffusers.utils import BaseOutput
+from diffusers.utils.import_utils import is_xformers_available
+from einops import rearrange, repeat
+from torch import nn
+
+
+def zero_module(module):
+    # Zero out the parameters of a module and return it.
+    for p in module.parameters():
+        p.detach().zero_()
+    return module
+
+
+@dataclass
+class TemporalTransformer3DModelOutput(BaseOutput):
+    sample: torch.FloatTensor
+
+
+if is_xformers_available():
+    import xformers
+    import xformers.ops
+else:
+    xformers = None
+
+
+def get_motion_module(in_channels, motion_module_type: str, motion_module_kwargs: dict):
+    if motion_module_type == "Vanilla":
+        return VanillaTemporalModule(
+            in_channels=in_channels,
+            **motion_module_kwargs,
+        )
+    else:
+        raise ValueError
+
+
+class VanillaTemporalModule(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        num_attention_heads=8,
+        num_transformer_block=2,
+        attention_block_types=("Temporal_Self", "Temporal_Self"),
+        cross_frame_attention_mode=None,
+        temporal_position_encoding=False,
+        temporal_position_encoding_max_len=24,
+        temporal_attention_dim_div=1,
+        zero_initialize=True,
+    ):
+        super().__init__()
+
+        self.temporal_transformer = TemporalTransformer3DModel(
+            in_channels=in_channels,
+            num_attention_heads=num_attention_heads,
+            attention_head_dim=in_channels
+            // num_attention_heads
+            // temporal_attention_dim_div,
+            num_layers=num_transformer_block,
+            attention_block_types=attention_block_types,
+            cross_frame_attention_mode=cross_frame_attention_mode,
+            temporal_position_encoding=temporal_position_encoding,
+            temporal_position_encoding_max_len=temporal_position_encoding_max_len,
+        )
+
+        if zero_initialize:
+            self.temporal_transformer.proj_out = zero_module(
+                self.temporal_transformer.proj_out
+            )
+
+    def forward(
+        self,
+        input_tensor,
+        temb,
+        encoder_hidden_states,
+        attention_mask=None,
+        anchor_frame_idx=None,
+    ):
+        hidden_states = input_tensor
+        hidden_states = self.temporal_transformer(
+            hidden_states, encoder_hidden_states, attention_mask
+        )
+
+        output = hidden_states
+        return output
+
+
+class TemporalTransformer3DModel(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        num_attention_heads,
+        attention_head_dim,
+        num_layers,
+        attention_block_types=(
+            "Temporal_Self",
+            "Temporal_Self",
+        ),
+        dropout=0.0,
+        norm_num_groups=32,
+        cross_attention_dim=768,
+        activation_fn="geglu",
+        attention_bias=False,
+        upcast_attention=False,
+        cross_frame_attention_mode=None,
+        temporal_position_encoding=False,
+        temporal_position_encoding_max_len=24,
+    ):
+        super().__init__()
+
+        inner_dim = num_attention_heads * attention_head_dim
+
+        self.norm = torch.nn.GroupNorm(
+            num_groups=norm_num_groups, num_channels=in_channels, eps=1e-6, affine=True
+        )
+        self.proj_in = nn.Linear(in_channels, inner_dim)
+
+        self.transformer_blocks = nn.ModuleList(
+            [
+                TemporalTransformerBlock(
+                    dim=inner_dim,
+                    num_attention_heads=num_attention_heads,
+                    attention_head_dim=attention_head_dim,
+                    attention_block_types=attention_block_types,
+                    dropout=dropout,
+                    norm_num_groups=norm_num_groups,
+                    cross_attention_dim=cross_attention_dim,
+                    activation_fn=activation_fn,
+                    attention_bias=attention_bias,
+                    upcast_attention=upcast_attention,
+                    cross_frame_attention_mode=cross_frame_attention_mode,
+                    temporal_position_encoding=temporal_position_encoding,
+                    temporal_position_encoding_max_len=temporal_position_encoding_max_len,
+                )
+                for d in range(num_layers)
+            ]
+        )
+        self.proj_out = nn.Linear(inner_dim, in_channels)
+
+    def forward(self, hidden_states, encoder_hidden_states=None, attention_mask=None):
+        assert (
+            hidden_states.dim() == 5
+        ), f"Expected hidden_states to have ndim=5, but got ndim={hidden_states.dim()}."
+        video_length = hidden_states.shape[2]
+        hidden_states = rearrange(hidden_states, "b c f h w -> (b f) c h w")
+
+        batch, channel, height, weight = hidden_states.shape
+        residual = hidden_states
+
+        hidden_states = self.norm(hidden_states)
+        inner_dim = hidden_states.shape[1]
+        hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(
+            batch, height * weight, inner_dim
+        )
+        hidden_states = self.proj_in(hidden_states)
+
+        # Transformer Blocks
+        for block in self.transformer_blocks:
+            hidden_states = block(
+                hidden_states,
+                encoder_hidden_states=encoder_hidden_states,
+                video_length=video_length,
+            )
+
+        # output
+        hidden_states = self.proj_out(hidden_states)
+        hidden_states = (
+            hidden_states.reshape(batch, height, weight, inner_dim)
+            .permute(0, 3, 1, 2)
+            .contiguous()
+        )
+
+        output = hidden_states + residual
+        output = rearrange(output, "(b f) c h w -> b c f h w", f=video_length)
+
+        return output
+
+
+class TemporalTransformerBlock(nn.Module):
+    def __init__(
+        self,
+        dim,
+        num_attention_heads,
+        attention_head_dim,
+        attention_block_types=(
+            "Temporal_Self",
+            "Temporal_Self",
+        ),
+        dropout=0.0,
+        norm_num_groups=32,
+        cross_attention_dim=768,
+        activation_fn="geglu",
+        attention_bias=False,
+        upcast_attention=False,
+        cross_frame_attention_mode=None,
+        temporal_position_encoding=False,
+        temporal_position_encoding_max_len=24,
+    ):
+        super().__init__()
+
+        attention_blocks = []
+        norms = []
+
+        for block_name in attention_block_types:
+            attention_blocks.append(
+                VersatileAttention(
+                    attention_mode=block_name.split("_")[0],
+                    cross_attention_dim=cross_attention_dim
+                    if block_name.endswith("_Cross")
+                    else None,
+                    query_dim=dim,
+                    heads=num_attention_heads,
+                    dim_head=attention_head_dim,
+                    dropout=dropout,
+                    bias=attention_bias,
+                    upcast_attention=upcast_attention,
+                    cross_frame_attention_mode=cross_frame_attention_mode,
+                    temporal_position_encoding=temporal_position_encoding,
+                    temporal_position_encoding_max_len=temporal_position_encoding_max_len,
+                )
+            )
+            norms.append(nn.LayerNorm(dim))
+
+        self.attention_blocks = nn.ModuleList(attention_blocks)
+        self.norms = nn.ModuleList(norms)
+
+        self.ff = FeedForward(dim, dropout=dropout, activation_fn=activation_fn)
+        self.ff_norm = nn.LayerNorm(dim)
+
+    def forward(
+        self,
+        hidden_states,
+        encoder_hidden_states=None,
+        attention_mask=None,
+        video_length=None,
+    ):
+        for attention_block, norm in zip(self.attention_blocks, self.norms):
+            norm_hidden_states = norm(hidden_states)
+            hidden_states = (
+                attention_block(
+                    norm_hidden_states,
+                    encoder_hidden_states=encoder_hidden_states
+                    if attention_block.is_cross_attention
+                    else None,
+                    video_length=video_length,
+                )
+                + hidden_states
+            )
+
+        hidden_states = self.ff(self.ff_norm(hidden_states)) + hidden_states
+
+        output = hidden_states
+        return output
+
+
+class PositionalEncoding(nn.Module):
+    def __init__(self, d_model, dropout=0.0, max_len=24):
+        super().__init__()
+        self.dropout = nn.Dropout(p=dropout)
+        position = torch.arange(max_len).unsqueeze(1)
+        div_term = torch.exp(
+            torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model)
+        )
+        pe = torch.zeros(1, max_len, d_model)
+        pe[0, :, 0::2] = torch.sin(position * div_term)
+        pe[0, :, 1::2] = torch.cos(position * div_term)
+        self.register_buffer("pe", pe)
+
+    def forward(self, x):
+        x = x + self.pe[:, : x.size(1)]
+        return self.dropout(x)
+
+
+class VersatileAttention(Attention):
+    def __init__(
+        self,
+        attention_mode=None,
+        cross_frame_attention_mode=None,
+        temporal_position_encoding=False,
+        temporal_position_encoding_max_len=24,
+        *args,
+        **kwargs,
+    ):
+        super().__init__(*args, **kwargs)
+        assert attention_mode == "Temporal"
+
+        self.attention_mode = attention_mode
+        self.is_cross_attention = kwargs["cross_attention_dim"] is not None
+
+        self.pos_encoder = (
+            PositionalEncoding(
+                kwargs["query_dim"],
+                dropout=0.0,
+                max_len=temporal_position_encoding_max_len,
+            )
+            if (temporal_position_encoding and attention_mode == "Temporal")
+            else None
+        )
+
+    def extra_repr(self):
+        return f"(Module Info) Attention_Mode: {self.attention_mode}, Is_Cross_Attention: {self.is_cross_attention}"
+
+    def set_use_memory_efficient_attention_xformers(
+        self,
+        use_memory_efficient_attention_xformers: bool,
+        attention_op: Optional[Callable] = None,
+    ):
+        if use_memory_efficient_attention_xformers:
+            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
+
+            # XFormersAttnProcessor corrupts video generation and work with Pytorch 1.13.
+            # Pytorch 2.0.1 AttnProcessor works the same as XFormersAttnProcessor in Pytorch 1.13.
+            # You don't need XFormersAttnProcessor here.
+            # processor = XFormersAttnProcessor(
+            #     attention_op=attention_op,
+            # )
+            processor = AttnProcessor()
+        else:
+            processor = AttnProcessor()
+
+        self.set_processor(processor)
+
+    def forward(
+        self,
+        hidden_states,
+        encoder_hidden_states=None,
+        attention_mask=None,
+        video_length=None,
+        **cross_attention_kwargs,
+    ):
+        if self.attention_mode == "Temporal":
+            d = hidden_states.shape[1]  # d means HxW
+            hidden_states = rearrange(
+                hidden_states, "(b f) d c -> (b d) f c", f=video_length
+            )
+
+            if self.pos_encoder is not None:
+                hidden_states = self.pos_encoder(hidden_states)
+
+            encoder_hidden_states = (
+                repeat(encoder_hidden_states, "b n c -> (b d) n c", d=d)
+                if encoder_hidden_states is not None
+                else encoder_hidden_states
+            )
+
+        else:
+            raise NotImplementedError
+
+        hidden_states = self.processor(
+            self,
+            hidden_states,
+            encoder_hidden_states=encoder_hidden_states,
+            attention_mask=attention_mask,
+            **cross_attention_kwargs,
+        )
+
+        if self.attention_mode == "Temporal":
+            hidden_states = rearrange(hidden_states, "(b d) f c -> (b f) d c", d=d)
+
+        return hidden_states

src/models/mutual_self_attention.py

@@ -0,0 +1,363 @@
+# Adapted from https://github.com/magic-research/magic-animate/blob/main/magicanimate/models/mutual_self_attention.py
+from typing import Any, Dict, Optional
+
+import torch
+from einops import rearrange
+
+from src.models.attention import TemporalBasicTransformerBlock
+
+from .attention import BasicTransformerBlock
+
+
+def torch_dfs(model: torch.nn.Module):
+    result = [model]
+    for child in model.children():
+        result += torch_dfs(child)
+    return result
+
+
+class ReferenceAttentionControl:
+    def __init__(
+        self,
+        unet,
+        mode="write",
+        do_classifier_free_guidance=False,
+        attention_auto_machine_weight=float("inf"),
+        gn_auto_machine_weight=1.0,
+        style_fidelity=1.0,
+        reference_attn=True,
+        reference_adain=False,
+        fusion_blocks="midup",
+        batch_size=1,
+    ) -> None:
+        # 10. Modify self attention and group norm
+        self.unet = unet
+        assert mode in ["read", "write"]
+        assert fusion_blocks in ["midup", "full"]
+        self.reference_attn = reference_attn
+        self.reference_adain = reference_adain
+        self.fusion_blocks = fusion_blocks
+        self.register_reference_hooks(
+            mode,
+            do_classifier_free_guidance,
+            attention_auto_machine_weight,
+            gn_auto_machine_weight,
+            style_fidelity,
+            reference_attn,
+            reference_adain,
+            fusion_blocks,
+            batch_size=batch_size,
+        )
+
+    def register_reference_hooks(
+        self,
+        mode,
+        do_classifier_free_guidance,
+        attention_auto_machine_weight,
+        gn_auto_machine_weight,
+        style_fidelity,
+        reference_attn,
+        reference_adain,
+        dtype=torch.float16,
+        batch_size=1,
+        num_images_per_prompt=1,
+        device=torch.device("cpu"),
+        fusion_blocks="midup",
+    ):
+        MODE = mode
+        do_classifier_free_guidance = do_classifier_free_guidance
+        attention_auto_machine_weight = attention_auto_machine_weight
+        gn_auto_machine_weight = gn_auto_machine_weight
+        style_fidelity = style_fidelity
+        reference_attn = reference_attn
+        reference_adain = reference_adain
+        fusion_blocks = fusion_blocks
+        num_images_per_prompt = num_images_per_prompt
+        dtype = dtype
+        if do_classifier_free_guidance:
+            uc_mask = (
+                torch.Tensor(
+                    [1] * batch_size * num_images_per_prompt * 16
+                    + [0] * batch_size * num_images_per_prompt * 16
+                )
+                .to(device)
+                .bool()
+            )
+        else:
+            uc_mask = (
+                torch.Tensor([0] * batch_size * num_images_per_prompt * 2)
+                .to(device)
+                .bool()
+            )
+
+        def hacked_basic_transformer_inner_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,
+            video_length=None,
+        ):
+            if self.use_ada_layer_norm:  # False
+                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)
+
+            # 1. Self-Attention
+            # self.only_cross_attention = False
+            cross_attention_kwargs = (
+                cross_attention_kwargs if cross_attention_kwargs is not None else {}
+            )
+            if self.only_cross_attention:
+                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,
+                )
+            else:
+                if MODE == "write":
+                    self.bank.append(norm_hidden_states.clone())
+                    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 MODE == "read":
+                    bank_fea = [
+                        rearrange(
+                            d.unsqueeze(1).repeat(1, video_length, 1, 1),
+                            "b t l c -> (b t) l c",
+                        )
+                        for d in self.bank
+                    ]
+                    modify_norm_hidden_states = torch.cat(
+                        [norm_hidden_states] + bank_fea, dim=1
+                    )
+                    hidden_states_uc = (
+                        self.attn1(
+                            norm_hidden_states,
+                            encoder_hidden_states=modify_norm_hidden_states,
+                            attention_mask=attention_mask,
+                        )
+                        + hidden_states
+                    )
+                    if do_classifier_free_guidance:
+                        hidden_states_c = hidden_states_uc.clone()
+                        _uc_mask = uc_mask.clone()
+                        if hidden_states.shape[0] != _uc_mask.shape[0]:
+                            _uc_mask = (
+                                torch.Tensor(
+                                    [1] * (hidden_states.shape[0] // 2)
+                                    + [0] * (hidden_states.shape[0] // 2)
+                                )
+                                .to(device)
+                                .bool()
+                            )
+                        hidden_states_c[_uc_mask] = (
+                            self.attn1(
+                                norm_hidden_states[_uc_mask],
+                                encoder_hidden_states=norm_hidden_states[_uc_mask],
+                                attention_mask=attention_mask,
+                            )
+                            + hidden_states[_uc_mask]
+                        )
+                        hidden_states = hidden_states_c.clone()
+                    else:
+                        hidden_states = hidden_states_uc
+
+                    # self.bank.clear()
+                    if self.attn2 is not None:
+                        # Cross-Attention
+                        norm_hidden_states = (
+                            self.norm2(hidden_states, timestep)
+                            if self.use_ada_layer_norm
+                            else self.norm2(hidden_states)
+                        )
+                        hidden_states = (
+                            self.attn2(
+                                norm_hidden_states,
+                                encoder_hidden_states=encoder_hidden_states,
+                                attention_mask=attention_mask,
+                            )
+                            + hidden_states
+                        )
+
+                    # Feed-forward
+                    hidden_states = self.ff(self.norm3(hidden_states)) + hidden_states
+
+                    # Temporal-Attention
+                    if self.unet_use_temporal_attention:
+                        d = hidden_states.shape[1]
+                        hidden_states = rearrange(
+                            hidden_states, "(b f) d c -> (b d) f c", f=video_length
+                        )
+                        norm_hidden_states = (
+                            self.norm_temp(hidden_states, timestep)
+                            if self.use_ada_layer_norm
+                            else self.norm_temp(hidden_states)
+                        )
+                        hidden_states = (
+                            self.attn_temp(norm_hidden_states) + hidden_states
+                        )
+                        hidden_states = rearrange(
+                            hidden_states, "(b d) f c -> (b f) d c", d=d
+                        )
+
+                    return hidden_states
+
+            if self.use_ada_layer_norm_zero:
+                attn_output = gate_msa.unsqueeze(1) * attn_output
+            hidden_states = attn_output + hidden_states
+
+            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)
+                )
+
+                # 2. Cross-Attention
+                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
+
+            # 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]
+                )
+
+            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
+
+        if self.reference_attn:
+            if self.fusion_blocks == "midup":
+                attn_modules = [
+                    module
+                    for module in (
+                        torch_dfs(self.unet.mid_block) + torch_dfs(self.unet.up_blocks)
+                    )
+                    if isinstance(module, BasicTransformerBlock)
+                    or isinstance(module, TemporalBasicTransformerBlock)
+                ]
+            elif self.fusion_blocks == "full":
+                attn_modules = [
+                    module
+                    for module in torch_dfs(self.unet)
+                    if isinstance(module, BasicTransformerBlock)
+                    or isinstance(module, TemporalBasicTransformerBlock)
+                ]
+            attn_modules = sorted(
+                attn_modules, key=lambda x: -x.norm1.normalized_shape[0]
+            )
+
+            for i, module in enumerate(attn_modules):
+                module._original_inner_forward = module.forward
+                if isinstance(module, BasicTransformerBlock):
+                    module.forward = hacked_basic_transformer_inner_forward.__get__(
+                        module, BasicTransformerBlock
+                    )
+                if isinstance(module, TemporalBasicTransformerBlock):
+                    module.forward = hacked_basic_transformer_inner_forward.__get__(
+                        module, TemporalBasicTransformerBlock
+                    )
+
+                module.bank = []
+                module.attn_weight = float(i) / float(len(attn_modules))
+
+    def update(self, writer, dtype=torch.float16):
+        if self.reference_attn:
+            if self.fusion_blocks == "midup":
+                reader_attn_modules = [
+                    module
+                    for module in (
+                        torch_dfs(self.unet.mid_block) + torch_dfs(self.unet.up_blocks)
+                    )
+                    if isinstance(module, TemporalBasicTransformerBlock)
+                ]
+                writer_attn_modules = [
+                    module
+                    for module in (
+                        torch_dfs(writer.unet.mid_block)
+                        + torch_dfs(writer.unet.up_blocks)
+                    )
+                    if isinstance(module, BasicTransformerBlock)
+                ]
+            elif self.fusion_blocks == "full":
+                reader_attn_modules = [
+                    module
+                    for module in torch_dfs(self.unet)
+                    if isinstance(module, TemporalBasicTransformerBlock)
+                ]
+                writer_attn_modules = [
+                    module
+                    for module in torch_dfs(writer.unet)
+                    if isinstance(module, BasicTransformerBlock)
+                ]
+            reader_attn_modules = sorted(
+                reader_attn_modules, key=lambda x: -x.norm1.normalized_shape[0]
+            )
+            writer_attn_modules = sorted(
+                writer_attn_modules, key=lambda x: -x.norm1.normalized_shape[0]
+            )
+            for r, w in zip(reader_attn_modules, writer_attn_modules):
+                r.bank = [v.clone().to(dtype) for v in w.bank]
+                # w.bank.clear()
+
+    def clear(self):
+        if self.reference_attn:
+            if self.fusion_blocks == "midup":
+                reader_attn_modules = [
+                    module
+                    for module in (
+                        torch_dfs(self.unet.mid_block) + torch_dfs(self.unet.up_blocks)
+                    )
+                    if isinstance(module, BasicTransformerBlock)
+                    or isinstance(module, TemporalBasicTransformerBlock)
+                ]
+            elif self.fusion_blocks == "full":
+                reader_attn_modules = [
+                    module
+                    for module in torch_dfs(self.unet)
+                    if isinstance(module, BasicTransformerBlock)
+                    or isinstance(module, TemporalBasicTransformerBlock)
+                ]
+            reader_attn_modules = sorted(
+                reader_attn_modules, key=lambda x: -x.norm1.normalized_shape[0]
+            )
+            for r in reader_attn_modules:
+                r.bank.clear()

src/models/pose_guider.py

@@ -0,0 +1,329 @@
+import os
+import torch
+import torch.nn as nn
+import torch.nn.init as init
+from einops import rearrange
+import numpy as np
+from diffusers.models.modeling_utils import ModelMixin
+
+from typing import Any, Dict, Optional
+from src.models.attention import BasicTransformerBlock
+
+
+class PoseGuider(ModelMixin):
+    def __init__(self, noise_latent_channels=320, use_ca=True):
+        super(PoseGuider, self).__init__()
+
+        self.use_ca = use_ca
+
+        self.conv_layers = nn.Sequential(
+            nn.Conv2d(in_channels=3, out_channels=3, kernel_size=3, padding=1),
+            nn.BatchNorm2d(3),
+            nn.ReLU(),
+            nn.Conv2d(in_channels=3, out_channels=16, kernel_size=4, stride=2, padding=1),
+            nn.BatchNorm2d(16),
+            nn.ReLU(),
+
+            nn.Conv2d(in_channels=16, out_channels=16, kernel_size=3, padding=1),
+            nn.BatchNorm2d(16),
+            nn.ReLU(),
+            nn.Conv2d(in_channels=16, out_channels=32, kernel_size=4, stride=2, padding=1),
+            nn.BatchNorm2d(32),
+            nn.ReLU(),
+
+            nn.Conv2d(in_channels=32, out_channels=32, kernel_size=3, padding=1),
+            nn.BatchNorm2d(32),
+            nn.ReLU(),
+            nn.Conv2d(in_channels=32, out_channels=64, kernel_size=4, stride=2, padding=1),
+            nn.BatchNorm2d(64),
+            nn.ReLU(),
+
+            nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, padding=1),
+            nn.BatchNorm2d(64),
+            nn.ReLU(),
+            nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(128),
+            nn.ReLU()
+        )
+
+        # Final projection layer
+        self.final_proj = nn.Conv2d(in_channels=128, out_channels=noise_latent_channels, kernel_size=1)
+        
+        self.conv_layers_1 = nn.Sequential(
+            nn.Conv2d(in_channels=noise_latent_channels, out_channels=noise_latent_channels, kernel_size=3, padding=1),
+            nn.BatchNorm2d(noise_latent_channels),
+            nn.ReLU(),
+            nn.Conv2d(in_channels=noise_latent_channels, out_channels=noise_latent_channels, kernel_size=3, stride=2, padding=1),
+            nn.BatchNorm2d(noise_latent_channels),
+            nn.ReLU(),
+        )
+        
+        self.conv_layers_2 = nn.Sequential(
+            nn.Conv2d(in_channels=noise_latent_channels, out_channels=noise_latent_channels, kernel_size=3, padding=1),
+            nn.BatchNorm2d(noise_latent_channels),
+            nn.ReLU(),
+            nn.Conv2d(in_channels=noise_latent_channels, out_channels=noise_latent_channels*2, kernel_size=3, stride=2, padding=1),
+            nn.BatchNorm2d(noise_latent_channels*2),
+            nn.ReLU(),
+        )
+
+        self.conv_layers_3 = nn.Sequential(
+            nn.Conv2d(in_channels=noise_latent_channels*2, out_channels=noise_latent_channels*2, kernel_size=3, padding=1),
+            nn.BatchNorm2d(noise_latent_channels*2),
+            nn.ReLU(),
+            nn.Conv2d(in_channels=noise_latent_channels*2, out_channels=noise_latent_channels*4, kernel_size=3, stride=2, padding=1),
+            nn.BatchNorm2d(noise_latent_channels*4),
+            nn.ReLU(),
+        )
+        
+        self.conv_layers_4 = nn.Sequential(
+            nn.Conv2d(in_channels=noise_latent_channels*4, out_channels=noise_latent_channels*4, kernel_size=3, padding=1),
+            nn.BatchNorm2d(noise_latent_channels*4),
+            nn.ReLU(),
+        )
+        
+        if self.use_ca:
+            self.cross_attn1 = Transformer2DModel(in_channels=noise_latent_channels)
+            self.cross_attn2 = Transformer2DModel(in_channels=noise_latent_channels*2)
+            self.cross_attn3 = Transformer2DModel(in_channels=noise_latent_channels*4)
+            self.cross_attn4 = Transformer2DModel(in_channels=noise_latent_channels*4)
+
+        # Initialize layers
+        self._initialize_weights()
+
+        self.scale = nn.Parameter(torch.ones(1) * 2)
+
+    # def _initialize_weights(self):
+    #     # Initialize weights with Gaussian distribution and zero out the final layer
+    #     for m in self.conv_layers:
+    #         if isinstance(m, nn.Conv2d):
+    #             init.normal_(m.weight, mean=0.0, std=0.02)
+    #             if m.bias is not None:
+    #                 init.zeros_(m.bias)
+
+    #     init.zeros_(self.final_proj.weight)
+    #     if self.final_proj.bias is not None:
+    #         init.zeros_(self.final_proj.bias)
+    
+    def _initialize_weights(self):
+        # Initialize weights with He initialization and zero out the biases
+        conv_blocks = [self.conv_layers, self.conv_layers_1, self.conv_layers_2, self.conv_layers_3, self.conv_layers_4]
+        for block_item in conv_blocks:
+            for m in block_item:
+                if isinstance(m, nn.Conv2d):
+                    n = m.kernel_size[0] * m.kernel_size[1] * m.in_channels
+                    init.normal_(m.weight, mean=0.0, std=np.sqrt(2. / n))
+                    if m.bias is not None:
+                        init.zeros_(m.bias)
+
+        # For the final projection layer, initialize weights to zero (or you may choose to use He initialization here as well)
+        init.zeros_(self.final_proj.weight)
+        if self.final_proj.bias is not None:
+            init.zeros_(self.final_proj.bias)
+
+    def forward(self, x, ref_x):
+        fea = []
+        b = x.shape[0]
+        
+        x = rearrange(x, "b c f h w -> (b f) c h w")
+        x = self.conv_layers(x)
+        x = self.final_proj(x)
+        x = x * self.scale
+        # x = rearrange(x, "(b f) c h w -> b c f h w", b=b)
+        fea.append(rearrange(x, "(b f) c h w -> b c f h w", b=b))
+        
+        x = self.conv_layers_1(x)
+        if self.use_ca:
+            ref_x = self.conv_layers(ref_x)
+            ref_x = self.final_proj(ref_x)
+            ref_x = ref_x * self.scale
+            ref_x = self.conv_layers_1(ref_x)
+            x = self.cross_attn1(x, ref_x)
+        fea.append(rearrange(x, "(b f) c h w -> b c f h w", b=b))
+        
+        x = self.conv_layers_2(x)
+        if self.use_ca:
+            ref_x = self.conv_layers_2(ref_x)
+            x = self.cross_attn2(x, ref_x)
+        fea.append(rearrange(x, "(b f) c h w -> b c f h w", b=b))
+        
+        x = self.conv_layers_3(x)
+        if self.use_ca:
+            ref_x = self.conv_layers_3(ref_x)
+            x = self.cross_attn3(x, ref_x)
+        fea.append(rearrange(x, "(b f) c h w -> b c f h w", b=b))
+        
+        x = self.conv_layers_4(x)
+        if self.use_ca:
+            ref_x = self.conv_layers_4(ref_x)
+            x = self.cross_attn4(x, ref_x)
+        fea.append(rearrange(x, "(b f) c h w -> b c f h w", b=b))
+
+        return fea
+
+    # @classmethod
+    # def from_pretrained(cls,pretrained_model_path):
+    #     if not os.path.exists(pretrained_model_path):
+    #         print(f"There is no model file in {pretrained_model_path}")
+    #     print(f"loaded PoseGuider's pretrained weights from {pretrained_model_path} ...")
+
+    #     state_dict = torch.load(pretrained_model_path, map_location="cpu")
+    #     model = Hack_PoseGuider(noise_latent_channels=320)
+                
+    #     m, u = model.load_state_dict(state_dict, strict=True)
+    #     # print(f"### missing keys: {len(m)}; \n### unexpected keys: {len(u)};")        
+    #     params = [p.numel() for n, p in model.named_parameters()]
+    #     print(f"### PoseGuider's Parameters: {sum(params) / 1e6} M")
+        
+    #     return model
+    
+
+class Transformer2DModel(ModelMixin):
+    _supports_gradient_checkpointing = True
+    def __init__(
+        self,
+        num_attention_heads: int = 16,
+        attention_head_dim: int = 88,
+        in_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,
+        activation_fn: str = "geglu",
+        num_embeds_ada_norm: Optional[int] = None,
+        use_linear_projection: bool = False,
+        only_cross_attention: bool = False,
+        double_self_attention: bool = False,
+        upcast_attention: bool = False,
+        norm_type: str = "layer_norm",
+        norm_elementwise_affine: bool = True,
+        norm_eps: float = 1e-5,
+        attention_type: str = "default",
+    ):
+        super().__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
+
+        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 = nn.Linear(in_channels, inner_dim)
+        else:
+            self.proj_in = nn.Conv2d(
+                in_channels, inner_dim, kernel_size=1, stride=1, padding=0
+            )
+
+        # 3. Define transformers blocks
+        self.transformer_blocks = nn.ModuleList(
+            [
+                BasicTransformerBlock(
+                    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,
+                    double_self_attention=double_self_attention,
+                    upcast_attention=upcast_attention,
+                    norm_type=norm_type,
+                    norm_elementwise_affine=norm_elementwise_affine,
+                    norm_eps=norm_eps,
+                    attention_type=attention_type,
+                )
+                for d in range(num_layers)
+            ]
+        )
+
+        if use_linear_projection:
+            self.proj_out = nn.Linear(inner_dim, in_channels)
+        else:
+            self.proj_out = nn.Conv2d(
+                inner_dim, in_channels, kernel_size=1, stride=1, padding=0
+            )
+
+        self.gradient_checkpointing = False
+
+    def _set_gradient_checkpointing(self, module, value=False):
+        if hasattr(module, "gradient_checkpointing"):
+            module.gradient_checkpointing = value
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: Optional[torch.Tensor] = None,
+        timestep: Optional[torch.LongTensor] = None,
+    ):
+        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)
+
+        for block in self.transformer_blocks:
+            hidden_states = block(
+                hidden_states,
+                encoder_hidden_states=encoder_hidden_states,
+                timestep=timestep,
+            )
+
+        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
+        return output
+
+
+if __name__ == '__main__':
+    model = PoseGuider(noise_latent_channels=320).to(device="cuda")
+    
+    input_data = torch.randn(1,3,1,512,512).to(device="cuda")
+    input_data1 = torch.randn(1,3,512,512).to(device="cuda")
+    
+    output = model(input_data, input_data1)
+    for item in output:
+        print(item.shape)
+    
+    # tf_model = Transformer2DModel(
+    #     in_channels=320
+    #     ).to('cuda')
+    
+    # input_data = torch.randn(4,320,32,32).to(device="cuda")
+    # # input_emb = torch.randn(4,1,768).to(device="cuda")
+    # input_emb = torch.randn(4,320,32,32).to(device="cuda")
+    # o1 = tf_model(input_data, input_emb)
+    # print(o1.shape)

src/models/resnet.py

@@ -0,0 +1,252 @@
+# Adapted from https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/resnet.py
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+from typing import Dict, Optional
+
+
+class InflatedConv3d(nn.Conv2d):
+    def forward(self, x):
+        video_length = x.shape[2]
+
+        x = rearrange(x, "b c f h w -> (b f) c h w")
+        x = super().forward(x)
+        x = rearrange(x, "(b f) c h w -> b c f h w", f=video_length)
+
+        return x
+
+
+class InflatedGroupNorm(nn.GroupNorm):
+    def forward(self, x):
+        video_length = x.shape[2]
+
+        x = rearrange(x, "b c f h w -> (b f) c h w")
+        x = super().forward(x)
+        x = rearrange(x, "(b f) c h w -> b c f h w", f=video_length)
+
+        return x
+
+
+class Upsample3D(nn.Module):
+    def __init__(
+        self,
+        channels,
+        use_conv=False,
+        use_conv_transpose=False,
+        out_channels=None,
+        name="conv",
+    ):
+        super().__init__()
+        self.channels = channels
+        self.out_channels = out_channels or channels
+        self.use_conv = use_conv
+        self.use_conv_transpose = use_conv_transpose
+        self.name = name
+
+        conv = None
+        if use_conv_transpose:
+            raise NotImplementedError
+        elif use_conv:
+            self.conv = InflatedConv3d(self.channels, self.out_channels, 3, padding=1)
+
+    def forward(self, hidden_states, output_size=None):
+        assert hidden_states.shape[1] == self.channels
+
+        if self.use_conv_transpose:
+            raise NotImplementedError
+
+        # Cast to float32 to as 'upsample_nearest2d_out_frame' op does not support bfloat16
+        dtype = hidden_states.dtype
+        if dtype == torch.bfloat16:
+            hidden_states = hidden_states.to(torch.float32)
+
+        # upsample_nearest_nhwc fails with large batch sizes. see https://github.com/huggingface/diffusers/issues/984
+        if hidden_states.shape[0] >= 64:
+            hidden_states = hidden_states.contiguous()
+
+        # if `output_size` is passed we force the interpolation output
+        # size and do not make use of `scale_factor=2`
+        if output_size is None:
+            hidden_states = F.interpolate(
+                hidden_states, scale_factor=[1.0, 2.0, 2.0], mode="nearest"
+            )
+        else:
+            hidden_states = F.interpolate(
+                hidden_states, size=output_size, mode="nearest"
+            )
+
+        # If the input is bfloat16, we cast back to bfloat16
+        if dtype == torch.bfloat16:
+            hidden_states = hidden_states.to(dtype)
+
+        # if self.use_conv:
+        #     if self.name == "conv":
+        #         hidden_states = self.conv(hidden_states)
+        #     else:
+        #         hidden_states = self.Conv2d_0(hidden_states)
+        hidden_states = self.conv(hidden_states)
+
+        return hidden_states
+
+
+class Downsample3D(nn.Module):
+    def __init__(
+        self, channels, use_conv=False, out_channels=None, padding=1, name="conv"
+    ):
+        super().__init__()
+        self.channels = channels
+        self.out_channels = out_channels or channels
+        self.use_conv = use_conv
+        self.padding = padding
+        stride = 2
+        self.name = name
+
+        if use_conv:
+            self.conv = InflatedConv3d(
+                self.channels, self.out_channels, 3, stride=stride, padding=padding
+            )
+        else:
+            raise NotImplementedError
+
+    def forward(self, hidden_states):
+        assert hidden_states.shape[1] == self.channels
+        if self.use_conv and self.padding == 0:
+            raise NotImplementedError
+
+        assert hidden_states.shape[1] == self.channels
+        hidden_states = self.conv(hidden_states)
+
+        return hidden_states
+
+
+class ResnetBlock3D(nn.Module):
+    def __init__(
+        self,
+        *,
+        in_channels,
+        out_channels=None,
+        conv_shortcut=False,
+        dropout=0.0,
+        temb_channels=512,
+        groups=32,
+        groups_out=None,
+        pre_norm=True,
+        eps=1e-6,
+        non_linearity="swish",
+        time_embedding_norm="default",
+        output_scale_factor=1.0,
+        use_in_shortcut=None,
+        use_inflated_groupnorm=None,
+    ):
+        super().__init__()
+        self.pre_norm = pre_norm
+        self.pre_norm = True
+        self.in_channels = in_channels
+        out_channels = in_channels if out_channels is None else out_channels
+        self.out_channels = out_channels
+        self.use_conv_shortcut = conv_shortcut
+        self.time_embedding_norm = time_embedding_norm
+        self.output_scale_factor = output_scale_factor
+
+        if groups_out is None:
+            groups_out = groups
+
+        assert use_inflated_groupnorm != None
+        if use_inflated_groupnorm:
+            self.norm1 = InflatedGroupNorm(
+                num_groups=groups, num_channels=in_channels, eps=eps, affine=True
+            )
+        else:
+            self.norm1 = torch.nn.GroupNorm(
+                num_groups=groups, num_channels=in_channels, eps=eps, affine=True
+            )
+
+        self.conv1 = InflatedConv3d(
+            in_channels, out_channels, kernel_size=3, stride=1, padding=1
+        )
+
+        if temb_channels is not None:
+            if self.time_embedding_norm == "default":
+                time_emb_proj_out_channels = out_channels
+            elif self.time_embedding_norm == "scale_shift":
+                time_emb_proj_out_channels = out_channels * 2
+            else:
+                raise ValueError(
+                    f"unknown time_embedding_norm : {self.time_embedding_norm} "
+                )
+
+            self.time_emb_proj = torch.nn.Linear(
+                temb_channels, time_emb_proj_out_channels
+            )
+        else:
+            self.time_emb_proj = None
+
+        if use_inflated_groupnorm:
+            self.norm2 = InflatedGroupNorm(
+                num_groups=groups_out, num_channels=out_channels, eps=eps, affine=True
+            )
+        else:
+            self.norm2 = torch.nn.GroupNorm(
+                num_groups=groups_out, num_channels=out_channels, eps=eps, affine=True
+            )
+        self.dropout = torch.nn.Dropout(dropout)
+        self.conv2 = InflatedConv3d(
+            out_channels, out_channels, kernel_size=3, stride=1, padding=1
+        )
+
+        if non_linearity == "swish":
+            self.nonlinearity = lambda x: F.silu(x)
+        elif non_linearity == "mish":
+            self.nonlinearity = Mish()
+        elif non_linearity == "silu":
+            self.nonlinearity = nn.SiLU()
+
+        self.use_in_shortcut = (
+            self.in_channels != self.out_channels
+            if use_in_shortcut is None
+            else use_in_shortcut
+        )
+
+        self.conv_shortcut = None
+        if self.use_in_shortcut:
+            self.conv_shortcut = InflatedConv3d(
+                in_channels, out_channels, kernel_size=1, stride=1, padding=0
+            )
+
+    def forward(self, input_tensor, temb):
+        hidden_states = input_tensor
+
+        hidden_states = self.norm1(hidden_states)
+        hidden_states = self.nonlinearity(hidden_states)
+
+        hidden_states = self.conv1(hidden_states)
+
+        if temb is not None:
+            temb = self.time_emb_proj(self.nonlinearity(temb))[:, :, None, None, None]
+
+        if temb is not None and self.time_embedding_norm == "default":
+            hidden_states = hidden_states + temb
+
+        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)
+
+        output_tensor = (input_tensor + hidden_states) / self.output_scale_factor
+
+        return output_tensor
+
+class Mish(torch.nn.Module):
+    def forward(self, hidden_states):
+        return hidden_states * torch.tanh(torch.nn.functional.softplus(hidden_states))

src/models/transformer_2d.py

@@ -0,0 +1,396 @@
+# Adapted from https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformer_2d.py
+from dataclasses import dataclass
+from typing import Any, Dict, Optional
+
+import torch
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.embeddings import CaptionProjection
+from diffusers.models.lora import LoRACompatibleConv, LoRACompatibleLinear
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.normalization import AdaLayerNormSingle
+from diffusers.utils import USE_PEFT_BACKEND, BaseOutput, deprecate, is_torch_version
+from torch import nn
+
+from .attention import BasicTransformerBlock
+
+
+@dataclass
+class Transformer2DModelOutput(BaseOutput):
+    """
+    The output of [`Transformer2DModel`].
+
+    Args:
+        sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` or `(batch size, num_vector_embeds - 1, num_latent_pixels)` if [`Transformer2DModel`] is discrete):
+            The hidden states output conditioned on the `encoder_hidden_states` input. If discrete, returns probability
+            distributions for the unnoised latent pixels.
+    """
+
+    sample: torch.FloatTensor
+    ref_feature: torch.FloatTensor
+
+
+class Transformer2DModel(ModelMixin, ConfigMixin):
+    """
+    A 2D Transformer model for image-like data.
+
+    Parameters:
+        num_attention_heads (`int`, *optional*, defaults to 16): The number of heads to use for multi-head attention.
+        attention_head_dim (`int`, *optional*, defaults to 88): The number of channels in each head.
+        in_channels (`int`, *optional*):
+            The number of channels in the input and output (specify if the input is **continuous**).
+        num_layers (`int`, *optional*, defaults to 1): The number of layers of Transformer blocks to use.
+        dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
+        cross_attention_dim (`int`, *optional*): The number of `encoder_hidden_states` dimensions to use.
+        sample_size (`int`, *optional*): The width of the latent images (specify if the input is **discrete**).
+            This is fixed during training since it is used to learn a number of position embeddings.
+        num_vector_embeds (`int`, *optional*):
+            The number of classes of the vector embeddings of the latent pixels (specify if the input is **discrete**).
+            Includes the class for the masked latent pixel.
+        activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to use in feed-forward.
+        num_embeds_ada_norm ( `int`, *optional*):
+            The number of diffusion steps used during training. Pass if at least one of the norm_layers is
+            `AdaLayerNorm`. This is fixed during training since it is used to learn a number of embeddings that are
+            added to the hidden states.
+
+            During inference, you can denoise for up to but not more steps than `num_embeds_ada_norm`.
+        attention_bias (`bool`, *optional*):
+            Configure if the `TransformerBlocks` attention should contain a bias parameter.
+    """
+
+    _supports_gradient_checkpointing = True
+
+    @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,
+        double_self_attention: bool = False,
+        upcast_attention: bool = False,
+        norm_type: str = "layer_norm",
+        norm_elementwise_affine: bool = True,
+        norm_eps: float = 1e-5,
+        attention_type: str = "default",
+        caption_channels: int = None,
+    ):
+        super().__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
+
+        conv_cls = nn.Conv2d if USE_PEFT_BACKEND else LoRACompatibleConv
+        linear_cls = nn.Linear if USE_PEFT_BACKEND else LoRACompatibleLinear
+
+        # 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
+        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 = linear_cls(in_channels, inner_dim)
+        else:
+            self.proj_in = conv_cls(
+                in_channels, inner_dim, kernel_size=1, stride=1, padding=0
+            )
+
+        # 3. Define transformers blocks
+        self.transformer_blocks = nn.ModuleList(
+            [
+                BasicTransformerBlock(
+                    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,
+                    double_self_attention=double_self_attention,
+                    upcast_attention=upcast_attention,
+                    norm_type=norm_type,
+                    norm_elementwise_affine=norm_elementwise_affine,
+                    norm_eps=norm_eps,
+                    attention_type=attention_type,
+                )
+                for d in range(num_layers)
+            ]
+        )
+
+        # 4. Define output layers
+        self.out_channels = in_channels if out_channels is None else out_channels
+        # TODO: should use out_channels for continuous projections
+        if use_linear_projection:
+            self.proj_out = linear_cls(inner_dim, in_channels)
+        else:
+            self.proj_out = conv_cls(
+                inner_dim, in_channels, kernel_size=1, stride=1, padding=0
+            )
+
+        # 5. PixArt-Alpha blocks.
+        self.adaln_single = None
+        self.use_additional_conditions = False
+        if norm_type == "ada_norm_single":
+            self.use_additional_conditions = self.config.sample_size == 128
+            # TODO(Sayak, PVP) clean this, for now we use sample size to determine whether to use
+            # additional conditions until we find better name
+            self.adaln_single = AdaLayerNormSingle(
+                inner_dim, use_additional_conditions=self.use_additional_conditions
+            )
+
+        self.caption_projection = None
+        if caption_channels is not None:
+            self.caption_projection = CaptionProjection(
+                in_features=caption_channels, hidden_size=inner_dim
+            )
+
+        self.gradient_checkpointing = False
+
+    def _set_gradient_checkpointing(self, module, value=False):
+        if hasattr(module, "gradient_checkpointing"):
+            module.gradient_checkpointing = value
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: Optional[torch.Tensor] = None,
+        timestep: Optional[torch.LongTensor] = None,
+        added_cond_kwargs: Dict[str, torch.Tensor] = 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,
+        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`.
+            cross_attention_kwargs ( `Dict[str, Any]`, *optional*):
+                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
+                `self.processor` in
+                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
+            attention_mask ( `torch.Tensor`, *optional*):
+                An attention mask of shape `(batch, key_tokens)` is applied to `encoder_hidden_states`. If `1` the mask
+                is kept, otherwise if `0` it is discarded. Mask will be converted into a bias, which adds large
+                negative values to the attention scores corresponding to "discard" tokens.
+            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)
+
+        # Retrieve lora scale.
+        lora_scale = (
+            cross_attention_kwargs.get("scale", 1.0)
+            if cross_attention_kwargs is not None
+            else 1.0
+        )
+
+        # 1. Input
+        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, scale=lora_scale)
+                if not USE_PEFT_BACKEND
+                else 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, scale=lora_scale)
+                if not USE_PEFT_BACKEND
+                else self.proj_in(hidden_states)
+            )
+
+        # 2. Blocks
+        if self.caption_projection is not None:
+            batch_size = hidden_states.shape[0]
+            encoder_hidden_states = self.caption_projection(encoder_hidden_states)
+            encoder_hidden_states = encoder_hidden_states.view(
+                batch_size, -1, hidden_states.shape[-1]
+            )
+
+        ref_feature = hidden_states.reshape(batch, height, width, inner_dim)
+        for block in self.transformer_blocks:
+            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(block),
+                    hidden_states,
+                    attention_mask,
+                    encoder_hidden_states,
+                    encoder_attention_mask,
+                    timestep,
+                    cross_attention_kwargs,
+                    class_labels,
+                    **ckpt_kwargs,
+                )
+            else:
+                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,
+                )
+
+        # 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, scale=lora_scale)
+                    if not USE_PEFT_BACKEND
+                    else self.proj_out(hidden_states)
+                )
+            else:
+                hidden_states = (
+                    self.proj_out(hidden_states, scale=lora_scale)
+                    if not USE_PEFT_BACKEND
+                    else 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
+        if not return_dict:
+            return (output, ref_feature)
+
+        return Transformer2DModelOutput(sample=output, ref_feature=ref_feature)

src/models/transformer_3d.py

@@ -0,0 +1,169 @@
+from dataclasses import dataclass
+from typing import Optional, Dict
+
+import torch
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models import ModelMixin
+from diffusers.utils import BaseOutput
+from diffusers.utils.import_utils import is_xformers_available
+from einops import rearrange, repeat
+from torch import nn
+
+from .attention import TemporalBasicTransformerBlock, ResidualTemporalBasicTransformerBlock
+
+
+@dataclass
+class Transformer3DModelOutput(BaseOutput):
+    sample: torch.FloatTensor
+
+
+if is_xformers_available():
+    import xformers
+    import xformers.ops
+else:
+    xformers = None
+
+
+class Transformer3DModel(ModelMixin, ConfigMixin):
+    _supports_gradient_checkpointing = True
+
+    @register_to_config
+    def __init__(
+        self,
+        num_attention_heads: int = 16,
+        attention_head_dim: int = 88,
+        in_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,
+        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,
+        unet_use_cross_frame_attention=None,
+        unet_use_temporal_attention=None,
+    ):
+        super().__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
+
+        # Define input layers
+        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 = nn.Linear(in_channels, inner_dim)
+        else:
+            self.proj_in = nn.Conv2d(
+                in_channels, inner_dim, kernel_size=1, stride=1, padding=0
+            )
+
+        # Define transformers blocks
+        self.transformer_blocks = nn.ModuleList(
+            [
+                TemporalBasicTransformerBlock(
+                    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,
+                    unet_use_cross_frame_attention=unet_use_cross_frame_attention,
+                    unet_use_temporal_attention=unet_use_temporal_attention,
+                )
+                for d in range(num_layers)
+            ]
+        )
+
+        # 4. Define output layers
+        if use_linear_projection:
+            self.proj_out = nn.Linear(in_channels, inner_dim)
+        else:
+            self.proj_out = nn.Conv2d(
+                inner_dim, in_channels, kernel_size=1, stride=1, padding=0
+            )
+
+        self.gradient_checkpointing = False
+
+    def _set_gradient_checkpointing(self, module, value=False):
+        if hasattr(module, "gradient_checkpointing"):
+            module.gradient_checkpointing = value
+
+    def forward(
+        self,
+        hidden_states,
+        encoder_hidden_states=None,
+        timestep=None,
+        return_dict: bool = True,
+    ):
+        # Input
+        assert (
+            hidden_states.dim() == 5
+        ), f"Expected hidden_states to have ndim=5, but got ndim={hidden_states.dim()}."
+        video_length = hidden_states.shape[2]
+        hidden_states = rearrange(hidden_states, "b c f h w -> (b f) c h w")
+        if encoder_hidden_states.shape[0] != hidden_states.shape[0]:
+            encoder_hidden_states = repeat(
+                encoder_hidden_states, "b n c -> (b f) n c", f=video_length
+            )
+
+        batch, channel, height, weight = 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 * weight, inner_dim
+            )
+        else:
+            inner_dim = hidden_states.shape[1]
+            hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(
+                batch, height * weight, inner_dim
+            )
+            hidden_states = self.proj_in(hidden_states)
+
+        # Blocks
+        for i, block in enumerate(self.transformer_blocks):
+            hidden_states = block(
+                hidden_states,
+                encoder_hidden_states=encoder_hidden_states,
+                timestep=timestep,
+                video_length=video_length,
+            )
+
+        # Output
+        if not self.use_linear_projection:
+            hidden_states = (
+                hidden_states.reshape(batch, height, weight, 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, weight, inner_dim)
+                .permute(0, 3, 1, 2)
+                .contiguous()
+            )
+
+        output = hidden_states + residual
+
+        output = rearrange(output, "(b f) c h w -> b c f h w", f=video_length)
+        if not return_dict:
+            return (output,)
+
+        return Transformer3DModelOutput(sample=output)

src/models/unet_2d_blocks.py

@@ -0,0 +1,1074 @@
+# Adapted from https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/unet_2d_blocks.py
+from typing import Any, Dict, Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from diffusers.models.activations import get_activation
+from diffusers.models.attention_processor import Attention
+from diffusers.models.dual_transformer_2d import DualTransformer2DModel
+from diffusers.models.resnet import Downsample2D, ResnetBlock2D, Upsample2D
+from diffusers.utils import is_torch_version, logging
+from diffusers.utils.torch_utils import apply_freeu
+from torch import nn
+
+from .transformer_2d import Transformer2DModel
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+def get_down_block(
+    down_block_type: str,
+    num_layers: int,
+    in_channels: int,
+    out_channels: int,
+    temb_channels: int,
+    add_downsample: bool,
+    resnet_eps: float,
+    resnet_act_fn: str,
+    transformer_layers_per_block: int = 1,
+    num_attention_heads: Optional[int] = None,
+    resnet_groups: Optional[int] = None,
+    cross_attention_dim: Optional[int] = None,
+    downsample_padding: Optional[int] = None,
+    dual_cross_attention: bool = False,
+    use_linear_projection: bool = False,
+    only_cross_attention: bool = False,
+    upcast_attention: bool = False,
+    resnet_time_scale_shift: str = "default",
+    attention_type: str = "default",
+    resnet_skip_time_act: bool = False,
+    resnet_out_scale_factor: float = 1.0,
+    cross_attention_norm: Optional[str] = None,
+    attention_head_dim: Optional[int] = None,
+    downsample_type: Optional[str] = None,
+    dropout: float = 0.0,
+):
+    # 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_down_block`. Defaulting `attention_head_dim` to {num_attention_heads}."
+        )
+        attention_head_dim = num_attention_heads
+
+    down_block_type = (
+        down_block_type[7:]
+        if down_block_type.startswith("UNetRes")
+        else down_block_type
+    )
+    if down_block_type == "DownBlock2D":
+        return DownBlock2D(
+            num_layers=num_layers,
+            in_channels=in_channels,
+            out_channels=out_channels,
+            temb_channels=temb_channels,
+            dropout=dropout,
+            add_downsample=add_downsample,
+            resnet_eps=resnet_eps,
+            resnet_act_fn=resnet_act_fn,
+            resnet_groups=resnet_groups,
+            downsample_padding=downsample_padding,
+            resnet_time_scale_shift=resnet_time_scale_shift,
+        )
+    elif down_block_type == "CrossAttnDownBlock2D":
+        if cross_attention_dim is None:
+            raise ValueError(
+                "cross_attention_dim must be specified for CrossAttnDownBlock2D"
+            )
+        return CrossAttnDownBlock2D(
+            num_layers=num_layers,
+            transformer_layers_per_block=transformer_layers_per_block,
+            in_channels=in_channels,
+            out_channels=out_channels,
+            temb_channels=temb_channels,
+            dropout=dropout,
+            add_downsample=add_downsample,
+            resnet_eps=resnet_eps,
+            resnet_act_fn=resnet_act_fn,
+            resnet_groups=resnet_groups,
+            downsample_padding=downsample_padding,
+            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,
+            attention_type=attention_type,
+        )
+    raise ValueError(f"{down_block_type} does not exist.")
+
+
+def get_up_block(
+    up_block_type: str,
+    num_layers: int,
+    in_channels: int,
+    out_channels: int,
+    prev_output_channel: int,
+    temb_channels: int,
+    add_upsample: bool,
+    resnet_eps: float,
+    resnet_act_fn: str,
+    resolution_idx: Optional[int] = None,
+    transformer_layers_per_block: int = 1,
+    num_attention_heads: Optional[int] = None,
+    resnet_groups: Optional[int] = None,
+    cross_attention_dim: Optional[int] = None,
+    dual_cross_attention: bool = False,
+    use_linear_projection: bool = False,
+    only_cross_attention: bool = False,
+    upcast_attention: bool = False,
+    resnet_time_scale_shift: str = "default",
+    attention_type: str = "default",
+    resnet_skip_time_act: bool = False,
+    resnet_out_scale_factor: float = 1.0,
+    cross_attention_norm: Optional[str] = None,
+    attention_head_dim: Optional[int] = None,
+    upsample_type: Optional[str] = None,
+    dropout: float = 0.0,
+) -> nn.Module:
+    # 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 UpBlock2D(
+            num_layers=num_layers,
+            in_channels=in_channels,
+            out_channels=out_channels,
+            prev_output_channel=prev_output_channel,
+            temb_channels=temb_channels,
+            resolution_idx=resolution_idx,
+            dropout=dropout,
+            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 == "CrossAttnUpBlock2D":
+        if cross_attention_dim is None:
+            raise ValueError(
+                "cross_attention_dim must be specified for CrossAttnUpBlock2D"
+            )
+        return CrossAttnUpBlock2D(
+            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,
+            resolution_idx=resolution_idx,
+            dropout=dropout,
+            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,
+            attention_type=attention_type,
+        )
+
+    raise ValueError(f"{up_block_type} does not exist.")
+
+
+class AutoencoderTinyBlock(nn.Module):
+    """
+    Tiny Autoencoder block used in [`AutoencoderTiny`]. It is a mini residual module consisting of plain conv + ReLU
+    blocks.
+
+    Args:
+        in_channels (`int`): The number of input channels.
+        out_channels (`int`): The number of output channels.
+        act_fn (`str`):
+            ` The activation function to use. Supported values are `"swish"`, `"mish"`, `"gelu"`, and `"relu"`.
+
+    Returns:
+        `torch.FloatTensor`: A tensor with the same shape as the input tensor, but with the number of channels equal to
+        `out_channels`.
+    """
+
+    def __init__(self, in_channels: int, out_channels: int, act_fn: str):
+        super().__init__()
+        act_fn = get_activation(act_fn)
+        self.conv = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
+            act_fn,
+            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
+            act_fn,
+            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
+        )
+        self.skip = (
+            nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False)
+            if in_channels != out_channels
+            else nn.Identity()
+        )
+        self.fuse = nn.ReLU()
+
+    def forward(self, x: torch.FloatTensor) -> torch.FloatTensor:
+        return self.fuse(self.conv(x) + self.skip(x))
+
+
+class UNetMidBlock2D(nn.Module):
+    """
+    A 2D UNet mid-block [`UNetMidBlock2D`] with multiple residual blocks and optional attention blocks.
+
+    Args:
+        in_channels (`int`): The number of input channels.
+        temb_channels (`int`): The number of temporal embedding channels.
+        dropout (`float`, *optional*, defaults to 0.0): The dropout rate.
+        num_layers (`int`, *optional*, defaults to 1): The number of residual blocks.
+        resnet_eps (`float`, *optional*, 1e-6 ): The epsilon value for the resnet blocks.
+        resnet_time_scale_shift (`str`, *optional*, defaults to `default`):
+            The type of normalization to apply to the time embeddings. This can help to improve the performance of the
+            model on tasks with long-range temporal dependencies.
+        resnet_act_fn (`str`, *optional*, defaults to `swish`): The activation function for the resnet blocks.
+        resnet_groups (`int`, *optional*, defaults to 32):
+            The number of groups to use in the group normalization layers of the resnet blocks.
+        attn_groups (`Optional[int]`, *optional*, defaults to None): The number of groups for the attention blocks.
+        resnet_pre_norm (`bool`, *optional*, defaults to `True`):
+            Whether to use pre-normalization for the resnet blocks.
+        add_attention (`bool`, *optional*, defaults to `True`): Whether to add attention blocks.
+        attention_head_dim (`int`, *optional*, defaults to 1):
+            Dimension of a single attention head. The number of attention heads is determined based on this value and
+            the number of input channels.
+        output_scale_factor (`float`, *optional*, defaults to 1.0): The output scale factor.
+
+    Returns:
+        `torch.FloatTensor`: The output of the last residual block, which is a tensor of shape `(batch_size,
+        in_channels, height, width)`.
+
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        temb_channels: int,
+        dropout: float = 0.0,
+        num_layers: int = 1,
+        resnet_eps: float = 1e-6,
+        resnet_time_scale_shift: str = "default",  # default, spatial
+        resnet_act_fn: str = "swish",
+        resnet_groups: int = 32,
+        attn_groups: Optional[int] = None,
+        resnet_pre_norm: bool = True,
+        add_attention: bool = True,
+        attention_head_dim: int = 1,
+        output_scale_factor: float = 1.0,
+    ):
+        super().__init__()
+        resnet_groups = (
+            resnet_groups if resnet_groups is not None else min(in_channels // 4, 32)
+        )
+        self.add_attention = add_attention
+
+        if attn_groups is None:
+            attn_groups = (
+                resnet_groups if resnet_time_scale_shift == "default" else None
+            )
+
+        # there is always at least one resnet
+        resnets = [
+            ResnetBlock2D(
+                in_channels=in_channels,
+                out_channels=in_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,
+            )
+        ]
+        attentions = []
+
+        if attention_head_dim is None:
+            logger.warn(
+                f"It is not recommend to pass `attention_head_dim=None`. Defaulting `attention_head_dim` to `in_channels`: {in_channels}."
+            )
+            attention_head_dim = in_channels
+
+        for _ in range(num_layers):
+            if self.add_attention:
+                attentions.append(
+                    Attention(
+                        in_channels,
+                        heads=in_channels // attention_head_dim,
+                        dim_head=attention_head_dim,
+                        rescale_output_factor=output_scale_factor,
+                        eps=resnet_eps,
+                        norm_num_groups=attn_groups,
+                        spatial_norm_dim=temb_channels
+                        if resnet_time_scale_shift == "spatial"
+                        else None,
+                        residual_connection=True,
+                        bias=True,
+                        upcast_softmax=True,
+                        _from_deprecated_attn_block=True,
+                    )
+                )
+            else:
+                attentions.append(None)
+
+            resnets.append(
+                ResnetBlock2D(
+                    in_channels=in_channels,
+                    out_channels=in_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.attentions = nn.ModuleList(attentions)
+        self.resnets = nn.ModuleList(resnets)
+
+    def forward(
+        self, hidden_states: torch.FloatTensor, temb: Optional[torch.FloatTensor] = None
+    ) -> torch.FloatTensor:
+        hidden_states = self.resnets[0](hidden_states, temb)
+        for attn, resnet in zip(self.attentions, self.resnets[1:]):
+            if attn is not None:
+                hidden_states = attn(hidden_states, temb=temb)
+            hidden_states = resnet(hidden_states, temb)
+
+        return hidden_states
+
+
+class UNetMidBlock2DCrossAttn(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        temb_channels: int,
+        dropout: float = 0.0,
+        num_layers: int = 1,
+        transformer_layers_per_block: Union[int, Tuple[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: int = 1,
+        output_scale_factor: float = 1.0,
+        cross_attention_dim: int = 1280,
+        dual_cross_attention: bool = False,
+        use_linear_projection: bool = False,
+        upcast_attention: bool = False,
+        attention_type: str = "default",
+    ):
+        super().__init__()
+
+        self.has_cross_attention = True
+        self.num_attention_heads = num_attention_heads
+        resnet_groups = (
+            resnet_groups if resnet_groups is not None else min(in_channels // 4, 32)
+        )
+
+        # support for variable transformer layers per block
+        if isinstance(transformer_layers_per_block, int):
+            transformer_layers_per_block = [transformer_layers_per_block] * num_layers
+
+        # there is always at least one resnet
+        resnets = [
+            ResnetBlock2D(
+                in_channels=in_channels,
+                out_channels=in_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,
+            )
+        ]
+        attentions = []
+
+        for i in range(num_layers):
+            if not dual_cross_attention:
+                attentions.append(
+                    Transformer2DModel(
+                        num_attention_heads,
+                        in_channels // num_attention_heads,
+                        in_channels=in_channels,
+                        num_layers=transformer_layers_per_block[i],
+                        cross_attention_dim=cross_attention_dim,
+                        norm_num_groups=resnet_groups,
+                        use_linear_projection=use_linear_projection,
+                        upcast_attention=upcast_attention,
+                        attention_type=attention_type,
+                    )
+                )
+            else:
+                attentions.append(
+                    DualTransformer2DModel(
+                        num_attention_heads,
+                        in_channels // num_attention_heads,
+                        in_channels=in_channels,
+                        num_layers=1,
+                        cross_attention_dim=cross_attention_dim,
+                        norm_num_groups=resnet_groups,
+                    )
+                )
+            resnets.append(
+                ResnetBlock2D(
+                    in_channels=in_channels,
+                    out_channels=in_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.attentions = nn.ModuleList(attentions)
+        self.resnets = nn.ModuleList(resnets)
+
+        self.gradient_checkpointing = False
+
+    def forward(
+        self,
+        hidden_states: torch.FloatTensor,
+        temb: Optional[torch.FloatTensor] = None,
+        encoder_hidden_states: Optional[torch.FloatTensor] = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+        encoder_attention_mask: Optional[torch.FloatTensor] = None,
+    ) -> torch.FloatTensor:
+        lora_scale = (
+            cross_attention_kwargs.get("scale", 1.0)
+            if cross_attention_kwargs is not None
+            else 1.0
+        )
+        hidden_states = self.resnets[0](hidden_states, temb, scale=lora_scale)
+        for attn, resnet in zip(self.attentions, self.resnets[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, ref_feature = 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,
+                )
+                hidden_states = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(resnet),
+                    hidden_states,
+                    temb,
+                    **ckpt_kwargs,
+                )
+            else:
+                hidden_states, ref_feature = 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,
+                )
+                hidden_states = resnet(hidden_states, temb, scale=lora_scale)
+
+        return hidden_states
+
+
+class CrossAttnDownBlock2D(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        temb_channels: int,
+        dropout: float = 0.0,
+        num_layers: int = 1,
+        transformer_layers_per_block: Union[int, Tuple[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: int = 1,
+        cross_attention_dim: int = 1280,
+        output_scale_factor: float = 1.0,
+        downsample_padding: int = 1,
+        add_downsample: bool = True,
+        dual_cross_attention: bool = False,
+        use_linear_projection: bool = False,
+        only_cross_attention: bool = False,
+        upcast_attention: bool = False,
+        attention_type: str = "default",
+    ):
+        super().__init__()
+        resnets = []
+        attentions = []
+
+        self.has_cross_attention = True
+        self.num_attention_heads = num_attention_heads
+        if isinstance(transformer_layers_per_block, int):
+            transformer_layers_per_block = [transformer_layers_per_block] * num_layers
+
+        for i in range(num_layers):
+            in_channels = in_channels if i == 0 else out_channels
+            resnets.append(
+                ResnetBlock2D(
+                    in_channels=in_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(
+                    Transformer2DModel(
+                        num_attention_heads,
+                        out_channels // num_attention_heads,
+                        in_channels=out_channels,
+                        num_layers=transformer_layers_per_block[i],
+                        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,
+                        attention_type=attention_type,
+                    )
+                )
+            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_downsample:
+            self.downsamplers = nn.ModuleList(
+                [
+                    Downsample2D(
+                        out_channels,
+                        use_conv=True,
+                        out_channels=out_channels,
+                        padding=downsample_padding,
+                        name="op",
+                    )
+                ]
+            )
+        else:
+            self.downsamplers = None
+
+        self.gradient_checkpointing = False
+
+    def forward(
+        self,
+        hidden_states: torch.FloatTensor,
+        temb: Optional[torch.FloatTensor] = None,
+        encoder_hidden_states: Optional[torch.FloatTensor] = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+        encoder_attention_mask: Optional[torch.FloatTensor] = None,
+        additional_residuals: Optional[torch.FloatTensor] = None,
+    ) -> Tuple[torch.FloatTensor, Tuple[torch.FloatTensor, ...]]:
+        output_states = ()
+
+        lora_scale = (
+            cross_attention_kwargs.get("scale", 1.0)
+            if cross_attention_kwargs is not None
+            else 1.0
+        )
+
+        blocks = list(zip(self.resnets, self.attentions))
+
+        for i, (resnet, attn) in enumerate(blocks):
+            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,
+                    **ckpt_kwargs,
+                )
+                hidden_states, ref_feature = 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,
+                )
+            else:
+                hidden_states = resnet(hidden_states, temb, scale=lora_scale)
+                hidden_states, ref_feature = 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,
+                )
+
+            # apply additional residuals to the output of the last pair of resnet and attention blocks
+            if i == len(blocks) - 1 and additional_residuals is not None:
+                hidden_states = hidden_states + additional_residuals
+
+            output_states = output_states + (hidden_states,)
+
+        if self.downsamplers is not None:
+            for downsampler in self.downsamplers:
+                hidden_states = downsampler(hidden_states, scale=lora_scale)
+
+            output_states = output_states + (hidden_states,)
+
+        return hidden_states, output_states
+
+
+class DownBlock2D(nn.Module):
+    def __init__(
+        self,
+        in_channels: 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: float = 1.0,
+        add_downsample: bool = True,
+        downsample_padding: int = 1,
+    ):
+        super().__init__()
+        resnets = []
+
+        for i in range(num_layers):
+            in_channels = in_channels if i == 0 else out_channels
+            resnets.append(
+                ResnetBlock2D(
+                    in_channels=in_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_downsample:
+            self.downsamplers = nn.ModuleList(
+                [
+                    Downsample2D(
+                        out_channels,
+                        use_conv=True,
+                        out_channels=out_channels,
+                        padding=downsample_padding,
+                        name="op",
+                    )
+                ]
+            )
+        else:
+            self.downsamplers = None
+
+        self.gradient_checkpointing = False
+
+    def forward(
+        self,
+        hidden_states: torch.FloatTensor,
+        temb: Optional[torch.FloatTensor] = None,
+        scale: float = 1.0,
+    ) -> Tuple[torch.FloatTensor, Tuple[torch.FloatTensor, ...]]:
+        output_states = ()
+
+        for resnet in self.resnets:
+            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, scale=scale)
+
+            output_states = output_states + (hidden_states,)
+
+        if self.downsamplers is not None:
+            for downsampler in self.downsamplers:
+                hidden_states = downsampler(hidden_states, scale=scale)
+
+            output_states = output_states + (hidden_states,)
+
+        return hidden_states, output_states
+
+
+class CrossAttnUpBlock2D(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        prev_output_channel: int,
+        temb_channels: int,
+        resolution_idx: Optional[int] = None,
+        dropout: float = 0.0,
+        num_layers: int = 1,
+        transformer_layers_per_block: Union[int, Tuple[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: int = 1,
+        cross_attention_dim: int = 1280,
+        output_scale_factor: float = 1.0,
+        add_upsample: bool = True,
+        dual_cross_attention: bool = False,
+        use_linear_projection: bool = False,
+        only_cross_attention: bool = False,
+        upcast_attention: bool = False,
+        attention_type: str = "default",
+    ):
+        super().__init__()
+        resnets = []
+        attentions = []
+
+        self.has_cross_attention = True
+        self.num_attention_heads = num_attention_heads
+
+        if isinstance(transformer_layers_per_block, int):
+            transformer_layers_per_block = [transformer_layers_per_block] * num_layers
+
+        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(
+                ResnetBlock2D(
+                    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(
+                    Transformer2DModel(
+                        num_attention_heads,
+                        out_channels // num_attention_heads,
+                        in_channels=out_channels,
+                        num_layers=transformer_layers_per_block[i],
+                        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,
+                        attention_type=attention_type,
+                    )
+                )
+            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
+        self.resolution_idx = resolution_idx
+
+    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,
+    ) -> torch.FloatTensor:
+        lora_scale = (
+            cross_attention_kwargs.get("scale", 1.0)
+            if cross_attention_kwargs is not None
+            else 1.0
+        )
+        is_freeu_enabled = (
+            getattr(self, "s1", None)
+            and getattr(self, "s2", None)
+            and getattr(self, "b1", None)
+            and getattr(self, "b2", None)
+        )
+
+        for resnet, attn in 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]
+
+            # FreeU: Only operate on the first two stages
+            if is_freeu_enabled:
+                hidden_states, res_hidden_states = apply_freeu(
+                    self.resolution_idx,
+                    hidden_states,
+                    res_hidden_states,
+                    s1=self.s1,
+                    s2=self.s2,
+                    b1=self.b1,
+                    b2=self.b2,
+                )
+
+            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,
+                    **ckpt_kwargs,
+                )
+                hidden_states, ref_feature = 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,
+                )
+            else:
+                hidden_states = resnet(hidden_states, temb, scale=lora_scale)
+                hidden_states, ref_feature = 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,
+                )
+
+        if self.upsamplers is not None:
+            for upsampler in self.upsamplers:
+                hidden_states = upsampler(
+                    hidden_states, upsample_size, scale=lora_scale
+                )
+
+        return hidden_states
+
+
+class UpBlock2D(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        prev_output_channel: int,
+        out_channels: int,
+        temb_channels: int,
+        resolution_idx: Optional[int] = None,
+        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: float = 1.0,
+        add_upsample: bool = True,
+    ):
+        super().__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(
+                ResnetBlock2D(
+                    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
+        self.resolution_idx = resolution_idx
+
+    def forward(
+        self,
+        hidden_states: torch.FloatTensor,
+        res_hidden_states_tuple: Tuple[torch.FloatTensor, ...],
+        temb: Optional[torch.FloatTensor] = None,
+        upsample_size: Optional[int] = None,
+        scale: float = 1.0,
+    ) -> torch.FloatTensor:
+        is_freeu_enabled = (
+            getattr(self, "s1", None)
+            and getattr(self, "s2", None)
+            and getattr(self, "b1", None)
+            and getattr(self, "b2", None)
+        )
+
+        for resnet in self.resnets:
+            # pop res hidden states
+            res_hidden_states = res_hidden_states_tuple[-1]
+            res_hidden_states_tuple = res_hidden_states_tuple[:-1]
+
+            # FreeU: Only operate on the first two stages
+            if is_freeu_enabled:
+                hidden_states, res_hidden_states = apply_freeu(
+                    self.resolution_idx,
+                    hidden_states,
+                    res_hidden_states,
+                    s1=self.s1,
+                    s2=self.s2,
+                    b1=self.b1,
+                    b2=self.b2,
+                )
+
+            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, scale=scale)
+
+        if self.upsamplers is not None:
+            for upsampler in self.upsamplers:
+                hidden_states = upsampler(hidden_states, upsample_size, scale=scale)
+
+        return hidden_states

src/models/unet_2d_condition.py

@@ -0,0 +1,1308 @@
+# Adapted from https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/unet_2d_condition.py
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.utils.checkpoint
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.loaders import UNet2DConditionLoadersMixin
+from diffusers.models.activations import get_activation
+from diffusers.models.attention_processor import (
+    ADDED_KV_ATTENTION_PROCESSORS,
+    CROSS_ATTENTION_PROCESSORS,
+    AttentionProcessor,
+    AttnAddedKVProcessor,
+    AttnProcessor,
+)
+from diffusers.models.embeddings import (
+    GaussianFourierProjection,
+    ImageHintTimeEmbedding,
+    ImageProjection,
+    ImageTimeEmbedding,
+    PositionNet,
+    TextImageProjection,
+    TextImageTimeEmbedding,
+    TextTimeEmbedding,
+    TimestepEmbedding,
+    Timesteps,
+)
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.utils import (
+    USE_PEFT_BACKEND,
+    BaseOutput,
+    deprecate,
+    logging,
+    scale_lora_layers,
+    unscale_lora_layers,
+)
+
+from .unet_2d_blocks import (
+    UNetMidBlock2D,
+    UNetMidBlock2DCrossAttn,
+    get_down_block,
+    get_up_block,
+)
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+@dataclass
+class UNet2DConditionOutput(BaseOutput):
+    """
+    The output of [`UNet2DConditionModel`].
+
+    Args:
+        sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
+            The hidden states output conditioned on `encoder_hidden_states` input. Output of last layer of model.
+    """
+
+    sample: torch.FloatTensor = None
+    ref_features: Tuple[torch.FloatTensor] = None
+
+
+class UNet2DConditionModel(ModelMixin, ConfigMixin, UNet2DConditionLoadersMixin):
+    r"""
+    A conditional 2D UNet model that takes a noisy sample, conditional state, and a timestep and returns a sample
+    shaped output.
+
+    This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented
+    for all models (such as downloading or saving).
+
+    Parameters:
+        sample_size (`int` or `Tuple[int, int]`, *optional*, defaults to `None`):
+            Height and width of input/output sample.
+        in_channels (`int`, *optional*, defaults to 4): Number of channels in the input sample.
+        out_channels (`int`, *optional*, defaults to 4): Number of channels in the output.
+        center_input_sample (`bool`, *optional*, defaults to `False`): Whether to center the input sample.
+        flip_sin_to_cos (`bool`, *optional*, defaults to `False`):
+            Whether to flip the sin to cos in the time embedding.
+        freq_shift (`int`, *optional*, defaults to 0): The frequency shift to apply to the time embedding.
+        down_block_types (`Tuple[str]`, *optional*, defaults to `("CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "DownBlock2D")`):
+            The tuple of downsample blocks to use.
+        mid_block_type (`str`, *optional*, defaults to `"UNetMidBlock2DCrossAttn"`):
+            Block type for middle of UNet, it can be one of `UNetMidBlock2DCrossAttn`, `UNetMidBlock2D`, or
+            `UNetMidBlock2DSimpleCrossAttn`. If `None`, the mid block layer is skipped.
+        up_block_types (`Tuple[str]`, *optional*, defaults to `("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")`):
+            The tuple of upsample blocks to use.
+        only_cross_attention(`bool` or `Tuple[bool]`, *optional*, default to `False`):
+            Whether to include self-attention in the basic transformer blocks, see
+            [`~models.attention.BasicTransformerBlock`].
+        block_out_channels (`Tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`):
+            The tuple of output channels for each block.
+        layers_per_block (`int`, *optional*, defaults to 2): The number of layers per block.
+        downsample_padding (`int`, *optional*, defaults to 1): The padding to use for the downsampling convolution.
+        mid_block_scale_factor (`float`, *optional*, defaults to 1.0): The scale factor to use for the mid block.
+        dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
+        act_fn (`str`, *optional*, defaults to `"silu"`): The activation function to use.
+        norm_num_groups (`int`, *optional*, defaults to 32): The number of groups to use for the normalization.
+            If `None`, normalization and activation layers is skipped in post-processing.
+        norm_eps (`float`, *optional*, defaults to 1e-5): The epsilon to use for the normalization.
+        cross_attention_dim (`int` or `Tuple[int]`, *optional*, defaults to 1280):
+            The dimension of the cross attention features.
+        transformer_layers_per_block (`int`, `Tuple[int]`, or `Tuple[Tuple]` , *optional*, defaults to 1):
+            The number of transformer blocks of type [`~models.attention.BasicTransformerBlock`]. Only relevant for
+            [`~models.unet_2d_blocks.CrossAttnDownBlock2D`], [`~models.unet_2d_blocks.CrossAttnUpBlock2D`],
+            [`~models.unet_2d_blocks.UNetMidBlock2DCrossAttn`].
+       reverse_transformer_layers_per_block : (`Tuple[Tuple]`, *optional*, defaults to None):
+            The number of transformer blocks of type [`~models.attention.BasicTransformerBlock`], in the upsampling
+            blocks of the U-Net. Only relevant if `transformer_layers_per_block` is of type `Tuple[Tuple]` and for
+            [`~models.unet_2d_blocks.CrossAttnDownBlock2D`], [`~models.unet_2d_blocks.CrossAttnUpBlock2D`],
+            [`~models.unet_2d_blocks.UNetMidBlock2DCrossAttn`].
+        encoder_hid_dim (`int`, *optional*, defaults to None):
+            If `encoder_hid_dim_type` is defined, `encoder_hidden_states` will be projected from `encoder_hid_dim`
+            dimension to `cross_attention_dim`.
+        encoder_hid_dim_type (`str`, *optional*, defaults to `None`):
+            If given, the `encoder_hidden_states` and potentially other embeddings are down-projected to text
+            embeddings of dimension `cross_attention` according to `encoder_hid_dim_type`.
+        attention_head_dim (`int`, *optional*, defaults to 8): The dimension of the attention heads.
+        num_attention_heads (`int`, *optional*):
+            The number of attention heads. If not defined, defaults to `attention_head_dim`
+        resnet_time_scale_shift (`str`, *optional*, defaults to `"default"`): Time scale shift config
+            for ResNet blocks (see [`~models.resnet.ResnetBlock2D`]). Choose from `default` or `scale_shift`.
+        class_embed_type (`str`, *optional*, defaults to `None`):
+            The type of class embedding to use which is ultimately summed with the time embeddings. Choose from `None`,
+            `"timestep"`, `"identity"`, `"projection"`, or `"simple_projection"`.
+        addition_embed_type (`str`, *optional*, defaults to `None`):
+            Configures an optional embedding which will be summed with the time embeddings. Choose from `None` or
+            "text". "text" will use the `TextTimeEmbedding` layer.
+        addition_time_embed_dim: (`int`, *optional*, defaults to `None`):
+            Dimension for the timestep embeddings.
+        num_class_embeds (`int`, *optional*, defaults to `None`):
+            Input dimension of the learnable embedding matrix to be projected to `time_embed_dim`, when performing
+            class conditioning with `class_embed_type` equal to `None`.
+        time_embedding_type (`str`, *optional*, defaults to `positional`):
+            The type of position embedding to use for timesteps. Choose from `positional` or `fourier`.
+        time_embedding_dim (`int`, *optional*, defaults to `None`):
+            An optional override for the dimension of the projected time embedding.
+        time_embedding_act_fn (`str`, *optional*, defaults to `None`):
+            Optional activation function to use only once on the time embeddings before they are passed to the rest of
+            the UNet. Choose from `silu`, `mish`, `gelu`, and `swish`.
+        timestep_post_act (`str`, *optional*, defaults to `None`):
+            The second activation function to use in timestep embedding. Choose from `silu`, `mish` and `gelu`.
+        time_cond_proj_dim (`int`, *optional*, defaults to `None`):
+            The dimension of `cond_proj` layer in the timestep embedding.
+        conv_in_kernel (`int`, *optional*, default to `3`): The kernel size of `conv_in` layer. conv_out_kernel (`int`,
+        *optional*, default to `3`): The kernel size of `conv_out` layer. projection_class_embeddings_input_dim (`int`,
+        *optional*): The dimension of the `class_labels` input when
+            `class_embed_type="projection"`. Required when `class_embed_type="projection"`.
+        class_embeddings_concat (`bool`, *optional*, defaults to `False`): Whether to concatenate the time
+            embeddings with the class embeddings.
+        mid_block_only_cross_attention (`bool`, *optional*, defaults to `None`):
+            Whether to use cross attention with the mid block when using the `UNetMidBlock2DSimpleCrossAttn`. If
+            `only_cross_attention` is given as a single boolean and `mid_block_only_cross_attention` is `None`, the
+            `only_cross_attention` value is used as the value for `mid_block_only_cross_attention`. Default to `False`
+            otherwise.
+    """
+
+    _supports_gradient_checkpointing = True
+
+    @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,
+        dropout: float = 0.0,
+        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], Tuple[Tuple]] = 1,
+        reverse_transformer_layers_per_block: Optional[Tuple[Tuple[int]]] = None,
+        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,
+        attention_type: str = "default",
+        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().__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}."
+            )
+        if (
+            isinstance(transformer_layers_per_block, list)
+            and reverse_transformer_layers_per_block is None
+        ):
+            for layer_number_per_block in transformer_layers_per_block:
+                if isinstance(layer_number_per_block, list):
+                    raise ValueError(
+                        "Must provide 'reverse_transformer_layers_per_block` if using asymmetrical UNet."
+                    )
+
+        # 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,
+                attention_type=attention_type,
+                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,
+                dropout=dropout,
+            )
+            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,
+                dropout=dropout,
+                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,
+                attention_type=attention_type,
+            )
+        elif mid_block_type == "UNetMidBlock2DSimpleCrossAttn":
+            raise NotImplementedError(f"Unsupport mid_block_type: {mid_block_type}")
+        elif mid_block_type == "UNetMidBlock2D":
+            self.mid_block = UNetMidBlock2D(
+                in_channels=block_out_channels[-1],
+                temb_channels=blocks_time_embed_dim,
+                dropout=dropout,
+                num_layers=0,
+                resnet_eps=norm_eps,
+                resnet_act_fn=act_fn,
+                output_scale_factor=mid_block_scale_factor,
+                resnet_groups=norm_num_groups,
+                resnet_time_scale_shift=resnet_time_scale_shift,
+                add_attention=False,
+            )
+        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))
+            if reverse_transformer_layers_per_block is None
+            else reverse_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_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,
+                resolution_idx=i,
+                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,
+                attention_type=attention_type,
+                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,
+                dropout=dropout,
+            )
+            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
+        self.conv_norm_out = 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,
+        # )
+
+        if attention_type in ["gated", "gated-text-image"]:
+            positive_len = 768
+            if isinstance(cross_attention_dim, int):
+                positive_len = cross_attention_dim
+            elif isinstance(cross_attention_dim, tuple) or isinstance(
+                cross_attention_dim, list
+            ):
+                positive_len = cross_attention_dim[0]
+
+            feature_type = "text-only" if attention_type == "gated" else "text-image"
+            self.position_net = PositionNet(
+                positive_len=positive_len,
+                out_dim=cross_attention_dim,
+                feature_type=feature_type,
+            )
+
+    @property
+    def attn_processors(self) -> Dict[str, AttentionProcessor]:
+        r"""
+        Returns:
+            `dict` of attention processors: A dictionary containing all attention processors used in the model with
+            indexed by its weight name.
+        """
+        # set recursively
+        processors = {}
+
+        def fn_recursive_add_processors(
+            name: str,
+            module: torch.nn.Module,
+            processors: Dict[str, AttentionProcessor],
+        ):
+            if hasattr(module, "get_processor"):
+                processors[f"{name}.processor"] = module.get_processor(
+                    return_deprecated_lora=True
+                )
+
+            for sub_name, child in module.named_children():
+                fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
+
+            return processors
+
+        for name, module in self.named_children():
+            fn_recursive_add_processors(name, module, processors)
+
+        return processors
+
+    def set_attn_processor(
+        self,
+        processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]],
+        _remove_lora=False,
+    ):
+        r"""
+        Sets the attention processor to use to compute attention.
+
+        Parameters:
+            processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`):
+                The instantiated processor class or a dictionary of processor classes that will be set as the processor
+                for **all** `Attention` layers.
+
+                If `processor` is a dict, the key needs to define the path to the corresponding cross attention
+                processor. This is strongly recommended when setting trainable attention processors.
+
+        """
+        count = len(self.attn_processors.keys())
+
+        if isinstance(processor, dict) and len(processor) != count:
+            raise ValueError(
+                f"A dict of processors was passed, but the number of processors {len(processor)} does not match the"
+                f" number of attention layers: {count}. Please make sure to pass {count} processor classes."
+            )
+
+        def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):
+            if hasattr(module, "set_processor"):
+                if not isinstance(processor, dict):
+                    module.set_processor(processor, _remove_lora=_remove_lora)
+                else:
+                    module.set_processor(
+                        processor.pop(f"{name}.processor"), _remove_lora=_remove_lora
+                    )
+
+            for sub_name, child in module.named_children():
+                fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
+
+        for name, module in self.named_children():
+            fn_recursive_attn_processor(name, module, processor)
+
+    def set_default_attn_processor(self):
+        """
+        Disables custom attention processors and sets the default attention implementation.
+        """
+        if all(
+            proc.__class__ in ADDED_KV_ATTENTION_PROCESSORS
+            for proc in self.attn_processors.values()
+        ):
+            processor = AttnAddedKVProcessor()
+        elif all(
+            proc.__class__ in CROSS_ATTENTION_PROCESSORS
+            for proc in self.attn_processors.values()
+        ):
+            processor = AttnProcessor()
+        else:
+            raise ValueError(
+                f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}"
+            )
+
+        self.set_attn_processor(processor, _remove_lora=True)
+
+    def set_attention_slice(self, slice_size):
+        r"""
+        Enable sliced attention computation.
+
+        When this option is enabled, the attention module splits the input tensor in slices to compute attention in
+        several steps. This is useful for saving some memory in exchange for a small decrease in speed.
+
+        Args:
+            slice_size (`str` or `int` or `list(int)`, *optional*, defaults to `"auto"`):
+                When `"auto"`, input to the attention heads is halved, so attention is computed in two steps. If
+                `"max"`, maximum amount of memory is saved by running only one slice at a time. If a number is
+                provided, uses as many slices as `attention_head_dim // slice_size`. In this case, `attention_head_dim`
+                must be a multiple of `slice_size`.
+        """
+        sliceable_head_dims = []
+
+        def fn_recursive_retrieve_sliceable_dims(module: torch.nn.Module):
+            if hasattr(module, "set_attention_slice"):
+                sliceable_head_dims.append(module.sliceable_head_dim)
+
+            for child in module.children():
+                fn_recursive_retrieve_sliceable_dims(child)
+
+        # retrieve number of attention layers
+        for module in self.children():
+            fn_recursive_retrieve_sliceable_dims(module)
+
+        num_sliceable_layers = len(sliceable_head_dims)
+
+        if slice_size == "auto":
+            # half the attention head size is usually a good trade-off between
+            # speed and memory
+            slice_size = [dim // 2 for dim in sliceable_head_dims]
+        elif slice_size == "max":
+            # make smallest slice possible
+            slice_size = num_sliceable_layers * [1]
+
+        slice_size = (
+            num_sliceable_layers * [slice_size]
+            if not isinstance(slice_size, list)
+            else slice_size
+        )
+
+        if len(slice_size) != len(sliceable_head_dims):
+            raise ValueError(
+                f"You have provided {len(slice_size)}, but {self.config} has {len(sliceable_head_dims)} different"
+                f" attention layers. Make sure to match `len(slice_size)` to be {len(sliceable_head_dims)}."
+            )
+
+        for i in range(len(slice_size)):
+            size = slice_size[i]
+            dim = sliceable_head_dims[i]
+            if size is not None and size > dim:
+                raise ValueError(f"size {size} has to be smaller or equal to {dim}.")
+
+        # Recursively walk through all the children.
+        # Any children which exposes the set_attention_slice method
+        # gets the message
+        def fn_recursive_set_attention_slice(
+            module: torch.nn.Module, slice_size: List[int]
+        ):
+            if hasattr(module, "set_attention_slice"):
+                module.set_attention_slice(slice_size.pop())
+
+            for child in module.children():
+                fn_recursive_set_attention_slice(child, slice_size)
+
+        reversed_slice_size = list(reversed(slice_size))
+        for module in self.children():
+            fn_recursive_set_attention_slice(module, reversed_slice_size)
+
+    def _set_gradient_checkpointing(self, module, value=False):
+        if hasattr(module, "gradient_checkpointing"):
+            module.gradient_checkpointing = value
+
+    def enable_freeu(self, s1, s2, b1, b2):
+        r"""Enables the FreeU mechanism from https://arxiv.org/abs/2309.11497.
+
+        The suffixes after the scaling factors represent the stage blocks where they are being applied.
+
+        Please refer to the [official repository](https://github.com/ChenyangSi/FreeU) for combinations of values that
+        are known to work well for different pipelines such as Stable Diffusion v1, v2, and Stable Diffusion XL.
+
+        Args:
+            s1 (`float`):
+                Scaling factor for stage 1 to attenuate the contributions of the skip features. This is done to
+                mitigate the "oversmoothing effect" in the enhanced denoising process.
+            s2 (`float`):
+                Scaling factor for stage 2 to attenuate the contributions of the skip features. This is done to
+                mitigate the "oversmoothing effect" in the enhanced denoising process.
+            b1 (`float`): Scaling factor for stage 1 to amplify the contributions of backbone features.
+            b2 (`float`): Scaling factor for stage 2 to amplify the contributions of backbone features.
+        """
+        for i, upsample_block in enumerate(self.up_blocks):
+            setattr(upsample_block, "s1", s1)
+            setattr(upsample_block, "s2", s2)
+            setattr(upsample_block, "b1", b1)
+            setattr(upsample_block, "b2", b2)
+
+    def disable_freeu(self):
+        """Disables the FreeU mechanism."""
+        freeu_keys = {"s1", "s2", "b1", "b2"}
+        for i, upsample_block in enumerate(self.up_blocks):
+            for k in freeu_keys:
+                if (
+                    hasattr(upsample_block, k)
+                    or getattr(upsample_block, k, None) is not None
+                ):
+                    setattr(upsample_block, k, None)
+
+    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,
+        down_intrablock_additional_residuals: Optional[Tuple[torch.Tensor]] = None,
+        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)`.
+            class_labels (`torch.Tensor`, *optional*, defaults to `None`):
+                Optional class labels for conditioning. Their embeddings will be summed with the timestep embeddings.
+            timestep_cond: (`torch.Tensor`, *optional*, defaults to `None`):
+                Conditional embeddings for timestep. If provided, the embeddings will be summed with the samples passed
+                through the `self.time_embedding` layer to obtain the timestep embeddings.
+            attention_mask (`torch.Tensor`, *optional*, defaults to `None`):
+                An attention mask of shape `(batch, key_tokens)` is applied to `encoder_hidden_states`. If `1` the mask
+                is kept, otherwise if `0` it is discarded. Mask will be converted into a bias, which adds large
+                negative values to the attention scores corresponding to "discard" tokens.
+            cross_attention_kwargs (`dict`, *optional*):
+                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
+                `self.processor` in
+                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
+            added_cond_kwargs: (`dict`, *optional*):
+                A kwargs dictionary containing additional embeddings that if specified are added to the embeddings that
+                are passed along to the UNet blocks.
+            down_block_additional_residuals: (`tuple` of `torch.Tensor`, *optional*):
+                A tuple of tensors that if specified are added to the residuals of down unet blocks.
+            mid_block_additional_residual: (`torch.Tensor`, *optional*):
+                A tensor that if specified is added to the residual of the middle unet block.
+            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.
+            down_block_additional_residuals (`tuple` of `torch.Tensor`, *optional*):
+                additional residuals to be added to UNet long skip connections from down blocks to up blocks for
+                example from ControlNet side model(s)
+            mid_block_additional_residual (`torch.Tensor`, *optional*):
+                additional residual to be added to UNet mid block output, for example from ControlNet side model
+            down_intrablock_additional_residuals (`tuple` of `torch.Tensor`, *optional*):
+                additional residuals to be added within UNet down blocks, for example from T2I-Adapter side model(s)
+
+        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
+
+        for dim in sample.shape[-2:]:
+            if dim % default_overall_up_factor != 0:
+                # Forward upsample size to force interpolation output size.
+                forward_upsample_size = True
+                break
+
+        # 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)
+        elif (
+            self.encoder_hid_proj is not None
+            and self.config.encoder_hid_dim_type == "ip_image_proj"
+        ):
+            if "image_embeds" not in added_cond_kwargs:
+                raise ValueError(
+                    f"{self.__class__} has the config param `encoder_hid_dim_type` set to 'ip_image_proj' which requires the keyword argument `image_embeds` to be passed in  `added_conditions`"
+                )
+            image_embeds = added_cond_kwargs.get("image_embeds")
+            image_embeds = self.encoder_hid_proj(image_embeds).to(
+                encoder_hidden_states.dtype
+            )
+            encoder_hidden_states = torch.cat(
+                [encoder_hidden_states, image_embeds], dim=1
+            )
+
+        # 2. pre-process
+        sample = self.conv_in(sample)
+
+        # 2.5 GLIGEN position net
+        if (
+            cross_attention_kwargs is not None
+            and cross_attention_kwargs.get("gligen", None) is not None
+        ):
+            cross_attention_kwargs = cross_attention_kwargs.copy()
+            gligen_args = cross_attention_kwargs.pop("gligen")
+            cross_attention_kwargs["gligen"] = {
+                "objs": self.position_net(**gligen_args)
+            }
+
+        # 3. down
+        lora_scale = (
+            cross_attention_kwargs.get("scale", 1.0)
+            if cross_attention_kwargs is not None
+            else 1.0
+        )
+        if USE_PEFT_BACKEND:
+            # weight the lora layers by setting `lora_scale` for each PEFT layer
+            scale_lora_layers(self, lora_scale)
+
+        is_controlnet = (
+            mid_block_additional_residual is not None
+            and down_block_additional_residuals is not None
+        )
+        # using new arg down_intrablock_additional_residuals for T2I-Adapters, to distinguish from controlnets
+        is_adapter = down_intrablock_additional_residuals is not None
+        # maintain backward compatibility for legacy usage, where
+        #       T2I-Adapter and ControlNet both use down_block_additional_residuals arg
+        #       but can only use one or the other
+        if (
+            not is_adapter
+            and mid_block_additional_residual is None
+            and down_block_additional_residuals is not None
+        ):
+            deprecate(
+                "T2I should not use down_block_additional_residuals",
+                "1.3.0",
+                "Passing intrablock residual connections with `down_block_additional_residuals` is deprecated \
+                       and will be removed in diffusers 1.3.0.  `down_block_additional_residuals` should only be used \
+                       for ControlNet. Please make sure use `down_intrablock_additional_residuals` instead. ",
+                standard_warn=False,
+            )
+            down_intrablock_additional_residuals = down_block_additional_residuals
+            is_adapter = True
+
+        down_block_res_samples = (sample,)
+        tot_referece_features = ()
+        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_intrablock_additional_residuals) > 0:
+                    additional_residuals[
+                        "additional_residuals"
+                    ] = down_intrablock_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, scale=lora_scale
+                )
+                if is_adapter and len(down_intrablock_additional_residuals) > 0:
+                    sample += down_intrablock_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:
+            if (
+                hasattr(self.mid_block, "has_cross_attention")
+                and self.mid_block.has_cross_attention
+            ):
+                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,
+                )
+            else:
+                sample = self.mid_block(sample, emb)
+
+            # To support T2I-Adapter-XL
+            if (
+                is_adapter
+                and len(down_intrablock_additional_residuals) > 0
+                and sample.shape == down_intrablock_additional_residuals[0].shape
+            ):
+                sample += down_intrablock_additional_residuals.pop(0)
+
+        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,
+                )
+            else:
+                sample = upsample_block(
+                    hidden_states=sample,
+                    temb=emb,
+                    res_hidden_states_tuple=res_samples,
+                    upsample_size=upsample_size,
+                    scale=lora_scale,
+                )
+
+        # 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 USE_PEFT_BACKEND:
+            # remove `lora_scale` from each PEFT layer
+            unscale_lora_layers(self, lora_scale)
+
+        if not return_dict:
+            return (sample,)
+
+        return UNet2DConditionOutput(sample=sample)

src/models/unet_3d.py

@@ -0,0 +1,673 @@
+# Adapted from https://github.com/guoyww/AnimateDiff/blob/main/animatediff/models/unet_blocks.py
+
+from collections import OrderedDict
+from dataclasses import dataclass
+import pdb
+from os import PathLike
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.utils.checkpoint
+import torch.nn.functional as F
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.attention_processor import AttentionProcessor
+from diffusers.models.embeddings import TimestepEmbedding, Timesteps
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.utils import SAFETENSORS_WEIGHTS_NAME, WEIGHTS_NAME, BaseOutput, logging
+from safetensors.torch import load_file
+
+from .resnet import InflatedConv3d, InflatedGroupNorm
+from .unet_3d_blocks import UNetMidBlock3DCrossAttn, get_down_block, get_up_block
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+@dataclass
+class UNet3DConditionOutput(BaseOutput):
+    sample: torch.FloatTensor
+
+
+class UNet3DConditionModel(ModelMixin, ConfigMixin):
+    _supports_gradient_checkpointing = True
+
+    @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] = (
+            "CrossAttnDownBlock3D",
+            "CrossAttnDownBlock3D",
+            "CrossAttnDownBlock3D",
+            "DownBlock3D",
+        ),
+        mid_block_type: str = "UNetMidBlock3DCrossAttn",
+        up_block_types: Tuple[str] = (
+            "UpBlock3D",
+            "CrossAttnUpBlock3D",
+            "CrossAttnUpBlock3D",
+            "CrossAttnUpBlock3D",
+        ),
+        only_cross_attention: Union[bool, Tuple[bool]] = False,
+        block_out_channels: Tuple[int] = (320, 640, 1280, 1280),
+        layers_per_block: int = 2,
+        downsample_padding: int = 1,
+        mid_block_scale_factor: float = 1,
+        act_fn: str = "silu",
+        norm_num_groups: int = 32,
+        norm_eps: float = 1e-5,
+        cross_attention_dim: int = 1280,
+        attention_head_dim: Union[int, Tuple[int]] = 8,
+        dual_cross_attention: bool = False,
+        use_linear_projection: bool = False,
+        class_embed_type: Optional[str] = None,
+        num_class_embeds: Optional[int] = None,
+        upcast_attention: bool = False,
+        resnet_time_scale_shift: str = "default",
+        use_inflated_groupnorm=False,
+        # Additional
+        use_motion_module=False,
+        motion_module_resolutions=(1, 2, 4, 8),
+        motion_module_mid_block=False,
+        motion_module_decoder_only=False,
+        motion_module_type=None,
+        motion_module_kwargs={},
+        unet_use_cross_frame_attention=None,
+        unet_use_temporal_attention=None,
+    ):
+        super().__init__()
+
+        self.sample_size = sample_size
+        time_embed_dim = block_out_channels[0] * 4
+
+        # input
+        self.conv_in = InflatedConv3d(
+            in_channels, block_out_channels[0], kernel_size=3, padding=(1, 1)
+        )
+
+        # time
+        self.time_proj = Timesteps(block_out_channels[0], flip_sin_to_cos, freq_shift)
+        timestep_input_dim = block_out_channels[0]
+
+        self.time_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim)
+
+        # 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)
+        elif class_embed_type == "identity":
+            self.class_embedding = nn.Identity(time_embed_dim, time_embed_dim)
+        else:
+            self.class_embedding = None
+
+        self.down_blocks = nn.ModuleList([])
+        self.mid_block = None
+        self.up_blocks = nn.ModuleList([])
+
+        if isinstance(only_cross_attention, bool):
+            only_cross_attention = [only_cross_attention] * len(down_block_types)
+
+        if isinstance(attention_head_dim, int):
+            attention_head_dim = (attention_head_dim,) * len(down_block_types)
+
+        # down
+        output_channel = block_out_channels[0]
+        for i, down_block_type in enumerate(down_block_types):
+            res = 2**i
+            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,
+                in_channels=input_channel,
+                out_channels=output_channel,
+                temb_channels=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,
+                attn_num_head_channels=attention_head_dim[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,
+                unet_use_cross_frame_attention=unet_use_cross_frame_attention,
+                unet_use_temporal_attention=unet_use_temporal_attention,
+                use_inflated_groupnorm=use_inflated_groupnorm,
+                use_motion_module=use_motion_module
+                and (res in motion_module_resolutions)
+                and (not motion_module_decoder_only),
+                motion_module_type=motion_module_type,
+                motion_module_kwargs=motion_module_kwargs,
+            )
+            self.down_blocks.append(down_block)
+
+        # mid
+        if mid_block_type == "UNetMidBlock3DCrossAttn":
+            self.mid_block = UNetMidBlock3DCrossAttn(
+                in_channels=block_out_channels[-1],
+                temb_channels=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,
+                attn_num_head_channels=attention_head_dim[-1],
+                resnet_groups=norm_num_groups,
+                dual_cross_attention=dual_cross_attention,
+                use_linear_projection=use_linear_projection,
+                upcast_attention=upcast_attention,
+                unet_use_cross_frame_attention=unet_use_cross_frame_attention,
+                unet_use_temporal_attention=unet_use_temporal_attention,
+                use_inflated_groupnorm=use_inflated_groupnorm,
+                use_motion_module=use_motion_module and motion_module_mid_block,
+                motion_module_type=motion_module_type,
+                motion_module_kwargs=motion_module_kwargs,
+            )
+        else:
+            raise ValueError(f"unknown mid_block_type : {mid_block_type}")
+
+        # count how many layers upsample the videos
+        self.num_upsamplers = 0
+
+        # up
+        reversed_block_out_channels = list(reversed(block_out_channels))
+        reversed_attention_head_dim = list(reversed(attention_head_dim))
+        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):
+            res = 2 ** (3 - i)
+            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_up_block(
+                up_block_type,
+                num_layers=layers_per_block + 1,
+                in_channels=input_channel,
+                out_channels=output_channel,
+                prev_output_channel=prev_output_channel,
+                temb_channels=time_embed_dim,
+                add_upsample=add_upsample,
+                resnet_eps=norm_eps,
+                resnet_act_fn=act_fn,
+                resnet_groups=norm_num_groups,
+                cross_attention_dim=cross_attention_dim,
+                attn_num_head_channels=reversed_attention_head_dim[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,
+                unet_use_cross_frame_attention=unet_use_cross_frame_attention,
+                unet_use_temporal_attention=unet_use_temporal_attention,
+                use_inflated_groupnorm=use_inflated_groupnorm,
+                use_motion_module=use_motion_module
+                and (res in motion_module_resolutions),
+                motion_module_type=motion_module_type,
+                motion_module_kwargs=motion_module_kwargs,
+            )
+            self.up_blocks.append(up_block)
+            prev_output_channel = output_channel
+
+        # out
+        if use_inflated_groupnorm:
+            self.conv_norm_out = InflatedGroupNorm(
+                num_channels=block_out_channels[0],
+                num_groups=norm_num_groups,
+                eps=norm_eps,
+            )
+        else:
+            self.conv_norm_out = nn.GroupNorm(
+                num_channels=block_out_channels[0],
+                num_groups=norm_num_groups,
+                eps=norm_eps,
+            )
+        self.conv_act = nn.SiLU()
+        self.conv_out = InflatedConv3d(
+            block_out_channels[0], out_channels, kernel_size=3, padding=1
+        )
+
+    @property
+    # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.attn_processors
+    def attn_processors(self) -> Dict[str, AttentionProcessor]:
+        r"""
+        Returns:
+            `dict` of attention processors: A dictionary containing all attention processors used in the model with
+            indexed by its weight name.
+        """
+        # set recursively
+        processors = {}
+
+        def fn_recursive_add_processors(
+            name: str,
+            module: torch.nn.Module,
+            processors: Dict[str, AttentionProcessor],
+        ):
+            if hasattr(module, "set_processor"):
+                processors[f"{name}.processor"] = module.processor
+
+            for sub_name, child in module.named_children():
+                if "temporal_transformer" not in sub_name:
+                    fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
+
+            return processors
+
+        for name, module in self.named_children():
+            if "temporal_transformer" not in name:
+                fn_recursive_add_processors(name, module, processors)
+
+        return processors
+
+    def set_attention_slice(self, slice_size):
+        r"""
+        Enable sliced attention computation.
+
+        When this option is enabled, the attention module will split the input tensor in slices, to compute attention
+        in several steps. This is useful to save some memory in exchange for a small speed decrease.
+
+        Args:
+            slice_size (`str` or `int` or `list(int)`, *optional*, defaults to `"auto"`):
+                When `"auto"`, halves the input to the attention heads, so attention will be computed in two steps. If
+                `"max"`, maxium amount of memory will be saved by running only one slice at a time. If a number is
+                provided, uses as many slices as `attention_head_dim // slice_size`. In this case, `attention_head_dim`
+                must be a multiple of `slice_size`.
+        """
+        sliceable_head_dims = []
+
+        def fn_recursive_retrieve_slicable_dims(module: torch.nn.Module):
+            if hasattr(module, "set_attention_slice"):
+                sliceable_head_dims.append(module.sliceable_head_dim)
+
+            for child in module.children():
+                fn_recursive_retrieve_slicable_dims(child)
+
+        # retrieve number of attention layers
+        for module in self.children():
+            fn_recursive_retrieve_slicable_dims(module)
+
+        num_slicable_layers = len(sliceable_head_dims)
+
+        if slice_size == "auto":
+            # half the attention head size is usually a good trade-off between
+            # speed and memory
+            slice_size = [dim // 2 for dim in sliceable_head_dims]
+        elif slice_size == "max":
+            # make smallest slice possible
+            slice_size = num_slicable_layers * [1]
+
+        slice_size = (
+            num_slicable_layers * [slice_size]
+            if not isinstance(slice_size, list)
+            else slice_size
+        )
+
+        if len(slice_size) != len(sliceable_head_dims):
+            raise ValueError(
+                f"You have provided {len(slice_size)}, but {self.config} has {len(sliceable_head_dims)} different"
+                f" attention layers. Make sure to match `len(slice_size)` to be {len(sliceable_head_dims)}."
+            )
+
+        for i in range(len(slice_size)):
+            size = slice_size[i]
+            dim = sliceable_head_dims[i]
+            if size is not None and size > dim:
+                raise ValueError(f"size {size} has to be smaller or equal to {dim}.")
+
+        # Recursively walk through all the children.
+        # Any children which exposes the set_attention_slice method
+        # gets the message
+        def fn_recursive_set_attention_slice(
+            module: torch.nn.Module, slice_size: List[int]
+        ):
+            if hasattr(module, "set_attention_slice"):
+                module.set_attention_slice(slice_size.pop())
+
+            for child in module.children():
+                fn_recursive_set_attention_slice(child, slice_size)
+
+        reversed_slice_size = list(reversed(slice_size))
+        for module in self.children():
+            fn_recursive_set_attention_slice(module, reversed_slice_size)
+
+    def _set_gradient_checkpointing(self, module, value=False):
+        if hasattr(module, "gradient_checkpointing"):
+            module.gradient_checkpointing = value
+
+    # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.set_attn_processor
+    def set_attn_processor(
+        self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]
+    ):
+        r"""
+        Sets the attention processor to use to compute attention.
+
+        Parameters:
+            processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`):
+                The instantiated processor class or a dictionary of processor classes that will be set as the processor
+                for **all** `Attention` layers.
+
+                If `processor` is a dict, the key needs to define the path to the corresponding cross attention
+                processor. This is strongly recommended when setting trainable attention processors.
+
+        """
+        count = len(self.attn_processors.keys())
+
+        if isinstance(processor, dict) and len(processor) != count:
+            raise ValueError(
+                f"A dict of processors was passed, but the number of processors {len(processor)} does not match the"
+                f" number of attention layers: {count}. Please make sure to pass {count} processor classes."
+            )
+
+        def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):
+            if hasattr(module, "set_processor"):
+                if not isinstance(processor, dict):
+                    module.set_processor(processor)
+                else:
+                    module.set_processor(processor.pop(f"{name}.processor"))
+
+            for sub_name, child in module.named_children():
+                if "temporal_transformer" not in sub_name:
+                    fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
+
+        for name, module in self.named_children():
+            if "temporal_transformer" not in name:
+                fn_recursive_attn_processor(name, module, processor)
+
+    def forward(
+        self,
+        sample: torch.FloatTensor,
+        timestep: Union[torch.Tensor, float, int],
+        encoder_hidden_states: torch.Tensor,
+        class_labels: Optional[torch.Tensor] = None,
+        pose_cond_fea = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        down_block_additional_residuals: Optional[Tuple[torch.Tensor]] = None,
+        mid_block_additional_residual: Optional[torch.Tensor] = None,
+        return_dict: bool = True,
+    ) -> Union[UNet3DConditionOutput, Tuple]:
+        r"""
+        Args:
+            sample (`torch.FloatTensor`): (batch, channel, height, width) noisy inputs tensor
+            timestep (`torch.FloatTensor` or `float` or `int`): (batch) timesteps
+            encoder_hidden_states (`torch.FloatTensor`): (batch, sequence_length, feature_dim) encoder hidden states
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`models.unet_2d_condition.UNet2DConditionOutput`] instead of a plain tuple.
+
+        Returns:
+            [`~models.unet_2d_condition.UNet2DConditionOutput`] or `tuple`:
+            [`~models.unet_2d_condition.UNet2DConditionOutput`] if `return_dict` is True, otherwise a `tuple`. When
+            returning a tuple, 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 layears).
+        # 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
+
+        # prepare attention_mask
+        if attention_mask is not None:
+            attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0
+            attention_mask = attention_mask.unsqueeze(1)
+
+        # center input if necessary
+        if self.config.center_input_sample:
+            sample = 2 * sample - 1.0
+
+        # time
+        timesteps = timestep
+        if not torch.is_tensor(timesteps):
+            # 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=self.dtype)
+        emb = self.time_embedding(t_emb)
+
+        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)
+
+            class_emb = self.class_embedding(class_labels).to(dtype=self.dtype)
+            emb = emb + class_emb
+
+        # pre-process
+        sample = self.conv_in(sample)
+        if pose_cond_fea is not None:
+            sample = sample + pose_cond_fea[0]
+
+        # down
+        down_block_res_samples = (sample,)
+        block_count = 1
+        for downsample_block in self.down_blocks:
+            if (
+                hasattr(downsample_block, "has_cross_attention")
+                and downsample_block.has_cross_attention
+            ):
+                sample, res_samples = downsample_block(
+                    hidden_states=sample,
+                    temb=emb,
+                    encoder_hidden_states=encoder_hidden_states,
+                    attention_mask=attention_mask,
+                )
+            else:
+                sample, res_samples = downsample_block(
+                    hidden_states=sample,
+                    temb=emb,
+                    encoder_hidden_states=encoder_hidden_states,
+                )
+            if pose_cond_fea is not None:
+                sample = sample + pose_cond_fea[block_count]
+                block_count += 1
+            down_block_res_samples += res_samples
+
+        if down_block_additional_residuals is not None:
+            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 += (down_block_res_sample,)
+
+            down_block_res_samples = new_down_block_res_samples
+
+        # mid
+        sample = self.mid_block(
+            sample,
+            emb,
+            encoder_hidden_states=encoder_hidden_states,
+            attention_mask=attention_mask,
+        )
+
+        if mid_block_additional_residual is not None:
+            sample = sample + mid_block_additional_residual
+
+        # 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,
+                    upsample_size=upsample_size,
+                    attention_mask=attention_mask,
+                )
+            else:
+                sample = upsample_block(
+                    hidden_states=sample,
+                    temb=emb,
+                    res_hidden_states_tuple=res_samples,
+                    upsample_size=upsample_size,
+                    encoder_hidden_states=encoder_hidden_states,
+                )
+
+        # post-process
+        sample = self.conv_norm_out(sample)
+        sample = self.conv_act(sample)
+        sample = self.conv_out(sample)
+
+        if not return_dict:
+            return (sample,)
+
+        return UNet3DConditionOutput(sample=sample)
+
+    @classmethod
+    def from_pretrained_2d(
+        cls,
+        pretrained_model_path: PathLike,
+        motion_module_path: PathLike,
+        subfolder=None,
+        unet_additional_kwargs=None,
+        mm_zero_proj_out=False,
+    ):
+        pretrained_model_path = Path(pretrained_model_path)
+        motion_module_path = Path(motion_module_path)
+        if subfolder is not None:
+            pretrained_model_path = pretrained_model_path.joinpath(subfolder)
+        logger.info(
+            f"loaded temporal unet's pretrained weights from {pretrained_model_path} ..."
+        )
+
+        config_file = pretrained_model_path / "config.json"
+        if not (config_file.exists() and config_file.is_file()):
+            raise RuntimeError(f"{config_file} does not exist or is not a file")
+
+        unet_config = cls.load_config(config_file)
+        unet_config["_class_name"] = cls.__name__
+        unet_config["down_block_types"] = [
+            "CrossAttnDownBlock3D",
+            "CrossAttnDownBlock3D",
+            "CrossAttnDownBlock3D",
+            "DownBlock3D",
+        ]
+        unet_config["up_block_types"] = [
+            "UpBlock3D",
+            "CrossAttnUpBlock3D",
+            "CrossAttnUpBlock3D",
+            "CrossAttnUpBlock3D",
+        ]
+        unet_config["mid_block_type"] = "UNetMidBlock3DCrossAttn"
+
+        model = cls.from_config(unet_config, **unet_additional_kwargs)
+        # load the vanilla weights
+        if pretrained_model_path.joinpath(SAFETENSORS_WEIGHTS_NAME).exists():
+            logger.debug(
+                f"loading safeTensors weights from {pretrained_model_path} ..."
+            )
+            state_dict = load_file(
+                pretrained_model_path.joinpath(SAFETENSORS_WEIGHTS_NAME), device="cpu"
+            )
+
+        elif pretrained_model_path.joinpath(WEIGHTS_NAME).exists():
+            logger.debug(f"loading weights from {pretrained_model_path} ...")
+            state_dict = torch.load(
+                pretrained_model_path.joinpath(WEIGHTS_NAME),
+                map_location="cpu",
+                weights_only=True,
+            )
+        else:
+            raise FileNotFoundError(f"no weights file found in {pretrained_model_path}")
+
+        # load the motion module weights
+        if motion_module_path.exists() and motion_module_path.is_file():
+            if motion_module_path.suffix.lower() in [".pth", ".pt", ".ckpt"]:
+                logger.info(f"Load motion module params from {motion_module_path}")
+                motion_state_dict = torch.load(
+                    motion_module_path, map_location="cpu", weights_only=True
+                )
+            elif motion_module_path.suffix.lower() == ".safetensors":
+                motion_state_dict = load_file(motion_module_path, device="cpu")
+            else:
+                raise RuntimeError(
+                    f"unknown file format for motion module weights: {motion_module_path.suffix}"
+                )
+            if mm_zero_proj_out:
+                logger.info(f"Zero initialize proj_out layers in motion module...")
+                new_motion_state_dict = OrderedDict()
+                for k in motion_state_dict:
+                    if "proj_out" in k:
+                        continue
+                    new_motion_state_dict[k] = motion_state_dict[k]
+                motion_state_dict = new_motion_state_dict
+
+            # merge the state dicts
+            state_dict.update(motion_state_dict)
+
+        # load the weights into the model
+        m, u = model.load_state_dict(state_dict, strict=False)
+        logger.debug(f"### missing keys: {len(m)}; \n### unexpected keys: {len(u)};")
+
+        params = [
+            p.numel() if "temporal" in n else 0 for n, p in model.named_parameters()
+        ]
+        logger.info(f"Loaded {sum(params) / 1e6}M-parameter motion module")
+
+        return model

src/models/unet_3d_blocks.py

@@ -0,0 +1,861 @@
+# Adapted from https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/unet_2d_blocks.py
+
+import pdb
+from typing import Dict, Optional
+import torch
+from torch import nn
+
+from .motion_module import get_motion_module
+
+# from .motion_module import get_motion_module
+from .resnet import Downsample3D, ResnetBlock3D, Upsample3D
+from .transformer_3d import Transformer3DModel
+
+
+def get_down_block(
+    down_block_type,
+    num_layers,
+    in_channels,
+    out_channels,
+    temb_channels,
+    add_downsample,
+    resnet_eps,
+    resnet_act_fn,
+    attn_num_head_channels,
+    resnet_groups=None,
+    cross_attention_dim=None,
+    downsample_padding=None,
+    dual_cross_attention=False,
+    use_linear_projection=False,
+    only_cross_attention=False,
+    upcast_attention=False,
+    resnet_time_scale_shift="default",
+    unet_use_cross_frame_attention=None,
+    unet_use_temporal_attention=None,
+    use_inflated_groupnorm=None,
+    use_motion_module=None,
+    motion_module_type=None,
+    motion_module_kwargs=None,
+):
+    down_block_type = (
+        down_block_type[7:]
+        if down_block_type.startswith("UNetRes")
+        else down_block_type
+    )
+    if down_block_type == "DownBlock3D":
+        return DownBlock3D(
+            num_layers=num_layers,
+            in_channels=in_channels,
+            out_channels=out_channels,
+            temb_channels=temb_channels,
+            add_downsample=add_downsample,
+            resnet_eps=resnet_eps,
+            resnet_act_fn=resnet_act_fn,
+            resnet_groups=resnet_groups,
+            downsample_padding=downsample_padding,
+            resnet_time_scale_shift=resnet_time_scale_shift,
+            use_inflated_groupnorm=use_inflated_groupnorm,
+            use_motion_module=use_motion_module,
+            motion_module_type=motion_module_type,
+            motion_module_kwargs=motion_module_kwargs,
+        )
+    elif down_block_type == "CrossAttnDownBlock3D":
+        if cross_attention_dim is None:
+            raise ValueError(
+                "cross_attention_dim must be specified for CrossAttnDownBlock3D"
+            )
+        return CrossAttnDownBlock3D(
+            num_layers=num_layers,
+            in_channels=in_channels,
+            out_channels=out_channels,
+            temb_channels=temb_channels,
+            add_downsample=add_downsample,
+            resnet_eps=resnet_eps,
+            resnet_act_fn=resnet_act_fn,
+            resnet_groups=resnet_groups,
+            downsample_padding=downsample_padding,
+            cross_attention_dim=cross_attention_dim,
+            attn_num_head_channels=attn_num_head_channels,
+            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,
+            unet_use_cross_frame_attention=unet_use_cross_frame_attention,
+            unet_use_temporal_attention=unet_use_temporal_attention,
+            use_inflated_groupnorm=use_inflated_groupnorm,
+            use_motion_module=use_motion_module,
+            motion_module_type=motion_module_type,
+            motion_module_kwargs=motion_module_kwargs,
+        )
+    raise ValueError(f"{down_block_type} does not exist.")
+
+
+def get_up_block(
+    up_block_type,
+    num_layers,
+    in_channels,
+    out_channels,
+    prev_output_channel,
+    temb_channels,
+    add_upsample,
+    resnet_eps,
+    resnet_act_fn,
+    attn_num_head_channels,
+    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",
+    unet_use_cross_frame_attention=None,
+    unet_use_temporal_attention=None,
+    use_inflated_groupnorm=None,
+    use_motion_module=None,
+    motion_module_type=None,
+    motion_module_kwargs=None,
+):
+    up_block_type = (
+        up_block_type[7:] if up_block_type.startswith("UNetRes") else up_block_type
+    )
+    if up_block_type == "UpBlock3D":
+        return UpBlock3D(
+            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,
+            use_inflated_groupnorm=use_inflated_groupnorm,
+            use_motion_module=use_motion_module,
+            motion_module_type=motion_module_type,
+            motion_module_kwargs=motion_module_kwargs,
+        )
+    elif up_block_type == "CrossAttnUpBlock3D":
+        if cross_attention_dim is None:
+            raise ValueError(
+                "cross_attention_dim must be specified for CrossAttnUpBlock3D"
+            )
+        return CrossAttnUpBlock3D(
+            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,
+            attn_num_head_channels=attn_num_head_channels,
+            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,
+            unet_use_cross_frame_attention=unet_use_cross_frame_attention,
+            unet_use_temporal_attention=unet_use_temporal_attention,
+            use_inflated_groupnorm=use_inflated_groupnorm,
+            use_motion_module=use_motion_module,
+            motion_module_type=motion_module_type,
+            motion_module_kwargs=motion_module_kwargs,
+        )
+    raise ValueError(f"{up_block_type} does not exist.")
+
+
+class UNetMidBlock3DCrossAttn(nn.Module):
+    def __init__(
+        self,
+        in_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,
+        attn_num_head_channels=1,
+        output_scale_factor=1.0,
+        cross_attention_dim=1280,
+        dual_cross_attention=False,
+        use_linear_projection=False,
+        upcast_attention=False,
+        unet_use_cross_frame_attention=None,
+        unet_use_temporal_attention=None,
+        use_inflated_groupnorm=None,
+        use_motion_module=None,
+        motion_module_type=None,
+        motion_module_kwargs=None,
+    ):
+        super().__init__()
+
+        self.has_cross_attention = True
+        self.attn_num_head_channels = attn_num_head_channels
+        resnet_groups = (
+            resnet_groups if resnet_groups is not None else min(in_channels // 4, 32)
+        )
+
+        # there is always at least one resnet
+        resnets = [
+            ResnetBlock3D(
+                in_channels=in_channels,
+                out_channels=in_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,
+                use_inflated_groupnorm=use_inflated_groupnorm,
+            )
+        ]
+        attentions = []
+        motion_modules = []
+
+        for _ in range(num_layers):
+            if dual_cross_attention:
+                raise NotImplementedError
+            attentions.append(
+                Transformer3DModel(
+                    attn_num_head_channels,
+                    in_channels // attn_num_head_channels,
+                    in_channels=in_channels,
+                    num_layers=1,
+                    cross_attention_dim=cross_attention_dim,
+                    norm_num_groups=resnet_groups,
+                    use_linear_projection=use_linear_projection,
+                    upcast_attention=upcast_attention,
+                    unet_use_cross_frame_attention=unet_use_cross_frame_attention,
+                    unet_use_temporal_attention=unet_use_temporal_attention,
+                )
+            )
+            motion_modules.append(
+                get_motion_module(
+                    in_channels=in_channels,
+                    motion_module_type=motion_module_type,
+                    motion_module_kwargs=motion_module_kwargs,
+                )
+                if use_motion_module
+                else None
+            )
+            resnets.append(
+                ResnetBlock3D(
+                    in_channels=in_channels,
+                    out_channels=in_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,
+                    use_inflated_groupnorm=use_inflated_groupnorm,
+                )
+            )
+
+        self.attentions = nn.ModuleList(attentions)
+        self.resnets = nn.ModuleList(resnets)
+        self.motion_modules = nn.ModuleList(motion_modules)
+
+    def forward(
+        self,
+        hidden_states,
+        temb=None,
+        encoder_hidden_states=None,
+        attention_mask=None,
+    ):
+        hidden_states = self.resnets[0](hidden_states, temb)
+        for attn, resnet, motion_module in zip(
+            self.attentions, self.resnets[1:], self.motion_modules
+        ):
+            hidden_states = attn(
+                hidden_states,
+                encoder_hidden_states=encoder_hidden_states,
+            ).sample
+            hidden_states = (
+                motion_module(
+                    hidden_states, temb, encoder_hidden_states=encoder_hidden_states
+                )
+                if motion_module is not None
+                else hidden_states
+            )
+            hidden_states = resnet(hidden_states, temb)
+
+        return hidden_states
+
+
+class CrossAttnDownBlock3D(nn.Module):
+    def __init__(
+        self,
+        in_channels: 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,
+        attn_num_head_channels=1,
+        cross_attention_dim=1280,
+        output_scale_factor=1.0,
+        downsample_padding=1,
+        add_downsample=True,
+        dual_cross_attention=False,
+        use_linear_projection=False,
+        only_cross_attention=False,
+        upcast_attention=False,
+        unet_use_cross_frame_attention=None,
+        unet_use_temporal_attention=None,
+        use_inflated_groupnorm=None,
+        use_motion_module=None,
+        motion_module_type=None,
+        motion_module_kwargs=None,
+    ):
+        super().__init__()
+        resnets = []
+        attentions = []
+        motion_modules = []
+
+        self.has_cross_attention = True
+        self.attn_num_head_channels = attn_num_head_channels
+
+        for i in range(num_layers):
+            in_channels = in_channels if i == 0 else out_channels
+            resnets.append(
+                ResnetBlock3D(
+                    in_channels=in_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,
+                    use_inflated_groupnorm=use_inflated_groupnorm,
+                )
+            )
+            if dual_cross_attention:
+                raise NotImplementedError
+            attentions.append(
+                Transformer3DModel(
+                    attn_num_head_channels,
+                    out_channels // attn_num_head_channels,
+                    in_channels=out_channels,
+                    num_layers=1,
+                    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,
+                    unet_use_cross_frame_attention=unet_use_cross_frame_attention,
+                    unet_use_temporal_attention=unet_use_temporal_attention,
+                )
+            )
+            motion_modules.append(
+                get_motion_module(
+                    in_channels=out_channels,
+                    motion_module_type=motion_module_type,
+                    motion_module_kwargs=motion_module_kwargs,
+                )
+                if use_motion_module
+                else None
+            )
+
+        self.attentions = nn.ModuleList(attentions)
+        self.resnets = nn.ModuleList(resnets)
+        self.motion_modules = nn.ModuleList(motion_modules)
+
+        if add_downsample:
+            self.downsamplers = nn.ModuleList(
+                [
+                    Downsample3D(
+                        out_channels,
+                        use_conv=True,
+                        out_channels=out_channels,
+                        padding=downsample_padding,
+                        name="op",
+                    )
+                ]
+            )
+        else:
+            self.downsamplers = None
+
+        self.gradient_checkpointing = False
+
+    def forward(
+        self,
+        hidden_states,
+        temb=None,
+        encoder_hidden_states=None,
+        attention_mask=None,
+    ):
+        output_states = ()
+
+        for i, (resnet, attn, motion_module) in enumerate(
+            zip(self.resnets, self.attentions, self.motion_modules)
+        ):
+            # self.gradient_checkpointing = False
+            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
+
+                hidden_states = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(resnet), hidden_states, temb
+                )
+                hidden_states = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(attn, return_dict=False),
+                    hidden_states,
+                    encoder_hidden_states,
+                )[0]
+
+                # add motion module
+                hidden_states = (
+                    motion_module(
+                        hidden_states, temb, encoder_hidden_states=encoder_hidden_states
+                    )
+                    if motion_module is not None
+                    else hidden_states
+                )
+
+            else:
+                hidden_states = resnet(hidden_states, temb)
+                hidden_states = attn(
+                    hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                ).sample
+
+                # add motion module
+                hidden_states = (
+                    motion_module(
+                        hidden_states, temb, encoder_hidden_states=encoder_hidden_states
+                    )
+                    if motion_module is not None
+                    else hidden_states
+                )
+
+            output_states += (hidden_states,)
+
+        if self.downsamplers is not None:
+            for downsampler in self.downsamplers:
+                hidden_states = downsampler(hidden_states)
+
+            output_states += (hidden_states,)
+
+        return hidden_states, output_states
+
+
+class DownBlock3D(nn.Module):
+    def __init__(
+        self,
+        in_channels: 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_downsample=True,
+        downsample_padding=1,
+        use_inflated_groupnorm=None,
+        use_motion_module=None,
+        motion_module_type=None,
+        motion_module_kwargs=None,
+    ):
+        super().__init__()
+        resnets = []
+        motion_modules = []
+
+        # use_motion_module = False
+        for i in range(num_layers):
+            in_channels = in_channels if i == 0 else out_channels
+            resnets.append(
+                ResnetBlock3D(
+                    in_channels=in_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,
+                    use_inflated_groupnorm=use_inflated_groupnorm,
+                )
+            )
+            motion_modules.append(
+                get_motion_module(
+                    in_channels=out_channels,
+                    motion_module_type=motion_module_type,
+                    motion_module_kwargs=motion_module_kwargs,
+                )
+                if use_motion_module
+                else None
+            )
+
+        self.resnets = nn.ModuleList(resnets)
+        self.motion_modules = nn.ModuleList(motion_modules)
+
+        if add_downsample:
+            self.downsamplers = nn.ModuleList(
+                [
+                    Downsample3D(
+                        out_channels,
+                        use_conv=True,
+                        out_channels=out_channels,
+                        padding=downsample_padding,
+                        name="op",
+                    )
+                ]
+            )
+        else:
+            self.downsamplers = None
+
+        self.gradient_checkpointing = False
+
+    def forward(self, hidden_states, temb=None, encoder_hidden_states=None):
+        output_states = ()
+
+        for resnet, motion_module in zip(self.resnets, self.motion_modules):
+            # print(f"DownBlock3D {self.gradient_checkpointing = }")
+            if self.training and self.gradient_checkpointing:
+
+                def create_custom_forward(module):
+                    def custom_forward(*inputs):
+                        return module(*inputs)
+
+                    return custom_forward
+
+                hidden_states = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(resnet), hidden_states, temb
+                )
+                if motion_module is not None:
+                    hidden_states = torch.utils.checkpoint.checkpoint(
+                        create_custom_forward(motion_module),
+                        hidden_states.requires_grad_(),
+                        temb,
+                        encoder_hidden_states,
+                    )
+            else:
+                hidden_states = resnet(hidden_states, temb)
+
+                # add motion module
+                hidden_states = (
+                    motion_module(
+                        hidden_states, temb, encoder_hidden_states=encoder_hidden_states
+                    )
+                    if motion_module is not None
+                    else hidden_states
+                )
+
+            output_states += (hidden_states,)
+
+        if self.downsamplers is not None:
+            for downsampler in self.downsamplers:
+                hidden_states = downsampler(hidden_states)
+
+            output_states += (hidden_states,)
+
+        return hidden_states, output_states
+
+
+class CrossAttnUpBlock3D(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        prev_output_channel: 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,
+        attn_num_head_channels=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,
+        unet_use_cross_frame_attention=None,
+        unet_use_temporal_attention=None,
+        use_motion_module=None,
+        use_inflated_groupnorm=None,
+        motion_module_type=None,
+        motion_module_kwargs=None,
+    ):
+        super().__init__()
+        resnets = []
+        attentions = []
+        motion_modules = []
+
+        self.has_cross_attention = True
+        self.attn_num_head_channels = attn_num_head_channels
+
+        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(
+                ResnetBlock3D(
+                    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,
+                    use_inflated_groupnorm=use_inflated_groupnorm,
+                )
+            )
+            if dual_cross_attention:
+                raise NotImplementedError
+            attentions.append(
+                Transformer3DModel(
+                    attn_num_head_channels,
+                    out_channels // attn_num_head_channels,
+                    in_channels=out_channels,
+                    num_layers=1,
+                    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,
+                    unet_use_cross_frame_attention=unet_use_cross_frame_attention,
+                    unet_use_temporal_attention=unet_use_temporal_attention,
+                )
+            )
+            motion_modules.append(
+                get_motion_module(
+                    in_channels=out_channels,
+                    motion_module_type=motion_module_type,
+                    motion_module_kwargs=motion_module_kwargs,
+                )
+                if use_motion_module
+                else None
+            )
+
+        self.attentions = nn.ModuleList(attentions)
+        self.resnets = nn.ModuleList(resnets)
+        self.motion_modules = nn.ModuleList(motion_modules)
+
+        if add_upsample:
+            self.upsamplers = nn.ModuleList(
+                [Upsample3D(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,
+        encoder_hidden_states=None,
+        upsample_size=None,
+        attention_mask=None,
+    ):
+        for i, (resnet, attn, motion_module) in enumerate(
+            zip(self.resnets, self.attentions, self.motion_modules)
+        ):
+            # 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
+
+                hidden_states = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(resnet), hidden_states, temb
+                )
+                hidden_states = attn(
+                    hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                ).sample
+                if motion_module is not None:
+                    hidden_states = torch.utils.checkpoint.checkpoint(
+                        create_custom_forward(motion_module),
+                        hidden_states.requires_grad_(),
+                        temb,
+                        encoder_hidden_states,
+                    )
+
+            else:
+                hidden_states = resnet(hidden_states, temb)
+                hidden_states = attn(
+                    hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                ).sample
+
+                # add motion module
+                hidden_states = (
+                    motion_module(
+                        hidden_states, temb, encoder_hidden_states=encoder_hidden_states
+                    )
+                    if motion_module is not None
+                    else hidden_states
+                )
+
+        if self.upsamplers is not None:
+            for upsampler in self.upsamplers:
+                hidden_states = upsampler(hidden_states, upsample_size)
+
+        return hidden_states
+
+class UpBlock3D(nn.Module):
+    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,
+        use_inflated_groupnorm=None,
+        use_motion_module=None,
+        motion_module_type=None,
+        motion_module_kwargs=None,
+    ):
+        super().__init__()
+        resnets = []
+        motion_modules = []
+
+        # use_motion_module = False
+        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(
+                ResnetBlock3D(
+                    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,
+                    use_inflated_groupnorm=use_inflated_groupnorm,
+                )
+            )
+            motion_modules.append(
+                get_motion_module(
+                    in_channels=out_channels,
+                    motion_module_type=motion_module_type,
+                    motion_module_kwargs=motion_module_kwargs,
+                )
+                if use_motion_module
+                else None
+            )
+
+        self.resnets = nn.ModuleList(resnets)
+        self.motion_modules = nn.ModuleList(motion_modules)
+
+        if add_upsample:
+            self.upsamplers = nn.ModuleList(
+                [Upsample3D(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,
+        encoder_hidden_states=None,
+    ):
+        for resnet, motion_module in zip(self.resnets, self.motion_modules):
+            # 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)
+
+            # print(f"UpBlock3D {self.gradient_checkpointing = }")
+            if self.training and self.gradient_checkpointing:
+
+                def create_custom_forward(module):
+                    def custom_forward(*inputs):
+                        return module(*inputs)
+
+                    return custom_forward
+
+                hidden_states = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(resnet), hidden_states, temb
+                )
+                if motion_module is not None:
+                    hidden_states = torch.utils.checkpoint.checkpoint(
+                        create_custom_forward(motion_module),
+                        hidden_states.requires_grad_(),
+                        temb,
+                        encoder_hidden_states,
+                    )
+            else:
+                hidden_states = resnet(hidden_states, temb)
+                hidden_states = (
+                    motion_module(
+                        hidden_states, temb, encoder_hidden_states=encoder_hidden_states
+                    )
+                    if motion_module is not None
+                    else hidden_states
+                )
+
+        if self.upsamplers is not None:
+            for upsampler in self.upsamplers:
+                hidden_states = upsampler(hidden_states, upsample_size)
+
+        return hidden_states

src/pipelines/context.py

@@ -0,0 +1,76 @@
+# TODO: Adapted from cli
+from typing import Callable, List, Optional
+
+import numpy as np
+
+
+def ordered_halving(val):
+    bin_str = f"{val:064b}"
+    bin_flip = bin_str[::-1]
+    as_int = int(bin_flip, 2)
+
+    return as_int / (1 << 64)
+
+
+def uniform(
+    step: int = ...,
+    num_steps: Optional[int] = None,
+    num_frames: int = ...,
+    context_size: Optional[int] = None,
+    context_stride: int = 3,
+    context_overlap: int = 4,
+    closed_loop: bool = True,
+):
+    if num_frames <= context_size:
+        yield list(range(num_frames))
+        return
+
+    context_stride = min(
+        context_stride, int(np.ceil(np.log2(num_frames / context_size))) + 1
+    )
+
+    for context_step in 1 << np.arange(context_stride):
+        pad = int(round(num_frames * ordered_halving(step)))
+        for j in range(
+            int(ordered_halving(step) * context_step) + pad,
+            num_frames + pad + (0 if closed_loop else -context_overlap),
+            (context_size * context_step - context_overlap),
+        ):
+            yield [
+                e % num_frames
+                for e in range(j, j + context_size * context_step, context_step)
+            ]
+
+
+def get_context_scheduler(name: str) -> Callable:
+    if name == "uniform":
+        return uniform
+    else:
+        raise ValueError(f"Unknown context_overlap policy {name}")
+
+
+def get_total_steps(
+    scheduler,
+    timesteps: List[int],
+    num_steps: Optional[int] = None,
+    num_frames: int = ...,
+    context_size: Optional[int] = None,
+    context_stride: int = 3,
+    context_overlap: int = 4,
+    closed_loop: bool = True,
+):
+    return sum(
+        len(
+            list(
+                scheduler(
+                    i,
+                    num_steps,
+                    num_frames,
+                    context_size,
+                    context_stride,
+                    context_overlap,
+                )
+            )
+        )
+        for i in range(len(timesteps))
+    )

src/pipelines/pipeline_pose2vid_long.py

@@ -0,0 +1,584 @@
+# Adapted from https://github.com/magic-research/magic-animate/blob/main/magicanimate/pipelines/pipeline_animation.py
+import inspect
+import math
+from dataclasses import dataclass
+from typing import Callable, List, Optional, Union
+
+import numpy as np
+import torch
+import torchvision.transforms as transforms
+from diffusers import DiffusionPipeline
+from diffusers.image_processor import VaeImageProcessor
+from diffusers.schedulers import (
+    DDIMScheduler,
+    DPMSolverMultistepScheduler,
+    EulerAncestralDiscreteScheduler,
+    EulerDiscreteScheduler,
+    LMSDiscreteScheduler,
+    PNDMScheduler,
+)
+from diffusers.utils import BaseOutput, deprecate, is_accelerate_available, logging
+from diffusers.utils.torch_utils import randn_tensor
+from einops import rearrange
+from tqdm import tqdm
+from transformers import CLIPImageProcessor
+
+from src.models.mutual_self_attention import ReferenceAttentionControl
+from src.pipelines.context import get_context_scheduler
+from src.pipelines.utils import get_tensor_interpolation_method
+
+
+@dataclass
+class Pose2VideoPipelineOutput(BaseOutput):
+    videos: Union[torch.Tensor, np.ndarray]
+
+
+class Pose2VideoPipeline(DiffusionPipeline):
+    _optional_components = []
+
+    def __init__(
+        self,
+        vae,
+        image_encoder,
+        reference_unet,
+        denoising_unet,
+        pose_guider,
+        scheduler: Union[
+            DDIMScheduler,
+            PNDMScheduler,
+            LMSDiscreteScheduler,
+            EulerDiscreteScheduler,
+            EulerAncestralDiscreteScheduler,
+            DPMSolverMultistepScheduler,
+        ],
+        image_proj_model=None,
+        tokenizer=None,
+        text_encoder=None,
+    ):
+        super().__init__()
+
+        self.register_modules(
+            vae=vae,
+            image_encoder=image_encoder,
+            reference_unet=reference_unet,
+            denoising_unet=denoising_unet,
+            pose_guider=pose_guider,
+            scheduler=scheduler,
+            image_proj_model=image_proj_model,
+            tokenizer=tokenizer,
+            text_encoder=text_encoder,
+        )
+        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
+        self.clip_image_processor = CLIPImageProcessor()
+        self.ref_image_processor = VaeImageProcessor(
+            vae_scale_factor=self.vae_scale_factor, do_convert_rgb=True
+        )
+        self.cond_image_processor = VaeImageProcessor(
+            vae_scale_factor=self.vae_scale_factor,
+            do_convert_rgb=True,
+            do_normalize=True,
+        )
+
+    def enable_vae_slicing(self):
+        self.vae.enable_slicing()
+
+    def disable_vae_slicing(self):
+        self.vae.disable_slicing()
+
+    def enable_sequential_cpu_offload(self, gpu_id=0):
+        if is_accelerate_available():
+            from accelerate import cpu_offload
+        else:
+            raise ImportError("Please install accelerate via `pip install accelerate`")
+
+        device = torch.device(f"cuda:{gpu_id}")
+
+        for cpu_offloaded_model in [self.unet, self.text_encoder, self.vae]:
+            if cpu_offloaded_model is not None:
+                cpu_offload(cpu_offloaded_model, device)
+
+    @property
+    def _execution_device(self):
+        if self.device != torch.device("meta") or not hasattr(self.unet, "_hf_hook"):
+            return self.device
+        for module in self.unet.modules():
+            if (
+                hasattr(module, "_hf_hook")
+                and hasattr(module._hf_hook, "execution_device")
+                and module._hf_hook.execution_device is not None
+            ):
+                return torch.device(module._hf_hook.execution_device)
+        return self.device
+
+    def decode_latents(self, latents):
+        video_length = latents.shape[2]
+        latents = 1 / 0.18215 * latents
+        latents = rearrange(latents, "b c f h w -> (b f) c h w")
+        # video = self.vae.decode(latents).sample
+        video = []
+        for frame_idx in tqdm(range(latents.shape[0])):
+            video.append(self.vae.decode(latents[frame_idx : frame_idx + 1]).sample)
+        video = torch.cat(video)
+        video = rearrange(video, "(b f) c h w -> b c f h w", f=video_length)
+        video = (video / 2 + 0.5).clamp(0, 1)
+        # we always cast to float32 as this does not cause significant overhead and is compatible with bfloa16
+        video = video.cpu().float().numpy()
+        return video
+
+    def prepare_extra_step_kwargs(self, generator, eta):
+        # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
+        # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
+        # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
+        # and should be between [0, 1]
+
+        accepts_eta = "eta" in set(
+            inspect.signature(self.scheduler.step).parameters.keys()
+        )
+        extra_step_kwargs = {}
+        if accepts_eta:
+            extra_step_kwargs["eta"] = eta
+
+        # check if the scheduler accepts generator
+        accepts_generator = "generator" in set(
+            inspect.signature(self.scheduler.step).parameters.keys()
+        )
+        if accepts_generator:
+            extra_step_kwargs["generator"] = generator
+        return extra_step_kwargs
+
+    def prepare_latents(
+        self,
+        batch_size,
+        num_channels_latents,
+        width,
+        height,
+        video_length,
+        dtype,
+        device,
+        generator,
+        latents=None,
+    ):
+        shape = (
+            batch_size,
+            num_channels_latents,
+            video_length,
+            height // self.vae_scale_factor,
+            width // self.vae_scale_factor,
+        )
+        if isinstance(generator, list) and len(generator) != batch_size:
+            raise ValueError(
+                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+            )
+
+        if latents is None:
+            latents = randn_tensor(
+                shape, generator=generator, device=device, dtype=dtype
+            )
+        else:
+            latents = latents.to(device)
+
+        # scale the initial noise by the standard deviation required by the scheduler
+        latents = latents * self.scheduler.init_noise_sigma
+        return latents
+
+    def _encode_prompt(
+        self,
+        prompt,
+        device,
+        num_videos_per_prompt,
+        do_classifier_free_guidance,
+        negative_prompt,
+    ):
+        batch_size = len(prompt) if isinstance(prompt, list) else 1
+
+        text_inputs = self.tokenizer(
+            prompt,
+            padding="max_length",
+            max_length=self.tokenizer.model_max_length,
+            truncation=True,
+            return_tensors="pt",
+        )
+        text_input_ids = text_inputs.input_ids
+        untruncated_ids = self.tokenizer(
+            prompt, padding="longest", return_tensors="pt"
+        ).input_ids
+
+        if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(
+            text_input_ids, untruncated_ids
+        ):
+            removed_text = self.tokenizer.batch_decode(
+                untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1]
+            )
+
+        if (
+            hasattr(self.text_encoder.config, "use_attention_mask")
+            and self.text_encoder.config.use_attention_mask
+        ):
+            attention_mask = text_inputs.attention_mask.to(device)
+        else:
+            attention_mask = None
+
+        text_embeddings = self.text_encoder(
+            text_input_ids.to(device),
+            attention_mask=attention_mask,
+        )
+        text_embeddings = text_embeddings[0]
+
+        # duplicate text embeddings for each generation per prompt, using mps friendly method
+        bs_embed, seq_len, _ = text_embeddings.shape
+        text_embeddings = text_embeddings.repeat(1, num_videos_per_prompt, 1)
+        text_embeddings = text_embeddings.view(
+            bs_embed * num_videos_per_prompt, seq_len, -1
+        )
+
+        # get unconditional embeddings for classifier free guidance
+        if do_classifier_free_guidance:
+            uncond_tokens: List[str]
+            if negative_prompt is None:
+                uncond_tokens = [""] * batch_size
+            elif type(prompt) is not type(negative_prompt):
+                raise TypeError(
+                    f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
+                    f" {type(prompt)}."
+                )
+            elif isinstance(negative_prompt, str):
+                uncond_tokens = [negative_prompt]
+            elif batch_size != len(negative_prompt):
+                raise ValueError(
+                    f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
+                    f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
+                    " the batch size of `prompt`."
+                )
+            else:
+                uncond_tokens = negative_prompt
+
+            max_length = text_input_ids.shape[-1]
+            uncond_input = self.tokenizer(
+                uncond_tokens,
+                padding="max_length",
+                max_length=max_length,
+                truncation=True,
+                return_tensors="pt",
+            )
+
+            if (
+                hasattr(self.text_encoder.config, "use_attention_mask")
+                and self.text_encoder.config.use_attention_mask
+            ):
+                attention_mask = uncond_input.attention_mask.to(device)
+            else:
+                attention_mask = None
+
+            uncond_embeddings = self.text_encoder(
+                uncond_input.input_ids.to(device),
+                attention_mask=attention_mask,
+            )
+            uncond_embeddings = uncond_embeddings[0]
+
+            # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
+            seq_len = uncond_embeddings.shape[1]
+            uncond_embeddings = uncond_embeddings.repeat(1, num_videos_per_prompt, 1)
+            uncond_embeddings = uncond_embeddings.view(
+                batch_size * num_videos_per_prompt, seq_len, -1
+            )
+
+            # For classifier free guidance, we need to do two forward passes.
+            # Here we concatenate the unconditional and text embeddings into a single batch
+            # to avoid doing two forward passes
+            text_embeddings = torch.cat([uncond_embeddings, text_embeddings])
+
+        return text_embeddings
+
+    def interpolate_latents(
+        self, latents: torch.Tensor, interpolation_factor: int, device
+    ):
+        if interpolation_factor < 2:
+            return latents
+
+        new_latents = torch.zeros(
+            (
+                latents.shape[0],
+                latents.shape[1],
+                ((latents.shape[2] - 1) * interpolation_factor) + 1,
+                latents.shape[3],
+                latents.shape[4],
+            ),
+            device=latents.device,
+            dtype=latents.dtype,
+        )
+
+        org_video_length = latents.shape[2]
+        rate = [i / interpolation_factor for i in range(interpolation_factor)][1:]
+
+        new_index = 0
+
+        v0 = None
+        v1 = None
+
+        for i0, i1 in zip(range(org_video_length), range(org_video_length)[1:]):
+            v0 = latents[:, :, i0, :, :]
+            v1 = latents[:, :, i1, :, :]
+
+            new_latents[:, :, new_index, :, :] = v0
+            new_index += 1
+
+            for f in rate:
+                v = get_tensor_interpolation_method()(
+                    v0.to(device=device), v1.to(device=device), f
+                )
+                new_latents[:, :, new_index, :, :] = v.to(latents.device)
+                new_index += 1
+
+        new_latents[:, :, new_index, :, :] = v1
+        new_index += 1
+
+        return new_latents
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        ref_image,
+        pose_images,
+        ref_pose_image,
+        width,
+        height,
+        video_length,
+        num_inference_steps,
+        guidance_scale,
+        num_images_per_prompt=1,
+        eta: float = 0.0,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        output_type: Optional[str] = "tensor",
+        return_dict: bool = True,
+        callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
+        callback_steps: Optional[int] = 1,
+        context_schedule="uniform",
+        context_frames=16,
+        context_stride=1,
+        context_overlap=4,
+        context_batch_size=1,
+        interpolation_factor=1,
+        **kwargs,
+    ):
+        # Default height and width to unet
+        height = height or self.unet.config.sample_size * self.vae_scale_factor
+        width = width or self.unet.config.sample_size * self.vae_scale_factor
+
+        device = self._execution_device
+
+        do_classifier_free_guidance = guidance_scale > 1.0
+
+        # Prepare timesteps
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+        timesteps = self.scheduler.timesteps
+
+        batch_size = 1
+
+        # Prepare clip image embeds
+        clip_image = self.clip_image_processor.preprocess(
+            ref_image.resize((224, 224)), return_tensors="pt"
+        ).pixel_values
+        clip_image_embeds = self.image_encoder(
+            clip_image.to(device, dtype=self.image_encoder.dtype)
+        ).image_embeds
+        encoder_hidden_states = clip_image_embeds.unsqueeze(1)
+        uncond_encoder_hidden_states = torch.zeros_like(encoder_hidden_states)
+
+        if do_classifier_free_guidance:
+            encoder_hidden_states = torch.cat(
+                [uncond_encoder_hidden_states, encoder_hidden_states], dim=0
+            )
+
+        reference_control_writer = ReferenceAttentionControl(
+            self.reference_unet,
+            do_classifier_free_guidance=do_classifier_free_guidance,
+            mode="write",
+            batch_size=batch_size,
+            fusion_blocks="full",
+        )
+        reference_control_reader = ReferenceAttentionControl(
+            self.denoising_unet,
+            do_classifier_free_guidance=do_classifier_free_guidance,
+            mode="read",
+            batch_size=batch_size,
+            fusion_blocks="full",
+        )
+
+        num_channels_latents = self.denoising_unet.in_channels
+        latents = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            width,
+            height,
+            video_length,
+            clip_image_embeds.dtype,
+            device,
+            generator,
+        )
+
+        # Prepare extra step kwargs.
+        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
+
+        # Prepare ref image latents
+        ref_image_tensor = self.ref_image_processor.preprocess(
+            ref_image, height=height, width=width
+        )  # (bs, c, width, height)
+        ref_image_tensor = ref_image_tensor.to(
+            dtype=self.vae.dtype, device=self.vae.device
+        )
+        ref_image_latents = self.vae.encode(ref_image_tensor).latent_dist.mean
+        ref_image_latents = ref_image_latents * 0.18215  # (b, 4, h, w)
+
+        # Prepare a list of pose condition images
+        pose_cond_tensor_list = []
+        for pose_image in pose_images:
+            pose_cond_tensor = self.cond_image_processor.preprocess(
+                pose_image, height=height, width=width
+            )
+            pose_cond_tensor = pose_cond_tensor.unsqueeze(2)  # (bs, c, 1, h, w)
+            pose_cond_tensor_list.append(pose_cond_tensor)
+        pose_cond_tensor = torch.cat(pose_cond_tensor_list, dim=2)  # (bs, c, t, h, w)
+        
+        pose_cond_tensor = pose_cond_tensor.to(
+            device=device, dtype=self.pose_guider.dtype
+        )
+        
+        ref_pose_tensor = self.cond_image_processor.preprocess(
+            ref_pose_image, height=height, width=width
+        )
+        ref_pose_tensor = ref_pose_tensor.to(
+            device=device, dtype=self.pose_guider.dtype
+        )
+
+        context_scheduler = get_context_scheduler(context_schedule)
+
+        # denoising loop
+        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                noise_pred = torch.zeros(
+                    (
+                        latents.shape[0] * (2 if do_classifier_free_guidance else 1),
+                        *latents.shape[1:],
+                    ),
+                    device=latents.device,
+                    dtype=latents.dtype,
+                )
+                counter = torch.zeros(
+                    (1, 1, latents.shape[2], 1, 1),
+                    device=latents.device,
+                    dtype=latents.dtype,
+                )
+
+                # 1. Forward reference image
+                if i == 0:
+                    self.reference_unet(
+                        ref_image_latents.repeat(
+                            (2 if do_classifier_free_guidance else 1), 1, 1, 1
+                        ),
+                        torch.zeros_like(t),
+                        # t,
+                        encoder_hidden_states=encoder_hidden_states,
+                        return_dict=False,
+                    )
+                    reference_control_reader.update(reference_control_writer)
+
+                context_queue = list(
+                    context_scheduler(
+                        0,
+                        num_inference_steps,
+                        latents.shape[2],
+                        context_frames,
+                        context_stride,
+                        0,
+                    )
+                )
+                num_context_batches = math.ceil(len(context_queue) / context_batch_size)
+
+                context_queue = list(
+                    context_scheduler(
+                        0,
+                        num_inference_steps,
+                        latents.shape[2],
+                        context_frames,
+                        context_stride,
+                        context_overlap,
+                    )
+                )
+
+                num_context_batches = math.ceil(len(context_queue) / context_batch_size)
+                global_context = []
+                for i in range(num_context_batches):
+                    global_context.append(
+                        context_queue[
+                            i * context_batch_size : (i + 1) * context_batch_size
+                        ]
+                    )
+
+                for context in global_context:
+                    # 3.1 expand the latents if we are doing classifier free guidance
+                    latent_model_input = (
+                        torch.cat([latents[:, :, c] for c in context])
+                        .to(device)
+                        .repeat(2 if do_classifier_free_guidance else 1, 1, 1, 1, 1)
+                    )
+                    latent_model_input = self.scheduler.scale_model_input(
+                        latent_model_input, t
+                    )
+                    b, c, f, h, w = latent_model_input.shape
+                    
+                    pose_cond_input = (
+                        torch.cat([pose_cond_tensor[:, :, c] for c in context])
+                        .to(device)
+                        .repeat(2 if do_classifier_free_guidance else 1, 1, 1, 1, 1)
+                    )
+                    pose_fea = self.pose_guider(pose_cond_input, ref_pose_tensor)
+
+                    pred = self.denoising_unet(
+                        latent_model_input,
+                        t,
+                        encoder_hidden_states=encoder_hidden_states[:b],
+                        pose_cond_fea=pose_fea,
+                        return_dict=False,
+                    )[0]
+
+                    for j, c in enumerate(context):
+                        noise_pred[:, :, c] = noise_pred[:, :, c] + pred
+                        counter[:, :, c] = counter[:, :, c] + 1
+
+                # perform guidance
+                if do_classifier_free_guidance:
+                    noise_pred_uncond, noise_pred_text = (noise_pred / counter).chunk(2)
+                    noise_pred = noise_pred_uncond + guidance_scale * (
+                        noise_pred_text - noise_pred_uncond
+                    )
+
+                latents = self.scheduler.step(
+                    noise_pred, t, latents, **extra_step_kwargs
+                ).prev_sample
+
+                if i == len(timesteps) - 1 or (
+                    (i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0
+                ):
+                    progress_bar.update()
+                    if callback is not None and i % callback_steps == 0:
+                        step_idx = i // getattr(self.scheduler, "order", 1)
+                        callback(step_idx, t, latents)
+
+            reference_control_reader.clear()
+            reference_control_writer.clear()
+
+        if interpolation_factor > 0:
+            latents = self.interpolate_latents(latents, interpolation_factor, device)
+        # Post-processing
+        images = self.decode_latents(latents)  # (b, c, f, h, w)
+
+        # Convert to tensor
+        if output_type == "tensor":
+            images = torch.from_numpy(images)
+
+        if not return_dict:
+            return images
+
+        return Pose2VideoPipelineOutput(videos=images)

src/pipelines/utils.py

@@ -0,0 +1,29 @@
+import torch
+
+tensor_interpolation = None
+
+
+def get_tensor_interpolation_method():
+    return tensor_interpolation
+
+
+def set_tensor_interpolation_method(is_slerp):
+    global tensor_interpolation
+    tensor_interpolation = slerp if is_slerp else linear
+
+
+def linear(v1, v2, t):
+    return (1.0 - t) * v1 + t * v2
+
+
+def slerp(
+    v0: torch.Tensor, v1: torch.Tensor, t: float, DOT_THRESHOLD: float = 0.9995
+) -> torch.Tensor:
+    u0 = v0 / v0.norm()
+    u1 = v1 / v1.norm()
+    dot = (u0 * u1).sum()
+    if dot.abs() > DOT_THRESHOLD:
+        # logger.info(f'warning: v0 and v1 close to parallel, using linear interpolation instead.')
+        return (1.0 - t) * v0 + t * v1
+    omega = dot.acos()
+    return (((1.0 - t) * omega).sin() * v0 + (t * omega).sin() * v1) / omega.sin()

src/utils/audio_util.py

@@ -0,0 +1,30 @@
+import os
+import math
+
+import librosa
+import numpy as np
+from transformers import Wav2Vec2FeatureExtractor
+
+
+class DataProcessor:
+    def __init__(self, sampling_rate, wav2vec_model_path):
+        self._processor = Wav2Vec2FeatureExtractor.from_pretrained(wav2vec_model_path, local_files_only=True)
+        self._sampling_rate = sampling_rate
+
+    def extract_feature(self, audio_path):
+        speech_array, sampling_rate = librosa.load(audio_path, sr=self._sampling_rate)
+        input_value = np.squeeze(self._processor(speech_array, sampling_rate=sampling_rate).input_values)
+        return input_value
+
+
+def prepare_audio_feature(wav_file, fps=30, sampling_rate=16000, wav2vec_model_path=None):
+    data_preprocessor = DataProcessor(sampling_rate, wav2vec_model_path)
+
+    input_value = data_preprocessor.extract_feature(wav_file)
+    seq_len = math.ceil(len(input_value)/sampling_rate*fps)
+    return {
+        "audio_feature": input_value,
+        "seq_len": seq_len
+    }
+
+

src/utils/draw_util.py

@@ -0,0 +1,149 @@
+import cv2
+import mediapipe as mp
+import numpy as np
+from mediapipe.framework.formats import landmark_pb2
+
+class FaceMeshVisualizer:
+    def __init__(self, forehead_edge=False):
+        self.mp_drawing = mp.solutions.drawing_utils
+        mp_face_mesh = mp.solutions.face_mesh
+        self.mp_face_mesh = mp_face_mesh
+        self.forehead_edge = forehead_edge
+
+        DrawingSpec = mp.solutions.drawing_styles.DrawingSpec
+        f_thick = 2
+        f_rad = 1
+        right_iris_draw = DrawingSpec(color=(10, 200, 250), thickness=f_thick, circle_radius=f_rad)
+        right_eye_draw = DrawingSpec(color=(10, 200, 180), thickness=f_thick, circle_radius=f_rad)
+        right_eyebrow_draw = DrawingSpec(color=(10, 220, 180), thickness=f_thick, circle_radius=f_rad)
+        left_iris_draw = DrawingSpec(color=(250, 200, 10), thickness=f_thick, circle_radius=f_rad)
+        left_eye_draw = DrawingSpec(color=(180, 200, 10), thickness=f_thick, circle_radius=f_rad)
+        left_eyebrow_draw = DrawingSpec(color=(180, 220, 10), thickness=f_thick, circle_radius=f_rad)
+        head_draw = DrawingSpec(color=(10, 200, 10), thickness=f_thick, circle_radius=f_rad)
+        
+        mouth_draw_obl = DrawingSpec(color=(10, 180, 20), thickness=f_thick, circle_radius=f_rad)
+        mouth_draw_obr = DrawingSpec(color=(20, 10, 180), thickness=f_thick, circle_radius=f_rad)
+        
+        mouth_draw_ibl = DrawingSpec(color=(100, 100, 30), thickness=f_thick, circle_radius=f_rad)
+        mouth_draw_ibr = DrawingSpec(color=(100, 150, 50), thickness=f_thick, circle_radius=f_rad)
+        
+        mouth_draw_otl = DrawingSpec(color=(20, 80, 100), thickness=f_thick, circle_radius=f_rad)
+        mouth_draw_otr = DrawingSpec(color=(80, 100, 20), thickness=f_thick, circle_radius=f_rad)
+        
+        mouth_draw_itl = DrawingSpec(color=(120, 100, 200), thickness=f_thick, circle_radius=f_rad)
+        mouth_draw_itr = DrawingSpec(color=(150 ,120, 100), thickness=f_thick, circle_radius=f_rad)
+        
+        FACEMESH_LIPS_OUTER_BOTTOM_LEFT = [(61,146),(146,91),(91,181),(181,84),(84,17)]
+        FACEMESH_LIPS_OUTER_BOTTOM_RIGHT = [(17,314),(314,405),(405,321),(321,375),(375,291)]
+        
+        FACEMESH_LIPS_INNER_BOTTOM_LEFT = [(78,95),(95,88),(88,178),(178,87),(87,14)]
+        FACEMESH_LIPS_INNER_BOTTOM_RIGHT = [(14,317),(317,402),(402,318),(318,324),(324,308)]
+        
+        FACEMESH_LIPS_OUTER_TOP_LEFT = [(61,185),(185,40),(40,39),(39,37),(37,0)]
+        FACEMESH_LIPS_OUTER_TOP_RIGHT = [(0,267),(267,269),(269,270),(270,409),(409,291)]
+        
+        FACEMESH_LIPS_INNER_TOP_LEFT = [(78,191),(191,80),(80,81),(81,82),(82,13)]
+        FACEMESH_LIPS_INNER_TOP_RIGHT = [(13,312),(312,311),(311,310),(310,415),(415,308)]
+        
+        FACEMESH_CUSTOM_FACE_OVAL = [(176, 149), (150, 136), (356, 454), (58, 132), (152, 148), (361, 288), (251, 389), (132, 93), (389, 356), (400, 377), (136, 172), (377, 152), (323, 361), (172, 58), (454, 323), (365, 379), (379, 378), (148, 176), (93, 234), (397, 365), (149, 150), (288, 397), (234, 127), (378, 400), (127, 162), (162, 21)]
+
+        # mp_face_mesh.FACEMESH_CONTOURS has all the items we care about.
+        face_connection_spec = {}
+        if self.forehead_edge:
+            for edge in mp_face_mesh.FACEMESH_FACE_OVAL:
+                face_connection_spec[edge] = head_draw
+        else:
+            for edge in FACEMESH_CUSTOM_FACE_OVAL:
+                face_connection_spec[edge] = head_draw
+        for edge in mp_face_mesh.FACEMESH_LEFT_EYE:
+            face_connection_spec[edge] = left_eye_draw
+        for edge in mp_face_mesh.FACEMESH_LEFT_EYEBROW:
+            face_connection_spec[edge] = left_eyebrow_draw
+        # for edge in mp_face_mesh.FACEMESH_LEFT_IRIS:
+        #    face_connection_spec[edge] = left_iris_draw
+        for edge in mp_face_mesh.FACEMESH_RIGHT_EYE:
+            face_connection_spec[edge] = right_eye_draw
+        for edge in mp_face_mesh.FACEMESH_RIGHT_EYEBROW:
+            face_connection_spec[edge] = right_eyebrow_draw
+        # for edge in mp_face_mesh.FACEMESH_RIGHT_IRIS:
+        #    face_connection_spec[edge] = right_iris_draw
+        # for edge in mp_face_mesh.FACEMESH_LIPS:
+        #     face_connection_spec[edge] = mouth_draw
+        
+        for edge in FACEMESH_LIPS_OUTER_BOTTOM_LEFT:
+            face_connection_spec[edge] = mouth_draw_obl
+        for edge in FACEMESH_LIPS_OUTER_BOTTOM_RIGHT:
+            face_connection_spec[edge] = mouth_draw_obr
+        for edge in FACEMESH_LIPS_INNER_BOTTOM_LEFT:
+            face_connection_spec[edge] = mouth_draw_ibl
+        for edge in FACEMESH_LIPS_INNER_BOTTOM_RIGHT:
+            face_connection_spec[edge] = mouth_draw_ibr
+        for edge in FACEMESH_LIPS_OUTER_TOP_LEFT:
+            face_connection_spec[edge] = mouth_draw_otl
+        for edge in FACEMESH_LIPS_OUTER_TOP_RIGHT:
+            face_connection_spec[edge] = mouth_draw_otr
+        for edge in FACEMESH_LIPS_INNER_TOP_LEFT:
+            face_connection_spec[edge] = mouth_draw_itl
+        for edge in FACEMESH_LIPS_INNER_TOP_RIGHT:
+            face_connection_spec[edge] = mouth_draw_itr
+
+        
+        iris_landmark_spec = {468: right_iris_draw, 473: left_iris_draw}
+        
+        self.face_connection_spec = face_connection_spec
+    def draw_pupils(self, image, landmark_list, drawing_spec, halfwidth: int = 2):
+        """We have a custom function to draw the pupils because the mp.draw_landmarks method requires a parameter for all
+        landmarks.  Until our PR is merged into mediapipe, we need this separate method."""
+        if len(image.shape) != 3:
+            raise ValueError("Input image must be H,W,C.")
+        image_rows, image_cols, image_channels = image.shape
+        if image_channels != 3:  # BGR channels
+            raise ValueError('Input image must contain three channel bgr data.')
+        for idx, landmark in enumerate(landmark_list.landmark):
+            if (
+                    (landmark.HasField('visibility') and landmark.visibility < 0.9) or
+                    (landmark.HasField('presence') and landmark.presence < 0.5)
+            ):
+                continue
+            if landmark.x >= 1.0 or landmark.x < 0 or landmark.y >= 1.0 or landmark.y < 0:
+                continue
+            image_x = int(image_cols*landmark.x)
+            image_y = int(image_rows*landmark.y)
+            draw_color = None
+            if isinstance(drawing_spec, Mapping):
+                if drawing_spec.get(idx) is None:
+                    continue
+                else:
+                    draw_color = drawing_spec[idx].color
+            elif isinstance(drawing_spec, DrawingSpec):
+                draw_color = drawing_spec.color
+            image[image_y-halfwidth:image_y+halfwidth, image_x-halfwidth:image_x+halfwidth, :] = draw_color
+        
+
+
+    def draw_landmarks(self, image_size, keypoints, normed=False):
+        ini_size = [512, 512]
+        image = np.zeros([ini_size[1], ini_size[0], 3], dtype=np.uint8)
+        new_landmarks = landmark_pb2.NormalizedLandmarkList()
+        for i in range(keypoints.shape[0]):
+            landmark = new_landmarks.landmark.add()
+            if normed:
+                landmark.x = keypoints[i, 0]
+                landmark.y = keypoints[i, 1]
+            else:
+                landmark.x = keypoints[i, 0] / image_size[0]
+                landmark.y = keypoints[i, 1] / image_size[1]
+            landmark.z = 1.0
+
+        self.mp_drawing.draw_landmarks(
+            image=image,
+            landmark_list=new_landmarks,
+            connections=self.face_connection_spec.keys(),
+            landmark_drawing_spec=None,
+            connection_drawing_spec=self.face_connection_spec
+        )
+        # draw_pupils(image, face_landmarks, iris_landmark_spec, 2)
+        image = cv2.resize(image, (image_size[0], image_size[1]))
+        
+        return image
+

src/utils/face_landmark.py

@@ -0,0 +1,3305 @@
+# Copyright 2023 The MediaPipe Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""MediaPipe face landmarker task."""
+
+import dataclasses
+import enum
+from typing import Callable, Mapping, Optional, List
+
+import numpy as np
+
+from mediapipe.framework.formats import classification_pb2
+from mediapipe.framework.formats import landmark_pb2
+from mediapipe.framework.formats import matrix_data_pb2
+from mediapipe.python import packet_creator
+from mediapipe.python import packet_getter
+from mediapipe.python._framework_bindings import image as image_module
+from mediapipe.python._framework_bindings import packet as packet_module
+# pylint: disable=unused-import
+from mediapipe.tasks.cc.vision.face_geometry.proto import face_geometry_pb2
+# pylint: enable=unused-import
+from mediapipe.tasks.cc.vision.face_landmarker.proto import face_landmarker_graph_options_pb2
+from mediapipe.tasks.python.components.containers import category as category_module
+from mediapipe.tasks.python.components.containers import landmark as landmark_module
+from mediapipe.tasks.python.core import base_options as base_options_module
+from mediapipe.tasks.python.core import task_info as task_info_module
+from mediapipe.tasks.python.core.optional_dependencies import doc_controls
+from mediapipe.tasks.python.vision.core import base_vision_task_api
+from mediapipe.tasks.python.vision.core import image_processing_options as image_processing_options_module
+from mediapipe.tasks.python.vision.core import vision_task_running_mode as running_mode_module
+
+_BaseOptions = base_options_module.BaseOptions
+_FaceLandmarkerGraphOptionsProto = (
+    face_landmarker_graph_options_pb2.FaceLandmarkerGraphOptions
+)
+_LayoutEnum = matrix_data_pb2.MatrixData.Layout
+_RunningMode = running_mode_module.VisionTaskRunningMode
+_ImageProcessingOptions = image_processing_options_module.ImageProcessingOptions
+_TaskInfo = task_info_module.TaskInfo
+
+_IMAGE_IN_STREAM_NAME = 'image_in'
+_IMAGE_OUT_STREAM_NAME = 'image_out'
+_IMAGE_TAG = 'IMAGE'
+_NORM_RECT_STREAM_NAME = 'norm_rect_in'
+_NORM_RECT_TAG = 'NORM_RECT'
+_NORM_LANDMARKS_STREAM_NAME = 'norm_landmarks'
+_NORM_LANDMARKS_TAG = 'NORM_LANDMARKS'
+_BLENDSHAPES_STREAM_NAME = 'blendshapes'
+_BLENDSHAPES_TAG = 'BLENDSHAPES'
+_FACE_GEOMETRY_STREAM_NAME = 'face_geometry'
+_FACE_GEOMETRY_TAG = 'FACE_GEOMETRY'
+_TASK_GRAPH_NAME = 'mediapipe.tasks.vision.face_landmarker.FaceLandmarkerGraph'
+_MICRO_SECONDS_PER_MILLISECOND = 1000
+
+
+class Blendshapes(enum.IntEnum):
+  """The 52 blendshape coefficients."""
+
+  NEUTRAL = 0
+  BROW_DOWN_LEFT = 1
+  BROW_DOWN_RIGHT = 2
+  BROW_INNER_UP = 3
+  BROW_OUTER_UP_LEFT = 4
+  BROW_OUTER_UP_RIGHT = 5
+  CHEEK_PUFF = 6
+  CHEEK_SQUINT_LEFT = 7
+  CHEEK_SQUINT_RIGHT = 8
+  EYE_BLINK_LEFT = 9
+  EYE_BLINK_RIGHT = 10
+  EYE_LOOK_DOWN_LEFT = 11
+  EYE_LOOK_DOWN_RIGHT = 12
+  EYE_LOOK_IN_LEFT = 13
+  EYE_LOOK_IN_RIGHT = 14
+  EYE_LOOK_OUT_LEFT = 15
+  EYE_LOOK_OUT_RIGHT = 16
+  EYE_LOOK_UP_LEFT = 17
+  EYE_LOOK_UP_RIGHT = 18
+  EYE_SQUINT_LEFT = 19
+  EYE_SQUINT_RIGHT = 20
+  EYE_WIDE_LEFT = 21
+  EYE_WIDE_RIGHT = 22
+  JAW_FORWARD = 23
+  JAW_LEFT = 24
+  JAW_OPEN = 25
+  JAW_RIGHT = 26
+  MOUTH_CLOSE = 27
+  MOUTH_DIMPLE_LEFT = 28
+  MOUTH_DIMPLE_RIGHT = 29
+  MOUTH_FROWN_LEFT = 30
+  MOUTH_FROWN_RIGHT = 31
+  MOUTH_FUNNEL = 32
+  MOUTH_LEFT = 33
+  MOUTH_LOWER_DOWN_LEFT = 34
+  MOUTH_LOWER_DOWN_RIGHT = 35
+  MOUTH_PRESS_LEFT = 36
+  MOUTH_PRESS_RIGHT = 37
+  MOUTH_PUCKER = 38
+  MOUTH_RIGHT = 39
+  MOUTH_ROLL_LOWER = 40
+  MOUTH_ROLL_UPPER = 41
+  MOUTH_SHRUG_LOWER = 42
+  MOUTH_SHRUG_UPPER = 43
+  MOUTH_SMILE_LEFT = 44
+  MOUTH_SMILE_RIGHT = 45
+  MOUTH_STRETCH_LEFT = 46
+  MOUTH_STRETCH_RIGHT = 47
+  MOUTH_UPPER_UP_LEFT = 48
+  MOUTH_UPPER_UP_RIGHT = 49
+  NOSE_SNEER_LEFT = 50
+  NOSE_SNEER_RIGHT = 51
+
+
+class FaceLandmarksConnections:
+  """The connections between face landmarks."""
+
+  @dataclasses.dataclass
+  class Connection:
+    """The connection class for face landmarks."""
+
+    start: int
+    end: int
+
+  FACE_LANDMARKS_LIPS: List[Connection] = [
+      Connection(61, 146),
+      Connection(146, 91),
+      Connection(91, 181),
+      Connection(181, 84),
+      Connection(84, 17),
+      Connection(17, 314),
+      Connection(314, 405),
+      Connection(405, 321),
+      Connection(321, 375),
+      Connection(375, 291),
+      Connection(61, 185),
+      Connection(185, 40),
+      Connection(40, 39),
+      Connection(39, 37),
+      Connection(37, 0),
+      Connection(0, 267),
+      Connection(267, 269),
+      Connection(269, 270),
+      Connection(270, 409),
+      Connection(409, 291),
+      Connection(78, 95),
+      Connection(95, 88),
+      Connection(88, 178),
+      Connection(178, 87),
+      Connection(87, 14),
+      Connection(14, 317),
+      Connection(317, 402),
+      Connection(402, 318),
+      Connection(318, 324),
+      Connection(324, 308),
+      Connection(78, 191),
+      Connection(191, 80),
+      Connection(80, 81),
+      Connection(81, 82),
+      Connection(82, 13),
+      Connection(13, 312),
+      Connection(312, 311),
+      Connection(311, 310),
+      Connection(310, 415),
+      Connection(415, 308),
+  ]
+
+  FACE_LANDMARKS_LEFT_EYE: List[Connection] = [
+      Connection(263, 249),
+      Connection(249, 390),
+      Connection(390, 373),
+      Connection(373, 374),
+      Connection(374, 380),
+      Connection(380, 381),
+      Connection(381, 382),
+      Connection(382, 362),
+      Connection(263, 466),
+      Connection(466, 388),
+      Connection(388, 387),
+      Connection(387, 386),
+      Connection(386, 385),
+      Connection(385, 384),
+      Connection(384, 398),
+      Connection(398, 362),
+  ]
+
+  FACE_LANDMARKS_LEFT_EYEBROW: List[Connection] = [
+      Connection(276, 283),
+      Connection(283, 282),
+      Connection(282, 295),
+      Connection(295, 285),
+      Connection(300, 293),
+      Connection(293, 334),
+      Connection(334, 296),
+      Connection(296, 336),
+  ]
+
+  FACE_LANDMARKS_LEFT_IRIS: List[Connection] = [
+      Connection(474, 475),
+      Connection(475, 476),
+      Connection(476, 477),
+      Connection(477, 474),
+  ]
+
+  FACE_LANDMARKS_RIGHT_EYE: List[Connection] = [
+      Connection(33, 7),
+      Connection(7, 163),
+      Connection(163, 144),
+      Connection(144, 145),
+      Connection(145, 153),
+      Connection(153, 154),
+      Connection(154, 155),
+      Connection(155, 133),
+      Connection(33, 246),
+      Connection(246, 161),
+      Connection(161, 160),
+      Connection(160, 159),
+      Connection(159, 158),
+      Connection(158, 157),
+      Connection(157, 173),
+      Connection(173, 133),
+  ]
+
+  FACE_LANDMARKS_RIGHT_EYEBROW: List[Connection] = [
+      Connection(46, 53),
+      Connection(53, 52),
+      Connection(52, 65),
+      Connection(65, 55),
+      Connection(70, 63),
+      Connection(63, 105),
+      Connection(105, 66),
+      Connection(66, 107),
+  ]
+
+  FACE_LANDMARKS_RIGHT_IRIS: List[Connection] = [
+      Connection(469, 470),
+      Connection(470, 471),
+      Connection(471, 472),
+      Connection(472, 469),
+  ]
+
+  FACE_LANDMARKS_FACE_OVAL: List[Connection] = [
+      Connection(10, 338),
+      Connection(338, 297),
+      Connection(297, 332),
+      Connection(332, 284),
+      Connection(284, 251),
+      Connection(251, 389),
+      Connection(389, 356),
+      Connection(356, 454),
+      Connection(454, 323),
+      Connection(323, 361),
+      Connection(361, 288),
+      Connection(288, 397),
+      Connection(397, 365),
+      Connection(365, 379),
+      Connection(379, 378),
+      Connection(378, 400),
+      Connection(400, 377),
+      Connection(377, 152),
+      Connection(152, 148),
+      Connection(148, 176),
+      Connection(176, 149),
+      Connection(149, 150),
+      Connection(150, 136),
+      Connection(136, 172),
+      Connection(172, 58),
+      Connection(58, 132),
+      Connection(132, 93),
+      Connection(93, 234),
+      Connection(234, 127),
+      Connection(127, 162),
+      Connection(162, 21),
+      Connection(21, 54),
+      Connection(54, 103),
+      Connection(103, 67),
+      Connection(67, 109),
+      Connection(109, 10),
+  ]
+
+  FACE_LANDMARKS_CONTOURS: List[Connection] = (
+      FACE_LANDMARKS_LIPS
+      + FACE_LANDMARKS_LEFT_EYE
+      + FACE_LANDMARKS_LEFT_EYEBROW
+      + FACE_LANDMARKS_RIGHT_EYE
+      + FACE_LANDMARKS_RIGHT_EYEBROW
+      + FACE_LANDMARKS_FACE_OVAL
+  )
+
+  FACE_LANDMARKS_TESSELATION: List[Connection] = [
+      Connection(127, 34),
+      Connection(34, 139),
+      Connection(139, 127),
+      Connection(11, 0),
+      Connection(0, 37),
+      Connection(37, 11),
+      Connection(232, 231),
+      Connection(231, 120),
+      Connection(120, 232),
+      Connection(72, 37),
+      Connection(37, 39),
+      Connection(39, 72),
+      Connection(128, 121),
+      Connection(121, 47),
+      Connection(47, 128),
+      Connection(232, 121),
+      Connection(121, 128),
+      Connection(128, 232),
+      Connection(104, 69),
+      Connection(69, 67),
+      Connection(67, 104),
+      Connection(175, 171),
+      Connection(171, 148),
+      Connection(148, 175),
+      Connection(118, 50),
+      Connection(50, 101),
+      Connection(101, 118),
+      Connection(73, 39),
+      Connection(39, 40),
+      Connection(40, 73),
+      Connection(9, 151),
+      Connection(151, 108),
+      Connection(108, 9),
+      Connection(48, 115),
+      Connection(115, 131),
+      Connection(131, 48),
+      Connection(194, 204),
+      Connection(204, 211),
+      Connection(211, 194),
+      Connection(74, 40),
+      Connection(40, 185),
+      Connection(185, 74),
+      Connection(80, 42),
+      Connection(42, 183),
+      Connection(183, 80),
+      Connection(40, 92),
+      Connection(92, 186),
+      Connection(186, 40),
+      Connection(230, 229),
+      Connection(229, 118),
+      Connection(118, 230),
+      Connection(202, 212),
+      Connection(212, 214),
+      Connection(214, 202),
+      Connection(83, 18),
+      Connection(18, 17),
+      Connection(17, 83),
+      Connection(76, 61),
+      Connection(61, 146),
+      Connection(146, 76),
+      Connection(160, 29),
+      Connection(29, 30),
+      Connection(30, 160),
+      Connection(56, 157),
+      Connection(157, 173),
+      Connection(173, 56),
+      Connection(106, 204),
+      Connection(204, 194),
+      Connection(194, 106),
+      Connection(135, 214),
+      Connection(214, 192),
+      Connection(192, 135),
+      Connection(203, 165),
+      Connection(165, 98),
+      Connection(98, 203),
+      Connection(21, 71),
+      Connection(71, 68),
+      Connection(68, 21),
+      Connection(51, 45),
+      Connection(45, 4),
+      Connection(4, 51),
+      Connection(144, 24),
+      Connection(24, 23),
+      Connection(23, 144),
+      Connection(77, 146),
+      Connection(146, 91),
+      Connection(91, 77),
+      Connection(205, 50),
+      Connection(50, 187),
+      Connection(187, 205),
+      Connection(201, 200),
+      Connection(200, 18),
+      Connection(18, 201),
+      Connection(91, 106),
+      Connection(106, 182),
+      Connection(182, 91),
+      Connection(90, 91),
+      Connection(91, 181),
+      Connection(181, 90),
+      Connection(85, 84),
+      Connection(84, 17),
+      Connection(17, 85),
+      Connection(206, 203),
+      Connection(203, 36),
+      Connection(36, 206),
+      Connection(148, 171),
+      Connection(171, 140),
+      Connection(140, 148),
+      Connection(92, 40),
+      Connection(40, 39),
+      Connection(39, 92),
+      Connection(193, 189),
+      Connection(189, 244),
+      Connection(244, 193),
+      Connection(159, 158),
+      Connection(158, 28),
+      Connection(28, 159),
+      Connection(247, 246),
+      Connection(246, 161),
+      Connection(161, 247),
+      Connection(236, 3),
+      Connection(3, 196),
+      Connection(196, 236),
+      Connection(54, 68),
+      Connection(68, 104),
+      Connection(104, 54),
+      Connection(193, 168),
+      Connection(168, 8),
+      Connection(8, 193),
+      Connection(117, 228),
+      Connection(228, 31),
+      Connection(31, 117),
+      Connection(189, 193),
+      Connection(193, 55),
+      Connection(55, 189),
+      Connection(98, 97),
+      Connection(97, 99),
+      Connection(99, 98),
+      Connection(126, 47),
+      Connection(47, 100),
+      Connection(100, 126),
+      Connection(166, 79),
+      Connection(79, 218),
+      Connection(218, 166),
+      Connection(155, 154),
+      Connection(154, 26),
+      Connection(26, 155),
+      Connection(209, 49),
+      Connection(49, 131),
+      Connection(131, 209),
+      Connection(135, 136),
+      Connection(136, 150),
+      Connection(150, 135),
+      Connection(47, 126),
+      Connection(126, 217),
+      Connection(217, 47),
+      Connection(223, 52),
+      Connection(52, 53),
+      Connection(53, 223),
+      Connection(45, 51),
+      Connection(51, 134),
+      Connection(134, 45),
+      Connection(211, 170),
+      Connection(170, 140),
+      Connection(140, 211),
+      Connection(67, 69),
+      Connection(69, 108),
+      Connection(108, 67),
+      Connection(43, 106),
+      Connection(106, 91),
+      Connection(91, 43),
+      Connection(230, 119),
+      Connection(119, 120),
+      Connection(120, 230),
+      Connection(226, 130),
+      Connection(130, 247),
+      Connection(247, 226),
+      Connection(63, 53),
+      Connection(53, 52),
+      Connection(52, 63),
+      Connection(238, 20),
+      Connection(20, 242),
+      Connection(242, 238),
+      Connection(46, 70),
+      Connection(70, 156),
+      Connection(156, 46),
+      Connection(78, 62),
+      Connection(62, 96),
+      Connection(96, 78),
+      Connection(46, 53),
+      Connection(53, 63),
+      Connection(63, 46),
+      Connection(143, 34),
+      Connection(34, 227),
+      Connection(227, 143),
+      Connection(123, 117),
+      Connection(117, 111),
+      Connection(111, 123),
+      Connection(44, 125),
+      Connection(125, 19),
+      Connection(19, 44),
+      Connection(236, 134),
+      Connection(134, 51),
+      Connection(51, 236),
+      Connection(216, 206),
+      Connection(206, 205),
+      Connection(205, 216),
+      Connection(154, 153),
+      Connection(153, 22),
+      Connection(22, 154),
+      Connection(39, 37),
+      Connection(37, 167),
+      Connection(167, 39),
+      Connection(200, 201),
+      Connection(201, 208),
+      Connection(208, 200),
+      Connection(36, 142),
+      Connection(142, 100),
+      Connection(100, 36),
+      Connection(57, 212),
+      Connection(212, 202),
+      Connection(202, 57),
+      Connection(20, 60),
+      Connection(60, 99),
+      Connection(99, 20),
+      Connection(28, 158),
+      Connection(158, 157),
+      Connection(157, 28),
+      Connection(35, 226),
+      Connection(226, 113),
+      Connection(113, 35),
+      Connection(160, 159),
+      Connection(159, 27),
+      Connection(27, 160),
+      Connection(204, 202),
+      Connection(202, 210),
+      Connection(210, 204),
+      Connection(113, 225),
+      Connection(225, 46),
+      Connection(46, 113),
+      Connection(43, 202),
+      Connection(202, 204),
+      Connection(204, 43),
+      Connection(62, 76),
+      Connection(76, 77),
+      Connection(77, 62),
+      Connection(137, 123),
+      Connection(123, 116),
+      Connection(116, 137),
+      Connection(41, 38),
+      Connection(38, 72),
+      Connection(72, 41),
+      Connection(203, 129),
+      Connection(129, 142),
+      Connection(142, 203),
+      Connection(64, 98),
+      Connection(98, 240),
+      Connection(240, 64),
+      Connection(49, 102),
+      Connection(102, 64),
+      Connection(64, 49),
+      Connection(41, 73),
+      Connection(73, 74),
+      Connection(74, 41),
+      Connection(212, 216),
+      Connection(216, 207),
+      Connection(207, 212),
+      Connection(42, 74),
+      Connection(74, 184),
+      Connection(184, 42),
+      Connection(169, 170),
+      Connection(170, 211),
+      Connection(211, 169),
+      Connection(170, 149),
+      Connection(149, 176),
+      Connection(176, 170),
+      Connection(105, 66),
+      Connection(66, 69),
+      Connection(69, 105),
+      Connection(122, 6),
+      Connection(6, 168),
+      Connection(168, 122),
+      Connection(123, 147),
+      Connection(147, 187),
+      Connection(187, 123),
+      Connection(96, 77),
+      Connection(77, 90),
+      Connection(90, 96),
+      Connection(65, 55),
+      Connection(55, 107),
+      Connection(107, 65),
+      Connection(89, 90),
+      Connection(90, 180),
+      Connection(180, 89),
+      Connection(101, 100),
+      Connection(100, 120),
+      Connection(120, 101),
+      Connection(63, 105),
+      Connection(105, 104),
+      Connection(104, 63),
+      Connection(93, 137),
+      Connection(137, 227),
+      Connection(227, 93),
+      Connection(15, 86),
+      Connection(86, 85),
+      Connection(85, 15),
+      Connection(129, 102),
+      Connection(102, 49),
+      Connection(49, 129),
+      Connection(14, 87),
+      Connection(87, 86),
+      Connection(86, 14),
+      Connection(55, 8),
+      Connection(8, 9),
+      Connection(9, 55),
+      Connection(100, 47),
+      Connection(47, 121),
+      Connection(121, 100),
+      Connection(145, 23),
+      Connection(23, 22),
+      Connection(22, 145),
+      Connection(88, 89),
+      Connection(89, 179),
+      Connection(179, 88),
+      Connection(6, 122),
+      Connection(122, 196),
+      Connection(196, 6),
+      Connection(88, 95),
+      Connection(95, 96),
+      Connection(96, 88),
+      Connection(138, 172),
+      Connection(172, 136),
+      Connection(136, 138),
+      Connection(215, 58),
+      Connection(58, 172),
+      Connection(172, 215),
+      Connection(115, 48),
+      Connection(48, 219),
+      Connection(219, 115),
+      Connection(42, 80),
+      Connection(80, 81),
+      Connection(81, 42),
+      Connection(195, 3),
+      Connection(3, 51),
+      Connection(51, 195),
+      Connection(43, 146),
+      Connection(146, 61),
+      Connection(61, 43),
+      Connection(171, 175),
+      Connection(175, 199),
+      Connection(199, 171),
+      Connection(81, 82),
+      Connection(82, 38),
+      Connection(38, 81),
+      Connection(53, 46),
+      Connection(46, 225),
+      Connection(225, 53),
+      Connection(144, 163),
+      Connection(163, 110),
+      Connection(110, 144),
+      Connection(52, 65),
+      Connection(65, 66),
+      Connection(66, 52),
+      Connection(229, 228),
+      Connection(228, 117),
+      Connection(117, 229),
+      Connection(34, 127),
+      Connection(127, 234),
+      Connection(234, 34),
+      Connection(107, 108),
+      Connection(108, 69),
+      Connection(69, 107),
+      Connection(109, 108),
+      Connection(108, 151),
+      Connection(151, 109),
+      Connection(48, 64),
+      Connection(64, 235),
+      Connection(235, 48),
+      Connection(62, 78),
+      Connection(78, 191),
+      Connection(191, 62),
+      Connection(129, 209),
+      Connection(209, 126),
+      Connection(126, 129),
+      Connection(111, 35),
+      Connection(35, 143),
+      Connection(143, 111),
+      Connection(117, 123),
+      Connection(123, 50),
+      Connection(50, 117),
+      Connection(222, 65),
+      Connection(65, 52),
+      Connection(52, 222),
+      Connection(19, 125),
+      Connection(125, 141),
+      Connection(141, 19),
+      Connection(221, 55),
+      Connection(55, 65),
+      Connection(65, 221),
+      Connection(3, 195),
+      Connection(195, 197),
+      Connection(197, 3),
+      Connection(25, 7),
+      Connection(7, 33),
+      Connection(33, 25),
+      Connection(220, 237),
+      Connection(237, 44),
+      Connection(44, 220),
+      Connection(70, 71),
+      Connection(71, 139),
+      Connection(139, 70),
+      Connection(122, 193),
+      Connection(193, 245),
+      Connection(245, 122),
+      Connection(247, 130),
+      Connection(130, 33),
+      Connection(33, 247),
+      Connection(71, 21),
+      Connection(21, 162),
+      Connection(162, 71),
+      Connection(170, 169),
+      Connection(169, 150),
+      Connection(150, 170),
+      Connection(188, 174),
+      Connection(174, 196),
+      Connection(196, 188),
+      Connection(216, 186),
+      Connection(186, 92),
+      Connection(92, 216),
+      Connection(2, 97),
+      Connection(97, 167),
+      Connection(167, 2),
+      Connection(141, 125),
+      Connection(125, 241),
+      Connection(241, 141),
+      Connection(164, 167),
+      Connection(167, 37),
+      Connection(37, 164),
+      Connection(72, 38),
+      Connection(38, 12),
+      Connection(12, 72),
+      Connection(38, 82),
+      Connection(82, 13),
+      Connection(13, 38),
+      Connection(63, 68),
+      Connection(68, 71),
+      Connection(71, 63),
+      Connection(226, 35),
+      Connection(35, 111),
+      Connection(111, 226),
+      Connection(101, 50),
+      Connection(50, 205),
+      Connection(205, 101),
+      Connection(206, 92),
+      Connection(92, 165),
+      Connection(165, 206),
+      Connection(209, 198),
+      Connection(198, 217),
+      Connection(217, 209),
+      Connection(165, 167),
+      Connection(167, 97),
+      Connection(97, 165),
+      Connection(220, 115),
+      Connection(115, 218),
+      Connection(218, 220),
+      Connection(133, 112),
+      Connection(112, 243),
+      Connection(243, 133),
+      Connection(239, 238),
+      Connection(238, 241),
+      Connection(241, 239),
+      Connection(214, 135),
+      Connection(135, 169),
+      Connection(169, 214),
+      Connection(190, 173),
+      Connection(173, 133),
+      Connection(133, 190),
+      Connection(171, 208),
+      Connection(208, 32),
+      Connection(32, 171),
+      Connection(125, 44),
+      Connection(44, 237),
+      Connection(237, 125),
+      Connection(86, 87),
+      Connection(87, 178),
+      Connection(178, 86),
+      Connection(85, 86),
+      Connection(86, 179),
+      Connection(179, 85),
+      Connection(84, 85),
+      Connection(85, 180),
+      Connection(180, 84),
+      Connection(83, 84),
+      Connection(84, 181),
+      Connection(181, 83),
+      Connection(201, 83),
+      Connection(83, 182),
+      Connection(182, 201),
+      Connection(137, 93),
+      Connection(93, 132),
+      Connection(132, 137),
+      Connection(76, 62),
+      Connection(62, 183),
+      Connection(183, 76),
+      Connection(61, 76),
+      Connection(76, 184),
+      Connection(184, 61),
+      Connection(57, 61),
+      Connection(61, 185),
+      Connection(185, 57),
+      Connection(212, 57),
+      Connection(57, 186),
+      Connection(186, 212),
+      Connection(214, 207),
+      Connection(207, 187),
+      Connection(187, 214),
+      Connection(34, 143),
+      Connection(143, 156),
+      Connection(156, 34),
+      Connection(79, 239),
+      Connection(239, 237),
+      Connection(237, 79),
+      Connection(123, 137),
+      Connection(137, 177),
+      Connection(177, 123),
+      Connection(44, 1),
+      Connection(1, 4),
+      Connection(4, 44),
+      Connection(201, 194),
+      Connection(194, 32),
+      Connection(32, 201),
+      Connection(64, 102),
+      Connection(102, 129),
+      Connection(129, 64),
+      Connection(213, 215),
+      Connection(215, 138),
+      Connection(138, 213),
+      Connection(59, 166),
+      Connection(166, 219),
+      Connection(219, 59),
+      Connection(242, 99),
+      Connection(99, 97),
+      Connection(97, 242),
+      Connection(2, 94),
+      Connection(94, 141),
+      Connection(141, 2),
+      Connection(75, 59),
+      Connection(59, 235),
+      Connection(235, 75),
+      Connection(24, 110),
+      Connection(110, 228),
+      Connection(228, 24),
+      Connection(25, 130),
+      Connection(130, 226),
+      Connection(226, 25),
+      Connection(23, 24),
+      Connection(24, 229),
+      Connection(229, 23),
+      Connection(22, 23),
+      Connection(23, 230),
+      Connection(230, 22),
+      Connection(26, 22),
+      Connection(22, 231),
+      Connection(231, 26),
+      Connection(112, 26),
+      Connection(26, 232),
+      Connection(232, 112),
+      Connection(189, 190),
+      Connection(190, 243),
+      Connection(243, 189),
+      Connection(221, 56),
+      Connection(56, 190),
+      Connection(190, 221),
+      Connection(28, 56),
+      Connection(56, 221),
+      Connection(221, 28),
+      Connection(27, 28),
+      Connection(28, 222),
+      Connection(222, 27),
+      Connection(29, 27),
+      Connection(27, 223),
+      Connection(223, 29),
+      Connection(30, 29),
+      Connection(29, 224),
+      Connection(224, 30),
+      Connection(247, 30),
+      Connection(30, 225),
+      Connection(225, 247),
+      Connection(238, 79),
+      Connection(79, 20),
+      Connection(20, 238),
+      Connection(166, 59),
+      Connection(59, 75),
+      Connection(75, 166),
+      Connection(60, 75),
+      Connection(75, 240),
+      Connection(240, 60),
+      Connection(147, 177),
+      Connection(177, 215),
+      Connection(215, 147),
+      Connection(20, 79),
+      Connection(79, 166),
+      Connection(166, 20),
+      Connection(187, 147),
+      Connection(147, 213),
+      Connection(213, 187),
+      Connection(112, 233),
+      Connection(233, 244),
+      Connection(244, 112),
+      Connection(233, 128),
+      Connection(128, 245),
+      Connection(245, 233),
+      Connection(128, 114),
+      Connection(114, 188),
+      Connection(188, 128),
+      Connection(114, 217),
+      Connection(217, 174),
+      Connection(174, 114),
+      Connection(131, 115),
+      Connection(115, 220),
+      Connection(220, 131),
+      Connection(217, 198),
+      Connection(198, 236),
+      Connection(236, 217),
+      Connection(198, 131),
+      Connection(131, 134),
+      Connection(134, 198),
+      Connection(177, 132),
+      Connection(132, 58),
+      Connection(58, 177),
+      Connection(143, 35),
+      Connection(35, 124),
+      Connection(124, 143),
+      Connection(110, 163),
+      Connection(163, 7),
+      Connection(7, 110),
+      Connection(228, 110),
+      Connection(110, 25),
+      Connection(25, 228),
+      Connection(356, 389),
+      Connection(389, 368),
+      Connection(368, 356),
+      Connection(11, 302),
+      Connection(302, 267),
+      Connection(267, 11),
+      Connection(452, 350),
+      Connection(350, 349),
+      Connection(349, 452),
+      Connection(302, 303),
+      Connection(303, 269),
+      Connection(269, 302),
+      Connection(357, 343),
+      Connection(343, 277),
+      Connection(277, 357),
+      Connection(452, 453),
+      Connection(453, 357),
+      Connection(357, 452),
+      Connection(333, 332),
+      Connection(332, 297),
+      Connection(297, 333),
+      Connection(175, 152),
+      Connection(152, 377),
+      Connection(377, 175),
+      Connection(347, 348),
+      Connection(348, 330),
+      Connection(330, 347),
+      Connection(303, 304),
+      Connection(304, 270),
+      Connection(270, 303),
+      Connection(9, 336),
+      Connection(336, 337),
+      Connection(337, 9),
+      Connection(278, 279),
+      Connection(279, 360),
+      Connection(360, 278),
+      Connection(418, 262),
+      Connection(262, 431),
+      Connection(431, 418),
+      Connection(304, 408),
+      Connection(408, 409),
+      Connection(409, 304),
+      Connection(310, 415),
+      Connection(415, 407),
+      Connection(407, 310),
+      Connection(270, 409),
+      Connection(409, 410),
+      Connection(410, 270),
+      Connection(450, 348),
+      Connection(348, 347),
+      Connection(347, 450),
+      Connection(422, 430),
+      Connection(430, 434),
+      Connection(434, 422),
+      Connection(313, 314),
+      Connection(314, 17),
+      Connection(17, 313),
+      Connection(306, 307),
+      Connection(307, 375),
+      Connection(375, 306),
+      Connection(387, 388),
+      Connection(388, 260),
+      Connection(260, 387),
+      Connection(286, 414),
+      Connection(414, 398),
+      Connection(398, 286),
+      Connection(335, 406),
+      Connection(406, 418),
+      Connection(418, 335),
+      Connection(364, 367),
+      Connection(367, 416),
+      Connection(416, 364),
+      Connection(423, 358),
+      Connection(358, 327),
+      Connection(327, 423),
+      Connection(251, 284),
+      Connection(284, 298),
+      Connection(298, 251),
+      Connection(281, 5),
+      Connection(5, 4),
+      Connection(4, 281),
+      Connection(373, 374),
+      Connection(374, 253),
+      Connection(253, 373),
+      Connection(307, 320),
+      Connection(320, 321),
+      Connection(321, 307),
+      Connection(425, 427),
+      Connection(427, 411),
+      Connection(411, 425),
+      Connection(421, 313),
+      Connection(313, 18),
+      Connection(18, 421),
+      Connection(321, 405),
+      Connection(405, 406),
+      Connection(406, 321),
+      Connection(320, 404),
+      Connection(404, 405),
+      Connection(405, 320),
+      Connection(315, 16),
+      Connection(16, 17),
+      Connection(17, 315),
+      Connection(426, 425),
+      Connection(425, 266),
+      Connection(266, 426),
+      Connection(377, 400),
+      Connection(400, 369),
+      Connection(369, 377),
+      Connection(322, 391),
+      Connection(391, 269),
+      Connection(269, 322),
+      Connection(417, 465),
+      Connection(465, 464),
+      Connection(464, 417),
+      Connection(386, 257),
+      Connection(257, 258),
+      Connection(258, 386),
+      Connection(466, 260),
+      Connection(260, 388),
+      Connection(388, 466),
+      Connection(456, 399),
+      Connection(399, 419),
+      Connection(419, 456),
+      Connection(284, 332),
+      Connection(332, 333),
+      Connection(333, 284),
+      Connection(417, 285),
+      Connection(285, 8),
+      Connection(8, 417),
+      Connection(346, 340),
+      Connection(340, 261),
+      Connection(261, 346),
+      Connection(413, 441),
+      Connection(441, 285),
+      Connection(285, 413),
+      Connection(327, 460),
+      Connection(460, 328),
+      Connection(328, 327),
+      Connection(355, 371),
+      Connection(371, 329),
+      Connection(329, 355),
+      Connection(392, 439),
+      Connection(439, 438),
+      Connection(438, 392),
+      Connection(382, 341),
+      Connection(341, 256),
+      Connection(256, 382),
+      Connection(429, 420),
+      Connection(420, 360),
+      Connection(360, 429),
+      Connection(364, 394),
+      Connection(394, 379),
+      Connection(379, 364),
+      Connection(277, 343),
+      Connection(343, 437),
+      Connection(437, 277),
+      Connection(443, 444),
+      Connection(444, 283),
+      Connection(283, 443),
+      Connection(275, 440),
+      Connection(440, 363),
+      Connection(363, 275),
+      Connection(431, 262),
+      Connection(262, 369),
+      Connection(369, 431),
+      Connection(297, 338),
+      Connection(338, 337),
+      Connection(337, 297),
+      Connection(273, 375),
+      Connection(375, 321),
+      Connection(321, 273),
+      Connection(450, 451),
+      Connection(451, 349),
+      Connection(349, 450),
+      Connection(446, 342),
+      Connection(342, 467),
+      Connection(467, 446),
+      Connection(293, 334),
+      Connection(334, 282),
+      Connection(282, 293),
+      Connection(458, 461),
+      Connection(461, 462),
+      Connection(462, 458),
+      Connection(276, 353),
+      Connection(353, 383),
+      Connection(383, 276),
+      Connection(308, 324),
+      Connection(324, 325),
+      Connection(325, 308),
+      Connection(276, 300),
+      Connection(300, 293),
+      Connection(293, 276),
+      Connection(372, 345),
+      Connection(345, 447),
+      Connection(447, 372),
+      Connection(352, 345),
+      Connection(345, 340),
+      Connection(340, 352),
+      Connection(274, 1),
+      Connection(1, 19),
+      Connection(19, 274),
+      Connection(456, 248),
+      Connection(248, 281),
+      Connection(281, 456),
+      Connection(436, 427),
+      Connection(427, 425),
+      Connection(425, 436),
+      Connection(381, 256),
+      Connection(256, 252),
+      Connection(252, 381),
+      Connection(269, 391),
+      Connection(391, 393),
+      Connection(393, 269),
+      Connection(200, 199),
+      Connection(199, 428),
+      Connection(428, 200),
+      Connection(266, 330),
+      Connection(330, 329),
+      Connection(329, 266),
+      Connection(287, 273),
+      Connection(273, 422),
+      Connection(422, 287),
+      Connection(250, 462),
+      Connection(462, 328),
+      Connection(328, 250),
+      Connection(258, 286),
+      Connection(286, 384),
+      Connection(384, 258),
+      Connection(265, 353),
+      Connection(353, 342),
+      Connection(342, 265),
+      Connection(387, 259),
+      Connection(259, 257),
+      Connection(257, 387),
+      Connection(424, 431),
+      Connection(431, 430),
+      Connection(430, 424),
+      Connection(342, 353),
+      Connection(353, 276),
+      Connection(276, 342),
+      Connection(273, 335),
+      Connection(335, 424),
+      Connection(424, 273),
+      Connection(292, 325),
+      Connection(325, 307),
+      Connection(307, 292),
+      Connection(366, 447),
+      Connection(447, 345),
+      Connection(345, 366),
+      Connection(271, 303),
+      Connection(303, 302),
+      Connection(302, 271),
+      Connection(423, 266),
+      Connection(266, 371),
+      Connection(371, 423),
+      Connection(294, 455),
+      Connection(455, 460),
+      Connection(460, 294),
+      Connection(279, 278),
+      Connection(278, 294),
+      Connection(294, 279),
+      Connection(271, 272),
+      Connection(272, 304),
+      Connection(304, 271),
+      Connection(432, 434),
+      Connection(434, 427),
+      Connection(427, 432),
+      Connection(272, 407),
+      Connection(407, 408),
+      Connection(408, 272),
+      Connection(394, 430),
+      Connection(430, 431),
+      Connection(431, 394),
+      Connection(395, 369),
+      Connection(369, 400),
+      Connection(400, 395),
+      Connection(334, 333),
+      Connection(333, 299),
+      Connection(299, 334),
+      Connection(351, 417),
+      Connection(417, 168),
+      Connection(168, 351),
+      Connection(352, 280),
+      Connection(280, 411),
+      Connection(411, 352),
+      Connection(325, 319),
+      Connection(319, 320),
+      Connection(320, 325),
+      Connection(295, 296),
+      Connection(296, 336),
+      Connection(336, 295),
+      Connection(319, 403),
+      Connection(403, 404),
+      Connection(404, 319),
+      Connection(330, 348),
+      Connection(348, 349),
+      Connection(349, 330),
+      Connection(293, 298),
+      Connection(298, 333),
+      Connection(333, 293),
+      Connection(323, 454),
+      Connection(454, 447),
+      Connection(447, 323),
+      Connection(15, 16),
+      Connection(16, 315),
+      Connection(315, 15),
+      Connection(358, 429),
+      Connection(429, 279),
+      Connection(279, 358),
+      Connection(14, 15),
+      Connection(15, 316),
+      Connection(316, 14),
+      Connection(285, 336),
+      Connection(336, 9),
+      Connection(9, 285),
+      Connection(329, 349),
+      Connection(349, 350),
+      Connection(350, 329),
+      Connection(374, 380),
+      Connection(380, 252),
+      Connection(252, 374),
+      Connection(318, 402),
+      Connection(402, 403),
+      Connection(403, 318),
+      Connection(6, 197),
+      Connection(197, 419),
+      Connection(419, 6),
+      Connection(318, 319),
+      Connection(319, 325),
+      Connection(325, 318),
+      Connection(367, 364),
+      Connection(364, 365),
+      Connection(365, 367),
+      Connection(435, 367),
+      Connection(367, 397),
+      Connection(397, 435),
+      Connection(344, 438),
+      Connection(438, 439),
+      Connection(439, 344),
+      Connection(272, 271),
+      Connection(271, 311),
+      Connection(311, 272),
+      Connection(195, 5),
+      Connection(5, 281),
+      Connection(281, 195),
+      Connection(273, 287),
+      Connection(287, 291),
+      Connection(291, 273),
+      Connection(396, 428),
+      Connection(428, 199),
+      Connection(199, 396),
+      Connection(311, 271),
+      Connection(271, 268),
+      Connection(268, 311),
+      Connection(283, 444),
+      Connection(444, 445),
+      Connection(445, 283),
+      Connection(373, 254),
+      Connection(254, 339),
+      Connection(339, 373),
+      Connection(282, 334),
+      Connection(334, 296),
+      Connection(296, 282),
+      Connection(449, 347),
+      Connection(347, 346),
+      Connection(346, 449),
+      Connection(264, 447),
+      Connection(447, 454),
+      Connection(454, 264),
+      Connection(336, 296),
+      Connection(296, 299),
+      Connection(299, 336),
+      Connection(338, 10),
+      Connection(10, 151),
+      Connection(151, 338),
+      Connection(278, 439),
+      Connection(439, 455),
+      Connection(455, 278),
+      Connection(292, 407),
+      Connection(407, 415),
+      Connection(415, 292),
+      Connection(358, 371),
+      Connection(371, 355),
+      Connection(355, 358),
+      Connection(340, 345),
+      Connection(345, 372),
+      Connection(372, 340),
+      Connection(346, 347),
+      Connection(347, 280),
+      Connection(280, 346),
+      Connection(442, 443),
+      Connection(443, 282),
+      Connection(282, 442),
+      Connection(19, 94),
+      Connection(94, 370),
+      Connection(370, 19),
+      Connection(441, 442),
+      Connection(442, 295),
+      Connection(295, 441),
+      Connection(248, 419),
+      Connection(419, 197),
+      Connection(197, 248),
+      Connection(263, 255),
+      Connection(255, 359),
+      Connection(359, 263),
+      Connection(440, 275),
+      Connection(275, 274),
+      Connection(274, 440),
+      Connection(300, 383),
+      Connection(383, 368),
+      Connection(368, 300),
+      Connection(351, 412),
+      Connection(412, 465),
+      Connection(465, 351),
+      Connection(263, 467),
+      Connection(467, 466),
+      Connection(466, 263),
+      Connection(301, 368),
+      Connection(368, 389),
+      Connection(389, 301),
+      Connection(395, 378),
+      Connection(378, 379),
+      Connection(379, 395),
+      Connection(412, 351),
+      Connection(351, 419),
+      Connection(419, 412),
+      Connection(436, 426),
+      Connection(426, 322),
+      Connection(322, 436),
+      Connection(2, 164),
+      Connection(164, 393),
+      Connection(393, 2),
+      Connection(370, 462),
+      Connection(462, 461),
+      Connection(461, 370),
+      Connection(164, 0),
+      Connection(0, 267),
+      Connection(267, 164),
+      Connection(302, 11),
+      Connection(11, 12),
+      Connection(12, 302),
+      Connection(268, 12),
+      Connection(12, 13),
+      Connection(13, 268),
+      Connection(293, 300),
+      Connection(300, 301),
+      Connection(301, 293),
+      Connection(446, 261),
+      Connection(261, 340),
+      Connection(340, 446),
+      Connection(330, 266),
+      Connection(266, 425),
+      Connection(425, 330),
+      Connection(426, 423),
+      Connection(423, 391),
+      Connection(391, 426),
+      Connection(429, 355),
+      Connection(355, 437),
+      Connection(437, 429),
+      Connection(391, 327),
+      Connection(327, 326),
+      Connection(326, 391),
+      Connection(440, 457),
+      Connection(457, 438),
+      Connection(438, 440),
+      Connection(341, 382),
+      Connection(382, 362),
+      Connection(362, 341),
+      Connection(459, 457),
+      Connection(457, 461),
+      Connection(461, 459),
+      Connection(434, 430),
+      Connection(430, 394),
+      Connection(394, 434),
+      Connection(414, 463),
+      Connection(463, 362),
+      Connection(362, 414),
+      Connection(396, 369),
+      Connection(369, 262),
+      Connection(262, 396),
+      Connection(354, 461),
+      Connection(461, 457),
+      Connection(457, 354),
+      Connection(316, 403),
+      Connection(403, 402),
+      Connection(402, 316),
+      Connection(315, 404),
+      Connection(404, 403),
+      Connection(403, 315),
+      Connection(314, 405),
+      Connection(405, 404),
+      Connection(404, 314),
+      Connection(313, 406),
+      Connection(406, 405),
+      Connection(405, 313),
+      Connection(421, 418),
+      Connection(418, 406),
+      Connection(406, 421),
+      Connection(366, 401),
+      Connection(401, 361),
+      Connection(361, 366),
+      Connection(306, 408),
+      Connection(408, 407),
+      Connection(407, 306),
+      Connection(291, 409),
+      Connection(409, 408),
+      Connection(408, 291),
+      Connection(287, 410),
+      Connection(410, 409),
+      Connection(409, 287),
+      Connection(432, 436),
+      Connection(436, 410),
+      Connection(410, 432),
+      Connection(434, 416),
+      Connection(416, 411),
+      Connection(411, 434),
+      Connection(264, 368),
+      Connection(368, 383),
+      Connection(383, 264),
+      Connection(309, 438),
+      Connection(438, 457),
+      Connection(457, 309),
+      Connection(352, 376),
+      Connection(376, 401),
+      Connection(401, 352),
+      Connection(274, 275),
+      Connection(275, 4),
+      Connection(4, 274),
+      Connection(421, 428),
+      Connection(428, 262),
+      Connection(262, 421),
+      Connection(294, 327),
+      Connection(327, 358),
+      Connection(358, 294),
+      Connection(433, 416),
+      Connection(416, 367),
+      Connection(367, 433),
+      Connection(289, 455),
+      Connection(455, 439),
+      Connection(439, 289),
+      Connection(462, 370),
+      Connection(370, 326),
+      Connection(326, 462),
+      Connection(2, 326),
+      Connection(326, 370),
+      Connection(370, 2),
+      Connection(305, 460),
+      Connection(460, 455),
+      Connection(455, 305),
+      Connection(254, 449),
+      Connection(449, 448),
+      Connection(448, 254),
+      Connection(255, 261),
+      Connection(261, 446),
+      Connection(446, 255),
+      Connection(253, 450),
+      Connection(450, 449),
+      Connection(449, 253),
+      Connection(252, 451),
+      Connection(451, 450),
+      Connection(450, 252),
+      Connection(256, 452),
+      Connection(452, 451),
+      Connection(451, 256),
+      Connection(341, 453),
+      Connection(453, 452),
+      Connection(452, 341),
+      Connection(413, 464),
+      Connection(464, 463),
+      Connection(463, 413),
+      Connection(441, 413),
+      Connection(413, 414),
+      Connection(414, 441),
+      Connection(258, 442),
+      Connection(442, 441),
+      Connection(441, 258),
+      Connection(257, 443),
+      Connection(443, 442),
+      Connection(442, 257),
+      Connection(259, 444),
+      Connection(444, 443),
+      Connection(443, 259),
+      Connection(260, 445),
+      Connection(445, 444),
+      Connection(444, 260),
+      Connection(467, 342),
+      Connection(342, 445),
+      Connection(445, 467),
+      Connection(459, 458),
+      Connection(458, 250),
+      Connection(250, 459),
+      Connection(289, 392),
+      Connection(392, 290),
+      Connection(290, 289),
+      Connection(290, 328),
+      Connection(328, 460),
+      Connection(460, 290),
+      Connection(376, 433),
+      Connection(433, 435),
+      Connection(435, 376),
+      Connection(250, 290),
+      Connection(290, 392),
+      Connection(392, 250),
+      Connection(411, 416),
+      Connection(416, 433),
+      Connection(433, 411),
+      Connection(341, 463),
+      Connection(463, 464),
+      Connection(464, 341),
+      Connection(453, 464),
+      Connection(464, 465),
+      Connection(465, 453),
+      Connection(357, 465),
+      Connection(465, 412),
+      Connection(412, 357),
+      Connection(343, 412),
+      Connection(412, 399),
+      Connection(399, 343),
+      Connection(360, 363),
+      Connection(363, 440),
+      Connection(440, 360),
+      Connection(437, 399),
+      Connection(399, 456),
+      Connection(456, 437),
+      Connection(420, 456),
+      Connection(456, 363),
+      Connection(363, 420),
+      Connection(401, 435),
+      Connection(435, 288),
+      Connection(288, 401),
+      Connection(372, 383),
+      Connection(383, 353),
+      Connection(353, 372),
+      Connection(339, 255),
+      Connection(255, 249),
+      Connection(249, 339),
+      Connection(448, 261),
+      Connection(261, 255),
+      Connection(255, 448),
+      Connection(133, 243),
+      Connection(243, 190),
+      Connection(190, 133),
+      Connection(133, 155),
+      Connection(155, 112),
+      Connection(112, 133),
+      Connection(33, 246),
+      Connection(246, 247),
+      Connection(247, 33),
+      Connection(33, 130),
+      Connection(130, 25),
+      Connection(25, 33),
+      Connection(398, 384),
+      Connection(384, 286),
+      Connection(286, 398),
+      Connection(362, 398),
+      Connection(398, 414),
+      Connection(414, 362),
+      Connection(362, 463),
+      Connection(463, 341),
+      Connection(341, 362),
+      Connection(263, 359),
+      Connection(359, 467),
+      Connection(467, 263),
+      Connection(263, 249),
+      Connection(249, 255),
+      Connection(255, 263),
+      Connection(466, 467),
+      Connection(467, 260),
+      Connection(260, 466),
+      Connection(75, 60),
+      Connection(60, 166),
+      Connection(166, 75),
+      Connection(238, 239),
+      Connection(239, 79),
+      Connection(79, 238),
+      Connection(162, 127),
+      Connection(127, 139),
+      Connection(139, 162),
+      Connection(72, 11),
+      Connection(11, 37),
+      Connection(37, 72),
+      Connection(121, 232),
+      Connection(232, 120),
+      Connection(120, 121),
+      Connection(73, 72),
+      Connection(72, 39),
+      Connection(39, 73),
+      Connection(114, 128),
+      Connection(128, 47),
+      Connection(47, 114),
+      Connection(233, 232),
+      Connection(232, 128),
+      Connection(128, 233),
+      Connection(103, 104),
+      Connection(104, 67),
+      Connection(67, 103),
+      Connection(152, 175),
+      Connection(175, 148),
+      Connection(148, 152),
+      Connection(119, 118),
+      Connection(118, 101),
+      Connection(101, 119),
+      Connection(74, 73),
+      Connection(73, 40),
+      Connection(40, 74),
+      Connection(107, 9),
+      Connection(9, 108),
+      Connection(108, 107),
+      Connection(49, 48),
+      Connection(48, 131),
+      Connection(131, 49),
+      Connection(32, 194),
+      Connection(194, 211),
+      Connection(211, 32),
+      Connection(184, 74),
+      Connection(74, 185),
+      Connection(185, 184),
+      Connection(191, 80),
+      Connection(80, 183),
+      Connection(183, 191),
+      Connection(185, 40),
+      Connection(40, 186),
+      Connection(186, 185),
+      Connection(119, 230),
+      Connection(230, 118),
+      Connection(118, 119),
+      Connection(210, 202),
+      Connection(202, 214),
+      Connection(214, 210),
+      Connection(84, 83),
+      Connection(83, 17),
+      Connection(17, 84),
+      Connection(77, 76),
+      Connection(76, 146),
+      Connection(146, 77),
+      Connection(161, 160),
+      Connection(160, 30),
+      Connection(30, 161),
+      Connection(190, 56),
+      Connection(56, 173),
+      Connection(173, 190),
+      Connection(182, 106),
+      Connection(106, 194),
+      Connection(194, 182),
+      Connection(138, 135),
+      Connection(135, 192),
+      Connection(192, 138),
+      Connection(129, 203),
+      Connection(203, 98),
+      Connection(98, 129),
+      Connection(54, 21),
+      Connection(21, 68),
+      Connection(68, 54),
+      Connection(5, 51),
+      Connection(51, 4),
+      Connection(4, 5),
+      Connection(145, 144),
+      Connection(144, 23),
+      Connection(23, 145),
+      Connection(90, 77),
+      Connection(77, 91),
+      Connection(91, 90),
+      Connection(207, 205),
+      Connection(205, 187),
+      Connection(187, 207),
+      Connection(83, 201),
+      Connection(201, 18),
+      Connection(18, 83),
+      Connection(181, 91),
+      Connection(91, 182),
+      Connection(182, 181),
+      Connection(180, 90),
+      Connection(90, 181),
+      Connection(181, 180),
+      Connection(16, 85),
+      Connection(85, 17),
+      Connection(17, 16),
+      Connection(205, 206),
+      Connection(206, 36),
+      Connection(36, 205),
+      Connection(176, 148),
+      Connection(148, 140),
+      Connection(140, 176),
+      Connection(165, 92),
+      Connection(92, 39),
+      Connection(39, 165),
+      Connection(245, 193),
+      Connection(193, 244),
+      Connection(244, 245),
+      Connection(27, 159),
+      Connection(159, 28),
+      Connection(28, 27),
+      Connection(30, 247),
+      Connection(247, 161),
+      Connection(161, 30),
+      Connection(174, 236),
+      Connection(236, 196),
+      Connection(196, 174),
+      Connection(103, 54),
+      Connection(54, 104),
+      Connection(104, 103),
+      Connection(55, 193),
+      Connection(193, 8),
+      Connection(8, 55),
+      Connection(111, 117),
+      Connection(117, 31),
+      Connection(31, 111),
+      Connection(221, 189),
+      Connection(189, 55),
+      Connection(55, 221),
+      Connection(240, 98),
+      Connection(98, 99),
+      Connection(99, 240),
+      Connection(142, 126),
+      Connection(126, 100),
+      Connection(100, 142),
+      Connection(219, 166),
+      Connection(166, 218),
+      Connection(218, 219),
+      Connection(112, 155),
+      Connection(155, 26),
+      Connection(26, 112),
+      Connection(198, 209),
+      Connection(209, 131),
+      Connection(131, 198),
+      Connection(169, 135),
+      Connection(135, 150),
+      Connection(150, 169),
+      Connection(114, 47),
+      Connection(47, 217),
+      Connection(217, 114),
+      Connection(224, 223),
+      Connection(223, 53),
+      Connection(53, 224),
+      Connection(220, 45),
+      Connection(45, 134),
+      Connection(134, 220),
+      Connection(32, 211),
+      Connection(211, 140),
+      Connection(140, 32),
+      Connection(109, 67),
+      Connection(67, 108),
+      Connection(108, 109),
+      Connection(146, 43),
+      Connection(43, 91),
+      Connection(91, 146),
+      Connection(231, 230),
+      Connection(230, 120),
+      Connection(120, 231),
+      Connection(113, 226),
+      Connection(226, 247),
+      Connection(247, 113),
+      Connection(105, 63),
+      Connection(63, 52),
+      Connection(52, 105),
+      Connection(241, 238),
+      Connection(238, 242),
+      Connection(242, 241),
+      Connection(124, 46),
+      Connection(46, 156),
+      Connection(156, 124),
+      Connection(95, 78),
+      Connection(78, 96),
+      Connection(96, 95),
+      Connection(70, 46),
+      Connection(46, 63),
+      Connection(63, 70),
+      Connection(116, 143),
+      Connection(143, 227),
+      Connection(227, 116),
+      Connection(116, 123),
+      Connection(123, 111),
+      Connection(111, 116),
+      Connection(1, 44),
+      Connection(44, 19),
+      Connection(19, 1),
+      Connection(3, 236),
+      Connection(236, 51),
+      Connection(51, 3),
+      Connection(207, 216),
+      Connection(216, 205),
+      Connection(205, 207),
+      Connection(26, 154),
+      Connection(154, 22),
+      Connection(22, 26),
+      Connection(165, 39),
+      Connection(39, 167),
+      Connection(167, 165),
+      Connection(199, 200),
+      Connection(200, 208),
+      Connection(208, 199),
+      Connection(101, 36),
+      Connection(36, 100),
+      Connection(100, 101),
+      Connection(43, 57),
+      Connection(57, 202),
+      Connection(202, 43),
+      Connection(242, 20),
+      Connection(20, 99),
+      Connection(99, 242),
+      Connection(56, 28),
+      Connection(28, 157),
+      Connection(157, 56),
+      Connection(124, 35),
+      Connection(35, 113),
+      Connection(113, 124),
+      Connection(29, 160),
+      Connection(160, 27),
+      Connection(27, 29),
+      Connection(211, 204),
+      Connection(204, 210),
+      Connection(210, 211),
+      Connection(124, 113),
+      Connection(113, 46),
+      Connection(46, 124),
+      Connection(106, 43),
+      Connection(43, 204),
+      Connection(204, 106),
+      Connection(96, 62),
+      Connection(62, 77),
+      Connection(77, 96),
+      Connection(227, 137),
+      Connection(137, 116),
+      Connection(116, 227),
+      Connection(73, 41),
+      Connection(41, 72),
+      Connection(72, 73),
+      Connection(36, 203),
+      Connection(203, 142),
+      Connection(142, 36),
+      Connection(235, 64),
+      Connection(64, 240),
+      Connection(240, 235),
+      Connection(48, 49),
+      Connection(49, 64),
+      Connection(64, 48),
+      Connection(42, 41),
+      Connection(41, 74),
+      Connection(74, 42),
+      Connection(214, 212),
+      Connection(212, 207),
+      Connection(207, 214),
+      Connection(183, 42),
+      Connection(42, 184),
+      Connection(184, 183),
+      Connection(210, 169),
+      Connection(169, 211),
+      Connection(211, 210),
+      Connection(140, 170),
+      Connection(170, 176),
+      Connection(176, 140),
+      Connection(104, 105),
+      Connection(105, 69),
+      Connection(69, 104),
+      Connection(193, 122),
+      Connection(122, 168),
+      Connection(168, 193),
+      Connection(50, 123),
+      Connection(123, 187),
+      Connection(187, 50),
+      Connection(89, 96),
+      Connection(96, 90),
+      Connection(90, 89),
+      Connection(66, 65),
+      Connection(65, 107),
+      Connection(107, 66),
+      Connection(179, 89),
+      Connection(89, 180),
+      Connection(180, 179),
+      Connection(119, 101),
+      Connection(101, 120),
+      Connection(120, 119),
+      Connection(68, 63),
+      Connection(63, 104),
+      Connection(104, 68),
+      Connection(234, 93),
+      Connection(93, 227),
+      Connection(227, 234),
+      Connection(16, 15),
+      Connection(15, 85),
+      Connection(85, 16),
+      Connection(209, 129),
+      Connection(129, 49),
+      Connection(49, 209),
+      Connection(15, 14),
+      Connection(14, 86),
+      Connection(86, 15),
+      Connection(107, 55),
+      Connection(55, 9),
+      Connection(9, 107),
+      Connection(120, 100),
+      Connection(100, 121),
+      Connection(121, 120),
+      Connection(153, 145),
+      Connection(145, 22),
+      Connection(22, 153),
+      Connection(178, 88),
+      Connection(88, 179),
+      Connection(179, 178),
+      Connection(197, 6),
+      Connection(6, 196),
+      Connection(196, 197),
+      Connection(89, 88),
+      Connection(88, 96),
+      Connection(96, 89),
+      Connection(135, 138),
+      Connection(138, 136),
+      Connection(136, 135),
+      Connection(138, 215),
+      Connection(215, 172),
+      Connection(172, 138),
+      Connection(218, 115),
+      Connection(115, 219),
+      Connection(219, 218),
+      Connection(41, 42),
+      Connection(42, 81),
+      Connection(81, 41),
+      Connection(5, 195),
+      Connection(195, 51),
+      Connection(51, 5),
+      Connection(57, 43),
+      Connection(43, 61),
+      Connection(61, 57),
+      Connection(208, 171),
+      Connection(171, 199),
+      Connection(199, 208),
+      Connection(41, 81),
+      Connection(81, 38),
+      Connection(38, 41),
+      Connection(224, 53),
+      Connection(53, 225),
+      Connection(225, 224),
+      Connection(24, 144),
+      Connection(144, 110),
+      Connection(110, 24),
+      Connection(105, 52),
+      Connection(52, 66),
+      Connection(66, 105),
+      Connection(118, 229),
+      Connection(229, 117),
+      Connection(117, 118),
+      Connection(227, 34),
+      Connection(34, 234),
+      Connection(234, 227),
+      Connection(66, 107),
+      Connection(107, 69),
+      Connection(69, 66),
+      Connection(10, 109),
+      Connection(109, 151),
+      Connection(151, 10),
+      Connection(219, 48),
+      Connection(48, 235),
+      Connection(235, 219),
+      Connection(183, 62),
+      Connection(62, 191),
+      Connection(191, 183),
+      Connection(142, 129),
+      Connection(129, 126),
+      Connection(126, 142),
+      Connection(116, 111),
+      Connection(111, 143),
+      Connection(143, 116),
+      Connection(118, 117),
+      Connection(117, 50),
+      Connection(50, 118),
+      Connection(223, 222),
+      Connection(222, 52),
+      Connection(52, 223),
+      Connection(94, 19),
+      Connection(19, 141),
+      Connection(141, 94),
+      Connection(222, 221),
+      Connection(221, 65),
+      Connection(65, 222),
+      Connection(196, 3),
+      Connection(3, 197),
+      Connection(197, 196),
+      Connection(45, 220),
+      Connection(220, 44),
+      Connection(44, 45),
+      Connection(156, 70),
+      Connection(70, 139),
+      Connection(139, 156),
+      Connection(188, 122),
+      Connection(122, 245),
+      Connection(245, 188),
+      Connection(139, 71),
+      Connection(71, 162),
+      Connection(162, 139),
+      Connection(149, 170),
+      Connection(170, 150),
+      Connection(150, 149),
+      Connection(122, 188),
+      Connection(188, 196),
+      Connection(196, 122),
+      Connection(206, 216),
+      Connection(216, 92),
+      Connection(92, 206),
+      Connection(164, 2),
+      Connection(2, 167),
+      Connection(167, 164),
+      Connection(242, 141),
+      Connection(141, 241),
+      Connection(241, 242),
+      Connection(0, 164),
+      Connection(164, 37),
+      Connection(37, 0),
+      Connection(11, 72),
+      Connection(72, 12),
+      Connection(12, 11),
+      Connection(12, 38),
+      Connection(38, 13),
+      Connection(13, 12),
+      Connection(70, 63),
+      Connection(63, 71),
+      Connection(71, 70),
+      Connection(31, 226),
+      Connection(226, 111),
+      Connection(111, 31),
+      Connection(36, 101),
+      Connection(101, 205),
+      Connection(205, 36),
+      Connection(203, 206),
+      Connection(206, 165),
+      Connection(165, 203),
+      Connection(126, 209),
+      Connection(209, 217),
+      Connection(217, 126),
+      Connection(98, 165),
+      Connection(165, 97),
+      Connection(97, 98),
+      Connection(237, 220),
+      Connection(220, 218),
+      Connection(218, 237),
+      Connection(237, 239),
+      Connection(239, 241),
+      Connection(241, 237),
+      Connection(210, 214),
+      Connection(214, 169),
+      Connection(169, 210),
+      Connection(140, 171),
+      Connection(171, 32),
+      Connection(32, 140),
+      Connection(241, 125),
+      Connection(125, 237),
+      Connection(237, 241),
+      Connection(179, 86),
+      Connection(86, 178),
+      Connection(178, 179),
+      Connection(180, 85),
+      Connection(85, 179),
+      Connection(179, 180),
+      Connection(181, 84),
+      Connection(84, 180),
+      Connection(180, 181),
+      Connection(182, 83),
+      Connection(83, 181),
+      Connection(181, 182),
+      Connection(194, 201),
+      Connection(201, 182),
+      Connection(182, 194),
+      Connection(177, 137),
+      Connection(137, 132),
+      Connection(132, 177),
+      Connection(184, 76),
+      Connection(76, 183),
+      Connection(183, 184),
+      Connection(185, 61),
+      Connection(61, 184),
+      Connection(184, 185),
+      Connection(186, 57),
+      Connection(57, 185),
+      Connection(185, 186),
+      Connection(216, 212),
+      Connection(212, 186),
+      Connection(186, 216),
+      Connection(192, 214),
+      Connection(214, 187),
+      Connection(187, 192),
+      Connection(139, 34),
+      Connection(34, 156),
+      Connection(156, 139),
+      Connection(218, 79),
+      Connection(79, 237),
+      Connection(237, 218),
+      Connection(147, 123),
+      Connection(123, 177),
+      Connection(177, 147),
+      Connection(45, 44),
+      Connection(44, 4),
+      Connection(4, 45),
+      Connection(208, 201),
+      Connection(201, 32),
+      Connection(32, 208),
+      Connection(98, 64),
+      Connection(64, 129),
+      Connection(129, 98),
+      Connection(192, 213),
+      Connection(213, 138),
+      Connection(138, 192),
+      Connection(235, 59),
+      Connection(59, 219),
+      Connection(219, 235),
+      Connection(141, 242),
+      Connection(242, 97),
+      Connection(97, 141),
+      Connection(97, 2),
+      Connection(2, 141),
+      Connection(141, 97),
+      Connection(240, 75),
+      Connection(75, 235),
+      Connection(235, 240),
+      Connection(229, 24),
+      Connection(24, 228),
+      Connection(228, 229),
+      Connection(31, 25),
+      Connection(25, 226),
+      Connection(226, 31),
+      Connection(230, 23),
+      Connection(23, 229),
+      Connection(229, 230),
+      Connection(231, 22),
+      Connection(22, 230),
+      Connection(230, 231),
+      Connection(232, 26),
+      Connection(26, 231),
+      Connection(231, 232),
+      Connection(233, 112),
+      Connection(112, 232),
+      Connection(232, 233),
+      Connection(244, 189),
+      Connection(189, 243),
+      Connection(243, 244),
+      Connection(189, 221),
+      Connection(221, 190),
+      Connection(190, 189),
+      Connection(222, 28),
+      Connection(28, 221),
+      Connection(221, 222),
+      Connection(223, 27),
+      Connection(27, 222),
+      Connection(222, 223),
+      Connection(224, 29),
+      Connection(29, 223),
+      Connection(223, 224),
+      Connection(225, 30),
+      Connection(30, 224),
+      Connection(224, 225),
+      Connection(113, 247),
+      Connection(247, 225),
+      Connection(225, 113),
+      Connection(99, 60),
+      Connection(60, 240),
+      Connection(240, 99),
+      Connection(213, 147),
+      Connection(147, 215),
+      Connection(215, 213),
+      Connection(60, 20),
+      Connection(20, 166),
+      Connection(166, 60),
+      Connection(192, 187),
+      Connection(187, 213),
+      Connection(213, 192),
+      Connection(243, 112),
+      Connection(112, 244),
+      Connection(244, 243),
+      Connection(244, 233),
+      Connection(233, 245),
+      Connection(245, 244),
+      Connection(245, 128),
+      Connection(128, 188),
+      Connection(188, 245),
+      Connection(188, 114),
+      Connection(114, 174),
+      Connection(174, 188),
+      Connection(134, 131),
+      Connection(131, 220),
+      Connection(220, 134),
+      Connection(174, 217),
+      Connection(217, 236),
+      Connection(236, 174),
+      Connection(236, 198),
+      Connection(198, 134),
+      Connection(134, 236),
+      Connection(215, 177),
+      Connection(177, 58),
+      Connection(58, 215),
+      Connection(156, 143),
+      Connection(143, 124),
+      Connection(124, 156),
+      Connection(25, 110),
+      Connection(110, 7),
+      Connection(7, 25),
+      Connection(31, 228),
+      Connection(228, 25),
+      Connection(25, 31),
+      Connection(264, 356),
+      Connection(356, 368),
+      Connection(368, 264),
+      Connection(0, 11),
+      Connection(11, 267),
+      Connection(267, 0),
+      Connection(451, 452),
+      Connection(452, 349),
+      Connection(349, 451),
+      Connection(267, 302),
+      Connection(302, 269),
+      Connection(269, 267),
+      Connection(350, 357),
+      Connection(357, 277),
+      Connection(277, 350),
+      Connection(350, 452),
+      Connection(452, 357),
+      Connection(357, 350),
+      Connection(299, 333),
+      Connection(333, 297),
+      Connection(297, 299),
+      Connection(396, 175),
+      Connection(175, 377),
+      Connection(377, 396),
+      Connection(280, 347),
+      Connection(347, 330),
+      Connection(330, 280),
+      Connection(269, 303),
+      Connection(303, 270),
+      Connection(270, 269),
+      Connection(151, 9),
+      Connection(9, 337),
+      Connection(337, 151),
+      Connection(344, 278),
+      Connection(278, 360),
+      Connection(360, 344),
+      Connection(424, 418),
+      Connection(418, 431),
+      Connection(431, 424),
+      Connection(270, 304),
+      Connection(304, 409),
+      Connection(409, 270),
+      Connection(272, 310),
+      Connection(310, 407),
+      Connection(407, 272),
+      Connection(322, 270),
+      Connection(270, 410),
+      Connection(410, 322),
+      Connection(449, 450),
+      Connection(450, 347),
+      Connection(347, 449),
+      Connection(432, 422),
+      Connection(422, 434),
+      Connection(434, 432),
+      Connection(18, 313),
+      Connection(313, 17),
+      Connection(17, 18),
+      Connection(291, 306),
+      Connection(306, 375),
+      Connection(375, 291),
+      Connection(259, 387),
+      Connection(387, 260),
+      Connection(260, 259),
+      Connection(424, 335),
+      Connection(335, 418),
+      Connection(418, 424),
+      Connection(434, 364),
+      Connection(364, 416),
+      Connection(416, 434),
+      Connection(391, 423),
+      Connection(423, 327),
+      Connection(327, 391),
+      Connection(301, 251),
+      Connection(251, 298),
+      Connection(298, 301),
+      Connection(275, 281),
+      Connection(281, 4),
+      Connection(4, 275),
+      Connection(254, 373),
+      Connection(373, 253),
+      Connection(253, 254),
+      Connection(375, 307),
+      Connection(307, 321),
+      Connection(321, 375),
+      Connection(280, 425),
+      Connection(425, 411),
+      Connection(411, 280),
+      Connection(200, 421),
+      Connection(421, 18),
+      Connection(18, 200),
+      Connection(335, 321),
+      Connection(321, 406),
+      Connection(406, 335),
+      Connection(321, 320),
+      Connection(320, 405),
+      Connection(405, 321),
+      Connection(314, 315),
+      Connection(315, 17),
+      Connection(17, 314),
+      Connection(423, 426),
+      Connection(426, 266),
+      Connection(266, 423),
+      Connection(396, 377),
+      Connection(377, 369),
+      Connection(369, 396),
+      Connection(270, 322),
+      Connection(322, 269),
+      Connection(269, 270),
+      Connection(413, 417),
+      Connection(417, 464),
+      Connection(464, 413),
+      Connection(385, 386),
+      Connection(386, 258),
+      Connection(258, 385),
+      Connection(248, 456),
+      Connection(456, 419),
+      Connection(419, 248),
+      Connection(298, 284),
+      Connection(284, 333),
+      Connection(333, 298),
+      Connection(168, 417),
+      Connection(417, 8),
+      Connection(8, 168),
+      Connection(448, 346),
+      Connection(346, 261),
+      Connection(261, 448),
+      Connection(417, 413),
+      Connection(413, 285),
+      Connection(285, 417),
+      Connection(326, 327),
+      Connection(327, 328),
+      Connection(328, 326),
+      Connection(277, 355),
+      Connection(355, 329),
+      Connection(329, 277),
+      Connection(309, 392),
+      Connection(392, 438),
+      Connection(438, 309),
+      Connection(381, 382),
+      Connection(382, 256),
+      Connection(256, 381),
+      Connection(279, 429),
+      Connection(429, 360),
+      Connection(360, 279),
+      Connection(365, 364),
+      Connection(364, 379),
+      Connection(379, 365),
+      Connection(355, 277),
+      Connection(277, 437),
+      Connection(437, 355),
+      Connection(282, 443),
+      Connection(443, 283),
+      Connection(283, 282),
+      Connection(281, 275),
+      Connection(275, 363),
+      Connection(363, 281),
+      Connection(395, 431),
+      Connection(431, 369),
+      Connection(369, 395),
+      Connection(299, 297),
+      Connection(297, 337),
+      Connection(337, 299),
+      Connection(335, 273),
+      Connection(273, 321),
+      Connection(321, 335),
+      Connection(348, 450),
+      Connection(450, 349),
+      Connection(349, 348),
+      Connection(359, 446),
+      Connection(446, 467),
+      Connection(467, 359),
+      Connection(283, 293),
+      Connection(293, 282),
+      Connection(282, 283),
+      Connection(250, 458),
+      Connection(458, 462),
+      Connection(462, 250),
+      Connection(300, 276),
+      Connection(276, 383),
+      Connection(383, 300),
+      Connection(292, 308),
+      Connection(308, 325),
+      Connection(325, 292),
+      Connection(283, 276),
+      Connection(276, 293),
+      Connection(293, 283),
+      Connection(264, 372),
+      Connection(372, 447),
+      Connection(447, 264),
+      Connection(346, 352),
+      Connection(352, 340),
+      Connection(340, 346),
+      Connection(354, 274),
+      Connection(274, 19),
+      Connection(19, 354),
+      Connection(363, 456),
+      Connection(456, 281),
+      Connection(281, 363),
+      Connection(426, 436),
+      Connection(436, 425),
+      Connection(425, 426),
+      Connection(380, 381),
+      Connection(381, 252),
+      Connection(252, 380),
+      Connection(267, 269),
+      Connection(269, 393),
+      Connection(393, 267),
+      Connection(421, 200),
+      Connection(200, 428),
+      Connection(428, 421),
+      Connection(371, 266),
+      Connection(266, 329),
+      Connection(329, 371),
+      Connection(432, 287),
+      Connection(287, 422),
+      Connection(422, 432),
+      Connection(290, 250),
+      Connection(250, 328),
+      Connection(328, 290),
+      Connection(385, 258),
+      Connection(258, 384),
+      Connection(384, 385),
+      Connection(446, 265),
+      Connection(265, 342),
+      Connection(342, 446),
+      Connection(386, 387),
+      Connection(387, 257),
+      Connection(257, 386),
+      Connection(422, 424),
+      Connection(424, 430),
+      Connection(430, 422),
+      Connection(445, 342),
+      Connection(342, 276),
+      Connection(276, 445),
+      Connection(422, 273),
+      Connection(273, 424),
+      Connection(424, 422),
+      Connection(306, 292),
+      Connection(292, 307),
+      Connection(307, 306),
+      Connection(352, 366),
+      Connection(366, 345),
+      Connection(345, 352),
+      Connection(268, 271),
+      Connection(271, 302),
+      Connection(302, 268),
+      Connection(358, 423),
+      Connection(423, 371),
+      Connection(371, 358),
+      Connection(327, 294),
+      Connection(294, 460),
+      Connection(460, 327),
+      Connection(331, 279),
+      Connection(279, 294),
+      Connection(294, 331),
+      Connection(303, 271),
+      Connection(271, 304),
+      Connection(304, 303),
+      Connection(436, 432),
+      Connection(432, 427),
+      Connection(427, 436),
+      Connection(304, 272),
+      Connection(272, 408),
+      Connection(408, 304),
+      Connection(395, 394),
+      Connection(394, 431),
+      Connection(431, 395),
+      Connection(378, 395),
+      Connection(395, 400),
+      Connection(400, 378),
+      Connection(296, 334),
+      Connection(334, 299),
+      Connection(299, 296),
+      Connection(6, 351),
+      Connection(351, 168),
+      Connection(168, 6),
+      Connection(376, 352),
+      Connection(352, 411),
+      Connection(411, 376),
+      Connection(307, 325),
+      Connection(325, 320),
+      Connection(320, 307),
+      Connection(285, 295),
+      Connection(295, 336),
+      Connection(336, 285),
+      Connection(320, 319),
+      Connection(319, 404),
+      Connection(404, 320),
+      Connection(329, 330),
+      Connection(330, 349),
+      Connection(349, 329),
+      Connection(334, 293),
+      Connection(293, 333),
+      Connection(333, 334),
+      Connection(366, 323),
+      Connection(323, 447),
+      Connection(447, 366),
+      Connection(316, 15),
+      Connection(15, 315),
+      Connection(315, 316),
+      Connection(331, 358),
+      Connection(358, 279),
+      Connection(279, 331),
+      Connection(317, 14),
+      Connection(14, 316),
+      Connection(316, 317),
+      Connection(8, 285),
+      Connection(285, 9),
+      Connection(9, 8),
+      Connection(277, 329),
+      Connection(329, 350),
+      Connection(350, 277),
+      Connection(253, 374),
+      Connection(374, 252),
+      Connection(252, 253),
+      Connection(319, 318),
+      Connection(318, 403),
+      Connection(403, 319),
+      Connection(351, 6),
+      Connection(6, 419),
+      Connection(419, 351),
+      Connection(324, 318),
+      Connection(318, 325),
+      Connection(325, 324),
+      Connection(397, 367),
+      Connection(367, 365),
+      Connection(365, 397),
+      Connection(288, 435),
+      Connection(435, 397),
+      Connection(397, 288),
+      Connection(278, 344),
+      Connection(344, 439),
+      Connection(439, 278),
+      Connection(310, 272),
+      Connection(272, 311),
+      Connection(311, 310),
+      Connection(248, 195),
+      Connection(195, 281),
+      Connection(281, 248),
+      Connection(375, 273),
+      Connection(273, 291),
+      Connection(291, 375),
+      Connection(175, 396),
+      Connection(396, 199),
+      Connection(199, 175),
+      Connection(312, 311),
+      Connection(311, 268),
+      Connection(268, 312),
+      Connection(276, 283),
+      Connection(283, 445),
+      Connection(445, 276),
+      Connection(390, 373),
+      Connection(373, 339),
+      Connection(339, 390),
+      Connection(295, 282),
+      Connection(282, 296),
+      Connection(296, 295),
+      Connection(448, 449),
+      Connection(449, 346),
+      Connection(346, 448),
+      Connection(356, 264),
+      Connection(264, 454),
+      Connection(454, 356),
+      Connection(337, 336),
+      Connection(336, 299),
+      Connection(299, 337),
+      Connection(337, 338),
+      Connection(338, 151),
+      Connection(151, 337),
+      Connection(294, 278),
+      Connection(278, 455),
+      Connection(455, 294),
+      Connection(308, 292),
+      Connection(292, 415),
+      Connection(415, 308),
+      Connection(429, 358),
+      Connection(358, 355),
+      Connection(355, 429),
+      Connection(265, 340),
+      Connection(340, 372),
+      Connection(372, 265),
+      Connection(352, 346),
+      Connection(346, 280),
+      Connection(280, 352),
+      Connection(295, 442),
+      Connection(442, 282),
+      Connection(282, 295),
+      Connection(354, 19),
+      Connection(19, 370),
+      Connection(370, 354),
+      Connection(285, 441),
+      Connection(441, 295),
+      Connection(295, 285),
+      Connection(195, 248),
+      Connection(248, 197),
+      Connection(197, 195),
+      Connection(457, 440),
+      Connection(440, 274),
+      Connection(274, 457),
+      Connection(301, 300),
+      Connection(300, 368),
+      Connection(368, 301),
+      Connection(417, 351),
+      Connection(351, 465),
+      Connection(465, 417),
+      Connection(251, 301),
+      Connection(301, 389),
+      Connection(389, 251),
+      Connection(394, 395),
+      Connection(395, 379),
+      Connection(379, 394),
+      Connection(399, 412),
+      Connection(412, 419),
+      Connection(419, 399),
+      Connection(410, 436),
+      Connection(436, 322),
+      Connection(322, 410),
+      Connection(326, 2),
+      Connection(2, 393),
+      Connection(393, 326),
+      Connection(354, 370),
+      Connection(370, 461),
+      Connection(461, 354),
+      Connection(393, 164),
+      Connection(164, 267),
+      Connection(267, 393),
+      Connection(268, 302),
+      Connection(302, 12),
+      Connection(12, 268),
+      Connection(312, 268),
+      Connection(268, 13),
+      Connection(13, 312),
+      Connection(298, 293),
+      Connection(293, 301),
+      Connection(301, 298),
+      Connection(265, 446),
+      Connection(446, 340),
+      Connection(340, 265),
+      Connection(280, 330),
+      Connection(330, 425),
+      Connection(425, 280),
+      Connection(322, 426),
+      Connection(426, 391),
+      Connection(391, 322),
+      Connection(420, 429),
+      Connection(429, 437),
+      Connection(437, 420),
+      Connection(393, 391),
+      Connection(391, 326),
+      Connection(326, 393),
+      Connection(344, 440),
+      Connection(440, 438),
+      Connection(438, 344),
+      Connection(458, 459),
+      Connection(459, 461),
+      Connection(461, 458),
+      Connection(364, 434),
+      Connection(434, 394),
+      Connection(394, 364),
+      Connection(428, 396),
+      Connection(396, 262),
+      Connection(262, 428),
+      Connection(274, 354),
+      Connection(354, 457),
+      Connection(457, 274),
+      Connection(317, 316),
+      Connection(316, 402),
+      Connection(402, 317),
+      Connection(316, 315),
+      Connection(315, 403),
+      Connection(403, 316),
+      Connection(315, 314),
+      Connection(314, 404),
+      Connection(404, 315),
+      Connection(314, 313),
+      Connection(313, 405),
+      Connection(405, 314),
+      Connection(313, 421),
+      Connection(421, 406),
+      Connection(406, 313),
+      Connection(323, 366),
+      Connection(366, 361),
+      Connection(361, 323),
+      Connection(292, 306),
+      Connection(306, 407),
+      Connection(407, 292),
+      Connection(306, 291),
+      Connection(291, 408),
+      Connection(408, 306),
+      Connection(291, 287),
+      Connection(287, 409),
+      Connection(409, 291),
+      Connection(287, 432),
+      Connection(432, 410),
+      Connection(410, 287),
+      Connection(427, 434),
+      Connection(434, 411),
+      Connection(411, 427),
+      Connection(372, 264),
+      Connection(264, 383),
+      Connection(383, 372),
+      Connection(459, 309),
+      Connection(309, 457),
+      Connection(457, 459),
+      Connection(366, 352),
+      Connection(352, 401),
+      Connection(401, 366),
+      Connection(1, 274),
+      Connection(274, 4),
+      Connection(4, 1),
+      Connection(418, 421),
+      Connection(421, 262),
+      Connection(262, 418),
+      Connection(331, 294),
+      Connection(294, 358),
+      Connection(358, 331),
+      Connection(435, 433),
+      Connection(433, 367),
+      Connection(367, 435),
+      Connection(392, 289),
+      Connection(289, 439),
+      Connection(439, 392),
+      Connection(328, 462),
+      Connection(462, 326),
+      Connection(326, 328),
+      Connection(94, 2),
+      Connection(2, 370),
+      Connection(370, 94),
+      Connection(289, 305),
+      Connection(305, 455),
+      Connection(455, 289),
+      Connection(339, 254),
+      Connection(254, 448),
+      Connection(448, 339),
+      Connection(359, 255),
+      Connection(255, 446),
+      Connection(446, 359),
+      Connection(254, 253),
+      Connection(253, 449),
+      Connection(449, 254),
+      Connection(253, 252),
+      Connection(252, 450),
+      Connection(450, 253),
+      Connection(252, 256),
+      Connection(256, 451),
+      Connection(451, 252),
+      Connection(256, 341),
+      Connection(341, 452),
+      Connection(452, 256),
+      Connection(414, 413),
+      Connection(413, 463),
+      Connection(463, 414),
+      Connection(286, 441),
+      Connection(441, 414),
+      Connection(414, 286),
+      Connection(286, 258),
+      Connection(258, 441),
+      Connection(441, 286),
+      Connection(258, 257),
+      Connection(257, 442),
+      Connection(442, 258),
+      Connection(257, 259),
+      Connection(259, 443),
+      Connection(443, 257),
+      Connection(259, 260),
+      Connection(260, 444),
+      Connection(444, 259),
+      Connection(260, 467),
+      Connection(467, 445),
+      Connection(445, 260),
+      Connection(309, 459),
+      Connection(459, 250),
+      Connection(250, 309),
+      Connection(305, 289),
+      Connection(289, 290),
+      Connection(290, 305),
+      Connection(305, 290),
+      Connection(290, 460),
+      Connection(460, 305),
+      Connection(401, 376),
+      Connection(376, 435),
+      Connection(435, 401),
+      Connection(309, 250),
+      Connection(250, 392),
+      Connection(392, 309),
+      Connection(376, 411),
+      Connection(411, 433),
+      Connection(433, 376),
+      Connection(453, 341),
+      Connection(341, 464),
+      Connection(464, 453),
+      Connection(357, 453),
+      Connection(453, 465),
+      Connection(465, 357),
+      Connection(343, 357),
+      Connection(357, 412),
+      Connection(412, 343),
+      Connection(437, 343),
+      Connection(343, 399),
+      Connection(399, 437),
+      Connection(344, 360),
+      Connection(360, 440),
+      Connection(440, 344),
+      Connection(420, 437),
+      Connection(437, 456),
+      Connection(456, 420),
+      Connection(360, 420),
+      Connection(420, 363),
+      Connection(363, 360),
+      Connection(361, 401),
+      Connection(401, 288),
+      Connection(288, 361),
+      Connection(265, 372),
+      Connection(372, 353),
+      Connection(353, 265),
+      Connection(390, 339),
+      Connection(339, 249),
+      Connection(249, 390),
+      Connection(339, 448),
+      Connection(448, 255),
+      Connection(255, 339),
+  ]
+
+
+@dataclasses.dataclass
+class FaceLandmarkerResult:
+  """The face landmarks detection result from FaceLandmarker, where each vector element represents a single face detected in the image.
+
+  Attributes:
+    face_landmarks: Detected face landmarks in normalized image coordinates.
+    face_blendshapes: Optional face blendshapes results.
+    facial_transformation_matrixes: Optional facial transformation matrix.
+  """
+
+  face_landmarks: List[List[landmark_module.NormalizedLandmark]]
+  face_blendshapes: List[List[category_module.Category]]
+  facial_transformation_matrixes: List[np.ndarray]
+
+
+def _build_landmarker_result(
+    output_packets: Mapping[str, packet_module.Packet]
+) -> FaceLandmarkerResult:
+  """Constructs a `FaceLandmarkerResult` from output packets."""
+  face_landmarks_proto_list = packet_getter.get_proto_list(
+      output_packets[_NORM_LANDMARKS_STREAM_NAME]
+  )
+
+  face_landmarks_results = []
+  for proto in face_landmarks_proto_list:
+    face_landmarks = landmark_pb2.NormalizedLandmarkList()
+    face_landmarks.MergeFrom(proto)
+    face_landmarks_list = []
+    for face_landmark in face_landmarks.landmark:
+      face_landmarks_list.append(
+          landmark_module.NormalizedLandmark.create_from_pb2(face_landmark)
+      )
+    face_landmarks_results.append(face_landmarks_list)
+
+  face_blendshapes_results = []
+  if _BLENDSHAPES_STREAM_NAME in output_packets:
+    face_blendshapes_proto_list = packet_getter.get_proto_list(
+        output_packets[_BLENDSHAPES_STREAM_NAME]
+    )
+    for proto in face_blendshapes_proto_list:
+      face_blendshapes_categories = []
+      face_blendshapes_classifications = classification_pb2.ClassificationList()
+      face_blendshapes_classifications.MergeFrom(proto)
+      for face_blendshapes in face_blendshapes_classifications.classification:
+        face_blendshapes_categories.append(
+            category_module.Category(
+                index=face_blendshapes.index,
+                score=face_blendshapes.score,
+                display_name=face_blendshapes.display_name,
+                category_name=face_blendshapes.label,
+            )
+        )
+      face_blendshapes_results.append(face_blendshapes_categories)
+
+  facial_transformation_matrixes_results = []
+  if _FACE_GEOMETRY_STREAM_NAME in output_packets:
+    facial_transformation_matrixes_proto_list = packet_getter.get_proto_list(
+        output_packets[_FACE_GEOMETRY_STREAM_NAME]
+    )
+    for proto in facial_transformation_matrixes_proto_list:
+      if hasattr(proto, 'pose_transform_matrix'):
+        matrix_data = matrix_data_pb2.MatrixData()
+        matrix_data.MergeFrom(proto.pose_transform_matrix)
+        matrix = np.array(matrix_data.packed_data)
+        matrix = matrix.reshape((matrix_data.rows, matrix_data.cols))
+        matrix = (
+            matrix if matrix_data.layout == _LayoutEnum.ROW_MAJOR else matrix.T
+        )
+        facial_transformation_matrixes_results.append(matrix)
+
+  return FaceLandmarkerResult(
+      face_landmarks_results,
+      face_blendshapes_results,
+      facial_transformation_matrixes_results,
+  )
+
+def _build_landmarker_result2(
+    output_packets: Mapping[str, packet_module.Packet]
+) -> FaceLandmarkerResult:
+  """Constructs a `FaceLandmarkerResult` from output packets."""
+  face_landmarks_proto_list = packet_getter.get_proto_list(
+      output_packets[_NORM_LANDMARKS_STREAM_NAME]
+  )
+
+  face_landmarks_results = []
+  for proto in face_landmarks_proto_list:
+    face_landmarks = landmark_pb2.NormalizedLandmarkList()
+    face_landmarks.MergeFrom(proto)
+    face_landmarks_list = []
+    for face_landmark in face_landmarks.landmark:
+      face_landmarks_list.append(
+          landmark_module.NormalizedLandmark.create_from_pb2(face_landmark)
+      )
+    face_landmarks_results.append(face_landmarks_list)
+
+  face_blendshapes_results = []
+  if _BLENDSHAPES_STREAM_NAME in output_packets:
+    face_blendshapes_proto_list = packet_getter.get_proto_list(
+        output_packets[_BLENDSHAPES_STREAM_NAME]
+    )
+    for proto in face_blendshapes_proto_list:
+      face_blendshapes_categories = []
+      face_blendshapes_classifications = classification_pb2.ClassificationList()
+      face_blendshapes_classifications.MergeFrom(proto)
+      for face_blendshapes in face_blendshapes_classifications.classification:
+        face_blendshapes_categories.append(
+            category_module.Category(
+                index=face_blendshapes.index,
+                score=face_blendshapes.score,
+                display_name=face_blendshapes.display_name,
+                category_name=face_blendshapes.label,
+            )
+        )
+      face_blendshapes_results.append(face_blendshapes_categories)
+
+  facial_transformation_matrixes_results = []
+  if _FACE_GEOMETRY_STREAM_NAME in output_packets:
+    facial_transformation_matrixes_proto_list = packet_getter.get_proto_list(
+        output_packets[_FACE_GEOMETRY_STREAM_NAME]
+    )
+    for proto in facial_transformation_matrixes_proto_list:
+      if hasattr(proto, 'pose_transform_matrix'):
+        matrix_data = matrix_data_pb2.MatrixData()
+        matrix_data.MergeFrom(proto.pose_transform_matrix)
+        matrix = np.array(matrix_data.packed_data)
+        matrix = matrix.reshape((matrix_data.rows, matrix_data.cols))
+        matrix = (
+            matrix if matrix_data.layout == _LayoutEnum.ROW_MAJOR else matrix.T
+        )
+        facial_transformation_matrixes_results.append(matrix)
+
+  return FaceLandmarkerResult(
+      face_landmarks_results,
+      face_blendshapes_results,
+      facial_transformation_matrixes_results,
+  ), facial_transformation_matrixes_proto_list[0].mesh
+
+@dataclasses.dataclass
+class FaceLandmarkerOptions:
+  """Options for the face landmarker task.
+
+  Attributes:
+    base_options: Base options for the face landmarker task.
+    running_mode: The running mode of the task. Default to the image mode.
+      FaceLandmarker has three running modes: 1) The image mode for detecting
+      face landmarks on single image inputs. 2) The video mode for detecting
+      face landmarks on the decoded frames of a video. 3) The live stream mode
+      for detecting face landmarks on the live stream of input data, such as
+      from camera. In this mode, the "result_callback" below must be specified
+      to receive the detection results asynchronously.
+    num_faces: The maximum number of faces that can be detected by the
+      FaceLandmarker.
+    min_face_detection_confidence: The minimum confidence score for the face
+      detection to be considered successful.
+    min_face_presence_confidence: The minimum confidence score of face presence
+      score in the face landmark detection.
+    min_tracking_confidence: The minimum confidence score for the face tracking
+      to be considered successful.
+    output_face_blendshapes: Whether FaceLandmarker outputs face blendshapes
+      classification. Face blendshapes are used for rendering the 3D face model.
+    output_facial_transformation_matrixes: Whether FaceLandmarker outputs facial
+      transformation_matrix. Facial transformation matrix is used to transform
+      the face landmarks in canonical face to the detected face, so that users
+      can apply face effects on the detected landmarks.
+    result_callback: The user-defined result callback for processing live stream
+      data. The result callback should only be specified when the running mode
+      is set to the live stream mode.
+  """
+
+  base_options: _BaseOptions
+  running_mode: _RunningMode = _RunningMode.IMAGE
+  num_faces: int = 1
+  min_face_detection_confidence: float = 0.5
+  min_face_presence_confidence: float = 0.5
+  min_tracking_confidence: float = 0.5
+  output_face_blendshapes: bool = False
+  output_facial_transformation_matrixes: bool = False
+  result_callback: Optional[
+      Callable[[FaceLandmarkerResult, image_module.Image, int], None]
+  ] = None
+
+  @doc_controls.do_not_generate_docs
+  def to_pb2(self) -> _FaceLandmarkerGraphOptionsProto:
+    """Generates an FaceLandmarkerGraphOptions protobuf object."""
+    base_options_proto = self.base_options.to_pb2()
+    base_options_proto.use_stream_mode = (
+        False if self.running_mode == _RunningMode.IMAGE else True
+    )
+
+    # Initialize the face landmarker options from base options.
+    face_landmarker_options_proto = _FaceLandmarkerGraphOptionsProto(
+        base_options=base_options_proto
+    )
+
+    # Configure face detector options.
+    face_landmarker_options_proto.face_detector_graph_options.num_faces = (
+        self.num_faces
+    )
+    face_landmarker_options_proto.face_detector_graph_options.min_detection_confidence = (
+        self.min_face_detection_confidence
+    )
+
+    # Configure face landmark detector options.
+    face_landmarker_options_proto.min_tracking_confidence = (
+        self.min_tracking_confidence
+    )
+    face_landmarker_options_proto.face_landmarks_detector_graph_options.min_detection_confidence = (
+        self.min_face_detection_confidence
+    )
+    return face_landmarker_options_proto
+
+
+class FaceLandmarker(base_vision_task_api.BaseVisionTaskApi):
+  """Class that performs face landmarks detection on images."""
+
+  @classmethod
+  def create_from_model_path(cls, model_path: str) -> 'FaceLandmarker':
+    """Creates an `FaceLandmarker` object from a TensorFlow Lite model and the default `FaceLandmarkerOptions`.
+
+    Note that the created `FaceLandmarker` instance is in image mode, for
+    detecting face landmarks on single image inputs.
+
+    Args:
+      model_path: Path to the model.
+
+    Returns:
+      `FaceLandmarker` object that's created from the model file and the
+      default `FaceLandmarkerOptions`.
+
+    Raises:
+      ValueError: If failed to create `FaceLandmarker` object from the
+        provided file such as invalid file path.
+      RuntimeError: If other types of error occurred.
+    """
+    base_options = _BaseOptions(model_asset_path=model_path)
+    options = FaceLandmarkerOptions(
+        base_options=base_options, running_mode=_RunningMode.IMAGE
+    )
+    return cls.create_from_options(options)
+
+  @classmethod
+  def create_from_options(
+      cls, options: FaceLandmarkerOptions
+  ) -> 'FaceLandmarker':
+    """Creates the `FaceLandmarker` object from face landmarker options.
+
+    Args:
+      options: Options for the face landmarker task.
+
+    Returns:
+      `FaceLandmarker` object that's created from `options`.
+
+    Raises:
+      ValueError: If failed to create `FaceLandmarker` object from
+        `FaceLandmarkerOptions` such as missing the model.
+      RuntimeError: If other types of error occurred.
+    """
+
+    def packets_callback(output_packets: Mapping[str, packet_module.Packet]):
+      if output_packets[_IMAGE_OUT_STREAM_NAME].is_empty():
+        return
+
+      image = packet_getter.get_image(output_packets[_IMAGE_OUT_STREAM_NAME])
+      if output_packets[_IMAGE_OUT_STREAM_NAME].is_empty():
+        return
+
+      if output_packets[_NORM_LANDMARKS_STREAM_NAME].is_empty():
+        empty_packet = output_packets[_NORM_LANDMARKS_STREAM_NAME]
+        options.result_callback(
+            FaceLandmarkerResult([], [], []),
+            image,
+            empty_packet.timestamp.value // _MICRO_SECONDS_PER_MILLISECOND,
+        )
+        return
+
+      face_landmarks_result = _build_landmarker_result(output_packets)
+      timestamp = output_packets[_NORM_LANDMARKS_STREAM_NAME].timestamp
+      options.result_callback(
+          face_landmarks_result,
+          image,
+          timestamp.value // _MICRO_SECONDS_PER_MILLISECOND,
+      )
+
+    output_streams = [
+        ':'.join([_NORM_LANDMARKS_TAG, _NORM_LANDMARKS_STREAM_NAME]),
+        ':'.join([_IMAGE_TAG, _IMAGE_OUT_STREAM_NAME]),
+    ]
+
+    if options.output_face_blendshapes:
+      output_streams.append(
+          ':'.join([_BLENDSHAPES_TAG, _BLENDSHAPES_STREAM_NAME])
+      )
+    if options.output_facial_transformation_matrixes:
+      output_streams.append(
+          ':'.join([_FACE_GEOMETRY_TAG, _FACE_GEOMETRY_STREAM_NAME])
+      )
+
+    task_info = _TaskInfo(
+        task_graph=_TASK_GRAPH_NAME,
+        input_streams=[
+            ':'.join([_IMAGE_TAG, _IMAGE_IN_STREAM_NAME]),
+            ':'.join([_NORM_RECT_TAG, _NORM_RECT_STREAM_NAME]),
+        ],
+        output_streams=output_streams,
+        task_options=options,
+    )
+    return cls(
+        task_info.generate_graph_config(
+            enable_flow_limiting=options.running_mode
+            == _RunningMode.LIVE_STREAM
+        ),
+        options.running_mode,
+        packets_callback if options.result_callback else None,
+    )
+
+  def detect(
+      self,
+      image: image_module.Image,
+      image_processing_options: Optional[_ImageProcessingOptions] = None,
+  ) -> FaceLandmarkerResult:
+    """Performs face landmarks detection on the given image.
+
+    Only use this method when the FaceLandmarker is created with the image
+    running mode.
+
+    The image can be of any size with format RGB or RGBA.
+    TODO: Describes how the input image will be preprocessed after the yuv
+    support is implemented.
+
+    Args:
+      image: MediaPipe Image.
+      image_processing_options: Options for image processing.
+
+    Returns:
+      The face landmarks detection results.
+
+    Raises:
+      ValueError: If any of the input arguments is invalid.
+      RuntimeError: If face landmarker detection failed to run.
+    """
+
+    normalized_rect = self.convert_to_normalized_rect(
+        image_processing_options, image, roi_allowed=False
+    )
+    output_packets = self._process_image_data({
+        _IMAGE_IN_STREAM_NAME: packet_creator.create_image(image),
+        _NORM_RECT_STREAM_NAME: packet_creator.create_proto(
+            normalized_rect.to_pb2()
+        ),
+    })
+
+    if output_packets[_NORM_LANDMARKS_STREAM_NAME].is_empty():
+      return FaceLandmarkerResult([], [], [])
+
+    return _build_landmarker_result2(output_packets)
+
+  def detect_for_video(
+      self,
+      image: image_module.Image,
+      timestamp_ms: int,
+      image_processing_options: Optional[_ImageProcessingOptions] = None,
+  ):
+    """Performs face landmarks detection on the provided video frame.
+
+    Only use this method when the FaceLandmarker is created with the video
+    running mode.
+
+    Only use this method when the FaceLandmarker is created with the video
+    running mode. It's required to provide the video frame's timestamp (in
+    milliseconds) along with the video frame. The input timestamps should be
+    monotonically increasing for adjacent calls of this method.
+
+    Args:
+      image: MediaPipe Image.
+      timestamp_ms: The timestamp of the input video frame in milliseconds.
+      image_processing_options: Options for image processing.
+
+    Returns:
+      The face landmarks detection results.
+
+    Raises:
+      ValueError: If any of the input arguments is invalid.
+      RuntimeError: If face landmarker detection failed to run.
+    """
+    normalized_rect = self.convert_to_normalized_rect(
+        image_processing_options, image, roi_allowed=False
+    )
+    output_packets = self._process_video_data({
+        _IMAGE_IN_STREAM_NAME: packet_creator.create_image(image).at(
+            timestamp_ms * _MICRO_SECONDS_PER_MILLISECOND
+        ),
+        _NORM_RECT_STREAM_NAME: packet_creator.create_proto(
+            normalized_rect.to_pb2()
+        ).at(timestamp_ms * _MICRO_SECONDS_PER_MILLISECOND),
+    })
+
+    if output_packets[_NORM_LANDMARKS_STREAM_NAME].is_empty():
+      return FaceLandmarkerResult([], [], [])
+
+    return _build_landmarker_result2(output_packets)
+
+  def detect_async(
+      self,
+      image: image_module.Image,
+      timestamp_ms: int,
+      image_processing_options: Optional[_ImageProcessingOptions] = None,
+  ) -> None:
+    """Sends live image data to perform face landmarks detection.
+
+    The results will be available via the "result_callback" provided in the
+    FaceLandmarkerOptions. Only use this method when the FaceLandmarker is
+    created with the live stream running mode.
+
+    Only use this method when the FaceLandmarker is created with the live
+    stream running mode. The input timestamps should be monotonically increasing
+    for adjacent calls of this method. This method will return immediately after
+    the input image is accepted. The results will be available via the
+    `result_callback` provided in the `FaceLandmarkerOptions`. The
+    `detect_async` method is designed to process live stream data such as
+    camera input. To lower the overall latency, face landmarker may drop the
+    input images if needed. In other words, it's not guaranteed to have output
+    per input image.
+
+    The `result_callback` provides:
+      - The face landmarks detection results.
+      - The input image that the face landmarker runs on.
+      - The input timestamp in milliseconds.
+
+    Args:
+      image: MediaPipe Image.
+      timestamp_ms: The timestamp of the input image in milliseconds.
+      image_processing_options: Options for image processing.
+
+    Raises:
+      ValueError: If the current input timestamp is smaller than what the
+      face landmarker has already processed.
+    """
+    normalized_rect = self.convert_to_normalized_rect(
+        image_processing_options, image, roi_allowed=False
+    )
+    self._send_live_stream_data({
+        _IMAGE_IN_STREAM_NAME: packet_creator.create_image(image).at(
+            timestamp_ms * _MICRO_SECONDS_PER_MILLISECOND
+        ),
+        _NORM_RECT_STREAM_NAME: packet_creator.create_proto(
+            normalized_rect.to_pb2()
+        ).at(timestamp_ms * _MICRO_SECONDS_PER_MILLISECOND),
+    })

src/utils/mp_utils.py

@@ -0,0 +1,95 @@
+import os
+import numpy as np
+import cv2
+import time
+from tqdm import tqdm
+import multiprocessing
+import glob
+
+import mediapipe as mp
+from mediapipe import solutions
+from mediapipe.framework.formats import landmark_pb2
+from mediapipe.tasks import python
+from mediapipe.tasks.python import vision
+from . import face_landmark
+
+CUR_DIR = os.path.dirname(__file__)
+
+
+class LMKExtractor():
+    def __init__(self, FPS=25):
+        # Create an FaceLandmarker object.
+        self.mode = mp.tasks.vision.FaceDetectorOptions.running_mode.IMAGE
+        base_options = python.BaseOptions(model_asset_path=os.path.join(CUR_DIR, 'mp_models/face_landmarker_v2_with_blendshapes.task'))
+        base_options.delegate = mp.tasks.BaseOptions.Delegate.CPU
+        options = vision.FaceLandmarkerOptions(base_options=base_options,
+                                            running_mode=self.mode,
+                                            output_face_blendshapes=True,
+                                            output_facial_transformation_matrixes=True,
+                                            num_faces=1)
+        self.detector = face_landmark.FaceLandmarker.create_from_options(options)
+        self.last_ts = 0
+        self.frame_ms = int(1000 / FPS)
+
+        det_base_options = python.BaseOptions(model_asset_path=os.path.join(CUR_DIR, 'mp_models/blaze_face_short_range.tflite'))
+        det_options = vision.FaceDetectorOptions(base_options=det_base_options)
+        self.det_detector = vision.FaceDetector.create_from_options(det_options)
+                
+
+    def __call__(self, img):
+        frame = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+        image = mp.Image(image_format=mp.ImageFormat.SRGB, data=frame)
+        t0 = time.time()
+        if self.mode == mp.tasks.vision.FaceDetectorOptions.running_mode.VIDEO:
+            det_result = self.det_detector.detect(image)
+            if len(det_result.detections) != 1:
+                return None
+            self.last_ts += self.frame_ms
+            try:
+                detection_result, mesh3d = self.detector.detect_for_video(image, timestamp_ms=self.last_ts)
+            except:
+                return None
+        elif self.mode == mp.tasks.vision.FaceDetectorOptions.running_mode.IMAGE:
+            # det_result = self.det_detector.detect(image)
+
+            # if len(det_result.detections) != 1:
+            #     return None
+            try:
+                detection_result, mesh3d = self.detector.detect(image)
+            except:
+                return None
+            
+        
+        bs_list = detection_result.face_blendshapes
+        if len(bs_list) == 1:
+            bs = bs_list[0]
+            bs_values = []
+            for index in range(len(bs)):
+                bs_values.append(bs[index].score)
+            bs_values = bs_values[1:] # remove neutral
+            trans_mat = detection_result.facial_transformation_matrixes[0]
+            face_landmarks_list = detection_result.face_landmarks
+            face_landmarks = face_landmarks_list[0]
+            lmks = []
+            for index in range(len(face_landmarks)):
+                x = face_landmarks[index].x
+                y = face_landmarks[index].y
+                z = face_landmarks[index].z
+                lmks.append([x, y, z])
+            lmks = np.array(lmks)
+            
+            lmks3d = np.array(mesh3d.vertex_buffer)
+            lmks3d = lmks3d.reshape(-1, 5)[:, :3]
+            mp_tris = np.array(mesh3d.index_buffer).reshape(-1, 3) + 1
+
+            return {
+                "lmks": lmks,
+                'lmks3d': lmks3d,
+                "trans_mat": trans_mat,
+                'faces': mp_tris,
+                "bs": bs_values
+            }
+        else:
+            # print('multiple faces in the image: {}'.format(img_path))
+            return None
+        

src/utils/pose_util.py

@@ -0,0 +1,78 @@
+import math
+
+import numpy as np
+from scipy.spatial.transform import Rotation as R
+
+
+def create_perspective_matrix(aspect_ratio):
+    kDegreesToRadians = np.pi / 180.
+    near = 1
+    far = 10000
+    perspective_matrix = np.zeros(16, dtype=np.float32)
+
+    # Standard perspective projection matrix calculations.
+    f = 1.0 / np.tan(kDegreesToRadians * 63 / 2.)
+
+    denom = 1.0 / (near - far)
+    perspective_matrix[0] = f / aspect_ratio
+    perspective_matrix[5] = f
+    perspective_matrix[10] = (near + far) * denom
+    perspective_matrix[11] = -1.
+    perspective_matrix[14] = 1. * far * near * denom
+
+    # If the environment's origin point location is in the top left corner,
+    # then skip additional flip along Y-axis is required to render correctly.
+
+    perspective_matrix[5] *= -1.
+    return perspective_matrix
+
+
+def project_points(points_3d, transformation_matrix, pose_vectors, image_shape):
+    P = create_perspective_matrix(image_shape[1] / image_shape[0]).reshape(4, 4).T
+    L, N, _ = points_3d.shape
+    projected_points = np.zeros((L, N, 2))
+    for i in range(L):
+        points_3d_frame = points_3d[i]
+        ones = np.ones((points_3d_frame.shape[0], 1))
+        points_3d_homogeneous = np.hstack([points_3d_frame, ones])  
+        transformed_points = points_3d_homogeneous @ (transformation_matrix @ euler_and_translation_to_matrix(pose_vectors[i][:3], pose_vectors[i][3:])).T @ P
+        projected_points_frame = transformed_points[:, :2] / transformed_points[:, 3, np.newaxis] # -1 ~ 1
+        projected_points_frame[:, 0] = (projected_points_frame[:, 0] + 1) * 0.5 * image_shape[1] 
+        projected_points_frame[:, 1]  = (projected_points_frame[:, 1] + 1) * 0.5 * image_shape[0]
+        projected_points[i] = projected_points_frame
+    return projected_points
+
+
+def project_points_with_trans(points_3d, transformation_matrix, image_shape):
+    P = create_perspective_matrix(image_shape[1] / image_shape[0]).reshape(4, 4).T
+    L, N, _ = points_3d.shape
+    projected_points = np.zeros((L, N, 2))
+    for i in range(L):
+        points_3d_frame = points_3d[i]
+        ones = np.ones((points_3d_frame.shape[0], 1))
+        points_3d_homogeneous = np.hstack([points_3d_frame, ones])  
+        transformed_points = points_3d_homogeneous @ transformation_matrix[i].T @ P
+        projected_points_frame = transformed_points[:, :2] / transformed_points[:, 3, np.newaxis] # -1 ~ 1
+        projected_points_frame[:, 0] = (projected_points_frame[:, 0] + 1) * 0.5 * image_shape[1] 
+        projected_points_frame[:, 1]  = (projected_points_frame[:, 1] + 1) * 0.5 * image_shape[0]
+        projected_points[i] = projected_points_frame
+    return projected_points
+
+
+def euler_and_translation_to_matrix(euler_angles, translation_vector):
+    rotation = R.from_euler('xyz', euler_angles, degrees=True)
+    rotation_matrix = rotation.as_matrix()
+    
+    matrix = np.eye(4)
+    matrix[:3, :3] = rotation_matrix
+    matrix[:3, 3] = translation_vector
+    
+    return matrix
+
+
+def matrix_to_euler_and_translation(matrix):
+    rotation_matrix = matrix[:3, :3]
+    translation_vector = matrix[:3, 3]
+    rotation = R.from_matrix(rotation_matrix)
+    euler_angles = rotation.as_euler('xyz', degrees=True)
+    return euler_angles, translation_vector

src/utils/util.py

@@ -0,0 +1,128 @@
+import importlib
+import os
+import os.path as osp
+import shutil
+import sys
+from pathlib import Path
+
+import av
+import numpy as np
+import torch
+import torchvision
+from einops import rearrange
+from PIL import Image
+
+
+def seed_everything(seed):
+    import random
+
+    import numpy as np
+
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    np.random.seed(seed % (2**32))
+    random.seed(seed)
+
+
+def import_filename(filename):
+    spec = importlib.util.spec_from_file_location("mymodule", filename)
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[spec.name] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+def delete_additional_ckpt(base_path, num_keep):
+    dirs = []
+    for d in os.listdir(base_path):
+        if d.startswith("checkpoint-"):
+            dirs.append(d)
+    num_tot = len(dirs)
+    if num_tot <= num_keep:
+        return
+    # ensure ckpt is sorted and delete the ealier!
+    del_dirs = sorted(dirs, key=lambda x: int(x.split("-")[-1]))[: num_tot - num_keep]
+    for d in del_dirs:
+        path_to_dir = osp.join(base_path, d)
+        if osp.exists(path_to_dir):
+            shutil.rmtree(path_to_dir)
+
+
+def save_videos_from_pil(pil_images, path, fps=8):
+    import av
+
+    save_fmt = Path(path).suffix
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+    width, height = pil_images[0].size
+
+    if save_fmt == ".mp4":
+        codec = "libx264"
+        container = av.open(path, "w")
+        stream = container.add_stream(codec, rate=fps)
+
+        stream.width = width
+        stream.height = height
+
+        for pil_image in pil_images:
+            # pil_image = Image.fromarray(image_arr).convert("RGB")
+            av_frame = av.VideoFrame.from_image(pil_image)
+            container.mux(stream.encode(av_frame))
+        container.mux(stream.encode())
+        container.close()
+
+    elif save_fmt == ".gif":
+        pil_images[0].save(
+            fp=path,
+            format="GIF",
+            append_images=pil_images[1:],
+            save_all=True,
+            duration=(1 / fps * 1000),
+            loop=0,
+        )
+    else:
+        raise ValueError("Unsupported file type. Use .mp4 or .gif.")
+
+
+def save_videos_grid(videos: torch.Tensor, path: str, rescale=False, n_rows=6, fps=8):
+    videos = rearrange(videos, "b c t h w -> t b c h w")
+    height, width = videos.shape[-2:]
+    outputs = []
+
+    for x in videos:
+        x = torchvision.utils.make_grid(x, nrow=n_rows)  # (c h w)
+        x = x.transpose(0, 1).transpose(1, 2).squeeze(-1)  # (h w c)
+        if rescale:
+            x = (x + 1.0) / 2.0  # -1,1 -> 0,1
+        x = (x * 255).numpy().astype(np.uint8)
+        x = Image.fromarray(x)
+
+        outputs.append(x)
+
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+
+    save_videos_from_pil(outputs, path, fps)
+
+
+def read_frames(video_path):
+    container = av.open(video_path)
+
+    video_stream = next(s for s in container.streams if s.type == "video")
+    frames = []
+    for packet in container.demux(video_stream):
+        for frame in packet.decode():
+            image = Image.frombytes(
+                "RGB",
+                (frame.width, frame.height),
+                frame.to_rgb().to_ndarray(),
+            )
+            frames.append(image)
+
+    return frames
+
+
+def get_fps(video_path):
+    container = av.open(video_path)
+    video_stream = next(s for s in container.streams if s.type == "video")
+    fps = video_stream.average_rate
+    container.close()
+    return fps

src/vid2vid.py

@@ -0,0 +1,233 @@
+import argparse
+import os
+import shutil
+import ffmpeg
+from datetime import datetime
+from pathlib import Path
+import numpy as np
+import cv2
+import torch
+# import spaces
+from diffusers import AutoencoderKL, DDIMScheduler
+from einops import repeat
+from omegaconf import OmegaConf
+from PIL import Image
+from torchvision import transforms
+from transformers import CLIPVisionModelWithProjection
+
+from src.models.pose_guider import PoseGuider
+from src.models.unet_2d_condition import UNet2DConditionModel
+from src.models.unet_3d import UNet3DConditionModel
+from src.pipelines.pipeline_pose2vid_long import Pose2VideoPipeline
+from src.utils.util import get_fps, read_frames, save_videos_grid
+
+from src.utils.mp_utils  import LMKExtractor
+from src.utils.draw_util import FaceMeshVisualizer
+from src.utils.pose_util import project_points_with_trans,  matrix_to_euler_and_translation
+from src.audio2vid import smooth_pose_seq
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", type=str, default='./configs/prompts/animation_facereenac.yaml')
+    parser.add_argument("-W", type=int, default=512)
+    parser.add_argument("-H", type=int, default=512)
+    parser.add_argument("-L", type=int)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--cfg", type=float, default=3.5)
+    parser.add_argument("--steps", type=int, default=25)
+    parser.add_argument("--fps", type=int)
+    args = parser.parse_args()
+
+    return args
+
+# @spaces.GPU
+def video2video(ref_img, source_video, size=512, steps=25, length=150, seed=42):
+    cfg = 3.5
+    
+    config = OmegaConf.load('./configs/prompts/animation_facereenac.yaml')
+
+    if config.weight_dtype == "fp16":
+        weight_dtype = torch.float16
+    else:
+        weight_dtype = torch.float32
+
+    vae = AutoencoderKL.from_pretrained(
+        config.pretrained_vae_path,
+    ).to("cuda", dtype=weight_dtype)
+
+    reference_unet = UNet2DConditionModel.from_pretrained(
+        config.pretrained_base_model_path,
+        subfolder="unet",
+    ).to(dtype=weight_dtype, device="cuda")
+
+    inference_config_path = config.inference_config
+    infer_config = OmegaConf.load(inference_config_path)
+    denoising_unet = UNet3DConditionModel.from_pretrained_2d(
+        config.pretrained_base_model_path,
+        config.motion_module_path,
+        subfolder="unet",
+        unet_additional_kwargs=infer_config.unet_additional_kwargs,
+    ).to(dtype=weight_dtype, device="cuda")
+
+    pose_guider = PoseGuider(noise_latent_channels=320, use_ca=True).to(device="cuda", dtype=weight_dtype) # not use cross attention
+
+    image_enc = CLIPVisionModelWithProjection.from_pretrained(
+        config.image_encoder_path
+    ).to(dtype=weight_dtype, device="cuda")
+
+    sched_kwargs = OmegaConf.to_container(infer_config.noise_scheduler_kwargs)
+    scheduler = DDIMScheduler(**sched_kwargs)
+
+    generator = torch.manual_seed(seed)
+
+    width, height = size, size
+
+    # load pretrained weights
+    denoising_unet.load_state_dict(
+        torch.load(config.denoising_unet_path, map_location="cpu"),
+        strict=False,
+    )
+    reference_unet.load_state_dict(
+        torch.load(config.reference_unet_path, map_location="cpu"),
+    )
+    pose_guider.load_state_dict(
+        torch.load(config.pose_guider_path, map_location="cpu"),
+    )
+
+    pipe = Pose2VideoPipeline(
+        vae=vae,
+        image_encoder=image_enc,
+        reference_unet=reference_unet,
+        denoising_unet=denoising_unet,
+        pose_guider=pose_guider,
+        scheduler=scheduler,
+    )
+    pipe = pipe.to("cuda", dtype=weight_dtype)
+
+    date_str = datetime.now().strftime("%Y%m%d")
+    time_str = datetime.now().strftime("%H%M")
+    save_dir_name = f"{time_str}--seed_{seed}-{size}x{size}"
+
+    save_dir = Path(f"output/{date_str}/{save_dir_name}")
+    save_dir.mkdir(exist_ok=True, parents=True)
+
+
+    lmk_extractor = LMKExtractor()
+    vis = FaceMeshVisualizer(forehead_edge=False)
+    
+
+
+    ref_image_np = cv2.cvtColor(ref_img, cv2.COLOR_RGB2BGR)
+    # TODO: 人脸检测+裁剪
+    ref_image_np = cv2.resize(ref_image_np, (size, size))
+    ref_image_pil = Image.fromarray(cv2.cvtColor(ref_image_np, cv2.COLOR_BGR2RGB))
+    
+    face_result = lmk_extractor(ref_image_np)
+    if face_result is None: 
+        return None
+    
+    lmks = face_result['lmks'].astype(np.float32)
+    ref_pose = vis.draw_landmarks((ref_image_np.shape[1], ref_image_np.shape[0]), lmks, normed=True)
+    
+    
+
+    source_images = read_frames(source_video)
+    src_fps = get_fps(source_video)
+    pose_transform = transforms.Compose(
+        [transforms.Resize((height, width)), transforms.ToTensor()]
+    )
+    
+    step = 1
+    if src_fps == 60:
+        src_fps = 30
+        step = 2
+    
+    pose_trans_list = []
+    verts_list = []
+    bs_list = []
+    src_tensor_list = []
+    args_L = len(source_images) if length==0 or length*step > len(source_images) else length*step
+    for src_image_pil in source_images[: args_L: step]:
+        src_tensor_list.append(pose_transform(src_image_pil))
+        src_img_np = cv2.cvtColor(np.array(src_image_pil), cv2.COLOR_RGB2BGR)
+        frame_height, frame_width, _ = src_img_np.shape
+        src_img_result = lmk_extractor(src_img_np)
+        if src_img_result is None:
+            break
+        pose_trans_list.append(src_img_result['trans_mat'])
+        verts_list.append(src_img_result['lmks3d'])
+        bs_list.append(src_img_result['bs'])
+
+    
+    # pose_arr = np.array(pose_trans_list)
+    trans_mat_arr = np.array(pose_trans_list)
+    verts_arr = np.array(verts_list)
+    bs_arr = np.array(bs_list)
+    min_bs_idx = np.argmin(bs_arr.sum(1))
+    
+    # compute delta pose
+    trans_mat_inv_frame_0 = np.linalg.inv(trans_mat_arr[0])
+    pose_arr = np.zeros([trans_mat_arr.shape[0], 6])
+
+    for i in range(pose_arr.shape[0]):
+        pose_mat = trans_mat_inv_frame_0 @ trans_mat_arr[i]
+        euler_angles, translation_vector = matrix_to_euler_and_translation(pose_mat)
+        pose_arr[i, :3] =  euler_angles
+        pose_arr[i, 3:6] =  translation_vector
+
+    pose_arr = smooth_pose_seq(pose_arr)
+
+    # face retarget
+    verts_arr = verts_arr - verts_arr[min_bs_idx] + face_result['lmks3d']
+    # project 3D mesh to 2D landmark
+    projected_vertices = project_points_with_trans(verts_arr, pose_arr, [frame_height, frame_width])
+    
+    pose_list = []
+    for i, verts in enumerate(projected_vertices):
+        lmk_img = vis.draw_landmarks((frame_width, frame_height), verts, normed=False)
+        pose_image_np = cv2.resize(lmk_img,  (width, height))
+        pose_list.append(pose_image_np)
+    
+    pose_list = np.array(pose_list)
+    
+    video_length = len(pose_list)
+
+    video = pipe(
+        ref_image_pil,
+        pose_list,
+        ref_pose,
+        width,
+        height,
+        video_length,
+        steps,
+        cfg,
+        generator=generator,
+    ).videos
+
+    save_path = f"{save_dir}/{size}x{size}_{time_str}_noaudio.mp4"
+    save_videos_grid(
+        video,
+        save_path,
+        n_rows=1,
+        fps=src_fps,
+    )
+    
+    audio_output = f'{save_dir}/audio_from_video.aac'
+    # extract audio
+    try:
+        ffmpeg.input(source_video).output(audio_output, acodec='copy').run()
+        # merge audio and video
+        stream = ffmpeg.input(save_path)
+        audio = ffmpeg.input(audio_output)
+        ffmpeg.output(stream.video, audio.audio, save_path.replace('_noaudio.mp4', '.mp4'), vcodec='copy', acodec='aac').run()
+    
+        os.remove(save_path)
+        os.remove(audio_output)
+    except:
+        shutil.move(
+            save_path,
+            save_path.replace('_noaudio.mp4', '.mp4')
+        )
+    
+    return save_path.replace('_noaudio.mp4', '.mp4')