page demo

.gitignore

@@ -0,0 +1 @@
+results/*

app.py

@@ -14,7 +14,7 @@ from inference import inference as inference_main
 def train_model(video, config):
     output_dir = 'results'
     os.makedirs(output_dir, exist_ok=True)
-    cur_save_dir = os.path.join(output_dir, str(len(os.listdir(output_dir))).zfill(2))
+    cur_save_dir = os.path.join(output_dir, 'custom')
 
     config.dataset.single_video_path = video
     config.train.output_dir = cur_save_dir
@@ -100,6 +100,12 @@ def update_preview_video(checkpoint_dir):
 
 
 if __name__ == "__main__":
+
+    if os.path.exists('results/custom'):
+        os.system('rm -rf results/custom')
+    if os.path.exists('outputs'):
+        os.system('rm -rf outputs')
+
     inject_motion_embeddings_combinations = ['down 1280','up 1280','down 640','up 640']
     default_motion_embeddings_combinations = ['down 1280','up 1280']
 

attn_ctrl.py

@@ -0,0 +1,264 @@
+import abc
+
+LOW_RESOURCE = False
+import torch 
+import cv2 
+import torch
+import os
+import numpy as np
+from collections import defaultdict
+from functools import partial
+from typing import Any, Dict, Optional
+
+def register_attention_control(unet, config=None):
+
+    def BasicTransformerBlock_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,
+        added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = 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.norm_type == "ada_norm":
+            norm_hidden_states = self.norm1(hidden_states, timestep)
+        elif self.norm_type == "ada_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.norm_type in ["layer_norm", "layer_norm_i2vgen"]:
+            norm_hidden_states = self.norm1(hidden_states)
+        elif self.norm_type == "ada_norm_continuous":
+            norm_hidden_states = self.norm1(hidden_states, added_cond_kwargs["pooled_text_emb"])
+        elif self.norm_type == "ada_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")
+
+        # save the origin_hidden_states w/o pos_embed, for the use of motion v embedding
+        origin_hidden_states = None
+        if self.pos_embed is not None or hasattr(self.attn1,'vSpatial'):
+            origin_hidden_states = norm_hidden_states.clone()
+            if cross_attention_kwargs is None:
+                cross_attention_kwargs = {}
+            cross_attention_kwargs["origin_hidden_states"] = origin_hidden_states
+
+        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.norm_type == "ada_norm_zero":
+            attn_output = gate_msa.unsqueeze(1) * attn_output
+        elif self.norm_type == "ada_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.norm_type == "ada_norm":
+                norm_hidden_states = self.norm2(hidden_states, timestep)
+            elif self.norm_type in ["ada_norm_zero", "layer_norm", "layer_norm_i2vgen"]:
+                norm_hidden_states = self.norm2(hidden_states)
+            elif self.norm_type == "ada_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
+            elif self.norm_type == "ada_norm_continuous":
+                norm_hidden_states = self.norm2(hidden_states, added_cond_kwargs["pooled_text_emb"])
+            else:
+                raise ValueError("Incorrect norm")
+
+            if self.pos_embed is not None and self.norm_type != "ada_norm_single":
+                # save the origin_hidden_states
+                origin_hidden_states = norm_hidden_states.clone()
+                norm_hidden_states = self.pos_embed(norm_hidden_states)
+                cross_attention_kwargs["origin_hidden_states"] = origin_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
+            # delete the origin_hidden_states
+            if cross_attention_kwargs is not None and "origin_hidden_states" in cross_attention_kwargs:
+                cross_attention_kwargs.pop("origin_hidden_states")
+
+        # 4. Feed-forward
+        # i2vgen doesn't have this norm 🤷‍♂️
+        if self.norm_type == "ada_norm_continuous":
+            norm_hidden_states = self.norm3(hidden_states, added_cond_kwargs["pooled_text_emb"])
+        elif not self.norm_type == "ada_norm_single":
+            norm_hidden_states = self.norm3(hidden_states)
+
+        if self.norm_type == "ada_norm_zero":
+            norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
+
+        if self.norm_type == "ada_norm_single":
+            norm_hidden_states = self.norm2(hidden_states)
+            norm_hidden_states = norm_hidden_states * (1 + scale_mlp) + shift_mlp
+
+        if self._chunk_size is not None:
+            # "feed_forward_chunk_size" can be used to save memory
+            ff_output = _chunked_feed_forward(
+                self.ff, norm_hidden_states, self._chunk_dim, self._chunk_size, lora_scale=lora_scale
+            )
+        else:
+            ff_output = self.ff(norm_hidden_states, scale=lora_scale)
+
+        if self.norm_type == "ada_norm_zero":
+            ff_output = gate_mlp.unsqueeze(1) * ff_output
+        elif self.norm_type == "ada_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
+
+
+    def temp_attn_forward(self, additional_info=None):
+        to_out = self.to_out
+        if type(to_out) is torch.nn.modules.container.ModuleList:
+            to_out = self.to_out[0]
+        else:
+            to_out = self.to_out
+
+        def forward(hidden_states, encoder_hidden_states=None, attention_mask=None,temb=None,origin_hidden_states=None):
+            
+            residual = hidden_states
+
+            if self.spatial_norm is not None:
+                hidden_states = self.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, sequence_length, _ = (
+                hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+            )
+
+            attention_mask = self.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+
+            if self.group_norm is not None:
+                hidden_states = self.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
+
+            if encoder_hidden_states is None:
+                encoder_hidden_states = hidden_states
+            elif self.norm_cross:
+                encoder_hidden_states = self.norm_encoder_hidden_states(encoder_hidden_states)
+
+            query = self.to_q(hidden_states)
+            key = self.to_k(encoder_hidden_states)
+            
+            # strategies to manipulate the motion value embedding
+            if additional_info is not None:
+                # empirically, in the inference stage of camera motion
+                # discarding the motion value embedding improves the text similarity of the generated video
+                if additional_info['removeMFromV']:
+                    value = self.to_v(origin_hidden_states)
+                elif hasattr(self,'vSpatial'):
+                    # during inference, the debiasing operation helps to generate more diverse videos
+                    # refer to the 'Figure.3 Right' in the paper for more details
+                    if additional_info['vSpatial_frameSubtraction']:
+                        value = self.to_v(self.vSpatial.forward_frameSubtraction(origin_hidden_states))
+                    # during training, do not apply debias operation for motion learning
+                    else:
+                        value = self.to_v(self.vSpatial(origin_hidden_states))
+                else:
+                    value = self.to_v(origin_hidden_states)
+            else:
+                value = self.to_v(encoder_hidden_states)
+
+
+            query = self.head_to_batch_dim(query)
+            key = self.head_to_batch_dim(key)
+            value = self.head_to_batch_dim(value)
+
+            attention_probs = self.get_attention_scores(query, key, attention_mask)
+
+            hidden_states = torch.bmm(attention_probs, value)
+            hidden_states = self.batch_to_head_dim(hidden_states)
+
+            # linear proj
+            hidden_states = to_out(hidden_states)
+
+            if input_ndim == 4:
+                hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
+
+            if self.residual_connection:
+                hidden_states = hidden_states + residual
+
+            hidden_states = hidden_states / self.rescale_output_factor
+
+            return hidden_states
+        return forward
+
+    def register_recr(net_, count, name, config=None):
+
+        if net_.__class__.__name__ == 'BasicTransformerBlock':
+            BasicTransformerBlock_forward_ = partial(BasicTransformerBlock_forward, net_)
+            net_.forward = BasicTransformerBlock_forward_
+
+        if net_.__class__.__name__ == 'Attention':
+            block_name = name.split('.attn')[0] 
+            if config is not None and block_name in set([l.split('.attn')[0].split('.pos_embed')[0] for l in config.model.embedding_layers]):
+                additional_info = {}
+                additional_info['layer_name'] = name
+                additional_info['removeMFromV'] = config.strategy.get('removeMFromV', False)
+                additional_info['vSpatial_frameSubtraction'] = config.strategy.get('vSpatial_frameSubtraction', False)
+                net_.forward = temp_attn_forward(net_, additional_info)
+                print('register Motion V embedding at ', block_name)
+                return count + 1
+            else:
+                return count
+
+        elif hasattr(net_, 'children'):
+            for net_name, net__ in dict(net_.named_children()).items():
+                count = register_recr(net__, count, name = name + '.' + net_name, config=config)
+        return count
+
+    sub_nets = unet.named_children()
+    
+    for net in sub_nets:
+        register_recr(net[1], 0,name = net[0], config=config)
+
+    
+

configs/config.yaml

@@ -0,0 +1,67 @@
+model:
+  type: unet
+  pretrained_model_path: cerspense/zeroscope_v2_576w
+  motion_embeddings:
+    combinations:
+    - - down
+      - 1280
+    - - up
+      - 1280
+  # unet can be either 'videoCrafter2' or 'zeroscope_v2_576w', the former produces better video quality
+  unet: videoCrafter2
+      
+train:
+  output_dir: ./results
+  validation_steps: 2000
+  checkpointing_steps: 50
+  checkpointing_start: 200
+  train_batch_size: 1
+  max_train_steps: 400
+  gradient_accumulation_steps: 1
+  cache_latents: true
+  cached_latent_dir: null
+  logger_type: tensorboard
+  mixed_precision: fp16
+  use_8bit_adam: false
+  resume_from_checkpoint: null
+  resume_step: null
+
+dataset:
+  type:
+  - single_video
+  single_video_path: ./assets/car-roundabout-24.mp4
+  single_video_prompt: 'A car turnaround in a city street'
+  width: 576
+  height: 320
+  n_sample_frames: 24
+  fps: 8
+  sample_start_idx: 1
+  frame_step: 1
+  use_bucketing: false
+  use_caption: false
+
+loss:
+  type: BaseLoss
+  learning_rate: 0.02
+  lr_scheduler: constant
+  lr_warmup_steps: 0
+ 
+noise_init:
+  type: BlendInit
+  noise_prior: 0.5
+
+val:
+  prompt:
+    - "A skateboard slides along a city lane"
+  negative_prompt: ""
+  sample_preview: true
+  width: 576
+  height: 320
+  num_frames: 24
+  num_inference_steps: 30
+  guidance_scale: 12.0
+  seeds: [0]
+
+strategy:
+  vSpatial_frameSubtraction: false
+  removeMFromV: false

dataset/__init__.py

@@ -0,0 +1,5 @@
+from .cached_dataset import CachedDataset
+from .image_dataset import ImageDataset
+from .single_video_dataset import SingleVideoDataset
+from .video_folder_dataset import VideoFolderDataset
+from .video_json_dataset import VideoJsonDataset

dataset/cached_dataset.py

@@ -0,0 +1,17 @@
+from utils.dataset_utils import *
+
+class CachedDataset(Dataset):
+    def __init__(self,cache_dir: str = ''):
+        self.cache_dir = cache_dir
+        self.cached_data_list = self.get_files_list()
+
+    def get_files_list(self):
+        tensors_list = [f"{self.cache_dir}/{x}" for x in os.listdir(self.cache_dir) if x.endswith('.pt')]
+        return sorted(tensors_list)
+
+    def __len__(self):
+        return len(self.cached_data_list)
+
+    def __getitem__(self, index):
+        cached_latent = torch.load(self.cached_data_list[index], map_location='cuda:0')
+        return cached_latent

dataset/image_dataset.py

@@ -0,0 +1,95 @@
+from utils.dataset_utils import *
+
+class ImageDataset(Dataset):
+    
+    def __init__(
+        self,
+        tokenizer = None,
+        width: int = 256,
+        height: int = 256,
+        base_width: int = 256,
+        base_height: int = 256,
+        use_caption:     bool = False,
+        image_dir: str = '',
+        single_img_prompt: str = '',
+        use_bucketing: bool = False,
+        fallback_prompt: str = '',
+        **kwargs
+    ):
+        self.tokenizer = tokenizer
+        self.img_types = (".png", ".jpg", ".jpeg", '.bmp')
+        self.use_bucketing = use_bucketing
+
+        self.image_dir = self.get_images_list(image_dir)
+        self.fallback_prompt = fallback_prompt
+
+        self.use_caption = use_caption
+        self.single_img_prompt = single_img_prompt
+
+        self.width = width
+        self.height = height
+
+    def get_images_list(self, image_dir):
+        if os.path.exists(image_dir):
+            imgs = [x for x in os.listdir(image_dir) if x.endswith(self.img_types)]
+            full_img_dir = []
+
+            for img in imgs: 
+                full_img_dir.append(f"{image_dir}/{img}")
+
+            return sorted(full_img_dir)
+
+        return ['']
+
+    def image_batch(self, index):
+        train_data = self.image_dir[index]
+        img = train_data
+
+        try:
+            img = torchvision.io.read_image(img, mode=torchvision.io.ImageReadMode.RGB)
+        except:
+            img = T.transforms.PILToTensor()(Image.open(img).convert("RGB"))
+
+        width = self.width
+        height = self.height
+
+        if self.use_bucketing:
+            _, h, w = img.shape
+            width, height = sensible_buckets(width, height, w, h)
+              
+        resize = T.transforms.Resize((height, width), antialias=True)
+
+        img = resize(img) 
+        img = repeat(img, 'c h w -> f c h w', f=16)
+
+        prompt = get_text_prompt(
+            file_path=train_data,
+            text_prompt=self.single_img_prompt,
+            fallback_prompt=self.fallback_prompt,
+            ext_types=self.img_types,  
+            use_caption=True
+        )
+        prompt_ids = get_prompt_ids(prompt, self.tokenizer)
+
+        return img, prompt, prompt_ids
+
+    @staticmethod
+    def __getname__(): return 'image'
+    
+    def __len__(self):
+        # Image directory
+        if os.path.exists(self.image_dir[0]):
+            return len(self.image_dir)
+        else:
+            return 0
+
+    def __getitem__(self, index):
+        img, prompt, prompt_ids = self.image_batch(index)
+        example = {
+            "pixel_values": (img / 127.5 - 1.0),
+            "prompt_ids": prompt_ids[0],
+            "text_prompt": prompt, 
+            'dataset': self.__getname__()
+        }
+
+        return example

dataset/single_video_dataset.py

@@ -0,0 +1,106 @@
+from utils.dataset_utils import *
+
+class SingleVideoDataset(Dataset):
+    def __init__(
+        self,
+            tokenizer = None,
+            width: int = 256,
+            height: int = 256,
+            n_sample_frames: int = 4,
+            frame_step: int = 1,
+            single_video_path: str = "",
+            single_video_prompt: str = "",
+            use_caption: bool = False,
+            use_bucketing: bool = False,
+            **kwargs
+    ):
+        self.tokenizer = tokenizer
+        self.use_bucketing = use_bucketing
+        self.frames = []
+        self.index = 1
+
+        self.vid_types = (".mp4", ".avi", ".mov", ".webm", ".flv", ".mjpeg")
+        self.n_sample_frames = n_sample_frames
+        self.frame_step = frame_step
+
+        self.single_video_path = single_video_path
+        self.single_video_prompt = single_video_prompt
+
+        self.width = width
+        self.height = height
+        
+    def create_video_chunks(self):
+        vr = decord.VideoReader(self.single_video_path)
+        vr_range = range(0, len(vr), self.frame_step)
+
+        self.frames = list(self.chunk(vr_range, self.n_sample_frames))
+        return self.frames
+
+    def chunk(self, it, size):
+        it = iter(it)
+        return iter(lambda: tuple(islice(it, size)), ())
+
+    def get_frame_batch(self, vr, resize=None):
+        index = self.index
+        frames = vr.get_batch(self.frames[self.index])
+
+        if type(frames) == decord.ndarray.NDArray: 
+            frames = torch.from_numpy(frames.asnumpy())
+
+        video = rearrange(frames, "f h w c -> f c h w")
+
+        if resize is not None: video = resize(video)
+        return video
+
+    def get_frame_buckets(self, vr):
+        h, w, c = vr[0].shape
+        width, height = sensible_buckets(self.width, self.height, w, h)
+        resize = T.transforms.Resize((height, width), antialias=True)
+
+        return resize
+    
+    def process_video_wrapper(self, vid_path):
+        video, vr = process_video(
+                vid_path,
+                self.use_bucketing,
+                self.width, 
+                self.height, 
+                self.get_frame_buckets, 
+                self.get_frame_batch
+            )
+        
+        return video, vr 
+
+    def single_video_batch(self, index):
+        train_data = self.single_video_path
+        self.index = index
+
+        if train_data.endswith(self.vid_types):
+            video, _ = self.process_video_wrapper(train_data)
+
+            prompt = self.single_video_prompt
+            prompt_ids = get_prompt_ids(prompt, self.tokenizer)
+
+            return video, prompt, prompt_ids
+        else:
+            raise ValueError(f"Single video is not a video type. Types: {self.vid_types}")
+    
+    @staticmethod
+    def __getname__(): return 'single_video'
+
+    def __len__(self):
+        
+        return len(self.create_video_chunks())
+
+    def __getitem__(self, index):
+
+        video, prompt, prompt_ids = self.single_video_batch(index)
+
+        example = {
+            "pixel_values": (video / 127.5 - 1.0),
+            "prompt_ids": prompt_ids[0],
+            "text_prompt": prompt,
+            'dataset': self.__getname__()
+        }
+
+        return example

dataset/video_folder_dataset.py

@@ -0,0 +1,90 @@
+from utils.dataset_utils import *
+
+class VideoFolderDataset(Dataset):
+    def __init__(
+        self,
+        tokenizer=None,
+        width: int = 256,
+        height: int = 256,
+        n_sample_frames: int = 16,
+        fps: int = 8,
+        path: str = "./data",
+        fallback_prompt: str = "",
+        use_bucketing: bool = False,
+        **kwargs
+    ):
+        self.tokenizer = tokenizer
+        self.use_bucketing = use_bucketing
+
+        self.fallback_prompt = fallback_prompt
+
+        self.video_files = glob(f"{path}/*.mp4")
+
+        self.width = width
+        self.height = height
+
+        self.n_sample_frames = n_sample_frames
+        self.fps = fps
+
+    def get_frame_buckets(self, vr):
+        h, w, c = vr[0].shape
+        width, height = sensible_buckets(self.width, self.height, w, h)
+        resize = T.transforms.Resize((height, width), antialias=True)
+
+        return resize
+
+    def get_frame_batch(self, vr, resize=None):
+        n_sample_frames = self.n_sample_frames
+        native_fps = vr.get_avg_fps()
+        
+        every_nth_frame = max(1, round(native_fps / self.fps))
+        every_nth_frame = min(len(vr), every_nth_frame)
+        
+        effective_length = len(vr) // every_nth_frame
+        if effective_length < n_sample_frames:
+            n_sample_frames = effective_length
+
+        effective_idx = random.randint(0, (effective_length - n_sample_frames))
+        idxs = every_nth_frame * np.arange(effective_idx, effective_idx + n_sample_frames)
+
+        video = vr.get_batch(idxs)
+        video = rearrange(video, "f h w c -> f c h w")
+
+        if resize is not None: video = resize(video)
+        return video, vr
+        
+    def process_video_wrapper(self, vid_path):
+        video, vr = process_video(
+                vid_path,
+                self.use_bucketing,
+                self.width, 
+                self.height, 
+                self.get_frame_buckets, 
+                self.get_frame_batch
+            )
+        return video, vr
+    
+    def get_prompt_ids(self, prompt):
+        return self.tokenizer(
+            prompt,
+            truncation=True,
+            padding="max_length",
+            max_length=self.tokenizer.model_max_length,
+            return_tensors="pt",
+        ).input_ids
+
+    @staticmethod
+    def __getname__(): return 'folder'
+
+    def __len__(self):
+        return len(self.video_files)
+
+    def __getitem__(self, index):
+
+        video, _ = self.process_video_wrapper(self.video_files[index])
+
+        prompt = self.fallback_prompt
+
+        prompt_ids = self.get_prompt_ids(prompt)
+
+        return {"pixel_values": (video[0] / 127.5 - 1.0), "prompt_ids": prompt_ids[0], "text_prompt": prompt, 'dataset': self.__getname__()}

dataset/video_json_dataset.py

@@ -0,0 +1,183 @@
+from utils.dataset_utils import *
+
+# https://github.com/ExponentialML/Video-BLIP2-Preprocessor
+class VideoJsonDataset(Dataset):
+    def __init__(
+            self,
+            tokenizer = None,
+            width: int = 256,
+            height: int = 256,
+            n_sample_frames: int = 4,
+            sample_start_idx: int = 1,
+            frame_step: int = 1,
+            json_path: str ="",
+            json_data = None,
+            vid_data_key: str = "video_path",
+            preprocessed: bool = False,
+            use_bucketing: bool = False,
+            **kwargs
+    ):
+        self.vid_types = (".mp4", ".avi", ".mov", ".webm", ".flv", ".mjpeg")
+        self.use_bucketing = use_bucketing
+        self.tokenizer = tokenizer
+        self.preprocessed = preprocessed
+        
+        self.vid_data_key = vid_data_key
+        self.train_data = self.load_from_json(json_path, json_data)
+
+        self.width = width
+        self.height = height
+
+        self.n_sample_frames = n_sample_frames
+        self.sample_start_idx = sample_start_idx
+        self.frame_step = frame_step
+
+    def build_json(self, json_data):
+        extended_data = []
+        for data in json_data['data']:
+            for nested_data in data['data']:
+                self.build_json_dict(
+                    data, 
+                    nested_data, 
+                    extended_data
+                )
+        json_data = extended_data
+        return json_data
+
+    def build_json_dict(self, data, nested_data, extended_data):
+        clip_path = nested_data['clip_path'] if 'clip_path' in nested_data else None
+        
+        extended_data.append({
+            self.vid_data_key: data[self.vid_data_key],
+            'frame_index': nested_data['frame_index'],
+            'prompt': nested_data['prompt'],
+            'clip_path': clip_path
+        })
+        
+    def load_from_json(self, path, json_data):
+        try:
+            with open(path) as jpath:
+                print(f"Loading JSON from {path}")
+                json_data = json.load(jpath)
+
+                return self.build_json(json_data)
+
+        except:
+            self.train_data = []
+            print("Non-existant JSON path. Skipping.")
+            
+    def validate_json(self, base_path, path):
+        return os.path.exists(f"{base_path}/{path}")
+
+    def get_frame_range(self, vr):
+        return get_video_frames(
+            vr, 
+            self.sample_start_idx, 
+            self.frame_step, 
+            self.n_sample_frames
+        )
+    
+    def get_vid_idx(self, vr, vid_data=None):
+        frames = self.n_sample_frames
+
+        if vid_data is not None:
+            idx = vid_data['frame_index']
+        else:
+            idx = self.sample_start_idx
+
+        return idx
+
+    def get_frame_buckets(self, vr):
+        _, h, w = vr[0].shape        
+        width, height = sensible_buckets(self.width, self.height, h, w)
+        # width, height = self.width, self.height
+        resize = T.transforms.Resize((height, width), antialias=True)
+
+        return resize
+
+    def get_frame_batch(self, vr, resize=None):
+        frame_range = self.get_frame_range(vr)
+        frames = vr.get_batch(frame_range)
+        video = rearrange(frames, "f h w c -> f c h w")
+
+        if resize is not None: video = resize(video)
+        return video
+
+    def process_video_wrapper(self, vid_path):
+        video, vr = process_video(
+                vid_path,
+                self.use_bucketing,
+                self.width, 
+                self.height, 
+                self.get_frame_buckets, 
+                self.get_frame_batch
+            )
+        
+        return video, vr 
+
+    def train_data_batch(self, index):
+
+        # If we are training on individual clips.
+        if 'clip_path' in self.train_data[index] and \
+            self.train_data[index]['clip_path'] is not None:
+
+            vid_data = self.train_data[index]
+
+            clip_path = vid_data['clip_path']
+            
+            # Get video prompt
+            prompt = vid_data['prompt']
+
+            video, _ = self.process_video_wrapper(clip_path)
+
+            prompt_ids = get_prompt_ids(prompt, self.tokenizer)
+
+            return video, prompt, prompt_ids
+
+         # Assign train data
+        train_data = self.train_data[index]
+        
+        # Get the frame of the current index.
+        self.sample_start_idx = train_data['frame_index']
+        
+        # Initialize resize
+        resize = None
+
+        video, vr = self.process_video_wrapper(train_data[self.vid_data_key])
+
+        # Get video prompt
+        prompt = train_data['prompt']
+        vr.seek(0)
+
+        prompt_ids = get_prompt_ids(prompt, self.tokenizer)
+
+        return video, prompt, prompt_ids
+
+    @staticmethod
+    def __getname__(): return 'json'
+
+    def __len__(self):
+        if self.train_data is not None:
+            return len(self.train_data)
+        else: 
+            return 0
+
+    def __getitem__(self, index):
+        
+        # Initialize variables
+        video = None
+        prompt = None
+        prompt_ids = None
+
+        # Use default JSON training
+        if self.train_data is not None:
+            video, prompt, prompt_ids = self.train_data_batch(index)
+
+        example = {
+            "pixel_values": (video / 127.5 - 1.0),
+            "prompt_ids": prompt_ids[0],
+            "text_prompt": prompt,
+            'dataset': self.__getname__()
+        }
+
+        return example

inference.py

@@ -0,0 +1,133 @@
+import torch
+from diffusers import DiffusionPipeline, DPMSolverMultistepScheduler
+from train import export_to_video
+from models.unet.motion_embeddings import load_motion_embeddings
+from noise_init.blend_init import BlendInit
+from noise_init.blend_freq_init import BlendFreqInit
+from noise_init.fft_init import FFTInit
+from noise_init.freq_init import FreqInit
+from attn_ctrl import register_attention_control
+import numpy as np
+import os
+from omegaconf import OmegaConf
+
+def get_pipe(embedding_dir='baseline',config=None,noisy_latent=None, video_round=None):
+
+    # load video generation model
+    pipe = DiffusionPipeline.from_pretrained(config.model.pretrained_model_path,torch_dtype=torch.float16)
+
+    # use videocrafterv2 unet
+    if config.model.unet == 'videoCrafter2':
+        from models.unet.unet_3d_condition import UNet3DConditionModel
+        # unet = UNet3DConditionModel.from_pretrained("adamdad/videocrafterv2_diffusers",subfolder='unet',torch_dtype=torch.float16)
+        unet = UNet3DConditionModel.from_pretrained("adamdad/videocrafterv2_diffusers",torch_dtype=torch.float16)
+        pipe.unet = unet
+
+    # pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
+    pipe.enable_model_cpu_offload()
+
+    # memory optimization
+    pipe.enable_vae_slicing()
+
+    # if 'vanilla' not in embedding_dir:
+
+    noisy_latent = torch.load(f'{embedding_dir}/cached_latents/cached_0.pt')['inversion_noise'][None,]
+    if video_round is None:
+        motion_embed = torch.load(f'{embedding_dir}/motion_embed.pt')
+    else:
+        motion_embed = torch.load(f'{embedding_dir}/{video_round}/motion_embed.pt')
+    load_motion_embeddings(
+        pipe.unet, 
+        motion_embed, 
+    )
+    config.model['embedding_layers'] = list(motion_embed.keys())
+
+    return pipe, config, noisy_latent
+
+def inference(embedding_dir='vanilla',
+              video_round=None, 
+              prompt=None,
+              save_dir=None,
+              seed=None,
+              motion_type=None,
+              inference_steps=30
+              ):
+    
+    # check motion type is valid
+    if motion_type != 'camera' and \
+        motion_type != 'object' and \
+            motion_type != 'hybrid':
+        raise ValueError('Invalid motion type')
+
+    if seed is None:
+        seed = 0
+
+    # load motion embedding
+    noisy_latent = None
+
+    config = OmegaConf.load(f'{embedding_dir}/config.yaml') 
+    
+
+    # different motion type assigns different strategy
+    if motion_type == 'camera':
+        config['strategy']['removeMFromV'] = True
+        
+    elif motion_type == 'object' or motion_type == 'hybrid':
+        config['strategy']['vSpatial_frameSubtraction'] = True
+
+    
+    pipe, config, noisy_latent = get_pipe(embedding_dir=embedding_dir,config=config,noisy_latent=noisy_latent,video_round=video_round)
+    n_frames = config.val.num_frames
+
+    shape = (config.val.height,config.val.width)
+    os.makedirs(save_dir,exist_ok=True)
+    
+
+    cur_save_dir = f'{save_dir}/{"_".join(prompt.split())}.mp4'
+
+    register_attention_control(pipe.unet,config=config)
+
+    if noisy_latent is not None:
+        torch.manual_seed(seed)
+        noise = torch.randn_like(noisy_latent)
+        init_noise = BlendInit(noisy_latent, noise, noise_prior=0.5)
+    else:
+        init_noise = None
+
+    input_init_noise = init_noise.clone() if not init_noise is None else None
+    video_frames = pipe(
+        prompt=prompt,
+        num_inference_steps=inference_steps,
+        guidance_scale=12,
+        height=shape[0],
+        width=shape[1],
+        num_frames=n_frames,
+        generator=torch.Generator("cuda").manual_seed(seed),
+        latents=input_init_noise,
+    ).frames[0]
+
+    video_path = export_to_video(video_frames,output_video_path=cur_save_dir,fps=8)
+    
+    return video_path
+
+
+if __name__ =="__main__":
+       
+    prompts = ["A skateboard slides along a city lane",
+                "A tank is running in the desert.",
+                "A toy train chugs around a roundabout tree"]
+
+    
+    embedding_dir = './results'
+    video_round = 'checkpoint-250'
+    save_dir = f'outputs'
+
+    inference(
+        embedding_dir=embedding_dir,
+        prompt=prompts, 
+        video_round=video_round,
+        save_dir=save_dir,
+        motion_type='hybrid',
+        seed=100
+        )
+    

loss/__init__.py

@@ -0,0 +1,4 @@
+from .base_loss import BaseLoss
+from .debiased_hybrid_loss import DebiasedHybridLoss
+from .debiased_temporal_loss import DebiasedTemporalLoss
+from .motion_distillation_loss import MotionDistillationLoss

loss/base_loss.py

@@ -0,0 +1,75 @@
+import torch
+import torch.nn.functional as F
+from utils.func_utils import tensor_to_vae_latent, sample_noise
+
+def BaseLoss(
+        train_loss_temporal, 
+        accelerator,
+        optimizers,
+        lr_schedulers,
+        unet, 
+        vae, 
+        text_encoder,
+        noise_scheduler, 
+        batch,
+        step,
+        config
+    ):
+    cache_latents = config.train.cache_latents
+
+    if not cache_latents:
+        latents = tensor_to_vae_latent(batch["pixel_values"], vae)
+    else:
+        latents = batch["latents"]
+
+    # Sample noise that we'll add to the latents
+    # use_offset_noise = use_offset_noise and not rescale_schedule
+        
+    noise = sample_noise(latents, 0.1, False)
+    bsz = latents.shape[0]
+
+    # Sample a random timestep for each video
+    timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (bsz,), device=latents.device)
+    timesteps = timesteps.long()
+
+    # Add noise to the latents according to the noise magnitude at each timestep
+    # (this is the forward diffusion process)
+    noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)
+
+    # *Potentially* Fixes gradient checkpointing training.
+    # See: https://github.com/prigoyal/pytorch_memonger/blob/master/tutorial/Checkpointing_for_PyTorch_models.ipynb
+    # if kwargs.get('eval_train', False):
+    #     unet.eval()
+    #     text_encoder.eval()
+
+    # Encode text embeddings
+    token_ids = batch['prompt_ids']
+    encoder_hidden_states = text_encoder(token_ids)[0]
+    detached_encoder_state = encoder_hidden_states.clone().detach()
+
+    # Get the target for loss depending on the prediction type
+    if noise_scheduler.config.prediction_type == "epsilon":
+        target = noise
+
+    elif noise_scheduler.config.prediction_type == "v_prediction":
+        target = noise_scheduler.get_velocity(latents, noise, timesteps)
+
+    else:
+        raise ValueError(f"Unknown prediction type {noise_scheduler.config.prediction_type}")
+
+    encoder_hidden_states = detached_encoder_state
+
+
+    # optimization
+    model_pred = unet(noisy_latents, timesteps, encoder_hidden_states=encoder_hidden_states).sample
+    loss_temporal = F.mse_loss(model_pred.float(), target.float(), reduction="mean")
+
+    avg_loss_temporal = accelerator.gather(loss_temporal.repeat(config.train.train_batch_size)).mean()
+    train_loss_temporal += avg_loss_temporal.item() / config.train.gradient_accumulation_steps
+
+    accelerator.backward(loss_temporal)
+    optimizers[0].step()
+    lr_schedulers[0].step()
+
+    return loss_temporal, train_loss_temporal
+

loss/debiased_hybrid_loss.py

@@ -0,0 +1,149 @@
+import torch
+from torchvision import transforms
+import torch.nn.functional as F
+import random
+
+from utils.lora import extract_lora_child_module
+from utils.func_utils import tensor_to_vae_latent, sample_noise
+
+def DebiasedHybridLoss(
+        train_loss_temporal,
+        accelerator,
+        optimizers,
+        lr_schedulers,
+        unet, 
+        vae, 
+        text_encoder,
+        noise_scheduler, 
+        batch,
+        step,
+        config,
+        random_hflip_img=False,
+        spatial_lora_num=1
+    ):
+    mask_spatial_lora = random.uniform(0, 1) < 0.2
+    cache_latents = config.train.cache_latents
+
+
+
+    if not cache_latents:
+        latents = tensor_to_vae_latent(batch["pixel_values"], vae)
+    else:
+        latents = batch["latents"]
+
+    # Sample noise that we'll add to the latents
+    # use_offset_noise = use_offset_noise and not rescale_schedule
+        
+    noise = sample_noise(latents, 0.1, False)
+    bsz = latents.shape[0]
+
+    # Sample a random timestep for each video
+    timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (bsz,), device=latents.device)
+    timesteps = timesteps.long()
+
+    # Add noise to the latents according to the noise magnitude at each timestep
+    # (this is the forward diffusion process)
+    noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)
+
+    # *Potentially* Fixes gradient checkpointing training.
+    # See: https://github.com/prigoyal/pytorch_memonger/blob/master/tutorial/Checkpointing_for_PyTorch_models.ipynb
+    # if kwargs.get('eval_train', False):
+    #     unet.eval()
+    #     text_encoder.eval()
+
+    # Encode text embeddings
+    token_ids = batch['prompt_ids']
+    encoder_hidden_states = text_encoder(token_ids)[0]
+    detached_encoder_state = encoder_hidden_states.clone().detach()
+
+    # Get the target for loss depending on the prediction type
+    if noise_scheduler.config.prediction_type == "epsilon":
+        target = noise
+
+    elif noise_scheduler.config.prediction_type == "v_prediction":
+        target = noise_scheduler.get_velocity(latents, noise, timesteps)
+
+    else:
+        raise ValueError(f"Unknown prediction type {noise_scheduler.config.prediction_type}")
+
+    encoder_hidden_states = detached_encoder_state
+
+
+    # optimization
+    if mask_spatial_lora:
+        loras = extract_lora_child_module(unet, target_replace_module=["Transformer2DModel"])
+        for lora_i in loras:
+            lora_i.scale = 0.
+        loss_spatial = None
+    else:
+        loras = extract_lora_child_module(unet, target_replace_module=["Transformer2DModel"])
+
+        if spatial_lora_num == 1:
+            for lora_i in loras:
+                lora_i.scale = 1.
+        else:
+            for lora_i in loras:
+                lora_i.scale = 0.
+
+            for lora_idx in range(0, len(loras), spatial_lora_num):
+                loras[lora_idx + step].scale = 1.
+
+        loras = extract_lora_child_module(unet, target_replace_module=["TransformerTemporalModel"])
+        if len(loras) > 0:
+            for lora_i in loras:
+                lora_i.scale = 0.
+
+        ran_idx = torch.randint(0, noisy_latents.shape[2], (1,)).item()
+
+        if random.uniform(0, 1) < random_hflip_img:
+            pixel_values_spatial = transforms.functional.hflip(
+                batch["pixel_values"][:, ran_idx, :, :, :]).unsqueeze(1)
+            latents_spatial = tensor_to_vae_latent(pixel_values_spatial, vae)
+            noise_spatial = sample_noise(latents_spatial, 0.1, False)
+            noisy_latents_input = noise_scheduler.add_noise(latents_spatial, noise_spatial, timesteps)
+            target_spatial = noise_spatial
+            model_pred_spatial = unet(noisy_latents_input, timesteps,
+                                    encoder_hidden_states=encoder_hidden_states).sample
+            loss_spatial = F.mse_loss(model_pred_spatial[:, :, 0, :, :].float(),
+                                    target_spatial[:, :, 0, :, :].float(), reduction="mean")
+        else:
+            noisy_latents_input = noisy_latents[:, :, ran_idx, :, :]
+            target_spatial = target[:, :, ran_idx, :, :]
+            model_pred_spatial = unet(noisy_latents_input.unsqueeze(2), timesteps,
+                                    encoder_hidden_states=encoder_hidden_states).sample
+            loss_spatial = F.mse_loss(model_pred_spatial[:, :, 0, :, :].float(),
+                                    target_spatial.float(), reduction="mean")
+
+
+    model_pred = unet(noisy_latents, timesteps, encoder_hidden_states=encoder_hidden_states).sample
+    loss_temporal = F.mse_loss(model_pred.float(), target.float(), reduction="mean")
+
+    beta = 1
+    alpha = (beta ** 2 + 1) ** 0.5
+    ran_idx = torch.randint(0, model_pred.shape[2], (1,)).item()
+    model_pred_decent = alpha * model_pred - beta * model_pred[:, :, ran_idx, :, :].unsqueeze(2)
+    target_decent = alpha * target - beta * target[:, :, ran_idx, :, :].unsqueeze(2)
+    loss_ad_temporal = F.mse_loss(model_pred_decent.float(), target_decent.float(), reduction="mean")
+    loss_temporal = loss_temporal + loss_ad_temporal
+
+    avg_loss_temporal = accelerator.gather(loss_temporal.repeat(config.train.train_batch_size)).mean()
+    train_loss_temporal += avg_loss_temporal.item() / config.train.gradient_accumulation_steps
+
+    if not mask_spatial_lora:
+        accelerator.backward(loss_spatial, retain_graph=True)
+        if spatial_lora_num == 1:
+            optimizers[1].step()
+        else:
+            optimizers[step+1].step()
+
+    accelerator.backward(loss_temporal)
+    optimizers[0].step()
+
+    if spatial_lora_num == 1:
+        lr_schedulers[1].step()
+    else:
+        lr_schedulers[1 + step].step()
+
+    lr_schedulers[0].step()
+
+    return loss_temporal, train_loss_temporal

loss/debiased_temporal_loss.py

@@ -0,0 +1,86 @@
+import torch
+import torch.nn.functional as F
+
+from utils.func_utils import tensor_to_vae_latent, sample_noise
+
+def DebiasedTemporalLoss(
+        train_loss_temporal,
+        accelerator,
+        optimizers,
+        lr_schedulers,
+        unet, 
+        vae, 
+        text_encoder,
+        noise_scheduler, 
+        batch,
+        step,
+        config
+    ):
+    cache_latents = config.train.cache_latents
+
+
+
+    if not cache_latents:
+        latents = tensor_to_vae_latent(batch["pixel_values"], vae)
+    else:
+        latents = batch["latents"]
+
+    # Sample noise that we'll add to the latents
+    # use_offset_noise = use_offset_noise and not rescale_schedule
+        
+    noise = sample_noise(latents, 0.1, False)
+    bsz = latents.shape[0]
+
+    # Sample a random timestep for each video
+    timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (bsz,), device=latents.device)
+    timesteps = timesteps.long()
+
+    # Add noise to the latents according to the noise magnitude at each timestep
+    # (this is the forward diffusion process)
+    noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)
+
+    # *Potentially* Fixes gradient checkpointing training.
+    # See: https://github.com/prigoyal/pytorch_memonger/blob/master/tutorial/Checkpointing_for_PyTorch_models.ipynb
+    # if kwargs.get('eval_train', False):
+    #     unet.eval()
+    #     text_encoder.eval()
+
+    # Encode text embeddings
+    token_ids = batch['prompt_ids']
+    encoder_hidden_states = text_encoder(token_ids)[0]
+    detached_encoder_state = encoder_hidden_states.clone().detach()
+
+    # Get the target for loss depending on the prediction type
+    if noise_scheduler.config.prediction_type == "epsilon":
+        target = noise
+
+    elif noise_scheduler.config.prediction_type == "v_prediction":
+        target = noise_scheduler.get_velocity(latents, noise, timesteps)
+
+    else:
+        raise ValueError(f"Unknown prediction type {noise_scheduler.config.prediction_type}")
+
+    encoder_hidden_states = detached_encoder_state
+
+
+    # optimization
+    model_pred = unet(noisy_latents, timesteps, encoder_hidden_states=encoder_hidden_states).sample
+    loss_temporal = F.mse_loss(model_pred.float(), target.float(), reduction="mean")
+
+    beta = 1
+    alpha = (beta ** 2 + 1) ** 0.5
+    ran_idx = torch.randint(0, model_pred.shape[2], (1,)).item()
+    model_pred_decent = alpha * model_pred - beta * model_pred[:, :, ran_idx, :, :].unsqueeze(2)
+    target_decent = alpha * target - beta * target[:, :, ran_idx, :, :].unsqueeze(2)
+    loss_ad_temporal = F.mse_loss(model_pred_decent.float(), target_decent.float(), reduction="mean")
+    loss_temporal = loss_temporal + loss_ad_temporal
+
+    avg_loss_temporal = accelerator.gather(loss_temporal.repeat(config.train.train_batch_size)).mean()
+    train_loss_temporal += avg_loss_temporal.item() / config.train.gradient_accumulation_steps
+
+    accelerator.backward(loss_temporal)
+    optimizers[0].step()
+
+    lr_schedulers[0].step()
+
+    return loss_temporal, train_loss_temporal

loss/motion_distillation_loss.py

@@ -0,0 +1,79 @@
+import torch
+import torch.nn.functional as F
+from utils.func_utils import tensor_to_vae_latent, sample_noise
+
+def MotionDistillationLoss(
+        train_loss_temporal, 
+        accelerator,
+        optimizers,
+        lr_schedulers,
+        unet, 
+        vae, 
+        text_encoder,
+        noise_scheduler, 
+        batch,
+        step,
+        config
+    ):
+    cache_latents = config.train.cache_latents
+
+    if not cache_latents:
+        latents = tensor_to_vae_latent(batch["pixel_values"], vae)
+    else:
+        latents = batch["latents"]
+
+    # Sample noise that we'll add to the latents
+    # use_offset_noise = use_offset_noise and not rescale_schedule
+        
+    noise = sample_noise(latents, 0.1, False)
+    bsz = latents.shape[0]
+
+    # Sample a random timestep for each video
+    timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (bsz,), device=latents.device)
+    timesteps = timesteps.long()
+
+    # Add noise to the latents according to the noise magnitude at each timestep
+    # (this is the forward diffusion process)
+    noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)
+
+    # *Potentially* Fixes gradient checkpointing training.
+    # See: https://github.com/prigoyal/pytorch_memonger/blob/master/tutorial/Checkpointing_for_PyTorch_models.ipynb
+    # if kwargs.get('eval_train', False):
+    #     unet.eval()
+    #     text_encoder.eval()
+
+    # Encode text embeddings
+    token_ids = batch['prompt_ids']
+    encoder_hidden_states = text_encoder(token_ids)[0]
+    detached_encoder_state = encoder_hidden_states.clone().detach()
+
+    # Get the target for loss depending on the prediction type
+    if noise_scheduler.config.prediction_type == "epsilon":
+        target = noise
+
+    elif noise_scheduler.config.prediction_type == "v_prediction":
+        target = noise_scheduler.get_velocity(latents, noise, timesteps)
+
+    else:
+        raise ValueError(f"Unknown prediction type {noise_scheduler.config.prediction_type}")
+
+    encoder_hidden_states = detached_encoder_state
+
+
+    # optimization
+    model_pred = unet(noisy_latents, timesteps, encoder_hidden_states=encoder_hidden_states).sample
+    
+    loss_temporal = 0
+    model_pred_reidual = torch.abs(model_pred[:,:,1:,:,:] - model_pred[:,:,:-1,:,:])
+    target_residual = torch.abs(target[:, :, 1:, :, :] - target[:, :, :-1, :, :])
+    loss_temporal = loss_temporal + (1 - F.cosine_similarity(model_pred_reidual, target_residual, dim=2).mean)
+
+    avg_loss_temporal = accelerator.gather(loss_temporal.repeat(config.train.train_batch_size)).mean()
+    train_loss_temporal += avg_loss_temporal.item() / config.train.gradient_accumulation_steps
+
+    accelerator.backward(loss_temporal)
+    optimizers[0].step()
+    lr_schedulers[0].step()
+
+    return loss_temporal, train_loss_temporal
+

models/dit/latte_t2v.py

@@ -0,0 +1,990 @@
+import torch
+
+import os
+import json
+
+from dataclasses import dataclass
+from einops import rearrange, repeat
+from typing import Any, Dict, Optional, Tuple
+from diffusers.models import Transformer2DModel
+from diffusers.utils import USE_PEFT_BACKEND, BaseOutput, deprecate
+from diffusers.models.embeddings import get_1d_sincos_pos_embed_from_grid, ImagePositionalEmbeddings, CaptionProjection, PatchEmbed, CombinedTimestepSizeEmbeddings
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.attention import BasicTransformerBlock
+from diffusers.models.lora import LoRACompatibleConv, LoRACompatibleLinear
+from diffusers.utils.torch_utils import maybe_allow_in_graph
+from diffusers.models.embeddings import SinusoidalPositionalEmbedding
+from diffusers.models.normalization import AdaLayerNorm, AdaLayerNormZero
+from diffusers.models.attention_processor import Attention
+from diffusers.models.activations import GEGLU, GELU, ApproximateGELU
+
+from dataclasses import dataclass
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+@maybe_allow_in_graph
+class GatedSelfAttentionDense(nn.Module):
+    r"""
+    A gated self-attention dense layer that combines visual features and object features.
+
+    Parameters:
+        query_dim (`int`): The number of channels in the query.
+        context_dim (`int`): The number of channels in the context.
+        n_heads (`int`): The number of heads to use for attention.
+        d_head (`int`): The number of channels in each head.
+    """
+
+    def __init__(self, query_dim: int, context_dim: int, n_heads: int, d_head: int):
+        super().__init__()
+
+        # we need a linear projection since we need cat visual feature and obj feature
+        self.linear = nn.Linear(context_dim, query_dim)
+
+        self.attn = Attention(query_dim=query_dim, heads=n_heads, dim_head=d_head)
+        self.ff = FeedForward(query_dim, activation_fn="geglu")
+
+        self.norm1 = nn.LayerNorm(query_dim)
+        self.norm2 = nn.LayerNorm(query_dim)
+
+        self.register_parameter("alpha_attn", nn.Parameter(torch.tensor(0.0)))
+        self.register_parameter("alpha_dense", nn.Parameter(torch.tensor(0.0)))
+
+        self.enabled = True
+
+    def forward(self, x: torch.Tensor, objs: torch.Tensor) -> torch.Tensor:
+        if not self.enabled:
+            return x
+
+        n_visual = x.shape[1]
+        objs = self.linear(objs)
+
+        x = x + self.alpha_attn.tanh() * self.attn(self.norm1(torch.cat([x, objs], dim=1)))[:, :n_visual, :]
+        x = x + self.alpha_dense.tanh() * self.ff(self.norm2(x))
+
+        return x
+
+class FeedForward(nn.Module):
+    r"""
+    A feed-forward layer.
+
+    Parameters:
+        dim (`int`): The number of channels in the input.
+        dim_out (`int`, *optional*): The number of channels in the output. If not given, defaults to `dim`.
+        mult (`int`, *optional*, defaults to 4): The multiplier to use for the hidden dimension.
+        dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
+        activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward.
+        final_dropout (`bool` *optional*, defaults to False): Apply a final dropout.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        dim_out: Optional[int] = None,
+        mult: int = 4,
+        dropout: float = 0.0,
+        activation_fn: str = "geglu",
+        final_dropout: bool = False,
+    ):
+        super().__init__()
+        inner_dim = int(dim * mult)
+        dim_out = dim_out if dim_out is not None else dim
+        linear_cls = LoRACompatibleLinear if not USE_PEFT_BACKEND else nn.Linear
+
+        if activation_fn == "gelu":
+            act_fn = GELU(dim, inner_dim)
+        if activation_fn == "gelu-approximate":
+            act_fn = GELU(dim, inner_dim, approximate="tanh")
+        elif activation_fn == "geglu":
+            act_fn = GEGLU(dim, inner_dim)
+        elif activation_fn == "geglu-approximate":
+            act_fn = ApproximateGELU(dim, inner_dim)
+
+        self.net = nn.ModuleList([])
+        # project in
+        self.net.append(act_fn)
+        # project dropout
+        self.net.append(nn.Dropout(dropout))
+        # project out
+        self.net.append(linear_cls(inner_dim, dim_out))
+        # FF as used in Vision Transformer, MLP-Mixer, etc. have a final dropout
+        if final_dropout:
+            self.net.append(nn.Dropout(dropout))
+
+    def forward(self, hidden_states: torch.Tensor, scale: float = 1.0) -> torch.Tensor:
+        compatible_cls = (GEGLU,) if USE_PEFT_BACKEND else (GEGLU, LoRACompatibleLinear)
+        for module in self.net:
+            if isinstance(module, compatible_cls):
+                hidden_states = module(hidden_states, scale)
+            else:
+                hidden_states = module(hidden_states)
+        return hidden_states
+
+@maybe_allow_in_graph
+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.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):
+        # 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 = self.norm3(hidden_states)
+            norm_hidden_states = norm_hidden_states * (1 + scale_mlp) + shift_mlp
+
+        if self._chunk_size is not None:
+            # "feed_forward_chunk_size" can be used to save memory
+            if norm_hidden_states.shape[self._chunk_dim] % self._chunk_size != 0:
+                raise ValueError(
+                    f"`hidden_states` dimension to be chunked: {norm_hidden_states.shape[self._chunk_dim]} has to be divisible by chunk size: {self._chunk_size}. Make sure to set an appropriate `chunk_size` when calling `unet.enable_forward_chunking`."
+                )
+
+            num_chunks = norm_hidden_states.shape[self._chunk_dim] // self._chunk_size
+            ff_output = torch.cat(
+                [
+                    self.ff(hid_slice, scale=lora_scale)
+                    for hid_slice in norm_hidden_states.chunk(num_chunks, dim=self._chunk_dim)
+                ],
+                dim=self._chunk_dim,
+            )
+        else:
+            ff_output = self.ff(norm_hidden_states, 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 AdaLayerNormSingle(nn.Module):
+    r"""
+    Norm layer adaptive layer norm single (adaLN-single).
+
+    As proposed in PixArt-Alpha (see: https://arxiv.org/abs/2310.00426; Section 2.3).
+
+    Parameters:
+        embedding_dim (`int`): The size of each embedding vector.
+        use_additional_conditions (`bool`): To use additional conditions for normalization or not.
+    """
+
+    def __init__(self, embedding_dim: int, use_additional_conditions: bool = False):
+        super().__init__()
+
+        self.emb = CombinedTimestepSizeEmbeddings(
+            embedding_dim, size_emb_dim=embedding_dim // 3, use_additional_conditions=use_additional_conditions
+        )
+
+        self.silu = nn.SiLU()
+        self.linear = nn.Linear(embedding_dim, 6 * embedding_dim, bias=True)
+
+    def forward(
+        self,
+        timestep: torch.Tensor,
+        added_cond_kwargs: Dict[str, torch.Tensor] = None,
+        batch_size: int = None,
+        hidden_dtype: Optional[torch.dtype] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        # No modulation happening here.
+        embedded_timestep = self.emb(timestep, batch_size=batch_size, hidden_dtype=hidden_dtype, resolution=None, aspect_ratio=None)
+        return self.linear(self.silu(embedded_timestep)), embedded_timestep
+
+@dataclass
+class Transformer3DModelOutput(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
+
+
+class LatteT2V(ModelMixin, ConfigMixin):
+    _supports_gradient_checkpointing = True
+
+    """
+    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.
+    """
+
+    @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,
+        video_length: int = 16,
+    ):
+        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.video_length = video_length
+
+        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
+        if self.is_input_continuous:
+            self.in_channels = in_channels
+
+            self.norm = torch.nn.GroupNorm(num_groups=norm_num_groups, num_channels=in_channels, eps=1e-6, affine=True)
+            if use_linear_projection:
+                self.proj_in = linear_cls(in_channels, inner_dim)
+            else:
+                self.proj_in = conv_cls(in_channels, inner_dim, kernel_size=1, stride=1, padding=0)
+        elif self.is_input_vectorized:
+            assert sample_size is not None, "Transformer2DModel over discrete input must provide sample_size"
+            assert num_vector_embeds is not None, "Transformer2DModel over discrete input must provide num_embed"
+
+            self.height = sample_size
+            self.width = sample_size
+            self.num_vector_embeds = num_vector_embeds
+            self.num_latent_pixels = self.height * self.width
+
+            self.latent_image_embedding = ImagePositionalEmbeddings(
+                num_embed=num_vector_embeds, embed_dim=inner_dim, height=self.height, width=self.width
+            )
+        elif self.is_input_patches:
+            assert sample_size is not None, "Transformer2DModel over patched input must provide sample_size"
+
+            self.height = sample_size
+            self.width = sample_size
+
+            self.patch_size = patch_size
+            interpolation_scale = self.config.sample_size // 64  # => 64 (= 512 pixart) has interpolation scale 1
+            interpolation_scale = max(interpolation_scale, 1)
+            self.pos_embed = PatchEmbed(
+                height=sample_size,
+                width=sample_size,
+                patch_size=patch_size,
+                in_channels=in_channels,
+                embed_dim=inner_dim,
+                interpolation_scale=interpolation_scale,
+            )
+
+        # 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)
+            ]
+        )
+
+        # Define temporal transformers blocks
+        self.temporal_transformer_blocks = nn.ModuleList(
+            [
+                BasicTransformerBlock_( # one attention
+                    inner_dim,
+                    num_attention_heads, # num_attention_heads
+                    attention_head_dim, # attention_head_dim 72
+                    dropout=dropout,
+                    cross_attention_dim=None,
+                    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=False,
+                    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
+        if self.is_input_continuous:
+            # 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)
+        elif self.is_input_vectorized:
+            self.norm_out = nn.LayerNorm(inner_dim)
+            self.out = nn.Linear(inner_dim, self.num_vector_embeds - 1)
+        elif self.is_input_patches and norm_type != "ada_norm_single":
+            self.norm_out = nn.LayerNorm(inner_dim, elementwise_affine=False, eps=1e-6)
+            self.proj_out_1 = nn.Linear(inner_dim, 2 * inner_dim)
+            self.proj_out_2 = nn.Linear(inner_dim, patch_size * patch_size * self.out_channels)
+        elif self.is_input_patches and norm_type == "ada_norm_single":
+            self.norm_out = nn.LayerNorm(inner_dim, elementwise_affine=False, eps=1e-6)
+            self.scale_shift_table = nn.Parameter(torch.randn(2, inner_dim) / inner_dim**0.5)
+            self.proj_out = nn.Linear(inner_dim, patch_size * patch_size * self.out_channels)
+
+        # 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 # False, 128 -> 1024
+            # 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
+
+        # define temporal positional embedding
+        temp_pos_embed = self.get_1d_sincos_temp_embed(inner_dim, video_length) # 1152 hidden size
+        self.register_buffer("temp_pos_embed", torch.from_numpy(temp_pos_embed).float().unsqueeze(0), persistent=False)
+
+
+    def _set_gradient_checkpointing(self, module, value=False):
+        self.gradient_checkpointing = value
+
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        timestep: Optional[torch.LongTensor] = None,
+        encoder_hidden_states: Optional[torch.Tensor] = 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,
+        use_image_num: int = 0,
+        enable_temporal_attentions: bool = True,
+        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, frame, 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.
+        """
+        input_batch_size, c, frame, h, w = hidden_states.shape
+        frame = frame - use_image_num
+        hidden_states = rearrange(hidden_states, 'b c f h w -> (b f) c h w').contiguous()
+       
+        # 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: # ndim == 2 means no image joint
+            encoder_attention_mask = (1 - encoder_attention_mask.to(hidden_states.dtype)) * -10000.0
+            encoder_attention_mask = encoder_attention_mask.unsqueeze(1)
+            encoder_attention_mask = repeat(encoder_attention_mask, 'b 1 l -> (b f) 1 l', f=frame).contiguous()
+        elif encoder_attention_mask is not None and encoder_attention_mask.ndim == 3: # ndim == 3 means image joint
+            encoder_attention_mask = (1 - encoder_attention_mask.to(hidden_states.dtype)) * -10000.0
+            encoder_attention_mask_video = encoder_attention_mask[:, :1, ...]
+            encoder_attention_mask_video = repeat(encoder_attention_mask_video, 'b 1 l -> b (1 f) l', f=frame).contiguous()
+            encoder_attention_mask_image = encoder_attention_mask[:, 1:, ...]
+            encoder_attention_mask = torch.cat([encoder_attention_mask_video, encoder_attention_mask_image], dim=1)
+            encoder_attention_mask = rearrange(encoder_attention_mask, 'b n l -> (b n) l').contiguous().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
+        if self.is_input_patches: # here
+            height, width = hidden_states.shape[-2] // self.patch_size, hidden_states.shape[-1] // self.patch_size
+            num_patches = height * width
+
+            hidden_states = self.pos_embed(hidden_states) # alrady add positional embeddings
+
+            if self.adaln_single is not None:
+                if self.use_additional_conditions and added_cond_kwargs is None:
+                    raise ValueError(
+                        "`added_cond_kwargs` cannot be None when using additional conditions for `adaln_single`."
+                    )
+                # batch_size = hidden_states.shape[0]
+                batch_size = input_batch_size
+                timestep, embedded_timestep = self.adaln_single(
+                    timestep, added_cond_kwargs, batch_size=batch_size, hidden_dtype=hidden_states.dtype
+                )
+
+        # 2. Blocks
+        if self.caption_projection is not None:
+            batch_size = hidden_states.shape[0]
+            encoder_hidden_states = self.caption_projection(encoder_hidden_states) # 3 120 1152
+
+            if use_image_num != 0 and self.training:
+                encoder_hidden_states_video = encoder_hidden_states[:, :1, ...]
+                encoder_hidden_states_video = repeat(encoder_hidden_states_video, 'b 1 t d -> b (1 f) t d', f=frame).contiguous()
+                encoder_hidden_states_image = encoder_hidden_states[:, 1:, ...]
+                encoder_hidden_states = torch.cat([encoder_hidden_states_video, encoder_hidden_states_image], dim=1)
+                encoder_hidden_states_spatial = rearrange(encoder_hidden_states, 'b f t d -> (b f) t d').contiguous()
+            else:
+                encoder_hidden_states_spatial = repeat(encoder_hidden_states, 'b t d -> (b f) t d', f=frame).contiguous()
+
+        # prepare timesteps for spatial and temporal block
+        timestep_spatial = repeat(timestep, 'b d -> (b f) d', f=frame + use_image_num).contiguous()
+        timestep_temp = repeat(timestep, 'b d -> (b p) d', p=num_patches).contiguous()
+
+        for i, (spatial_block, temp_block) in enumerate(zip(self.transformer_blocks, self.temporal_transformer_blocks)):
+
+            if self.training and self.gradient_checkpointing:
+                hidden_states = torch.utils.checkpoint.checkpoint(
+                    spatial_block,
+                    hidden_states,
+                    attention_mask,
+                    encoder_hidden_states_spatial,
+                    encoder_attention_mask,
+                    timestep_spatial,
+                    cross_attention_kwargs,
+                    class_labels,
+                    use_reentrant=False,
+                )
+
+                if enable_temporal_attentions:
+                    hidden_states = rearrange(hidden_states, '(b f) t d -> (b t) f d', b=input_batch_size).contiguous()
+
+                    if use_image_num != 0: # image-video joitn training
+                        hidden_states_video = hidden_states[:, :frame, ...]
+                        hidden_states_image = hidden_states[:, frame:, ...]
+
+                        if i == 0:
+                            hidden_states_video = hidden_states_video + self.temp_pos_embed
+
+                        hidden_states_video = torch.utils.checkpoint.checkpoint(
+                            temp_block,
+                            hidden_states_video,
+                            None, # attention_mask
+                            None, # encoder_hidden_states
+                            None, # encoder_attention_mask
+                            timestep_temp,
+                            cross_attention_kwargs,
+                            class_labels,
+                            use_reentrant=False,
+                        )
+
+                        hidden_states = torch.cat([hidden_states_video, hidden_states_image], dim=1)
+                        hidden_states = rearrange(hidden_states, '(b t) f d -> (b f) t d', b=input_batch_size).contiguous()
+
+                    else:
+                        if i == 0:
+                            hidden_states = hidden_states + self.temp_pos_embed
+                        
+                        hidden_states = torch.utils.checkpoint.checkpoint(
+                            temp_block,
+                            hidden_states,
+                            None, # attention_mask
+                            None, # encoder_hidden_states
+                            None, # encoder_attention_mask
+                            timestep_temp,
+                            cross_attention_kwargs,
+                            class_labels,
+                            use_reentrant=False,
+                        )
+
+                        hidden_states = rearrange(hidden_states, '(b t) f d -> (b f) t d', b=input_batch_size).contiguous()
+            else:
+                hidden_states = spatial_block(
+                    hidden_states,
+                    attention_mask,
+                    encoder_hidden_states_spatial,
+                    encoder_attention_mask,
+                    timestep_spatial,
+                    cross_attention_kwargs,
+                    class_labels,
+                )
+
+                if enable_temporal_attentions:
+
+                    hidden_states = rearrange(hidden_states, '(b f) t d -> (b t) f d', b=input_batch_size).contiguous()
+
+                    if use_image_num != 0 and self.training:
+                        hidden_states_video = hidden_states[:, :frame, ...]
+                        hidden_states_image = hidden_states[:, frame:, ...]
+
+                        hidden_states_video = temp_block(
+                            hidden_states_video,
+                            None, # attention_mask
+                            None, # encoder_hidden_states
+                            None, # encoder_attention_mask
+                            timestep_temp,
+                            cross_attention_kwargs,
+                            class_labels,
+                        )
+
+                        hidden_states = torch.cat([hidden_states_video, hidden_states_image], dim=1)
+                        hidden_states = rearrange(hidden_states, '(b t) f d -> (b f) t d', b=input_batch_size).contiguous()
+
+                    else:
+                        if i == 0:
+                            hidden_states = hidden_states + self.temp_pos_embed
+                        
+                        hidden_states = temp_block(
+                            hidden_states,
+                            None, # attention_mask
+                            None, # encoder_hidden_states
+                            None, # encoder_attention_mask
+                            timestep_temp,
+                            cross_attention_kwargs,
+                            class_labels,
+                        )
+
+                        hidden_states = rearrange(hidden_states, '(b t) f d -> (b f) t d', b=input_batch_size).contiguous()
+
+
+        if self.is_input_patches:
+            if self.config.norm_type != "ada_norm_single":
+                conditioning = self.transformer_blocks[0].norm1.emb(
+                    timestep, class_labels, hidden_dtype=hidden_states.dtype
+                )
+                shift, scale = self.proj_out_1(F.silu(conditioning)).chunk(2, dim=1)
+                hidden_states = self.norm_out(hidden_states) * (1 + scale[:, None]) + shift[:, None]
+                hidden_states = self.proj_out_2(hidden_states)
+            elif self.config.norm_type == "ada_norm_single":
+                embedded_timestep = repeat(embedded_timestep, 'b d -> (b f) d', f=frame + use_image_num).contiguous()
+                shift, scale = (self.scale_shift_table[None] + embedded_timestep[:, None]).chunk(2, dim=1)
+                hidden_states = self.norm_out(hidden_states)
+                # Modulation
+                hidden_states = hidden_states * (1 + scale) + shift
+                hidden_states = self.proj_out(hidden_states)
+
+            # unpatchify
+            if self.adaln_single is None:
+                height = width = int(hidden_states.shape[1] ** 0.5)
+            hidden_states = hidden_states.reshape(
+                shape=(-1, height, width, self.patch_size, self.patch_size, self.out_channels)
+            )
+            hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states)
+            output = hidden_states.reshape(
+                shape=(-1, self.out_channels, height * self.patch_size, width * self.patch_size)
+            )
+            output = rearrange(output, '(b f) c h w -> b c f h w', b=input_batch_size).contiguous()
+
+        if not return_dict:
+            return (output,)
+
+        return Transformer3DModelOutput(sample=output)
+    
+    def get_1d_sincos_temp_embed(self, embed_dim, length):
+        pos = torch.arange(0, length).unsqueeze(1)
+        return get_1d_sincos_pos_embed_from_grid(embed_dim, pos)
+    
+    @classmethod
+    def from_pretrained_2d(cls, pretrained_model_path, subfolder=None, **kwargs):
+        if subfolder is not None:
+            pretrained_model_path = os.path.join(pretrained_model_path, subfolder)
+
+
+        config_file = os.path.join(pretrained_model_path, 'config.json')
+        if not os.path.isfile(config_file):
+            raise RuntimeError(f"{config_file} does not exist")
+        with open(config_file, "r") as f:
+            config = json.load(f)
+        
+        model = cls.from_config(config, **kwargs)
+        
+        # model_files = [
+        #     os.path.join(pretrained_model_path, 'diffusion_pytorch_model.bin'),
+        #     os.path.join(pretrained_model_path, 'diffusion_pytorch_model.safetensors')
+        # ]
+
+        # model_file = None
+
+        # for fp in model_files:
+        #     if os.path.exists(fp):
+        #         model_file = fp
+
+        # if not model_file:
+        #     raise RuntimeError(f"{model_file} does not exist")
+
+        # if model_file.split(".")[-1] == "safetensors":
+        #     from safetensors import safe_open
+        #     state_dict = {}
+        #     with safe_open(model_file, framework="pt", device="cpu") as f:
+        #         for key in f.keys():
+        #             state_dict[key] = f.get_tensor(key)
+        # else:
+        #     state_dict = torch.load(model_file, map_location="cpu")
+        
+        # for k, v in model.state_dict().items():
+        #     if 'temporal_transformer_blocks' in k:
+        #         state_dict.update({k: v})
+
+        # model.load_state_dict(state_dict)
+
+        return model

models/unet/motion_embeddings.py

@@ -0,0 +1,283 @@
+import re
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+class MotionEmbedding(nn.Module):
+
+    def __init__(self, embed_dim: int = None, max_seq_length: int = 32, wh: int = 1):
+        super().__init__()
+        self.embed = nn.Parameter(torch.zeros(wh, max_seq_length, embed_dim))
+        print('register spatial motion embedding with', wh)
+
+        self.scale = 1.0
+        self.trained_length = -1
+
+    def set_scale(self, scale: float):
+        self.scale = scale
+
+    def set_lengths(self, trained_length: int):
+        if trained_length > self.embed.shape[1] or trained_length <= 0:
+            raise ValueError("Trained length is out of bounds")
+        self.trained_length = trained_length
+
+    def forward(self, x):
+        _, seq_length, _ = x.shape  # seq_length here is the target sequence length for x
+        # print('seq_length',seq_length)
+        # Assuming self.embed is [batch, frames, dim]
+        embeddings = self.embed[:, :seq_length]  # Initial slice, may not be necessary depending on the interpolation logic
+
+        # Check if interpolation is needed
+        if self.trained_length != -1 and seq_length != self.trained_length:
+            # Interpolate embeddings to match x's sequence length
+            # Ensure embeddings is [batch, dim, frames] for 1D interpolation across frames
+            embeddings = embeddings.permute(0, 2, 1)  # Now [batch, dim, frames]
+            embeddings = F.interpolate(embeddings, size=(seq_length,), mode='linear', align_corners=False)
+            embeddings = embeddings.permute(0, 2, 1)  # Revert to [batch, frames, dim]
+
+        # Ensure the interpolated embeddings match the sequence length of x
+        if embeddings.shape[1] != seq_length:
+            raise ValueError(f"Interpolated embeddings sequence length {embeddings.shape[1]} does not match x's sequence length {seq_length}")
+
+        if x.shape[0] != embeddings.shape[0]:
+            x = x + embeddings.repeat(x.shape[0]//embeddings.shape[0],1,1) * self.scale
+        else:
+            # Now embeddings should have the shape [batch, seq_length, dim] matching x
+            x = x + embeddings * self.scale  # Assuming broadcasting is desired over the batch and dim dimensions
+
+        return x
+
+
+    def forward_average(self, x):
+        _, seq_length, _ = x.shape  # seq_length here is the target sequence length for x
+        # print('seq_length',seq_length)
+        # Assuming self.embed is [batch, frames, dim]
+        embeddings = self.embed[:, :seq_length]  # Initial slice, may not be necessary depending on the interpolation logic
+
+        # Check if interpolation is needed
+        if self.trained_length != -1 and seq_length != self.trained_length:
+            # Interpolate embeddings to match x's sequence length
+            # Ensure embeddings is [batch, dim, frames] for 1D interpolation across frames
+            embeddings = embeddings.permute(0, 2, 1)  # Now [batch, dim, frames]
+            embeddings = F.interpolate(embeddings, size=(seq_length,), mode='linear', align_corners=False)
+            embeddings = embeddings.permute(0, 2, 1)  # Revert to [batch, frames, dim]
+
+        # Ensure the interpolated embeddings match the sequence length of x
+        if embeddings.shape[1] != seq_length:
+            raise ValueError(f"Interpolated embeddings sequence length {embeddings.shape[1]} does not match x's sequence length {seq_length}")
+
+        embeddings_mean = embeddings.mean(dim=1, keepdim=True)
+        embeddings = embeddings - embeddings_mean
+        if x.shape[0] != embeddings.shape[0]:
+            x = x + embeddings.repeat(x.shape[0]//embeddings.shape[0],1,1) * self.scale
+        else:
+            # Now embeddings should have the shape [batch, seq_length, dim] matching x
+            x = x + embeddings * self.scale  # Assuming broadcasting is desired over the batch and dim dimensions
+
+        return x
+
+    def forward_frameSubtraction(self, x):
+        _, seq_length, _ = x.shape  # seq_length here is the target sequence length for x
+        # print('seq_length',seq_length)
+        # Assuming self.embed is [batch, frames, dim]
+        embeddings = self.embed[:, :seq_length]  # Initial slice, may not be necessary depending on the interpolation logic
+
+        # Check if interpolation is needed
+        if self.trained_length != -1 and seq_length != self.trained_length:
+            # Interpolate embeddings to match x's sequence length
+            # Ensure embeddings is [batch, dim, frames] for 1D interpolation across frames
+            embeddings = embeddings.permute(0, 2, 1)  # Now [batch, dim, frames]
+            embeddings = F.interpolate(embeddings, size=(seq_length,), mode='linear', align_corners=False)
+            embeddings = embeddings.permute(0, 2, 1)  # Revert to [batch, frames, dim]
+
+        # Ensure the interpolated embeddings match the sequence length of x
+        if embeddings.shape[1] != seq_length:
+            raise ValueError(f"Interpolated embeddings sequence length {embeddings.shape[1]} does not match x's sequence length {seq_length}")
+
+        embeddings_subtraction = embeddings[:,1:] - embeddings[:,:-1]
+        
+        embeddings = embeddings.clone().detach()
+        embeddings[:,1:] = embeddings_subtraction
+        
+        # first frame minus mean
+        # embeddings[:,0:1] = embeddings[:,0:1] - embeddings.mean(dim=1, keepdim=True)
+
+        if x.shape[0] != embeddings.shape[0]:
+            x = x + embeddings.repeat(x.shape[0]//embeddings.shape[0],1,1) * self.scale
+        else:
+            # Now embeddings should have the shape [batch, seq_length, dim] matching x
+            x = x + embeddings * self.scale  # Assuming broadcasting is desired over the batch and dim dimensions
+
+        return x
+
+class MotionEmbeddingSpatial(nn.Module):
+
+    def __init__(self, h: int = None, w: int = None, embed_dim: int = None, max_seq_length: int = 32):
+        super().__init__()
+        self.embed = nn.Parameter(torch.zeros(h*w, max_seq_length, embed_dim))
+        self.scale = 1.0
+        self.trained_length = -1
+
+    def set_scale(self, scale: float):
+        self.scale = scale
+
+    def set_lengths(self, trained_length: int):
+        if trained_length > self.embed.shape[1] or trained_length <= 0:
+            raise ValueError("Trained length is out of bounds")
+        self.trained_length = trained_length
+
+    def forward(self, x):
+        _, seq_length, _ = x.shape  # seq_length here is the target sequence length for x
+
+        # Assuming self.embed is [batch, frames, dim]
+        embeddings = self.embed[:, :seq_length]  # Initial slice, may not be necessary depending on the interpolation logic
+
+        # Check if interpolation is needed
+        if self.trained_length != -1 and seq_length != self.trained_length:
+            # Interpolate embeddings to match x's sequence length
+            # Ensure embeddings is [batch, dim, frames] for 1D interpolation across frames
+            embeddings = embeddings.permute(0, 2, 1)  # Now [batch, dim, frames]
+            embeddings = F.interpolate(embeddings, size=(seq_length,), mode='linear', align_corners=False)
+            embeddings = embeddings.permute(0, 2, 1)  # Revert to [batch, frames, dim]
+
+        # Ensure the interpolated embeddings match the sequence length of x
+        if embeddings.shape[1] != seq_length:
+            raise ValueError(f"Interpolated embeddings sequence length {embeddings.shape[1]} does not match x's sequence length {seq_length}")
+
+        if x.shape[0] != embeddings.shape[0]:
+            x = x + embeddings.repeat(x.shape[0]//embeddings.shape[0],1,1) * self.scale
+        else:
+            # Now embeddings should have the shape [batch, seq_length, dim] matching x
+            x = x + embeddings * self.scale  # Assuming broadcasting is desired over the batch and dim dimensions
+
+        return x
+
+
+def inject_motion_embeddings(model, combinations=None, config=None):
+    spatial_shape=np.array([config.dataset.height,config.dataset.width])
+    shape32 = np.ceil(spatial_shape/32).astype(int)
+    shape16 = np.ceil(spatial_shape/16).astype(int)
+    spatial_name = 'vSpatial'
+    replacement_dict = {}
+    # support for 32 frames
+    max_seq_length = 32
+    inject_layers = []
+    for name, module in model.named_modules():
+        
+        # check if the module is temp_attention
+        PETemporal = '.temp_attentions.' in name
+
+        if not(PETemporal and re.search(r'transformer_blocks\.\d+$', name)):
+            continue
+
+        if not ([name.split('_')[0], module.norm1.normalized_shape[0]] in combinations):
+            continue
+        
+        replacement_dict[f'{name}.pos_embed'] = MotionEmbedding(max_seq_length=max_seq_length, embed_dim=module.norm1.normalized_shape[0]).to(dtype=model.dtype, device=model.device)
+         
+    replacement_keys = list(set(replacement_dict.keys()))
+    temp_attn_list =    [name.replace('pos_embed','attn1') for name in replacement_keys] + \
+                        [name.replace('pos_embed','attn2') for name in replacement_keys]
+    embed_dims = [replacement_dict[replacement_keys[i]].embed.shape[2] for i in range(len(replacement_keys))]
+    
+    for temp_attn_index,temp_attn in enumerate(temp_attn_list):
+        place_in_net = temp_attn.split('_')[0]
+        pattern = r'(\d+)\.temp_attentions'
+        match = re.search(pattern, temp_attn)
+        place_in_net = temp_attn.split('_')[0]
+        index_in_net = match.group(1)
+        h,w = None,None
+        if place_in_net == 'up':
+            if index_in_net == "1":
+                h, w = shape32
+            elif index_in_net == "2":
+                h, w = shape16
+        elif place_in_net == 'down':
+            if index_in_net == "1":
+                h, w = shape16
+            elif index_in_net == "2":
+                h, w = shape32
+        
+        replacement_dict[temp_attn+'.'+spatial_name] = \
+            MotionEmbedding(
+                wh=h*w,
+                embed_dim=embed_dims[temp_attn_index%len(replacement_keys)]
+                ).to(dtype=model.dtype, device=model.device)
+
+    for name, new_module in replacement_dict.items():
+        parent_name = name.rsplit('.', 1)[0] if '.' in name else ''
+        module_name = name.rsplit('.', 1)[-1]
+        parent_module = model
+        if parent_name:
+            parent_module = dict(model.named_modules())[parent_name]
+
+        if [parent_name.split('_')[0], new_module.embed.shape[-1]] in combinations:
+            inject_layers.append(name)
+            setattr(parent_module, module_name, new_module)
+
+    inject_layers = list(set(inject_layers))
+    for name in inject_layers:
+        print(f"Injecting motion embedding at {name}")
+
+    parameters_list = []
+    for name, para in model.named_parameters():
+        if 'pos_embed' in name or spatial_name in name:
+            parameters_list.append(para)
+            para.requires_grad = True
+        else:
+            para.requires_grad = False
+
+    return parameters_list, inject_layers
+
+def save_motion_embeddings(model, file_path):
+    # Extract motion embedding from all instances of MotionEmbedding
+    motion_embeddings = {
+        name: module.embed
+        for name, module in model.named_modules()
+        if isinstance(module, MotionEmbedding) or isinstance(module, MotionEmbeddingSpatial)
+    }
+    # Save the motion embeddings to the specified file path
+    torch.save(motion_embeddings, file_path)
+
+def load_motion_embeddings(model, saved_embeddings):
+    for key, embedding in saved_embeddings.items():
+        # Extract parent module and module name from the key
+        parent_name = key.rsplit('.', 1)[0] if '.' in key else ''
+        module_name = key.rsplit('.', 1)[-1]
+
+        # Retrieve the parent module
+        parent_module = model
+        if parent_name:
+            parent_module = dict(model.named_modules())[parent_name]
+
+        # Create a new MotionEmbedding instance with the correct dimensions
+
+        new_module = MotionEmbedding(wh = embedding.shape[0],embed_dim=embedding.shape[-1], max_seq_length=embedding.shape[-2])
+
+        # Properly assign the loaded embeddings to the 'embed' parameter wrapped in nn.Parameter
+        # Ensure the embedding is on the correct device and has the correct dtype
+        new_module.embed = nn.Parameter(embedding.to(dtype=model.dtype, device=model.device))
+
+        # Replace the corresponding module in the model with the new MotionEmbedding instance
+        setattr(parent_module, module_name, new_module)
+
+def set_motion_embedding_scale(model, scale_value):
+    # Iterate over all modules in the model
+    for _, module in model.named_modules():
+        # Check if the module is an instance of MotionEmbedding
+        if isinstance(module, MotionEmbedding):
+            # Set the scale attribute to the specified value
+            module.scale = scale_value
+
+def set_motion_embedding_length(model, trained_length):
+    # Iterate over all modules in the model
+    for _, module in model.named_modules():
+        # Check if the module is an instance of MotionEmbedding
+        if isinstance(module, MotionEmbedding):
+            # Set the length to the specified value
+            module.trained_length = trained_length
+
+
+
+
+

models/unet/unet_3d_blocks.py

@@ -0,0 +1,842 @@
+# Copyright 2023 The HuggingFace Team. All rights reserved.
+#
+# 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.
+
+import torch
+import torch.utils.checkpoint as checkpoint
+from torch import nn
+from diffusers.models.resnet import Downsample2D, ResnetBlock2D, TemporalConvLayer, Upsample2D
+from diffusers.models.transformer_2d import Transformer2DModel
+from diffusers.models.transformer_temporal import TransformerTemporalModel
+
+# Assign gradient checkpoint function to simple variable for readability.
+g_c = checkpoint.checkpoint
+
+def use_temporal(module, num_frames, x):
+    if num_frames == 1:
+        if isinstance(module, TransformerTemporalModel):
+            return {"sample": x}
+        else:
+            return x
+
+def custom_checkpoint(module, mode=None):
+    if mode == None: raise ValueError('Mode for gradient checkpointing cannot be none.')
+    custom_forward = None
+
+    if mode == 'resnet':
+        def custom_forward(hidden_states, temb):
+            inputs = module(hidden_states, temb)
+            return inputs
+
+    if mode == 'attn':
+        def custom_forward(
+            hidden_states, 
+            encoder_hidden_states=None, 
+            cross_attention_kwargs=None
+        ):
+            inputs = module(
+                hidden_states,
+                encoder_hidden_states,
+                cross_attention_kwargs
+            )
+            return inputs
+
+    if mode == 'temp':
+         def custom_forward(hidden_states, num_frames=None):
+            inputs = use_temporal(module, num_frames, hidden_states)
+            if inputs is None: inputs = module(
+                hidden_states, 
+                num_frames=num_frames
+            )
+            return inputs
+
+    return custom_forward
+
+def transformer_g_c(transformer, sample, num_frames):
+    sample = g_c(custom_checkpoint(transformer, mode='temp'), 
+        sample, num_frames, use_reentrant=False
+    )['sample']
+
+    return sample
+
+def cross_attn_g_c(
+        attn, 
+        temp_attn, 
+        resnet, 
+        temp_conv, 
+        hidden_states, 
+        encoder_hidden_states, 
+        cross_attention_kwargs, 
+        temb, 
+        num_frames,
+        inverse_temp=False
+    ):
+    
+    def ordered_g_c(idx):
+
+        # Self and CrossAttention
+        if idx == 0: return g_c(custom_checkpoint(attn, mode='attn'),
+            hidden_states, encoder_hidden_states,cross_attention_kwargs, use_reentrant=False
+        )['sample']
+
+        # Temporal Self and CrossAttention
+        if idx == 1: return g_c(custom_checkpoint(temp_attn, mode='temp'), 
+            hidden_states, num_frames, use_reentrant=False)['sample']
+
+        # Resnets
+        if idx == 2: return g_c(custom_checkpoint(resnet, mode='resnet'), 
+            hidden_states, temb, use_reentrant=False)
+        
+        # Temporal Convolutions
+        if idx == 3: return g_c(custom_checkpoint(temp_conv, mode='temp'), 
+            hidden_states, num_frames, use_reentrant=False
+    )
+
+    # Here we call the function depending on the order in which they are called. 
+    # For some layers, the orders are different, so we access the appropriate one by index.
+    
+    if not inverse_temp:
+        for idx in [0,1,2,3]: hidden_states = ordered_g_c(idx) 
+    else:
+        for idx in [2,3,0,1]: hidden_states = ordered_g_c(idx)
+
+    return hidden_states
+
+def up_down_g_c(resnet, temp_conv, hidden_states, temb, num_frames):
+    hidden_states = g_c(custom_checkpoint(resnet, mode='resnet'), hidden_states, temb, use_reentrant=False)
+    hidden_states = g_c(custom_checkpoint(temp_conv, mode='temp'), 
+        hidden_states, num_frames,  use_reentrant=False
+    )
+    return hidden_states
+
+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=True,
+    only_cross_attention=False,
+    upcast_attention=False,
+    resnet_time_scale_shift="default",
+):
+    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,
+        )
+    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,
+        )
+    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=True,
+    only_cross_attention=False,
+    upcast_attention=False,
+    resnet_time_scale_shift="default",
+):
+    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,
+        )
+    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,
+        )
+    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=True,
+        upcast_attention=False,
+    ):
+        super().__init__()
+
+        self.gradient_checkpointing = False
+        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 = [
+            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,
+            )
+        ]
+        temp_convs = [
+            TemporalConvLayer(
+                in_channels,
+                in_channels,
+                dropout=0.1
+            )
+        ]
+        attentions = []
+        temp_attentions = []
+
+        for _ in range(num_layers):
+            attentions.append(
+                Transformer2DModel(
+                    in_channels // attn_num_head_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,
+                )
+            )
+            temp_attentions.append(
+                TransformerTemporalModel(
+                    in_channels // attn_num_head_channels,
+                    attn_num_head_channels,
+                    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,
+                )
+            )
+            temp_convs.append(
+                TemporalConvLayer(
+                    in_channels,
+                    in_channels,
+                    dropout=0.1
+                )
+            )
+
+        self.resnets = nn.ModuleList(resnets)
+        self.temp_convs = nn.ModuleList(temp_convs)
+        self.attentions = nn.ModuleList(attentions)
+        self.temp_attentions = nn.ModuleList(temp_attentions)
+
+    def forward(
+        self,
+        hidden_states,
+        temb=None,
+        encoder_hidden_states=None,
+        attention_mask=None,
+        num_frames=1,
+        cross_attention_kwargs=None,
+    ):
+        if self.gradient_checkpointing:
+            hidden_states = up_down_g_c(
+                    self.resnets[0], 
+                    self.temp_convs[0], 
+                    hidden_states, 
+                    temb, 
+                    num_frames
+                )
+        else:
+            hidden_states = self.resnets[0](hidden_states, temb)
+            hidden_states = self.temp_convs[0](hidden_states, num_frames=num_frames)
+            
+        for attn, temp_attn, resnet, temp_conv in zip(
+            self.attentions, self.temp_attentions, self.resnets[1:], self.temp_convs[1:]
+        ):
+            if self.gradient_checkpointing:
+                hidden_states = cross_attn_g_c(
+                        attn, 
+                        temp_attn, 
+                        resnet, 
+                        temp_conv, 
+                        hidden_states, 
+                        encoder_hidden_states, 
+                        cross_attention_kwargs, 
+                        temb, 
+                        num_frames
+                    )
+            else:
+                hidden_states = attn(
+                    hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    cross_attention_kwargs=cross_attention_kwargs,
+                ).sample
+                
+                if num_frames > 1:
+                    hidden_states = temp_attn(hidden_states, num_frames=num_frames).sample
+
+                hidden_states = resnet(hidden_states, temb)
+
+                if num_frames > 1:
+                    hidden_states = temp_conv(hidden_states, num_frames=num_frames)
+
+        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,
+    ):
+        super().__init__()
+        resnets = []
+        attentions = []
+        temp_attentions = []
+        temp_convs = []
+
+        self.gradient_checkpointing = False
+        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(
+                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,
+                )
+            )
+            temp_convs.append(
+                TemporalConvLayer(
+                    out_channels,
+                    out_channels,
+                    dropout=0.1
+                )
+            )
+            attentions.append(
+                Transformer2DModel(
+                    out_channels // attn_num_head_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,
+                )
+            )
+            temp_attentions.append(
+                TransformerTemporalModel(
+                    out_channels // attn_num_head_channels,
+                    attn_num_head_channels,
+                    in_channels=out_channels,
+                    num_layers=1,
+                    cross_attention_dim=cross_attention_dim,
+                    norm_num_groups=resnet_groups,
+                )
+            )
+        self.resnets = nn.ModuleList(resnets)
+        self.temp_convs = nn.ModuleList(temp_convs)
+        self.attentions = nn.ModuleList(attentions)
+        self.temp_attentions = nn.ModuleList(temp_attentions)
+
+        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
+
+    def forward(
+        self,
+        hidden_states,
+        temb=None,
+        encoder_hidden_states=None,
+        attention_mask=None,
+        num_frames=1,
+        cross_attention_kwargs=None,
+    ):
+        # TODO(Patrick, William) - attention mask is not used
+        output_states = ()
+
+        for resnet, temp_conv, attn, temp_attn in zip(
+            self.resnets, self.temp_convs, self.attentions, self.temp_attentions
+        ):
+        
+            if self.gradient_checkpointing:
+                hidden_states = cross_attn_g_c(
+                        attn, 
+                        temp_attn, 
+                        resnet, 
+                        temp_conv, 
+                        hidden_states, 
+                        encoder_hidden_states, 
+                        cross_attention_kwargs, 
+                        temb, 
+                        num_frames,
+                        inverse_temp=True
+                    )
+            else:
+                hidden_states = resnet(hidden_states, temb)
+
+                if num_frames > 1:
+                    hidden_states = temp_conv(hidden_states, num_frames=num_frames)
+
+                hidden_states = attn(
+                    hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    cross_attention_kwargs=cross_attention_kwargs,
+                ).sample
+
+                if num_frames > 1:
+                    hidden_states = temp_attn(hidden_states, num_frames=num_frames).sample
+
+            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,
+    ):
+        super().__init__()
+        resnets = []
+        temp_convs = []
+
+        self.gradient_checkpointing = False
+        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,
+                )
+            )
+            temp_convs.append(
+                TemporalConvLayer(
+                    out_channels,
+                    out_channels,
+                    dropout=0.1
+                )
+            )
+
+        self.resnets = nn.ModuleList(resnets)
+        self.temp_convs = nn.ModuleList(temp_convs)
+
+        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
+
+    def forward(self, hidden_states, temb=None, num_frames=1):
+        output_states = ()
+
+        for resnet, temp_conv in zip(self.resnets, self.temp_convs):
+            if self.gradient_checkpointing:
+                hidden_states = up_down_g_c(resnet, temp_conv, hidden_states, temb, num_frames)
+            else:
+                hidden_states = resnet(hidden_states, temb)
+
+                if num_frames > 1:
+                    hidden_states = temp_conv(hidden_states, num_frames=num_frames)
+
+            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,
+    ):
+        super().__init__()
+        resnets = []
+        temp_convs = []
+        attentions = []
+        temp_attentions = []
+
+        self.gradient_checkpointing = False
+        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(
+                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,
+                )
+            )
+            temp_convs.append(
+                TemporalConvLayer(
+                    out_channels,
+                    out_channels,
+                    dropout=0.1
+                )
+            )
+            attentions.append(
+                Transformer2DModel(
+                    out_channels // attn_num_head_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,
+                )
+            )
+            temp_attentions.append(
+                TransformerTemporalModel(
+                    out_channels // attn_num_head_channels,
+                    attn_num_head_channels,
+                    in_channels=out_channels,
+                    num_layers=1,
+                    cross_attention_dim=cross_attention_dim,
+                    norm_num_groups=resnet_groups,
+                )
+            )
+        self.resnets = nn.ModuleList(resnets)
+        self.temp_convs = nn.ModuleList(temp_convs)
+        self.attentions = nn.ModuleList(attentions)
+        self.temp_attentions = nn.ModuleList(temp_attentions)
+
+        if add_upsample:
+            self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)])
+        else:
+            self.upsamplers = None
+
+    def forward(
+        self,
+        hidden_states,
+        res_hidden_states_tuple,
+        temb=None,
+        encoder_hidden_states=None,
+        upsample_size=None,
+        attention_mask=None,
+        num_frames=1,
+        cross_attention_kwargs=None,
+    ):
+        # TODO(Patrick, William) - attention mask is not used
+        for resnet, temp_conv, attn, temp_attn in zip(
+            self.resnets, self.temp_convs, self.attentions, self.temp_attentions
+        ):
+            # pop res hidden states
+            res_hidden_states = res_hidden_states_tuple[-1]
+            res_hidden_states_tuple = res_hidden_states_tuple[:-1]
+            hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)
+
+            if self.gradient_checkpointing:
+                hidden_states = cross_attn_g_c(
+                        attn, 
+                        temp_attn, 
+                        resnet, 
+                        temp_conv, 
+                        hidden_states, 
+                        encoder_hidden_states, 
+                        cross_attention_kwargs, 
+                        temb, 
+                        num_frames,
+                        inverse_temp=True
+                    )
+            else:
+                hidden_states = resnet(hidden_states, temb)
+
+                if num_frames > 1:
+                    hidden_states = temp_conv(hidden_states, num_frames=num_frames)
+
+                hidden_states = attn(
+                    hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    cross_attention_kwargs=cross_attention_kwargs,
+                ).sample
+
+                if num_frames > 1:
+                    hidden_states = temp_attn(hidden_states, num_frames=num_frames).sample
+
+        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,
+    ):
+        super().__init__()
+        resnets = []
+        temp_convs = []
+        self.gradient_checkpointing = 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(
+                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,
+                )
+            )
+            temp_convs.append(
+                TemporalConvLayer(
+                    out_channels,
+                    out_channels,
+                    dropout=0.1
+                )
+            )
+
+        self.resnets = nn.ModuleList(resnets)
+        self.temp_convs = nn.ModuleList(temp_convs)
+
+        if add_upsample:
+            self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)])
+        else:
+            self.upsamplers = None
+
+    def forward(self, hidden_states, res_hidden_states_tuple, temb=None, upsample_size=None, num_frames=1):
+        for resnet, temp_conv in zip(self.resnets, self.temp_convs):
+            # 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.gradient_checkpointing:
+                hidden_states = up_down_g_c(resnet, temp_conv, hidden_states, temb, num_frames)
+            else:
+                hidden_states = resnet(hidden_states, temb)
+
+                if num_frames > 1:
+                    hidden_states = temp_conv(hidden_states, num_frames=num_frames)
+
+        if self.upsamplers is not None:
+            for upsampler in self.upsamplers:
+                hidden_states = upsampler(hidden_states, upsample_size)
+
+        return hidden_states

models/unet/unet_3d_condition.py

@@ -0,0 +1,500 @@
+# Copyright 2023 Alibaba DAMO-VILAB and The HuggingFace Team. All rights reserved.
+# Copyright 2023 The ModelScope Team.
+#
+# 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.
+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.utils import BaseOutput, logging
+from diffusers.models.embeddings import TimestepEmbedding, Timesteps
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.transformer_temporal import TransformerTemporalModel
+from .unet_3d_blocks import (
+    CrossAttnDownBlock3D,
+    CrossAttnUpBlock3D,
+    DownBlock3D,
+    UNetMidBlock3DCrossAttn,
+    UpBlock3D,
+    get_down_block,
+    get_up_block,
+    transformer_g_c
+)
+
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+@dataclass
+class UNet3DConditionOutput(BaseOutput):
+    """
+    Args:
+        sample (`torch.FloatTensor` of shape `(batch_size, num_frames, num_channels, height, width)`):
+            Hidden states conditioned on `encoder_hidden_states` input. Output of last layer of model.
+    """
+
+    sample: torch.FloatTensor
+
+
+class UNet3DConditionModel(ModelMixin, ConfigMixin):
+    r"""
+    UNet3DConditionModel is a conditional 2D UNet model that takes in a noisy sample, conditional state, and a timestep
+    and returns sample shaped output.
+
+    This model inherits from [`ModelMixin`]. Check the superclass documentation for the generic methods the library
+    implements for all the models (such as downloading or saving, etc.)
+
+    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): The number of channels in the input sample.
+        out_channels (`int`, *optional*, defaults to 4): The number of channels in the output.
+        down_block_types (`Tuple[str]`, *optional*, defaults to `("CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "DownBlock2D")`):
+            The tuple of downsample blocks to use.
+        up_block_types (`Tuple[str]`, *optional*, defaults to `("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D",)`):
+            The tuple of upsample blocks to use.
+        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.
+        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`, it will skip the normalization and activation layers in post-processing
+        norm_eps (`float`, *optional*, defaults to 1e-5): The epsilon to use for the normalization.
+        cross_attention_dim (`int`, *optional*, defaults to 1280): The dimension of the cross attention features.
+        attention_head_dim (`int`, *optional*, defaults to 8): The dimension of the attention heads.
+    """
+
+    _supports_gradient_checkpointing = True
+
+    @register_to_config
+    def __init__(
+        self,
+        sample_size: Optional[int] = None,
+        in_channels: int = 4,
+        out_channels: int = 4,
+        down_block_types: Tuple[str] = (
+            "CrossAttnDownBlock3D",
+            "CrossAttnDownBlock3D",
+            "CrossAttnDownBlock3D",
+            "DownBlock3D",
+        ),
+        up_block_types: Tuple[str] = ("UpBlock3D", "CrossAttnUpBlock3D", "CrossAttnUpBlock3D", "CrossAttnUpBlock3D"),
+        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: Optional[int] = 32,
+        norm_eps: float = 1e-5,
+        cross_attention_dim: int = 1024,
+        attention_head_dim: Union[int, Tuple[int]] = 64,
+    ):
+        super().__init__()
+
+        self.sample_size = sample_size
+        self.gradient_checkpointing = False
+        # 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(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}."
+            )
+
+        # input
+        conv_in_kernel = 3
+        conv_out_kernel = 3
+        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
+        time_embed_dim = block_out_channels[0] * 4
+        self.time_proj = Timesteps(block_out_channels[0], True, 0)
+        timestep_input_dim = block_out_channels[0]
+
+        self.time_embedding = TimestepEmbedding(
+            timestep_input_dim,
+            time_embed_dim,
+            act_fn=act_fn,
+        )
+
+        self.transformer_in = TransformerTemporalModel(
+            num_attention_heads=8,
+            attention_head_dim=attention_head_dim,
+            in_channels=block_out_channels[0],
+            num_layers=1,
+        )
+
+        # class embedding
+        self.down_blocks = nn.ModuleList([])
+        self.up_blocks = nn.ModuleList([])
+
+        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):
+            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=False,
+            )
+            self.down_blocks.append(down_block)
+
+        # mid
+        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,
+            cross_attention_dim=cross_attention_dim,
+            attn_num_head_channels=attention_head_dim[-1],
+            resnet_groups=norm_num_groups,
+            dual_cross_attention=False,
+        )
+
+        # count how many layers upsample the images
+        self.num_upsamplers = 0
+
+        # up
+        reversed_block_out_channels = list(reversed(block_out_channels))
+        reversed_attention_head_dim = list(reversed(attention_head_dim))
+
+        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=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=False,
+            )
+            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 = nn.SiLU()
+        else:
+            self.conv_norm_out = None
+            self.conv_act = None
+
+        conv_out_padding = (conv_out_kernel - 1) // 2
+        self.conv_out = nn.Conv2d(
+            block_out_channels[0], out_channels, kernel_size=conv_out_kernel, padding=conv_out_padding
+        )
+
+    def 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, value=False):
+        self.gradient_checkpointing = value
+        self.mid_block.gradient_checkpointing = value
+        for module in self.down_blocks + self.up_blocks:
+            if isinstance(module, (CrossAttnDownBlock3D, DownBlock3D, CrossAttnUpBlock3D, UpBlock3D)):
+                module.gradient_checkpointing = value
+
+    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,
+        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, num_frames, 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.UNet3DConditionOutput`] instead of a plain tuple.
+            cross_attention_kwargs (`dict`, *optional*):
+                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
+                `self.processor` in
+                [diffusers.cross_attention](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/cross_attention.py).
+
+        Returns:
+            [`~models.unet_2d_condition.UNet3DConditionOutput`] or `tuple`:
+            [`~models.unet_2d_condition.UNet3DConditionOutput`] 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)
+
+        # 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
+        num_frames = sample.shape[2]
+        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, timestep_cond)
+        emb = emb.repeat_interleave(repeats=num_frames, dim=0)
+        encoder_hidden_states = encoder_hidden_states.repeat_interleave(repeats=num_frames, dim=0)
+
+        # 2. pre-process
+        sample = sample.permute(0, 2, 1, 3, 4).reshape((sample.shape[0] * num_frames, -1) + sample.shape[3:])
+        sample = self.conv_in(sample)
+        
+        if num_frames > 1:
+            if self.gradient_checkpointing:
+                sample = transformer_g_c(self.transformer_in, sample, num_frames)
+            else:
+                sample = self.transformer_in(sample, num_frames=num_frames).sample
+
+        # 3. down
+        down_block_res_samples = (sample,)
+        for downsample_block in self.down_blocks:
+            if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention:
+                sample, res_samples = downsample_block(
+                    hidden_states=sample,
+                    temb=emb,
+                    encoder_hidden_states=encoder_hidden_states,
+                    attention_mask=attention_mask,
+                    num_frames=num_frames,
+                    cross_attention_kwargs=cross_attention_kwargs,
+                )
+            else:
+                sample, res_samples = downsample_block(hidden_states=sample, temb=emb, num_frames=num_frames)
+
+            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
+
+        # 4. mid
+        if self.mid_block is not None:
+            sample = self.mid_block(
+                sample,
+                emb,
+                encoder_hidden_states=encoder_hidden_states,
+                attention_mask=attention_mask,
+                num_frames=num_frames,
+                cross_attention_kwargs=cross_attention_kwargs,
+            )
+
+        if mid_block_additional_residual is not None:
+            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,
+                    upsample_size=upsample_size,
+                    attention_mask=attention_mask,
+                    num_frames=num_frames,
+                    cross_attention_kwargs=cross_attention_kwargs,
+                )
+            else:
+                sample = upsample_block(
+                    hidden_states=sample,
+                    temb=emb,
+                    res_hidden_states_tuple=res_samples,
+                    upsample_size=upsample_size,
+                    num_frames=num_frames,
+                )
+
+        # 6. post-process
+        if self.conv_norm_out:
+            sample = self.conv_norm_out(sample)
+            sample = self.conv_act(sample)
+
+        sample = self.conv_out(sample)
+
+        # reshape to (batch, channel, framerate, width, height)
+        sample = sample[None, :].reshape((-1, num_frames) + sample.shape[1:]).permute(0, 2, 1, 3, 4)
+
+        if not return_dict:
+            return (sample,)
+
+        return UNet3DConditionOutput(sample=sample)

noise_init/__init__.py

@@ -0,0 +1,4 @@
+from .freq_init import FreqInit
+from .blend_init import BlendInit
+from .blend_freq_init import BlendFreqInit
+from .fft_init import FFTInit