Generate synthetic dataset

Delete Open Sora Plan Modeling

amend log validation

Change name to fast video

Remove OSP modeling

clean up name changing

remove files

 Deleted unnecessary files

commit first

commit first

training

ok

Update generate_synthetic.sh and deepspeed_zero2_config.yaml

14

debug

overfitting ....

debugging ..

Debug successful!

Add zero3

OK

update optimizer

load

small bug

random seed args

sp enable & still has bug

rename

update inference sp code / can run / still has bug / don't output normal mp4

switch to deepspeed dummyoptim

fix some bugs; still output green

SP inference done!

fix typos in readme and latent dataset debug file

.gitignore

@@ -7,7 +7,7 @@ UCF-101/
 results/
 vae
 build/
-opensora.egg-info/
+fastvideo.egg-info/
 wandb/
 .idea
 *.ipynb

README.md

@@ -7,37 +7,30 @@ Change the index-url cuda version according to your system.
 conda create -n fastvideo python=3.10.12
 conda activate fastvideo
 pip3 install torch==2.5.0 torchvision  --index-url https://download.pytorch.org/whl/cu121
-pip3 install -U xformers --index-url https://download.pytorch.org/whl/cu121
+pip3 install -U xformers==0.0.28.post2 --index-url https://download.pytorch.org/whl/cu121
 cd .. && git clone  https://github.com/huggingface/diffusers
 cd diffusers && git checkout mochi && pip install -e . && cd ../FastVideo
 ```
 
 ```
 pip install -e . && pip install -e ".[train]"
-apt-get update && apt install screen && pip install watch gpustat
+sudo apt-get update && apt install screen && pip install watch gpustat
 ```
 
 ## Prepare Data & Models
 We've prepared some debug data to facilitate development. To make sure the training pipeline is correct, train on the debug data and make sure the model overfit on it (feed it the same text prompt and see if the output video is the same as the training data)
+
 ```
 python scripts/download_hf.py --repo_id=Stealths-Video/dummyVid --local_dir=data/dummyVid --repo_type=dataset
-python scripts/download_hf.py --repo_id=LanguageBind/Open-Sora-Plan-v1.2.0  --file_name vae/checkpoint.ckpt --local_dir=data/Open-Sora-Plan-v1.2.0 --repo_type=model 
-python scripts/download_hf.py --repo_id=LanguageBind/Open-Sora-Plan-v1.2.0  --file_name vae/config.json --local_dir=data/Open-Sora-Plan-v1.2.0 --repo_type=model 
-python scripts/download_hf.py --repo_id=LanguageBind/Open-Sora-Plan-v1.2.0  --file_name 29x480p/config.json --local_dir=data/Open-Sora-Plan-v1.2.0 --repo_type=model
-python scripts/download_hf.py --repo_id=LanguageBind/Open-Sora-Plan-v1.2.0  --file_name 29x480p/diffusion_pytorch_model.safetensors --local_dir=data/Open-Sora-Plan-v1.2.0 --repo_type=model
 python scripts/download_hf.py --repo_id=Stealths-Video/mochi --local_dir=data/mochi --repo_type=model
+python scripts/download_hf.py --repo_id=Stealths-Video/Mochi-Synthetic-Data --local_dir=data/Mochi-Synthetic-Data --repo_type=dataset
 ```
 
+## How to overfit
 ```
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
-tokenizer = AutoTokenizer.from_pretrained("google/mt5-xxl", cache_dir="data/.cache")
-model = AutoModelForSeq2SeqLM.from_pretrained("google/mt5-xxl", cache_dir="data/.cache")
+bash t2v_debug_multi.sh
 ```
-## Debug Training
-```
-bash t2v_debug_single.sh
-```
-
+Make sure to edit data/Mochi-Synthetic-Data/videos2caption.json such that this is only one video in the dataset (you can copy multiple annotations of the same video). Also make sure to edit the prompt in scripts/overfit.shto match the prompt in the training data. I observe the overfitting  after 50 steps. 
 
 ## TODO
 

fastvideo/dataset/__init__.py

@@ -2,10 +2,9 @@ from transformers import AutoTokenizer
 
 from torchvision import transforms
 
-from opensora.dataset.t2v_datasets import T2V_dataset
-from opensora.dataset.inpaint_datasets import Inpaint_dataset
-from opensora.dataset.transform import Normalize255, TemporalRandomCrop,CenterCropResizeVideo
-from opensora.models.causalvideovae import ae_norm, ae_denorm
+from fastvideo.dataset.t2v_datasets import T2V_dataset
+from fastvideo.dataset.latent_datasets import LatentDataset
+from fastvideo.dataset.transform import Normalize255, TemporalRandomCrop,CenterCropResizeVideo
 
 def getdataset(args):
     temporal_sample = TemporalRandomCrop(args.num_frames)  # 16 x
@@ -28,17 +27,14 @@ def getdataset(args):
     tokenizer = AutoTokenizer.from_pretrained(args.text_encoder_name, cache_dir=args.cache_dir)
     if args.dataset == 't2v':
         return T2V_dataset(args, transform=transform, temporal_sample=temporal_sample, tokenizer=tokenizer, 
-                           transform_topcrop=transform_topcrop, video_length_tolerance_range=args.video_length_tolerance_range)
-    elif args.dataset == 'inpaint' or args.dataset == 'i2v':
-        return Inpaint_dataset(args, transform=transform, temporal_sample=temporal_sample, tokenizer=tokenizer, 
-                            transform_topcrop=transform_topcrop)
+                           transform_topcrop=transform_topcrop)
     
     raise NotImplementedError(args.dataset)
 
 
 if __name__ == "__main__":
     from accelerate import Accelerator
-    from opensora.dataset.t2v_datasets import dataset_prog
+    from fastvideo.dataset.t2v_datasets import dataset_prog
     import random
     from tqdm import tqdm
     args = type('args', (), 
@@ -47,7 +43,7 @@ if __name__ == "__main__":
         'dataset': 't2v', 
         'attention_mode': 'xformers', 
         'use_rope': True, 
-        'model_max_length': 300, 
+        'text_max_length': 300, 
         'max_height': 320,
         'max_width': 240,
         'num_frames': 1,

fastvideo/dataset/latent_datasets.py

@@ -0,0 +1,67 @@
+import torch
+from torch.utils.data import Dataset
+import json
+import os
+
+class LatentDataset(Dataset):
+    def __init__(self, data_merge_path, num_latent_t):
+        # data_merge_path: video_dir, latent_dir, prompt_embed_dir, json_path
+        self.data_merge_path = data_merge_path
+        # read txt
+        with open(data_merge_path, 'r') as f:
+            merge_data = f.readlines()
+        video_dir, latent_dir, prompt_embed_dir, prompt_attention_mask_dir, json_path = merge_data[0].split(",")
+        self.video_dir = video_dir
+        self.latent_dir = latent_dir
+        self.prompt_embed_dir = prompt_embed_dir
+        self.prompt_attention_mask_dir = prompt_attention_mask_dir
+        self.json_path = json_path
+        with open(self.json_path, 'r') as f:
+            self.data_anno = json.load(f)
+            
+        self.num_latent_t = num_latent_t
+    def __getitem__(self, idx):
+        latent_file = self.data_anno[idx]["latent_path"]
+        prompt_embed_file = self.data_anno[idx]["prompt_embed_path"]
+        prompt_attention_mask_file = self.data_anno[idx]["prompt_attention_mask"]
+        # load 
+        latent = torch.load(os.path.join(self.latent_dir, latent_file), map_location="cpu", weights_only=True)[:, -self.num_latent_t:]
+        prompt_embed = torch.load(os.path.join(self.prompt_embed_dir, prompt_embed_file), map_location="cpu", weights_only=True)
+        prompt_attention_mask = torch.load(os.path.join(self.prompt_attention_mask_dir, prompt_attention_mask_file), map_location="cpu", weights_only=True)
+        return latent, prompt_embed, prompt_attention_mask
+    
+    def __len__(self):
+        return len(self.data_anno)
+    
+def latent_collate_function(batch):
+    # return latent, prompt, latent_attn_mask, text_attn_mask
+    # latent_attn_mask: # b t h w
+    # text_attn_mask: b 1 l
+    # needs to check if the latent/prompt' size and apply padding & attn mask
+    latents, prompt_embeds, prompt_attention_masks = zip(*batch)
+    # calculate max shape
+    max_t = max([latent.shape[1] for latent in latents])
+    max_h = max([latent.shape[2] for latent in latents])
+    max_w = max([latent.shape[3] for latent in latents])
+    
+    # padding
+    latents = [torch.nn.functional.pad(latent, (0, max_t - latent.shape[1], 0, max_h - latent.shape[2], 0, max_w - latent.shape[3])) for latent in latents]
+    # attn mask
+    latent_attn_mask = torch.ones(len(latents), max_t, max_h, max_w)
+    # set to 0 if padding
+    for i, latent in enumerate(latents):
+        latent_attn_mask[i, latent.shape[1]:, :, :] = 0
+        latent_attn_mask[i, :, latent.shape[2]:, :] = 0
+        latent_attn_mask[i, :, :, latent.shape[3]:] = 0
+    
+    prompt_embeds = torch.stack(prompt_embeds, dim=0)
+    prompt_attention_masks = torch.stack(prompt_attention_masks, dim=0)
+    latents = torch.stack(latents, dim=0)
+    return latents, prompt_embeds, latent_attn_mask, prompt_attention_masks
+
+if __name__ == "__main__":
+    dataset = LatentDataset("data/Mochi-Synthetic-Data/merge.txt", num_latent_t=28)
+    dataloader = torch.utils.data.DataLoader(dataset, batch_size=2, shuffle=False, collate_fn=latent_collate_function)
+    for latent, prompt_embed, latent_attn_mask, prompt_attention_mask in dataloader:
+        print(latent.shape, prompt_embed.shape, latent_attn_mask.shape, prompt_attention_mask.shape)
+        import pdb; pdb.set_trace()

fastvideo/dataset/t2v_datasets.py

@@ -19,8 +19,8 @@ from tqdm import tqdm
 from PIL import Image
 from accelerate.logging import get_logger
 
-from opensora.utils.dataset_utils import DecordInit
-from opensora.utils.utils import text_preprocessing
+from fastvideo.utils.dataset_utils import DecordInit
+from fastvideo.utils.utils import text_preprocessing
 logger = get_logger(__name__)
 
 def filter_json_by_existed_files(directory, data, postfixes=[".mp4", ".jpg"]):
@@ -113,7 +113,7 @@ def filter_resolution(h, w, max_h_div_w_ratio=17/16, min_h_div_w_ratio=8 / 16):
 
 
 class T2V_dataset(Dataset):
-    def __init__(self, args, transform, temporal_sample, tokenizer, transform_topcrop, video_length_tolerance_range):
+    def __init__(self, args, transform, temporal_sample, tokenizer, transform_topcrop):
         self.data = args.data
         self.num_frames = args.num_frames
         self.train_fps = args.train_fps
@@ -122,7 +122,7 @@ class T2V_dataset(Dataset):
         self.transform_topcrop = transform_topcrop
         self.temporal_sample = temporal_sample
         self.tokenizer = tokenizer
-        self.model_max_length = args.model_max_length
+        self.text_max_length = args.text_max_length
         self.cfg = args.cfg
         self.speed_factor = args.speed_factor
         self.max_height = args.max_height
@@ -130,7 +130,7 @@ class T2V_dataset(Dataset):
         self.drop_short_ratio = args.drop_short_ratio
         assert self.speed_factor >= 1
         self.v_decoder = DecordInit()
-        self.video_length_tolerance_range = video_length_tolerance_range
+        self.video_length_tolerance_range = args.video_length_tolerance_range
         self.support_Chinese = True
         if not ('mt5' in args.text_encoder_name):
             self.support_Chinese = False
@@ -193,7 +193,7 @@ class T2V_dataset(Dataset):
         text = text_preprocessing(text, support_Chinese=self.support_Chinese) if random.random() > self.cfg else ""
         text_tokens_and_mask = self.tokenizer(
             text,
-            max_length=self.model_max_length,
+            max_length=self.text_max_length,
             padding='max_length',
             truncation=True,
             return_attention_mask=True,
@@ -227,7 +227,7 @@ class T2V_dataset(Dataset):
         text = text if random.random() > self.cfg else ""
         text_tokens_and_mask = self.tokenizer(
             text,
-            max_length=self.model_max_length,
+            max_length=self.text_max_length,
             padding='max_length',
             truncation=True,
             return_attention_mask=True,

fastvideo/dataset/transform.py

@@ -426,7 +426,7 @@ class CenterCropVideo:
         return f"{self.__class__.__name__}(size={self.size}, interpolation_mode={self.interpolation_mode}"
 
 
-class Normalize255:
+class Normalize:
     """
     Normalize the video clip by mean subtraction and division by standard deviation
     Args:
@@ -588,4 +588,4 @@ if __name__ == '__main__':
 
     for i in range(target_video_len):
         save_image(select_vframes_trans[i], os.path.join('./test000', '%04d.png' % i), normalize=True,
-                   value_range=(-1, 1))
+                   value_range=(-1, 1))

fastvideo/sample/generate_synthetic.py

@@ -0,0 +1,81 @@
+import json
+
+
+import torch
+from fastvideo.sample.mochi_pipe import MochiPipeline
+import os
+from diffusers.utils import export_to_video
+import argparse
+
+def generate_video_and_latent(pipe, prompt, height, width, num_frames, num_inference_steps, guidance_scale):
+    # Set the random seed for reproducibility
+    generator = torch.Generator("cpu").manual_seed(12345)
+    # Generate videos from the input prompt
+    video, latent, prompt_embed, prompt_attention_mask = pipe(
+        prompt=prompt,
+        height=height,
+        width=width,
+        num_frames=num_frames,
+        generator=generator,
+        num_inference_steps=num_inference_steps,
+        guidance_scale=guidance_scale,
+        output_type="latent_and_video"
+    )
+
+    return video[0], latent[0], prompt_embed[0], prompt_attention_mask[0]
+    
+    # return dummy tensor to debug first
+    # return torch.zeros(1, 3, 480, 848), torch.zeros(1, 256, 16, 16)
+
+
+if __name__ == "__main__":
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--num_frames", type=int, default=163)
+    parser.add_argument("--height", type=int, default=480)
+    parser.add_argument("--width", type=int, default=848)
+    parser.add_argument("--num_inference_steps", type=int, default=64)
+    parser.add_argument("--guidance_scale", type=float, default=4.5)
+    parser.add_argument("--model_path", type=str, default="data/mochi")
+    parser.add_argument("--prompt_path", type=str, default="data/dummyVid/videos2caption.json")
+    parser.add_argument("--dataset_output_dir", type=str, default="data/dummySynthetic")
+    args = parser.parse_args()
+
+    with open(args.prompt_path, 'r') as f:
+        data = json.load(f)
+    prompt_list = []
+    for item in data:
+        prompt_list.append(item["cap"])
+        
+        
+    pipe = MochiPipeline.from_pretrained(args.model_path, torch_dtype=torch.bfloat16)
+    pipe.enable_vae_tiling()
+    pipe.enable_model_cpu_offload(gpu_id=0)
+    # make dir if not exist 
+    
+    os.makedirs(args.dataset_output_dir, exist_ok=True)
+    os.makedirs(os.path.join(args.dataset_output_dir, "video"), exist_ok=True)
+    os.makedirs(os.path.join(args.dataset_output_dir, "latent"), exist_ok=True)
+    os.makedirs(os.path.join(args.dataset_output_dir, "prompt_embed"), exist_ok=True)
+    os.makedirs(os.path.join(args.dataset_output_dir, "prompt_attention_mask"), exist_ok=True)
+    for i, prompt in enumerate(prompt_list):
+        video, latent, prompt_embed, prompt_attention_mask = generate_video_and_latent(pipe, prompt, args.height, args.width, args.num_frames, args.num_inference_steps, args.guidance_scale)
+        # save latent
+        video_name = data[i]["path"].split("/")[-1].split(".")[0]
+        latent_path = os.path.join(args.dataset_output_dir, "latent", video_name + ".pt")
+        prompt_embed_path = os.path.join(args.dataset_output_dir, "prompt_embed", video_name + ".pt")
+        video_path = os.path.join(args.dataset_output_dir, "video", video_name + ".mp4")
+        prompt_attention_mask_path = os.path.join(args.dataset_output_dir, "prompt_attention_mask", video_name + ".pt")
+        # save latent
+        torch.save(latent, latent_path)
+        torch.save(prompt_embed, prompt_embed_path)
+        torch.save(prompt_attention_mask, prompt_attention_mask_path)
+        export_to_video(video, video_path, fps=30)
+        data[i]["latent_path"] = video_name + ".pt"
+        data[i]["prompt_embed_path"] = video_name + ".pt"
+        data[i]["prompt_attention_mask"] = video_name + ".pt"
+    # save json
+    with open(os.path.join(args.dataset_output_dir, "videos2caption.json"), 'w') as f:
+        json.dump(data, f, indent=4)
+
+    

fastvideo/sample/monkey_patch/sp_attention.py

@@ -0,0 +1,147 @@
+import inspect
+import math
+from typing import Callable, List, Optional, Tuple, Union
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from einops import rearrange, repeat
+
+from diffusers.image_processor import IPAdapterMaskProcessor
+from diffusers.utils import deprecate, logging
+from diffusers.utils.import_utils import is_torch_npu_available, is_xformers_available
+from diffusers.utils.torch_utils import is_torch_version, maybe_allow_in_graph
+from diffusers.models.embeddings import apply_rotary_emb
+
+import diffusers
+from types import MethodType
+
+from diffusers.models.attention_processor import Attention
+from fastvideo.utils.parallel_states import get_sequence_parallel_state, nccl_info
+from fastvideo.utils.communications import all_gather_BHSD, all_to_all_SBH
+
+class NewMochiAttnProcessor2_0:
+    """Attention processor used in Mochi."""
+
+    def __call__(
+        self,
+        attn: Attention,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        image_rotary_emb: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        # [b, s, h * d]
+        query = attn.to_q(hidden_states)
+        key = attn.to_k(hidden_states)
+        value = attn.to_v(hidden_states)
+
+        # [b, s, h=24, d=128]
+        query = query.unflatten(2, (attn.heads, -1))
+        key = key.unflatten(2, (attn.heads, -1))
+        value = value.unflatten(2, (attn.heads, -1))
+
+        # [b, 256, h * d]
+        encoder_query = attn.add_q_proj(encoder_hidden_states)
+        encoder_key = attn.add_k_proj(encoder_hidden_states)
+        encoder_value = attn.add_v_proj(encoder_hidden_states)
+
+        # [b, 256, h=24, d=128]
+        encoder_query = encoder_query.unflatten(2, (attn.heads, -1))
+        encoder_key = encoder_key.unflatten(2, (attn.heads, -1))
+        encoder_value = encoder_value.unflatten(2, (attn.heads, -1))
+
+        freqs_cos, freqs_sin = image_rotary_emb[0], image_rotary_emb[1]
+        # shard the head dimension
+        if get_sequence_parallel_state():
+            # B, S, H, D to (S, B,) H, D
+            batch_size, seq_len, attn_heads, head_dim = query.shape
+            query = rearrange(query, 'b s h d -> (s b) h d')
+            key = rearrange(key, 'b s h d -> (s b) h d')
+            value = rearrange(value, 'b s h d -> (s b) h d')
+            
+            query = all_to_all_SBH(query, scatter_dim=1, gather_dim=0).reshape(-1, batch_size, attn_heads // nccl_info.world_size, head_dim)
+            key = all_to_all_SBH(key, scatter_dim=1, gather_dim=0).reshape(-1, batch_size, attn_heads // nccl_info.world_size, head_dim)
+            value = all_to_all_SBH(value, scatter_dim=1, gather_dim=0).reshape(-1, batch_size, attn_heads // nccl_info.world_size, head_dim)
+            # batch_size, S * world_size, H / world_size, D
+            query = query.transpose(0, 1)
+            key = key.transpose(0, 1)
+            value = value.transpose(0, 1)
+            
+            
+            def shrink_head(encoder_state, dim):
+                local_heads = encoder_state.shape[dim] // nccl_info.world_size
+                return encoder_state.narrow(dim, nccl_info.rank * local_heads, local_heads)
+            encoder_query = shrink_head(encoder_query, dim=2)
+            encoder_key = shrink_head(encoder_key, dim=2)
+            encoder_value = shrink_head(encoder_value, dim=2)
+            
+            freqs_cos = shrink_head(freqs_cos, dim=1)
+            freqs_sin = shrink_head(freqs_sin, dim=1)
+    
+        if attn.norm_q is not None:
+            query = attn.norm_q(query)
+        if attn.norm_k is not None:
+            key = attn.norm_k(key)
+            
+        if attn.norm_added_q is not None:
+            encoder_query = attn.norm_added_q(encoder_query)
+        if attn.norm_added_k is not None:
+            encoder_key = attn.norm_added_k(encoder_key)
+    
+        if image_rotary_emb is not None:
+            def apply_rotary_emb(x, freqs_cos, freqs_sin):
+                x_even = x[..., 0::2].float()
+                x_odd = x[..., 1::2].float()
+
+                cos = (x_even * freqs_cos - x_odd * freqs_sin).to(x.dtype)
+                sin = (x_even * freqs_sin + x_odd * freqs_cos).to(x.dtype)
+
+                return torch.stack([cos, sin], dim=-1).flatten(-2)
+            query = apply_rotary_emb(query, freqs_cos, freqs_sin)
+            key = apply_rotary_emb(key, freqs_cos, freqs_sin)
+            
+        query, key, value = query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2)
+        encoder_query, encoder_key, encoder_value = (
+            encoder_query.transpose(1, 2),
+            encoder_key.transpose(1, 2),
+            encoder_value.transpose(1, 2),
+        )
+        # [b, h, s, d]
+        sequence_length = query.size(2)
+        encoder_sequence_length = encoder_query.size(2)
+
+        # Hint: please check encoder_query.shape
+        query = torch.cat([query, encoder_query], dim=2)
+        key = torch.cat([key, encoder_key], dim=2)
+        value = torch.cat([value, encoder_value], dim=2)
+        
+        
+                
+        hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False)
+        hidden_states, encoder_hidden_states = hidden_states.split_with_sizes(
+            (sequence_length, encoder_sequence_length), dim=2
+        )
+        if get_sequence_parallel_state():
+            hidden_states = rearrange(hidden_states, 'b h s d -> s b h d').reshape(-1, attn_heads // nccl_info.world_size, head_dim)
+            hidden_states = all_to_all_SBH(hidden_states, scatter_dim=0, gather_dim=1).reshape(-1, batch_size, attn_heads, head_dim).transpose(0, 1)
+            encoder_hidden_states = all_gather_BHSD(encoder_hidden_states, dim=1).contiguous()
+            hidden_states = hidden_states.flatten(2, 3)
+        else: 
+            hidden_states = hidden_states.transpose(1,2).flatten(2, 3)
+        hidden_states = hidden_states.to(query.dtype)
+        encoder_hidden_states = encoder_hidden_states.transpose(1, 2).flatten(2, 3)
+        encoder_hidden_states = encoder_hidden_states.to(query.dtype)
+        
+
+
+        # linear proj
+        hidden_states = attn.to_out[0](hidden_states)
+        # dropout
+        hidden_states = attn.to_out[1](hidden_states)
+
+        if hasattr(attn, "to_add_out"):
+            encoder_hidden_states = attn.to_add_out(encoder_hidden_states)
+
+        return hidden_states, encoder_hidden_states
+

fastvideo/sample/monkey_patch/sp_rope.py

@@ -0,0 +1,29 @@
+import inspect
+import math
+from typing import Callable, List, Optional, Tuple, Union
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from einops import rearrange, repeat
+
+from fastvideo.utils.parallel_states import get_sequence_parallel_state, nccl_info
+
+class NewMochiRoPE(nn.Module):
+    def _get_positions(
+        self,
+        num_frames: int,
+        height: int,
+        width: int,
+        device: Optional[torch.device] = None,
+        dtype: Optional[torch.dtype] = None,
+    ) -> torch.Tensor:
+        scale = (self.target_area / (height * width)) ** 0.5
+        t = torch.arange(num_frames * nccl_info.world_size, device=device, dtype=dtype)
+        h = self._centers(-height * scale / 2, height * scale / 2, height, device, dtype)
+        w = self._centers(-width * scale / 2, width * scale / 2, width, device, dtype)
+
+        grid_t, grid_h, grid_w = torch.meshgrid(t, h, w, indexing="ij")
+
+        positions = torch.stack([grid_t, grid_h, grid_w], dim=-1).view(-1, 3)
+        return positions

fastvideo/sample/pipeline_mochi.py

@@ -32,7 +32,9 @@ from diffusers.utils.torch_utils import randn_tensor
 from diffusers.video_processor import VideoProcessor
 from diffusers.pipelines.pipeline_utils import DiffusionPipeline
 from diffusers.pipelines.mochi.pipeline_output import MochiPipelineOutput
-
+from einops import rearrange
+from fastvideo.utils.parallel_states import get_sequence_parallel_state, nccl_info
+from fastvideo.utils.communications import  all_gather_BHSD
 
 if is_torch_xla_available():
     import torch_xla.core.xla_model as xm
@@ -86,7 +88,7 @@ def linear_quadratic_schedule(num_steps, threshold_noise, linear_steps=None):
     quadratic_sigma_schedule = [
         quadratic_coef * (i**2) + linear_coef * i + const for i in range(linear_steps, num_steps)
     ]
-    sigma_schedule = linear_sigma_schedule + quadratic_sigma_schedule + [1.0]
+    sigma_schedule = linear_sigma_schedule + quadratic_sigma_schedule
     sigma_schedule = [1.0 - x for x in sigma_schedule]
     return sigma_schedule
 
@@ -629,7 +631,11 @@ class MochiPipeline(DiffusionPipeline):
             generator,
             latents,
         )
-
+        world_size, rank = nccl_info.world_size, nccl_info.rank
+        if get_sequence_parallel_state():
+            latents = rearrange(latents, "b t (n s) h w -> b t n s h w", n=world_size).contiguous()
+            latents = latents[:, :, rank, :, :, :]
+            
         # 5. Prepare timestep
         # from https://github.com/genmoai/models/blob/075b6e36db58f1242921deff83a1066887b9c9e1/src/mochi_preview/infer.py#L77
         threshold_noise = 0.025
@@ -693,6 +699,15 @@ class MochiPipeline(DiffusionPipeline):
                 if XLA_AVAILABLE:
                     xm.mark_step()
 
+        if get_sequence_parallel_state():
+            latents = all_gather_BHSD(latents, dim=2)
+            #latents_shape = list(latents.shape)
+            #full_shape = [latents_shape[0] * world_size] + latents_shape[1:]
+            #all_latents = torch.zeros(full_shape, dtype=latents.dtype, device=latents.device)
+            #torch.distributed.all_gather_into_tensor(all_latents, latents)
+            #latents_list = list(all_latents.chunk(world_size, dim=0))
+            #latents = torch.cat(latents_list, dim=2)
+
         if output_type == "latent":
             video = latents
         else:
@@ -720,4 +735,4 @@ class MochiPipeline(DiffusionPipeline):
         if not return_dict:
             return (video,)
 
-        return MochiPipelineOutput(frames=video)
+        return MochiPipelineOutput(frames=video)

fastvideo/sample/sample_t2v_mochi.py

@@ -0,0 +1,47 @@
+import torch
+from diffusers import MochiPipeline, MochiTransformer3DModel
+from diffusers.utils import export_to_video, load_image, load_video
+import argparse
+
+def main(args):
+    # Set the random seed for reproducibility
+    generator = torch.Generator("cuda").manual_seed(args.seed)
+
+    if args.transformer_path is not None:
+        transformer = MochiTransformer3DModel.from_pretrained(args.transformer_path, torch_dtype=torch.bfloat16)
+        pipe = MochiPipeline.from_pretrained(args.model_path, transformer = transformer, torch_dtype=torch.bfloat16)
+    else:
+        pipe = MochiPipeline.from_pretrained(args.model_path,  torch_dtype=torch.bfloat16)
+    pipe.enable_vae_tiling()
+    # pipe.to("cuda:1")
+    pipe.enable_model_cpu_offload()
+
+    # Generate videos from the input prompt
+    video = pipe(
+        prompt=args.prompt,
+        height=args.height,
+        width=args.width,
+        num_frames=args.num_frames,
+        generator=generator,
+        num_inference_steps=args.num_inference_steps,
+        guidance_scale=args.guidance_scale,
+    ).frames[0]
+
+
+    export_to_video(video, args.output_path, fps=30)
+
+if __name__ == "__main__":
+    # arg parse 
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--prompt", type=str, default="A hand with delicate fingers picks up a bright yellow lemon from a wooden bowl filled with lemons and sprigs of mint against a peach-colored background. The hand gently tosses the lemon up and catches it, showcasing its smooth texture. A beige string bag sits beside the bowl, adding a rustic touch to the scene. Additional lemons, one halved, are scattered around the base of the bowl. The even lighting enhances the vibrant colors and creates a fresh, inviting atmosphere.")
+    parser.add_argument("--num_frames", type=int, default=163)
+    parser.add_argument("--height", type=int, default=480)
+    parser.add_argument("--width", type=int, default=848)
+    parser.add_argument("--num_inference_steps", type=int, default=64)
+    parser.add_argument("--guidance_scale", type=float, default=4.5)
+    parser.add_argument("--model_path", type=str, default="data/mochi")
+    parser.add_argument("--seed", type=int, default=12345)
+    parser.add_argument("--transformer_path", type=str, default=None)
+    parser.add_argument("--output_path", type=str, default="./outputs.mp4")
+    args = parser.parse_args()
+    main(args)

fastvideo/sample/sample_t2v_mochi_sp.py

@@ -0,0 +1,67 @@
+import torch
+from fastvideo.sample.pipeline_mochi import MochiPipeline
+import torch.distributed as dist
+
+from diffusers.utils import export_to_video, load_image, load_video
+from fastvideo.utils.parallel_states import initialize_sequence_parallel_state, nccl_info
+import argparse
+import os
+import diffusers
+from fastvideo.sample.monkey_patch.sp_attention import NewMochiAttnProcessor2_0
+from fastvideo.sample.monkey_patch.sp_rope import NewMochiRoPE
+
+def initialize_distributed():
+    local_rank = int(os.getenv('RANK', 0))
+    world_size = int(os.getenv('WORLD_SIZE', 1))
+    print('world_size', world_size)
+    torch.cuda.set_device(local_rank)
+    dist.init_process_group(backend='nccl', init_method='env://', world_size=world_size, rank=local_rank)
+    initialize_sequence_parallel_state(world_size)
+    
+
+def hf_mochi_add_sp_monkey_patch():
+    initialize_distributed()
+    diffusers.models.attention_processor.MochiAttnProcessor2_0.__call__ = NewMochiAttnProcessor2_0.__call__
+    diffusers.models.transformers.transformer_mochi.MochiRoPE._get_positions = NewMochiRoPE._get_positions
+
+def main(args):
+    hf_mochi_add_sp_monkey_patch()
+    print(nccl_info.world_size)
+    device = torch.cuda.current_device()
+    generator = torch.Generator(device).manual_seed(args.seed)
+    weight_dtype = torch.bfloat16
+    
+    pipe = MochiPipeline.from_pretrained(args.model_path, torch_dtype=torch.bfloat16)
+    pipe.enable_vae_tiling()
+    pipe.to(device)
+    #pipe.enable_model_cpu_offload()
+
+    # Generate videos from the input prompt
+    video = pipe(
+        prompt=args.prompt,
+        height=args.height,
+        width=args.width,
+        num_frames=args.num_frames,
+        num_inference_steps=args.num_inference_steps,
+        guidance_scale=args.guidance_scale,
+        generator=generator,
+    ).frames[0]
+
+    dist.barrier()
+    if nccl_info.rank <= 0:
+        export_to_video(video, args.output_path, fps=30)
+
+if __name__ == "__main__":
+    # arg parse 
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--prompt", type=str, default="A hand with delicate fingers picks up a bright yellow lemon from a wooden bowl filled with lemons and sprigs of mint against a peach-colored background. The hand gently tosses the lemon up and catches it, showcasing its smooth texture. A beige string bag sits beside the bowl, adding a rustic touch to the scene. Additional lemons, one halved, are scattered around the base of the bowl. The even lighting enhances the vibrant colors and creates a fresh, inviting atmosphere.")
+    parser.add_argument("--num_frames", type=int, default=163)
+    parser.add_argument("--height", type=int, default=480)
+    parser.add_argument("--width", type=int, default=848)
+    parser.add_argument("--num_inference_steps", type=int, default=64)
+    parser.add_argument("--guidance_scale", type=float, default=4.5)
+    parser.add_argument("--model_path", type=str, default="data/mochi")
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--output_path", type=str, default="./outputs.mp4")
+    args = parser.parse_args()
+    main(args)

fastvideo/train.py

@@ -0,0 +1,684 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""
+A minimal training script for DiT using PyTorch DDP.
+"""
+import argparse
+import logging
+import math
+import os
+import shutil
+from pathlib import Path
+import gc
+import wandb
+from einops import rearrange
+from tqdm import tqdm
+from diffusers.training_utils import cast_training_params, compute_density_for_timestep_sampling, compute_loss_weighting_for_sd3
+from fastvideo.utils.parallel_states import initialize_sequence_parallel_state, \
+    destroy_sequence_parallel_group, get_sequence_parallel_state, set_sequence_parallel_state
+from fastvideo.utils.communications import prepare_parallel_data, broadcast
+import time
+from torch.utils.data import DataLoader
+from copy import deepcopy
+import accelerate
+import torch
+from torch.nn import functional as F
+import transformers
+from accelerate import Accelerator, init_empty_weights, DistributedType
+from diffusers.utils.torch_utils import is_compiled_module
+from accelerate.logging import get_logger
+from accelerate.utils import DistributedType, ProjectConfiguration, set_seed, DummyOptim, DummyScheduler
+from tqdm.auto import tqdm
+
+import diffusers
+from diffusers import (
+    AutoencoderKLMochi,
+    FlowMatchEulerDiscreteScheduler,
+    MochiTransformer3DModel,
+    MochiPipeline,
+)
+from diffusers.optimization import get_scheduler
+from diffusers.training_utils import compute_snr
+from diffusers.utils import check_min_version, is_wandb_available
+from transformers import AutoTokenizer
+from fastvideo.utils.ema import EMAModel
+from fastvideo.dataset.latent_datasets import LatentDataset, latent_collate_function
+
+
+# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
+check_min_version("0.31.0")
+logger = get_logger(__name__)
+
+
+import sys
+import pdb
+
+class ForkedPdb(pdb.Pdb):
+    """A Pdb subclass that may be used
+    from a forked multiprocessing child
+
+    """
+    def interaction(self, *args, **kwargs):
+        _stdin = sys.stdin
+        try:
+            sys.stdin = open('/dev/stdin')
+            pdb.Pdb.interaction(self, *args, **kwargs)
+        finally:
+            sys.stdin = _stdin
+            
+
+@torch.inference_mode()
+def log_validation(args, transformer,vae,  accelerator, weight_dtype, global_step,  ema=False):
+    logger.info(f"Running validation....\n")
+    transformer = accelerator.unwrap_model(transformer)
+    mochi_pipeline = MochiPipeline.from_pretrained(args.pretrained_model_name_or_path,vae=vae, transformer=transformer, torch_dtype=weight_dtype).to(device=accelerator.device)
+    videos = []
+    for prompt in args.validaiton_prompt:
+        logger.info('Processing the ({}) prompt'.format(prompt))
+        video = mochi_pipeline(
+                                args.prompt,
+                                num_frames=args.num_frames,
+                                height=args.max_height,
+                                width=args.max_width,
+                                num_inference_steps=args.validation_sampling_steps,
+                                guidance_scale=args.guidance_scale,
+                                ).images
+        videos.append(video[0])
+    # import ipdb;ipdb.set_trace()
+    gc.collect()
+    torch.cuda.empty_cache()
+    videos = torch.stack(videos).numpy()
+    videos = rearrange(videos, 'b t h w c -> b t c h w')
+    for tracker in accelerator.trackers:
+        logs = {
+            f"{'ema_' if ema else ''}validation": [
+                wandb.Video(video, caption=f"{i}: {prompt}", fps=24)
+                for i, (video, prompt) in enumerate(zip(videos, args.validaiton_prompt))
+            ]
+        }
+        tracker.log(logs, step=global_step)
+
+    del mochi_pipeline
+    gc.collect()
+    torch.cuda.empty_cache()
+
+
+class ProgressInfo:
+    def __init__(self, global_step, train_loss=0.0):
+        self.global_step = global_step
+        self.train_loss = train_loss
+
+
+#################################################################################
+#                                  Training Loop                                #
+#################################################################################
+
+def main(args):
+    logging_dir = Path(args.output_dir, args.logging_dir)
+
+    # use LayerNorm, GeLu, SiLu always as fp32 mode
+    # TODO: 
+    if args.enable_stable_fp32:
+        raise NotImplementedError("enable_stable_fp32 is not supported now.")
+
+    accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, logging_dir=logging_dir)
+
+    accelerator = Accelerator(
+        gradient_accumulation_steps=args.gradient_accumulation_steps,
+        mixed_precision=args.mixed_precision,
+        log_with="wandb",
+        project_config=accelerator_project_config,
+    )
+
+    initialize_sequence_parallel_state(args.sp_size)
+
+    if not is_wandb_available():
+        raise ImportError("Make sure to install wandb if you want to use it for logging during training.")
+    # Make one log on every process with the configuration for debugging.
+    logging.basicConfig(
+        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+        datefmt="%m/%d/%Y %H:%M:%S",
+        level=logging.INFO,
+    )
+    logger.info(accelerator.state, main_process_only=False)
+    if accelerator.is_local_main_process:
+        transformers.utils.logging.set_verbosity_warning()
+        diffusers.utils.logging.set_verbosity_info()
+    else:
+        transformers.utils.logging.set_verbosity_error()
+        diffusers.utils.logging.set_verbosity_error()
+
+    # If passed along, set the training seed now. On GPU...
+    if args.seed is not None:
+        set_seed(args.seed + accelerator.process_index)
+
+    # Handle the repository creation
+    if accelerator.is_main_process:
+        if args.output_dir is not None:
+            os.makedirs(args.output_dir, exist_ok=True)
+
+    # For mixed precision training we cast all non-trainable weigths to half-precision
+    # as these weights are only used for inference, keeping weights in full precision is not required.
+    weight_dtype = torch.float32
+    if accelerator.mixed_precision == "fp16":
+        weight_dtype = torch.float16
+    elif accelerator.mixed_precision == "bf16":
+        weight_dtype = torch.bfloat16
+
+    # Create model:
+    
+    vae = AutoencoderKLMochi.from_pretrained(args.pretrained_model_name_or_path, subfolder="vae", torch_dtype=weight_dtype)
+    if args.enable_tiling:
+        vae.enable_tiling()
+        vae.tile_sample_stride_height = args.tile_sample_stride
+        vae.tile_sample_stride_width = args.tile_sample_stride
+    
+    load_dtype = torch.bfloat16 # TODO
+    transformer = MochiTransformer3DModel.from_pretrained(
+        args.pretrained_model_name_or_path,
+        subfolder="transformer",
+        torch_dtype=load_dtype,
+    )
+    
+    if args.gradient_checkpointing:
+        transformer.enable_gradient_checkpointing()
+
+    # Freeze vae and text encoders.
+    vae.requires_grad_(False)
+    # Set model as trainable.
+    transformer.train()
+
+    noise_scheduler = FlowMatchEulerDiscreteScheduler()
+    # Move unet, vae and text_encoder to device and cast to weight_dtype
+    # The VAE is in float32 to avoid NaN losses.
+    vae.to(accelerator.device, dtype=weight_dtype)
+    # ae.vae.to(accelerator.device, dtype=weight_dtype)
+    # transformer.to(accelerator.device, dtype=weight_dtype)
+
+    # Create EMA for the unet.
+    if args.use_ema:
+        ema_transformer = deepcopy(transformer)
+        ema_transformer = EMAModel(ema_transformer.parameters(), decay=args.ema_decay, update_after_step=args.ema_start_step,
+                             model_cls=MochiTransformer3DModel, model_config=ema_transformer.config)
+
+    def unwrap_model(model):
+        model = accelerator.unwrap_model(model)
+        model = model._orig_mod if is_compiled_module(model) else model
+        return model
+    # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format
+    def save_model_hook(models, weights, output_dir):
+        if accelerator.is_main_process:
+            if args.use_ema:
+                ema_transformer.save_pretrained(os.path.join(output_dir, "model_ema"))
+
+            for model in models:
+                if isinstance(unwrap_model(model), type(unwrap_model(transformer))):
+                    model: MochiTransformer3DModel = unwrap_model(model)
+                    model.save_pretrained(
+                        os.path.join(output_dir, "transformer"), safe_serialization=True, max_shard_size="5GB"
+                    )
+                else:
+                    raise ValueError(f"Unexpected save model: {model.__class__}")
+                # make sure to pop weight so that corresponding model is not saved again
+                if weights:
+                    weights.pop()
+
+
+    def load_model_hook(models, input_dir):
+        if args.use_ema:
+            load_model = EMAModel.from_pretrained(os.path.join(input_dir, "model_ema"), MochiTransformer3DModel)
+            ema_transformer.load_state_dict(load_model.state_dict())
+            ema_transformer.to(accelerator.device)
+            del load_model
+
+        if not accelerator.distributed_type == DistributedType.DEEPSPEED:
+            while len(models) > 0:
+                model = models.pop()
+
+                if isinstance(unwrap_model(model), type(unwrap_model(transformer))):
+                    transformer_ = unwrap_model(model)
+                else:
+                    raise ValueError(f"Unexpected save model: {unwrap_model(model).__class__}")
+        else:
+            with init_empty_weights():
+                transformer_ = MochiTransformer3DModel.from_config(
+                    args.pretrained_model_name_or_path, subfolder="transformer"
+                )
+                init_under_meta = True
+        
+        load_model = MochiTransformer3DModel.from_pretrained(os.path.join(input_dir, "transformer"))
+        transformer_.register_to_config(**load_model.config)
+        transformer_.load_state_dict(load_model.state_dict(), assign=init_under_meta)
+        del load_model
+        
+        if args.mixed_precision == "fp16":
+            cast_training_params([transformer_])
+
+    accelerator.register_save_state_pre_hook(save_model_hook)
+    accelerator.register_load_state_pre_hook(load_model_hook)
+
+    # Enable TF32 for faster training on Ampere GPUs,
+    # cf https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices
+    if args.allow_tf32 and torch.cuda.is_available():
+        torch.backends.cuda.matmul.allow_tf32 = True
+
+    if args.scale_lr:
+        args.learning_rate = (
+                args.learning_rate * args.gradient_accumulation_steps * args.train_batch_size * accelerator.num_processes
+        )
+
+    # Make sure the trainable params are in float32.
+    # if args.mixed_precision == "fp16" or args.mixed_precision == "bf16":
+    #     cast_training_params([transformer], dtype=torch.float32)
+        
+
+    params_to_optimize = transformer.parameters()
+
+    optimizer = DummyOptim(params_to_optimize, lr=args.learning_rate)
+    lr_scheduler = DummyScheduler(
+        optimizer, 
+        warmup_num_steps=args.lr_warmup_steps * args.gradient_accumulation_steps, 
+        total_num_steps=args.max_train_steps * args.gradient_accumulation_steps
+    )
+
+    logger.info(f"optimizer: {optimizer}")
+    
+    train_dataset = LatentDataset(args.data_merge_path, args.num_latent_t)
+    train_dataloader = DataLoader(
+        train_dataset,
+        shuffle=True,
+        collate_fn=latent_collate_function,
+        pin_memory=True,
+        batch_size=args.train_batch_size,
+        num_workers=args.dataloader_num_workers,
+        drop_last=True, 
+    )
+
+    # Scheduler and math around the number of training steps.
+    overrode_max_train_steps = False
+    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps * args.sp_size / args.train_sp_batch_size)
+    if args.max_train_steps is None:
+        args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
+        overrode_max_train_steps = True
+
+
+
+    # Prepare everything with our `accelerator`.
+    # model.requires_grad_(False)
+    # model.pos_embed.requires_grad_(True)
+    logger.info(f'before accelerator.prepare')
+    transformer, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
+        transformer, optimizer, train_dataloader, lr_scheduler
+    )
+    logger.info(f'after accelerator.prepare')
+    if args.use_ema:
+        ema_transformer.to(accelerator.device)
+
+    # We need to recalculate our total training steps as the size of the training dataloader may have changed.
+    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps * args.sp_size / args.train_sp_batch_size)
+    if overrode_max_train_steps:
+        args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
+    # Afterwards we recalculate our number of training epochs
+    args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)
+
+    # We need to initialize the trackers we use, and also store our configuration.
+    # The trackers initializes automatically on the main process.
+    if accelerator.is_main_process:
+        tracker_project_name = args.tracker_project_name or "fastvideo"
+        accelerator.init_trackers(tracker_project_name, config=vars(args))
+
+    # Train!
+    total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
+    total_batch_size = total_batch_size // args.sp_size * args.train_sp_batch_size
+    logger.info("***** Running training *****")
+    logger.info(f"  Num examples = {len(train_dataset)}")
+    logger.info(f"  Num Epochs = {args.num_train_epochs}")
+    logger.info(f"  Instantaneous batch size per device = {args.train_batch_size}")
+    logger.info(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")
+    logger.info(f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")
+    logger.info(f"  Total optimization steps = {args.max_train_steps}")
+    logger.info(f"  Total training parameters = {sum(p.numel() for p in transformer.parameters() if p.requires_grad) / 1e9} B")
+    global_step = 0
+    first_epoch = 0
+
+    # Potentially load in the weights and states from a previous save
+    if args.resume_from_checkpoint:
+        if args.resume_from_checkpoint != "latest":
+            path = os.path.basename(args.resume_from_checkpoint)
+        else:
+            # Get the most recent checkpoint
+            dirs = os.listdir(args.output_dir)
+            dirs = [d for d in dirs if d.startswith("checkpoint")]
+            dirs = sorted(dirs, key=lambda x: int(x.split("-")[1]))
+            path = dirs[-1] if len(dirs) > 0 else None
+
+        if path is None:
+            accelerator.print(
+                f"Checkpoint '{args.resume_from_checkpoint}' does not exist. Starting a new training run."
+            )
+            args.resume_from_checkpoint = None
+            initial_global_step = 0
+        else:
+            accelerator.print(f"Resuming from checkpoint {path}")
+            accelerator.load_state(os.path.join(args.output_dir, path))
+            global_step = int(path.split("-")[1])
+
+            initial_global_step = global_step
+            first_epoch = global_step // num_update_steps_per_epoch
+    else:
+        initial_global_step = 0
+
+    progress_bar = tqdm(
+        range(0, args.max_train_steps),
+        initial=initial_global_step,
+        desc="Steps",
+        # Only show the progress bar once on each machine.
+        disable=not accelerator.is_local_main_process,
+    )
+    progress_info = ProgressInfo(global_step, train_loss=0.0)
+
+    def sync_gradients_info(loss):
+        # Checks if the accelerator has performed an optimization step behind the scenes
+        if args.use_ema:
+            ema_transformer.step(transformer.parameters())
+        progress_bar.update(1)
+        progress_info.global_step += 1
+        end_time = time.time()
+        one_step_duration = end_time - start_time
+        accelerator.log({"train_loss": progress_info.train_loss}, step=progress_info.global_step)
+        logs = {"step_loss": progress_info.train_loss, "lr": lr_scheduler.get_last_lr()[0]}
+        progress_bar.set_postfix(**logs)
+        
+        progress_info.train_loss = 0.0
+
+        # DeepSpeed requires saving weights on every device; saving weights only on the main process would cause issues.
+        if accelerator.distributed_type == DistributedType.DEEPSPEED or accelerator.is_main_process:
+            if progress_info.global_step % args.checkpointing_steps == 0:
+                # _before_ saving state, check if this save would set us over the `checkpoints_total_limit`
+                if accelerator.is_main_process and args.checkpoints_total_limit is not None:
+                    checkpoints = os.listdir(args.output_dir)
+                    checkpoints = [d for d in checkpoints if d.startswith("checkpoint")]
+                    checkpoints = sorted(checkpoints, key=lambda x: int(x.split("-")[1]))
+
+                    # before we save the new checkpoint, we need to have at _most_ `checkpoints_total_limit - 1` checkpoints
+                    if len(checkpoints) >= args.checkpoints_total_limit:
+                        num_to_remove = len(checkpoints) - args.checkpoints_total_limit + 1
+                        removing_checkpoints = checkpoints[0:num_to_remove]
+
+                        logger.info(
+                            f"{len(checkpoints)} checkpoints already exist, removing {len(removing_checkpoints)} checkpoints"
+                        )
+                        logger.info(f"removing checkpoints: {', '.join(removing_checkpoints)}")
+
+                        for removing_checkpoint in removing_checkpoints:
+                            removing_checkpoint = os.path.join(args.output_dir, removing_checkpoint)
+                            shutil.rmtree(removing_checkpoint)
+
+                save_path = os.path.join(args.output_dir, f"checkpoint-{progress_info.global_step}")
+                accelerator.save_state(save_path)
+                logger.info(f"Saved state to {save_path}")
+                
+    def get_sigmas(timesteps, n_dim=4, dtype=torch.float32):
+        sigmas = noise_scheduler.sigmas.to(device=accelerator.device, dtype=dtype)
+        schedule_timesteps = noise_scheduler.timesteps.to(accelerator.device)
+        timesteps = timesteps.to(accelerator.device)
+        step_indices = [(schedule_timesteps == t).nonzero().item() for t in timesteps]
+
+        sigma = sigmas[step_indices].flatten()
+        while len(sigma.shape) < n_dim:
+            sigma = sigma.unsqueeze(-1)
+        return sigma
+    def run(latents, encoder_hidden_states,encoder_attention_mask, prof):
+        global start_time
+        start_time = time.time()
+        bsz = latents.shape[0]
+        noise = torch.randn_like(latents)
+        u = compute_density_for_timestep_sampling(
+            weighting_scheme=args.weighting_scheme,
+            batch_size=bsz,
+            logit_mean=args.logit_mean,
+            logit_std=args.logit_std,
+            mode_scale=args.mode_scale,
+        )
+        
+        indices = (u * noise_scheduler.config.num_train_timesteps).long()
+        timesteps = noise_scheduler.timesteps[indices].to(device=latents.device)
+
+        # Add noise according to flow matching.
+        # zt = (1 - texp) * x + texp * z1
+        sigmas = get_sigmas(timesteps, n_dim=latents.ndim, dtype=latents.dtype)
+        noisy_model_input = (1.0 - sigmas) * latents + sigmas * noise
+        model_pred = transformer(
+            noisy_model_input,
+            encoder_hidden_states,
+            noise_scheduler.config.num_train_timesteps - timesteps, # the timestep for mochi is invereted
+            encoder_attention_mask, # B, L
+            return_dict= False
+        )[0]
+        
+        if args.precondition_outputs:
+            model_pred = model_pred * sigmas + noisy_model_input
+
+        # these weighting schemes use a uniform timestep sampling
+        # and instead post-weight the loss
+        weighting = compute_loss_weighting_for_sd3(weighting_scheme=args.weighting_scheme, sigmas=sigmas)
+
+                # flow matching loss
+        if args.precondition_outputs:
+            target = latents
+        else:
+            target = latents - noise
+
+        loss = torch.mean(weighting.float() * (model_pred.float() - target.float()) ** 2)
+
+        # Gather the losses across all processes for logging (if we use distributed training).
+        # TODO Why repeat here? Weird
+        avg_loss = accelerator.reduce(loss.clone().detach(), "mean")
+        progress_info.train_loss += avg_loss.item() / args.gradient_accumulation_steps
+
+        # Backpropagate
+        accelerator.backward(loss)
+        if accelerator.sync_gradients:
+            params_to_clip = transformer.parameters()
+            accelerator.clip_grad_norm_(params_to_clip, args.max_grad_norm)
+        optimizer.step()
+        lr_scheduler.step()
+        optimizer.zero_grad()
+
+        if accelerator.sync_gradients:
+            sync_gradients_info(loss)
+
+        if accelerator.is_main_process:
+
+            if progress_info.global_step % args.checkpointing_steps == 0:
+
+                if args.enable_tracker:
+                    log_validation(args, transformer, vae,  accelerator,
+                                   weight_dtype, progress_info.global_step)
+
+                    if args.use_ema:
+                        # Store the UNet parameters temporarily and load the EMA parameters to perform inference.
+                        ema_transformer.store(transformer.parameters())
+                        ema_transformer.copy_to(transformer.parameters())
+                        log_validation(args, transformer,vae,  accelerator,
+                                    weight_dtype, progress_info.global_step, ema=True)
+                        # Switch back to the original UNet parameters.
+                        ema_transformer.restore(transformer.parameters())
+
+        if prof is not None:
+            prof.step()
+
+
+        return loss
+
+    def train_one_step( data_item_, prof_=None):
+        latents, cond,attn_mask, cond_mask = data_item_    
+        # TODO
+        latents = latents.to(dtype=torch.bfloat16)
+        cond = cond.to(dtype=torch.bfloat16)
+        latents_mean = (
+            torch.tensor(vae.config.latents_mean).view(1, 12, 1, 1, 1).to(latents.device, latents.dtype)
+        )
+        latents_std = (
+            torch.tensor(vae.config.latents_std).view(1, 12, 1, 1, 1).to(latents.device, latents.dtype)
+        )
+
+        latents = (latents - latents_mean) * vae.config.scaling_factor / latents_std
+        
+            
+        if get_sequence_parallel_state():
+            latents, cond, attn_mask, cond_mask = prepare_parallel_data(latents, cond, attn_mask, cond_mask)
+            for iter in range(args.train_batch_size * args.sp_size // args.train_sp_batch_size):
+                with accelerator.accumulate(transformer):
+                    st_idx = iter * args.train_sp_batch_size
+                    ed_idx = (iter + 1) * args.train_sp_batch_size
+                    encoder_hidden_states=cond[st_idx: ed_idx]
+                    # TODO
+                    attention_mask=attn_mask[st_idx: ed_idx]
+                    encoder_attention_mask=cond_mask[st_idx: ed_idx]
+                    run(latents[st_idx: ed_idx], encoder_hidden_states, encoder_attention_mask, prof_)
+        else:
+            with accelerator.accumulate(transformer):
+                encoder_hidden_states=cond
+                attention_mask=attn_mask,
+                encoder_attention_mask=cond_mask
+                run(latents, encoder_hidden_states, encoder_attention_mask, prof_)
+
+        if progress_info.global_step >= args.max_train_steps:
+            return True
+
+        return False
+
+    def train_all_epoch(prof_=None):
+        for epoch in range(first_epoch, args.num_train_epochs):
+            progress_info.train_loss = 0.0
+            if progress_info.global_step >= args.max_train_steps:
+                return True
+
+            for step, data_item in enumerate(train_dataloader):
+                if train_one_step(data_item, prof_):
+                    break
+
+
+    train_all_epoch()
+    accelerator.wait_for_everyone()
+    accelerator.end_training()
+    if get_sequence_parallel_state():
+        destroy_sequence_parallel_group()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+
+    # dataset & dataloader
+    parser.add_argument("--data_merge_path", type=str, required=True)
+    parser.add_argument("--num_frames", type=int, default=65)
+    parser.add_argument("--dataloader_num_workers", type=int, default=10, help="Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process.")
+    parser.add_argument("--train_batch_size", type=int, default=16, help="Batch size (per device) for the training dataloader.")
+    parser.add_argument("--num_latent_t", type=int, default=28, help="Number of latent timesteps.")
+    parser.add_argument("--group_frame", action="store_true") # TODO
+    parser.add_argument("--group_resolution", action="store_true") # TODO
+
+    # text encoder & vae & diffusion model
+    parser.add_argument("--pretrained_model_name_or_path", type=str)
+    parser.add_argument('--tile_sample_stride', type=int, default=192)
+    parser.add_argument('--enable_tiling', action='store_true')
+    parser.add_argument("--downsampler", type=str, default=None)
+    parser.add_argument("--cache_dir", type=str, default='./cache_dir')
+    parser.add_argument("--pretrained", type=str, default=None)
+    parser.add_argument('--enable_stable_fp32', action='store_true') # TODO
+
+    # diffusion setting
+    parser.add_argument("--use_ema", action="store_true", help="Whether to use EMA model.")
+    parser.add_argument("--ema_decay", type=float, default=0.999)
+    parser.add_argument("--ema_start_step", type=int, default=0)
+    parser.add_argument('--cfg', type=float, default=0.1)
+    parser.add_argument("--prediction_type", type=str, default=None, help="The prediction_type that shall be used for training. Choose between 'epsilon' or 'v_prediction' or leave `None`. If left to `None` the default prediction type of the scheduler: `noise_scheduler.config.prediciton_type` is chosen.")
+            
+    parser.add_argument(
+        "--weighting_scheme",
+        type=str,
+        default="logit_normal",
+        choices=["sigma_sqrt", "logit_normal", "mode", "cosmap"],
+    )
+    parser.add_argument(
+        "--logit_mean", type=float, default=0.0, help="mean to use when using the `'logit_normal'` weighting scheme."
+    )
+    parser.add_argument(
+        "--logit_std", type=float, default=1.0, help="std to use when using the `'logit_normal'` weighting scheme."
+    )
+    parser.add_argument(
+        "--mode_scale",
+        type=float,
+        default=1.29,
+        help="Scale of mode weighting scheme. Only effective when using the `'mode'` as the `weighting_scheme`.",
+    )
+    parser.add_argument("--precondition_outputs", action="store_true", help="Whether to precondition the outputs of the model.")
+    
+    # validation & logs
+    parser.add_argument("--validaiton_prompt", nargs='+', help="List of prompts to use for validation.")
+    parser.add_argument("--validation_sampling_steps", type=int, default=64)
+    parser.add_argument('--guidance_scale', type=float, default=4.5)
+    parser.add_argument("--enable_tracker", action="store_true")
+    parser.add_argument("--tracker_project_name", type=str, default=None)
+    parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.")
+    parser.add_argument("--output_dir", type=str, default=None, help="The output directory where the model predictions and checkpoints will be written.")
+    parser.add_argument("--checkpoints_total_limit", type=int, default=None, help=("Max number of checkpoints to store."))
+    parser.add_argument("--checkpointing_steps", type=int, default=500,
+                        help=(
+                            "Save a checkpoint of the training state every X updates. These checkpoints can be used both as final"
+                            " checkpoints in case they are better than the last checkpoint, and are also suitable for resuming"
+                            " training using `--resume_from_checkpoint`."
+                        ),
+                        )
+    parser.add_argument("--resume_from_checkpoint", type=str, default=None,
+                        help=(
+                            "Whether training should be resumed from a previous checkpoint. Use a path saved by"
+                            ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.'
+                        ),
+                        )
+    parser.add_argument("--logging_dir", type=str, default="logs",
+                        help=(
+                            "[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"
+                            " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."
+                        ),
+                        )
+
+    # optimizer & scheduler & Training
+    parser.add_argument("--num_train_epochs", type=int, default=100)
+    parser.add_argument("--max_train_steps", type=int, default=None, help="Total number of training steps to perform.  If provided, overrides num_train_epochs.")
+    parser.add_argument("--gradient_accumulation_steps", type=int, default=1, help="Number of updates steps to accumulate before performing a backward/update pass.")
+    parser.add_argument("--learning_rate", type=float, default=1e-4, help="Initial learning rate (after the potential warmup period) to use.")
+    parser.add_argument("--scale_lr", action="store_true", default=False, help="Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.")
+    parser.add_argument("--lr_warmup_steps", type=int, default=10, help="Number of steps for the warmup in the lr scheduler.")
+    parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.")
+    parser.add_argument("--gradient_checkpointing", action="store_true", help="Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.")
+    parser.add_argument("--lr_scheduler", type=str, default="constant",
+help=(
+                            'The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",'
+                            ' "constant", "constant_with_warmup"]'
+                        ),
+                        )
+    parser.add_argument("--allow_tf32", action="store_true",
+                        help=(
+                            "Whether or not to allow TF32 on Ampere GPUs. Can be used to speed up training. For more information, see"
+                            " https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices"
+                        ),
+                        )
+    parser.add_argument("--mixed_precision", type=str, default=None, choices=["no", "fp16", "bf16"],
+                        help=(
+                            "Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >="
+                            " 1.10.and an Nvidia Ampere GPU.  Default to the value of accelerate config of the current system or the"
+                            " flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config."
+                        ),
+                        )
+
+    parser.add_argument("--sp_size", type=int, default=1, help="For sequence parallel")
+    parser.add_argument("--train_sp_batch_size", type=int, default=1, help="Batch size for sequence parallel training")
+
+    args = parser.parse_args()
+    main(args)

fastvideo/utils/communications.py

@@ -1,7 +1,7 @@
 import torch
 import torch.distributed as dist
 from einops import rearrange
-from opensora.utils.parallel_states import nccl_info
+from fastvideo.utils.parallel_states import nccl_info
 
 def broadcast(input_: torch.Tensor):
     sp_size = nccl_info.world_size
@@ -9,25 +9,13 @@ def broadcast(input_: torch.Tensor):
     dist.broadcast(input_, src=src, group=nccl_info.group)
 
 _COUNT = 0
-def _all_to_all(
-    input_: torch.Tensor,
-    scatter_dim: int,
-    gather_dim: int,
-):
-    group = nccl_info.group
-    sp_size = nccl_info.world_size
-    input_list = [t.contiguous() for t in torch.tensor_split(input_, sp_size, scatter_dim)]
-    output_list = [torch.empty_like(input_list[0]) for _ in range(sp_size)]
-    dist.all_to_all(output_list, input_list, group=group)
-    return torch.cat(output_list, dim=gather_dim).contiguous()
-
 def _single_all_to_all(
     input_: torch.Tensor,
     scatter_dim: int,
     gather_dim: int,
     enable_HCCL=False,
 ):
-
+    assert scatter_dim == 0 or scatter_dim == 1, "scatter_dim should be 0 or 1"
     sp_size = nccl_info.world_size
     inp_shape = list(input_.shape)
     inp_shape[scatter_dim] = inp_shape[scatter_dim] // sp_size
@@ -85,6 +73,8 @@ class _AllToAll(torch.autograd.Function):
             None,
         )
 
+
+
 def all_to_all_SBH(
     input_: torch.Tensor,
     scatter_dim: int = 1,
@@ -92,15 +82,60 @@ def all_to_all_SBH(
 ):
     return _AllToAll.apply(input_, scatter_dim, gather_dim, _single_all_to_all)
 
-def all_to_all_BSND(
-    input_: torch.Tensor,
-    scatter_dim: int = 2,
-    gather_dim: int = 1,
-):
-    return _AllToAll.apply(input_, scatter_dim, gather_dim, _all_to_all)
 
 
-def prepare_parallel_data(hidden_states, encoder_hidden_states, attention_mask, encoder_attention_mask, use_image_num):
+class _AllGather(torch.autograd.Function):
+    """All-gather communication with autograd support.
+
+    Args:
+        input_: input tensor
+        dim: dimension along which to concatenate
+    """
+
+    @staticmethod
+    def forward(ctx, input_, dim):
+        ctx.dim = dim
+        world_size = nccl_info.world_size
+        group = nccl_info.group
+        input_size = list(input_.size())
+
+        ctx.input_size = input_size[dim]
+
+        tensor_list = [torch.empty_like(input_) for _ in range(world_size)]
+        input_ = input_.contiguous()
+        dist.all_gather(tensor_list, input_, group=group)
+
+        output = torch.cat(tensor_list, dim=dim)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        world_size = nccl_info.world_size
+        rank = nccl_info.rank
+        dim = ctx.dim
+        input_size = ctx.input_size
+
+        sizes = [input_size] * world_size
+
+        grad_input_list = torch.split(grad_output, sizes, dim=dim)
+        grad_input = grad_input_list[rank]
+
+        return grad_input, None
+
+def all_gather_BHSD(input_: torch.Tensor, dim: int = 1):
+    """Performs an all-gather operation on the input tensor along the specified dimension.
+
+    Args:
+        input_ (torch.Tensor): Input tensor of shape [B, H, S, D].
+        dim (int, optional): Dimension along which to concatenate. Defaults to 1.
+
+    Returns:
+        torch.Tensor: Output tensor after all-gather operation, concatenated along 'dim'.
+    """
+    return _AllGather.apply(input_, dim)
+
+
+def prepare_parallel_data(hidden_states, encoder_hidden_states, attention_mask, encoder_attention_mask):
     def all_to_all(hidden_states, encoder_hidden_states, attention_mask, encoder_attention_mask):
         hidden_states = _single_all_to_all(hidden_states, scatter_dim=2, gather_dim=0, enable_HCCL=True)
         encoder_hidden_states = _single_all_to_all(encoder_hidden_states, scatter_dim=1, gather_dim=0, enable_HCCL=True)
@@ -119,4 +154,4 @@ def prepare_parallel_data(hidden_states, encoder_hidden_states, attention_mask,
                                                                                             attention_mask.repeat(1, sp_size, 1, 1),
                                                                                             encoder_attention_mask.repeat(1, sp_size, 1))
 
-    return hidden_states, encoder_hidden_states, attention_mask, encoder_attention_mask, use_image_num
+    return hidden_states, encoder_hidden_states, attention_mask, encoder_attention_mask

fastvideo/utils/encode_vae_t5.py

@@ -14,15 +14,13 @@ import os
 import shutil
 from pathlib import Path
 import gc
-import numpy as np
+import wandb
 from einops import rearrange
 from tqdm import tqdm
 
-from opensora.adaptor.modules import replace_with_fp32_forwards
-
-from opensora.utils.parallel_states import initialize_sequence_parallel_state, \
+from fastvideo.utils.parallel_states import initialize_sequence_parallel_state, \
     destroy_sequence_parallel_group, get_sequence_parallel_state, set_sequence_parallel_state
-from opensora.utils.communications import prepare_parallel_data, broadcast
+from fastvideo.utils.communications import prepare_parallel_data, broadcast
 import time
 from torch.utils.data import DataLoader
 from copy import deepcopy
@@ -37,19 +35,19 @@ from packaging import version
 from tqdm.auto import tqdm
 
 import diffusers
-from diffusers import DDPMScheduler, PNDMScheduler, DPMSolverMultistepScheduler
+from diffusers import (
+    AutoencoderKL,
+    FlowMatchEulerDiscreteScheduler,
+    MochiTransformer3DModel,
+    MochiPipeline,
+)
 from diffusers.optimization import get_scheduler
 from diffusers.training_utils import compute_snr
 from diffusers.utils import check_min_version, is_wandb_available
 
-from opensora.utils.ema import EMAModel
-from opensora.dataset import getdataset
-from opensora.models import CausalVAEModelWrapper
-from opensora.models.text_encoder import  get_text_warpper
-from opensora.models.causalvideovae import ae_stride_config, ae_channel_config
-from opensora.models.diffusion import Diffusion_models, Diffusion_models_class
-from opensora.utils.dataset_utils import Collate, LengthGroupedSampler
-from opensora.sample.pipeline_opensora import OpenSoraPipeline
+from fastvideo.utils.ema import EMAModel
+from fastvideo.dataset import getdataset
+from fastvideo.utils.dataset_utils import Collate, LengthGroupedSampler
 
 
 # Will error if the minimal version of diffusers is not installed. Remove at your own risks.
@@ -59,42 +57,24 @@ logger = get_logger(__name__)
 
 @torch.inference_mode()
 def log_validation(args, model, vae, text_encoder, tokenizer, accelerator, weight_dtype, global_step, ema=False):
-    positive_prompt = "(masterpiece), (best quality), (ultra-detailed), {}. emotional, harmonious, vignette, 4k epic detailed, shot on kodak, 35mm photo, sharp focus, high budget, cinemascope, moody, epic, gorgeous"
-    negative_prompt = """nsfw, lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry, 
-                        """
-    validation_prompt = [
-        "A stylish woman walks down a Tokyo street filled with warm glowing neon and animated city signage. She wears a black leather jacket, a long red dress, and black boots, and carries a black purse. She wears sunglasses and red lipstick. She walks confidently and casually. The street is damp and reflective, creating a mirror effect of the colorful lights. Many pedestrians walk about.",
-        "A serene underwater scene featuring a sea turtle swimming through a coral reef. The turtle, with its greenish-brown shell, is the main focus of the video, swimming gracefully towards the right side of the frame. The coral reef, teeming with life, is visible in the background, providing a vibrant and colorful backdrop to the turtle's journey. Several small fish, darting around the turtle, add a sense of movement and dynamism to the scene."
-        ]
-    if 'mt5' in args.text_encoder_name:
-        validation_prompt_cn = [
-            "一只戴着墨镜在泳池当救生员的猫咪。",
-            "这是一个宁静的水下场景，一只海龟游过珊瑚礁。海龟带着绿褐色的龟壳，优雅地游向画面右侧，成为视频的焦点。背景中的珊瑚礁生机盎然，为海龟的旅程提供了生动多彩的背景。几条小鱼在海龟周围穿梭，为画面增添了动感和活力。"
-            ]
-        validation_prompt += validation_prompt_cn
     logger.info(f"Running validation....\n")
     model = accelerator.unwrap_model(model)
-    scheduler = DPMSolverMultistepScheduler()
-    opensora_pipeline = OpenSoraPipeline(vae=vae,
+    scheduler = FlowMatchEulerDiscreteScheduler()
+    mochi_pipeline = MochiPipeline(vae=vae,
                                          text_encoder=text_encoder,
                                          tokenizer=tokenizer,
                                          scheduler=scheduler,
                                          transformer=model).to(device=accelerator.device)
     videos = []
-    for prompt in validation_prompt:
+    for prompt in args.validaiton_prompt:
         logger.info('Processing the ({}) prompt'.format(prompt))
-        video = opensora_pipeline(
-                                positive_prompt.format(prompt),
-                                negative_prompt=negative_prompt, 
+        video = mochi_pipeline(
+                                args.prompt,
                                 num_frames=args.num_frames,
                                 height=args.max_height,
                                 width=args.max_width,
-                                num_inference_steps=args.num_sampling_steps,
+                                num_inference_steps=args.validation_sampling_steps,
                                 guidance_scale=args.guidance_scale,
-                                enable_temporal_attentions=True,
-                                num_images_per_prompt=1,
-                                mask_feature=True,
-                                max_sequence_length=args.model_max_length,
                                 ).images
         videos.append(video[0])
     # import ipdb;ipdb.set_trace()
@@ -103,35 +83,24 @@ def log_validation(args, model, vae, text_encoder, tokenizer, accelerator, weigh
     videos = torch.stack(videos).numpy()
     videos = rearrange(videos, 'b t h w c -> b t c h w')
     for tracker in accelerator.trackers:
-        if tracker.name == "tensorboard":
-            if videos.shape[1] == 1:
-                assert args.num_frames == 1
-                images = rearrange(videos, 'b 1 c h w -> (b 1) h w c')
-                np_images = np.stack([np.asarray(img) for img in images])
-                tracker.writer.add_images(f"{'ema_' if ema else ''}validation", np_images, global_step, dataformats="NHWC")
-            else:
-                np_videos = np.stack([np.asarray(vid) for vid in videos])
-                tracker.writer.add_video(f"{'ema_' if ema else ''}validation", np_videos, global_step, fps=24)
-        if tracker.name == "wandb":
-            import wandb
-            if videos.shape[1] == 1:
-                images = rearrange(videos, 'b 1 c h w -> (b 1) h w c')
-                logs = {
-                    f"{'ema_' if ema else ''}validation": [
-                        wandb.Image(image, caption=f"{i}: {prompt}")
-                        for i, (image, prompt) in enumerate(zip(images, validation_prompt))
-                    ]
-                }
-            else:
-                logs = {
-                    f"{'ema_' if ema else ''}validation": [
-                        wandb.Video(video, caption=f"{i}: {prompt}", fps=24)
-                        for i, (video, prompt) in enumerate(zip(videos, validation_prompt))
-                    ]
-                }
-            tracker.log(logs, step=global_step)
+        if videos.shape[1] == 1:
+            images = rearrange(videos, 'b 1 c h w -> (b 1) h w c')
+            logs = {
+                f"{'ema_' if ema else ''}validation": [
+                    wandb.Image(image, caption=f"{i}: {prompt}")
+                    for i, (image, prompt) in enumerate(zip(images, args.validaiton_prompt))
+                ]
+            }
+        else:
+            logs = {
+                f"{'ema_' if ema else ''}validation": [
+                    wandb.Video(video, caption=f"{i}: {prompt}", fps=24)
+                    for i, (video, prompt) in enumerate(zip(videos, args.validaiton_prompt))
+                ]
+            }
+        tracker.log(logs, step=global_step)
 
-    del opensora_pipeline
+    del mochi_pipeline
     gc.collect()
     torch.cuda.empty_cache()
 
@@ -150,6 +119,7 @@ def main(args):
     logging_dir = Path(args.output_dir, args.logging_dir)
 
     # use LayerNorm, GeLu, SiLu always as fp32 mode
+    # TODO: 
     if args.enable_stable_fp32:
         replace_with_fp32_forwards()
 
@@ -160,6 +130,7 @@ def main(args):
         mixed_precision=args.mixed_precision,
         log_with=args.report_to,
         project_config=accelerator_project_config,
+        log_with="wandb"
     )
 
     if args.num_frames != 1 and args.use_image_num == 0:
@@ -205,7 +176,7 @@ def main(args):
     ae = CausalVAEModelWrapper(args.ae_path, cache_dir=args.cache_dir, **kwargs).eval()
     if args.enable_tiling:
         ae.vae.enable_tiling()
-        ae.vae.tile_overlap_factor = args.tile_overlap_factor
+        ae.vae.tile_sample_stride = args.tile_sample_stride
 
     kwargs = {'load_in_8bit': args.enable_8bit_t5, 'torch_dtype': weight_dtype, 'low_cpu_mem_usage': True}
     text_enc = get_text_warpper(args.text_encoder_name)(args, **kwargs).eval()
@@ -258,7 +229,7 @@ def main(args):
         interpolation_scale_t=args.interpolation_scale_t,
         downsampler=args.downsampler,
         use_rope=args.use_rope,
-        # model_max_length=args.model_max_length,
+        # text_max_length=args.text_max_length,
         use_stable_fp32=args.enable_stable_fp32, 
     )
     model.gradient_checkpointing = args.gradient_checkpointing
@@ -707,27 +678,11 @@ def main(args):
 
         with torch.no_grad():
             B, N, L = input_ids.shape  # B 1+num_images L
-            # use batch inference
             input_ids_ = input_ids.reshape(-1, L)
             cond_mask_ = cond_mask.reshape(-1, L)
             cond = text_enc(input_ids_, cond_mask_)  # B 1+num_images L D
             cond = cond.reshape(B, N, L, -1)
-            # Map input images to latent space + normalize latents
             x = ae.encode(x)  # B C T H W
-
-            # def custom_to_video(x: torch.Tensor, fps: float = 2.0, output_file: str = 'output_video.mp4') -> None:
-            #     from examples.rec_video import array_to_video
-            #     x = x.detach().cpu()
-            #     x = torch.clamp(x, -1, 1)
-            #     x = (x + 1) / 2
-            #     x = x.permute(1, 2, 3, 0).numpy()
-            #     x = (255*x).astype(np.uint8)
-            #     array_to_video(x, fps=fps, output_file=output_file)
-            #     return
-            # videos = ae.decode(x)[0]
-            # videos = videos.transpose(0, 1)
-            # custom_to_video(videos.to(torch.float32), fps=24, output_file='tmp.mp4')
-            # import sys;sys.exit()
         current_step_frame = x.shape[2]
         current_step_sp_state = get_sequence_parallel_state()
         if args.sp_size != 1:  # enable sp
@@ -786,6 +741,7 @@ if __name__ == "__main__":
     # dataset & dataloader
     parser.add_argument("--dataset", type=str, required=True)
     parser.add_argument("--data", type=str, required='')
+    parser.add_argument("--cache_vae_latent", action="store_true")
     parser.add_argument("--train_fps", type=int, default=24)
     parser.add_argument("--drop_short_ratio", type=float, default=1.0)
     parser.add_argument("--speed_factor", type=float, default=1.0)
@@ -794,7 +750,7 @@ if __name__ == "__main__":
     parser.add_argument("--max_height", type=int, default=320)
     parser.add_argument("--max_width", type=int, default=240)
     parser.add_argument("--use_image_num", type=int, default=0)
-    parser.add_argument("--model_max_length", type=int, default=512)
+    parser.add_argument("--text_max_length", type=int, default=512)
     parser.add_argument('--cfg', type=float, default=0.1)
     parser.add_argument("--dataloader_num_workers", type=int, default=10, help="Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process.")
     parser.add_argument("--train_batch_size", type=int, default=16, help="Batch size (per device) for the training dataloader.")
@@ -804,7 +760,7 @@ if __name__ == "__main__":
     # text encoder & vae & diffusion model
     parser.add_argument("--model", type=str, choices=list(Diffusion_models.keys()), default="Latte-XL/122")
     parser.add_argument('--enable_8bit_t5', action='store_true')
-    parser.add_argument('--tile_overlap_factor', type=float, default=0.125)
+    parser.add_argument('--tile_sample_stride', type=float, default=0.125)
     parser.add_argument('--enable_tiling', action='store_true')
     parser.add_argument("--attention_mode", type=str, choices=['xformers', 'math', 'flash'], default="xformers")
     parser.add_argument('--use_rope', action='store_true')
@@ -829,7 +785,8 @@ if __name__ == "__main__":
     parser.add_argument("--prediction_type", type=str, default=None, help="The prediction_type that shall be used for training. Choose between 'epsilon' or 'v_prediction' or leave `None`. If left to `None` the default prediction type of the scheduler: `noise_scheduler.config.prediciton_type` is chosen.")
 
     # validation & logs
-    parser.add_argument("--num_sampling_steps", type=int, default=20)
+    parser.add_argument("--validaiton_prompt", nargs='+', help="List of prompts to use for validation.")
+    parser.add_argument("--validation_sampling_steps", type=int, default=64)
     parser.add_argument('--guidance_scale', type=float, default=4.5)
     parser.add_argument("--enable_tracker", action="store_true")
     parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.")

fastvideo/utils/optimizer.py

@@ -0,0 +1,71 @@
+from accelerate.logging import get_logger
+import torch
+logger = get_logger(__name__)
+
+def get_optimizer(args, params_to_optimize, use_deepspeed: bool = False):
+    # Optimizer creation
+    supported_optimizers = ["adam", "adamw", "prodigy"]
+    if args.optimizer not in supported_optimizers:
+        logger.warning(
+            f"Unsupported choice of optimizer: {args.optimizer}. Supported optimizers include {supported_optimizers}. Defaulting to AdamW"
+        )
+        args.optimizer = "adamw"
+
+    if args.use_8bit_adam and not (args.optimizer.lower() not in ["adam", "adamw"]):
+        logger.warning(
+            f"use_8bit_adam is ignored when optimizer is not set to 'Adam' or 'AdamW'. Optimizer was "
+            f"set to {args.optimizer.lower()}"
+        )
+
+    if args.use_8bit_adam:
+        try:
+            import bitsandbytes as bnb
+        except ImportError:
+            raise ImportError(
+                "To use 8-bit Adam, please install the bitsandbytes library: `pip install bitsandbytes`."
+            )
+
+    if args.optimizer.lower() == "adamw":
+        optimizer_class = bnb.optim.AdamW8bit if args.use_8bit_adam else torch.optim.AdamW
+
+        optimizer = optimizer_class(
+            params_to_optimize,
+            betas=(args.adam_beta1, args.adam_beta2),
+            eps=args.adam_epsilon,
+            weight_decay=args.adam_weight_decay,
+        )
+    elif args.optimizer.lower() == "adam":
+        optimizer_class = bnb.optim.Adam8bit if args.use_8bit_adam else torch.optim.Adam
+
+        optimizer = optimizer_class(
+            params_to_optimize,
+            betas=(args.adam_beta1, args.adam_beta2),
+            eps=args.adam_epsilon,
+            weight_decay=args.adam_weight_decay,
+        )
+    elif args.optimizer.lower() == "prodigy":
+        try:
+            import prodigyopt
+        except ImportError:
+            raise ImportError("To use Prodigy, please install the prodigyopt library: `pip install prodigyopt`")
+
+        optimizer_class = prodigyopt.Prodigy
+
+        if args.learning_rate <= 0.1:
+            logger.warning(
+                "Learning rate is too low. When using prodigy, it's generally better to set learning rate around 1.0"
+            )
+
+        optimizer = optimizer_class(
+            params_to_optimize,
+            lr=args.learning_rate,
+            betas=(args.adam_beta1, args.adam_beta2),
+            beta3=args.prodigy_beta3,
+            weight_decay=args.adam_weight_decay,
+            eps=args.adam_epsilon,
+            decouple=args.prodigy_decouple,
+            use_bias_correction=args.prodigy_use_bias_correction,
+            safeguard_warmup=args.prodigy_safeguard_warmup,
+        )
+
+    return optimizer

fastvideo/utils/parallel_states.py

@@ -5,8 +5,8 @@ import os
 class COMM_INFO:
     def __init__(self):
         self.group = None
-        self.world_size = 0
-        self.rank = -1
+        self.world_size = 1
+        self.rank = 0
 
 nccl_info = COMM_INFO()
 _SEQUENCE_PARALLEL_STATE = False
@@ -28,7 +28,6 @@ def initialize_sequence_parallel_group(sequence_parallel_size):
     rank = int(os.getenv('RANK', '0'))
     world_size = int(os.getenv("WORLD_SIZE", '1'))
     assert world_size % sequence_parallel_size == 0, "world_size must be divisible by sequence_parallel_size, but got world_size: {}, sequence_parallel_size: {}".format(world_size, sequence_parallel_size)
-    # hccl
     nccl_info.world_size = sequence_parallel_size
     nccl_info.rank = rank
     num_sequence_parallel_groups: int = world_size // sequence_parallel_size

generate_synthetic.sh

@@ -0,0 +1,11 @@
+python fastvideo/sample/generate_synthetic.py \
+    --model_path data/mochi \
+    --num_frames 163 \
+    --height 480 \
+    --width 848 \
+    --num_inference_steps 64 \
+    --guidance_scale 4.5 \
+    --prompt_path data/dummyVid/videos2caption_5video.json \
+    --dataset_output_dir data/Mochi-Synthetic-Data
+
+    

infenrece_mochi.sh

@@ -0,0 +1,10 @@
+python fastvideo/sample/sample_t2v_mochi.py \
+    --model_path data/mochi \
+    --prompt "A hand with delicate fingers picks up a bright yellow lemon from a wooden bowl filled with lemons and sprigs of mint against a peach-colored background. The hand gently tosses the lemon up and catches it, showcasing its smooth texture. A beige string bag sits beside the bowl, adding a rustic touch to the scene. Additional lemons, one halved, are scattered around the base of the bowl. The even lighting enhances the vibrant colors and creates a fresh, inviting atmosphere." \
+    --num_frames 163 \
+    --height 480 \
+    --width 848 \
+    --num_inference_steps 64 \
+    --guidance_scale 4.5 \
+    --seed 12345 \
+    --output_path outputs_2.mp4

infenrece_mochi_sp.sh

@@ -0,0 +1,14 @@
+#!/bin/bash
+
+num_gpus=4
+
+torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29503 \
+    fastvideo/sample/sample_t2v_mochi_sp.py \
+    --model_path data/mochi \
+    --prompt "A hand with delicate fingers picks up a bright yellow lemon from a wooden bowl filled with lemons and sprigs of mint against a peach-colored background. The hand gently tosses the lemon up and catches it, showcasing its smooth texture. A beige string bag sits beside the bowl, adding a rustic touch to the scene. Additional lemons, one halved, are scattered around the base of the bowl. The even lighting enhances the vibrant colors and creates a fresh, inviting atmosphere." \
+    --num_frames 163 \
+    --height 480 \
+    --width 848 \
+    --num_inference_steps 64 \
+    --guidance_scale 4.5 \
+    --seed 12345

opensora/__init__.py

@@ -1 +0,0 @@
-# 

opensora/adaptor/modules.py

@@ -1,23 +0,0 @@
-import torch
-from torch import nn
-from torch.nn import functional as F
-
-
-def fp32_layer_norm_forward(self, inputs: torch.Tensor) -> torch.Tensor:
-    origin_dtype = inputs.dtype
-    return F.layer_norm(inputs.float(), self.normalized_shape, self.weight.float() if self.weight is not None else None,
-                        self.bias.float() if self.bias is not None else None, self.eps).to(origin_dtype)
-
-
-def fp32_silu_forward(self, inputs: torch.Tensor) -> torch.Tensor:
-    return torch.nn.functional.silu(inputs.float(), inplace=self.inplace).to(inputs.dtype)
-
-
-def fp32_gelu_forward(self, inputs: torch.Tensor) -> torch.Tensor:
-    return torch.nn.functional.gelu(inputs.float(), approximate=self.approximate).to(inputs.dtype)
-
-
-def replace_with_fp32_forwards():
-    nn.GELU.forward = fp32_gelu_forward
-    nn.SiLU.forward = fp32_silu_forward
-    nn.LayerNorm.forward = fp32_layer_norm_forward

opensora/dataset/inpaint_datasets.py

@@ -1,116 +0,0 @@
-
-import os, random
-import numpy as np
-
-import torch
-from accelerate.logging import get_logger
-
-from opensora.utils.utils import text_preprocessing
-
-from opensora.dataset.t2v_datasets import SingletonMeta, DataSetProg
-from opensora.dataset.t2v_datasets import T2V_dataset
-
-logger = get_logger(__name__)
-
-dataset_prog = DataSetProg()
-
-
-class Meta_dataset(T2V_dataset):
-    def __init__(self, args, transform, temporal_sample, tokenizer, transform_topcrop):
-        super().__init__(args, transform, temporal_sample, tokenizer, transform_topcrop)
-
-        if self.num_frames != 1:
-            # inpaint
-            # The proportion of executing the i2v task.
-            self.i2v_ratio = args.i2v_ratio
-            self.transition_ratio = args.transition_ratio
-            self.v2v_ratio = args.v2v_ratio
-            self.clear_video_ratio = args.clear_video_ratio
-            assert self.i2v_ratio + self.transition_ratio + self.v2v_ratio + self.clear_video_ratio < 1, 'The sum of i2v_ratio, transition_ratio, v2v_ratio and clear video ratio should be less than 1.'
-        
-        self.default_text_ratio = args.default_text_ratio
-        self.default_text = f"The {'video' if self.num_frames != 1 else 'image'} showcases a scene with coherent and clear visuals."
-    
-    def get_mask_masked_video(self, video):
-        # video shape (T, C, H, W)
-        # 1 means masked, 0 means not masked
-        t, c, h, w = video.shape
-        mask = torch.ones_like(video, device=video.device, dtype=video.dtype)
-        
-        rand_num = random.random()
-        # i2v
-        if rand_num < self.i2v_ratio:
-            mask[0] = 0
-        # transition
-        elif rand_num < self.i2v_ratio + self.transition_ratio:
-            mask[0] = 0
-            mask[-1] = 0
-        # video continuation
-        elif rand_num < self.i2v_ratio + self.transition_ratio + self.v2v_ratio:
-            end_idx = random.randint(1, t)
-            mask[:end_idx] = 0
-        # clear video
-        elif rand_num < self.i2v_ratio + self.transition_ratio + self.v2v_ratio + self.clear_video_ratio:
-            mask[:] = 0
-        # random mask
-        else:
-            idx_to_select = random.randint(0, t - 1)
-            selected_indices = random.sample(range(0, t), idx_to_select)
-            mask[selected_indices] = 0
-        masked_video = video * (mask < 0.5)
-
-        # save_video(masked_video.permute(0, 2, 3, 1).cpu().numpy(), 'masked_video.mp4')
-        return dict(mask=mask, masked_video=masked_video)
-
-    def drop(self, text):
-        rand_num = random.random()
-        rand_num_text = random.random()
-
-        if rand_num < self.cfg:
-            text = self.default_text if rand_num_text < self.default_text_ratio else ''
-
-        return dict(text=text)
-
-    
-class Inpaint_dataset(Meta_dataset):
-    def get_video(self, idx):
-        video_path = dataset_prog.cap_list[idx]['path']
-        assert os.path.exists(video_path), f"file {video_path} do not exist!"
-        frame_indice = dataset_prog.cap_list[idx]['sample_frame_index']
-        video = self.decord_read(video_path, predefine_num_frames=len(frame_indice))
-
-        h, w = video.shape[-2:]
-        assert h / w <= 17 / 16 and h / w >= 8 / 16, f'Only videos with a ratio (h/w) less than 17/16 and more than 8/16 are supported. But video ({video_path}) found ratio is {round(h / w, 2)} with the shape of {video.shape}'
-        t = video.shape[0]
-        video = self.transform(video)  # T C H W -> T C H W
-
-        # inpaint
-        inpaint_cond_data = self.get_mask_masked_video(video)
-        mask, masked_video = inpaint_cond_data['mask'], inpaint_cond_data['masked_video']
-        video = torch.cat([video, masked_video, mask], dim=1) # T 3*C H W
-
-        # video = torch.rand(221, 3, 480, 640)
-
-        video = video.transpose(0, 1)  # T C H W -> C T H W
-        text = dataset_prog.cap_list[idx]['cap']
-        if not isinstance(text, list):
-            text = [text]
-        text = [random.choice(text)]
-
-        text = text_preprocessing(text, support_Chinese=self.support_Chinese)
-
-        text = self.drop(text)['text']
-
-        text_tokens_and_mask = self.tokenizer(
-            text,
-            max_length=self.model_max_length,
-            padding='max_length',
-            truncation=True,
-            return_attention_mask=True,
-            add_special_tokens=True,
-            return_tensors='pt'
-        )
-        input_ids = text_tokens_and_mask['input_ids']
-        cond_mask = text_tokens_and_mask['attention_mask']
-        return dict(pixel_values=video, input_ids=input_ids, cond_mask=cond_mask)
-    

opensora/eval/cal_flolpips.py

@@ -1,83 +0,0 @@
-import numpy as np
-import torch
-from tqdm import tqdm
-import math
-from einops import rearrange
-import sys
-sys.path.append(".")
-from opensora.eval.flolpips.pwcnet import Network as PWCNet
-from opensora.eval.flolpips.flolpips import FloLPIPS
-
-loss_fn = FloLPIPS(net='alex', version='0.1').eval().requires_grad_(False)
-flownet = PWCNet().eval().requires_grad_(False)
-
-def trans(x):
-    return x
-
-
-def calculate_flolpips(videos1, videos2, device):
-    global loss_fn, flownet
-    
-    print("calculate_flowlpips...")
-    loss_fn = loss_fn.to(device)
-    flownet = flownet.to(device)
-    
-    if videos1.shape != videos2.shape:
-        print("Warning: the shape of videos are not equal.")
-        min_frames = min(videos1.shape[1], videos2.shape[1])
-        videos1 = videos1[:, :min_frames]
-        videos2 = videos2[:, :min_frames]
-        
-    videos1 = trans(videos1)
-    videos2 = trans(videos2)
-
-    flolpips_results = []
-    for video_num in tqdm(range(videos1.shape[0])):
-        video1 = videos1[video_num].to(device)
-        video2 = videos2[video_num].to(device)
-        frames_rec = video1[:-1]
-        frames_rec_next = video1[1:]
-        frames_gt = video2[:-1]
-        frames_gt_next = video2[1:]
-        t, c, h, w = frames_gt.shape
-        flow_gt = flownet(frames_gt, frames_gt_next)
-        flow_dis = flownet(frames_rec, frames_rec_next)
-        flow_diff = flow_gt - flow_dis
-        flolpips = loss_fn.forward(frames_gt, frames_rec, flow_diff, normalize=True)
-        flolpips_results.append(flolpips.cpu().numpy().tolist())
-        
-    flolpips_results = np.array(flolpips_results) # [batch_size, num_frames]
-    flolpips = {}
-    flolpips_std = {}
-
-    for clip_timestamp in range(flolpips_results.shape[1]):
-        flolpips[clip_timestamp] = np.mean(flolpips_results[:,clip_timestamp], axis=-1)
-        flolpips_std[clip_timestamp] = np.std(flolpips_results[:,clip_timestamp], axis=-1)
-
-    result = {
-        "value": flolpips,
-        "value_std": flolpips_std,
-        "video_setting": video1.shape,
-        "video_setting_name": "time, channel, heigth, width",
-        "result": flolpips_results,
-        "details": flolpips_results.tolist()
-    }
-
-    return result
-
-# test code / using example
-
-def main():
-    NUMBER_OF_VIDEOS = 8
-    VIDEO_LENGTH = 50
-    CHANNEL = 3
-    SIZE = 64
-    videos1 = torch.zeros(NUMBER_OF_VIDEOS, VIDEO_LENGTH, CHANNEL, SIZE, SIZE, requires_grad=False)
-    videos2 = torch.zeros(NUMBER_OF_VIDEOS, VIDEO_LENGTH, CHANNEL, SIZE, SIZE, requires_grad=False)
-
-    import json
-    result = calculate_flolpips(videos1, videos2, "cuda:0")
-    print(json.dumps(result, indent=4))
-
-if __name__ == "__main__":
-    main()

opensora/eval/cal_fvd.py

@@ -1,85 +0,0 @@
-import numpy as np
-import torch
-from tqdm import tqdm
-
-def trans(x):
-    # if greyscale images add channel
-    if x.shape[-3] == 1:
-        x = x.repeat(1, 1, 3, 1, 1)
-
-    # permute BTCHW -> BCTHW
-    x = x.permute(0, 2, 1, 3, 4) 
-
-    return x
-
-def calculate_fvd(videos1, videos2, device, method='styleganv'):
-
-    if method == 'styleganv':
-        from fvd.styleganv.fvd import get_fvd_feats, frechet_distance, load_i3d_pretrained
-    elif method == 'videogpt':
-        from fvd.videogpt.fvd import load_i3d_pretrained
-        from fvd.videogpt.fvd import get_fvd_logits as get_fvd_feats
-        from fvd.videogpt.fvd import frechet_distance
-
-    print("calculate_fvd...")
-
-    # videos [batch_size, timestamps, channel, h, w]
-    
-    assert videos1.shape == videos2.shape
-
-    i3d = load_i3d_pretrained(device=device)
-    fvd_results = []
-
-    # support grayscale input, if grayscale -> channel*3
-    # BTCHW -> BCTHW
-    # videos -> [batch_size, channel, timestamps, h, w]
-
-    videos1 = trans(videos1)
-    videos2 = trans(videos2)
-
-    fvd_results = {}
-
-    # for calculate FVD, each clip_timestamp must >= 10
-    for clip_timestamp in tqdm(range(10, videos1.shape[-3]+1)):
-       
-        # get a video clip
-        # videos_clip [batch_size, channel, timestamps[:clip], h, w]
-        videos_clip1 = videos1[:, :, : clip_timestamp]
-        videos_clip2 = videos2[:, :, : clip_timestamp]
-
-        # get FVD features
-        feats1 = get_fvd_feats(videos_clip1, i3d=i3d, device=device)
-        feats2 = get_fvd_feats(videos_clip2, i3d=i3d, device=device)
-      
-        # calculate FVD when timestamps[:clip]
-        fvd_results[clip_timestamp] = frechet_distance(feats1, feats2)
-
-    result = {
-        "value": fvd_results,
-        "video_setting": videos1.shape,
-        "video_setting_name": "batch_size, channel, time, heigth, width",
-    }
-
-    return result
-
-# test code / using example
-
-def main():
-    NUMBER_OF_VIDEOS = 8
-    VIDEO_LENGTH = 50
-    CHANNEL = 3
-    SIZE = 64
-    videos1 = torch.zeros(NUMBER_OF_VIDEOS, VIDEO_LENGTH, CHANNEL, SIZE, SIZE, requires_grad=False)
-    videos2 = torch.ones(NUMBER_OF_VIDEOS, VIDEO_LENGTH, CHANNEL, SIZE, SIZE, requires_grad=False)
-    device = torch.device("cuda")
-    # device = torch.device("cpu")
-
-    import json
-    result = calculate_fvd(videos1, videos2, device, method='videogpt')
-    print(json.dumps(result, indent=4))
-
-    result = calculate_fvd(videos1, videos2, device, method='styleganv')
-    print(json.dumps(result, indent=4))
-
-if __name__ == "__main__":
-    main()

opensora/eval/cal_lpips.py

@@ -1,97 +0,0 @@
-import numpy as np
-import torch
-from tqdm import tqdm
-import math
-
-import torch
-import lpips
-
-spatial = True         # Return a spatial map of perceptual distance.
-
-# Linearly calibrated models (LPIPS)
-loss_fn = lpips.LPIPS(net='alex', spatial=spatial) # Can also set net = 'squeeze' or 'vgg'
-# loss_fn = lpips.LPIPS(net='alex', spatial=spatial, lpips=False) # Can also set net = 'squeeze' or 'vgg'
-
-def trans(x):
-    # if greyscale images add channel
-    if x.shape[-3] == 1:
-        x = x.repeat(1, 1, 3, 1, 1)
-
-    # value range [0, 1] -> [-1, 1]
-    x = x * 2 - 1
-
-    return x
-
-def calculate_lpips(videos1, videos2, device):
-    # image should be RGB, IMPORTANT: normalized to [-1,1]
-    print("calculate_lpips...")
-
-    assert videos1.shape == videos2.shape
-
-    # videos [batch_size, timestamps, channel, h, w]
-
-    # support grayscale input, if grayscale -> channel*3
-    # value range [0, 1] -> [-1, 1]
-    videos1 = trans(videos1)
-    videos2 = trans(videos2)
-
-    lpips_results = []
-
-    for video_num in tqdm(range(videos1.shape[0])):
-        # get a video
-        # video [timestamps, channel, h, w]
-        video1 = videos1[video_num]
-        video2 = videos2[video_num]
-
-        lpips_results_of_a_video = []
-        for clip_timestamp in range(len(video1)):
-            # get a img
-            # img [timestamps[x], channel, h, w]
-            # img [channel, h, w] tensor
-
-            img1 = video1[clip_timestamp].unsqueeze(0).to(device)
-            img2 = video2[clip_timestamp].unsqueeze(0).to(device)
-            
-            loss_fn.to(device)
-
-            # calculate lpips of a video
-            lpips_results_of_a_video.append(loss_fn.forward(img1, img2).mean().detach().cpu().tolist())
-        lpips_results.append(lpips_results_of_a_video)
-    
-    lpips_results = np.array(lpips_results)
-    
-    lpips = {}
-    lpips_std = {}
-
-    for clip_timestamp in range(len(video1)):
-        lpips[clip_timestamp] = np.mean(lpips_results[:,clip_timestamp])
-        lpips_std[clip_timestamp] = np.std(lpips_results[:,clip_timestamp])
-
-
-    result = {
-        "value": lpips,
-        "value_std": lpips_std,
-        "video_setting": video1.shape,
-        "video_setting_name": "time, channel, heigth, width",
-    }
-
-    return result
-
-# test code / using example
-
-def main():
-    NUMBER_OF_VIDEOS = 8
-    VIDEO_LENGTH = 50
-    CHANNEL = 3
-    SIZE = 64
-    videos1 = torch.zeros(NUMBER_OF_VIDEOS, VIDEO_LENGTH, CHANNEL, SIZE, SIZE, requires_grad=False)
-    videos2 = torch.ones(NUMBER_OF_VIDEOS, VIDEO_LENGTH, CHANNEL, SIZE, SIZE, requires_grad=False)
-    device = torch.device("cuda")
-    # device = torch.device("cpu")
-
-    import json
-    result = calculate_lpips(videos1, videos2, device)
-    print(json.dumps(result, indent=4))
-
-if __name__ == "__main__":
-    main()

opensora/eval/cal_psnr.py

@@ -1,84 +0,0 @@
-import numpy as np
-import torch
-from tqdm import tqdm
-import math
-
-def img_psnr(img1, img2):
-    # [0,1]
-    # compute mse
-    # mse = np.mean((img1-img2)**2)
-    mse = np.mean((img1 / 1.0 - img2 / 1.0) ** 2)
-    # compute psnr
-    if mse < 1e-10:
-        return 100
-    psnr = 20 * math.log10(1 / math.sqrt(mse))
-    return psnr
-
-def trans(x):
-    return x
-
-def calculate_psnr(videos1, videos2):
-    print("calculate_psnr...")
-
-    # videos [batch_size, timestamps, channel, h, w]
-    
-    assert videos1.shape == videos2.shape
-
-    videos1 = trans(videos1)
-    videos2 = trans(videos2)
-
-    psnr_results = []
-    
-    for video_num in tqdm(range(videos1.shape[0])):
-        # get a video
-        # video [timestamps, channel, h, w]
-        video1 = videos1[video_num]
-        video2 = videos2[video_num]
-
-        psnr_results_of_a_video = []
-        for clip_timestamp in range(len(video1)):
-            # get a img
-            # img [timestamps[x], channel, h, w]
-            # img [channel, h, w] numpy
-
-            img1 = video1[clip_timestamp].numpy()
-            img2 = video2[clip_timestamp].numpy()
-            
-            # calculate psnr of a video
-            psnr_results_of_a_video.append(img_psnr(img1, img2))
-
-        psnr_results.append(psnr_results_of_a_video)
-    
-    psnr_results = np.array(psnr_results) # [batch_size, num_frames]
-    psnr = {}
-    psnr_std = {}
-
-    for clip_timestamp in range(len(video1)):
-        psnr[clip_timestamp] = np.mean(psnr_results[:,clip_timestamp])
-        psnr_std[clip_timestamp] = np.std(psnr_results[:,clip_timestamp])
-
-    result = {
-        "value": psnr,
-        "value_std": psnr_std,
-        "video_setting": video1.shape,
-        "video_setting_name": "time, channel, heigth, width",
-    }
-
-    return result
-
-# test code / using example
-
-def main():
-    NUMBER_OF_VIDEOS = 8
-    VIDEO_LENGTH = 50
-    CHANNEL = 3
-    SIZE = 64
-    videos1 = torch.zeros(NUMBER_OF_VIDEOS, VIDEO_LENGTH, CHANNEL, SIZE, SIZE, requires_grad=False)
-    videos2 = torch.zeros(NUMBER_OF_VIDEOS, VIDEO_LENGTH, CHANNEL, SIZE, SIZE, requires_grad=False)
-
-    import json
-    result = calculate_psnr(videos1, videos2)
-    print(json.dumps(result, indent=4))
-
-if __name__ == "__main__":
-    main()

opensora/eval/cal_ssim.py

@@ -1,113 +0,0 @@
-import numpy as np
-import torch
-from tqdm import tqdm
-import cv2
- 
-def ssim(img1, img2):
-    C1 = 0.01 ** 2
-    C2 = 0.03 ** 2
-    img1 = img1.astype(np.float64)
-    img2 = img2.astype(np.float64)
-    kernel = cv2.getGaussianKernel(11, 1.5)
-    window = np.outer(kernel, kernel.transpose())
-    mu1 = cv2.filter2D(img1, -1, window)[5:-5, 5:-5]  # valid
-    mu2 = cv2.filter2D(img2, -1, window)[5:-5, 5:-5]
-    mu1_sq = mu1 ** 2
-    mu2_sq = mu2 ** 2
-    mu1_mu2 = mu1 * mu2
-    sigma1_sq = cv2.filter2D(img1 ** 2, -1, window)[5:-5, 5:-5] - mu1_sq
-    sigma2_sq = cv2.filter2D(img2 ** 2, -1, window)[5:-5, 5:-5] - mu2_sq
-    sigma12 = cv2.filter2D(img1 * img2, -1, window)[5:-5, 5:-5] - mu1_mu2
-    ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) *
-                                                            (sigma1_sq + sigma2_sq + C2))
-    return ssim_map.mean()
- 
- 
-def calculate_ssim_function(img1, img2):
-    # [0,1]
-    # ssim is the only metric extremely sensitive to gray being compared to b/w 
-    if not img1.shape == img2.shape:
-        raise ValueError('Input images must have the same dimensions.')
-    if img1.ndim == 2:
-        return ssim(img1, img2)
-    elif img1.ndim == 3:
-        if img1.shape[0] == 3:
-            ssims = []
-            for i in range(3):
-                ssims.append(ssim(img1[i], img2[i]))
-            return np.array(ssims).mean()                   
-        elif img1.shape[0] == 1:
-            return ssim(np.squeeze(img1), np.squeeze(img2))
-    else:
-        raise ValueError('Wrong input image dimensions.')
-
-def trans(x):
-    return x
-
-def calculate_ssim(videos1, videos2):
-    print("calculate_ssim...")
-
-    # videos [batch_size, timestamps, channel, h, w]
-    
-    assert videos1.shape == videos2.shape
-
-    videos1 = trans(videos1)
-    videos2 = trans(videos2)
-
-    ssim_results = []
-    
-    for video_num in tqdm(range(videos1.shape[0])):
-        # get a video
-        # video [timestamps, channel, h, w]
-        video1 = videos1[video_num]
-        video2 = videos2[video_num]
-
-        ssim_results_of_a_video = []
-        for clip_timestamp in range(len(video1)):
-            # get a img
-            # img [timestamps[x], channel, h, w]
-            # img [channel, h, w] numpy
-
-            img1 = video1[clip_timestamp].numpy()
-            img2 = video2[clip_timestamp].numpy()
-            
-            # calculate ssim of a video
-            ssim_results_of_a_video.append(calculate_ssim_function(img1, img2))
-
-        ssim_results.append(ssim_results_of_a_video)
-
-    ssim_results = np.array(ssim_results)
-
-    ssim = {}
-    ssim_std = {}
-
-    for clip_timestamp in range(len(video1)):
-        ssim[clip_timestamp] = np.mean(ssim_results[:,clip_timestamp])
-        ssim_std[clip_timestamp] = np.std(ssim_results[:,clip_timestamp])
-
-    result = {
-        "value": ssim,
-        "value_std": ssim_std,
-        "video_setting": video1.shape,
-        "video_setting_name": "time, channel, heigth, width",
-    }
-
-    return result
-
-# test code / using example
-
-def main():
-    NUMBER_OF_VIDEOS = 8
-    VIDEO_LENGTH = 50
-    CHANNEL = 3
-    SIZE = 64
-    videos1 = torch.zeros(NUMBER_OF_VIDEOS, VIDEO_LENGTH, CHANNEL, SIZE, SIZE, requires_grad=False)
-    videos2 = torch.zeros(NUMBER_OF_VIDEOS, VIDEO_LENGTH, CHANNEL, SIZE, SIZE, requires_grad=False)
-    device = torch.device("cuda")
-
-    import json
-    result = calculate_ssim(videos1, videos2)
-    print(json.dumps(result, indent=4))
-
-if __name__ == "__main__":
-    main()

opensora/eval/eval_clip_score.py

@@ -1,225 +0,0 @@
-"""Calculates the CLIP Scores
-
-The CLIP model is a contrasitively learned language-image model. There is
-an image encoder and a text encoder. It is believed that the CLIP model could 
-measure the similarity of cross modalities. Please find more information from 
-https://github.com/openai/CLIP.
-
-The CLIP Score measures the Cosine Similarity between two embedded features.
-This repository utilizes the pretrained CLIP Model to calculate 
-the mean average of cosine similarities. 
-
-See --help to see further details.
-
-Code apapted from https://github.com/mseitzer/pytorch-fid and https://github.com/openai/CLIP.
-
-Copyright 2023 The Hong Kong Polytechnic University
-
-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 os
-import os.path as osp
-from argparse import ArgumentDefaultsHelpFormatter, ArgumentParser
-
-import clip
-import torch
-from PIL import Image
-from torch.utils.data import Dataset, DataLoader
-
-try:
-    from tqdm import tqdm
-except ImportError:
-    # If tqdm is not available, provide a mock version of it
-    def tqdm(x):
-        return x
-
-
-IMAGE_EXTENSIONS = {'bmp', 'jpg', 'jpeg', 'pgm', 'png', 'ppm',
-                    'tif', 'tiff', 'webp'}
-
-TEXT_EXTENSIONS = {'txt'}
-
-
-class DummyDataset(Dataset):
-    
-    FLAGS = ['img', 'txt']
-    def __init__(self, real_path, generated_path,
-                 real_flag: str = 'img',
-                 generated_flag: str = 'img',
-                 transform = None,
-                 tokenizer = None) -> None:
-        super().__init__()
-        assert real_flag in self.FLAGS and generated_flag in self.FLAGS, \
-            'CLIP Score only support modality of {}. However, get {} and {}'.format(
-                self.FLAGS, real_flag, generated_flag
-            )
-        self.real_folder = self._combine_without_prefix(real_path)
-        self.real_flag = real_flag
-        self.fake_foler = self._combine_without_prefix(generated_path)
-        self.generated_flag = generated_flag
-        self.transform = transform
-        self.tokenizer = tokenizer
-        # assert self._check()
-
-    def __len__(self):
-        return len(self.real_folder)
-
-    def __getitem__(self, index):
-        if index >= len(self):
-            raise IndexError
-        real_path = self.real_folder[index]
-        generated_path = self.fake_foler[index]
-        real_data = self._load_modality(real_path, self.real_flag)
-        fake_data = self._load_modality(generated_path, self.generated_flag)
-
-        sample = dict(real=real_data, fake=fake_data)
-        return sample
-    
-    def _load_modality(self, path, modality):
-        if modality == 'img':
-            data = self._load_img(path)
-        elif modality == 'txt':
-            data = self._load_txt(path)
-        else:
-            raise TypeError("Got unexpected modality: {}".format(modality))
-        return data
-
-    def _load_img(self, path):
-        img = Image.open(path)
-        if self.transform is not None:
-            img = self.transform(img)
-        return img
-    
-    def _load_txt(self, path):
-        with open(path, 'r') as fp:
-            data = fp.read()
-            fp.close()
-        if self.tokenizer is not None:
-            data = self.tokenizer(data).squeeze()
-        return data
-
-    def _check(self):
-        for idx in range(len(self)):
-            real_name = self.real_folder[idx].split('.')
-            fake_name = self.fake_folder[idx].split('.')
-            if fake_name != real_name:
-                return False
-        return True
-    
-    def _combine_without_prefix(self, folder_path, prefix='.'):
-        folder = []
-        for name in os.listdir(folder_path):
-            if name[0] == prefix:
-                continue
-            folder.append(osp.join(folder_path, name))
-        folder.sort()
-        return folder
-
-
-@torch.no_grad()
-def calculate_clip_score(dataloader, model, real_flag, generated_flag):
-    score_acc = 0.
-    sample_num = 0.
-    logit_scale = model.logit_scale.exp()
-    for batch_data in tqdm(dataloader):
-        real = batch_data['real']
-        real_features = forward_modality(model, real, real_flag)
-        fake = batch_data['fake']
-        fake_features = forward_modality(model, fake, generated_flag)
-        
-        # normalize features
-        real_features = real_features / real_features.norm(dim=1, keepdim=True).to(torch.float32)
-        fake_features = fake_features / fake_features.norm(dim=1, keepdim=True).to(torch.float32)
-        
-        # calculate scores
-        # score = logit_scale * real_features @ fake_features.t()
-        # score_acc += torch.diag(score).sum()
-        score = logit_scale * (fake_features * real_features).sum()
-        score_acc += score
-        sample_num += real.shape[0]
-    
-    return score_acc / sample_num
-
-        
-def forward_modality(model, data, flag):
-    device = next(model.parameters()).device
-    if flag == 'img':
-        features = model.encode_image(data.to(device))
-    elif flag == 'txt':
-        features = model.encode_text(data.to(device))
-    else:
-        raise TypeError
-    return features
-
-
-def main():
-    parser = ArgumentParser(formatter_class=ArgumentDefaultsHelpFormatter)
-    parser.add_argument('--batch-size', type=int, default=50,
-                    help='Batch size to use')
-    parser.add_argument('--clip-model', type=str, default='ViT-B/32',
-                    help='CLIP model to use')
-    parser.add_argument('--num-workers', type=int, default=8,
-                    help=('Number of processes to use for data loading. '
-                          'Defaults to `min(8, num_cpus)`'))
-    parser.add_argument('--device', type=str, default=None,
-                    help='Device to use. Like cuda, cuda:0 or cpu')
-    parser.add_argument('--real_flag', type=str, default='img',
-                    help=('The modality of real path. '
-                          'Default to img'))
-    parser.add_argument('--generated_flag', type=str, default='txt',
-                    help=('The modality of generated path. '
-                          'Default to txt'))
-    parser.add_argument('--real_path', type=str, 
-                    help=('Paths to the real images or '
-                          'to .npz statistic files'))
-    parser.add_argument('--generated_path', type=str,
-                    help=('Paths to the generated images or '
-                          'to .npz statistic files'))
-    args = parser.parse_args()
-
-    if args.device is None:
-        device = torch.device('cuda' if (torch.cuda.is_available()) else 'cpu')
-    else:
-        device = torch.device(args.device)
-
-    if args.num_workers is None:
-        try:
-            num_cpus = len(os.sched_getaffinity(0))
-        except AttributeError:
-            # os.sched_getaffinity is not available under Windows, use
-            # os.cpu_count instead (which may not return the *available* number
-            # of CPUs).
-            num_cpus = os.cpu_count()
-
-        num_workers = min(num_cpus, 8) if num_cpus is not None else 0
-    else:
-        num_workers = args.num_workers
-
-    print('Loading CLIP model: {}'.format(args.clip_model))
-    model, preprocess = clip.load(args.clip_model, device=device)
-    
-    dataset = DummyDataset(args.real_path, args.generated_path,
-                           args.real_flag, args.generated_flag,
-                           transform=preprocess, tokenizer=clip.tokenize)
-    dataloader = DataLoader(dataset, args.batch_size, 
-                            num_workers=num_workers, pin_memory=True)
-    
-    print('Calculating CLIP Score:')
-    clip_score = calculate_clip_score(dataloader, model,
-                                      args.real_flag, args.generated_flag)
-    clip_score = clip_score.cpu().item()
-    print('CLIP Score: ', clip_score)
-
-
-if __name__ == '__main__':
-    main()

opensora/eval/eval_common_metric.py

@@ -1,224 +0,0 @@
-"""Calculates the CLIP Scores
-
-The CLIP model is a contrasitively learned language-image model. There is
-an image encoder and a text encoder. It is believed that the CLIP model could 
-measure the similarity of cross modalities. Please find more information from 
-https://github.com/openai/CLIP.
-
-The CLIP Score measures the Cosine Similarity between two embedded features.
-This repository utilizes the pretrained CLIP Model to calculate 
-the mean average of cosine similarities. 
-
-See --help to see further details.
-
-Code apapted from https://github.com/mseitzer/pytorch-fid and https://github.com/openai/CLIP.
-
-Copyright 2023 The Hong Kong Polytechnic University
-
-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 os
-import os.path as osp
-from argparse import ArgumentDefaultsHelpFormatter, ArgumentParser
-import numpy as np
-import torch
-from torch.utils.data import Dataset, DataLoader, Subset
-from decord import VideoReader, cpu
-import random
-from pytorchvideo.transforms import ShortSideScale
-from torchvision.io import read_video
-from torchvision.transforms import Lambda, Compose
-from torchvision.transforms.transforms import CenterCropVideo
-import sys
-sys.path.append(".")
-from opensora.eval.cal_lpips import calculate_lpips
-from opensora.eval.cal_fvd import calculate_fvd
-from opensora.eval.cal_psnr import calculate_psnr
-from opensora.eval.cal_flolpips import calculate_flolpips
-from opensora.eval.cal_ssim import calculate_ssim
-
-try:
-    from tqdm import tqdm
-except ImportError:
-    # If tqdm is not available, provide a mock version of it
-    def tqdm(x):
-        return x
-
-class VideoDataset(Dataset):
-    def __init__(self, 
-                 real_video_dir,
-                 generated_video_dir,
-                 num_frames,
-                 sample_rate = 1,
-                 crop_size=None,
-                 resolution=128,
-                 ) -> None:
-        super().__init__()
-        self.real_video_files = self._combine_without_prefix(real_video_dir)
-        self.generated_video_files = self._combine_without_prefix(generated_video_dir)
-        self.num_frames = num_frames
-        self.sample_rate = sample_rate
-        self.crop_size = crop_size
-        self.short_size = resolution
-
-
-    def __len__(self):
-        return len(self.real_video_files)
-
-    def __getitem__(self, index):
-        if index >= len(self):
-            raise IndexError
-        real_video_file = self.real_video_files[index]
-        generated_video_file = self.generated_video_files[index]
-        print(real_video_file, generated_video_file)
-        real_video_tensor  = self._load_video(real_video_file)
-        generated_video_tensor  = self._load_video(generated_video_file)
-        return {'real': real_video_tensor, 'generated':generated_video_tensor }
-
-
-    def _load_video(self, video_path):
-        num_frames = self.num_frames
-        sample_rate = self.sample_rate
-        decord_vr = VideoReader(video_path, ctx=cpu(0))
-        total_frames = len(decord_vr)
-        sample_frames_len = sample_rate * num_frames
-
-        if total_frames >= sample_frames_len:
-            s = 0
-            e = s + sample_frames_len
-            num_frames = num_frames
-        else:
-            s = 0
-            e = total_frames
-            num_frames = int(total_frames / sample_frames_len * num_frames)
-            print(f'sample_frames_len {sample_frames_len}, only can sample {num_frames * sample_rate}', video_path,
-                total_frames)
-
-
-        frame_id_list = np.linspace(s, e - 1, num_frames, dtype=int)
-        video_data = decord_vr.get_batch(frame_id_list).asnumpy()
-        video_data = torch.from_numpy(video_data)
-        video_data = video_data.permute(0, 3, 1, 2) # (T, H, W, C) -> (C, T, H, W)
-        return _preprocess(video_data, short_size=self.short_size, crop_size = self.crop_size)
-
-
-    def _combine_without_prefix(self, folder_path, prefix='.'):
-        folder = []
-        os.makedirs(folder_path, exist_ok=True)
-        for name in os.listdir(folder_path):
-            if name[0] == prefix:
-                continue
-            if osp.isfile(osp.join(folder_path, name)):
-                folder.append(osp.join(folder_path, name))
-        folder.sort()
-        return folder
-
-def _preprocess(video_data, short_size=128, crop_size=None):
-    transform = Compose(
-        [
-            Lambda(lambda x: x / 255.0),
-            ShortSideScale(size=short_size),
-            CenterCropVideo(crop_size=crop_size),
-        ]
-    )
-    video_outputs = transform(video_data)
-    # video_outputs = torch.unsqueeze(video_outputs, 0) # (bz,c,t,h,w)
-    return video_outputs
-
-
-def calculate_common_metric(args, dataloader, device):
-
-    score_list = []
-    for batch_data in tqdm(dataloader): # {'real': real_video_tensor, 'generated':generated_video_tensor }
-        real_videos = batch_data['real'] 
-        generated_videos = batch_data['generated']
-        assert real_videos.shape[2] == generated_videos.shape[2]
-        if args.metric == 'fvd':
-            tmp_list = list(calculate_fvd(real_videos, generated_videos, args.device, method=args.fvd_method)['value'].values())
-        elif args.metric == 'ssim':
-            tmp_list = list(calculate_ssim(real_videos, generated_videos)['value'].values())
-        elif args.metric == 'psnr':
-            tmp_list = list(calculate_psnr(real_videos, generated_videos)['value'].values())
-        elif args.metric == 'flolpips':
-            result = calculate_flolpips(real_videos, generated_videos, args.device)
-            tmp_list = list(result['value'].values())
-        else:
-            tmp_list  = list(calculate_lpips(real_videos, generated_videos, args.device)['value'].values())
-        score_list += tmp_list
-    return np.mean(score_list)
-        
-def main():
-    parser = ArgumentParser(formatter_class=ArgumentDefaultsHelpFormatter)
-    parser.add_argument('--batch_size', type=int, default=2,
-                    help='Batch size to use')
-    parser.add_argument('--real_video_dir', type=str,
-                    help=('the path of real videos`'))
-    parser.add_argument('--generated_video_dir', type=str,
-                    help=('the path of generated videos`'))
-    parser.add_argument('--device', type=str, default=None,
-                    help='Device to use. Like cuda, cuda:0 or cpu')
-    parser.add_argument('--num_workers', type=int, default=8,
-                    help=('Number of processes to use for data loading. '
-                          'Defaults to `min(8, num_cpus)`'))
-    parser.add_argument('--sample_fps', type=int, default=30)
-    parser.add_argument('--resolution', type=int, default=336)
-    parser.add_argument('--crop_size', type=int, default=None)
-    parser.add_argument('--num_frames', type=int, default=100)
-    parser.add_argument('--sample_rate', type=int, default=1)
-    parser.add_argument('--subset_size', type=int, default=None)
-    parser.add_argument("--metric", type=str, default="fvd",choices=['fvd','psnr','ssim','lpips', 'flolpips'])
-    parser.add_argument("--fvd_method", type=str, default='styleganv',choices=['styleganv','videogpt'])
-
-
-    args = parser.parse_args()
-
-    if args.device is None:
-        device = torch.device('cuda' if (torch.cuda.is_available()) else 'cpu')
-    else:
-        device = torch.device(args.device)
-
-    if args.num_workers is None:
-        try:
-            num_cpus = len(os.sched_getaffinity(0))
-        except AttributeError:
-            # os.sched_getaffinity is not available under Windows, use
-            # os.cpu_count instead (which may not return the *available* number
-            # of CPUs).
-            num_cpus = os.cpu_count()
-
-        num_workers = min(num_cpus, 8) if num_cpus is not None else 0
-    else:
-        num_workers = args.num_workers
-
-    
-    dataset = VideoDataset(args.real_video_dir,
-                           args.generated_video_dir,
-                            num_frames = args.num_frames,
-                            sample_rate = args.sample_rate,
-                            crop_size=args.crop_size,
-                            resolution=args.resolution) 
-
-    if args.subset_size:
-        indices = range(args.subset_size)
-        dataset = Subset(dataset, indices=indices)
-    
-    dataloader = DataLoader(dataset, args.batch_size, 
-                            num_workers=num_workers, pin_memory=True)
-    
-
-    metric_score = calculate_common_metric(args, dataloader,device)
-    print('metric: ', args.metric, " ",metric_score)
-
-if __name__ == '__main__':
-    main()

opensora/eval/flolpips/correlation/correlation.py

@@ -1,397 +0,0 @@
-#!/usr/bin/env python
-
-import torch
-
-import cupy
-import re
-
-kernel_Correlation_rearrange = '''
-	extern "C" __global__ void kernel_Correlation_rearrange(
-		const int n,
-		const float* input,
-		float* output
-	) {
-	  int intIndex = (blockIdx.x * blockDim.x) + threadIdx.x;
-
-	  if (intIndex >= n) {
-	    return;
-	  }
-
-	  int intSample = blockIdx.z;
-	  int intChannel = blockIdx.y;
-
-	  float fltValue = input[(((intSample * SIZE_1(input)) + intChannel) * SIZE_2(input) * SIZE_3(input)) + intIndex];
-
-	  __syncthreads();
-
-	  int intPaddedY = (intIndex / SIZE_3(input)) + 4;
-	  int intPaddedX = (intIndex % SIZE_3(input)) + 4;
-	  int intRearrange = ((SIZE_3(input) + 8) * intPaddedY) + intPaddedX;
-
-	  output[(((intSample * SIZE_1(output) * SIZE_2(output)) + intRearrange) * SIZE_1(input)) + intChannel] = fltValue;
-	}
-'''
-
-kernel_Correlation_updateOutput = '''
-	extern "C" __global__ void kernel_Correlation_updateOutput(
-	  const int n,
-	  const float* rbot0,
-	  const float* rbot1,
-	  float* top
-	) {
-	  extern __shared__ char patch_data_char[];
-	  
-	  float *patch_data = (float *)patch_data_char;
-	  
-	  // First (upper left) position of kernel upper-left corner in current center position of neighborhood in image 1
-	  int x1 = blockIdx.x + 4;
-	  int y1 = blockIdx.y + 4;
-	  int item = blockIdx.z;
-	  int ch_off = threadIdx.x;
-	  
-	  // Load 3D patch into shared shared memory
-	  for (int j = 0; j < 1; j++) { // HEIGHT
-	    for (int i = 0; i < 1; i++) { // WIDTH
-	      int ji_off = (j + i) * SIZE_3(rbot0);
-	      for (int ch = ch_off; ch < SIZE_3(rbot0); ch += 32) { // CHANNELS
-	        int idx1 = ((item * SIZE_1(rbot0) + y1+j) * SIZE_2(rbot0) + x1+i) * SIZE_3(rbot0) + ch;
-	        int idxPatchData = ji_off + ch;
-	        patch_data[idxPatchData] = rbot0[idx1];
-	      }
-	    }
-	  }
-	  
-	  __syncthreads();
-	  
-	  __shared__ float sum[32];
-	  
-	  // Compute correlation
-	  for (int top_channel = 0; top_channel < SIZE_1(top); top_channel++) {
-	    sum[ch_off] = 0;
-	  
-	    int s2o = top_channel % 9 - 4;
-	    int s2p = top_channel / 9 - 4;
-	    
-	    for (int j = 0; j < 1; j++) { // HEIGHT
-	      for (int i = 0; i < 1; i++) { // WIDTH
-	        int ji_off = (j + i) * SIZE_3(rbot0);
-	        for (int ch = ch_off; ch < SIZE_3(rbot0); ch += 32) { // CHANNELS
-	          int x2 = x1 + s2o;
-	          int y2 = y1 + s2p;
-	          
-	          int idxPatchData = ji_off + ch;
-	          int idx2 = ((item * SIZE_1(rbot0) + y2+j) * SIZE_2(rbot0) + x2+i) * SIZE_3(rbot0) + ch;
-	          
-	          sum[ch_off] += patch_data[idxPatchData] * rbot1[idx2];
-	        }
-	      }
-	    }
-	    
-	    __syncthreads();
-	    
-	    if (ch_off == 0) {
-	      float total_sum = 0;
-	      for (int idx = 0; idx < 32; idx++) {
-	        total_sum += sum[idx];
-	      }
-	      const int sumelems = SIZE_3(rbot0);
-	      const int index = ((top_channel*SIZE_2(top) + blockIdx.y)*SIZE_3(top))+blockIdx.x;
-	      top[index + item*SIZE_1(top)*SIZE_2(top)*SIZE_3(top)] = total_sum / (float)sumelems;
-	    }
-	  }
-	}
-'''
-
-kernel_Correlation_updateGradFirst = '''
-	#define ROUND_OFF 50000
-
-	extern "C" __global__ void kernel_Correlation_updateGradFirst(
-	  const int n,
-	  const int intSample,
-	  const float* rbot0,
-	  const float* rbot1,
-	  const float* gradOutput,
-	  float* gradFirst,
-	  float* gradSecond
-	) { for (int intIndex = (blockIdx.x * blockDim.x) + threadIdx.x; intIndex < n; intIndex += blockDim.x * gridDim.x) {
-	  int n = intIndex % SIZE_1(gradFirst); // channels
-	  int l = (intIndex / SIZE_1(gradFirst)) % SIZE_3(gradFirst) + 4; // w-pos
-	  int m = (intIndex / SIZE_1(gradFirst) / SIZE_3(gradFirst)) % SIZE_2(gradFirst) + 4; // h-pos
-	  
-	  // round_off is a trick to enable integer division with ceil, even for negative numbers
-	  // We use a large offset, for the inner part not to become negative.
-	  const int round_off = ROUND_OFF;
-	  const int round_off_s1 = round_off;
-	  
-	  // We add round_off before_s1 the int division and subtract round_off after it, to ensure the formula matches ceil behavior:
-	  int xmin = (l - 4 + round_off_s1 - 1) + 1 - round_off; // ceil (l - 4)
-	  int ymin = (m - 4 + round_off_s1 - 1) + 1 - round_off; // ceil (l - 4)
-	  
-	  // Same here:
-	  int xmax = (l - 4 + round_off_s1) - round_off; // floor (l - 4)
-	  int ymax = (m - 4 + round_off_s1) - round_off; // floor (m - 4)
-	  
-	  float sum = 0;
-	  if (xmax>=0 && ymax>=0 && (xmin<=SIZE_3(gradOutput)-1) && (ymin<=SIZE_2(gradOutput)-1)) {
-	    xmin = max(0,xmin);
-	    xmax = min(SIZE_3(gradOutput)-1,xmax);
-	    
-	    ymin = max(0,ymin);
-	    ymax = min(SIZE_2(gradOutput)-1,ymax);
-	    
-	    for (int p = -4; p <= 4; p++) {
-	      for (int o = -4; o <= 4; o++) {
-	        // Get rbot1 data:
-	        int s2o = o;
-	        int s2p = p;
-	        int idxbot1 = ((intSample * SIZE_1(rbot0) + (m+s2p)) * SIZE_2(rbot0) + (l+s2o)) * SIZE_3(rbot0) + n;
-	        float bot1tmp = rbot1[idxbot1]; // rbot1[l+s2o,m+s2p,n]
-	        
-	        // Index offset for gradOutput in following loops:
-	        int op = (p+4) * 9 + (o+4); // index[o,p]
-	        int idxopoffset = (intSample * SIZE_1(gradOutput) + op);
-	        
-	        for (int y = ymin; y <= ymax; y++) {
-	          for (int x = xmin; x <= xmax; x++) {
-	            int idxgradOutput = (idxopoffset * SIZE_2(gradOutput) + y) * SIZE_3(gradOutput) + x; // gradOutput[x,y,o,p]
-	            sum += gradOutput[idxgradOutput] * bot1tmp;
-	          }
-	        }
-	      }
-	    }
-	  }
-	  const int sumelems = SIZE_1(gradFirst);
-	  const int bot0index = ((n * SIZE_2(gradFirst)) + (m-4)) * SIZE_3(gradFirst) + (l-4);
-	  gradFirst[bot0index + intSample*SIZE_1(gradFirst)*SIZE_2(gradFirst)*SIZE_3(gradFirst)] = sum / (float)sumelems;
-	} }
-'''
-
-kernel_Correlation_updateGradSecond = '''
-	#define ROUND_OFF 50000
-
-	extern "C" __global__ void kernel_Correlation_updateGradSecond(
-	  const int n,
-	  const int intSample,
-	  const float* rbot0,
-	  const float* rbot1,
-	  const float* gradOutput,
-	  float* gradFirst,
-	  float* gradSecond
-	) { for (int intIndex = (blockIdx.x * blockDim.x) + threadIdx.x; intIndex < n; intIndex += blockDim.x * gridDim.x) {
-	  int n = intIndex % SIZE_1(gradSecond); // channels
-	  int l = (intIndex / SIZE_1(gradSecond)) % SIZE_3(gradSecond) + 4; // w-pos
-	  int m = (intIndex / SIZE_1(gradSecond) / SIZE_3(gradSecond)) % SIZE_2(gradSecond) + 4; // h-pos
-	  
-	  // round_off is a trick to enable integer division with ceil, even for negative numbers
-	  // We use a large offset, for the inner part not to become negative.
-	  const int round_off = ROUND_OFF;
-	  const int round_off_s1 = round_off;
-	  
-	  float sum = 0;
-	  for (int p = -4; p <= 4; p++) {
-	    for (int o = -4; o <= 4; o++) {
-	      int s2o = o;
-	      int s2p = p;
-	      
-	      //Get X,Y ranges and clamp
-	      // We add round_off before_s1 the int division and subtract round_off after it, to ensure the formula matches ceil behavior:
-	      int xmin = (l - 4 - s2o + round_off_s1 - 1) + 1 - round_off; // ceil (l - 4 - s2o)
-	      int ymin = (m - 4 - s2p + round_off_s1 - 1) + 1 - round_off; // ceil (l - 4 - s2o)
-	      
-	      // Same here:
-	      int xmax = (l - 4 - s2o + round_off_s1) - round_off; // floor (l - 4 - s2o)
-	      int ymax = (m - 4 - s2p + round_off_s1) - round_off; // floor (m - 4 - s2p)
-          
-	      if (xmax>=0 && ymax>=0 && (xmin<=SIZE_3(gradOutput)-1) && (ymin<=SIZE_2(gradOutput)-1)) {
-	        xmin = max(0,xmin);
-	        xmax = min(SIZE_3(gradOutput)-1,xmax);
-	        
-	        ymin = max(0,ymin);
-	        ymax = min(SIZE_2(gradOutput)-1,ymax);
-	        
-	        // Get rbot0 data:
-	        int idxbot0 = ((intSample * SIZE_1(rbot0) + (m-s2p)) * SIZE_2(rbot0) + (l-s2o)) * SIZE_3(rbot0) + n;
-	        float bot0tmp = rbot0[idxbot0]; // rbot1[l+s2o,m+s2p,n]
-	        
-	        // Index offset for gradOutput in following loops:
-	        int op = (p+4) * 9 + (o+4); // index[o,p]
-	        int idxopoffset = (intSample * SIZE_1(gradOutput) + op);
-	        
-	        for (int y = ymin; y <= ymax; y++) {
-	          for (int x = xmin; x <= xmax; x++) {
-	            int idxgradOutput = (idxopoffset * SIZE_2(gradOutput) + y) * SIZE_3(gradOutput) + x; // gradOutput[x,y,o,p]
-	            sum += gradOutput[idxgradOutput] * bot0tmp;
-	          }
-	        }
-	      }
-	    }
-	  }
-	  const int sumelems = SIZE_1(gradSecond);
-	  const int bot1index = ((n * SIZE_2(gradSecond)) + (m-4)) * SIZE_3(gradSecond) + (l-4);
-	  gradSecond[bot1index + intSample*SIZE_1(gradSecond)*SIZE_2(gradSecond)*SIZE_3(gradSecond)] = sum / (float)sumelems;
-	} }
-'''
-
-def cupy_kernel(strFunction, objVariables):
-	strKernel = globals()[strFunction]
-
-	while True:
-		objMatch = re.search('(SIZE_)([0-4])(\()([^\)]*)(\))', strKernel)
-
-		if objMatch is None:
-			break
-		# end
-
-		intArg = int(objMatch.group(2))
-
-		strTensor = objMatch.group(4)
-		intSizes = objVariables[strTensor].size()
-
-		strKernel = strKernel.replace(objMatch.group(), str(intSizes[intArg]))
-	# end
-
-	while True:
-		objMatch = re.search('(VALUE_)([0-4])(\()([^\)]+)(\))', strKernel)
-
-		if objMatch is None:
-			break
-		# end
-
-		intArgs = int(objMatch.group(2))
-		strArgs = objMatch.group(4).split(',')
-
-		strTensor = strArgs[0]
-		intStrides = objVariables[strTensor].stride()
-		strIndex = [ '((' + strArgs[intArg + 1].replace('{', '(').replace('}', ')').strip() + ')*' + str(intStrides[intArg]) + ')' for intArg in range(intArgs) ]
-
-		strKernel = strKernel.replace(objMatch.group(0), strTensor + '[' + str.join('+', strIndex) + ']')
-	# end
-
-	return strKernel
-# end
-
-@cupy.memoize(for_each_device=True)
-def cupy_launch(strFunction, strKernel):
-    return cupy.RawKernel(strKernel, strFunction)
-# end
-
-class _FunctionCorrelation(torch.autograd.Function):
-	@staticmethod
-	def forward(self, first, second):
-		rbot0 = first.new_zeros([ first.shape[0], first.shape[2] + 8, first.shape[3] + 8, first.shape[1] ])
-		rbot1 = first.new_zeros([ first.shape[0], first.shape[2] + 8, first.shape[3] + 8, first.shape[1] ])
-
-		self.save_for_backward(first, second, rbot0, rbot1)
-
-		first = first.contiguous();	assert(first.is_cuda == True)
-		second = second.contiguous(); assert(second.is_cuda == True)
-
-		output = first.new_zeros([ first.shape[0], 81, first.shape[2], first.shape[3] ])
-
-		if first.is_cuda == True:
-			n = first.shape[2] * first.shape[3]
-			cupy_launch('kernel_Correlation_rearrange', cupy_kernel('kernel_Correlation_rearrange', {
-				'input': first,
-				'output': rbot0
-			}))(
-				grid=tuple([ int((n + 16 - 1) / 16), first.shape[1], first.shape[0] ]),
-				block=tuple([ 16, 1, 1 ]),
-				args=[ n, first.data_ptr(), rbot0.data_ptr() ]
-			)
-
-			n = second.shape[2] * second.shape[3]
-			cupy_launch('kernel_Correlation_rearrange', cupy_kernel('kernel_Correlation_rearrange', {
-				'input': second,
-				'output': rbot1
-			}))(
-				grid=tuple([ int((n + 16 - 1) / 16), second.shape[1], second.shape[0] ]),
-				block=tuple([ 16, 1, 1 ]),
-				args=[ n, second.data_ptr(), rbot1.data_ptr() ]
-			)
-
-			n = output.shape[1] * output.shape[2] * output.shape[3]
-			cupy_launch('kernel_Correlation_updateOutput', cupy_kernel('kernel_Correlation_updateOutput', {
-				'rbot0': rbot0,
-				'rbot1': rbot1,
-				'top': output
-			}))(
-				grid=tuple([ output.shape[3], output.shape[2], output.shape[0] ]),
-				block=tuple([ 32, 1, 1 ]),
-				shared_mem=first.shape[1] * 4,
-				args=[ n, rbot0.data_ptr(), rbot1.data_ptr(), output.data_ptr() ]
-			)
-
-		elif first.is_cuda == False:
-			raise NotImplementedError()
-
-		# end
-
-		return output
-	# end
-
-	@staticmethod
-	def backward(self, gradOutput):
-		first, second, rbot0, rbot1 = self.saved_tensors
-
-		gradOutput = gradOutput.contiguous(); assert(gradOutput.is_cuda == True)
-
-		gradFirst = first.new_zeros([ first.shape[0], first.shape[1], first.shape[2], first.shape[3] ]) if self.needs_input_grad[0] == True else None
-		gradSecond = first.new_zeros([ first.shape[0], first.shape[1], first.shape[2], first.shape[3] ]) if self.needs_input_grad[1] == True else None
-
-		if first.is_cuda == True:
-			if gradFirst is not None:
-				for intSample in range(first.shape[0]):
-					n = first.shape[1] * first.shape[2] * first.shape[3]
-					cupy_launch('kernel_Correlation_updateGradFirst', cupy_kernel('kernel_Correlation_updateGradFirst', {
-						'rbot0': rbot0,
-						'rbot1': rbot1,
-						'gradOutput': gradOutput,
-						'gradFirst': gradFirst,
-						'gradSecond': None
-					}))(
-						grid=tuple([ int((n + 512 - 1) / 512), 1, 1 ]),
-						block=tuple([ 512, 1, 1 ]),
-						args=[ n, intSample, rbot0.data_ptr(), rbot1.data_ptr(), gradOutput.data_ptr(), gradFirst.data_ptr(), None ]
-					)
-				# end
-			# end
-
-			if gradSecond is not None:
-				for intSample in range(first.shape[0]):
-					n = first.shape[1] * first.shape[2] * first.shape[3]
-					cupy_launch('kernel_Correlation_updateGradSecond', cupy_kernel('kernel_Correlation_updateGradSecond', {
-						'rbot0': rbot0,
-						'rbot1': rbot1,
-						'gradOutput': gradOutput,
-						'gradFirst': None,
-						'gradSecond': gradSecond
-					}))(
-						grid=tuple([ int((n + 512 - 1) / 512), 1, 1 ]),
-						block=tuple([ 512, 1, 1 ]),
-						args=[ n, intSample, rbot0.data_ptr(), rbot1.data_ptr(), gradOutput.data_ptr(), None, gradSecond.data_ptr() ]
-					)
-				# end
-			# end
-
-		elif first.is_cuda == False:
-			raise NotImplementedError()
-
-		# end
-
-		return gradFirst, gradSecond
-	# end
-# end
-
-def FunctionCorrelation(tenFirst, tenSecond):
-	return _FunctionCorrelation.apply(tenFirst, tenSecond)
-# end
-
-class ModuleCorrelation(torch.nn.Module):
-	def __init__(self):
-		super(ModuleCorrelation, self).__init__()
-	# end
-
-	def forward(self, tenFirst, tenSecond):
-		return _FunctionCorrelation.apply(tenFirst, tenSecond)
-	# end
-# end

opensora/eval/flolpips/flolpips.py

@@ -1,308 +0,0 @@
-
-from __future__ import absolute_import
-import os
-import numpy as np
-import torch
-import torch.nn as nn
-from torch.autograd import Variable
-from .pretrained_networks import vgg16, alexnet, squeezenet
-import torch.nn
-import torch.nn.functional as F
-import torchvision.transforms.functional as TF
-import cv2
-
-from .pwcnet import Network as PWCNet
-from .utils import *
-
-def spatial_average(in_tens, keepdim=True):
-    return in_tens.mean([2,3],keepdim=keepdim)
-
-def mw_spatial_average(in_tens, flow, keepdim=True):
-    _,_,h,w = in_tens.shape
-    flow = F.interpolate(flow, (h,w), align_corners=False, mode='bilinear')
-    flow_mag = torch.sqrt(flow[:,0:1]**2 + flow[:,1:2]**2)
-    flow_mag = flow_mag / torch.sum(flow_mag, dim=[1,2,3], keepdim=True)
-    return torch.sum(in_tens*flow_mag, dim=[2,3],keepdim=keepdim)
-
-
-def mtw_spatial_average(in_tens, flow, texture, keepdim=True):
-    _,_,h,w = in_tens.shape
-    flow = F.interpolate(flow, (h,w), align_corners=False, mode='bilinear')
-    texture = F.interpolate(texture, (h,w), align_corners=False, mode='bilinear')
-    flow_mag = torch.sqrt(flow[:,0:1]**2 + flow[:,1:2]**2)
-    flow_mag = (flow_mag - flow_mag.min()) / (flow_mag.max() - flow_mag.min()) + 1e-6
-    texture = (texture - texture.min()) / (texture.max() - texture.min()) + 1e-6
-    weight = flow_mag / texture
-    weight /= torch.sum(weight)
-    return torch.sum(in_tens*weight, dim=[2,3],keepdim=keepdim)
-
-
-
-def m2w_spatial_average(in_tens, flow, keepdim=True):
-    _,_,h,w = in_tens.shape
-    flow = F.interpolate(flow, (h,w), align_corners=False, mode='bilinear')
-    flow_mag = flow[:,0:1]**2 + flow[:,1:2]**2 # B,1,H,W
-    flow_mag = flow_mag / torch.sum(flow_mag)
-    return torch.sum(in_tens*flow_mag, dim=[2,3],keepdim=keepdim)
-
-def upsample(in_tens, out_HW=(64,64)): # assumes scale factor is same for H and W
-    in_H, in_W = in_tens.shape[2], in_tens.shape[3]
-    return nn.Upsample(size=out_HW, mode='bilinear', align_corners=False)(in_tens)
-
-# Learned perceptual metric
-class LPIPS(nn.Module):
-    def __init__(self, pretrained=True, net='alex', version='0.1', lpips=True, spatial=False, 
-        pnet_rand=False, pnet_tune=False, use_dropout=True, model_path=None, eval_mode=True, verbose=False):
-        # lpips - [True] means with linear calibration on top of base network
-        # pretrained - [True] means load linear weights
-
-        super(LPIPS, self).__init__()
-        if(verbose):
-            print('Setting up [%s] perceptual loss: trunk [%s], v[%s], spatial [%s]'%
-                ('LPIPS' if lpips else 'baseline', net, version, 'on' if spatial else 'off'))
-
-        self.pnet_type = net
-        self.pnet_tune = pnet_tune
-        self.pnet_rand = pnet_rand
-        self.spatial = spatial
-        self.lpips = lpips # false means baseline of just averaging all layers
-        self.version = version
-        self.scaling_layer = ScalingLayer()
-
-        if(self.pnet_type in ['vgg','vgg16']):
-            net_type = vgg16
-            self.chns = [64,128,256,512,512]
-        elif(self.pnet_type=='alex'):
-            net_type = alexnet
-            self.chns = [64,192,384,256,256]
-        elif(self.pnet_type=='squeeze'):
-            net_type = squeezenet
-            self.chns = [64,128,256,384,384,512,512]
-        self.L = len(self.chns)
-
-        self.net = net_type(pretrained=not self.pnet_rand, requires_grad=self.pnet_tune)
-
-        if(lpips):
-            self.lin0 = NetLinLayer(self.chns[0], use_dropout=use_dropout)
-            self.lin1 = NetLinLayer(self.chns[1], use_dropout=use_dropout)
-            self.lin2 = NetLinLayer(self.chns[2], use_dropout=use_dropout)
-            self.lin3 = NetLinLayer(self.chns[3], use_dropout=use_dropout)
-            self.lin4 = NetLinLayer(self.chns[4], use_dropout=use_dropout)
-            self.lins = [self.lin0,self.lin1,self.lin2,self.lin3,self.lin4]
-            if(self.pnet_type=='squeeze'): # 7 layers for squeezenet
-                self.lin5 = NetLinLayer(self.chns[5], use_dropout=use_dropout)
-                self.lin6 = NetLinLayer(self.chns[6], use_dropout=use_dropout)
-                self.lins+=[self.lin5,self.lin6]
-            self.lins = nn.ModuleList(self.lins)
-
-            if(pretrained):
-                if(model_path is None):
-                    import inspect
-                    import os
-                    model_path = os.path.abspath(os.path.join(inspect.getfile(self.__init__), '..', 'weights/v%s/%s.pth'%(version,net)))
-
-                if(verbose):
-                    print('Loading model from: %s'%model_path)
-                self.load_state_dict(torch.load(model_path, map_location='cpu'), strict=False)          
-
-        if(eval_mode):
-            self.eval()
-
-    def forward(self, in0, in1, retPerLayer=False, normalize=False):
-        if normalize: # turn on this flag if input is [0,1] so it can be adjusted to [-1, +1]
-            in0 = 2 * in0  - 1
-            in1 = 2 * in1  - 1
-
-        # v0.0 - original release had a bug, where input was not scaled
-        in0_input, in1_input = (self.scaling_layer(in0), self.scaling_layer(in1)) if self.version=='0.1' else (in0, in1)
-        outs0, outs1 = self.net.forward(in0_input), self.net.forward(in1_input)
-        feats0, feats1, diffs = {}, {}, {}
-
-        for kk in range(self.L):
-            feats0[kk], feats1[kk] = normalize_tensor(outs0[kk]), normalize_tensor(outs1[kk])
-            diffs[kk] = (feats0[kk]-feats1[kk])**2
-
-        if(self.lpips):
-            if(self.spatial):
-                res = [upsample(self.lins[kk](diffs[kk]), out_HW=in0.shape[2:]) for kk in range(self.L)]
-            else:
-                res = [spatial_average(self.lins[kk](diffs[kk]), keepdim=True) for kk in range(self.L)]
-        else:
-            if(self.spatial):
-                res = [upsample(diffs[kk].sum(dim=1,keepdim=True), out_HW=in0.shape[2:]) for kk in range(self.L)]
-            else:
-                res = [spatial_average(diffs[kk].sum(dim=1,keepdim=True), keepdim=True) for kk in range(self.L)]
-
-        # val = res[0]
-        # for l in range(1,self.L):
-        #     val += res[l]
-        #     print(val)
-
-        # a = spatial_average(self.lins[kk](diffs[kk]), keepdim=True)
-        # b = torch.max(self.lins[kk](feats0[kk]**2))
-        # for kk in range(self.L):
-        #     a += spatial_average(self.lins[kk](diffs[kk]), keepdim=True)
-        #     b = torch.max(b,torch.max(self.lins[kk](feats0[kk]**2)))
-        # a = a/self.L
-        # from IPython import embed
-        # embed()
-        # return 10*torch.log10(b/a)
-        
-        # if(retPerLayer):
-        #     return (val, res)
-        # else:
-        return torch.sum(torch.cat(res, 1), dim=(1,2,3), keepdims=False)
-
-
-class ScalingLayer(nn.Module):
-    def __init__(self):
-        super(ScalingLayer, self).__init__()
-        self.register_buffer('shift', torch.Tensor([-.030,-.088,-.188])[None,:,None,None])
-        self.register_buffer('scale', torch.Tensor([.458,.448,.450])[None,:,None,None])
-
-    def forward(self, inp):
-        return (inp - self.shift) / self.scale
-
-
-class NetLinLayer(nn.Module):
-    ''' A single linear layer which does a 1x1 conv '''
-    def __init__(self, chn_in, chn_out=1, use_dropout=False):
-        super(NetLinLayer, self).__init__()
-
-        layers = [nn.Dropout(),] if(use_dropout) else []
-        layers += [nn.Conv2d(chn_in, chn_out, 1, stride=1, padding=0, bias=False),]
-        self.model = nn.Sequential(*layers)
-
-    def forward(self, x):
-        return self.model(x)
-
-class Dist2LogitLayer(nn.Module):
-    ''' takes 2 distances, puts through fc layers, spits out value between [0,1] (if use_sigmoid is True) '''
-    def __init__(self, chn_mid=32, use_sigmoid=True):
-        super(Dist2LogitLayer, self).__init__()
-
-        layers = [nn.Conv2d(5, chn_mid, 1, stride=1, padding=0, bias=True),]
-        layers += [nn.LeakyReLU(0.2,True),]
-        layers += [nn.Conv2d(chn_mid, chn_mid, 1, stride=1, padding=0, bias=True),]
-        layers += [nn.LeakyReLU(0.2,True),]
-        layers += [nn.Conv2d(chn_mid, 1, 1, stride=1, padding=0, bias=True),]
-        if(use_sigmoid):
-            layers += [nn.Sigmoid(),]
-        self.model = nn.Sequential(*layers)
-
-    def forward(self,d0,d1,eps=0.1):
-        return self.model.forward(torch.cat((d0,d1,d0-d1,d0/(d1+eps),d1/(d0+eps)),dim=1))
-
-class BCERankingLoss(nn.Module):
-    def __init__(self, chn_mid=32):
-        super(BCERankingLoss, self).__init__()
-        self.net = Dist2LogitLayer(chn_mid=chn_mid)
-        # self.parameters = list(self.net.parameters())
-        self.loss = torch.nn.BCELoss()
-
-    def forward(self, d0, d1, judge):
-        per = (judge+1.)/2.
-        self.logit = self.net.forward(d0,d1)
-        return self.loss(self.logit, per)
-
-# L2, DSSIM metrics
-class FakeNet(nn.Module):
-    def __init__(self, use_gpu=True, colorspace='Lab'):
-        super(FakeNet, self).__init__()
-        self.use_gpu = use_gpu
-        self.colorspace = colorspace
-
-class L2(FakeNet):
-    def forward(self, in0, in1, retPerLayer=None):
-        assert(in0.size()[0]==1) # currently only supports batchSize 1
-
-        if(self.colorspace=='RGB'):
-            (N,C,X,Y) = in0.size()
-            value = torch.mean(torch.mean(torch.mean((in0-in1)**2,dim=1).view(N,1,X,Y),dim=2).view(N,1,1,Y),dim=3).view(N)
-            return value
-        elif(self.colorspace=='Lab'):
-            value = l2(tensor2np(tensor2tensorlab(in0.data,to_norm=False)), 
-                tensor2np(tensor2tensorlab(in1.data,to_norm=False)), range=100.).astype('float')
-            ret_var = Variable( torch.Tensor((value,) ) )
-            if(self.use_gpu):
-                ret_var = ret_var.cuda()
-            return ret_var
-
-class DSSIM(FakeNet):
-
-    def forward(self, in0, in1, retPerLayer=None):
-        assert(in0.size()[0]==1) # currently only supports batchSize 1
-
-        if(self.colorspace=='RGB'):
-            value = dssim(1.*tensor2im(in0.data), 1.*tensor2im(in1.data), range=255.).astype('float')
-        elif(self.colorspace=='Lab'):
-            value = dssim(tensor2np(tensor2tensorlab(in0.data,to_norm=False)), 
-                tensor2np(tensor2tensorlab(in1.data,to_norm=False)), range=100.).astype('float')
-        ret_var = Variable( torch.Tensor((value,) ) )
-        if(self.use_gpu):
-            ret_var = ret_var.cuda()
-        return ret_var
-
-def print_network(net):
-    num_params = 0
-    for param in net.parameters():
-        num_params += param.numel()
-    print('Network',net)
-    print('Total number of parameters: %d' % num_params)
-
-
-class FloLPIPS(LPIPS):
-    def __init__(self, pretrained=True, net='alex', version='0.1', lpips=True, spatial=False, pnet_rand=False, pnet_tune=False, use_dropout=True, model_path=None, eval_mode=True, verbose=False):
-        super(FloLPIPS, self).__init__(pretrained, net, version, lpips, spatial, pnet_rand, pnet_tune, use_dropout, model_path, eval_mode, verbose)
-
-    def forward(self, in0, in1, flow, retPerLayer=False, normalize=False):
-        if normalize: # turn on this flag if input is [0,1] so it can be adjusted to [-1, +1]
-            in0 = 2 * in0  - 1
-            in1 = 2 * in1  - 1
-
-        in0_input, in1_input = (self.scaling_layer(in0), self.scaling_layer(in1)) if self.version=='0.1' else (in0, in1)
-        outs0, outs1 = self.net.forward(in0_input), self.net.forward(in1_input)
-        feats0, feats1, diffs = {}, {}, {}
-
-        for kk in range(self.L):
-            feats0[kk], feats1[kk] = normalize_tensor(outs0[kk]), normalize_tensor(outs1[kk])
-            diffs[kk] = (feats0[kk]-feats1[kk])**2
-
-        res = [mw_spatial_average(self.lins[kk](diffs[kk]), flow, keepdim=True) for kk in range(self.L)]
-
-        return torch.sum(torch.cat(res, 1), dim=(1,2,3), keepdims=False)
-
-
-
-
-
-class Flolpips(nn.Module):
-    def __init__(self):
-        super(Flolpips, self).__init__()
-        self.loss_fn = FloLPIPS(net='alex',version='0.1')
-        self.flownet = PWCNet()
-    
-    @torch.no_grad()
-    def forward(self, I0, I1, frame_dis, frame_ref):
-        """
-        args:
-            I0: first frame of the triplet, shape: [B, C, H, W]
-            I1: third frame of the triplet, shape: [B, C, H, W]
-            frame_dis: prediction of the intermediate frame, shape: [B, C, H, W]
-            frame_ref: ground-truth of the intermediate frame, shape: [B, C, H, W]
-        """
-        assert I0.size() == I1.size() == frame_dis.size() == frame_ref.size(), \
-                "the 4 input tensors should have same size"
-
-        flow_ref = self.flownet(frame_ref, I0)
-        flow_dis = self.flownet(frame_dis, I0)
-        flow_diff = flow_ref - flow_dis
-        flolpips_wrt_I0 = self.loss_fn.forward(frame_ref, frame_dis, flow_diff, normalize=True)
-
-        flow_ref = self.flownet(frame_ref, I1)
-        flow_dis = self.flownet(frame_dis, I1)
-        flow_diff = flow_ref - flow_dis
-        flolpips_wrt_I1 = self.loss_fn.forward(frame_ref, frame_dis, flow_diff, normalize=True)
-
-        flolpips = (flolpips_wrt_I0 + flolpips_wrt_I1) / 2
-        return flolpips

opensora/eval/flolpips/pretrained_networks.py

@@ -1,180 +0,0 @@
-from collections import namedtuple
-import torch
-from torchvision import models as tv
-
-class squeezenet(torch.nn.Module):
-    def __init__(self, requires_grad=False, pretrained=True):
-        super(squeezenet, self).__init__()
-        pretrained_features = tv.squeezenet1_1(pretrained=pretrained).features
-        self.slice1 = torch.nn.Sequential()
-        self.slice2 = torch.nn.Sequential()
-        self.slice3 = torch.nn.Sequential()
-        self.slice4 = torch.nn.Sequential()
-        self.slice5 = torch.nn.Sequential()
-        self.slice6 = torch.nn.Sequential()
-        self.slice7 = torch.nn.Sequential()
-        self.N_slices = 7
-        for x in range(2):
-            self.slice1.add_module(str(x), pretrained_features[x])
-        for x in range(2,5):
-            self.slice2.add_module(str(x), pretrained_features[x])
-        for x in range(5, 8):
-            self.slice3.add_module(str(x), pretrained_features[x])
-        for x in range(8, 10):
-            self.slice4.add_module(str(x), pretrained_features[x])
-        for x in range(10, 11):
-            self.slice5.add_module(str(x), pretrained_features[x])
-        for x in range(11, 12):
-            self.slice6.add_module(str(x), pretrained_features[x])
-        for x in range(12, 13):
-            self.slice7.add_module(str(x), pretrained_features[x])
-        if not requires_grad:
-            for param in self.parameters():
-                param.requires_grad = False
-
-    def forward(self, X):
-        h = self.slice1(X)
-        h_relu1 = h
-        h = self.slice2(h)
-        h_relu2 = h
-        h = self.slice3(h)
-        h_relu3 = h
-        h = self.slice4(h)
-        h_relu4 = h
-        h = self.slice5(h)
-        h_relu5 = h
-        h = self.slice6(h)
-        h_relu6 = h
-        h = self.slice7(h)
-        h_relu7 = h
-        vgg_outputs = namedtuple("SqueezeOutputs", ['relu1','relu2','relu3','relu4','relu5','relu6','relu7'])
-        out = vgg_outputs(h_relu1,h_relu2,h_relu3,h_relu4,h_relu5,h_relu6,h_relu7)
-
-        return out
-
-
-class alexnet(torch.nn.Module):
-    def __init__(self, requires_grad=False, pretrained=True):
-        super(alexnet, self).__init__()
-        alexnet_pretrained_features = tv.alexnet(pretrained=pretrained).features
-        self.slice1 = torch.nn.Sequential()
-        self.slice2 = torch.nn.Sequential()
-        self.slice3 = torch.nn.Sequential()
-        self.slice4 = torch.nn.Sequential()
-        self.slice5 = torch.nn.Sequential()
-        self.N_slices = 5
-        for x in range(2):
-            self.slice1.add_module(str(x), alexnet_pretrained_features[x])
-        for x in range(2, 5):
-            self.slice2.add_module(str(x), alexnet_pretrained_features[x])
-        for x in range(5, 8):
-            self.slice3.add_module(str(x), alexnet_pretrained_features[x])
-        for x in range(8, 10):
-            self.slice4.add_module(str(x), alexnet_pretrained_features[x])
-        for x in range(10, 12):
-            self.slice5.add_module(str(x), alexnet_pretrained_features[x])
-        if not requires_grad:
-            for param in self.parameters():
-                param.requires_grad = False
-
-    def forward(self, X):
-        h = self.slice1(X)
-        h_relu1 = h
-        h = self.slice2(h)
-        h_relu2 = h
-        h = self.slice3(h)
-        h_relu3 = h
-        h = self.slice4(h)
-        h_relu4 = h
-        h = self.slice5(h)
-        h_relu5 = h
-        alexnet_outputs = namedtuple("AlexnetOutputs", ['relu1', 'relu2', 'relu3', 'relu4', 'relu5'])
-        out = alexnet_outputs(h_relu1, h_relu2, h_relu3, h_relu4, h_relu5)
-
-        return out
-
-class vgg16(torch.nn.Module):
-    def __init__(self, requires_grad=False, pretrained=True):
-        super(vgg16, self).__init__()
-        vgg_pretrained_features = tv.vgg16(pretrained=pretrained).features
-        self.slice1 = torch.nn.Sequential()
-        self.slice2 = torch.nn.Sequential()
-        self.slice3 = torch.nn.Sequential()
-        self.slice4 = torch.nn.Sequential()
-        self.slice5 = torch.nn.Sequential()
-        self.N_slices = 5
-        for x in range(4):
-            self.slice1.add_module(str(x), vgg_pretrained_features[x])
-        for x in range(4, 9):
-            self.slice2.add_module(str(x), vgg_pretrained_features[x])
-        for x in range(9, 16):
-            self.slice3.add_module(str(x), vgg_pretrained_features[x])
-        for x in range(16, 23):
-            self.slice4.add_module(str(x), vgg_pretrained_features[x])
-        for x in range(23, 30):
-            self.slice5.add_module(str(x), vgg_pretrained_features[x])
-        if not requires_grad:
-            for param in self.parameters():
-                param.requires_grad = False
-
-    def forward(self, X):
-        h = self.slice1(X)
-        h_relu1_2 = h
-        h = self.slice2(h)
-        h_relu2_2 = h
-        h = self.slice3(h)
-        h_relu3_3 = h
-        h = self.slice4(h)
-        h_relu4_3 = h
-        h = self.slice5(h)
-        h_relu5_3 = h
-        vgg_outputs = namedtuple("VggOutputs", ['relu1_2', 'relu2_2', 'relu3_3', 'relu4_3', 'relu5_3'])
-        out = vgg_outputs(h_relu1_2, h_relu2_2, h_relu3_3, h_relu4_3, h_relu5_3)
-
-        return out
-
-
-
-class resnet(torch.nn.Module):
-    def __init__(self, requires_grad=False, pretrained=True, num=18):
-        super(resnet, self).__init__()
-        if(num==18):
-            self.net = tv.resnet18(pretrained=pretrained)
-        elif(num==34):
-            self.net = tv.resnet34(pretrained=pretrained)
-        elif(num==50):
-            self.net = tv.resnet50(pretrained=pretrained)
-        elif(num==101):
-            self.net = tv.resnet101(pretrained=pretrained)
-        elif(num==152):
-            self.net = tv.resnet152(pretrained=pretrained)
-        self.N_slices = 5
-
-        self.conv1 = self.net.conv1
-        self.bn1 = self.net.bn1
-        self.relu = self.net.relu
-        self.maxpool = self.net.maxpool
-        self.layer1 = self.net.layer1
-        self.layer2 = self.net.layer2
-        self.layer3 = self.net.layer3
-        self.layer4 = self.net.layer4
-
-    def forward(self, X):
-        h = self.conv1(X)
-        h = self.bn1(h)
-        h = self.relu(h)
-        h_relu1 = h
-        h = self.maxpool(h)
-        h = self.layer1(h)
-        h_conv2 = h
-        h = self.layer2(h)
-        h_conv3 = h
-        h = self.layer3(h)
-        h_conv4 = h
-        h = self.layer4(h)
-        h_conv5 = h
-
-        outputs = namedtuple("Outputs", ['relu1','conv2','conv3','conv4','conv5'])
-        out = outputs(h_relu1, h_conv2, h_conv3, h_conv4, h_conv5)
-
-        return out

opensora/eval/flolpips/pwcnet.py

@@ -1,344 +0,0 @@
-#!/usr/bin/env python
-
-import torch
-
-import getopt
-import math
-import numpy
-import os
-import PIL
-import PIL.Image
-import sys
-
-# try:
-from .correlation import correlation # the custom cost volume layer
-# except:
-# 	sys.path.insert(0, './correlation'); import correlation # you should consider upgrading python
-# end
-
-##########################################################
-
-# assert(int(str('').join(torch.__version__.split('.')[0:2])) >= 13) # requires at least pytorch version 1.3.0
-
-# torch.set_grad_enabled(False) # make sure to not compute gradients for computational performance
-
-# torch.backends.cudnn.enabled = True # make sure to use cudnn for computational performance
-
-# ##########################################################
-
-# arguments_strModel = 'default' # 'default', or 'chairs-things'
-# arguments_strFirst = './images/first.png'
-# arguments_strSecond = './images/second.png'
-# arguments_strOut = './out.flo'
-
-# for strOption, strArgument in getopt.getopt(sys.argv[1:], '', [ strParameter[2:] + '=' for strParameter in sys.argv[1::2] ])[0]:
-# 	if strOption == '--model' and strArgument != '': arguments_strModel = strArgument # which model to use
-# 	if strOption == '--first' and strArgument != '': arguments_strFirst = strArgument # path to the first frame
-# 	if strOption == '--second' and strArgument != '': arguments_strSecond = strArgument # path to the second frame
-# 	if strOption == '--out' and strArgument != '': arguments_strOut = strArgument # path to where the output should be stored
-# end
-
-##########################################################
-
-
-
-def backwarp(tenInput, tenFlow):
-	backwarp_tenGrid = {}
-	backwarp_tenPartial = {}
-	if str(tenFlow.shape) not in backwarp_tenGrid:
-		tenHor = torch.linspace(-1.0 + (1.0 / tenFlow.shape[3]), 1.0 - (1.0 / tenFlow.shape[3]), tenFlow.shape[3]).view(1, 1, 1, -1).expand(-1, -1, tenFlow.shape[2], -1)
-		tenVer = torch.linspace(-1.0 + (1.0 / tenFlow.shape[2]), 1.0 - (1.0 / tenFlow.shape[2]), tenFlow.shape[2]).view(1, 1, -1, 1).expand(-1, -1, -1, tenFlow.shape[3])
-
-		backwarp_tenGrid[str(tenFlow.shape)] = torch.cat([ tenHor, tenVer ], 1).cuda()
-	# end
-
-	if str(tenFlow.shape) not in backwarp_tenPartial:
-		backwarp_tenPartial[str(tenFlow.shape)] = tenFlow.new_ones([ tenFlow.shape[0], 1, tenFlow.shape[2], tenFlow.shape[3] ])
-	# end
-
-	tenFlow = torch.cat([ tenFlow[:, 0:1, :, :] / ((tenInput.shape[3] - 1.0) / 2.0), tenFlow[:, 1:2, :, :] / ((tenInput.shape[2] - 1.0) / 2.0) ], 1)
-	tenInput = torch.cat([ tenInput, backwarp_tenPartial[str(tenFlow.shape)] ], 1)
-
-	tenOutput = torch.nn.functional.grid_sample(input=tenInput, grid=(backwarp_tenGrid[str(tenFlow.shape)] + tenFlow).permute(0, 2, 3, 1), mode='bilinear', padding_mode='zeros', align_corners=False)
-
-	tenMask = tenOutput[:, -1:, :, :]; tenMask[tenMask > 0.999] = 1.0; tenMask[tenMask < 1.0] = 0.0
-
-	return tenOutput[:, :-1, :, :] * tenMask
-# end
-
-##########################################################
-
-class Network(torch.nn.Module):
-	def __init__(self):
-		super(Network, self).__init__()
-
-		class Extractor(torch.nn.Module):
-			def __init__(self):
-				super(Extractor, self).__init__()
-
-				self.netOne = torch.nn.Sequential(
-					torch.nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, stride=2, padding=1),
-					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
-					torch.nn.Conv2d(in_channels=16, out_channels=16, kernel_size=3, stride=1, padding=1),
-					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
-					torch.nn.Conv2d(in_channels=16, out_channels=16, kernel_size=3, stride=1, padding=1),
-					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1)
-				)
-
-				self.netTwo = torch.nn.Sequential(
-					torch.nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, stride=2, padding=1),
-					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
-					torch.nn.Conv2d(in_channels=32, out_channels=32, kernel_size=3, stride=1, padding=1),
-					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
-					torch.nn.Conv2d(in_channels=32, out_channels=32, kernel_size=3, stride=1, padding=1),
-					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1)
-				)
-
-				self.netThr = torch.nn.Sequential(
-					torch.nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, stride=2, padding=1),
-					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
-					torch.nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1),
-					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
-					torch.nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1),
-					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1)
-				)
-
-				self.netFou = torch.nn.Sequential(
-					torch.nn.Conv2d(in_channels=64, out_channels=96, kernel_size=3, stride=2, padding=1),
-					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
-					torch.nn.Conv2d(in_channels=96, out_channels=96, kernel_size=3, stride=1, padding=1),
-					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
-					torch.nn.Conv2d(in_channels=96, out_channels=96, kernel_size=3, stride=1, padding=1),
-					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1)
-				)
-
-				self.netFiv = torch.nn.Sequential(
-					torch.nn.Conv2d(in_channels=96, out_channels=128, kernel_size=3, stride=2, padding=1),
-					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
-					torch.nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3, stride=1, padding=1),
-					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
-					torch.nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3, stride=1, padding=1),
-					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1)
-				)
-
-				self.netSix = torch.nn.Sequential(
-					torch.nn.Conv2d(in_channels=128, out_channels=196, kernel_size=3, stride=2, padding=1),
-					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
-					torch.nn.Conv2d(in_channels=196, out_channels=196, kernel_size=3, stride=1, padding=1),
-					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
-					torch.nn.Conv2d(in_channels=196, out_channels=196, kernel_size=3, stride=1, padding=1),
-					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1)
-				)
-			# end
-
-			def forward(self, tenInput):
-				tenOne = self.netOne(tenInput)
-				tenTwo = self.netTwo(tenOne)
-				tenThr = self.netThr(tenTwo)
-				tenFou = self.netFou(tenThr)
-				tenFiv = self.netFiv(tenFou)
-				tenSix = self.netSix(tenFiv)
-
-				return [ tenOne, tenTwo, tenThr, tenFou, tenFiv, tenSix ]
-			# end
-		# end
-
-		class Decoder(torch.nn.Module):
-			def __init__(self, intLevel):
-				super(Decoder, self).__init__()
-
-				intPrevious = [ None, None, 81 + 32 + 2 + 2, 81 + 64 + 2 + 2, 81 + 96 + 2 + 2, 81 + 128 + 2 + 2, 81, None ][intLevel + 1]
-				intCurrent = [ None, None, 81 + 32 + 2 + 2, 81 + 64 + 2 + 2, 81 + 96 + 2 + 2, 81 + 128 + 2 + 2, 81, None ][intLevel + 0]
-
-				if intLevel < 6: self.netUpflow = torch.nn.ConvTranspose2d(in_channels=2, out_channels=2, kernel_size=4, stride=2, padding=1)
-				if intLevel < 6: self.netUpfeat = torch.nn.ConvTranspose2d(in_channels=intPrevious + 128 + 128 + 96 + 64 + 32, out_channels=2, kernel_size=4, stride=2, padding=1)
-				if intLevel < 6: self.fltBackwarp = [ None, None, None, 5.0, 2.5, 1.25, 0.625, None ][intLevel + 1]
-
-				self.netOne = torch.nn.Sequential(
-					torch.nn.Conv2d(in_channels=intCurrent, out_channels=128, kernel_size=3, stride=1, padding=1),
-					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1)
-				)
-
-				self.netTwo = torch.nn.Sequential(
-					torch.nn.Conv2d(in_channels=intCurrent + 128, out_channels=128, kernel_size=3, stride=1, padding=1),
-					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1)
-				)
-
-				self.netThr = torch.nn.Sequential(
-					torch.nn.Conv2d(in_channels=intCurrent + 128 + 128, out_channels=96, kernel_size=3, stride=1, padding=1),
-					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1)
-				)
-
-				self.netFou = torch.nn.Sequential(
-					torch.nn.Conv2d(in_channels=intCurrent + 128 + 128 + 96, out_channels=64, kernel_size=3, stride=1, padding=1),
-					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1)
-				)
-
-				self.netFiv = torch.nn.Sequential(
-					torch.nn.Conv2d(in_channels=intCurrent + 128 + 128 + 96 + 64, out_channels=32, kernel_size=3, stride=1, padding=1),
-					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1)
-				)
-
-				self.netSix = torch.nn.Sequential(
-					torch.nn.Conv2d(in_channels=intCurrent + 128 + 128 + 96 + 64 + 32, out_channels=2, kernel_size=3, stride=1, padding=1)
-				)
-			# end
-
-			def forward(self, tenFirst, tenSecond, objPrevious):
-				tenFlow = None
-				tenFeat = None
-
-				if objPrevious is None:
-					tenFlow = None
-					tenFeat = None
-
-					tenVolume = torch.nn.functional.leaky_relu(input=correlation.FunctionCorrelation(tenFirst=tenFirst, tenSecond=tenSecond), negative_slope=0.1, inplace=False)
-
-					tenFeat = torch.cat([ tenVolume ], 1)
-
-				elif objPrevious is not None:
-					tenFlow = self.netUpflow(objPrevious['tenFlow'])
-					tenFeat = self.netUpfeat(objPrevious['tenFeat'])
-
-					tenVolume = torch.nn.functional.leaky_relu(input=correlation.FunctionCorrelation(tenFirst=tenFirst, tenSecond=backwarp(tenInput=tenSecond, tenFlow=tenFlow * self.fltBackwarp)), negative_slope=0.1, inplace=False)
-
-					tenFeat = torch.cat([ tenVolume, tenFirst, tenFlow, tenFeat ], 1)
-
-				# end
-
-				tenFeat = torch.cat([ self.netOne(tenFeat), tenFeat ], 1)
-				tenFeat = torch.cat([ self.netTwo(tenFeat), tenFeat ], 1)
-				tenFeat = torch.cat([ self.netThr(tenFeat), tenFeat ], 1)
-				tenFeat = torch.cat([ self.netFou(tenFeat), tenFeat ], 1)
-				tenFeat = torch.cat([ self.netFiv(tenFeat), tenFeat ], 1)
-
-				tenFlow = self.netSix(tenFeat)
-
-				return {
-					'tenFlow': tenFlow,
-					'tenFeat': tenFeat
-				}
-			# end
-		# end
-
-		class Refiner(torch.nn.Module):
-			def __init__(self):
-				super(Refiner, self).__init__()
-
-				self.netMain = torch.nn.Sequential(
-					torch.nn.Conv2d(in_channels=81 + 32 + 2 + 2 + 128 + 128 + 96 + 64 + 32, out_channels=128, kernel_size=3, stride=1, padding=1, dilation=1),
-					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
-					torch.nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3, stride=1, padding=2, dilation=2),
-					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
-					torch.nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3, stride=1, padding=4, dilation=4),
-					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
-					torch.nn.Conv2d(in_channels=128, out_channels=96, kernel_size=3, stride=1, padding=8, dilation=8),
-					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
-					torch.nn.Conv2d(in_channels=96, out_channels=64, kernel_size=3, stride=1, padding=16, dilation=16),
-					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
-					torch.nn.Conv2d(in_channels=64, out_channels=32, kernel_size=3, stride=1, padding=1, dilation=1),
-					torch.nn.LeakyReLU(inplace=False, negative_slope=0.1),
-					torch.nn.Conv2d(in_channels=32, out_channels=2, kernel_size=3, stride=1, padding=1, dilation=1)
-				)
-			# end
-
-			def forward(self, tenInput):
-				return self.netMain(tenInput)
-			# end
-		# end
-
-		self.netExtractor = Extractor()
-
-		self.netTwo = Decoder(2)
-		self.netThr = Decoder(3)
-		self.netFou = Decoder(4)
-		self.netFiv = Decoder(5)
-		self.netSix = Decoder(6)
-
-		self.netRefiner = Refiner()
-
-		self.load_state_dict({ strKey.replace('module', 'net'): tenWeight for strKey, tenWeight in torch.hub.load_state_dict_from_url(url='http://content.sniklaus.com/github/pytorch-pwc/network-' + 'default' + '.pytorch').items() })
-	# end
-
-	def forward(self, tenFirst, tenSecond):
-		intWidth = tenFirst.shape[3]
-		intHeight = tenFirst.shape[2]
-
-		intPreprocessedWidth = int(math.floor(math.ceil(intWidth / 64.0) * 64.0))
-		intPreprocessedHeight = int(math.floor(math.ceil(intHeight / 64.0) * 64.0))
-
-		tenPreprocessedFirst = torch.nn.functional.interpolate(input=tenFirst, size=(intPreprocessedHeight, intPreprocessedWidth), mode='bilinear', align_corners=False)
-		tenPreprocessedSecond = torch.nn.functional.interpolate(input=tenSecond, size=(intPreprocessedHeight, intPreprocessedWidth), mode='bilinear', align_corners=False)
-		
-		tenFirst = self.netExtractor(tenPreprocessedFirst)
-		tenSecond = self.netExtractor(tenPreprocessedSecond)
-
-
-		objEstimate = self.netSix(tenFirst[-1], tenSecond[-1], None)
-		objEstimate = self.netFiv(tenFirst[-2], tenSecond[-2], objEstimate)
-		objEstimate = self.netFou(tenFirst[-3], tenSecond[-3], objEstimate)
-		objEstimate = self.netThr(tenFirst[-4], tenSecond[-4], objEstimate)
-		objEstimate = self.netTwo(tenFirst[-5], tenSecond[-5], objEstimate)
-
-		tenFlow = objEstimate['tenFlow'] + self.netRefiner(objEstimate['tenFeat'])
-		tenFlow = 20.0 * torch.nn.functional.interpolate(input=tenFlow, size=(intHeight, intWidth), mode='bilinear', align_corners=False)
-		tenFlow[:, 0, :, :] *= float(intWidth) / float(intPreprocessedWidth)
-		tenFlow[:, 1, :, :] *= float(intHeight) / float(intPreprocessedHeight)
-
-		return tenFlow
-	# end
-# end
-
-netNetwork = None
-
-##########################################################
-
-def estimate(tenFirst, tenSecond):
-	global netNetwork
-
-	if netNetwork is None:
-		netNetwork = Network().cuda().eval()
-	# end
-
-	assert(tenFirst.shape[1] == tenSecond.shape[1])
-	assert(tenFirst.shape[2] == tenSecond.shape[2])
-
-	intWidth = tenFirst.shape[2]
-	intHeight = tenFirst.shape[1]
-
-	assert(intWidth == 1024) # remember that there is no guarantee for correctness, comment this line out if you acknowledge this and want to continue
-	assert(intHeight == 436) # remember that there is no guarantee for correctness, comment this line out if you acknowledge this and want to continue
-
-	tenPreprocessedFirst = tenFirst.cuda().view(1, 3, intHeight, intWidth)
-	tenPreprocessedSecond = tenSecond.cuda().view(1, 3, intHeight, intWidth)
-
-	intPreprocessedWidth = int(math.floor(math.ceil(intWidth / 64.0) * 64.0))
-	intPreprocessedHeight = int(math.floor(math.ceil(intHeight / 64.0) * 64.0))
-
-	tenPreprocessedFirst = torch.nn.functional.interpolate(input=tenPreprocessedFirst, size=(intPreprocessedHeight, intPreprocessedWidth), mode='bilinear', align_corners=False)
-	tenPreprocessedSecond = torch.nn.functional.interpolate(input=tenPreprocessedSecond, size=(intPreprocessedHeight, intPreprocessedWidth), mode='bilinear', align_corners=False)
-
-	tenFlow = 20.0 * torch.nn.functional.interpolate(input=netNetwork(tenPreprocessedFirst, tenPreprocessedSecond), size=(intHeight, intWidth), mode='bilinear', align_corners=False)
-
-	tenFlow[:, 0, :, :] *= float(intWidth) / float(intPreprocessedWidth)
-	tenFlow[:, 1, :, :] *= float(intHeight) / float(intPreprocessedHeight)
-
-	return tenFlow[0, :, :, :].cpu()
-# end
-
-##########################################################
-
-# if __name__ == '__main__':
-# 	tenFirst = torch.FloatTensor(numpy.ascontiguousarray(numpy.array(PIL.Image.open(arguments_strFirst))[:, :, ::-1].transpose(2, 0, 1).astype(numpy.float32) * (1.0 / 255.0)))
-# 	tenSecond = torch.FloatTensor(numpy.ascontiguousarray(numpy.array(PIL.Image.open(arguments_strSecond))[:, :, ::-1].transpose(2, 0, 1).astype(numpy.float32) * (1.0 / 255.0)))
-
-# 	tenOutput = estimate(tenFirst, tenSecond)
-
-# 	objOutput = open(arguments_strOut, 'wb')
-
-# 	numpy.array([ 80, 73, 69, 72 ], numpy.uint8).tofile(objOutput)
-# 	numpy.array([ tenOutput.shape[2], tenOutput.shape[1] ], numpy.int32).tofile(objOutput)
-# 	numpy.array(tenOutput.numpy().transpose(1, 2, 0), numpy.float32).tofile(objOutput)
-
-# 	objOutput.close()
-# end

opensora/eval/flolpips/utils.py

@@ -1,95 +0,0 @@
-import numpy as np
-import cv2
-import torch
-
-
-def normalize_tensor(in_feat,eps=1e-10):
-    norm_factor = torch.sqrt(torch.sum(in_feat**2,dim=1,keepdim=True))
-    return in_feat/(norm_factor+eps)
-
-def l2(p0, p1, range=255.):
-    return .5*np.mean((p0 / range - p1 / range)**2)
-
-def dssim(p0, p1, range=255.):
-    from skimage.measure import compare_ssim
-    return (1 - compare_ssim(p0, p1, data_range=range, multichannel=True)) / 2.
-
-def tensor2im(image_tensor, imtype=np.uint8, cent=1., factor=255./2.):
-    image_numpy = image_tensor[0].cpu().float().numpy()
-    image_numpy = (np.transpose(image_numpy, (1, 2, 0)) + cent) * factor
-    return image_numpy.astype(imtype)
-
-def tensor2np(tensor_obj):
-    # change dimension of a tensor object into a numpy array
-    return tensor_obj[0].cpu().float().numpy().transpose((1,2,0))
-
-def np2tensor(np_obj):
-     # change dimenion of np array into tensor array
-    return torch.Tensor(np_obj[:, :, :, np.newaxis].transpose((3, 2, 0, 1)))
-
-def tensor2tensorlab(image_tensor,to_norm=True,mc_only=False):
-    # image tensor to lab tensor
-    from skimage import color
-
-    img = tensor2im(image_tensor)
-    img_lab = color.rgb2lab(img)
-    if(mc_only):
-        img_lab[:,:,0] = img_lab[:,:,0]-50
-    if(to_norm and not mc_only):
-        img_lab[:,:,0] = img_lab[:,:,0]-50
-        img_lab = img_lab/100.
-
-    return np2tensor(img_lab)
-
-def read_frame_yuv2rgb(stream, width, height, iFrame, bit_depth, pix_fmt='420'):
-    if pix_fmt == '420':
-        multiplier = 1
-        uv_factor = 2
-    elif pix_fmt == '444':
-        multiplier = 2
-        uv_factor = 1
-    else:
-        print('Pixel format {} is not supported'.format(pix_fmt))
-        return
-
-    if bit_depth == 8:
-        datatype = np.uint8
-        stream.seek(iFrame*1.5*width*height*multiplier)
-        Y = np.fromfile(stream, dtype=datatype, count=width*height).reshape((height, width))
-        
-        # read chroma samples and upsample since original is 4:2:0 sampling
-        U = np.fromfile(stream, dtype=datatype, count=(width//uv_factor)*(height//uv_factor)).\
-                                reshape((height//uv_factor, width//uv_factor))
-        V = np.fromfile(stream, dtype=datatype, count=(width//uv_factor)*(height//uv_factor)).\
-                                reshape((height//uv_factor, width//uv_factor))
-
-    else:
-        datatype = np.uint16
-        stream.seek(iFrame*3*width*height*multiplier)
-        Y = np.fromfile(stream, dtype=datatype, count=width*height).reshape((height, width))
-                
-        U = np.fromfile(stream, dtype=datatype, count=(width//uv_factor)*(height//uv_factor)).\
-                                reshape((height//uv_factor, width//uv_factor))
-        V = np.fromfile(stream, dtype=datatype, count=(width//uv_factor)*(height//uv_factor)).\
-                                reshape((height//uv_factor, width//uv_factor))
-
-    if pix_fmt == '420':
-        yuv = np.empty((height*3//2, width), dtype=datatype)
-        yuv[0:height,:] = Y
-
-        yuv[height:height+height//4,:] = U.reshape(-1, width)
-        yuv[height+height//4:,:] = V.reshape(-1, width)
-
-        if bit_depth != 8:
-            yuv = (yuv/(2**bit_depth-1)*255).astype(np.uint8)
-
-        #convert to rgb
-        rgb = cv2.cvtColor(yuv, cv2.COLOR_YUV2RGB_I420)
-    
-    else:
-        yvu = np.stack([Y,V,U],axis=2)
-        if bit_depth != 8:
-            yvu = (yvu/(2**bit_depth-1)*255).astype(np.uint8)
-        rgb = cv2.cvtColor(yvu, cv2.COLOR_YCrCb2RGB)
-
-    return rgb

opensora/eval/fvd/styleganv/fvd.py

@@ -1,90 +0,0 @@
-import torch
-import os
-import math
-import torch.nn.functional as F
-
-# https://github.com/universome/fvd-comparison
-
-
-def load_i3d_pretrained(device=torch.device('cpu')):
-    i3D_WEIGHTS_URL = "https://www.dropbox.com/s/ge9e5ujwgetktms/i3d_torchscript.pt"
-    filepath = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'i3d_torchscript.pt')
-    print(filepath)
-    if not os.path.exists(filepath):
-        print(f"preparing for download {i3D_WEIGHTS_URL}, you can download it by yourself.")
-        os.system(f"wget {i3D_WEIGHTS_URL} -O {filepath}")
-    i3d = torch.jit.load(filepath).eval().to(device)
-    i3d = torch.nn.DataParallel(i3d)
-    return i3d
-    
-
-def get_feats(videos, detector, device, bs=10):
-    # videos : torch.tensor BCTHW [0, 1]
-    detector_kwargs = dict(rescale=False, resize=False, return_features=True) # Return raw features before the softmax layer.
-    feats = np.empty((0, 400))
-    with torch.no_grad():
-        for i in range((len(videos)-1)//bs + 1):
-            feats = np.vstack([feats, detector(torch.stack([preprocess_single(video) for video in videos[i*bs:(i+1)*bs]]).to(device), **detector_kwargs).detach().cpu().numpy()])
-    return feats
-
-
-def get_fvd_feats(videos, i3d, device, bs=10):
-    # videos in [0, 1] as torch tensor BCTHW
-    # videos = [preprocess_single(video) for video in videos]
-    embeddings = get_feats(videos, i3d, device, bs)
-    return embeddings
-
-
-def preprocess_single(video, resolution=224, sequence_length=None):
-    # video: CTHW, [0, 1]
-    c, t, h, w = video.shape
-
-    # temporal crop
-    if sequence_length is not None:
-        assert sequence_length <= t
-        video = video[:, :sequence_length]
-
-    # scale shorter side to resolution
-    scale = resolution / min(h, w)
-    if h < w:
-        target_size = (resolution, math.ceil(w * scale))
-    else:
-        target_size = (math.ceil(h * scale), resolution)
-    video = F.interpolate(video, size=target_size, mode='bilinear', align_corners=False)
-
-    # center crop
-    c, t, h, w = video.shape
-    w_start = (w - resolution) // 2
-    h_start = (h - resolution) // 2
-    video = video[:, :, h_start:h_start + resolution, w_start:w_start + resolution]
-
-    # [0, 1] -> [-1, 1]
-    video = (video - 0.5) * 2
-
-    return video.contiguous()
-
-
-"""
-Copy-pasted from https://github.com/cvpr2022-stylegan-v/stylegan-v/blob/main/src/metrics/frechet_video_distance.py
-"""
-from typing import Tuple
-from scipy.linalg import sqrtm
-import numpy as np
-
-
-def compute_stats(feats: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
-    mu = feats.mean(axis=0) # [d]
-    sigma = np.cov(feats, rowvar=False) # [d, d]
-    return mu, sigma
-
-
-def frechet_distance(feats_fake: np.ndarray, feats_real: np.ndarray) -> float:
-    mu_gen, sigma_gen = compute_stats(feats_fake)
-    mu_real, sigma_real = compute_stats(feats_real)
-    m = np.square(mu_gen - mu_real).sum()
-    if feats_fake.shape[0]>1:
-        s, _ = sqrtm(np.dot(sigma_gen, sigma_real), disp=False) # pylint: disable=no-member
-        fid = np.real(m + np.trace(sigma_gen + sigma_real - s * 2))
-    else:
-        fid = np.real(m)
-    return float(fid)

opensora/eval/fvd/videogpt/fvd.py

@@ -1,137 +0,0 @@
-import torch
-import os
-import math
-import torch.nn.functional as F
-import numpy as np
-import einops
-
-def load_i3d_pretrained(device=torch.device('cpu')):
-    i3D_WEIGHTS_URL = "https://onedrive.live.com/download?cid=78EEF3EB6AE7DBCB&resid=78EEF3EB6AE7DBCB%21199&authkey=AApKdFHPXzWLNyI"
-    filepath = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'i3d_pretrained_400.pt')
-    print(filepath)
-    if not os.path.exists(filepath):
-        print(f"preparing for download {i3D_WEIGHTS_URL}, you can download it by yourself.")
-        os.system(f"wget {i3D_WEIGHTS_URL} -O {filepath}")
-    from .pytorch_i3d import InceptionI3d
-    i3d = InceptionI3d(400, in_channels=3).eval().to(device)
-    i3d.load_state_dict(torch.load(filepath, map_location=device))
-    i3d = torch.nn.DataParallel(i3d)
-    return i3d
-
-def preprocess_single(video, resolution, sequence_length=None):
-    # video: THWC, {0, ..., 255}
-    video = video.permute(0, 3, 1, 2).float() / 255. # TCHW
-    t, c, h, w = video.shape
-
-    # temporal crop
-    if sequence_length is not None:
-        assert sequence_length <= t
-        video = video[:sequence_length]
-
-    # scale shorter side to resolution
-    scale = resolution / min(h, w)
-    if h < w:
-        target_size = (resolution, math.ceil(w * scale))
-    else:
-        target_size = (math.ceil(h * scale), resolution)
-    video = F.interpolate(video, size=target_size, mode='bilinear',
-                          align_corners=False)
-
-    # center crop
-    t, c, h, w = video.shape
-    w_start = (w - resolution) // 2
-    h_start = (h - resolution) // 2
-    video = video[:, :, h_start:h_start + resolution, w_start:w_start + resolution]
-    video = video.permute(1, 0, 2, 3).contiguous() # CTHW
-
-    video -= 0.5
-
-    return video
-
-def preprocess(videos, target_resolution=224):
-    # we should tras videos in [0-1] [b c t h w] as th.float 
-    # -> videos in {0, ..., 255} [b t h w c] as np.uint8 array
-    videos = einops.rearrange(videos, 'b c t h w -> b t h w c')
-    videos = (videos*255).numpy().astype(np.uint8)
-
-    b, t, h, w, c = videos.shape
-    videos = torch.from_numpy(videos)
-    videos = torch.stack([preprocess_single(video, target_resolution) for video in videos])
-    return videos * 2 # [-0.5, 0.5] -> [-1, 1]
-
-def get_fvd_logits(videos, i3d, device, bs=10):
-    videos = preprocess(videos)
-    embeddings = get_logits(i3d, videos, device, bs=10)
-    return embeddings
-
-# https://github.com/tensorflow/gan/blob/de4b8da3853058ea380a6152bd3bd454013bf619/tensorflow_gan/python/eval/classifier_metrics.py#L161
-def _symmetric_matrix_square_root(mat, eps=1e-10):
-    u, s, v = torch.svd(mat)
-    si = torch.where(s < eps, s, torch.sqrt(s))
-    return torch.matmul(torch.matmul(u, torch.diag(si)), v.t())
-
-# https://github.com/tensorflow/gan/blob/de4b8da3853058ea380a6152bd3bd454013bf619/tensorflow_gan/python/eval/classifier_metrics.py#L400
-def trace_sqrt_product(sigma, sigma_v):
-    sqrt_sigma = _symmetric_matrix_square_root(sigma)
-    sqrt_a_sigmav_a = torch.matmul(sqrt_sigma, torch.matmul(sigma_v, sqrt_sigma))
-    return torch.trace(_symmetric_matrix_square_root(sqrt_a_sigmav_a))
-
-# https://discuss.pytorch.org/t/covariance-and-gradient-support/16217/2
-def cov(m, rowvar=False):
-    '''Estimate a covariance matrix given data.
-
-    Covariance indicates the level to which two variables vary together.
-    If we examine N-dimensional samples, `X = [x_1, x_2, ... x_N]^T`,
-    then the covariance matrix element `C_{ij}` is the covariance of
-    `x_i` and `x_j`. The element `C_{ii}` is the variance of `x_i`.
-
-    Args:
-        m: A 1-D or 2-D array containing multiple variables and observations.
-            Each row of `m` represents a variable, and each column a single
-            observation of all those variables.
-        rowvar: If `rowvar` is True, then each row represents a
-            variable, with observations in the columns. Otherwise, the
-            relationship is transposed: each column represents a variable,
-            while the rows contain observations.
-
-    Returns:
-        The covariance matrix of the variables.
-    '''
-    if m.dim() > 2:
-        raise ValueError('m has more than 2 dimensions')
-    if m.dim() < 2:
-        m = m.view(1, -1)
-    if not rowvar and m.size(0) != 1:
-        m = m.t()
-
-    fact = 1.0 / (m.size(1) - 1) # unbiased estimate
-    m -= torch.mean(m, dim=1, keepdim=True)
-    mt = m.t()  # if complex: mt = m.t().conj()
-    return fact * m.matmul(mt).squeeze()
-
-
-def frechet_distance(x1, x2):
-    x1 = x1.flatten(start_dim=1)
-    x2 = x2.flatten(start_dim=1)
-    m, m_w = x1.mean(dim=0), x2.mean(dim=0)
-    sigma, sigma_w = cov(x1, rowvar=False), cov(x2, rowvar=False)
-    mean = torch.sum((m - m_w) ** 2)
-    if x1.shape[0]>1:
-        sqrt_trace_component = trace_sqrt_product(sigma, sigma_w)
-        trace = torch.trace(sigma + sigma_w) - 2.0 * sqrt_trace_component
-        fd = trace + mean
-    else:
-        fd = np.real(mean)
-    return float(fd)
-
-
-def get_logits(i3d, videos, device, bs=10):
-    # assert videos.shape[0] % 16 == 0
-    with torch.no_grad():
-        logits = []
-        for i in range(0, videos.shape[0], bs):
-            batch = videos[i:i + bs].to(device)
-            # logits.append(i3d.module.extract_features(batch)) # wrong
-            logits.append(i3d(batch)) # right
-        logits = torch.cat(logits, dim=0)
-        return logits

opensora/eval/fvd/videogpt/pytorch_i3d.py

@@ -1,322 +0,0 @@
-# Original code from https://github.com/piergiaj/pytorch-i3d
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import numpy as np
-
-class MaxPool3dSamePadding(nn.MaxPool3d):
-
-    def compute_pad(self, dim, s):
-        if s % self.stride[dim] == 0:
-            return max(self.kernel_size[dim] - self.stride[dim], 0)
-        else:
-            return max(self.kernel_size[dim] - (s % self.stride[dim]), 0)
-
-    def forward(self, x):
-        # compute 'same' padding
-        (batch, channel, t, h, w) = x.size()
-        out_t = np.ceil(float(t) / float(self.stride[0]))
-        out_h = np.ceil(float(h) / float(self.stride[1]))
-        out_w = np.ceil(float(w) / float(self.stride[2]))
-        pad_t = self.compute_pad(0, t)
-        pad_h = self.compute_pad(1, h)
-        pad_w = self.compute_pad(2, w)
-
-        pad_t_f = pad_t // 2
-        pad_t_b = pad_t - pad_t_f
-        pad_h_f = pad_h // 2
-        pad_h_b = pad_h - pad_h_f
-        pad_w_f = pad_w // 2
-        pad_w_b = pad_w - pad_w_f
-
-        pad = (pad_w_f, pad_w_b, pad_h_f, pad_h_b, pad_t_f, pad_t_b)
-        x = F.pad(x, pad)
-        return super(MaxPool3dSamePadding, self).forward(x)
-
-
-class Unit3D(nn.Module):
-
-    def __init__(self, in_channels,
-                 output_channels,
-                 kernel_shape=(1, 1, 1),
-                 stride=(1, 1, 1),
-                 padding=0,
-                 activation_fn=F.relu,
-                 use_batch_norm=True,
-                 use_bias=False,
-                 name='unit_3d'):
-
-        """Initializes Unit3D module."""
-        super(Unit3D, self).__init__()
-
-        self._output_channels = output_channels
-        self._kernel_shape = kernel_shape
-        self._stride = stride
-        self._use_batch_norm = use_batch_norm
-        self._activation_fn = activation_fn
-        self._use_bias = use_bias
-        self.name = name
-        self.padding = padding
-
-        self.conv3d = nn.Conv3d(in_channels=in_channels,
-                                out_channels=self._output_channels,
-                                kernel_size=self._kernel_shape,
-                                stride=self._stride,
-                                padding=0, # we always want padding to be 0 here. We will dynamically pad based on input size in forward function
-                                bias=self._use_bias)
-
-        if self._use_batch_norm:
-            self.bn = nn.BatchNorm3d(self._output_channels, eps=1e-5, momentum=0.001)
-
-    def compute_pad(self, dim, s):
-        if s % self._stride[dim] == 0:
-            return max(self._kernel_shape[dim] - self._stride[dim], 0)
-        else:
-            return max(self._kernel_shape[dim] - (s % self._stride[dim]), 0)
-
-
-    def forward(self, x):
-        # compute 'same' padding
-        (batch, channel, t, h, w) = x.size()
-        out_t = np.ceil(float(t) / float(self._stride[0]))
-        out_h = np.ceil(float(h) / float(self._stride[1]))
-        out_w = np.ceil(float(w) / float(self._stride[2]))
-        pad_t = self.compute_pad(0, t)
-        pad_h = self.compute_pad(1, h)
-        pad_w = self.compute_pad(2, w)
-
-        pad_t_f = pad_t // 2
-        pad_t_b = pad_t - pad_t_f
-        pad_h_f = pad_h // 2
-        pad_h_b = pad_h - pad_h_f
-        pad_w_f = pad_w // 2
-        pad_w_b = pad_w - pad_w_f
-
-        pad = (pad_w_f, pad_w_b, pad_h_f, pad_h_b, pad_t_f, pad_t_b)
-        x = F.pad(x, pad)
-
-        x = self.conv3d(x)
-        if self._use_batch_norm:
-            x = self.bn(x)
-        if self._activation_fn is not None:
-            x = self._activation_fn(x)
-        return x
-
-
-
-class InceptionModule(nn.Module):
-    def __init__(self, in_channels, out_channels, name):
-        super(InceptionModule, self).__init__()
-
-        self.b0 = Unit3D(in_channels=in_channels, output_channels=out_channels[0], kernel_shape=[1, 1, 1], padding=0,
-                         name=name+'/Branch_0/Conv3d_0a_1x1')
-        self.b1a = Unit3D(in_channels=in_channels, output_channels=out_channels[1], kernel_shape=[1, 1, 1], padding=0,
-                          name=name+'/Branch_1/Conv3d_0a_1x1')
-        self.b1b = Unit3D(in_channels=out_channels[1], output_channels=out_channels[2], kernel_shape=[3, 3, 3],
-                          name=name+'/Branch_1/Conv3d_0b_3x3')
-        self.b2a = Unit3D(in_channels=in_channels, output_channels=out_channels[3], kernel_shape=[1, 1, 1], padding=0,
-                          name=name+'/Branch_2/Conv3d_0a_1x1')
-        self.b2b = Unit3D(in_channels=out_channels[3], output_channels=out_channels[4], kernel_shape=[3, 3, 3],
-                          name=name+'/Branch_2/Conv3d_0b_3x3')
-        self.b3a = MaxPool3dSamePadding(kernel_size=[3, 3, 3],
-                                stride=(1, 1, 1), padding=0)
-        self.b3b = Unit3D(in_channels=in_channels, output_channels=out_channels[5], kernel_shape=[1, 1, 1], padding=0,
-                          name=name+'/Branch_3/Conv3d_0b_1x1')
-        self.name = name
-
-    def forward(self, x):
-        b0 = self.b0(x)
-        b1 = self.b1b(self.b1a(x))
-        b2 = self.b2b(self.b2a(x))
-        b3 = self.b3b(self.b3a(x))
-        return torch.cat([b0,b1,b2,b3], dim=1)
-
-
-class InceptionI3d(nn.Module):
-    """Inception-v1 I3D architecture.
-    The model is introduced in:
-        Quo Vadis, Action Recognition? A New Model and the Kinetics Dataset
-        Joao Carreira, Andrew Zisserman
-        https://arxiv.org/pdf/1705.07750v1.pdf.
-    See also the Inception architecture, introduced in:
-        Going deeper with convolutions
-        Christian Szegedy, Wei Liu, Yangqing Jia, Pierre Sermanet, Scott Reed,
-        Dragomir Anguelov, Dumitru Erhan, Vincent Vanhoucke, Andrew Rabinovich.
-        http://arxiv.org/pdf/1409.4842v1.pdf.
-    """
-
-    # Endpoints of the model in order. During construction, all the endpoints up
-    # to a designated `final_endpoint` are returned in a dictionary as the
-    # second return value.
-    VALID_ENDPOINTS = (
-        'Conv3d_1a_7x7',
-        'MaxPool3d_2a_3x3',
-        'Conv3d_2b_1x1',
-        'Conv3d_2c_3x3',
-        'MaxPool3d_3a_3x3',
-        'Mixed_3b',
-        'Mixed_3c',
-        'MaxPool3d_4a_3x3',
-        'Mixed_4b',
-        'Mixed_4c',
-        'Mixed_4d',
-        'Mixed_4e',
-        'Mixed_4f',
-        'MaxPool3d_5a_2x2',
-        'Mixed_5b',
-        'Mixed_5c',
-        'Logits',
-        'Predictions',
-    )
-
-    def __init__(self, num_classes=400, spatial_squeeze=True,
-                 final_endpoint='Logits', name='inception_i3d', in_channels=3, dropout_keep_prob=0.5):
-        """Initializes I3D model instance.
-        Args:
-          num_classes: The number of outputs in the logit layer (default 400, which
-              matches the Kinetics dataset).
-          spatial_squeeze: Whether to squeeze the spatial dimensions for the logits
-              before returning (default True).
-          final_endpoint: The model contains many possible endpoints.
-              `final_endpoint` specifies the last endpoint for the model to be built
-              up to. In addition to the output at `final_endpoint`, all the outputs
-              at endpoints up to `final_endpoint` will also be returned, in a
-              dictionary. `final_endpoint` must be one of
-              InceptionI3d.VALID_ENDPOINTS (default 'Logits').
-          name: A string (optional). The name of this module.
-        Raises:
-          ValueError: if `final_endpoint` is not recognized.
-        """
-
-        if final_endpoint not in self.VALID_ENDPOINTS:
-            raise ValueError('Unknown final endpoint %s' % final_endpoint)
-
-        super(InceptionI3d, self).__init__()
-        self._num_classes = num_classes
-        self._spatial_squeeze = spatial_squeeze
-        self._final_endpoint = final_endpoint
-        self.logits = None
-
-        if self._final_endpoint not in self.VALID_ENDPOINTS:
-            raise ValueError('Unknown final endpoint %s' % self._final_endpoint)
-
-        self.end_points = {}
-        end_point = 'Conv3d_1a_7x7'
-        self.end_points[end_point] = Unit3D(in_channels=in_channels, output_channels=64, kernel_shape=[7, 7, 7],
-                                            stride=(2, 2, 2), padding=(3,3,3),  name=name+end_point)
-        if self._final_endpoint == end_point: return
-
-        end_point = 'MaxPool3d_2a_3x3'
-        self.end_points[end_point] = MaxPool3dSamePadding(kernel_size=[1, 3, 3], stride=(1, 2, 2),
-                                                             padding=0)
-        if self._final_endpoint == end_point: return
-
-        end_point = 'Conv3d_2b_1x1'
-        self.end_points[end_point] = Unit3D(in_channels=64, output_channels=64, kernel_shape=[1, 1, 1], padding=0,
-                                       name=name+end_point)
-        if self._final_endpoint == end_point: return
-
-        end_point = 'Conv3d_2c_3x3'
-        self.end_points[end_point] = Unit3D(in_channels=64, output_channels=192, kernel_shape=[3, 3, 3], padding=1,
-                                       name=name+end_point)
-        if self._final_endpoint == end_point: return
-
-        end_point = 'MaxPool3d_3a_3x3'
-        self.end_points[end_point] = MaxPool3dSamePadding(kernel_size=[1, 3, 3], stride=(1, 2, 2),
-                                                             padding=0)
-        if self._final_endpoint == end_point: return
-
-        end_point = 'Mixed_3b'
-        self.end_points[end_point] = InceptionModule(192, [64,96,128,16,32,32], name+end_point)
-        if self._final_endpoint == end_point: return
-
-        end_point = 'Mixed_3c'
-        self.end_points[end_point] = InceptionModule(256, [128,128,192,32,96,64], name+end_point)
-        if self._final_endpoint == end_point: return
-
-        end_point = 'MaxPool3d_4a_3x3'
-        self.end_points[end_point] = MaxPool3dSamePadding(kernel_size=[3, 3, 3], stride=(2, 2, 2),
-                                                             padding=0)
-        if self._final_endpoint == end_point: return
-
-        end_point = 'Mixed_4b'
-        self.end_points[end_point] = InceptionModule(128+192+96+64, [192,96,208,16,48,64], name+end_point)
-        if self._final_endpoint == end_point: return
-
-        end_point = 'Mixed_4c'
-        self.end_points[end_point] = InceptionModule(192+208+48+64, [160,112,224,24,64,64], name+end_point)
-        if self._final_endpoint == end_point: return
-
-        end_point = 'Mixed_4d'
-        self.end_points[end_point] = InceptionModule(160+224+64+64, [128,128,256,24,64,64], name+end_point)
-        if self._final_endpoint == end_point: return
-
-        end_point = 'Mixed_4e'
-        self.end_points[end_point] = InceptionModule(128+256+64+64, [112,144,288,32,64,64], name+end_point)
-        if self._final_endpoint == end_point: return
-
-        end_point = 'Mixed_4f'
-        self.end_points[end_point] = InceptionModule(112+288+64+64, [256,160,320,32,128,128], name+end_point)
-        if self._final_endpoint == end_point: return
-
-        end_point = 'MaxPool3d_5a_2x2'
-        self.end_points[end_point] = MaxPool3dSamePadding(kernel_size=[2, 2, 2], stride=(2, 2, 2),
-                                                             padding=0)
-        if self._final_endpoint == end_point: return
-
-        end_point = 'Mixed_5b'
-        self.end_points[end_point] = InceptionModule(256+320+128+128, [256,160,320,32,128,128], name+end_point)
-        if self._final_endpoint == end_point: return
-
-        end_point = 'Mixed_5c'
-        self.end_points[end_point] = InceptionModule(256+320+128+128, [384,192,384,48,128,128], name+end_point)
-        if self._final_endpoint == end_point: return
-
-        end_point = 'Logits'
-        self.avg_pool = nn.AvgPool3d(kernel_size=[2, 7, 7],
-                                     stride=(1, 1, 1))
-        self.dropout = nn.Dropout(dropout_keep_prob)
-        self.logits = Unit3D(in_channels=384+384+128+128, output_channels=self._num_classes,
-                             kernel_shape=[1, 1, 1],
-                             padding=0,
-                             activation_fn=None,
-                             use_batch_norm=False,
-                             use_bias=True,
-                             name='logits')
-
-        self.build()
-
-
-    def replace_logits(self, num_classes):
-        self._num_classes = num_classes
-        self.logits = Unit3D(in_channels=384+384+128+128, output_channels=self._num_classes,
-                             kernel_shape=[1, 1, 1],
-                             padding=0,
-                             activation_fn=None,
-                             use_batch_norm=False,
-                             use_bias=True,
-                             name='logits')
-
-
-    def build(self):
-        for k in self.end_points.keys():
-            self.add_module(k, self.end_points[k])
-
-    def forward(self, x):
-        for end_point in self.VALID_ENDPOINTS:
-            if end_point in self.end_points:
-                x = self._modules[end_point](x) # use _modules to work with dataparallel
-
-        x = self.logits(self.dropout(self.avg_pool(x)))
-        if self._spatial_squeeze:
-            logits = x.squeeze(3).squeeze(3)
-        logits = logits.mean(dim=2)
-        # logits is batch X time X classes, which is what we want to work with
-        return logits
-
-
-    def extract_features(self, x):
-        for end_point in self.VALID_ENDPOINTS:
-            if end_point in self.end_points:
-                x = self._modules[end_point](x)
-        return self.avg_pool(x)

opensora/eval/script/cal_clip_score.sh

@@ -1,23 +0,0 @@
-# clip_score cross modality
-python eval_clip_score.py \
-    --real_path path/to/image \
-    --generated_path path/to/text \
-    --batch-size 50 \
-    --device "cuda"
-
-# clip_score within the same modality
-python eval_clip_score.py \
-    --real_path path/to/textA \
-    --generated_path path/to/textB \
-    --real_flag txt \
-    --generated_flag txt \
-    --batch-size 50 \
-    --device "cuda"
-
-python eval_clip_score.py \
-    --real_path path/to/imageA \
-    --generated_path path/to/imageB \
-    --real_flag img \
-    --generated_flag img \
-    --batch-size 50 \
-    --device "cuda"

opensora/eval/script/cal_fvd.sh

@@ -1,9 +0,0 @@
-python eval_common_metric.py \
-    --real_video_dir path/to/imageA\
-    --generated_video_dir path/to/imageB \
-    --batch_size 10 \
-    --crop_size 64 \
-    --num_frames 20 \
-    --device 'cuda' \
-    --metric 'fvd' \
-    --fvd_method 'styleganv'

opensora/eval/script/cal_lpips.sh

@@ -1,8 +0,0 @@
-python eval_common_metric.py \
-    --real_video_dir path/to/imageA\
-    --generated_video_dir path/to/imageB \
-    --batch_size 10 \
-    --num_frames 20 \
-    --crop_size 64 \
-    --device 'cuda' \
-    --metric 'lpips'

opensora/eval/script/cal_psnr.sh

@@ -1,9 +0,0 @@
-
-python eval_common_metric.py \
-    --real_video_dir /data/xiaogeng_liu/data/video1 \
-    --generated_video_dir /data/xiaogeng_liu/data/video2 \
-    --batch_size 10 \
-    --num_frames 20 \
-    --crop_size 64 \
-    --device 'cuda' \
-    --metric 'psnr'

opensora/eval/script/cal_ssim.sh

@@ -1,8 +0,0 @@
-python eval_common_metric.py \
-    --real_video_dir /data/xiaogeng_liu/data/video1 \
-    --generated_video_dir /data/xiaogeng_liu/data/video2 \
-    --batch_size 10 \
-    --num_frames 20 \
-    --crop_size 64 \
-    --device 'cuda' \
-    --metric 'ssim'

opensora/models/__init__.py

@@ -1 +0,0 @@
-from .causalvideovae.model import CausalVAEModelWrapper

opensora/models/causalvideovae/__init__.py

@@ -1,32 +0,0 @@
-from torchvision.transforms import Lambda
-from .model.causal_vae import CausalVAEModelWrapper
-
-ae_stride_config = {
-    'CausalVAEModel_D4_2x8x8': [2, 8, 8],
-    'CausalVAEModel_D8_2x8x8': [2, 8, 8],
-    'CausalVAEModel_D4_4x8x8': [4, 8, 8],
-    'CausalVAEModel_D8_4x8x8': [4, 8, 8],
-}
-
-
-ae_channel_config = {
-    'CausalVAEModel_D4_2x8x8': 4,
-    'CausalVAEModel_D8_2x8x8': 8,
-    'CausalVAEModel_D4_4x8x8': 4,
-    'CausalVAEModel_D8_4x8x8': 8,
-}
-
-
-ae_denorm = {
-    'CausalVAEModel_D4_2x8x8': lambda x: (x + 1.) / 2.,
-    'CausalVAEModel_D8_2x8x8': lambda x: (x + 1.) / 2.,
-    'CausalVAEModel_D4_4x8x8': lambda x: (x + 1.) / 2.,
-    'CausalVAEModel_D8_4x8x8': lambda x: (x + 1.) / 2.,
-}
-
-ae_norm = {
-    'CausalVAEModel_D4_2x8x8': Lambda(lambda x: 2. * x - 1.),
-    'CausalVAEModel_D8_2x8x8': Lambda(lambda x: 2. * x - 1.),
-    'CausalVAEModel_D4_4x8x8': Lambda(lambda x: 2. * x - 1.),
-    'CausalVAEModel_D8_4x8x8': Lambda(lambda x: 2. * x - 1.),
-}

opensora/models/causalvideovae/model/__init__.py

@@ -1,4 +0,0 @@
-from .causal_vae import (
-    CausalVAEModel, CausalVAEModelWrapper
-)
-from .ema_model import EMA

opensora/models/causalvideovae/model/causal_vae/__init__.py

@@ -1,26 +0,0 @@
-from .modeling_causalvae import CausalVAEModel
-
-from einops import rearrange
-from torch import nn
-
-class CausalVAEModelWrapper(nn.Module):
-    def __init__(self, model_path, subfolder=None, cache_dir=None, use_ema=False, **kwargs):
-        super(CausalVAEModelWrapper, self).__init__()
-        # if os.path.exists(ckpt):
-        # self.vae = CausalVAEModel.load_from_checkpoint(ckpt)
-        self.vae = CausalVAEModel.from_pretrained(model_path, subfolder=subfolder, cache_dir=cache_dir, **kwargs)
-        if use_ema:
-            self.vae.init_from_ema(model_path)
-            self.vae = self.vae.ema
-    def encode(self, x):  # b c t h w
-        # x = self.vae.encode(x).sample()
-        x = self.vae.encode(x).sample().mul_(0.18215)
-        return x
-    def decode(self, x):
-        # x = self.vae.decode(x)
-        x = self.vae.decode(x / 0.18215)
-        x = rearrange(x, 'b c t h w -> b t c h w').contiguous()
-        return x
-
-    def dtype(self):
-        return self.vae.dtype

opensora/models/causalvideovae/model/causal_vae/modeling_causalvae.py

@@ -1,603 +0,0 @@
-
-from ..modeling_videobase import VideoBaseAE
-from ..modules import Normalize
-from ..modules.ops import nonlinearity
-from typing import Tuple
-import torch.nn as nn
-
-from ..utils.module_utils import resolve_str_to_obj, Module
-from ..utils.distrib_utils import DiagonalGaussianDistribution
-
-
-import torch
-from diffusers.configuration_utils import register_to_config
-from copy import deepcopy
-import os
-
-
-class Encoder(nn.Module):
-    def __init__(
-        self,
-        z_channels: int,
-        hidden_size: int,
-        hidden_size_mult: Tuple[int] = (1, 2, 4, 4),
-        attn_resolutions: Tuple[int] = (16,),
-        conv_in: Module = "Conv2d",
-        conv_out: Module = "CasualConv3d",
-        attention: Module = "AttnBlock",
-        resnet_blocks: Tuple[Module] = (
-            "ResnetBlock2D",
-            "ResnetBlock2D",
-            "ResnetBlock2D",
-            "ResnetBlock3D",
-        ),
-        spatial_downsample: Tuple[Module] = (
-            "Downsample",
-            "Downsample",
-            "Downsample",
-            "",
-        ),
-        temporal_downsample: Tuple[Module] = ("", "", "TimeDownsampleRes2x", ""),
-        mid_resnet: Module = "ResnetBlock3D",
-        dropout: float = 0.0,
-        resolution: int = 256,
-        num_res_blocks: int = 2,
-        double_z: bool = True,
-    ) -> None:
-        super().__init__()
-        assert len(resnet_blocks) == len(hidden_size_mult), print(
-            hidden_size_mult, resnet_blocks
-        )
-        # ---- Config ----
-        self.num_resolutions = len(hidden_size_mult)
-        self.resolution = resolution
-        self.num_res_blocks = num_res_blocks
-        
-        # ---- In ----
-        self.conv_in = resolve_str_to_obj(conv_in)(
-            3, hidden_size, kernel_size=3, stride=1, padding=1
-        )
-
-        # ---- Downsample ----
-        curr_res = resolution
-        in_ch_mult = (1,) + tuple(hidden_size_mult)
-        self.in_ch_mult = in_ch_mult
-        self.down = nn.ModuleList()
-        for i_level in range(self.num_resolutions):
-            block = nn.ModuleList()
-            attn = nn.ModuleList()
-            block_in = hidden_size * in_ch_mult[i_level]
-            block_out = hidden_size * hidden_size_mult[i_level]
-            for i_block in range(self.num_res_blocks):
-                block.append(
-                    resolve_str_to_obj(resnet_blocks[i_level])(
-                        in_channels=block_in,
-                        out_channels=block_out,
-                        dropout=dropout,
-                    )
-                )
-                block_in = block_out
-                if curr_res in attn_resolutions:
-                    attn.append(resolve_str_to_obj(attention)(block_in))
-            down = nn.Module()
-            down.block = block
-            down.attn = attn
-            if spatial_downsample[i_level]:
-                down.downsample = resolve_str_to_obj(spatial_downsample[i_level])(
-                    block_in, block_in
-                )
-                curr_res = curr_res // 2
-            if temporal_downsample[i_level]:
-                down.time_downsample = resolve_str_to_obj(temporal_downsample[i_level])(
-                    block_in, block_in
-                )
-            self.down.append(down)
-
-        # ---- Mid ----
-        self.mid = nn.Module()
-        self.mid.block_1 = resolve_str_to_obj(mid_resnet)(
-            in_channels=block_in,
-            out_channels=block_in,
-            dropout=dropout,
-        )
-        self.mid.attn_1 = resolve_str_to_obj(attention)(block_in)
-        self.mid.block_2 = resolve_str_to_obj(mid_resnet)(
-            in_channels=block_in,
-            out_channels=block_in,
-            dropout=dropout,
-        )
-        # ---- Out ----
-        self.norm_out = Normalize(block_in)
-        self.conv_out = resolve_str_to_obj(conv_out)(
-            block_in,
-            2 * z_channels if double_z else z_channels,
-            kernel_size=3,
-            stride=1,
-            padding=1,
-        )
-
-    def forward(self, x):
-        hs = [self.conv_in(x)]
-        for i_level in range(self.num_resolutions):
-            for i_block in range(self.num_res_blocks):
-                h = self.down[i_level].block[i_block](hs[-1])
-                if len(self.down[i_level].attn) > 0:
-                    h = self.down[i_level].attn[i_block](h)
-                hs.append(h)
-            if hasattr(self.down[i_level], "downsample"):
-                hs.append(self.down[i_level].downsample(hs[-1]))
-            if hasattr(self.down[i_level], "time_downsample"):
-                hs_down = self.down[i_level].time_downsample(hs[-1])
-                hs.append(hs_down)
-
-        h = self.mid.block_1(h)
-        h = self.mid.attn_1(h)
-        h = self.mid.block_2(h)
-        h = self.norm_out(h)
-        h = nonlinearity(h)
-        h = self.conv_out(h)
-        return h
-
-
-class Decoder(nn.Module):
-    def __init__(
-        self,
-        z_channels: int,
-        hidden_size: int,
-        hidden_size_mult: Tuple[int] = (1, 2, 4, 4),
-        attn_resolutions: Tuple[int] = (16,),
-        conv_in: Module = "Conv2d",
-        conv_out: Module = "CasualConv3d",
-        attention: Module = "AttnBlock",
-        resnet_blocks: Tuple[Module] = (
-            "ResnetBlock3D",
-            "ResnetBlock3D",
-            "ResnetBlock3D",
-            "ResnetBlock3D",
-        ),
-        spatial_upsample: Tuple[Module] = (
-            "",
-            "SpatialUpsample2x",
-            "SpatialUpsample2x",
-            "SpatialUpsample2x",
-        ),
-        temporal_upsample: Tuple[Module] = ("", "", "", "TimeUpsampleRes2x"),
-        mid_resnet: Module = "ResnetBlock3D",
-        dropout: float = 0.0,
-        resolution: int = 256,
-        num_res_blocks: int = 2,
-    ):
-        super().__init__()
-        # ---- Config ----
-        self.num_resolutions = len(hidden_size_mult)
-        self.resolution = resolution
-        self.num_res_blocks = num_res_blocks
-
-        # ---- In ----
-        block_in = hidden_size * hidden_size_mult[self.num_resolutions - 1]
-        curr_res = resolution // 2 ** (self.num_resolutions - 1)
-        self.conv_in = resolve_str_to_obj(conv_in)(
-            z_channels, block_in, kernel_size=3, padding=1
-        )
-
-        # ---- Mid ----
-        self.mid = nn.Module()
-        self.mid.block_1 = resolve_str_to_obj(mid_resnet)(
-            in_channels=block_in,
-            out_channels=block_in,
-            dropout=dropout,
-        )
-        self.mid.attn_1 = resolve_str_to_obj(attention)(block_in)
-        self.mid.block_2 = resolve_str_to_obj(mid_resnet)(
-            in_channels=block_in,
-            out_channels=block_in,
-            dropout=dropout,
-        )
-
-        # ---- Upsample ----
-        self.up = nn.ModuleList()
-        for i_level in reversed(range(self.num_resolutions)):
-            block = nn.ModuleList()
-            attn = nn.ModuleList()
-            block_out = hidden_size * hidden_size_mult[i_level]
-            for i_block in range(self.num_res_blocks + 1):
-                block.append(
-                    resolve_str_to_obj(resnet_blocks[i_level])(
-                        in_channels=block_in,
-                        out_channels=block_out,
-                        dropout=dropout,
-                    )
-                )
-                block_in = block_out
-                if curr_res in attn_resolutions:
-                    attn.append(resolve_str_to_obj(attention)(block_in))
-            up = nn.Module()
-            up.block = block
-            up.attn = attn
-            if spatial_upsample[i_level]:
-                up.upsample = resolve_str_to_obj(spatial_upsample[i_level])(
-                    block_in, block_in
-                )
-                curr_res = curr_res * 2
-            if temporal_upsample[i_level]:
-                up.time_upsample = resolve_str_to_obj(temporal_upsample[i_level])(
-                    block_in, block_in
-                )
-            self.up.insert(0, up)
-
-        # ---- Out ----
-        self.norm_out = Normalize(block_in)
-        self.conv_out = resolve_str_to_obj(conv_out)(
-            block_in, 3, kernel_size=3, padding=1
-        )
-
-    def forward(self, z):
-        h = self.conv_in(z)
-        h = self.mid.block_1(h)
-        h = self.mid.attn_1(h)
-        h = self.mid.block_2(h)
-
-        for i_level in reversed(range(self.num_resolutions)):
-            for i_block in range(self.num_res_blocks + 1):
-                h = self.up[i_level].block[i_block](h)
-                if len(self.up[i_level].attn) > 0:
-                    h = self.up[i_level].attn[i_block](h)
-            if hasattr(self.up[i_level], "upsample"):
-                h = self.up[i_level].upsample(h)
-            if hasattr(self.up[i_level], "time_upsample"):
-                h = self.up[i_level].time_upsample(h)
-        h = self.norm_out(h)
-        h = nonlinearity(h)
-        h = self.conv_out(h)
-        return h
-
-
-class CausalVAEModel(VideoBaseAE):
-
-    @register_to_config
-    def __init__(
-        self,
-        hidden_size: int = 128,
-        z_channels: int = 4,
-        hidden_size_mult: Tuple[int] = (1, 2, 4, 4),
-        attn_resolutions: Tuple[int] = [],
-        dropout: float = 0.0,
-        resolution: int = 256,
-        double_z: bool = True,
-        embed_dim: int = 4,
-        num_res_blocks: int = 2,
-        q_conv: str = "CausalConv3d",
-        encoder_conv_in: Module = "CausalConv3d",
-        encoder_conv_out: Module = "CausalConv3d",
-        encoder_attention: Module = "AttnBlock3D",
-        encoder_resnet_blocks: Tuple[Module] = (
-            "ResnetBlock3D",
-            "ResnetBlock3D",
-            "ResnetBlock3D",
-            "ResnetBlock3D",
-        ),
-        encoder_spatial_downsample: Tuple[Module] = (
-            "SpatialDownsample2x",
-            "SpatialDownsample2x",
-            "SpatialDownsample2x",
-            "",
-        ),
-        encoder_temporal_downsample: Tuple[Module] = (
-            "",
-            "TimeDownsample2x",
-            "TimeDownsample2x",
-            "",
-        ),
-        encoder_mid_resnet: Module = "ResnetBlock3D",
-        decoder_conv_in: Module = "CausalConv3d",
-        decoder_conv_out: Module = "CausalConv3d",
-        decoder_attention: Module = "AttnBlock3D",
-        decoder_resnet_blocks: Tuple[Module] = (
-            "ResnetBlock3D",
-            "ResnetBlock3D",
-            "ResnetBlock3D",
-            "ResnetBlock3D",
-        ),
-        decoder_spatial_upsample: Tuple[Module] = (
-            "",
-            "SpatialUpsample2x",
-            "SpatialUpsample2x",
-            "SpatialUpsample2x",
-        ),
-        decoder_temporal_upsample: Tuple[Module] = ("", "", "TimeUpsample2x", "TimeUpsample2x"),
-        decoder_mid_resnet: Module = "ResnetBlock3D",
-        use_quant_layer: bool = True
-    ) -> None:
-        super().__init__()
-        
-        
-        self.tile_sample_min_size = 256
-        self.tile_sample_min_size_t = 33
-        self.tile_latent_min_size = int(self.tile_sample_min_size / (2 ** (len(hidden_size_mult) - 1)))
-        
-        # t_down_ratio = [i for i in encoder_temporal_downsample if len(i) > 0]
-        # self.tile_latent_min_size_t = int((self.tile_sample_min_size_t-1) / (2 ** len(t_down_ratio))) + 1
-        self.tile_latent_min_size_t = 16
-        self.tile_overlap_factor = 0.125
-        self.use_tiling = False
-        
-        self.use_quant_layer = use_quant_layer
-
-        self.encoder = Encoder(
-            z_channels=z_channels,
-            hidden_size=hidden_size,
-            hidden_size_mult=hidden_size_mult,
-            attn_resolutions=attn_resolutions,
-            conv_in=encoder_conv_in,
-            conv_out=encoder_conv_out,
-            attention=encoder_attention,
-            resnet_blocks=encoder_resnet_blocks,
-            spatial_downsample=encoder_spatial_downsample,
-            temporal_downsample=encoder_temporal_downsample,
-            mid_resnet=encoder_mid_resnet,
-            dropout=dropout,
-            resolution=resolution,
-            num_res_blocks=num_res_blocks,
-            double_z=double_z,
-        )
-
-        self.decoder = Decoder(
-            z_channels=z_channels,
-            hidden_size=hidden_size,
-            hidden_size_mult=hidden_size_mult,
-            attn_resolutions=attn_resolutions,
-            conv_in=decoder_conv_in,
-            conv_out=decoder_conv_out,
-            attention=decoder_attention,
-            resnet_blocks=decoder_resnet_blocks,
-            spatial_upsample=decoder_spatial_upsample,
-            temporal_upsample=decoder_temporal_upsample,
-            mid_resnet=decoder_mid_resnet,
-            dropout=dropout,
-            resolution=resolution,
-            num_res_blocks=num_res_blocks,
-        )
-        if self.use_quant_layer:
-            quant_conv_cls = resolve_str_to_obj(q_conv)
-            self.quant_conv = quant_conv_cls(2 * z_channels, 2 * embed_dim, 1)
-            self.post_quant_conv = quant_conv_cls(embed_dim, z_channels, 1)
-
-    def get_encoder(self):
-        if self.use_quant_layer:
-            return [self.quant_conv, self.encoder]
-        return [self.encoder]
-            
-    def get_decoder(self):
-        if self.use_quant_layer:
-            return [self.post_quant_conv, self.decoder]
-        return [self.decoder]
-    
-    def encode(self, x):
-        if self.use_tiling and (
-            x.shape[-1] > self.tile_sample_min_size
-            or x.shape[-2] > self.tile_sample_min_size
-            or x.shape[-3] > self.tile_sample_min_size_t
-        ):
-            return self.tiled_encode(x)
-        h = self.encoder(x)
-        if self.use_quant_layer:
-            h = self.quant_conv(h)
-        posterior = DiagonalGaussianDistribution(h)
-        return posterior
-
-    def decode(self, z):
-        if self.use_tiling and (
-            z.shape[-1] > self.tile_latent_min_size
-            or z.shape[-2] > self.tile_latent_min_size
-            or z.shape[-3] > self.tile_latent_min_size_t
-        ):
-            return self.tiled_decode(z)
-        if self.use_quant_layer:
-            z = self.post_quant_conv(z)
-        dec = self.decoder(z)
-        return dec
-
-    def forward(self, input, sample_posterior=True):
-        posterior = self.encode(input)
-        if sample_posterior:
-            z = posterior.sample()
-        else:
-            z = posterior.mode()
-        dec = self.decode(z)
-        return dec, posterior
-    
-    def on_train_start(self):
-        self.ema = deepcopy(self) if self.save_ema==True else None
-    
-    def get_last_layer(self):
-        if hasattr(self.decoder.conv_out, "conv"):
-            return self.decoder.conv_out.conv.weight
-        else:
-            return self.decoder.conv_out.weight
-
-    def blend_v(
-        self, a: torch.Tensor, b: torch.Tensor, blend_extent: int
-    ) -> torch.Tensor:
-        blend_extent = min(a.shape[3], b.shape[3], blend_extent)
-        for y in range(blend_extent):
-            b[:, :, :, y, :] = a[:, :, :, -blend_extent + y, :] * (
-                1 - y / blend_extent
-            ) + b[:, :, :, y, :] * (y / blend_extent)
-        return b
-
-    def blend_h(
-        self, a: torch.Tensor, b: torch.Tensor, blend_extent: int
-    ) -> torch.Tensor:
-        blend_extent = min(a.shape[4], b.shape[4], blend_extent)
-        for x in range(blend_extent):
-            b[:, :, :, :, x] = a[:, :, :, :, -blend_extent + x] * (
-                1 - x / blend_extent
-            ) + b[:, :, :, :, x] * (x / blend_extent)
-        return b
-
-    def tiled_encode(self, x):
-        t = x.shape[2]
-        t_chunk_idx = [i for i in range(0, t, self.tile_sample_min_size_t-1)]
-        if len(t_chunk_idx) == 1 and t_chunk_idx[0] == 0:
-            t_chunk_start_end = [[0, t]]
-        else:
-            t_chunk_start_end = [[t_chunk_idx[i], t_chunk_idx[i+1]+1] for i in range(len(t_chunk_idx)-1)]
-            if t_chunk_start_end[-1][-1] > t:
-                t_chunk_start_end[-1][-1] = t
-            elif t_chunk_start_end[-1][-1] < t:
-                last_start_end = [t_chunk_idx[-1], t]
-                t_chunk_start_end.append(last_start_end)
-        moments = []
-        for idx, (start, end) in enumerate(t_chunk_start_end):
-            chunk_x = x[:, :, start: end]
-            if idx != 0:
-                moment = self.tiled_encode2d(chunk_x, return_moments=True)[:, :, 1:]
-            else:
-                moment = self.tiled_encode2d(chunk_x, return_moments=True)
-            moments.append(moment)
-        moments = torch.cat(moments, dim=2)
-        posterior = DiagonalGaussianDistribution(moments)
-        return posterior
-    
-    def tiled_decode(self, x):
-        t = x.shape[2]
-        t_chunk_idx = [i for i in range(0, t, self.tile_latent_min_size_t-1)]
-        if len(t_chunk_idx) == 1 and t_chunk_idx[0] == 0:
-            t_chunk_start_end = [[0, t]]
-        else:
-            t_chunk_start_end = [[t_chunk_idx[i], t_chunk_idx[i+1]+1] for i in range(len(t_chunk_idx)-1)]
-            if t_chunk_start_end[-1][-1] > t:
-                t_chunk_start_end[-1][-1] = t
-            elif t_chunk_start_end[-1][-1] < t:
-                last_start_end = [t_chunk_idx[-1], t]
-                t_chunk_start_end.append(last_start_end)
-        dec_ = []
-        for idx, (start, end) in enumerate(t_chunk_start_end):
-            chunk_x = x[:, :, start: end]
-            if idx != 0:
-                dec = self.tiled_decode2d(chunk_x)[:, :, 1:]
-            else:
-                dec = self.tiled_decode2d(chunk_x)
-            dec_.append(dec)
-        dec_ = torch.cat(dec_, dim=2)
-        return dec_
-
-    def tiled_encode2d(self, x, return_moments=False):
-        overlap_size = int(self.tile_sample_min_size * (1 - self.tile_overlap_factor))
-        blend_extent = int(self.tile_latent_min_size * self.tile_overlap_factor)
-        row_limit = self.tile_latent_min_size - blend_extent
-
-        # Split the image into 512x512 tiles and encode them separately.
-        rows = []
-        for i in range(0, x.shape[3], overlap_size):
-            row = []
-            for j in range(0, x.shape[4], overlap_size):
-                tile = x[
-                    :,
-                    :,
-                    :,
-                    i : i + self.tile_sample_min_size,
-                    j : j + self.tile_sample_min_size,
-                ]
-                tile = self.encoder(tile)
-                if self.use_quant_layer:
-                    tile = self.quant_conv(tile)
-                row.append(tile)
-            rows.append(row)
-        result_rows = []
-        for i, row in enumerate(rows):
-            result_row = []
-            for j, tile in enumerate(row):
-                # blend the above tile and the left tile
-                # to the current tile and add the current tile to the result row
-                if i > 0:
-                    tile = self.blend_v(rows[i - 1][j], tile, blend_extent)
-                if j > 0:
-                    tile = self.blend_h(row[j - 1], tile, blend_extent)
-                result_row.append(tile[:, :, :, :row_limit, :row_limit])
-            result_rows.append(torch.cat(result_row, dim=4))
-
-        moments = torch.cat(result_rows, dim=3)
-        posterior = DiagonalGaussianDistribution(moments)
-        if return_moments:
-            return moments
-        return posterior
-
-    def tiled_decode2d(self, z):
-
-        overlap_size = int(self.tile_latent_min_size * (1 - self.tile_overlap_factor))
-        blend_extent = int(self.tile_sample_min_size * self.tile_overlap_factor)
-        row_limit = self.tile_sample_min_size - blend_extent
-
-        # Split z into overlapping 64x64 tiles and decode them separately.
-        # The tiles have an overlap to avoid seams between tiles.
-        rows = []
-        for i in range(0, z.shape[3], overlap_size):
-            row = []
-            for j in range(0, z.shape[4], overlap_size):
-                tile = z[
-                    :,
-                    :,
-                    :,
-                    i : i + self.tile_latent_min_size,
-                    j : j + self.tile_latent_min_size,
-                ]
-                if self.use_quant_layer:
-                    tile = self.post_quant_conv(tile)
-                decoded = self.decoder(tile)
-                row.append(decoded)
-            rows.append(row)
-        result_rows = []
-        for i, row in enumerate(rows):
-            result_row = []
-            for j, tile in enumerate(row):
-                # blend the above tile and the left tile
-                # to the current tile and add the current tile to the result row
-                if i > 0:
-                    tile = self.blend_v(rows[i - 1][j], tile, blend_extent)
-                if j > 0:
-                    tile = self.blend_h(row[j - 1], tile, blend_extent)
-                result_row.append(tile[:, :, :, :row_limit, :row_limit])
-            result_rows.append(torch.cat(result_row, dim=4))
-
-        dec = torch.cat(result_rows, dim=3)
-        return dec
-
-    def enable_tiling(self, use_tiling: bool = True):
-        self.use_tiling = use_tiling
-
-    def disable_tiling(self):
-        self.enable_tiling(False)
-
-    def init_from_ckpt(self, path, ignore_keys=list()):
-        sd = torch.load(path, map_location="cpu", weights_only=True)
-        print("init from " + path)
-        
-        if "ema_state_dict" in sd and len(sd['ema_state_dict']) > 0 and os.environ.get("NOT_USE_EMA_MODEL", 0) == 0:
-            print("Load from ema model!")
-            sd = sd["ema_state_dict"]
-            sd = {key.replace("module.", ""): value for key, value in sd.items()}
-        elif "state_dict" in sd:
-            print("Load from normal model!")
-            if "gen_model" in sd["state_dict"]:
-                sd = sd["state_dict"]["gen_model"]
-            else:
-                sd = sd["state_dict"]
-                
-        keys = list(sd.keys())
-        
-        for k in keys:
-            for ik in ignore_keys:
-                if k.startswith(ik):
-                    print("Deleting key {} from state_dict.".format(k))
-                    del sd[k]
-        
-        miss, unexpected = self.load_state_dict(sd, strict=False)
-        assert len(miss) == 0, f"miss key: {miss}"
-        if len(unexpected) > 0:
-            for i in unexpected:
-                assert 'loss' in i, "unexpected key: {i}"
-
-
-
-

opensora/models/causalvideovae/model/configuration_videobase.py

@@ -1,44 +0,0 @@
-import json
-import yaml
-from typing import TypeVar, Dict, Any
-from diffusers import ConfigMixin
-
-T = TypeVar('T', bound='VideoBaseConfiguration')
-class VideoBaseConfiguration(ConfigMixin):
-    config_name = "VideoBaseConfiguration"
-    _nested_config_fields: Dict[str, Any] = {}
-    
-    def __init__(self, **kwargs):
-        pass
-    
-    def to_dict(self) -> Dict[str, Any]:
-        d = {}
-        for key, value in vars(self).items():
-            if isinstance(value, VideoBaseConfiguration):
-                d[key] = value.to_dict()  # Serialize nested VideoBaseConfiguration instances
-            elif isinstance(value, tuple):
-                d[key] = list(value)
-            else:
-                d[key] = value
-        return d
-    
-    def to_yaml_file(self, yaml_path: str):
-        with open(yaml_path, 'w') as yaml_file:
-            yaml.dump(self.to_dict(), yaml_file, default_flow_style=False)
-    
-    @classmethod
-    def load_from_yaml(cls: T, yaml_path: str) -> T:
-        with open(yaml_path, 'r') as yaml_file:
-            config_dict = yaml.safe_load(yaml_file)
-        for field, field_type in cls._nested_config_fields.items():
-            if field in config_dict:
-                config_dict[field] = field_type.load_from_dict(config_dict[field])
-        return cls(**config_dict)
-
-    @classmethod
-    def load_from_dict(cls: T, config_dict: Dict[str, Any]) -> T:
-        # Process nested configuration objects
-        for field, field_type in cls._nested_config_fields.items():
-            if field in config_dict:
-                config_dict[field] = field_type.load_from_dict(config_dict[field])
-        return cls(**config_dict)

opensora/models/causalvideovae/model/dataset_videobase.py

@@ -1,107 +0,0 @@
-import os.path as osp
-import random
-from glob import glob
-
-from torchvision import transforms
-import numpy as np
-import torch
-import torch.utils.data as data
-import torch.nn.functional as F
-from torchvision.transforms import Lambda
-
-from ..dataset.transform import ToTensorVideo, CenterCropVideo
-from ..utils.dataset_utils import DecordInit
-
-def TemporalRandomCrop(total_frames, size):
-    """
-    Performs a random temporal crop on a video sequence.
-
-    This function randomly selects a continuous frame sequence of length `size` from a video sequence.
-    `total_frames` indicates the total number of frames in the video sequence, and `size` represents the length of the frame sequence to be cropped.
-
-    Parameters:
-    - total_frames (int): The total number of frames in the video sequence.
-    - size (int): The length of the frame sequence to be cropped.
-
-    Returns:
-    - (int, int): A tuple containing two integers. The first integer is the starting frame index of the cropped sequence,
-                  and the second integer is the ending frame index (inclusive) of the cropped sequence.
-    """
-    rand_end = max(0, total_frames - size - 1)
-    begin_index = random.randint(0, rand_end)
-    end_index = min(begin_index + size, total_frames)
-    return begin_index, end_index
-
-def resize(x, resolution):
-    height, width = x.shape[-2:]
-    resolution = min(2 * resolution, height, width)
-    aspect_ratio = width / height
-    if width <= height:
-        new_width = resolution
-        new_height = int(resolution / aspect_ratio)
-    else:
-        new_height = resolution
-        new_width = int(resolution * aspect_ratio)
-    resized_x = F.interpolate(x, size=(new_height, new_width), mode='bilinear', align_corners=True, antialias=True)
-    return resized_x
-
-class VideoDataset(data.Dataset):
-    """ Generic dataset for videos files stored in folders
-    Returns BCTHW videos in the range [-0.5, 0.5] """
-    video_exts = ['avi', 'mp4', 'webm']
-    def __init__(self, video_folder, sequence_length, image_folder=None, train=True, resolution=64, sample_rate=1, dynamic_sample=True):
-
-        self.train = train
-        self.sequence_length = sequence_length
-        self.sample_rate = sample_rate
-        self.resolution = resolution
-        self.v_decoder = DecordInit()
-        self.video_folder = video_folder
-        self.dynamic_sample = dynamic_sample
-
-        self.transform = transforms.Compose([
-            ToTensorVideo(),
-            # Lambda(lambda x: resize(x, self.resolution)),
-            CenterCropVideo(self.resolution),
-            Lambda(lambda x: 2.0 * x - 1.0)
-        ])
-        print('Building datasets...')
-        self.samples = self._make_dataset()
-
-    def _make_dataset(self):
-        samples = []
-        samples += sum([glob(osp.join(self.video_folder, '**', f'*.{ext}'), recursive=True)
-                            for ext in self.video_exts], [])
-        return samples
-
-    def __len__(self):
-        return len(self.samples)
-
-    def __getitem__(self, idx):
-        video_path = self.samples[idx]
-        try:
-            video = self.decord_read(video_path)
-            video = self.transform(video)  # T C H W -> T C H W
-            video = video.transpose(0, 1)  # T C H W -> C T H W
-            return dict(video=video, label="")
-        except Exception as e:
-            print(f'Error with {e}, {video_path}')
-            return self.__getitem__(random.randint(0, self.__len__()-1))
-
-    def decord_read(self, path):
-        decord_vr = self.v_decoder(path)
-        total_frames = len(decord_vr)
-        # Sampling video frames
-        if self.dynamic_sample:
-            sample_rate = random.randint(1, self.sample_rate)
-        else:
-            sample_rate = self.sample_rate
-        size = self.sequence_length * sample_rate
-        start_frame_ind, end_frame_ind = TemporalRandomCrop(total_frames, size)
-        # assert end_frame_ind - start_frame_ind >= self.num_frames
-        frame_indice = np.linspace(start_frame_ind, end_frame_ind - 1, self.sequence_length, dtype=int)
-
-        video_data = decord_vr.get_batch(frame_indice).asnumpy()
-        video_data = torch.from_numpy(video_data)
-        video_data = video_data.permute(0, 3, 1, 2)  # (T, H, W, C) -> (T C H W)
-        return video_data

opensora/models/causalvideovae/model/ema_model.py

@@ -1,31 +0,0 @@
-class EMA:
-    def __init__(self, model, decay):
-        self.model = model
-        self.decay = decay
-        self.shadow = {}
-        self.backup = {}
-        
-    def register(self):
-        for name, param in self.model.named_parameters():
-            if param.requires_grad:
-                self.shadow[name] = param.data.clone()
-
-    def update(self):
-        for name, param in self.model.named_parameters():
-            if name in self.shadow:
-                new_average = (1.0 - self.decay) * param.data + self.decay * self.shadow[name]
-                self.shadow[name] = new_average.clone()
-
-    def apply_shadow(self):
-        for name, param in self.model.named_parameters():
-            if name in self.shadow:
-                self.backup[name] = param.data
-                param.data = self.shadow[name]
-
-    def restore(self):
-        for name, param in self.model.named_parameters():
-            if name in self.shadow:
-                param.data = self.backup[name]
-        self.backup = {}
-        
-        

opensora/models/causalvideovae/model/modeling_videobase.py

@@ -1,53 +0,0 @@
-import torch
-from diffusers import ModelMixin, ConfigMixin
-from torch import nn
-import os
-import json
-import pytorch_lightning as pl
-from diffusers.configuration_utils import ConfigMixin
-from diffusers.models.modeling_utils import ModelMixin
-from typing import Optional, Union
-import glob
-
-
-class VideoBaseAE(ModelMixin, ConfigMixin):
-    config_name = "config.json"
-    
-    def __init__(self, *args, **kwargs) -> None:
-        super().__init__(*args, **kwargs)
-    
-    def encode(self, x: torch.Tensor, *args, **kwargs):
-        pass
-
-    def decode(self, encoding: torch.Tensor, *args, **kwargs):
-        pass
-    
-    @property
-    def num_training_steps(self) -> int:
-        """Total training steps inferred from datamodule and devices."""
-        if self.trainer.max_steps:
-            return self.trainer.max_steps
-    
-        limit_batches = self.trainer.limit_train_batches
-        batches = len(self.train_dataloader())
-        batches = min(batches, limit_batches) if isinstance(limit_batches, int) else int(limit_batches * batches)     
-    
-        num_devices = max(1, self.trainer.num_gpus, self.trainer.num_processes)
-        if self.trainer.tpu_cores:
-            num_devices = max(num_devices, self.trainer.tpu_cores)
-    
-        effective_accum = self.trainer.accumulate_grad_batches * num_devices
-        return (batches // effective_accum) * self.trainer.max_epochs
-    
-    @classmethod
-    def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], **kwargs):
-        ckpt_files = glob.glob(os.path.join(pretrained_model_name_or_path, '*.ckpt'))
-        if ckpt_files:
-            # Adapt to checkpoint
-            last_ckpt_file = ckpt_files[-1]
-            config_file = os.path.join(pretrained_model_name_or_path, cls.config_name)
-            model = cls.from_config(config_file)
-            model.init_from_ckpt(last_ckpt_file)
-            return model
-        else:
-            return super().from_pretrained(pretrained_model_name_or_path, **kwargs)

opensora/models/causalvideovae/model/modules/__init__.py

@@ -1,6 +0,0 @@
-from .block import Block
-from .attention import *
-from .conv import *
-from .normalize import *
-from .resnet_block import *
-from .updownsample import *

opensora/models/causalvideovae/model/modules/attention.py

@@ -1,230 +0,0 @@
-import torch.nn as nn
-from .normalize import Normalize
-from .conv import CausalConv3d
-import torch
-import numpy as np
-from einops import rearrange
-from .block import Block
-from .ops import video_to_image
-
-class LinearAttention(Block):
-    def __init__(self, dim, heads=4, dim_head=32):
-        super().__init__()
-        self.heads = heads
-        hidden_dim = dim_head * heads
-        self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias=False)
-        self.to_out = nn.Conv2d(hidden_dim, dim, 1)
-
-    def forward(self, x):
-        b, c, h, w = x.shape
-        qkv = self.to_qkv(x)
-        q, k, v = rearrange(
-            qkv, "b (qkv heads c) h w -> qkv b heads c (h w)", heads=self.heads, qkv=3
-        )
-        k = k.softmax(dim=-1)
-        context = torch.einsum("bhdn,bhen->bhde", k, v)
-        out = torch.einsum("bhde,bhdn->bhen", context, q)
-        out = rearrange(
-            out, "b heads c (h w) -> b (heads c) h w", heads=self.heads, h=h, w=w
-        )
-        return self.to_out(out)
-
-
-class LinAttnBlock(LinearAttention):
-    """to match AttnBlock usage"""
-
-    def __init__(self, in_channels):
-        super().__init__(dim=in_channels, heads=1, dim_head=in_channels)
-
-
-class AttnBlock3D(Block):
-    """Compatible with old versions, there are issues, use with caution."""
-    def __init__(self, in_channels):
-        super().__init__()
-        self.in_channels = in_channels
-
-        self.norm = Normalize(in_channels)
-        self.q = CausalConv3d(in_channels, in_channels, kernel_size=1, stride=1)
-        self.k = CausalConv3d(in_channels, in_channels, kernel_size=1, stride=1)
-        self.v = CausalConv3d(in_channels, in_channels, kernel_size=1, stride=1)
-        self.proj_out = CausalConv3d(in_channels, in_channels, kernel_size=1, stride=1)
-
-    def forward(self, x):
-        h_ = x
-        h_ = self.norm(h_)
-        q = self.q(h_)
-        k = self.k(h_)
-        v = self.v(h_)
-
-        # compute attention
-        b, c, t, h, w = q.shape
-        q = q.reshape(b * t, c, h * w)
-        q = q.permute(0, 2, 1)  # b,hw,c
-        k = k.reshape(b * t, c, h * w)  # b,c,hw
-        w_ = torch.bmm(q, k)  # b,hw,hw    w[b,i,j]=sum_c q[b,i,c]k[b,c,j]
-        w_ = w_ * (int(c) ** (-0.5))
-        w_ = torch.nn.functional.softmax(w_, dim=2)
-
-        # attend to values
-        v = v.reshape(b * t, c, h * w)
-        w_ = w_.permute(0, 2, 1)  # b,hw,hw (first hw of k, second of q)
-        h_ = torch.bmm(v, w_)  # b, c,hw (hw of q) h_[b,c,j] = sum_i v[b,c,i] w_[b,i,j]
-        h_ = h_.reshape(b, c, t, h, w)
-
-        h_ = self.proj_out(h_)
-
-        return x + h_
-
-class AttnBlock3DFix(nn.Module):
-    """
-    Thanks to https://github.com/PKU-YuanGroup/Open-Sora-Plan/pull/172.
-    """
-    def __init__(self, in_channels):
-        super().__init__()
-        self.in_channels = in_channels
-
-        self.norm = Normalize(in_channels)
-        self.q = CausalConv3d(in_channels, in_channels, kernel_size=1, stride=1)
-        self.k = CausalConv3d(in_channels, in_channels, kernel_size=1, stride=1)
-        self.v = CausalConv3d(in_channels, in_channels, kernel_size=1, stride=1)
-        self.proj_out = CausalConv3d(in_channels, in_channels, kernel_size=1, stride=1)
-
-    def forward(self, x):
-        h_ = x
-        h_ = self.norm(h_)
-        q = self.q(h_)
-        k = self.k(h_)
-        v = self.v(h_)
-
-        # compute attention
-        # q: (b c t h w) -> (b t c h w) -> (b*t c h*w) -> (b*t h*w c)
-        b, c, t, h, w = q.shape
-        q = q.permute(0, 2, 1, 3, 4)
-        q = q.reshape(b * t, c, h * w)
-        q = q.permute(0, 2, 1)
-
-        # k: (b c t h w) -> (b t c h w) -> (b*t c h*w)
-        k = k.permute(0, 2, 1, 3, 4)
-        k = k.reshape(b * t, c, h * w)
-
-        # w: (b*t hw hw)
-        w_ = torch.bmm(q, k)
-        w_ = w_ * (int(c) ** (-0.5))
-        w_ = torch.nn.functional.softmax(w_, dim=2)
-
-        # attend to values
-        # v: (b c t h w) -> (b t c h w) -> (bt c hw)
-        # w_: (bt hw hw) -> (bt hw hw)
-        v = v.permute(0, 2, 1, 3, 4)
-        v = v.reshape(b * t, c, h * w)
-        w_ = w_.permute(0, 2, 1)  # b,hw,hw (first hw of k, second of q)
-        h_ = torch.bmm(v, w_)  # b, c,hw (hw of q) h_[b,c,j] = sum_i v[b,c,i] w_[b,i,j]
-
-        # h_: (b*t c hw) -> (b t c h w) -> (b c t h w)
-        h_ = h_.reshape(b, t, c, h, w)
-        h_ = h_.permute(0, 2, 1, 3 ,4)
-
-        h_ = self.proj_out(h_)
-
-        return x + h_
-
-
-class AttnBlock(Block):
-    def __init__(self, in_channels):
-        super().__init__()
-        self.in_channels = in_channels
-
-        self.norm = Normalize(in_channels)
-        self.q = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
-        self.k = torch.nn.Conv2d(
-            in_channels, in_channels, kernel_size=1, stride=1, padding=0
-        )
-        self.v = torch.nn.Conv2d(
-            in_channels, in_channels, kernel_size=1, stride=1, padding=0
-        )
-        self.proj_out = torch.nn.Conv2d(
-            in_channels, in_channels, kernel_size=1, stride=1, padding=0
-        )
-        
-    @video_to_image
-    def forward(self, x):
-        h_ = x
-        h_ = self.norm(h_)
-        q = self.q(h_)
-        k = self.k(h_)
-        v = self.v(h_)
-
-        # compute attention
-        b, c, h, w = q.shape
-        q = q.reshape(b, c, h * w)
-        q = q.permute(0, 2, 1)  # b,hw,c
-        k = k.reshape(b, c, h * w)  # b,c,hw
-        w_ = torch.bmm(q, k)  # b,hw,hw    w[b,i,j]=sum_c q[b,i,c]k[b,c,j]
-        w_ = w_ * (int(c) ** (-0.5))
-        w_ = torch.nn.functional.softmax(w_, dim=2)
-
-        # attend to values
-        v = v.reshape(b, c, h * w)
-        w_ = w_.permute(0, 2, 1)  # b,hw,hw (first hw of k, second of q)
-        h_ = torch.bmm(v, w_)  # b, c,hw (hw of q) h_[b,c,j] = sum_i v[b,c,i] w_[b,i,j]
-        h_ = h_.reshape(b, c, h, w)
-
-        h_ = self.proj_out(h_)
-
-        return x + h_
-
-
-class TemporalAttnBlock(Block):
-    def __init__(self, in_channels):
-        super().__init__()
-        self.in_channels = in_channels
-
-        self.norm = Normalize(in_channels)
-        self.q = torch.nn.Conv3d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
-        self.k = torch.nn.Conv3d(
-            in_channels, in_channels, kernel_size=1, stride=1, padding=0
-        )
-        self.v = torch.nn.Conv3d(
-            in_channels, in_channels, kernel_size=1, stride=1, padding=0
-        )
-        self.proj_out = torch.nn.Conv3d(
-            in_channels, in_channels, kernel_size=1, stride=1, padding=0
-        )
-        
-    def forward(self, x):
-        h_ = x
-        h_ = self.norm(h_)
-        q = self.q(h_)
-        k = self.k(h_)
-        v = self.v(h_)
-
-        # compute attention
-        b, c, t, h, w = q.shape
-        q = rearrange(q, "b c t h w -> (b h w) t c")
-        k = rearrange(k, "b c t h w -> (b h w) c t")
-        v = rearrange(v, "b c t h w -> (b h w) c t")
-        w_ = torch.bmm(q, k)
-        w_ = w_ * (int(c) ** (-0.5))
-        w_ = torch.nn.functional.softmax(w_, dim=2)
-
-        # attend to values
-        w_ = w_.permute(0, 2, 1)
-        h_ = torch.bmm(v, w_)
-        h_ = rearrange(h_, "(b h w) c t -> b c t h w", h=h, w=w)
-        h_ = self.proj_out(h_)
-
-        return x + h_
-
-
-def make_attn(in_channels, attn_type="vanilla"):
-    assert attn_type in ["vanilla", "linear", "none", "vanilla3D"], f"attn_type {attn_type} unknown"
-    print(f"making attention of type '{attn_type}' with {in_channels} in_channels")
-    print(attn_type)
-    if attn_type == "vanilla":
-        return AttnBlock(in_channels)
-    elif attn_type == "vanilla3D":
-        return AttnBlock3D(in_channels)
-    elif attn_type == "none":
-        return nn.Identity(in_channels)
-    else:
-        return LinAttnBlock(in_channels)

opensora/models/causalvideovae/model/modules/block.py

@@ -1,5 +0,0 @@
-import torch.nn as nn
-
-class Block(nn.Module):
-    def __init__(self, *args, **kwargs) -> None:
-        super().__init__(*args, **kwargs)

opensora/models/causalvideovae/model/modules/conv.py

@@ -1,111 +0,0 @@
-import torch.nn as nn
-from typing import Union, Tuple
-import torch.nn.functional as F
-import torch
-from .ops import cast_tuple
-from .ops import video_to_image
-from torch.utils.checkpoint import checkpoint
-
-class Conv2d(nn.Conv2d):
-    def __init__(
-        self,
-        in_channels: int,
-        out_channels: int,
-        kernel_size: Union[int, Tuple[int]] = 3,
-        stride: Union[int, Tuple[int]] = 1,
-        padding: Union[str, int, Tuple[int]] = 0,
-        dilation: Union[int, Tuple[int]] = 1,
-        groups: int = 1,
-        bias: bool = True,
-        padding_mode: str = "zeros",
-        device=None,
-        dtype=None,
-    ) -> None:
-        super().__init__(
-            in_channels,
-            out_channels,
-            kernel_size,
-            stride,
-            padding,
-            dilation,
-            groups,
-            bias,
-            padding_mode,
-            device,
-            dtype,
-        )
-        
-    @video_to_image
-    def forward(self, x):
-        return super().forward(x)
-        
-
-
-class CausalConv3d(nn.Module):
-    def __init__(
-        self, chan_in, chan_out, kernel_size: Union[int, Tuple[int, int, int]], init_method="random", **kwargs
-    ):
-        super().__init__()
-        self.kernel_size = cast_tuple(kernel_size, 3)
-        self.time_kernel_size = self.kernel_size[0]
-        self.chan_in = chan_in
-        self.chan_out = chan_out
-        stride = kwargs.pop("stride", 1)
-        padding = kwargs.pop("padding", 0)
-        padding = list(cast_tuple(padding, 3))
-        padding[0] = 0
-        stride = cast_tuple(stride, 3)
-        self.conv = nn.Conv3d(chan_in, chan_out, self.kernel_size, stride=stride, padding=padding)
-        self.pad = nn.ReplicationPad2d((0, 0, self.time_kernel_size - 1, 0))
-        self._init_weights(init_method)
-        
-    def _init_weights(self, init_method):
-        ks = torch.tensor(self.kernel_size)
-        if init_method == "avg":
-            assert (
-                self.kernel_size[1] == 1 and self.kernel_size[2] == 1
-            ), "only support temporal up/down sample"
-            assert self.chan_in == self.chan_out, "chan_in must be equal to chan_out"
-            weight = torch.zeros((self.chan_out, self.chan_in, *self.kernel_size))
-
-            eyes = torch.concat(
-                [
-                    torch.eye(self.chan_in).unsqueeze(-1) * 1/3,
-                    torch.eye(self.chan_in).unsqueeze(-1) * 1/3,
-                    torch.eye(self.chan_in).unsqueeze(-1) * 1/3,
-                ],
-                dim=-1,
-            )
-            weight[:, :, :, 0, 0] = eyes
-
-            self.conv.weight = nn.Parameter(
-                weight,
-                requires_grad=True,
-            )
-        elif init_method == "zero":
-            self.conv.weight = nn.Parameter(
-                torch.zeros((self.chan_out, self.chan_in, *self.kernel_size)),
-                requires_grad=True,
-            )
-        if self.conv.bias is not None:
-            nn.init.constant_(self.conv.bias, 0)
-            
-    def forward(self, x):
-        # 1 + 16   16 as video, 1 as image
-        first_frame_pad = x[:, :, :1, :, :].repeat(
-            (1, 1, self.time_kernel_size - 1, 1, 1)
-        )  # b c t h w
-        x = torch.concatenate((first_frame_pad, x), dim=2)  # 3 + 16
-        return self.conv(x)
-    
-    
-class CausalConv3d_GC(CausalConv3d):
-    def __init__(self, chan_in, chan_out, kernel_size: Union[int, Tuple[int]], init_method="random", **kwargs):
-        super().__init__(chan_in, chan_out, kernel_size, init_method, **kwargs)
-    def forward(self, x):
-        # 1 + 16   16 as video, 1 as image
-        first_frame_pad = x[:, :, :1, :, :].repeat(
-            (1, 1, self.time_kernel_size - 1, 1, 1)
-        )   # b c t h w
-        x = torch.concatenate((first_frame_pad, x), dim=2)  # 3 + 16
-        return checkpoint(self.conv, x)

opensora/models/causalvideovae/model/modules/normalize.py

@@ -1,101 +0,0 @@
-import torch
-import torch.nn as nn
-from .block import Block
-
-class GroupNorm(Block):
-    def __init__(self, num_channels, num_groups=32, eps=1e-6, *args, **kwargs) -> None:
-        super().__init__(*args, **kwargs)
-        self.norm = torch.nn.GroupNorm(
-            num_groups=num_groups, num_channels=num_channels, eps=1e-6, affine=True
-        )
-    def forward(self, x):
-        return self.norm(x)
-
-def Normalize(in_channels, num_groups=32):
-    return torch.nn.GroupNorm(
-        num_groups=num_groups, num_channels=in_channels, eps=1e-6, affine=True
-    )
-
-class ActNorm(nn.Module):
-    def __init__(self, num_features, logdet=False, affine=True,
-                 allow_reverse_init=False):
-        assert affine
-        super().__init__()
-        self.logdet = logdet
-        self.loc = nn.Parameter(torch.zeros(1, num_features, 1, 1))
-        self.scale = nn.Parameter(torch.ones(1, num_features, 1, 1))
-        self.allow_reverse_init = allow_reverse_init
-
-        self.register_buffer('initialized', torch.tensor(0, dtype=torch.uint8))
-
-    def initialize(self, input):
-        with torch.no_grad():
-            flatten = input.permute(1, 0, 2, 3).contiguous().view(input.shape[1], -1)
-            mean = (
-                flatten.mean(1)
-                .unsqueeze(1)
-                .unsqueeze(2)
-                .unsqueeze(3)
-                .permute(1, 0, 2, 3)
-            )
-            std = (
-                flatten.std(1)
-                .unsqueeze(1)
-                .unsqueeze(2)
-                .unsqueeze(3)
-                .permute(1, 0, 2, 3)
-            )
-
-            self.loc.data.copy_(-mean)
-            self.scale.data.copy_(1 / (std + 1e-6))
-
-    def forward(self, input, reverse=False):
-        if reverse:
-            return self.reverse(input)
-        if len(input.shape) == 2:
-            input = input[:,:,None,None]
-            squeeze = True
-        else:
-            squeeze = False
-
-        _, _, height, width = input.shape
-
-        if self.training and self.initialized.item() == 0:
-            self.initialize(input)
-            self.initialized.fill_(1)
-
-        h = self.scale * (input + self.loc)
-
-        if squeeze:
-            h = h.squeeze(-1).squeeze(-1)
-
-        if self.logdet:
-            log_abs = torch.log(torch.abs(self.scale))
-            logdet = height*width*torch.sum(log_abs)
-            logdet = logdet * torch.ones(input.shape[0]).to(input)
-            return h, logdet
-
-        return h
-
-    def reverse(self, output):
-        if self.training and self.initialized.item() == 0:
-            if not self.allow_reverse_init:
-                raise RuntimeError(
-                    "Initializing ActNorm in reverse direction is "
-                    "disabled by default. Use allow_reverse_init=True to enable."
-                )
-            else:
-                self.initialize(output)
-                self.initialized.fill_(1)
-
-        if len(output.shape) == 2:
-            output = output[:,:,None,None]
-            squeeze = True
-        else:
-            squeeze = False
-
-        h = output / self.scale - self.loc
-
-        if squeeze:
-            h = h.squeeze(-1).squeeze(-1)
-        return h

opensora/models/causalvideovae/model/modules/ops.py

@@ -1,40 +0,0 @@
-import torch
-from einops import rearrange
-
-def video_to_image(func):
-    def wrapper(self, x, *args, **kwargs):
-        if x.dim() == 5:
-            t = x.shape[2]
-            x = rearrange(x, "b c t h w -> (b t) c h w")
-            x = func(self, x, *args, **kwargs)
-            x = rearrange(x, "(b t) c h w -> b c t h w", t=t)
-        return x
-    return wrapper
-
-def nonlinearity(x):
-    return x * torch.sigmoid(x)
-
-def cast_tuple(t, length=1):
-    return t if isinstance(t, tuple) or isinstance(t, list) else ((t,) * length)
-
-def shift_dim(x, src_dim=-1, dest_dim=-1, make_contiguous=True):
-    n_dims = len(x.shape)
-    if src_dim < 0:
-        src_dim = n_dims + src_dim
-    if dest_dim < 0:
-        dest_dim = n_dims + dest_dim
-    assert 0 <= src_dim < n_dims and 0 <= dest_dim < n_dims
-    dims = list(range(n_dims))
-    del dims[src_dim]
-    permutation = []
-    ctr = 0
-    for i in range(n_dims):
-        if i == dest_dim:
-            permutation.append(src_dim)
-        else:
-            permutation.append(dims[ctr])
-            ctr += 1
-    x = x.permute(permutation)
-    if make_contiguous:
-        x = x.contiguous()
-    return x

opensora/models/causalvideovae/model/modules/quant.py

@@ -1,100 +0,0 @@
-import torch
-import torch.nn as nn
-import torch.distributed as dist
-import numpy as np
-import torch.nn.functional as F
-from .ops import shift_dim
-
-class Codebook(nn.Module):
-    def __init__(self, n_codes, embedding_dim):
-        super().__init__()
-        self.register_buffer("embeddings", torch.randn(n_codes, embedding_dim))
-        self.register_buffer("N", torch.zeros(n_codes))
-        self.register_buffer("z_avg", self.embeddings.data.clone())
-
-        self.n_codes = n_codes
-        self.embedding_dim = embedding_dim
-        self._need_init = True
-
-    def _tile(self, x):
-        d, ew = x.shape
-        if d < self.n_codes:
-            n_repeats = (self.n_codes + d - 1) // d
-            std = 0.01 / np.sqrt(ew)
-            x = x.repeat(n_repeats, 1)
-            x = x + torch.randn_like(x) * std
-        return x
-
-    def _init_embeddings(self, z):
-        # z: [b, c, t, h, w]
-        self._need_init = False
-        flat_inputs = shift_dim(z, 1, -1).flatten(end_dim=-2)
-        y = self._tile(flat_inputs)
-
-        d = y.shape[0]
-        _k_rand = y[torch.randperm(y.shape[0])][: self.n_codes]
-        if dist.is_initialized():
-            dist.broadcast(_k_rand, 0)
-        self.embeddings.data.copy_(_k_rand)
-        self.z_avg.data.copy_(_k_rand)
-        self.N.data.copy_(torch.ones(self.n_codes))
-
-    def forward(self, z):
-        # z: [b, c, t, h, w]
-        if self._need_init and self.training:
-            self._init_embeddings(z)
-        flat_inputs = shift_dim(z, 1, -1).flatten(end_dim=-2)
-        distances = (
-            (flat_inputs**2).sum(dim=1, keepdim=True)
-            - 2 * flat_inputs @ self.embeddings.t()
-            + (self.embeddings.t() ** 2).sum(dim=0, keepdim=True)
-        )
-
-        encoding_indices = torch.argmin(distances, dim=1)
-        encode_onehot = F.one_hot(encoding_indices, self.n_codes).type_as(flat_inputs)
-        encoding_indices = encoding_indices.view(z.shape[0], *z.shape[2:])
-
-        embeddings = F.embedding(encoding_indices, self.embeddings)
-        embeddings = shift_dim(embeddings, -1, 1)
-
-        commitment_loss = 0.25 * F.mse_loss(z, embeddings.detach())
-
-        # EMA codebook update
-        if self.training:
-            n_total = encode_onehot.sum(dim=0)
-            encode_sum = flat_inputs.t() @ encode_onehot
-            if dist.is_initialized():
-                dist.all_reduce(n_total)
-                dist.all_reduce(encode_sum)
-
-            self.N.data.mul_(0.99).add_(n_total, alpha=0.01)
-            self.z_avg.data.mul_(0.99).add_(encode_sum.t(), alpha=0.01)
-
-            n = self.N.sum()
-            weights = (self.N + 1e-7) / (n + self.n_codes * 1e-7) * n
-            encode_normalized = self.z_avg / weights.unsqueeze(1)
-            self.embeddings.data.copy_(encode_normalized)
-
-            y = self._tile(flat_inputs)
-            _k_rand = y[torch.randperm(y.shape[0])][: self.n_codes]
-            if dist.is_initialized():
-                dist.broadcast(_k_rand, 0)
-
-            usage = (self.N.view(self.n_codes, 1) >= 1).float()
-            self.embeddings.data.mul_(usage).add_(_k_rand * (1 - usage))
-
-        embeddings_st = (embeddings - z).detach() + z
-
-        avg_probs = torch.mean(encode_onehot, dim=0)
-        perplexity = torch.exp(-torch.sum(avg_probs * torch.log(avg_probs + 1e-10)))
-
-        return dict(
-            embeddings=embeddings_st,
-            encodings=encoding_indices,
-            commitment_loss=commitment_loss,
-            perplexity=perplexity,
-        )
-
-    def dictionary_lookup(self, encodings):
-        embeddings = F.embedding(encodings, self.embeddings)
-        return embeddings

opensora/models/causalvideovae/model/modules/resnet_block.py

@@ -1,125 +0,0 @@
-import torch
-import torch.nn as nn
-from einops import rearrange, pack, unpack
-from .normalize import Normalize
-from .ops import nonlinearity, video_to_image
-from .conv import CausalConv3d, CausalConv3d_GC
-from .block import Block
-from torch.utils.checkpoint import checkpoint
-    
-class ResnetBlock2D(Block):
-    def __init__(self, *, in_channels, out_channels=None, conv_shortcut=False,
-                 dropout):
-        super().__init__()
-        self.in_channels = in_channels
-        self.out_channels = in_channels if out_channels is None else out_channels
-        self.use_conv_shortcut = conv_shortcut
-
-        self.norm1 = Normalize(in_channels)
-        self.conv1 = torch.nn.Conv2d(
-            in_channels, out_channels, kernel_size=3, stride=1, padding=1
-        )
-        self.norm2 = Normalize(out_channels)
-        self.dropout = torch.nn.Dropout(dropout)
-        self.conv2 = torch.nn.Conv2d(
-            out_channels, out_channels, kernel_size=3, stride=1, padding=1
-        )
-        if self.in_channels != self.out_channels:
-            if self.use_conv_shortcut:
-                self.conv_shortcut = torch.nn.Conv2d(
-                    in_channels, out_channels, kernel_size=3, stride=1, padding=1
-                )
-            else:
-                self.nin_shortcut = torch.nn.Conv2d(
-                    in_channels, out_channels, kernel_size=1, stride=1, padding=0
-                )
-                
-    @video_to_image
-    def forward(self, x):
-        h = x
-        h = self.norm1(h)
-        h = nonlinearity(h)
-        h = self.conv1(h)
-        h = self.norm2(h)
-        h = nonlinearity(h)
-        h = self.dropout(h)
-        h = self.conv2(h)
-        if self.in_channels != self.out_channels:
-            if self.use_conv_shortcut:
-                x = self.conv_shortcut(x)
-            else:
-                x = self.nin_shortcut(x)
-        x = x + h
-        return x
-
-class ResnetBlock3D(Block):
-    def __init__(self, *, in_channels, out_channels=None, conv_shortcut=False, dropout):
-        super().__init__()
-        self.in_channels = in_channels
-        self.out_channels = in_channels if out_channels is None else out_channels
-        self.use_conv_shortcut = conv_shortcut
-
-        self.norm1 = Normalize(in_channels)
-        self.conv1 = CausalConv3d(in_channels, out_channels, 3, padding=1)
-        self.norm2 = Normalize(out_channels)
-        self.dropout = torch.nn.Dropout(dropout)
-        self.conv2 = CausalConv3d(out_channels, out_channels, 3, padding=1)
-        if self.in_channels != self.out_channels:
-            if self.use_conv_shortcut:
-                self.conv_shortcut = CausalConv3d(in_channels, out_channels, 3, padding=1)
-            else:
-                self.nin_shortcut = CausalConv3d(in_channels, out_channels, 1, padding=0)
-
-    def forward(self, x):
-        h = x
-        h = self.norm1(h)
-        h = nonlinearity(h)
-        h = self.conv1(h)
-        h = self.norm2(h)
-        h = nonlinearity(h)
-        h = self.dropout(h)
-        h = self.conv2(h)
-        if self.in_channels != self.out_channels:
-            if self.use_conv_shortcut:
-                x = self.conv_shortcut(x)
-            else:
-                x = self.nin_shortcut(x)
-        return x + h
-    
-    
-class ResnetBlock3D_GC(Block):
-    def __init__(self, *, in_channels, out_channels=None, conv_shortcut=False, dropout):
-        super().__init__()
-        self.in_channels = in_channels
-        self.out_channels = in_channels if out_channels is None else out_channels
-        self.use_conv_shortcut = conv_shortcut
-
-        self.norm1 = Normalize(in_channels)
-        self.conv1 = CausalConv3d(in_channels, out_channels, 3, padding=1)
-        self.norm2 = Normalize(out_channels)
-        self.dropout = torch.nn.Dropout(dropout)
-        self.conv2 = CausalConv3d(out_channels, out_channels, 3, padding=1)
-        if self.in_channels != self.out_channels:
-            if self.use_conv_shortcut:
-                self.conv_shortcut = CausalConv3d(in_channels, out_channels, 3, padding=1)
-            else:
-                self.nin_shortcut = CausalConv3d(in_channels, out_channels, 1, padding=0)
-                
-    def forward(self, x):
-        return checkpoint(self._forward, x, use_reentrant=True)
-    
-    def _forward(self, x):
-        h = x
-        h = self.norm1(h)
-        h = nonlinearity(h)
-        h = self.conv1(h)
-        h = self.norm2(h)
-        h = nonlinearity(h)
-        h = self.dropout(h)
-        h = self.conv2(h)
-        if self.in_channels != self.out_channels:
-            if self.use_conv_shortcut:
-                x = self.conv_shortcut(x)
-            else:
-                x = self.nin_shortcut(x)
-        return x + h

opensora/models/causalvideovae/model/modules/updownsample.py

@@ -1,338 +0,0 @@
-from typing import Union, Tuple
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from .resnet_block import ResnetBlock3D
-from .attention import TemporalAttnBlock
-from .normalize import Normalize
-from .ops import cast_tuple, video_to_image
-from .conv import CausalConv3d, CausalConv3d_GC
-from einops import rearrange
-from .block import Block
-
-
-class Upsample(Block):
-    def __init__(self, in_channels, out_channels):
-        super().__init__()
-        self.with_conv = True
-        if self.with_conv:
-            self.conv = torch.nn.Conv2d(in_channels,
-                                        out_channels,
-                                        kernel_size=3,
-                                        stride=1,
-                                        padding=1)
-            
-    @video_to_image
-    def forward(self, x):
-        x = torch.nn.functional.interpolate(x, scale_factor=2.0, mode="nearest")
-        if self.with_conv:
-            x = self.conv(x)
-        return x
-
-class Downsample(Block):
-    def __init__(self, in_channels, out_channels, undown=False):
-        super().__init__()
-        self.with_conv = True
-        self.undown = undown
-        if self.with_conv:
-            # no asymmetric padding in torch conv, must do it ourselves
-            if self.undown:
-                self.conv = torch.nn.Conv2d(in_channels,
-                                        out_channels,
-                                        kernel_size=3,
-                                        stride=1,
-                                        padding=1)
-            else:
-                self.conv = torch.nn.Conv2d(in_channels,
-                                        out_channels,
-                                        kernel_size=3,
-                                        stride=2,
-                                        padding=0)
-    @video_to_image
-    def forward(self, x):
-        if self.with_conv:
-            if self.undown:
-
-                x = self.conv(x)
-            else:
-                pad = (0, 1, 0, 1)
-
-                x = torch.nn.functional.pad(x, pad, mode="constant", value=0)
-                x = self.conv(x)
-        else:
-            x = torch.nn.functional.avg_pool2d(x, kernel_size=2, stride=2)
-        return x
-
-class SpatialDownsample2x(Block):
-    def __init__(
-        self,
-        chan_in,
-        chan_out,
-        kernel_size: Union[int, Tuple[int]] = (3, 3),
-        stride: Union[int, Tuple[int]] = (2, 2),
-        **kwargs
-    ):
-        super().__init__()
-        kernel_size = cast_tuple(kernel_size, 2)
-        stride = cast_tuple(stride, 2)
-        self.chan_in = chan_in
-        self.chan_out = chan_out
-        self.kernel_size = kernel_size
-        self.conv = CausalConv3d(
-            self.chan_in,
-            self.chan_out,
-            (1,) + self.kernel_size,
-            stride=(1, ) + stride,
-            padding=0
-        )
-
-    def forward(self, x):
-        pad = (0,1,0,1,0,0)
-        x = torch.nn.functional.pad(x, pad, mode="constant", value=0)
-        x = self.conv(x)
-        return x
-
-class SpatialUpsample2x_GC(Block):
-    def __init__(
-        self,
-        chan_in,
-        chan_out,
-        kernel_size: Union[int, Tuple[int]] = (3, 3),
-        stride: Union[int, Tuple[int]] = (1, 1),
-        unup=False,
-    ):
-        super().__init__()
-        self.chan_in = chan_in
-        self.chan_out = chan_out
-        self.kernel_size = kernel_size
-        self.unup = unup
-        self.conv = CausalConv3d_GC(
-            self.chan_in,
-            self.chan_out,
-            (1,) + self.kernel_size,
-            stride=(1, ) + stride,
-            padding=1
-        )
-
-    def forward(self, x):
-        if not self.unup:
-            t = x.shape[2]
-            x = rearrange(x, "b c t h w -> b (c t) h w")
-            x = F.interpolate(x, scale_factor=(2,2), mode="nearest")
-            x = rearrange(x, "b (c t) h w -> b c t h w", t=t)
-        x = self.conv(x)
-        return x
-    
-
-class SpatialUpsample2x(Block):
-    def __init__(
-        self,
-        chan_in,
-        chan_out,
-        kernel_size: Union[int, Tuple[int]] = (3, 3),
-        stride: Union[int, Tuple[int]] = (1, 1),
-        unup=False,
-    ):
-        super().__init__()
-        self.chan_in = chan_in
-        self.chan_out = chan_out
-        self.kernel_size = kernel_size
-        self.unup = unup
-        self.conv = CausalConv3d(
-            self.chan_in,
-            self.chan_out,
-            (1,) + self.kernel_size,
-            stride=(1, ) + stride,
-            padding=1
-        )
-
-    def forward(self, x):
-        if not self.unup:
-            t = x.shape[2]
-            x = rearrange(x, "b c t h w -> b (c t) h w")
-            x = F.interpolate(x, scale_factor=(2,2), mode="nearest")
-            x = rearrange(x, "b (c t) h w -> b c t h w", t=t)
-        x = self.conv(x)
-        return x
-    
-class TimeDownsample2x(Block):
-    def __init__(
-        self,
-        chan_in,
-        chan_out,
-        kernel_size: int = 3
-    ):
-        super().__init__()
-        self.kernel_size = kernel_size
-
-        self.conv = nn.AvgPool3d((kernel_size, 1, 1), stride=(2, 1, 1))
-        
-    def forward(self, x):
-
-        first_frame_pad = x[:, :, :1, :, :].repeat(
-            (1, 1, self.kernel_size - 1, 1, 1)
-        )
-        x = torch.concatenate((first_frame_pad, x), dim=2)
-        return self.conv(x)
-
-class TimeUpsample2x(Block):
-    def __init__(
-        self,
-        chan_in,
-        chan_out
-    ):
-        super().__init__()
-    def forward(self, x):
-        if x.size(2) > 1:
-            x,x_= x[:,:,:1],x[:,:,1:]
-            x_= F.interpolate(x_, scale_factor=(2,1,1), mode='trilinear')
-            x = torch.concat([x, x_], dim=2)
-        return x
-    
-class TimeDownsampleRes2x(nn.Module):
-    def __init__(
-        self,
-        in_channels,
-        out_channels,
-        kernel_size: int = 3,
-        mix_factor: float = 2.0,
-    ):
-        super().__init__()
-        self.kernel_size = cast_tuple(kernel_size, 3)
-
-        self.avg_pool = nn.AvgPool3d((kernel_size, 1, 1), stride=(2, 1, 1))
-        self.conv = nn.Conv3d(
-            in_channels, out_channels, self.kernel_size, stride=(2,1,1), padding=(0,1,1)
-        )
-        self.mix_factor = torch.nn.Parameter(torch.Tensor([mix_factor]))
-    
-    def forward(self, x):
-        alpha = torch.sigmoid(self.mix_factor)
-
-        first_frame_pad = x[:, :, :1, :, :].repeat(
-            (1, 1, self.kernel_size[0] - 1, 1, 1)
-        )
-        x = torch.concatenate((first_frame_pad, x), dim=2)
-        return alpha * self.avg_pool(x) + (1 - alpha) * self.conv(x)
-
-class TimeUpsampleRes2x(nn.Module):
-    def __init__(
-        self,
-        in_channels,
-        out_channels,
-        kernel_size: int = 3,
-        mix_factor: float = 2.0,
-    ):
-        super().__init__()
-        self.conv = CausalConv3d(
-            in_channels, out_channels, kernel_size, padding=1
-        )
-        self.mix_factor = torch.nn.Parameter(torch.Tensor([mix_factor]))
-        
-    def forward(self, x):
-        alpha = torch.sigmoid(self.mix_factor)
-        if x.size(2) > 1:
-            x,x_= x[:,:,:1],x[:,:,1:]
-
-            x_= F.interpolate(x_, scale_factor=(2,1,1), mode='trilinear')
-            x = torch.concat([x, x_], dim=2)
-        return alpha * x + (1-alpha) * self.conv(x)
-
-class TimeDownsampleResAdv2x(nn.Module):
-    def __init__(
-        self,
-        in_channels,
-        out_channels,
-        kernel_size: int = 3,
-        mix_factor: float = 1.5,
-    ):
-        super().__init__()
-        self.kernel_size = cast_tuple(kernel_size, 3)
-        self.avg_pool = nn.AvgPool3d((kernel_size,1,1), stride=(2,1,1))
-        self.attn = TemporalAttnBlock(in_channels)
-        self.res = ResnetBlock3D(in_channels=in_channels, out_channels=in_channels, dropout=0.0)
-        self.conv = nn.Conv3d(
-            in_channels, out_channels, self.kernel_size, stride=(2,1,1), padding=(0,1,1)
-        )
-        self.mix_factor = torch.nn.Parameter(torch.Tensor([mix_factor]))
-    
-    def forward(self, x):
-        first_frame_pad = x[:, :, :1, :, :].repeat(
-            (1, 1, self.kernel_size[0] - 1, 1, 1)
-        )
-        x = torch.concatenate((first_frame_pad, x), dim=2)
-        alpha = torch.sigmoid(self.mix_factor)
-        return alpha * self.avg_pool(x) + (1 - alpha) * self.conv(self.attn((self.res(x))))
-
-class TimeUpsampleResAdv2x(nn.Module):
-    def __init__(
-        self,
-        in_channels,
-        out_channels,
-        kernel_size: int = 3,
-        mix_factor: float = 1.5,
-    ):
-        super().__init__()
-        self.res = ResnetBlock3D(in_channels=in_channels, out_channels=in_channels, dropout=0.0)
-        self.attn = TemporalAttnBlock(in_channels)
-        self.norm = Normalize(in_channels=in_channels)
-        self.conv = CausalConv3d(
-            in_channels, out_channels, kernel_size, padding=1
-        )
-        self.mix_factor = torch.nn.Parameter(torch.Tensor([mix_factor]))
-        
-    def forward(self, x):
-        if x.size(2) > 1:
-            x,x_= x[:,:,:1],x[:,:,1:]
-
-            x_= F.interpolate(x_, scale_factor=(2,1,1), mode='trilinear')
-            x = torch.concat([x, x_], dim=2)
-        alpha = torch.sigmoid(self.mix_factor)
-        return alpha * x + (1 - alpha) * self.conv(self.attn(self.res(x)))
-
-
-class Spatial2xTime2x3DDownsample(Block):
-    def __init__(self, in_channels, out_channels):
-        super().__init__()
-        self.conv = CausalConv3d(in_channels, out_channels, kernel_size=3, padding=0, stride=2)
-
-    def forward(self, x):
-        pad = (0,1,0,1,0,0)
-        x = torch.nn.functional.pad(x, pad, mode="constant", value=0)
-        x = self.conv(x)
-        return x
-
-class Spatial2x3DDownsample(Block):
-    def __init__(self, in_channels, out_channels):
-        super().__init__()
-        self.conv = CausalConv3d(in_channels, out_channels, kernel_size=3, padding=0, stride=(1,2,2))
-
-    def forward(self, x):
-        pad = (0,1,0,1,0,0)
-        x = torch.nn.functional.pad(x, pad, mode="constant", value=0)
-        x = self.conv(x)
-        return x
-    
-
-class Spatial2x3DUpsample(Block):
-    def __init__(self, in_channels, out_channels):
-        super().__init__()
-        self.conv = CausalConv3d(in_channels, out_channels, kernel_size=3, padding=1)
-
-    def forward(self, x):
-        x = F.interpolate(x, scale_factor=(1,2,2), mode='trilinear')
-        return self.conv(x)
-    
-class Spatial2xTime2x3DUpsample(Block):
-    def __init__(self, in_channels, out_channels):
-        super().__init__()
-        self.conv = CausalConv3d(in_channels, out_channels, kernel_size=3, padding=1)
-    def forward(self, x):
-        if x.size(2) > 1:
-            x,x_= x[:,:,:1],x[:,:,1:]
-            x_= F.interpolate(x_, scale_factor=(2,2,2), mode='trilinear')
-            x = F.interpolate(x, scale_factor=(1,2,2), mode='trilinear')
-            x = torch.concat([x, x_], dim=2)
-        else:
-             x = F.interpolate(x, scale_factor=(1,2,2), mode='trilinear')
-        return self.conv(x)

opensora/models/causalvideovae/model/trainer_videobase.py

@@ -1,26 +0,0 @@
-from transformers import Trainer
-import torch.nn.functional as F
-from typing import Optional
-import os
-import torch
-from transformers.utils import WEIGHTS_NAME
-import json
-
-class VideoBaseTrainer(Trainer):
-
-    def _save(self, output_dir: Optional[str] = None, state_dict=None):
-        output_dir = output_dir if output_dir is not None else self.args.output_dir
-        os.makedirs(output_dir, exist_ok=True)
-        if state_dict is None:
-            state_dict = self.model.state_dict()
-        
-        # get model config
-        model_config = self.model.config.to_dict()
-        
-        # add more information
-        model_config['model'] = self.model.__class__.__name__
-        
-        with open(os.path.join(output_dir, "config.json"), "w") as file:
-            json.dump(self.model.config.to_dict(), file)
-        torch.save(state_dict, os.path.join(output_dir, WEIGHTS_NAME))
-        torch.save(self.args, os.path.join(output_dir, "training_args.bin"))

opensora/models/causalvideovae/model/utils/distrib_utils.py

@@ -1,42 +0,0 @@
-import torch
-import numpy as np
-
-class DiagonalGaussianDistribution(object):
-    def __init__(self, parameters, deterministic=False):
-        self.parameters = parameters
-        self.mean, self.logvar = torch.chunk(parameters, 2, dim=1)
-        self.logvar = torch.clamp(self.logvar, -30.0, 20.0)
-        self.deterministic = deterministic
-        self.std = torch.exp(0.5 * self.logvar)
-        self.var = torch.exp(self.logvar)
-        if self.deterministic:
-            self.var = self.std = torch.zeros_like(self.mean).to(device=self.parameters.device)
-
-    def sample(self):
-        x = self.mean + self.std * torch.randn(self.mean.shape).to(device=self.parameters.device)
-        return x
-
-    def kl(self, other=None):
-        if self.deterministic:
-            return torch.Tensor([0.])
-        else:
-            if other is None:
-                return 0.5 * torch.sum(torch.pow(self.mean, 2)
-                                       + self.var - 1.0 - self.logvar,
-                                       dim=[1, 2, 3])
-            else:
-                return 0.5 * torch.sum(
-                    torch.pow(self.mean - other.mean, 2) / other.var
-                    + self.var / other.var - 1.0 - self.logvar + other.logvar,
-                    dim=[1, 2, 3])
-
-    def nll(self, sample, dims=[1,2,3]):
-        if self.deterministic:
-            return torch.Tensor([0.])
-        logtwopi = np.log(2.0 * np.pi)
-        return 0.5 * torch.sum(
-            logtwopi + self.logvar + torch.pow(sample - self.mean, 2) / self.var,
-            dim=dims)
-
-    def mode(self):
-        return self.mean

opensora/models/causalvideovae/model/utils/module_utils.py

@@ -1,17 +0,0 @@
-import importlib
-
-Module = str
-MODULES_BASE = "opensora.models.causalvideovae.model.modules."
-
-def resolve_str_to_obj(str_val, append=True):
-    if append:
-        str_val = MODULES_BASE + str_val
-    module_name, class_name = str_val.rsplit('.', 1)
-    module = importlib.import_module(module_name)
-    return getattr(module, class_name)
-
-def create_instance(module_class_str: str, **kwargs):
-    module_name, class_name = module_class_str.rsplit('.', 1)
-    module = importlib.import_module(module_name)
-    class_ = getattr(module, class_name)
-    return class_(**kwargs)

opensora/models/causalvideovae/model/utils/scheduler_utils.py

@@ -1,7 +0,0 @@
-import torch
-
-def cosine_scheduler(step, max_steps, value_base=1, value_end=0):
-    step = torch.tensor(step)
-    cosine_value = 0.5 * (1 + torch.cos(torch.pi * step / max_steps))
-    value = value_end + (value_base - value_end) * cosine_value
-    return value

opensora/models/causalvideovae/model/utils/video_utils.py

@@ -1,10 +0,0 @@
-import torch
-import numpy as np
-
-def tensor_to_video(x):
-    x = (x * 2 - 1).detach().cpu()
-    x = torch.clamp(x, -1, 1)
-    x = (x + 1) / 2
-    x = x.permute(1, 0, 2, 3).float().numpy() # c t h w -> t c h w
-    x = (255 * x).astype(np.uint8)
-    return x

opensora/models/causalvideovae/sample/rec_video_vae.py

@@ -1,277 +0,0 @@
-import random
-import argparse
-import cv2
-from tqdm import tqdm
-import numpy as np
-import numpy.typing as npt
-import torch
-from decord import VideoReader, cpu
-from torch.nn import functional as F
-from pytorchvideo.transforms import ShortSideScale
-from torchvision.transforms import Lambda, Compose
-from torchvision.transforms.transforms import CenterCropVideo
-import sys
-from torch.utils.data import Dataset, DataLoader, Subset
-import os
-import glob
-sys.path.append(".")
-from causalvideovae.model import CausalVAEModel
-import torch.nn as nn
-
-
-def array_to_video(
-    image_array: npt.NDArray, fps: float = 30.0, output_file: str = "output_video.mp4"
-) -> None:
-    height, width, channels = image_array[0].shape
-    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
-    video_writer = cv2.VideoWriter(output_file, fourcc, float(fps), (width, height))
-
-    for image in image_array:
-        image_rgb = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
-        video_writer.write(image_rgb)
-
-    video_writer.release()
-
-
-def custom_to_video(
-    x: torch.Tensor, fps: float = 2.0, output_file: str = "output_video.mp4"
-) -> None:
-    x = x.detach().cpu()
-    x = torch.clamp(x, -1, 1)
-    x = (x + 1) / 2
-    x = x.permute(1, 2, 3, 0).float().numpy()
-    x = (255 * x).astype(np.uint8)
-    array_to_video(x, fps=fps, output_file=output_file)
-    return
-
-
-def read_video(video_path: str, num_frames: int, sample_rate: int) -> torch.Tensor:
-    decord_vr = VideoReader(video_path, ctx=cpu(0), num_threads=8)
-    total_frames = len(decord_vr)
-    sample_frames_len = sample_rate * num_frames
-
-    if total_frames > sample_frames_len:
-        s = 0
-        e = s + sample_frames_len
-        num_frames = num_frames
-    else:
-        s = 0
-        e = total_frames
-        num_frames = int(total_frames / sample_frames_len * num_frames)
-        print(
-            f"sample_frames_len {sample_frames_len}, only can sample {num_frames * sample_rate}",
-            video_path,
-            total_frames,
-        )
-
-    frame_id_list = np.linspace(s, e - 1, num_frames, dtype=int)
-    video_data = decord_vr.get_batch(frame_id_list).asnumpy()
-    video_data = torch.from_numpy(video_data)
-    video_data = video_data.permute(3, 0, 1, 2)  # (T, H, W, C) -> (C, T, H, W)
-    return video_data
-
-
-class RealVideoDataset(Dataset):
-    video_exts = ['avi', 'mp4', 'webm']
-    
-    def __init__(
-        self,
-        real_video_dir,
-        num_frames,
-        sample_rate=1,
-        crop_size=None,
-        resolution=128,
-    ) -> None:
-        super().__init__()
-        self.real_video_files = self._combine_without_prefix(real_video_dir)
-        self.num_frames = num_frames
-        self.sample_rate = sample_rate
-        self.crop_size = crop_size
-        self.short_size = resolution
-
-    def __len__(self):
-        return len(self.real_video_files)
-
-    def __getitem__(self, index):
-        if index >= len(self):
-            raise IndexError
-        real_video_file = self.real_video_files[index]
-        real_video_tensor = self._load_video(real_video_file)
-        video_name = os.path.basename(real_video_file)
-        return {'video': real_video_tensor, 'file_name': video_name }
-
-    def _load_video(self, video_path):
-        num_frames = self.num_frames
-        sample_rate = self.sample_rate
-        decord_vr = VideoReader(video_path, ctx=cpu(0))
-        total_frames = len(decord_vr)
-        sample_frames_len = sample_rate * num_frames
-
-        if total_frames > sample_frames_len:
-            s = 0
-            e = s + sample_frames_len
-            num_frames = num_frames
-        else:
-            s = 0
-            e = total_frames
-            num_frames = int(total_frames / sample_frames_len * num_frames)
-            print(
-                f"sample_frames_len {sample_frames_len}, only can sample {num_frames * sample_rate}",
-                video_path,
-                total_frames,
-            )
-
-        frame_id_list = np.linspace(s, e - 1, num_frames, dtype=int)
-        video_data = decord_vr.get_batch(frame_id_list).asnumpy()
-        video_data = torch.from_numpy(video_data)
-        video_data = video_data.permute(3, 0, 1, 2)
-        return _preprocess(
-            video_data, short_size=self.short_size, crop_size=self.crop_size
-        )
-
-    def _combine_without_prefix(self, folder_path):
-        samples = []
-        samples += sum([glob.glob(os.path.join(folder_path, '**', f'*.{ext}'), recursive=True)
-                            for ext in self.video_exts], [])
-        samples.sort()
-        return samples
-
-def resize(x, resolution):
-    height, width = x.shape[-2:]
-    aspect_ratio = width / height
-    if width <= height:
-        new_width = resolution
-        new_height = int(resolution / aspect_ratio)
-    else:
-        new_height = resolution
-        new_width = int(resolution * aspect_ratio)
-    resized_x = F.interpolate(x, size=(new_height, new_width), mode='bilinear', align_corners=True, antialias=True)
-    return resized_x
-
-def _preprocess(video_data, short_size=128, crop_size=None):
-    transform = Compose(
-        [
-            Lambda(lambda x: ((x / 255.0) * 2 - 1)),
-            Lambda(lambda x: resize(x, short_size)),
-            (
-                CenterCropVideo(crop_size=crop_size)
-                if crop_size is not None
-                else Lambda(lambda x: x)
-            ),
-        ]
-    )
-    video_outputs = transform(video_data)
-    video_outputs = _format_video_shape(video_outputs)
-    return video_outputs
-
-
-def _format_video_shape(video, time_compress=4, spatial_compress=8):
-    time = video.shape[1]
-    height = video.shape[2]
-    width = video.shape[3]
-    new_time = (
-        (time - (time - 1) % time_compress)
-        if (time - 1) % time_compress != 0
-        else time
-    )
-    new_height = (
-        (height - (height) % spatial_compress)
-        if height % spatial_compress != 0
-        else height
-    )
-    new_width = (
-        (width - (width) % spatial_compress) if width % spatial_compress != 0 else width
-    )
-    return video[:, :new_time, :new_height, :new_width]
-
-
-@torch.no_grad()
-def main(args: argparse.Namespace):
-    real_video_dir = args.real_video_dir
-    generated_video_dir = args.generated_video_dir
-    ckpt = args.ckpt
-    sample_rate = args.sample_rate
-    resolution = args.resolution
-    crop_size = args.crop_size
-    num_frames = args.num_frames
-    sample_rate = args.sample_rate
-    device = args.device
-    sample_fps = args.sample_fps
-    batch_size = args.batch_size
-    num_workers = args.num_workers
-    subset_size = args.subset_size
-    
-    if not os.path.exists(args.generated_video_dir):
-        os.makedirs(args.generated_video_dir, exist_ok=True)
-    
-    data_type = torch.bfloat16
-    
-    # ---- Load Model ----
-    device = args.device
-    vqvae = CausalVAEModel.from_pretrained(args.ckpt)
-    print(vqvae)
-    vqvae = vqvae.to(device).to(data_type)
-    if args.enable_tiling:
-        vqvae.enable_tiling()
-        vqvae.tile_overlap_factor = args.tile_overlap_factor
-    # ---- Load Model ----
-
-    # ---- Prepare Dataset ----
-    dataset = RealVideoDataset(
-        real_video_dir=real_video_dir,
-        num_frames=num_frames,
-        sample_rate=sample_rate,
-        crop_size=crop_size,
-        resolution=resolution,
-    )
-    
-    if subset_size:
-        indices = range(subset_size)
-        dataset = Subset(dataset, indices=indices)
-        
-    dataloader = DataLoader(
-        dataset, batch_size=batch_size, pin_memory=True, num_workers=num_workers
-    )
-    # ---- Prepare Dataset
-
-    # ---- Inference ----
-    for batch in tqdm(dataloader):
-        x, file_names = batch['video'], batch['file_name']
-        
-        x = x.to(device=device, dtype=data_type)  # b c t h w
-        latents = vqvae.encode(x).sample().to(data_type)
-        video_recon = vqvae.decode(latents)
-        for idx, video in enumerate(video_recon):
-            output_path = os.path.join(generated_video_dir, file_names[idx])
-            if args.output_origin:
-                os.makedirs(os.path.join(generated_video_dir, "origin/"), exist_ok=True)
-                origin_output_path = os.path.join(generated_video_dir, "origin/", file_names[idx])
-                custom_to_video(
-                    x[idx], fps=sample_fps / sample_rate, output_file=origin_output_path
-                )
-            custom_to_video(
-                video, fps=sample_fps / sample_rate, output_file=output_path
-            )
-    # ---- Inference ----
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--real_video_dir", type=str, default="")
-    parser.add_argument("--generated_video_dir", type=str, default="")
-    parser.add_argument("--ckpt", type=str, default="")
-    parser.add_argument("--sample_fps", type=int, default=30)
-    parser.add_argument("--resolution", type=int, default=336)
-    parser.add_argument("--crop_size", type=int, default=None)
-    parser.add_argument("--num_frames", type=int, default=17)
-    parser.add_argument("--sample_rate", type=int, default=1)
-    parser.add_argument("--batch_size", type=int, default=1)
-    parser.add_argument("--num_workers", type=int, default=8)
-    parser.add_argument("--subset_size", type=int, default=None)
-    parser.add_argument("--tile_overlap_factor", type=float, default=0.25)
-    parser.add_argument('--enable_tiling', action='store_true')
-    parser.add_argument('--output_origin', action='store_true')
-    parser.add_argument("--device", type=str, default="cuda")
-
-    args = parser.parse_args()
-    main(args)
-    

opensora/models/causalvideovae/utils/dataset_utils.py

@@ -1,160 +0,0 @@
-import math
-from einops import rearrange
-import decord
-from torch.nn import functional as F
-import torch
-
-
-IMG_EXTENSIONS = ['.jpg', '.JPG', '.jpeg', '.JPEG', '.png', '.PNG']
-
-def is_image_file(filename):
-    return any(filename.endswith(extension) for extension in IMG_EXTENSIONS)
-
-class DecordInit(object):
-    """Using Decord(https://github.com/dmlc/decord) to initialize the video_reader."""
-
-    def __init__(self, num_threads=1):
-        self.num_threads = num_threads
-        self.ctx = decord.cpu(0)
-
-    def __call__(self, filename):
-        """Perform the Decord initialization.
-        Args:
-            results (dict): The resulting dict to be modified and passed
-                to the next transform in pipeline.
-        """
-        reader = decord.VideoReader(filename,
-                                    ctx=self.ctx,
-                                    num_threads=self.num_threads)
-        return reader
-
-    def __repr__(self):
-        repr_str = (f'{self.__class__.__name__}('
-                    f'sr={self.sr},'
-                    f'num_threads={self.num_threads})')
-        return repr_str
-
-def pad_to_multiple(number, ds_stride):
-    remainder = number % ds_stride
-    if remainder == 0:
-        return number
-    else:
-        padding = ds_stride - remainder
-        return number + padding
-
-class Collate:
-    def __init__(self, args):
-        self.max_image_size = args.max_image_size
-        self.ae_stride = args.ae_stride
-        self.ae_stride_t = args.ae_stride_t
-        self.ae_stride_thw = (self.ae_stride_t, self.ae_stride, self.ae_stride)
-        self.ae_stride_1hw = (1, self.ae_stride, self.ae_stride)
-
-        self.patch_size = args.patch_size
-        self.patch_size_t = args.patch_size_t
-        self.patch_size_thw = (self.patch_size_t, self.patch_size, self.patch_size)
-        self.patch_size_1hw = (1, self.patch_size, self.patch_size)
-
-        self.num_frames = args.num_frames
-        self.use_image_num = args.use_image_num
-        self.max_thw = (self.num_frames, self.max_image_size, self.max_image_size)
-        self.max_1hw = (1, self.max_image_size, self.max_image_size)
-
-    def package(self, batch):
-        # import ipdb;ipdb.set_trace()
-        batch_tubes_vid = [i['video_data']['video'] for i in batch]  # b [c t h w]
-        input_ids_vid = torch.stack([i['video_data']['input_ids'] for i in batch])  # b 1 l
-        cond_mask_vid = torch.stack([i['video_data']['cond_mask'] for i in batch])  # b 1 l
-        batch_tubes_img, input_ids_img, cond_mask_img = None, None, None
-        if self.use_image_num != 0: 
-            batch_tubes_img = [j for i in batch for j in i['image_data']['image']]  # b*num_img [c 1 h w]
-            input_ids_img = torch.stack([i['image_data']['input_ids'] for i in batch])  # b image_num l
-            cond_mask_img = torch.stack([i['image_data']['cond_mask'] for i in batch])  # b image_num l
-        return batch_tubes_vid, input_ids_vid, cond_mask_vid, batch_tubes_img, input_ids_img, cond_mask_img
-
-    def __call__(self, batch):
-        batch_tubes_vid, input_ids_vid, cond_mask_vid, batch_tubes_img, input_ids_img, cond_mask_img = self.package(batch)
-
-        # import ipdb;ipdb.set_trace()
-        ds_stride = self.ae_stride * self.patch_size
-        t_ds_stride = self.ae_stride_t * self.patch_size_t
-        if self.use_image_num == 0:
-            pad_batch_tubes, attention_mask = self.process(batch_tubes_vid, t_ds_stride, ds_stride, 
-                                                      self.max_thw, self.ae_stride_thw, self.patch_size_thw, extra_1=True)
-            # attention_mask: b t h w
-            input_ids, cond_mask = input_ids_vid.squeeze(1), cond_mask_vid.squeeze(1)  # b 1 l -> b l
-        else:
-            pad_batch_tubes_vid, attention_mask_vid = self.process(batch_tubes_vid, t_ds_stride, ds_stride, 
-                                                                   self.max_thw, self.ae_stride_thw, self.patch_size_thw, extra_1=True)
-            # attention_mask_vid: b t h w
-            pad_batch_tubes_img, attention_mask_img = self.process(batch_tubes_img, 1, ds_stride, 
-                                                                   self.max_1hw, self.ae_stride_1hw, self.patch_size_1hw, extra_1=False)
-            pad_batch_tubes_img = rearrange(pad_batch_tubes_img, '(b i) c 1 h w -> b c i h w', i=self.use_image_num)
-            attention_mask_img = rearrange(attention_mask_img, '(b i) 1 h w -> b i h w', i=self.use_image_num)
-            pad_batch_tubes = torch.cat([pad_batch_tubes_vid, pad_batch_tubes_img], dim=2)  # concat at temporal, video first
-            # attention_mask_img: b num_img h w
-            attention_mask = torch.cat([attention_mask_vid, attention_mask_img], dim=1)  # b t+num_img h w
-            input_ids = torch.cat([input_ids_vid, input_ids_img], dim=1)  # b 1+num_img hw
-            cond_mask = torch.cat([cond_mask_vid, cond_mask_img], dim=1)  # b 1+num_img hw
-        return pad_batch_tubes, attention_mask, input_ids, cond_mask
-
-    def process(self, batch_tubes, t_ds_stride, ds_stride, max_thw, ae_stride_thw, patch_size_thw, extra_1):
-        
-        # pad to max multiple of ds_stride
-        batch_input_size = [i.shape for i in batch_tubes]  # [(c t h w), (c t h w)]
-        max_t, max_h, max_w = max_thw
-        pad_max_t, pad_max_h, pad_max_w = pad_to_multiple(max_t-1 if extra_1 else max_t, t_ds_stride), \
-                                          pad_to_multiple(max_h, ds_stride), \
-                                          pad_to_multiple(max_w, ds_stride)
-        pad_max_t = pad_max_t + 1 if extra_1 else pad_max_t
-        each_pad_t_h_w = [[pad_max_t - i.shape[1],
-                           pad_max_h - i.shape[2],
-                           pad_max_w - i.shape[3]] for i in batch_tubes]
-        pad_batch_tubes = [F.pad(im,
-                                 (0, pad_w,
-                                  0, pad_h,
-                                  0, pad_t), value=0) for (pad_t, pad_h, pad_w), im in zip(each_pad_t_h_w, batch_tubes)]
-        pad_batch_tubes = torch.stack(pad_batch_tubes, dim=0)
-
-        # make attention_mask
-        # first_channel_first_frame, first_channel_other_frame = pad_batch_tubes[:, :1, :1], pad_batch_tubes[:, :1, 1:]  # first channel to make attention_mask
-        # attention_mask_first_frame = F.max_pool3d(first_channel_first_frame, kernel_size=(1, *ae_stride_thw[1:]), stride=(1, *ae_stride_thw[1:]))
-        # if first_channel_other_frame.numel() != 0:
-        #     attention_mask_other_frame = F.max_pool3d(first_channel_other_frame, kernel_size=ae_stride_thw, stride=ae_stride_thw)
-        #     attention_mask = torch.cat([attention_mask_first_frame, attention_mask_other_frame], dim=2)
-        # else:
-        #     attention_mask = attention_mask_first_frame
-        # attention_mask_ = attention_mask[:, 0].bool().float()  # b t h w, do not channel
-
-        # import ipdb;ipdb.set_trace()
-        max_tube_size = [pad_max_t, pad_max_h, pad_max_w]
-        max_latent_size = [((max_tube_size[0]-1) // ae_stride_thw[0] + 1) if extra_1 else (max_tube_size[0] // ae_stride_thw[0]),
-                           max_tube_size[1] // ae_stride_thw[1],
-                           max_tube_size[2] // ae_stride_thw[2]]
-        valid_latent_size = [[int(math.ceil((i[1]-1) / ae_stride_thw[0])) + 1 if extra_1 else int(math.ceil(i[1] / ae_stride_thw[0])),
-                            int(math.ceil(i[2] / ae_stride_thw[1])),
-                            int(math.ceil(i[3] / ae_stride_thw[2]))] for i in batch_input_size]
-        attention_mask = [F.pad(torch.ones(i),
-                                (0, max_latent_size[2] - i[2],
-                                 0, max_latent_size[1] - i[1],
-                                 0, max_latent_size[0] - i[0]), value=0) for i in valid_latent_size]
-        attention_mask = torch.stack(attention_mask)  # b t h w
-
-
-        # max_tube_size = [pad_max_t, pad_max_h, pad_max_w]
-        # max_latent_size = [((max_tube_size[0]-1) // ae_stride_thw[0] + 1) if extra_1 else (max_tube_size[0] // ae_stride_thw[0]),
-        #                    max_tube_size[1] // ae_stride_thw[1],
-        #                    max_tube_size[2] // ae_stride_thw[2]]
-        # max_patchify_latent_size = [((max_latent_size[0]-1) // patch_size_thw[0] + 1) if extra_1 else (max_latent_size[0] // patch_size_thw[0]),
-        #                             max_latent_size[1] // patch_size_thw[1],
-        #                             max_latent_size[2] // patch_size_thw[2]]
-        # valid_patchify_latent_size = [[int(math.ceil((i[1]-1) / t_ds_stride)) + 1 if extra_1 else int(math.ceil(i[1] / t_ds_stride)),
-        #                                int(math.ceil(i[2] / ds_stride)),
-        #                                int(math.ceil(i[3] / ds_stride))] for i in batch_input_size]
-        # attention_mask = [F.pad(torch.ones(i),
-        #                         (0, max_patchify_latent_size[2] - i[2],
-        #                          0, max_patchify_latent_size[1] - i[1],
-        #                          0, max_patchify_latent_size[0] - i[0]), value=0) for i in valid_patchify_latent_size]
-        # attention_mask = torch.stack(attention_mask)  # b t h w
-
-        return pad_batch_tubes, attention_mask

opensora/models/causalvideovae/utils/downloader.py

@@ -1,18 +0,0 @@
-import gdown
-import os
-
-opensora_cache_home = os.path.expanduser(
-    os.getenv("OPENSORA_HOME", os.path.join("~/.cache", "opensora"))
-)
-
-
-def gdown_download(id, fname, cache_dir=None):
-    cache_dir = opensora_cache_home if not cache_dir else cache_dir
-
-    os.makedirs(cache_dir, exist_ok=True)
-    destination = os.path.join(cache_dir, fname)
-    if os.path.exists(destination):
-        return destination
-
-    gdown.download(id=id, output=destination, quiet=False)
-    return destination

opensora/models/causalvideovae/utils/utils.py

@@ -1,627 +0,0 @@
-import os
-
-import torch
-
-import os
-import math
-import torch
-import logging
-import random
-import subprocess
-import numpy as np
-import torch.distributed as dist
-
-# from torch._six import inf
-from torch import inf
-from PIL import Image
-from typing import Union, Iterable
-from collections import OrderedDict
-from torch.utils.tensorboard import SummaryWriter
-
-from diffusers.utils import is_bs4_available, is_ftfy_available
-
-import html
-import re
-import urllib.parse as ul
-
-from torch.utils.data import Dataset
-from torchvision.transforms import Lambda, Compose
-from decord import VideoReader, cpu
-
-from torchvision.transforms.transforms import CenterCropVideo
-from torch.nn import functional as F
-import cv2
-import numpy.typing as npt
-
-if is_bs4_available():
-    from bs4 import BeautifulSoup
-
-if is_ftfy_available():
-    import ftfy
-
-_tensor_or_tensors = Union[torch.Tensor, Iterable[torch.Tensor]]
-
-def find_model(model_name):
-    """
-    Finds a pre-trained Latte model, downloading it if necessary. Alternatively, loads a model from a local path.
-    """
-    assert os.path.isfile(model_name), f'Could not find Latte checkpoint at {model_name}'
-    checkpoint = torch.load(model_name, map_location=lambda storage, loc: storage)
-
-    # if "ema" in checkpoint:  # supports checkpoints from train.py
-    #     print('Using Ema!')
-    #     checkpoint = checkpoint["ema"]
-    # else:
-    print('Using model!')
-    checkpoint = checkpoint['model']
-    return checkpoint
-
-#################################################################################
-#                             Training Clip Gradients                           #
-#################################################################################
-import deepspeed
-def print_grad_norm(model):
-    # 计算并打印梯度范数
-    # model_engine = accelerator.deepspeed_engine_wrapped.engine
-    # gradients = model_engine.get_gradients()
-    # grad_norm = get_grad_norm(gradients)
-    # 计算并打印梯度范数
-    grad_norm = 0
-    n_grad = 0
-    for name, param in model.named_parameters():
-        grad_data = deepspeed.utils.safe_get_full_grad(param)
-        # self.print_tensor_stats(grad_data, name=name)
-
-        if grad_data is not None:
-            param_norm = grad_data.norm(2)
-            grad_norm += param_norm.item() ** 2
-            n_grad += 1
-    grad_norm = (grad_norm / n_grad) ** (1. / 2)
-
-    # self.print_msg('=' * 50)
-    print(f'Gradient Norm is : {grad_norm}')
-
-def get_grad_norm(
-        parameters: _tensor_or_tensors, norm_type: float = 2.0) -> torch.Tensor:
-    r"""
-    Copy from torch.nn.utils.clip_grad_norm_
-
-    Clips gradient norm of an iterable of parameters.
-
-    The norm is computed over all gradients together, as if they were
-    concatenated into a single vector. Gradients are modified in-place.
-
-    Args:
-        parameters (Iterable[Tensor] or Tensor): an iterable of Tensors or a
-            single Tensor that will have gradients normalized
-        max_norm (float or int): max norm of the gradients
-        norm_type (float or int): type of the used p-norm. Can be ``'inf'`` for
-            infinity norm.
-        error_if_nonfinite (bool): if True, an error is thrown if the total
-            norm of the gradients from :attr:`parameters` is ``nan``,
-            ``inf``, or ``-inf``. Default: False (will switch to True in the future)
-
-    Returns:
-        Total norm of the parameter gradients (viewed as a single vector).
-    """
-    if isinstance(parameters, torch.Tensor):
-        parameters = [parameters]
-    grads = [p.grad for p in parameters if p.grad is not None]
-    norm_type = float(norm_type)
-    if len(grads) == 0:
-        return torch.tensor(0.)
-    device = grads[0].device
-    if norm_type == inf:
-        norms = [g.detach().abs().max().to(device) for g in grads]
-        total_norm = norms[0] if len(norms) == 1 else torch.max(torch.stack(norms))
-    else:
-        total_norm = torch.norm(torch.stack([torch.norm(g.detach(), norm_type).to(device) for g in grads]), norm_type)
-    return total_norm
-
-
-def clip_grad_norm_(
-        parameters: _tensor_or_tensors, max_norm: float, norm_type: float = 2.0,
-        error_if_nonfinite: bool = False, clip_grad=True) -> torch.Tensor:
-    r"""
-    Copy from torch.nn.utils.clip_grad_norm_
-
-    Clips gradient norm of an iterable of parameters.
-
-    The norm is computed over all gradients together, as if they were
-    concatenated into a single vector. Gradients are modified in-place.
-
-    Args:
-        parameters (Iterable[Tensor] or Tensor): an iterable of Tensors or a
-            single Tensor that will have gradients normalized
-        max_norm (float or int): max norm of the gradients
-        norm_type (float or int): type of the used p-norm. Can be ``'inf'`` for
-            infinity norm.
-        error_if_nonfinite (bool): if True, an error is thrown if the total
-            norm of the gradients from :attr:`parameters` is ``nan``,
-            ``inf``, or ``-inf``. Default: False (will switch to True in the future)
-
-    Returns:
-        Total norm of the parameter gradients (viewed as a single vector).
-    """
-    if isinstance(parameters, torch.Tensor):
-        parameters = [parameters]
-    grads = [p.grad for p in parameters if p.grad is not None]
-    max_norm = float(max_norm)
-    norm_type = float(norm_type)
-    if len(grads) == 0:
-        return torch.tensor(0.)
-    device = grads[0].device
-    if norm_type == inf:
-        norms = [g.detach().abs().max().to(device) for g in grads]
-        total_norm = norms[0] if len(norms) == 1 else torch.max(torch.stack(norms))
-    else:
-        total_norm = torch.norm(torch.stack([torch.norm(g.detach(), norm_type).to(device) for g in grads]), norm_type)
-
-    if clip_grad:
-        if error_if_nonfinite and torch.logical_or(total_norm.isnan(), total_norm.isinf()):
-            raise RuntimeError(
-                f'The total norm of order {norm_type} for gradients from '
-                '`parameters` is non-finite, so it cannot be clipped. To disable '
-                'this error and scale the gradients by the non-finite norm anyway, '
-                'set `error_if_nonfinite=False`')
-        clip_coef = max_norm / (total_norm + 1e-6)
-        # Note: multiplying by the clamped coef is redundant when the coef is clamped to 1, but doing so
-        # avoids a `if clip_coef < 1:` conditional which can require a CPU <=> device synchronization
-        # when the gradients do not reside in CPU memory.
-        clip_coef_clamped = torch.clamp(clip_coef, max=1.0)
-        for g in grads:
-            g.detach().mul_(clip_coef_clamped.to(g.device))
-        # gradient_cliped = torch.norm(torch.stack([torch.norm(g.detach(), norm_type).to(device) for g in grads]), norm_type)
-        # print(gradient_cliped)
-    return total_norm
-
-
-def get_experiment_dir(root_dir, args):
-    # if args.pretrained is not None and 'Latte-XL-2-256x256.pt' not in args.pretrained:
-    #     root_dir += '-WOPRE'
-    if args.use_compile:
-        root_dir += '-Compile'  # speedup by torch compile
-    if args.attention_mode:
-        root_dir += f'-{args.attention_mode.upper()}'
-    # if args.enable_xformers_memory_efficient_attention:
-    #     root_dir += '-Xfor'
-    if args.gradient_checkpointing:
-        root_dir += '-Gc'
-    if args.mixed_precision:
-        root_dir += f'-{args.mixed_precision.upper()}'
-    root_dir += f'-{args.max_image_size}'
-    return root_dir
-
-def get_precision(args):
-    if args.mixed_precision == "bf16":
-        dtype = torch.bfloat16
-    elif args.mixed_precision == "fp16":
-        dtype = torch.float16
-    else:
-        dtype = torch.float32
-    return dtype
-
-#################################################################################
-#                             Training Logger                                   #
-#################################################################################
-
-def create_logger(logging_dir):
-    """
-    Create a logger that writes to a log file and stdout.
-    """
-    if dist.get_rank() == 0:  # real logger
-        logging.basicConfig(
-            level=logging.INFO,
-            # format='[\033[34m%(asctime)s\033[0m] %(message)s',
-            format='[%(asctime)s] %(message)s',
-            datefmt='%Y-%m-%d %H:%M:%S',
-            handlers=[logging.StreamHandler(), logging.FileHandler(f"{logging_dir}/log.txt")]
-        )
-        logger = logging.getLogger(__name__)
-
-    else:  # dummy logger (does nothing)
-        logger = logging.getLogger(__name__)
-        logger.addHandler(logging.NullHandler())
-    return logger
-
-
-def create_tensorboard(tensorboard_dir):
-    """
-    Create a tensorboard that saves losses.
-    """
-    if dist.get_rank() == 0:  # real tensorboard
-        # tensorboard
-        writer = SummaryWriter(tensorboard_dir)
-
-    return writer
-
-
-def write_tensorboard(writer, *args):
-    '''
-    write the loss information to a tensorboard file.
-    Only for pytorch DDP mode.
-    '''
-    if dist.get_rank() == 0:  # real tensorboard
-        writer.add_scalar(args[0], args[1], args[2])
-
-
-#################################################################################
-#                      EMA Update/ DDP Training Utils                           #
-#################################################################################
-
-@torch.no_grad()
-def update_ema(ema_model, model, decay=0.9999):
-    """
-    Step the EMA model towards the current model.
-    """
-    ema_params = OrderedDict(ema_model.named_parameters())
-    model_params = OrderedDict(model.named_parameters())
-
-    for name, param in model_params.items():
-        # TODO: Consider applying only to params that require_grad to avoid small numerical changes of pos_embed
-        ema_params[name].mul_(decay).add_(param.data, alpha=1 - decay)
-
-
-def requires_grad(model, flag=True):
-    """
-    Set requires_grad flag for all parameters in a model.
-    """
-    for p in model.parameters():
-        p.requires_grad = flag
-
-
-def cleanup():
-    """
-    End DDP training.
-    """
-    dist.destroy_process_group()
-
-
-def setup_distributed(backend="nccl", port=None):
-    """Initialize distributed training environment.
-    support both slurm and torch.distributed.launch
-    see torch.distributed.init_process_group() for more details
-    """
-    num_gpus = torch.cuda.device_count()
-
-    if "SLURM_JOB_ID" in os.environ:
-        rank = int(os.environ["SLURM_PROCID"])
-        world_size = int(os.environ["SLURM_NTASKS"])
-        node_list = os.environ["SLURM_NODELIST"]
-        addr = subprocess.getoutput(f"scontrol show hostname {node_list} | head -n1")
-        # specify master port
-        if port is not None:
-            os.environ["MASTER_PORT"] = str(port)
-        elif "MASTER_PORT" not in os.environ:
-            # os.environ["MASTER_PORT"] = "29566"
-            os.environ["MASTER_PORT"] = str(29567 + num_gpus)
-        if "MASTER_ADDR" not in os.environ:
-            os.environ["MASTER_ADDR"] = addr
-        os.environ["WORLD_SIZE"] = str(world_size)
-        os.environ["LOCAL_RANK"] = str(rank % num_gpus)
-        os.environ["RANK"] = str(rank)
-    else:
-        rank = int(os.environ["RANK"])
-        world_size = int(os.environ["WORLD_SIZE"])
-
-    # torch.cuda.set_device(rank % num_gpus)
-
-    dist.init_process_group(
-        backend=backend,
-        world_size=world_size,
-        rank=rank,
-    )
-
-
-#################################################################################
-#                             Testing  Utils                                    #
-#################################################################################
-
-def save_video_grid(video, nrow=None):
-    b, t, h, w, c = video.shape
-
-    if nrow is None:
-        nrow = math.ceil(math.sqrt(b))
-    ncol = math.ceil(b / nrow)
-    padding = 1
-    video_grid = torch.zeros((t, (padding + h) * nrow + padding,
-                              (padding + w) * ncol + padding, c), dtype=torch.uint8)
-
-    print(video_grid.shape)
-    for i in range(b):
-        r = i // ncol
-        c = i % ncol
-        start_r = (padding + h) * r
-        start_c = (padding + w) * c
-        video_grid[:, start_r:start_r + h, start_c:start_c + w] = video[i]
-
-    return video_grid
-
-
-#################################################################################
-#                             MMCV  Utils                                    #
-#################################################################################
-
-
-def collect_env():
-    # Copyright (c) OpenMMLab. All rights reserved.
-    from mmcv.utils import collect_env as collect_base_env
-    from mmcv.utils import get_git_hash
-    """Collect the information of the running environments."""
-
-    env_info = collect_base_env()
-    env_info['MMClassification'] = get_git_hash()[:7]
-
-    for name, val in env_info.items():
-        print(f'{name}: {val}')
-
-    print(torch.cuda.get_arch_list())
-    print(torch.version.cuda)
-
-
-#################################################################################
-#                          Pixart-alpha  Utils                                  #
-#################################################################################
-
-
-bad_punct_regex = re.compile(r'['+'#®•©™&@·º½¾¿¡§~'+'\)'+'\('+'\]'+'\['+'\}'+'\{'+'\|'+'\\'+'\/'+'\*' + r']{1,}')  # noqa
-
-def text_preprocessing(text):
-    # The exact text cleaning as was in the training stage:
-    text = clean_caption(text)
-    text = clean_caption(text)
-    return text
-
-def basic_clean(text):
-    text = ftfy.fix_text(text)
-    text = html.unescape(html.unescape(text))
-    return text.strip()
-
-def clean_caption(caption):
-    caption = str(caption)
-    caption = ul.unquote_plus(caption)
-    caption = caption.strip().lower()
-    caption = re.sub('<person>', 'person', caption)
-    # urls:
-    caption = re.sub(
-        r'\b((?:https?:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))',  # noqa
-        '', caption)  # regex for urls
-    caption = re.sub(
-        r'\b((?:www:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))',  # noqa
-        '', caption)  # regex for urls
-    # html:
-    caption = BeautifulSoup(caption, features='html.parser').text
-
-    # @<nickname>
-    caption = re.sub(r'@[\w\d]+\b', '', caption)
-
-    # 31C0—31EF CJK Strokes
-    # 31F0—31FF Katakana Phonetic Extensions
-    # 3200—32FF Enclosed CJK Letters and Months
-    # 3300—33FF CJK Compatibility
-    # 3400—4DBF CJK Unified Ideographs Extension A
-    # 4DC0—4DFF Yijing Hexagram Symbols
-    # 4E00—9FFF CJK Unified Ideographs
-    caption = re.sub(r'[\u31c0-\u31ef]+', '', caption)
-    caption = re.sub(r'[\u31f0-\u31ff]+', '', caption)
-    caption = re.sub(r'[\u3200-\u32ff]+', '', caption)
-    caption = re.sub(r'[\u3300-\u33ff]+', '', caption)
-    caption = re.sub(r'[\u3400-\u4dbf]+', '', caption)
-    caption = re.sub(r'[\u4dc0-\u4dff]+', '', caption)
-    caption = re.sub(r'[\u4e00-\u9fff]+', '', caption)
-    #######################################################
-
-    # все виды тире / all types of dash --> "-"
-    caption = re.sub(
-        r'[\u002D\u058A\u05BE\u1400\u1806\u2010-\u2015\u2E17\u2E1A\u2E3A\u2E3B\u2E40\u301C\u3030\u30A0\uFE31\uFE32\uFE58\uFE63\uFF0D]+',  # noqa
-        '-', caption)
-
-    # кавычки к одному стандарту
-    caption = re.sub(r'[`´«»“”¨]', '"', caption)
-    caption = re.sub(r'[‘’]', "'", caption)
-
-    # &quot;
-    caption = re.sub(r'&quot;?', '', caption)
-    # &amp
-    caption = re.sub(r'&amp', '', caption)
-
-    # ip adresses:
-    caption = re.sub(r'\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}', ' ', caption)
-
-    # article ids:
-    caption = re.sub(r'\d:\d\d\s+$', '', caption)
-
-    # \n
-    caption = re.sub(r'\\n', ' ', caption)
-
-    # "#123"
-    caption = re.sub(r'#\d{1,3}\b', '', caption)
-    # "#12345.."
-    caption = re.sub(r'#\d{5,}\b', '', caption)
-    # "123456.."
-    caption = re.sub(r'\b\d{6,}\b', '', caption)
-    # filenames:
-    caption = re.sub(r'[\S]+\.(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)', '', caption)
-
-    #
-    caption = re.sub(r'[\"\']{2,}', r'"', caption)  # """AUSVERKAUFT"""
-    caption = re.sub(r'[\.]{2,}', r' ', caption)  # """AUSVERKAUFT"""
-
-    caption = re.sub(bad_punct_regex, r' ', caption)  # ***AUSVERKAUFT***, #AUSVERKAUFT
-    caption = re.sub(r'\s+\.\s+', r' ', caption)  # " . "
-
-    # this-is-my-cute-cat / this_is_my_cute_cat
-    regex2 = re.compile(r'(?:\-|\_)')
-    if len(re.findall(regex2, caption)) > 3:
-        caption = re.sub(regex2, ' ', caption)
-
-    caption = basic_clean(caption)
-
-    caption = re.sub(r'\b[a-zA-Z]{1,3}\d{3,15}\b', '', caption)  # jc6640
-    caption = re.sub(r'\b[a-zA-Z]+\d+[a-zA-Z]+\b', '', caption)  # jc6640vc
-    caption = re.sub(r'\b\d+[a-zA-Z]+\d+\b', '', caption)  # 6640vc231
-
-    caption = re.sub(r'(worldwide\s+)?(free\s+)?shipping', '', caption)
-    caption = re.sub(r'(free\s)?download(\sfree)?', '', caption)
-    caption = re.sub(r'\bclick\b\s(?:for|on)\s\w+', '', caption)
-    caption = re.sub(r'\b(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)(\simage[s]?)?', '', caption)
-    caption = re.sub(r'\bpage\s+\d+\b', '', caption)
-
-    caption = re.sub(r'\b\d*[a-zA-Z]+\d+[a-zA-Z]+\d+[a-zA-Z\d]*\b', r' ', caption)  # j2d1a2a...
-
-    caption = re.sub(r'\b\d+\.?\d*[xх×]\d+\.?\d*\b', '', caption)
-
-    caption = re.sub(r'\b\s+\:\s+', r': ', caption)
-    caption = re.sub(r'(\D[,\./])\b', r'\1 ', caption)
-    caption = re.sub(r'\s+', ' ', caption)
-
-    caption.strip()
-
-    caption = re.sub(r'^[\"\']([\w\W]+)[\"\']$', r'\1', caption)
-    caption = re.sub(r'^[\'\_,\-\:;]', r'', caption)
-    caption = re.sub(r'[\'\_,\-\:\-\+]$', r'', caption)
-    caption = re.sub(r'^\.\S+$', '', caption)
-
-    return caption.strip()
-
-
-#################################################################################
-#                          eval PSNR when training                                  #
-#################################################################################
-
-def resize(x, resolution):
-    height, width = x.shape[-2:]
-    aspect_ratio = width / height
-    if width <= height:
-        new_width = resolution
-        new_height = int(resolution / aspect_ratio)
-    else:
-        new_height = resolution
-        new_width = int(resolution * aspect_ratio)
-    resized_x = F.interpolate(x, size=(new_height, new_width), mode='bilinear', align_corners=True, antialias=True)
-    return resized_x
-
-def _preprocess(video_data, short_size=128, crop_size=None):
-    transform = Compose(
-        [
-            Lambda(lambda x: (x / 255.0)),
-            Lambda(lambda x: resize(x, short_size)),
-            (
-                CenterCropVideo(crop_size=crop_size)
-                if crop_size is not None
-                else Lambda(lambda x: x)
-            ),
-        ]
-    )
-    video_outputs = transform(video_data)
-    video_outputs = _format_video_shape(video_outputs)
-    return video_outputs
-
-def _format_video_shape(video, time_compress=4, spatial_compress=8):
-    time = video.shape[1]
-    height = video.shape[2]
-    width = video.shape[3]
-    new_time = (
-        (time - (time - 1) % time_compress)
-        if (time - 1) % time_compress != 0
-        else time
-    )
-    new_height = (
-        (height - (height) % spatial_compress)
-        if height % spatial_compress != 0
-        else height
-    )
-    new_width = (
-        (width - (width) % spatial_compress) if width % spatial_compress != 0 else width
-    )
-    return video[:, :new_time, :new_height, :new_width]
-
-def array_to_video(
-    image_array: npt.NDArray, fps: float = 30.0, output_file: str = "output_video.mp4"
-) -> None:
-    height, width, channels = image_array[0].shape
-    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
-    video_writer = cv2.VideoWriter(output_file, fourcc, float(fps), (width, height))
-
-    for image in image_array:
-        image_rgb = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
-        video_writer.write(image_rgb)
-
-    video_writer.release()
-
-
-def custom_to_video(
-    x: torch.Tensor, fps: float = 2.0, output_file: str = "output_video.mp4"
-) -> None:
-    x = x.detach().cpu()
-    x = torch.clamp(x, -1, 1)
-    x = (x + 1) / 2
-    x = x.permute(1, 2, 3, 0).float().numpy()
-    x = (255 * x).astype(np.uint8)
-    array_to_video(x, fps=fps, output_file=output_file)
-    return
-
-class RealVideoDataset(Dataset):
-    def __init__(
-        self,
-        real_video_dir,
-        num_frames,
-        sample_rate=1,
-        crop_size=None,
-        resolution=128,
-    ) -> None:
-        super().__init__()
-        self.real_video_files = self._combine_without_prefix(real_video_dir)
-        self.num_frames = num_frames
-        self.sample_rate = sample_rate
-        self.crop_size = crop_size
-        self.short_size = resolution
-
-    def __len__(self):
-        return len(self.real_video_files)
-
-    def __getitem__(self, index):
-        if index >= len(self):
-            raise IndexError
-        real_video_file = self.real_video_files[index]
-        real_video_tensor = self._load_video(real_video_file)
-        video_name = os.path.basename(real_video_file)
-        return {'video': real_video_tensor, 'file_name': video_name }
-
-    def _load_video(self, video_path):
-        num_frames = self.num_frames
-        sample_rate = self.sample_rate
-        decord_vr = VideoReader(video_path, ctx=cpu(0))
-        total_frames = len(decord_vr)
-        sample_frames_len = sample_rate * num_frames
-
-        if total_frames > sample_frames_len:
-            s = 0
-            e = s + sample_frames_len
-            num_frames = num_frames
-        else:
-            s = 0
-            e = total_frames
-            num_frames = int(total_frames / sample_frames_len * num_frames)
-            print(
-                f"sample_frames_len {sample_frames_len}, only can sample {num_frames * sample_rate}",
-                video_path,
-                total_frames,
-            )
-
-        frame_id_list = np.linspace(s, e - 1, num_frames, dtype=int)
-        video_data = decord_vr.get_batch(frame_id_list).asnumpy()
-        video_data = torch.from_numpy(video_data)
-        video_data = video_data.permute(3, 0, 1, 2)
-        return _preprocess(
-            video_data, short_size=self.short_size, crop_size=self.crop_size
-        )
-
-    def _combine_without_prefix(self, folder_path, prefix="."):
-        folder = []
-        for name in os.listdir(folder_path):
-            if name[0] == prefix:
-                continue
-            folder.append(os.path.join(folder_path, name))
-        folder.sort()
-        return folder

opensora/models/diffusion/__init__.py

@@ -1,11 +0,0 @@
-
-from .opensora.modeling_opensora import OpenSora_models
-
-
-Diffusion_models = {}
-Diffusion_models.update(OpenSora_models)
-
-from .opensora.modeling_opensora import OpenSora_models_class
-
-Diffusion_models_class = {}
-Diffusion_models_class.update(OpenSora_models_class)

opensora/models/diffusion/opensora/modeling_opensora.py

@@ -1,721 +0,0 @@
-import os
-import numpy as np
-from torch import nn
-import torch
-from einops import rearrange, repeat
-from typing import Any, Dict, Optional, Tuple
-from torch.nn import functional as F
-from diffusers.models.modeling_outputs import Transformer2DModelOutput
-from diffusers.utils import is_torch_version, deprecate
-from diffusers.configuration_utils import ConfigMixin, register_to_config
-from diffusers.models.modeling_utils import ModelMixin
-from diffusers.models.normalization import AdaLayerNormSingle
-from diffusers.models.embeddings import PixArtAlphaTextProjection
-from opensora.models.diffusion.opensora.modules import OverlapPatchEmbed3D, OverlapPatchEmbed2D, PatchEmbed2D, BasicTransformerBlock
-from opensora.utils.utils import to_2tuple
-
-from opensora.utils.parallel_states import get_sequence_parallel_state, nccl_info
-
-class OpenSoraT2V(ModelMixin, ConfigMixin):
-    """
-    A 2D Transformer model for image-like data.
-
-    Parameters:
-        num_attention_heads (`int`, *optional*, defaults to 16): The number of heads to use for multi-head attention.
-        attention_head_dim (`int`, *optional*, defaults to 88): The number of channels in each head.
-        in_channels (`int`, *optional*):
-            The number of channels in the input and output (specify if the input is **continuous**).
-        num_layers (`int`, *optional*, defaults to 1): The number of layers of Transformer blocks to use.
-        dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
-        cross_attention_dim (`int`, *optional*): The number of `encoder_hidden_states` dimensions to use.
-        sample_size (`int`, *optional*): The width of the latent images (specify if the input is **discrete**).
-            This is fixed during training since it is used to learn a number of position embeddings.
-        num_vector_embeds (`int`, *optional*):
-            The number of classes of the vector embeddings of the latent pixels (specify if the input is **discrete**).
-            Includes the class for the masked latent pixel.
-        activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to use in feed-forward.
-        num_embeds_ada_norm ( `int`, *optional*):
-            The number of diffusion steps used during training. Pass if at least one of the norm_layers is
-            `AdaLayerNorm`. This is fixed during training since it is used to learn a number of embeddings that are
-            added to the hidden states.
-
-            During inference, you can denoise for up to but not more steps than `num_embeds_ada_norm`.
-        attention_bias (`bool`, *optional*):
-            Configure if the `TransformerBlocks` attention should contain a bias parameter.
-    """
-
-    _supports_gradient_checkpointing = True
-
-    @register_to_config
-    def __init__(
-        self,
-        num_attention_heads: int = 16,
-        attention_head_dim: int = 88,
-        in_channels: Optional[int] = None,
-        out_channels: Optional[int] = None,
-        num_layers: int = 1,
-        dropout: float = 0.0,
-        norm_num_groups: int = 32,
-        cross_attention_dim: Optional[int] = None,
-        attention_bias: bool = False,
-        sample_size: Optional[int] = None,
-        sample_size_t: Optional[int] = None,
-        num_vector_embeds: Optional[int] = None,
-        patch_size: Optional[int] = None,
-        patch_size_t: 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",  # 'layer_norm', 'ada_norm', 'ada_norm_zero', 'ada_norm_single', 'ada_norm_continuous', 'layer_norm_i2vgen'
-        norm_elementwise_affine: bool = True,
-        norm_eps: float = 1e-5,
-        attention_type: str = "default",
-        caption_channels: int = None,
-        interpolation_scale_h: float = None,
-        interpolation_scale_w: float = None,
-        interpolation_scale_t: float = None,
-        use_additional_conditions: Optional[bool] = None,
-        attention_mode: str = 'xformers', 
-        downsampler: str = None, 
-        use_rope: bool = False,
-        use_stable_fp32: bool = False,
-    ):
-        super().__init__()
-
-        # Validate inputs.
-        if patch_size is not None:
-            if norm_type not in ["ada_norm", "ada_norm_zero", "ada_norm_single"]:
-                raise NotImplementedError(
-                    f"Forward pass is not implemented when `patch_size` is not None and `norm_type` is '{norm_type}'."
-                )
-            elif norm_type in ["ada_norm", "ada_norm_zero"] and num_embeds_ada_norm is None:
-                raise ValueError(
-                    f"When using a `patch_size` and this `norm_type` ({norm_type}), `num_embeds_ada_norm` cannot be None."
-                )
-
-        # Set some common variables used across the board.
-        self.use_rope = use_rope
-        self.use_linear_projection = use_linear_projection
-        self.interpolation_scale_t = interpolation_scale_t
-        self.interpolation_scale_h = interpolation_scale_h
-        self.interpolation_scale_w = interpolation_scale_w
-        self.downsampler = downsampler
-        self.caption_channels = caption_channels
-        self.num_attention_heads = num_attention_heads
-        self.attention_head_dim = attention_head_dim
-        self.inner_dim = self.config.num_attention_heads * self.config.attention_head_dim
-        self.in_channels = in_channels
-        self.out_channels = in_channels if out_channels is None else out_channels
-        self.gradient_checkpointing = False
-        self.config.hidden_size = self.inner_dim
-        use_additional_conditions = False
-        # if use_additional_conditions is None:
-            # if norm_type == "ada_norm_single" and sample_size == 128:
-            #     use_additional_conditions = True
-            # else:
-            # use_additional_conditions = False
-        self.use_additional_conditions = use_additional_conditions
-
-        # 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
-        assert 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"
-
-        # 2. Initialize the right blocks.
-        # Initialize the output blocks and other projection blocks when necessary.
-        self._init_patched_inputs(norm_type=norm_type)
-
-    def _init_patched_inputs(self, norm_type):
-        assert self.config.sample_size_t is not None, "OpenSoraT2V over patched input must provide sample_size_t"
-        assert self.config.sample_size is not None, "OpenSoraT2V over patched input must provide sample_size"
-        #assert not (self.config.sample_size_t == 1 and self.config.patch_size_t == 2), "Image do not need patchfy in t-dim"
-
-        self.num_frames = self.config.sample_size_t
-        self.config.sample_size = to_2tuple(self.config.sample_size)
-        self.height = self.config.sample_size[0]
-        self.width = self.config.sample_size[1]
-        self.patch_size_t = self.config.patch_size_t
-        self.patch_size = self.config.patch_size
-        interpolation_scale_t = ((self.config.sample_size_t - 1) // 16 + 1) if self.config.sample_size_t % 2 == 1 else self.config.sample_size_t / 16
-        interpolation_scale_t = (
-            self.config.interpolation_scale_t if self.config.interpolation_scale_t is not None else interpolation_scale_t
-        )
-        interpolation_scale = (
-            self.config.interpolation_scale_h if self.config.interpolation_scale_h is not None else self.config.sample_size[0] / 30, 
-            self.config.interpolation_scale_w if self.config.interpolation_scale_w is not None else self.config.sample_size[1] / 40, 
-        )
-        # if self.config.sample_size_t > 1:
-        #     self.pos_embed = PatchEmbed3D(
-        #         num_frames=self.config.sample_size_t,
-        #         height=self.config.sample_size[0],
-        #         width=self.config.sample_size[1],
-        #         patch_size_t=self.config.patch_size_t,
-        #         patch_size=self.config.patch_size,
-        #         in_channels=self.in_channels,
-        #         embed_dim=self.inner_dim,
-        #         interpolation_scale=interpolation_scale, 
-        #         interpolation_scale_t=interpolation_scale_t,
-        #     )
-        # else:
-        if self.config.downsampler is not None and len(self.config.downsampler) == 9:
-            self.pos_embed = OverlapPatchEmbed3D(
-                num_frames=self.config.sample_size_t,
-                height=self.config.sample_size[0],
-                width=self.config.sample_size[1],
-                patch_size_t=self.config.patch_size_t,
-                patch_size=self.config.patch_size,
-                in_channels=self.in_channels,
-                embed_dim=self.inner_dim,
-                interpolation_scale=interpolation_scale, 
-                interpolation_scale_t=interpolation_scale_t,
-                use_abs_pos=not self.config.use_rope, 
-            )
-        elif self.config.downsampler is not None and len(self.config.downsampler) == 7:
-            self.pos_embed = OverlapPatchEmbed2D(
-                num_frames=self.config.sample_size_t,
-                height=self.config.sample_size[0],
-                width=self.config.sample_size[1],
-                patch_size_t=self.config.patch_size_t,
-                patch_size=self.config.patch_size,
-                in_channels=self.in_channels,
-                embed_dim=self.inner_dim,
-                interpolation_scale=interpolation_scale, 
-                interpolation_scale_t=interpolation_scale_t,
-                use_abs_pos=not self.config.use_rope, 
-            )
-        
-        else:
-            self.pos_embed = PatchEmbed2D(
-                num_frames=self.config.sample_size_t,
-                height=self.config.sample_size[0],
-                width=self.config.sample_size[1],
-                patch_size_t=self.config.patch_size_t,
-                patch_size=self.config.patch_size,
-                in_channels=self.in_channels,
-                embed_dim=self.inner_dim,
-                interpolation_scale=interpolation_scale, 
-                interpolation_scale_t=interpolation_scale_t,
-                use_abs_pos=not self.config.use_rope, 
-            )
-        interpolation_scale_thw = (interpolation_scale_t, *interpolation_scale)
-        self.transformer_blocks = nn.ModuleList(
-            [
-                BasicTransformerBlock(
-                    self.inner_dim,
-                    self.config.num_attention_heads,
-                    self.config.attention_head_dim,
-                    dropout=self.config.dropout,
-                    cross_attention_dim=self.config.cross_attention_dim,
-                    activation_fn=self.config.activation_fn,
-                    num_embeds_ada_norm=self.config.num_embeds_ada_norm,
-                    attention_bias=self.config.attention_bias,
-                    only_cross_attention=self.config.only_cross_attention,
-                    double_self_attention=self.config.double_self_attention,
-                    upcast_attention=self.config.upcast_attention,
-                    norm_type=norm_type,
-                    norm_elementwise_affine=self.config.norm_elementwise_affine,
-                    norm_eps=self.config.norm_eps,
-                    attention_type=self.config.attention_type,
-                    attention_mode=self.config.attention_mode, 
-                    downsampler=self.config.downsampler, 
-                    use_rope=self.config.use_rope, 
-                    interpolation_scale_thw=interpolation_scale_thw, 
-                )
-                for _ in range(self.config.num_layers)
-            ]
-        )
-
-        if self.config.norm_type != "ada_norm_single":
-            self.norm_out = nn.LayerNorm(self.inner_dim, elementwise_affine=False, eps=1e-6)
-            self.proj_out_1 = nn.Linear(self.inner_dim, 2 * self.inner_dim)
-            self.proj_out_2 = nn.Linear(
-                self.inner_dim, self.config.patch_size_t * self.config.patch_size * self.config.patch_size * self.out_channels
-            )
-        elif self.config.norm_type == "ada_norm_single":
-            self.norm_out = nn.LayerNorm(self.inner_dim, elementwise_affine=False, eps=1e-6)
-            self.scale_shift_table = nn.Parameter(torch.randn(2, self.inner_dim) / self.inner_dim**0.5)
-            self.proj_out = nn.Linear(
-                self.inner_dim, self.config.patch_size_t * self.config.patch_size * self.config.patch_size * self.out_channels
-            )
-
-        # PixArt-Alpha blocks.
-        self.adaln_single = None
-        if self.config.norm_type == "ada_norm_single":
-            # 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(
-                self.inner_dim, use_additional_conditions=self.use_additional_conditions
-            )
-
-        self.caption_projection = None
-        if self.caption_channels is not None:
-            self.caption_projection = PixArtAlphaTextProjection(
-                in_features=self.caption_channels, hidden_size=self.inner_dim
-            )
-
-    def _set_gradient_checkpointing(self, module, value=False):
-        if hasattr(module, "gradient_checkpointing"):
-            module.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: Optional[int] = 0,
-        return_dict: bool = True,
-    ):
-        """
-        The [`Transformer2DModel`] forward method.
-
-        Args:
-            hidden_states (`torch.LongTensor` of shape `(batch size, num latent pixels)` if discrete, `torch.FloatTensor` of shape `(batch size, channel, height, width)` if continuous):
-                Input `hidden_states`.
-            encoder_hidden_states ( `torch.FloatTensor` of shape `(batch size, sequence len, embed dims)`, *optional*):
-                Conditional embeddings for cross attention layer. If not given, cross-attention defaults to
-                self-attention.
-            timestep ( `torch.LongTensor`, *optional*):
-                Used to indicate denoising step. Optional timestep to be applied as an embedding in `AdaLayerNorm`.
-            class_labels ( `torch.LongTensor` of shape `(batch size, num classes)`, *optional*):
-                Used to indicate class labels conditioning. Optional class labels to be applied as an embedding in
-                `AdaLayerZeroNorm`.
-            cross_attention_kwargs ( `Dict[str, Any]`, *optional*):
-                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
-                `self.processor` in
-                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
-            attention_mask ( `torch.Tensor`, *optional*):
-                An attention mask of shape `(batch, key_tokens)` is applied to `encoder_hidden_states`. If `1` the mask
-                is kept, otherwise if `0` it is discarded. Mask will be converted into a bias, which adds large
-                negative values to the attention scores corresponding to "discard" tokens.
-            encoder_attention_mask ( `torch.Tensor`, *optional*):
-                Cross-attention mask applied to `encoder_hidden_states`. Two formats supported:
-
-                    * Mask `(batch, sequence_length)` True = keep, False = discard.
-                    * Bias `(batch, 1, sequence_length)` 0 = keep, -10000 = discard.
-
-                If `ndim == 2`: will be interpreted as a mask, then converted into a bias consistent with the format
-                above. This bias will be added to the cross-attention scores.
-            return_dict (`bool`, *optional*, defaults to `True`):
-                Whether or not to return a [`~models.unets.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.
-        """
-        batch_size, c, frame, h, w = hidden_states.shape
-        # print('hidden_states.shape', hidden_states.shape)
-        frame = frame - use_image_num  # 21-4=17
-        if cross_attention_kwargs is not None:
-            if cross_attention_kwargs.get("scale", None) is not None:
-                print.warning("Passing `scale` to `cross_attention_kwargs` is deprecated. `scale` will be ignored.")
-        # 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)
-        attention_mask_vid, attention_mask_img = None, None
-        if attention_mask is not None and attention_mask.ndim == 4:
-            # 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)
-            # b, frame+use_image_num, h, w -> a video with images
-            # b, 1, h, w -> only images
-            attention_mask = attention_mask.to(self.dtype)
-            if get_sequence_parallel_state():
-
-                # print('before attention_mask.shape', attention_mask.shape)
-                attention_mask_vid = attention_mask[:, :frame * nccl_info.world_size]  # b, frame, h, w
-                attention_mask_img = attention_mask[:, frame * nccl_info.world_size:]  # b, use_image_num, h, w
-                # print('after attention_mask.shape', attention_mask_vid.shape)
-            else:
-                attention_mask_vid = attention_mask[:, :frame]  # b, frame, h, w
-                attention_mask_img = attention_mask[:, frame:]  # b, use_image_num, h, w
-
-            if attention_mask_vid.numel() > 0:
-                attention_mask_vid_first_frame = attention_mask_vid[:, :1].repeat(1, self.patch_size_t-1, 1, 1)
-                attention_mask_vid = torch.cat([attention_mask_vid_first_frame, attention_mask_vid], dim=1)
-                attention_mask_vid = attention_mask_vid.unsqueeze(1)  # b 1 t h w
-                attention_mask_vid = F.max_pool3d(attention_mask_vid, kernel_size=(self.patch_size_t, self.patch_size, self.patch_size), 
-                                                  stride=(self.patch_size_t, self.patch_size, self.patch_size))
-                attention_mask_vid = rearrange(attention_mask_vid, 'b 1 t h w -> (b 1) 1 (t h w)') 
-            if attention_mask_img.numel() > 0:
-                attention_mask_img = F.max_pool2d(attention_mask_img, kernel_size=(self.patch_size, self.patch_size), stride=(self.patch_size, self.patch_size))
-                attention_mask_img = rearrange(attention_mask_img, 'b i h w -> (b i) 1 (h w)') 
-
-            attention_mask_vid = (1 - attention_mask_vid.bool().to(self.dtype)) * -10000.0 if attention_mask_vid.numel() > 0 else None
-            attention_mask_img = (1 - attention_mask_img.bool().to(self.dtype)) * -10000.0 if attention_mask_img.numel() > 0 else None
-
-            if frame == 1 and use_image_num == 0 and not get_sequence_parallel_state():
-                attention_mask_img = attention_mask_vid
-                attention_mask_vid = None
-        # convert encoder_attention_mask to a bias the same way we do for attention_mask
-        # import ipdb;ipdb.set_trace()
-        if encoder_attention_mask is not None and encoder_attention_mask.ndim == 3:  
-            # b, 1+use_image_num, l -> a video with images
-            # b, 1, l -> only images
-            encoder_attention_mask = (1 - encoder_attention_mask.to(self.dtype)) * -10000.0
-            in_t = encoder_attention_mask.shape[1]
-            encoder_attention_mask_vid = encoder_attention_mask[:, :in_t-use_image_num]  # b, 1, l
-            encoder_attention_mask_vid = rearrange(encoder_attention_mask_vid, 'b 1 l -> (b 1) 1 l') if encoder_attention_mask_vid.numel() > 0 else None
-
-            encoder_attention_mask_img = encoder_attention_mask[:, in_t-use_image_num:]  # b, use_image_num, l
-            encoder_attention_mask_img = rearrange(encoder_attention_mask_img, 'b i l -> (b i) 1 l') if encoder_attention_mask_img.numel() > 0 else None
-
-            if frame == 1 and use_image_num == 0 and not get_sequence_parallel_state():
-                encoder_attention_mask_img = encoder_attention_mask_vid
-                encoder_attention_mask_vid = None
-
-
-
-
-        # 1. Input
-        frame = ((frame - 1) // self.patch_size_t + 1) if frame % 2 == 1 else frame // self.patch_size_t  # patchfy
-        # print('frame', frame)
-        height, width = hidden_states.shape[-2] // self.patch_size, hidden_states.shape[-1] // self.patch_size
-
-        added_cond_kwargs = {"resolution": None, "aspect_ratio": None}
-        hidden_states_vid, hidden_states_img, encoder_hidden_states_vid, encoder_hidden_states_img, \
-        timestep_vid, timestep_img, embedded_timestep_vid, embedded_timestep_img = self._operate_on_patched_inputs(
-            hidden_states, encoder_hidden_states, timestep, added_cond_kwargs, batch_size, frame, use_image_num
-        )
-        # 2. Blocks
-        # import ipdb;ipdb.set_trace()
-        if get_sequence_parallel_state():
-            if hidden_states_vid is not None:
-                # print(333333333333333)
-                hidden_states_vid = rearrange(hidden_states_vid, 'b s h -> s b h', b=batch_size).contiguous()
-                encoder_hidden_states_vid = rearrange(encoder_hidden_states_vid, 'b s h -> s b h',
-                                                      b=batch_size).contiguous()
-                timestep_vid = timestep_vid.view(batch_size, 6, -1).transpose(0, 1).contiguous()
-                # print('timestep_vid', timestep_vid.shape)
-
-                
-        for block in self.transformer_blocks:
-            if self.training and self.gradient_checkpointing:
-
-                def create_custom_forward(module, return_dict=None):
-                    def custom_forward(*inputs):
-                        if return_dict is not None:
-                            return module(*inputs, return_dict=return_dict)
-                        else:
-                            return module(*inputs)
-
-                    return custom_forward
-
-                ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
-                # import ipdb;ipdb.set_trace()
-                if hidden_states_vid is not None:
-                    hidden_states_vid = torch.utils.checkpoint.checkpoint(
-                        create_custom_forward(block),
-                        hidden_states_vid,
-                        attention_mask_vid,
-                        encoder_hidden_states_vid,
-                        encoder_attention_mask_vid,
-                        timestep_vid,
-                        cross_attention_kwargs,
-                        class_labels,
-                        frame, 
-                        height, 
-                        width, 
-                        **ckpt_kwargs,
-                    )
-                # import ipdb;ipdb.set_trace()
-                if hidden_states_img is not None:
-                    hidden_states_img = torch.utils.checkpoint.checkpoint(
-                        create_custom_forward(block),
-                        hidden_states_img,
-                        attention_mask_img,
-                        encoder_hidden_states_img,
-                        encoder_attention_mask_img,
-                        timestep_img,
-                        cross_attention_kwargs,
-                        class_labels,
-                        1, 
-                        height, 
-                        width, 
-                        **ckpt_kwargs,
-                    )
-            else:
-                if hidden_states_vid is not None:
-                    hidden_states_vid = block(
-                        hidden_states_vid,
-                        attention_mask=attention_mask_vid,
-                        encoder_hidden_states=encoder_hidden_states_vid,
-                        encoder_attention_mask=encoder_attention_mask_vid,
-                        timestep=timestep_vid,
-                        cross_attention_kwargs=cross_attention_kwargs,
-                        class_labels=class_labels,
-                        frame=frame, 
-                        height=height, 
-                        width=width, 
-                    )
-                if hidden_states_img is not None:
-                    hidden_states_img = block(
-                        hidden_states_img,
-                        attention_mask=attention_mask_img,
-                        encoder_hidden_states=encoder_hidden_states_img,
-                        encoder_attention_mask=encoder_attention_mask_img,
-                        timestep=timestep_img,
-                        cross_attention_kwargs=cross_attention_kwargs,
-                        class_labels=class_labels,
-                        frame=1, 
-                        height=height, 
-                        width=width, 
-                    )
-
-        if get_sequence_parallel_state():
-            if hidden_states_vid is not None:
-                hidden_states_vid = rearrange(hidden_states_vid, 's b h -> b s h', b=batch_size).contiguous()
-
-        # 3. Output
-        output_vid, output_img = None, None 
-        if hidden_states_vid is not None:
-            output_vid = self._get_output_for_patched_inputs(
-                hidden_states=hidden_states_vid,
-                timestep=timestep_vid,
-                class_labels=class_labels,
-                embedded_timestep=embedded_timestep_vid,
-                num_frames=frame, 
-                height=height,
-                width=width,
-            )  # b c t h w
-        if hidden_states_img is not None:
-            output_img = self._get_output_for_patched_inputs(
-                hidden_states=hidden_states_img,
-                timestep=timestep_img,
-                class_labels=class_labels,
-                embedded_timestep=embedded_timestep_img,
-                num_frames=1, 
-                height=height,
-                width=width,
-            )  # b c 1 h w
-            if use_image_num != 0:
-                output_img = rearrange(output_img, '(b i) c 1 h w -> b c i h w', i=use_image_num)
-
-        if output_vid is not None and output_img is not None:
-            output = torch.cat([output_vid, output_img], dim=2)
-        elif output_vid is not None:
-            output = output_vid
-        elif output_img is not None:
-            output = output_img
-
-        if not return_dict:
-            return (output,)
-
-        return Transformer2DModelOutput(sample=output)
-
-
-    def _operate_on_patched_inputs(self, hidden_states, encoder_hidden_states, timestep, added_cond_kwargs, batch_size, frame, use_image_num):
-        # batch_size = hidden_states.shape[0]
-        hidden_states_vid, hidden_states_img = self.pos_embed(hidden_states.to(self.dtype), frame)
-        timestep_vid, timestep_img = None, None
-        embedded_timestep_vid, embedded_timestep_img = None, None
-        encoder_hidden_states_vid, encoder_hidden_states_img = None, None
-
-        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`."
-                )
-            timestep, embedded_timestep = self.adaln_single(
-                timestep, added_cond_kwargs, batch_size=batch_size, hidden_dtype=self.dtype
-            )  # b 6d, b d
-            if hidden_states_vid is None:
-                timestep_img = timestep
-                embedded_timestep_img = embedded_timestep
-            else:
-                timestep_vid = timestep
-                embedded_timestep_vid = embedded_timestep
-                if hidden_states_img is not None:
-                    timestep_img = repeat(timestep, 'b d -> (b i) d', i=use_image_num).contiguous()
-                    embedded_timestep_img = repeat(embedded_timestep, 'b d -> (b i) d', i=use_image_num).contiguous()
-
-        if self.caption_projection is not None:
-            encoder_hidden_states = self.caption_projection(encoder_hidden_states)  # b, 1+use_image_num, l, d or b, 1, l, d
-            if hidden_states_vid is None:
-                encoder_hidden_states_img = rearrange(encoder_hidden_states, 'b 1 l d -> (b 1) l d')
-            else:
-                encoder_hidden_states_vid = rearrange(encoder_hidden_states[:, :1], 'b 1 l d -> (b 1) l d')
-                if hidden_states_img is not None:
-                    encoder_hidden_states_img = rearrange(encoder_hidden_states[:, 1:], 'b i l d -> (b i) l d')
-
-
-        return hidden_states_vid, hidden_states_img, encoder_hidden_states_vid, encoder_hidden_states_img, timestep_vid, timestep_img, embedded_timestep_vid, embedded_timestep_img
-
-    
-    
-    def _get_output_for_patched_inputs(
-        self, hidden_states, timestep, class_labels, embedded_timestep, num_frames, height=None, width=None
-    ):  
-        # import ipdb;ipdb.set_trace()
-        if self.config.norm_type != "ada_norm_single":
-            conditioning = self.transformer_blocks[0].norm1.emb(
-                timestep, class_labels, hidden_dtype=self.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":
-            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)
-            hidden_states = hidden_states.squeeze(1)
-
-        # unpatchify
-        if self.adaln_single is None:
-            height = width = int(hidden_states.shape[1] ** 0.5)
-        hidden_states = hidden_states.reshape(
-            shape=(-1, num_frames, height, width, self.patch_size_t, self.patch_size, self.patch_size, self.out_channels)
-        )
-        hidden_states = torch.einsum("nthwopqc->nctohpwq", hidden_states)
-        output = hidden_states.reshape(
-            shape=(-1, self.out_channels, num_frames * self.patch_size_t, height * self.patch_size, width * self.patch_size)
-        )
-        # import ipdb;ipdb.set_trace()
-        # if output.shape[2] % 2 == 0:
-        #     output = output[:, :, 1:]
-        return output
-     
-def OpenSoraT2V_mT5_Baes_S_122(**kwargs):
-    return OpenSoraT2V(num_layers=28, attention_head_dim=96, num_attention_heads=16, patch_size_t=1, patch_size=2,
-                       norm_type="ada_norm_single", caption_channels=768, cross_attention_dim=1536, **kwargs)
-    
-def OpenSoraT2V_S_122(**kwargs):
-    return OpenSoraT2V(num_layers=28, attention_head_dim=96, num_attention_heads=16, patch_size_t=1, patch_size=2,
-                       norm_type="ada_norm_single", caption_channels=4096, cross_attention_dim=1536, **kwargs)
-
-def OpenSoraT2V_B_122(**kwargs):
-    return OpenSoraT2V(num_layers=32, attention_head_dim=96, num_attention_heads=16, patch_size_t=1, patch_size=2,
-                       norm_type="ada_norm_single", caption_channels=4096, cross_attention_dim=1920, **kwargs)
-
-def OpenSoraT2V_L_122(**kwargs):
-    return OpenSoraT2V(num_layers=40, attention_head_dim=128, num_attention_heads=16, patch_size_t=1, patch_size=2,
-                       norm_type="ada_norm_single", caption_channels=4096, cross_attention_dim=2048, **kwargs)
-
-def OpenSoraT2V_ROPE_L_122(**kwargs):
-    return OpenSoraT2V(num_layers=32, attention_head_dim=96, num_attention_heads=24, patch_size_t=1, patch_size=2,
-                       norm_type="ada_norm_single", caption_channels=4096, cross_attention_dim=2304, **kwargs)
-
-OpenSora_models = {
-    "OpenSoraT2V-mT5-Base-S/122": OpenSoraT2V_mT5_Baes_S_122,
-    "OpenSoraT2V-S/122": OpenSoraT2V_S_122,  #       1.1B
-    "OpenSoraT2V-B/122": OpenSoraT2V_B_122,
-    "OpenSoraT2V-L/122": OpenSoraT2V_L_122,
-    "OpenSoraT2V-ROPE-L/122": OpenSoraT2V_ROPE_L_122,
-}
-
-OpenSora_models_class = {
-    "OpenSoraT2V-mT5-Base-S/122": OpenSoraT2V,
-    "OpenSoraT2V-S/122": OpenSoraT2V,
-    "OpenSoraT2V-B/122": OpenSoraT2V,
-    "OpenSoraT2V-L/122": OpenSoraT2V,
-    "OpenSoraT2V-ROPE-L/122": OpenSoraT2V,
-}
-
-if __name__ == '__main__':
-    from opensora.models.causalvideovae import ae_stride_config
-
-    args = type('args', (), 
-    {
-        'ae': 'CausalVAEModel_4x8x8', 
-        'attention_mode': 'xformers', 
-        'use_rope': True, 
-        'model_max_length': 300, 
-        'max_height': 320,
-        'max_width': 240,
-        'num_frames': 1,
-        'use_image_num': 0, 
-        'interpolation_scale_t': 1,
-        'interpolation_scale_h': 1,
-        'interpolation_scale_w': 1,
-    }
-    )
-    b = 16
-    c = 8
-    cond_c = 4096
-    num_timesteps = 1000
-    ae_stride_t, ae_stride_h, ae_stride_w = ae_stride_config[args.ae]
-    latent_size = (args.max_height // ae_stride_h, args.max_width // ae_stride_w)
-    num_frames = (args.num_frames - 1) // ae_stride_t + 1
-
-    device = torch.device('cuda:0')
-    model = OpenSoraT2V_ROPE_L_122(in_channels=c, 
-                              out_channels=c, 
-                              sample_size=latent_size, 
-                              sample_size_t=num_frames, 
-                              activation_fn="gelu-approximate",
-                            attention_bias=True,
-                            attention_type="default",
-                            double_self_attention=False,
-                            norm_elementwise_affine=False,
-                            norm_eps=1e-06,
-                            norm_num_groups=32,
-                            num_vector_embeds=None,
-                            only_cross_attention=False,
-                            upcast_attention=False,
-                            use_linear_projection=False,
-                            use_additional_conditions=False, 
-                            downsampler=None,
-                            interpolation_scale_t=args.interpolation_scale_t, 
-                            interpolation_scale_h=args.interpolation_scale_h, 
-                            interpolation_scale_w=args.interpolation_scale_w, 
-                            use_rope=args.use_rope, 
-                            ).to(device)
-    
-    try:
-        path = "PixArt-Alpha-XL-2-512.safetensors"
-        from safetensors.torch import load_file as safe_load
-        ckpt = safe_load(path, device="cpu")
-        # import ipdb;ipdb.set_trace()
-        if ckpt['pos_embed.proj.weight'].shape != model.pos_embed.proj.weight.shape and ckpt['pos_embed.proj.weight'].ndim == 4:
-            repeat = model.pos_embed.proj.weight.shape[2]
-            ckpt['pos_embed.proj.weight'] = ckpt['pos_embed.proj.weight'].unsqueeze(2).repeat(1, 1, repeat, 1, 1) / float(repeat)
-            del ckpt['proj_out.weight'], ckpt['proj_out.bias']
-        msg = model.load_state_dict(ckpt, strict=False)
-        print(msg)
-    except Exception as e:
-        print(e)
-    print(model)
-    print(f'{sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e9} B')
-    # import sys;sys.exit()
-    x = torch.randn(b, c,  1+(args.num_frames-1)//ae_stride_t+args.use_image_num, args.max_height//ae_stride_h, args.max_width//ae_stride_w).to(device)
-    cond = torch.randn(b, 1+args.use_image_num, args.model_max_length, cond_c).to(device)
-    attn_mask = torch.randint(0, 2, (b, 1+(args.num_frames-1)//ae_stride_t+args.use_image_num, args.max_height//ae_stride_h, args.max_width//ae_stride_w)).to(device)  # B L or B 1+num_images L
-    cond_mask = torch.randint(0, 2, (b, 1+args.use_image_num, args.model_max_length)).to(device)  # B L or B 1+num_images L
-    timestep = torch.randint(0, 1000, (b,), device=device)
-    model_kwargs = dict(hidden_states=x, encoder_hidden_states=cond, attention_mask=attn_mask, 
-                        encoder_attention_mask=cond_mask, use_image_num=args.use_image_num, timestep=timestep)
-    with torch.no_grad():
-        output = model(**model_kwargs)
-    print(output[0].shape)
-
-
-
-

opensora/models/diffusion/opensora/rope.py

@@ -1,91 +0,0 @@
-import torch
-
-from opensora.utils.parallel_states import get_sequence_parallel_state
-
-class PositionGetter3D(object):
-    """ return positions of patches """
-
-    def __init__(self, ):
-        self.cache_positions = {}
-        
-    def __call__(self, b, t, h, w, device):
-        if not (b, t,h,w) in self.cache_positions:
-            x = torch.arange(w, device=device)
-            y = torch.arange(h, device=device)
-            z = torch.arange(t, device=device)
-            pos = torch.cartesian_prod(z, y, x)
-            if get_sequence_parallel_state():
-                # print('PositionGetter3D', PositionGetter3D)
-                pos = pos.reshape(t * h * w, 3).transpose(0, 1).reshape(3, -1, 1).contiguous().expand(3, -1, b).clone()
-            else:
-                pos = pos.reshape(t * h * w, 3).transpose(0, 1).reshape(3, 1, -1).contiguous().expand(3, b, -1).clone()
-            poses = (pos[0].contiguous(), pos[1].contiguous(), pos[2].contiguous())
-            max_poses = (int(poses[0].max()), int(poses[1].max()), int(poses[2].max()))
-
-            self.cache_positions[b, t, h, w] = (poses, max_poses)
-        pos = self.cache_positions[b, t, h, w]
-
-        return pos
-    
-
-class RoPE3D(torch.nn.Module):
-
-    def __init__(self, freq=10000.0, F0=1.0, interpolation_scale_thw=(1, 1, 1)):
-        super().__init__()
-        self.base = freq
-        self.F0 = F0
-        self.interpolation_scale_t = interpolation_scale_thw[0]
-        self.interpolation_scale_h = interpolation_scale_thw[1]
-        self.interpolation_scale_w = interpolation_scale_thw[2]
-        self.cache = {}
-
-    def get_cos_sin(self, D, seq_len, device, dtype, interpolation_scale=1):
-        if (D, seq_len, device, dtype) not in self.cache:
-            inv_freq = 1.0 / (self.base ** (torch.arange(0, D, 2).float().to(device) / D))
-            t = torch.arange(seq_len, device=device, dtype=inv_freq.dtype) / interpolation_scale
-            freqs = torch.einsum("i,j->ij", t, inv_freq).to(dtype)
-            freqs = torch.cat((freqs, freqs), dim=-1)
-            cos = freqs.cos()  # (Seq, Dim)
-            sin = freqs.sin()
-            self.cache[D, seq_len, device, dtype] = (cos, sin)
-        return self.cache[D, seq_len, device, dtype]
-
-    @staticmethod
-    def rotate_half(x):
-        x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2:]
-        return torch.cat((-x2, x1), dim=-1)
-
-    def apply_rope1d(self, tokens, pos1d, cos, sin):
-        assert pos1d.ndim == 2
-        if not get_sequence_parallel_state():
-            # for (batch_size x nheads x ntokens x dim)
-            cos = torch.nn.functional.embedding(pos1d, cos)[:, None, :, :]
-            sin = torch.nn.functional.embedding(pos1d, sin)[:, None, :, :]
-        else: 
-            cos = torch.nn.functional.embedding(pos1d, cos)[:, :, None, :]
-            sin = torch.nn.functional.embedding(pos1d, sin)[:, :, None, :]
-
-        return (tokens * cos) + (self.rotate_half(tokens) * sin)
-
-    def forward(self, tokens, positions):
-        """
-        input:
-            * tokens: batch_size x nheads x ntokens x dim
-            * positions: batch_size x ntokens x 3 (t, y and x position of each token)
-        output:
-            * tokens after appplying RoPE3D (batch_size x nheads x ntokens x x dim)
-        """
-        assert tokens.size(3) % 3 == 0, "number of dimensions should be a multiple of three"
-        D = tokens.size(3) // 3
-        poses, max_poses = positions
-        assert len(poses) == 3 and poses[0].ndim == 2# Batch, Seq, 3
-        cos_t, sin_t = self.get_cos_sin(D, max_poses[0] + 1, tokens.device, tokens.dtype, self.interpolation_scale_t)
-        cos_y, sin_y = self.get_cos_sin(D, max_poses[1] + 1, tokens.device, tokens.dtype, self.interpolation_scale_h)
-        cos_x, sin_x = self.get_cos_sin(D, max_poses[2] + 1, tokens.device, tokens.dtype, self.interpolation_scale_w)
-        # split features into three along the feature dimension, and apply rope1d on each half
-        t, y, x = tokens.chunk(3, dim=-1)
-        t = self.apply_rope1d(t, poses[0], cos_t, sin_t)
-        y = self.apply_rope1d(y, poses[1], cos_y, sin_y)
-        x = self.apply_rope1d(x, poses[2], cos_x, sin_x)
-        tokens = torch.cat((t, y, x), dim=-1)
-        return tokens

opensora/models/diffusion/utils/curope/__init__.py

@@ -1,4 +0,0 @@
-# Copyright (C) 2022-present Naver Corporation. All rights reserved.
-# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
-
-from .curope2d import cuRoPE2D

opensora/models/diffusion/utils/curope/curope.cpp

@@ -1,69 +0,0 @@
-/* 
-  Copyright (C) 2022-present Naver Corporation. All rights reserved.
-  Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
-*/
-
-#include <torch/extension.h>
-
-// forward declaration
-void rope_2d_cuda( torch::Tensor tokens, const torch::Tensor pos, const float base, const float fwd );
-
-void rope_2d_cpu( torch::Tensor tokens, const torch::Tensor positions, const float base, const float fwd )
-{
-    const int B = tokens.size(0);
-    const int N = tokens.size(1);
-    const int H = tokens.size(2);
-    const int D = tokens.size(3) / 4;
-
-    auto tok = tokens.accessor<float, 4>();
-    auto pos = positions.accessor<int64_t, 3>();
-
-    for (int b = 0; b < B; b++) {
-      for (int x = 0; x < 2; x++) { // y and then x (2d)
-        for (int n = 0; n < N; n++) {
-        
-            // grab the token position
-            const int p = pos[b][n][x];
-
-            for (int h = 0; h < H; h++) {
-                for (int d = 0; d < D; d++) {
-                    // grab the two values
-                    float u = tok[b][n][h][d+0+x*2*D];
-                    float v = tok[b][n][h][d+D+x*2*D];
-
-                    // grab the cos,sin
-                    const float inv_freq = fwd * p / powf(base, d/float(D));
-                    float c = cosf(inv_freq);
-                    float s = sinf(inv_freq);
-
-                    // write the result
-                    tok[b][n][h][d+0+x*2*D] = u*c - v*s;
-                    tok[b][n][h][d+D+x*2*D] = v*c + u*s;
-                }
-            }
-        }
-      }
-    }
-}
-
-void rope_2d( torch::Tensor tokens,     // B,N,H,D
-        const torch::Tensor positions,  // B,N,2
-        const float base, 
-        const float fwd )
-{
-    TORCH_CHECK(tokens.dim() == 4, "tokens must have 4 dimensions");
-    TORCH_CHECK(positions.dim() == 3, "positions must have 3 dimensions");
-    TORCH_CHECK(tokens.size(0) == positions.size(0), "batch size differs between tokens & positions");
-    TORCH_CHECK(tokens.size(1) == positions.size(1), "seq_length differs between tokens & positions");
-    TORCH_CHECK(positions.size(2) == 2, "positions.shape[2] must be equal to 2");
-    TORCH_CHECK(tokens.is_cuda() == positions.is_cuda(), "tokens and positions are not on the same device" );
-
-    if (tokens.is_cuda())
-        rope_2d_cuda( tokens, positions, base, fwd );
-    else
-        rope_2d_cpu( tokens, positions, base, fwd );
-}
-
-PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
-  m.def("rope_2d", &rope_2d, "RoPE 2d forward/backward");
-}

opensora/models/diffusion/utils/curope/curope2d.py

@@ -1,40 +0,0 @@
-# Copyright (C) 2022-present Naver Corporation. All rights reserved.
-# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
-
-import torch
-
-try:
-    import curope as _kernels # run `python setup.py install`
-except ModuleNotFoundError:
-    from . import curope as _kernels # run `python setup.py build_ext --inplace`
-
-
-class cuRoPE2D_func (torch.autograd.Function):
-
-    @staticmethod
-    def forward(ctx, tokens, positions, base, F0=1):
-        ctx.save_for_backward(positions)
-        ctx.saved_base = base
-        ctx.saved_F0 = F0
-        # tokens = tokens.clone() # uncomment this if inplace doesn't work
-        _kernels.rope_2d( tokens, positions, base, F0 )
-        ctx.mark_dirty(tokens)
-        return tokens
-
-    @staticmethod
-    def backward(ctx, grad_res):
-        positions, base, F0 = ctx.saved_tensors[0], ctx.saved_base, ctx.saved_F0
-        _kernels.rope_2d( grad_res, positions, base, -F0 )
-        ctx.mark_dirty(grad_res)
-        return grad_res, None, None, None
-
-
-class cuRoPE2D(torch.nn.Module):
-    def __init__(self, freq=100.0, F0=1.0):
-        super().__init__()
-        self.base = freq 
-        self.F0 = F0
-
-    def forward(self, tokens, positions): 
-        cuRoPE2D_func.apply( tokens.transpose(1,2), positions, self.base, self.F0 )
-        return tokens

opensora/models/diffusion/utils/curope/kernels.cu

@@ -1,108 +0,0 @@
-/* 
-  Copyright (C) 2022-present Naver Corporation. All rights reserved.
-  Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
-*/
-
-#include <torch/extension.h>
-#include <cuda.h>
-#include <cuda_runtime.h>
-#include <vector>
-
-#define CHECK_CUDA(tensor) {\
-    TORCH_CHECK((tensor).is_cuda(), #tensor " is not in cuda memory"); \
-    TORCH_CHECK((tensor).is_contiguous(), #tensor " is not contiguous"); }
-void CHECK_KERNEL() {auto error = cudaGetLastError(); TORCH_CHECK( error == cudaSuccess, cudaGetErrorString(error));}
-
-
-template < typename scalar_t  >
-__global__ void rope_2d_cuda_kernel( 
-        //scalar_t* __restrict__ tokens, 
-        torch::PackedTensorAccessor32<scalar_t,4,torch::RestrictPtrTraits> tokens,
-        const int64_t* __restrict__ pos, 
-        const float base, 
-        const float fwd )
-        // const int N, const int H, const int D )
-{
-    // tokens shape = (B, N, H, D)
-    const int N = tokens.size(1);
-    const int H = tokens.size(2);
-    const int D = tokens.size(3);
-    
-    // each block update a single token, for all heads
-    // each thread takes care of a single output
-    extern __shared__ float shared[];
-    float* shared_inv_freq = shared + D;
-
-    const int b = blockIdx.x / N;
-    const int n = blockIdx.x % N;
-
-    const int Q = D / 4; 
-    // one token = [0..Q : Q..2Q : 2Q..3Q : 3Q..D]
-    //              u_Y     v_Y     u_X      v_X
-
-    // shared memory: first, compute inv_freq
-    if (threadIdx.x < Q)
-        shared_inv_freq[threadIdx.x] = fwd / powf(base, threadIdx.x/float(Q));
-    __syncthreads();
-
-    // start of X or Y part
-    const int X = threadIdx.x < D/2 ? 0 : 1; 
-    const int m = (X*D/2) + (threadIdx.x % Q);   // index of u_Y or u_X
-
-    // grab the cos,sin appropriate for me
-    const float freq = pos[blockIdx.x*2+X] * shared_inv_freq[threadIdx.x % Q];
-    const float cos = cosf(freq);
-    const float sin = sinf(freq);
-    /*
-    float* shared_cos_sin = shared + D + D/4;
-    if ((threadIdx.x % (D/2)) < Q)
-        shared_cos_sin[m+0] = cosf(freq);
-    else
-        shared_cos_sin[m+Q] = sinf(freq);
-    __syncthreads();
-    const float cos = shared_cos_sin[m+0];
-    const float sin = shared_cos_sin[m+Q];
-    */
-
-    for (int h = 0; h < H; h++)
-    {
-        // then, load all the token for this head in shared memory
-        shared[threadIdx.x] = tokens[b][n][h][threadIdx.x];
-        __syncthreads();
-
-        const float u = shared[m];
-        const float v = shared[m+Q];
-        
-        // write output
-        if ((threadIdx.x % (D/2)) < Q)
-            tokens[b][n][h][threadIdx.x] = u*cos - v*sin;
-        else
-            tokens[b][n][h][threadIdx.x] = v*cos + u*sin;
-    }
-}
-
-void rope_2d_cuda( torch::Tensor tokens, const torch::Tensor pos, const float base, const float fwd ) 
-{
-    const int B = tokens.size(0); // batch size
-    const int N = tokens.size(1); // sequence length
-    const int H = tokens.size(2); // number of heads
-    const int D = tokens.size(3); // dimension per head
-
-    TORCH_CHECK(tokens.stride(3) == 1 && tokens.stride(2) == D, "tokens are not contiguous");
-    TORCH_CHECK(pos.is_contiguous(), "positions are not contiguous");
-    TORCH_CHECK(pos.size(0) == B && pos.size(1) == N && pos.size(2) == 2, "bad pos.shape");
-    TORCH_CHECK(D % 4 == 0, "token dim must be multiple of 4");
-
-    // one block for each layer, one thread per local-max
-    const int THREADS_PER_BLOCK = D;
-    const int N_BLOCKS = B * N; // each block takes care of H*D values
-    const int SHARED_MEM = sizeof(float) * (D + D/4);
-
-    AT_DISPATCH_FLOATING_TYPES_AND_HALF(tokens.type(), "rope_2d_cuda", ([&] {
-        rope_2d_cuda_kernel<scalar_t> <<<N_BLOCKS, THREADS_PER_BLOCK, SHARED_MEM>>> (
-            //tokens.data_ptr<scalar_t>(), 
-            tokens.packed_accessor32<scalar_t,4,torch::RestrictPtrTraits>(),
-            pos.data_ptr<int64_t>(), 
-            base, fwd); //, N, H, D );
-    }));
-}

opensora/models/diffusion/utils/curope/setup.py

@@ -1,34 +0,0 @@
-# Copyright (C) 2022-present Naver Corporation. All rights reserved.
-# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
-
-from setuptools import setup
-from torch import cuda
-from torch.utils.cpp_extension import BuildExtension, CUDAExtension
-
-# compile for all possible CUDA architectures
-all_cuda_archs = cuda.get_gencode_flags().replace('compute=','arch=').split()
-# alternatively, you can list cuda archs that you want, eg:
-# all_cuda_archs = [
-    # '-gencode', 'arch=compute_70,code=sm_70',
-    # '-gencode', 'arch=compute_75,code=sm_75',
-    # '-gencode', 'arch=compute_80,code=sm_80',
-    # '-gencode', 'arch=compute_86,code=sm_86'
-# ]
-
-setup(
-    name = 'curope',
-    ext_modules = [
-        CUDAExtension(
-                name='curope',
-                sources=[
-                    "curope.cpp",
-                    "kernels.cu",
-                ],
-                extra_compile_args = dict(
-                    nvcc=['-O3','--ptxas-options=-v',"--use_fast_math"]+all_cuda_archs, 
-                    cxx=['-O3'])
-                )
-    ],
-    cmdclass = {
-        'build_ext': BuildExtension
-    })

opensora/models/diffusion/utils/pos_embed.py

@@ -1,243 +0,0 @@
-# Copyright (C) 2022-present Naver Corporation. All rights reserved.
-# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
-
-# croco: https://github.com/naver/croco
-# diffusers: https://github.com/huggingface/diffusers
-# --------------------------------------------------------
-# Position embedding utils
-# --------------------------------------------------------
-
-import numpy as np
-import torch
-
-
-def get_2d_sincos_pos_embed(
-        embed_dim, grid_size, cls_token=False, extra_tokens=0, interpolation_scale=1.0, base_size=16
-):
-    """
-    grid_size: int of the grid height and width return: pos_embed: [grid_size*grid_size, embed_dim] or
-    [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
-    """
-    if isinstance(grid_size, int):
-        grid_size = (grid_size, grid_size)
-
-    grid_h = np.arange(grid_size[0], dtype=np.float32) / (grid_size[0] / base_size) / interpolation_scale
-    grid_w = np.arange(grid_size[1], dtype=np.float32) / (grid_size[1] / base_size) / interpolation_scale
-    grid = np.meshgrid(grid_w, grid_h)  # here w goes first
-    grid = np.stack(grid, axis=0)
-
-    grid = grid.reshape([2, 1, grid_size[1], grid_size[0]])
-    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
-    if cls_token and extra_tokens > 0:
-        pos_embed = np.concatenate([np.zeros([extra_tokens, embed_dim]), pos_embed], axis=0)
-    return pos_embed
-
-
-def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
-    if embed_dim % 2 != 0:
-        raise ValueError("embed_dim must be divisible by 2")
-
-    # use half of dimensions to encode grid_h
-    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])  # (H*W, D/2)
-    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])  # (H*W, D/2)
-
-    emb = np.concatenate([emb_h, emb_w], axis=1)  # (H*W, D)
-    return emb
-
-
-def get_1d_sincos_pos_embed(
-        embed_dim, length, interpolation_scale=1.0, base_size=16
-):
-    pos = torch.arange(0, length).unsqueeze(1) / interpolation_scale
-    pos_embed = get_1d_sincos_pos_embed_from_grid(embed_dim, pos)
-    return pos_embed
-
-def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
-    """
-    embed_dim: output dimension for each position pos: a list of positions to be encoded: size (M,) out: (M, D)
-    """
-    if embed_dim % 2 != 0:
-        raise ValueError("embed_dim must be divisible by 2")
-
-    omega = np.arange(embed_dim // 2, dtype=np.float64)
-    omega /= embed_dim / 2.0
-    omega = 1.0 / 10000 ** omega  # (D/2,)
-
-    pos = pos.reshape(-1)  # (M,)
-    out = np.einsum("m,d->md", pos, omega)  # (M, D/2), outer product
-
-    emb_sin = np.sin(out)  # (M, D/2)
-    emb_cos = np.cos(out)  # (M, D/2)
-
-    emb = np.concatenate([emb_sin, emb_cos], axis=1)  # (M, D)
-    return emb
-
-
-# ----------------------------------------------------------
-# RoPE2D: RoPE implementation in 2D
-# ----------------------------------------------------------
-
-try:
-    from .curope import cuRoPE2D
-
-    RoPE2D = cuRoPE2D
-except ImportError:
-    print('Warning, cannot find cuda-compiled version of RoPE2D, using a slow pytorch version instead')
-
-
-    class RoPE2D(torch.nn.Module):
-
-        def __init__(self, freq=10000.0, F0=1.0, scaling_factor=1.0):
-            super().__init__()
-            self.base = freq
-            self.F0 = F0
-            self.scaling_factor = scaling_factor
-            self.cache = {}
-
-        def get_cos_sin(self, D, seq_len, device, dtype):
-            if (D, seq_len, device, dtype) not in self.cache:
-                inv_freq = 1.0 / (self.base ** (torch.arange(0, D, 2).float().to(device) / D))
-                t = torch.arange(seq_len, device=device, dtype=inv_freq.dtype)
-                freqs = torch.einsum("i,j->ij", t, inv_freq).to(dtype)
-                freqs = torch.cat((freqs, freqs), dim=-1)
-                cos = freqs.cos()  # (Seq, Dim)
-                sin = freqs.sin()
-                self.cache[D, seq_len, device, dtype] = (cos, sin)
-            return self.cache[D, seq_len, device, dtype]
-
-        @staticmethod
-        def rotate_half(x):
-            x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2:]
-            return torch.cat((-x2, x1), dim=-1)
-
-        def apply_rope1d(self, tokens, pos1d, cos, sin):
-            assert pos1d.ndim == 2
-            
-            cos = torch.nn.functional.embedding(pos1d, cos)[:, None, :, :]
-            sin = torch.nn.functional.embedding(pos1d, sin)[:, None, :, :]
-            return (tokens * cos) + (self.rotate_half(tokens) * sin)
-
-        def forward(self, tokens, positions):
-            """
-            input:
-                * tokens: batch_size x nheads x ntokens x dim
-                * positions: batch_size x ntokens x 2 (y and x position of each token)
-            output:
-                * tokens after appplying RoPE2D (batch_size x nheads x ntokens x dim)
-            """
-            assert tokens.size(3) % 2 == 0, "number of dimensions should be a multiple of two"
-            D = tokens.size(3) // 2
-            assert positions.ndim == 3 and positions.shape[-1] == 2  # Batch, Seq, 2
-            cos, sin = self.get_cos_sin(D, int(positions.max()) + 1, tokens.device, tokens.dtype)
-            # split features into two along the feature dimension, and apply rope1d on each half
-            y, x = tokens.chunk(2, dim=-1)
-            y = self.apply_rope1d(y, positions[:, :, 0], cos, sin)
-            x = self.apply_rope1d(x, positions[:, :, 1], cos, sin)
-            tokens = torch.cat((y, x), dim=-1)
-            return tokens
-
-class LinearScalingRoPE2D(RoPE2D):
-    """Code from https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L148"""
-
-    def forward(self, tokens, positions):
-        # difference to the original RoPE: a scaling factor is aplied to the position ids
-        dtype = positions.dtype
-        positions = positions.float() / self.scaling_factor
-        positions = positions.to(dtype)
-        tokens = super().forward(tokens, positions)
-        return tokens
-
-
-try:
-    from .curope import cuRoPE1D
-
-    RoPE1D = cuRoPE1D
-except ImportError:
-    print('Warning, cannot find cuda-compiled version of RoPE2D, using a slow pytorch version instead')
-
-
-    class RoPE1D(torch.nn.Module):
-
-        def __init__(self, freq=10000.0, F0=1.0, scaling_factor=1.0):
-            super().__init__()
-            self.base = freq
-            self.F0 = F0
-            self.scaling_factor = scaling_factor
-            self.cache = {}
-
-        def get_cos_sin(self, D, seq_len, device, dtype):
-            if (D, seq_len, device, dtype) not in self.cache:
-                inv_freq = 1.0 / (self.base ** (torch.arange(0, D, 2).float().to(device) / D))
-                t = torch.arange(seq_len, device=device, dtype=inv_freq.dtype)
-                freqs = torch.einsum("i,j->ij", t, inv_freq).to(dtype)
-                freqs = torch.cat((freqs, freqs), dim=-1)
-                cos = freqs.cos()  # (Seq, Dim)
-                sin = freqs.sin()
-                self.cache[D, seq_len, device, dtype] = (cos, sin)
-            return self.cache[D, seq_len, device, dtype]
-
-        @staticmethod
-        def rotate_half(x):
-            x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2:]
-            return torch.cat((-x2, x1), dim=-1)
-
-        def apply_rope1d(self, tokens, pos1d, cos, sin):
-            assert pos1d.ndim == 2
-            cos = torch.nn.functional.embedding(pos1d, cos)[:, None, :, :]
-            sin = torch.nn.functional.embedding(pos1d, sin)[:, None, :, :]
-            return (tokens * cos) + (self.rotate_half(tokens) * sin)
-
-        def forward(self, tokens, positions):
-            """
-            input:
-                * tokens: batch_size x nheads x ntokens x dim
-                * positions: batch_size x ntokens (t position of each token)
-            output:
-                * tokens after appplying RoPE2D (batch_size x nheads x ntokens x dim)
-            """
-            D = tokens.size(3)
-            assert positions.ndim == 2  # Batch, Seq
-            cos, sin = self.get_cos_sin(D, int(positions.max()) + 1, tokens.device, tokens.dtype)
-            tokens = self.apply_rope1d(tokens, positions, cos, sin)
-            return tokens
-
-class LinearScalingRoPE1D(RoPE1D):
-    """Code from https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L148"""
-    
-    def forward(self, tokens, positions):
-        # difference to the original RoPE: a scaling factor is aplied to the position ids
-        dtype = positions.dtype
-        positions = positions.float() / self.scaling_factor
-        positions = positions.to(dtype)
-        tokens = super().forward(tokens, positions)
-        return tokens
-
-
-class PositionGetter2D(object):
-    """ return positions of patches """
-
-    def __init__(self):
-        self.cache_positions = {}
-        
-    def __call__(self, b, h, w, device):
-        if not (h,w) in self.cache_positions:
-            x = torch.arange(w, device=device)
-            y = torch.arange(h, device=device)
-            self.cache_positions[h,w] = torch.cartesian_prod(y, x) # (h, w, 2)
-        pos = self.cache_positions[h,w].view(1, h*w, 2).expand(b, -1, 2).clone()
-        return pos
-
-
-
-class PositionGetter1D(object):
-    """ return positions of patches """
-
-    def __init__(self):
-        self.cache_positions = {}
-        
-    def __call__(self, b, l, device):
-        if not (l) in self.cache_positions:
-            x = torch.arange(l, device=device)
-            self.cache_positions[l] = x # (l, )
-        pos = self.cache_positions[l].view(1, l).expand(b, -1).clone()
-        return pos

opensora/models/diffusion/utils/transport/__init__.py

@@ -1,63 +0,0 @@
-from .transport import Transport, ModelType, WeightType, PathType, Sampler
-
-def create_transport(
-    path_type='Linear',
-    prediction="velocity",
-    loss_weight=None,
-    train_eps=None,
-    sample_eps=None,
-):
-    """function for creating Transport object
-    **Note**: model prediction defaults to velocity
-    Args:
-    - path_type: type of path to use; default to linear
-    - learn_score: set model prediction to score
-    - learn_noise: set model prediction to noise
-    - velocity_weighted: weight loss by velocity weight
-    - likelihood_weighted: weight loss by likelihood weight
-    - train_eps: small epsilon for avoiding instability during training
-    - sample_eps: small epsilon for avoiding instability during sampling
-    """
-
-    if prediction == "noise":
-        model_type = ModelType.NOISE
-    elif prediction == "score":
-        model_type = ModelType.SCORE
-    else:
-        model_type = ModelType.VELOCITY
-
-    if loss_weight == "velocity":
-        loss_type = WeightType.VELOCITY
-    elif loss_weight == "likelihood":
-        loss_type = WeightType.LIKELIHOOD
-    else:
-        loss_type = WeightType.NONE
-
-    path_choice = {
-        "Linear": PathType.LINEAR,
-        "GVP": PathType.GVP,
-        "VP": PathType.VP,
-    }
-
-    path_type = path_choice[path_type]
-
-    if (path_type in [PathType.VP]):
-        train_eps = 1e-5 if train_eps is None else train_eps
-        sample_eps = 1e-3 if train_eps is None else sample_eps
-    elif (path_type in [PathType.GVP, PathType.LINEAR] and model_type != ModelType.VELOCITY):
-        train_eps = 1e-3 if train_eps is None else train_eps
-        sample_eps = 1e-3 if train_eps is None else sample_eps
-    else: # velocity & [GVP, LINEAR] is stable everywhere
-        train_eps = 0
-        sample_eps = 0
-    
-    # create flow state
-    state = Transport(
-        model_type=model_type,
-        path_type=path_type,
-        loss_type=loss_type,
-        train_eps=train_eps,
-        sample_eps=sample_eps,
-    )
-    
-    return state

opensora/models/diffusion/utils/transport/integrators.py

@@ -1,117 +0,0 @@
-import numpy as np
-import torch as th
-import torch.nn as nn
-from torchdiffeq import odeint
-from functools import partial
-from tqdm import tqdm
-
-class sde:
-    """SDE solver class"""
-    def __init__(
-        self, 
-        drift,
-        diffusion,
-        *,
-        t0,
-        t1,
-        num_steps,
-        sampler_type,
-    ):
-        assert t0 < t1, "SDE sampler has to be in forward time"
-
-        self.num_timesteps = num_steps
-        self.t = th.linspace(t0, t1, num_steps)
-        self.dt = self.t[1] - self.t[0]
-        self.drift = drift
-        self.diffusion = diffusion
-        self.sampler_type = sampler_type
-
-    def __Euler_Maruyama_step(self, x, mean_x, t, model, **model_kwargs):
-        w_cur = th.randn(x.size()).to(x)
-        t = th.ones(x.size(0)).to(x) * t
-        dw = w_cur * th.sqrt(self.dt)
-        drift = self.drift(x, t, model, **model_kwargs)
-        diffusion = self.diffusion(x, t)
-        mean_x = x + drift * self.dt
-        x = mean_x + th.sqrt(2 * diffusion) * dw
-        return x, mean_x
-    
-    def __Heun_step(self, x, _, t, model, **model_kwargs):
-        w_cur = th.randn(x.size()).to(x)
-        dw = w_cur * th.sqrt(self.dt)
-        t_cur = th.ones(x.size(0)).to(x) * t
-        diffusion = self.diffusion(x, t_cur)
-        xhat = x + th.sqrt(2 * diffusion) * dw
-        K1 = self.drift(xhat, t_cur, model, **model_kwargs)
-        xp = xhat + self.dt * K1
-        K2 = self.drift(xp, t_cur + self.dt, model, **model_kwargs)
-        return xhat + 0.5 * self.dt * (K1 + K2), xhat # at last time point we do not perform the heun step
-
-    def __forward_fn(self):
-        """TODO: generalize here by adding all private functions ending with steps to it"""
-        sampler_dict = {
-            "Euler": self.__Euler_Maruyama_step,
-            "Heun": self.__Heun_step,
-        }
-
-        try:
-            sampler = sampler_dict[self.sampler_type]
-        except:
-            raise NotImplementedError("Smapler type not implemented.")
-    
-        return sampler
-
-    def sample(self, init, model, **model_kwargs):
-        """forward loop of sde"""
-        x = init
-        mean_x = init 
-        samples = []
-        sampler = self.__forward_fn()
-        for ti in self.t[:-1]:
-            with th.no_grad():
-                x, mean_x = sampler(x, mean_x, ti, model, **model_kwargs)
-                samples.append(x)
-
-        return samples
-
-class ode:
-    """ODE solver class"""
-    def __init__(
-        self,
-        drift,
-        *,
-        t0,
-        t1,
-        sampler_type,
-        num_steps,
-        atol,
-        rtol,
-    ):
-        assert t0 < t1, "ODE sampler has to be in forward time"
-
-        self.drift = drift
-        self.t = th.linspace(t0, t1, num_steps)
-        self.atol = atol
-        self.rtol = rtol
-        self.sampler_type = sampler_type
-
-    def sample(self, x, model, **model_kwargs):
-        
-        device = x[0].device if isinstance(x, tuple) else x.device
-        def _fn(t, x):
-            t = th.ones(x[0].size(0)).to(device) * t if isinstance(x, tuple) else th.ones(x.size(0)).to(device) * t
-            model_output = self.drift(x, t, model, **model_kwargs)
-            return model_output
-
-        t = self.t.to(device)
-        atol = [self.atol] * len(x) if isinstance(x, tuple) else [self.atol]
-        rtol = [self.rtol] * len(x) if isinstance(x, tuple) else [self.rtol]
-        samples = odeint(
-            _fn,
-            x,
-            t,
-            method=self.sampler_type,
-            atol=atol,
-            rtol=rtol
-        )
-        return samples

opensora/models/diffusion/utils/transport/path.py

@@ -1,192 +0,0 @@
-import torch as th
-import numpy as np
-from functools import partial
-
-def expand_t_like_x(t, x):
-    """Function to reshape time t to broadcastable dimension of x
-    Args:
-      t: [batch_dim,], time vector
-      x: [batch_dim,...], data point
-    """
-    dims = [1] * (len(x.size()) - 1)
-    t = t.view(t.size(0), *dims)
-    return t
-
-
-#################### Coupling Plans ####################
-
-class ICPlan:
-    """Linear Coupling Plan"""
-    def __init__(self, sigma=0.0):
-        self.sigma = sigma
-
-    def compute_alpha_t(self, t):
-        """Compute the data coefficient along the path"""
-        return t, 1
-    
-    def compute_sigma_t(self, t):
-        """Compute the noise coefficient along the path"""
-        return 1 - t, -1
-    
-    def compute_d_alpha_alpha_ratio_t(self, t):
-        """Compute the ratio between d_alpha and alpha"""
-        return 1 / t
-
-    def compute_drift(self, x, t):
-        """We always output sde according to score parametrization; """
-        t = expand_t_like_x(t, x)
-        alpha_ratio = self.compute_d_alpha_alpha_ratio_t(t)
-        sigma_t, d_sigma_t = self.compute_sigma_t(t)
-        drift = alpha_ratio * x
-        diffusion = alpha_ratio * (sigma_t ** 2) - sigma_t * d_sigma_t
-
-        return -drift, diffusion
-
-    def compute_diffusion(self, x, t, form="constant", norm=1.0):
-        """Compute the diffusion term of the SDE
-        Args:
-          x: [batch_dim, ...], data point
-          t: [batch_dim,], time vector
-          form: str, form of the diffusion term
-          norm: float, norm of the diffusion term
-        """
-        t = expand_t_like_x(t, x)
-        choices = {
-            "constant": norm,
-            "SBDM": norm * self.compute_drift(x, t)[1],
-            "sigma": norm * self.compute_sigma_t(t)[0],
-            "linear": norm * (1 - t),
-            "decreasing": 0.25 * (norm * th.cos(np.pi * t) + 1) ** 2,
-            "inccreasing-decreasing": norm * th.sin(np.pi * t) ** 2,
-        }
-
-        try:
-            diffusion = choices[form]
-        except KeyError:
-            raise NotImplementedError(f"Diffusion form {form} not implemented")
-        
-        return diffusion
-
-    def get_score_from_velocity(self, velocity, x, t):
-        """Wrapper function: transfrom velocity prediction model to score
-        Args:
-            velocity: [batch_dim, ...] shaped tensor; velocity model output
-            x: [batch_dim, ...] shaped tensor; x_t data point
-            t: [batch_dim,] time tensor
-        """
-        t = expand_t_like_x(t, x)
-        alpha_t, d_alpha_t = self.compute_alpha_t(t)
-        sigma_t, d_sigma_t = self.compute_sigma_t(t)
-        mean = x
-        reverse_alpha_ratio = alpha_t / d_alpha_t
-        var = sigma_t**2 - reverse_alpha_ratio * d_sigma_t * sigma_t
-        score = (reverse_alpha_ratio * velocity - mean) / var
-        return score
-    
-    def get_noise_from_velocity(self, velocity, x, t):
-        """Wrapper function: transfrom velocity prediction model to denoiser
-        Args:
-            velocity: [batch_dim, ...] shaped tensor; velocity model output
-            x: [batch_dim, ...] shaped tensor; x_t data point
-            t: [batch_dim,] time tensor
-        """
-        t = expand_t_like_x(t, x)
-        alpha_t, d_alpha_t = self.compute_alpha_t(t)
-        sigma_t, d_sigma_t = self.compute_sigma_t(t)
-        mean = x
-        reverse_alpha_ratio = alpha_t / d_alpha_t
-        var = reverse_alpha_ratio * d_sigma_t - sigma_t
-        noise = (reverse_alpha_ratio * velocity - mean) / var
-        return noise
-
-    def get_velocity_from_score(self, score, x, t):
-        """Wrapper function: transfrom score prediction model to velocity
-        Args:
-            score: [batch_dim, ...] shaped tensor; score model output
-            x: [batch_dim, ...] shaped tensor; x_t data point
-            t: [batch_dim,] time tensor
-        """
-        t = expand_t_like_x(t, x)
-        drift, var = self.compute_drift(x, t)
-        velocity = var * score - drift
-        return velocity
-
-    def compute_mu_t(self, t, x0, x1):
-        """Compute the mean of time-dependent density p_t"""
-        t = expand_t_like_x(t, x1)
-        alpha_t, _ = self.compute_alpha_t(t)
-        sigma_t, _ = self.compute_sigma_t(t)
-        return alpha_t * x1 + sigma_t * x0
-    
-    def compute_xt(self, t, x0, x1):
-        """Sample xt from time-dependent density p_t; rng is required"""
-        xt = self.compute_mu_t(t, x0, x1)
-        return xt
-    
-    def compute_ut(self, t, x0, x1, xt):
-        """Compute the vector field corresponding to p_t"""
-        t = expand_t_like_x(t, x1)
-        _, d_alpha_t = self.compute_alpha_t(t)
-        _, d_sigma_t = self.compute_sigma_t(t)
-        return d_alpha_t * x1 + d_sigma_t * x0
-    
-    def plan(self, t, x0, x1):
-        xt = self.compute_xt(t, x0, x1)
-        ut = self.compute_ut(t, x0, x1, xt)
-        return t, xt, ut
-    
-
-class VPCPlan(ICPlan):
-    """class for VP path flow matching"""
-
-    def __init__(self, sigma_min=0.1, sigma_max=20.0):
-        self.sigma_min = sigma_min
-        self.sigma_max = sigma_max
-        self.log_mean_coeff = lambda t: -0.25 * ((1 - t) ** 2) * (self.sigma_max - self.sigma_min) - 0.5 * (1 - t) * self.sigma_min 
-        self.d_log_mean_coeff = lambda t: 0.5 * (1 - t) * (self.sigma_max - self.sigma_min) + 0.5 * self.sigma_min
-
-
-    def compute_alpha_t(self, t):
-        """Compute coefficient of x1"""
-        alpha_t = self.log_mean_coeff(t)
-        alpha_t = th.exp(alpha_t)
-        d_alpha_t = alpha_t * self.d_log_mean_coeff(t)
-        return alpha_t, d_alpha_t
-    
-    def compute_sigma_t(self, t):
-        """Compute coefficient of x0"""
-        p_sigma_t = 2 * self.log_mean_coeff(t)
-        sigma_t = th.sqrt(1 - th.exp(p_sigma_t))
-        d_sigma_t = th.exp(p_sigma_t) * (2 * self.d_log_mean_coeff(t)) / (-2 * sigma_t)
-        return sigma_t, d_sigma_t
-    
-    def compute_d_alpha_alpha_ratio_t(self, t):
-        """Special purposed function for computing numerical stabled d_alpha_t / alpha_t"""
-        return self.d_log_mean_coeff(t)
-
-    def compute_drift(self, x, t):
-        """Compute the drift term of the SDE"""
-        t = expand_t_like_x(t, x)
-        beta_t = self.sigma_min + (1 - t) * (self.sigma_max - self.sigma_min)
-        return -0.5 * beta_t * x, beta_t / 2
-    
-
-class GVPCPlan(ICPlan):
-    def __init__(self, sigma=0.0):
-        super().__init__(sigma)
-    
-    def compute_alpha_t(self, t):
-        """Compute coefficient of x1"""
-        alpha_t = th.sin(t * np.pi / 2)
-        d_alpha_t = np.pi / 2 * th.cos(t * np.pi / 2)
-        return alpha_t, d_alpha_t
-    
-    def compute_sigma_t(self, t):
-        """Compute coefficient of x0"""
-        sigma_t = th.cos(t * np.pi / 2)
-        d_sigma_t = -np.pi / 2 * th.sin(t * np.pi / 2)
-        return sigma_t, d_sigma_t
-    
-    def compute_d_alpha_alpha_ratio_t(self, t):
-        """Special purposed function for computing numerical stabled d_alpha_t / alpha_t"""
-        return np.pi / (2 * th.tan(t * np.pi / 2))

opensora/models/diffusion/utils/transport/transport.py

@@ -1,443 +0,0 @@
-import torch as th
-import numpy as np
-import logging
-
-import enum
-
-from . import path
-from .utils import EasyDict, log_state, mean_flat
-from .integrators import ode, sde
-
-class ModelType(enum.Enum):
-    """
-    Which type of output the model predicts.
-    """
-
-    NOISE = enum.auto()  # the model predicts epsilon
-    SCORE = enum.auto()  # the model predicts \nabla \log p(x)
-    VELOCITY = enum.auto()  # the model predicts v(x)
-
-class PathType(enum.Enum):
-    """
-    Which type of path to use.
-    """
-
-    LINEAR = enum.auto()
-    GVP = enum.auto()
-    VP = enum.auto()
-
-class WeightType(enum.Enum):
-    """
-    Which type of weighting to use.
-    """
-
-    NONE = enum.auto()
-    VELOCITY = enum.auto()
-    LIKELIHOOD = enum.auto()
-
-
-class Transport:
-
-    def __init__(
-        self,
-        *,
-        model_type,
-        path_type,
-        loss_type,
-        train_eps,
-        sample_eps,
-    ):
-        path_options = {
-            PathType.LINEAR: path.ICPlan,
-            PathType.GVP: path.GVPCPlan,
-            PathType.VP: path.VPCPlan,
-        }
-
-        self.loss_type = loss_type
-        self.model_type = model_type
-        self.path_sampler = path_options[path_type]()
-        self.train_eps = train_eps
-        self.sample_eps = sample_eps
-
-    def prior_logp(self, z):
-        '''
-            Standard multivariate normal prior
-            Assume z is batched
-        '''
-        shape = th.tensor(z.size())
-        N = th.prod(shape[1:])
-        _fn = lambda x: -N / 2. * np.log(2 * np.pi) - th.sum(x ** 2) / 2.
-        return th.vmap(_fn)(z)
-    
-
-    def check_interval(
-        self, 
-        train_eps, 
-        sample_eps, 
-        *, 
-        diffusion_form="SBDM",
-        sde=False, 
-        reverse=False, 
-        eval=False,
-        last_step_size=0.0,
-    ):
-        t0 = 0
-        t1 = 1
-        eps = train_eps if not eval else sample_eps
-        if (type(self.path_sampler) in [path.VPCPlan]):
-
-            t1 = 1 - eps if (not sde or last_step_size == 0) else 1 - last_step_size
-
-        elif (type(self.path_sampler) in [path.ICPlan, path.GVPCPlan]) \
-            and (self.model_type != ModelType.VELOCITY or sde): # avoid numerical issue by taking a first semi-implicit step
-
-            t0 = eps if (diffusion_form == "SBDM" and sde) or self.model_type != ModelType.VELOCITY else 0
-            t1 = 1 - eps if (not sde or last_step_size == 0) else 1 - last_step_size
-        
-        if reverse:
-            t0, t1 = 1 - t0, 1 - t1
-
-        return t0, t1
-
-
-    def sample(self, x1):
-        """Sampling x0 & t based on shape of x1 (if needed)
-          Args:
-            x1 - data point; [batch, *dim]
-        """
-        
-        x0 = th.randn_like(x1)
-        t0, t1 = self.check_interval(self.train_eps, self.sample_eps)
-        t = th.rand((x1.shape[0],)) * (t1 - t0) + t0
-        t = t.to(x1)
-        return t, x0, x1
-    
-
-    def training_losses(
-        self, 
-        model,  
-        x1, 
-        model_kwargs=None
-    ):
-        """Loss for training the score model
-        Args:
-        - model: backbone model; could be score, noise, or velocity
-        - x1: datapoint
-        - model_kwargs: additional arguments for the model
-        """
-        if model_kwargs == None:
-            model_kwargs = {}
-        
-        t, x0, x1 = self.sample(x1)
-        t, xt, ut = self.path_sampler.plan(t, x0, x1)
-        model_output = model(xt, t, **model_kwargs)
-        B, *_, C = xt.shape
-        assert model_output.size() == (B, *xt.size()[1:-1], C)
-
-        terms = {}
-        terms['pred'] = model_output
-        if self.model_type == ModelType.VELOCITY:
-            terms['loss'] = mean_flat(((model_output - ut) ** 2))
-        else: 
-            _, drift_var = self.path_sampler.compute_drift(xt, t)
-            sigma_t, _ = self.path_sampler.compute_sigma_t(path.expand_t_like_x(t, xt))
-            if self.loss_type in [WeightType.VELOCITY]:
-                weight = (drift_var / sigma_t) ** 2
-            elif self.loss_type in [WeightType.LIKELIHOOD]:
-                weight = drift_var / (sigma_t ** 2)
-            elif self.loss_type in [WeightType.NONE]:
-                weight = 1
-            else:
-                raise NotImplementedError()
-            
-            if self.model_type == ModelType.NOISE:
-                terms['loss'] = mean_flat(weight * ((model_output - x0) ** 2))
-            else:
-                terms['loss'] = mean_flat(weight * ((model_output * sigma_t + x0) ** 2))
-                
-        return terms
-    
-
-    def get_drift(
-        self
-    ):
-        """member function for obtaining the drift of the probability flow ODE"""
-        def score_ode(x, t, model, **model_kwargs):
-            drift_mean, drift_var = self.path_sampler.compute_drift(x, t)
-            model_output = model(x, t, **model_kwargs)
-            return (-drift_mean + drift_var * model_output) # by change of variable
-        
-        def noise_ode(x, t, model, **model_kwargs):
-            drift_mean, drift_var = self.path_sampler.compute_drift(x, t)
-            sigma_t, _ = self.path_sampler.compute_sigma_t(path.expand_t_like_x(t, x))
-            model_output = model(x, t, **model_kwargs)
-            score = model_output / -sigma_t
-            return (-drift_mean + drift_var * score)
-        
-        def velocity_ode(x, t, model, **model_kwargs):
-            model_output = model(x, t, **model_kwargs)
-            return model_output
-
-        if self.model_type == ModelType.NOISE:
-            drift_fn = noise_ode
-        elif self.model_type == ModelType.SCORE:
-            drift_fn = score_ode
-        else:
-            drift_fn = velocity_ode
-        
-        def body_fn(x, t, model, **model_kwargs):
-            model_output = drift_fn(x, t, model, **model_kwargs)
-            assert model_output.shape == x.shape, "Output shape from ODE solver must match input shape"
-            return model_output
-
-        return body_fn
-    
-
-    def get_score(
-        self,
-    ):
-        """member function for obtaining score of 
-            x_t = alpha_t * x + sigma_t * eps"""
-        if self.model_type == ModelType.NOISE:
-            score_fn = lambda x, t, model, **kwargs: model(x, t, **kwargs) / -self.path_sampler.compute_sigma_t(path.expand_t_like_x(t, x))[0]
-        elif self.model_type == ModelType.SCORE:
-            score_fn = lambda x, t, model, **kwagrs: model(x, t, **kwagrs)
-        elif self.model_type == ModelType.VELOCITY:
-            score_fn = lambda x, t, model, **kwargs: self.path_sampler.get_score_from_velocity(model(x, t, **kwargs), x, t)
-        else:
-            raise NotImplementedError()
-        
-        return score_fn
-
-
-class Sampler:
-    """Sampler class for the transport model"""
-    def __init__(
-        self,
-        transport,
-    ):
-        """Constructor for a general sampler; supporting different sampling methods
-        Args:
-        - transport: an tranport object specify model prediction & interpolant type
-        """
-        
-        self.transport = transport
-        self.drift = self.transport.get_drift()
-        self.score = self.transport.get_score()
-    
-    def __get_sde_diffusion_and_drift(
-        self,
-        *,
-        diffusion_form="SBDM",
-        diffusion_norm=1.0,
-    ):
-
-        def diffusion_fn(x, t):
-            diffusion = self.transport.path_sampler.compute_diffusion(x, t, form=diffusion_form, norm=diffusion_norm)
-            return diffusion
-        
-        sde_drift = \
-            lambda x, t, model, **kwargs: \
-                self.drift(x, t, model, **kwargs) + diffusion_fn(x, t) * self.score(x, t, model, **kwargs)
-    
-        sde_diffusion = diffusion_fn
-
-        return sde_drift, sde_diffusion
-    
-    def __get_last_step(
-        self,
-        sde_drift,
-        *,
-        last_step,
-        last_step_size,
-    ):
-        """Get the last step function of the SDE solver"""
-    
-        if last_step is None:
-            last_step_fn = \
-                lambda x, t, model, **model_kwargs: \
-                    x
-        elif last_step == "Mean":
-            last_step_fn = \
-                lambda x, t, model, **model_kwargs: \
-                    x + sde_drift(x, t, model, **model_kwargs) * last_step_size
-        elif last_step == "Tweedie":
-            alpha = self.transport.path_sampler.compute_alpha_t # simple aliasing; the original name was too long
-            sigma = self.transport.path_sampler.compute_sigma_t
-            last_step_fn = \
-                lambda x, t, model, **model_kwargs: \
-                    x / alpha(t)[0][0] + (sigma(t)[0][0] ** 2) / alpha(t)[0][0] * self.score(x, t, model, **model_kwargs)
-        elif last_step == "Euler":
-            last_step_fn = \
-                lambda x, t, model, **model_kwargs: \
-                    x + self.drift(x, t, model, **model_kwargs) * last_step_size
-        else:
-            raise NotImplementedError()
-
-        return last_step_fn
-
-    def sample_sde(
-        self,
-        *,
-        sampling_method="Euler",
-        diffusion_form="SBDM",
-        diffusion_norm=1.0,
-        last_step="Mean",
-        last_step_size=0.04,
-        num_steps=250,
-    ):
-        """returns a sampling function with given SDE settings
-        Args:
-        - sampling_method: type of sampler used in solving the SDE; default to be Euler-Maruyama
-        - diffusion_form: function form of diffusion coefficient; default to be matching SBDM
-        - diffusion_norm: function magnitude of diffusion coefficient; default to 1
-        - last_step: type of the last step; default to identity
-        - last_step_size: size of the last step; default to match the stride of 250 steps over [0,1]
-        - num_steps: total integration step of SDE
-        """
-
-        if last_step is None:
-            last_step_size = 0.0
-
-        sde_drift, sde_diffusion = self.__get_sde_diffusion_and_drift(
-            diffusion_form=diffusion_form,
-            diffusion_norm=diffusion_norm,
-        )
-
-        t0, t1 = self.transport.check_interval(
-            self.transport.train_eps,
-            self.transport.sample_eps,
-            diffusion_form=diffusion_form,
-            sde=True,
-            eval=True,
-            reverse=False,
-            last_step_size=last_step_size,
-        )
-
-        _sde = sde(
-            sde_drift,
-            sde_diffusion,
-            t0=t0,
-            t1=t1,
-            num_steps=num_steps,
-            sampler_type=sampling_method
-        )
-
-        last_step_fn = self.__get_last_step(sde_drift, last_step=last_step, last_step_size=last_step_size)
-            
-
-        def _sample(init, model, **model_kwargs):
-            xs = _sde.sample(init, model, **model_kwargs)
-            ts = th.ones(init.size(0), device=init.device) * t1
-            x = last_step_fn(xs[-1], ts, model, **model_kwargs)
-            xs.append(x)
-
-            assert len(xs) == num_steps, "Samples does not match the number of steps"
-
-            return xs
-
-        return _sample
-    
-    def sample_ode(
-        self,
-        *,
-        sampling_method="dopri5",
-        num_steps=50,
-        atol=1e-6,
-        rtol=1e-3,
-        reverse=False,
-    ):
-        """returns a sampling function with given ODE settings
-        Args:
-        - sampling_method: type of sampler used in solving the ODE; default to be Dopri5
-        - num_steps: 
-            - fixed solver (Euler, Heun): the actual number of integration steps performed
-            - adaptive solver (Dopri5): the number of datapoints saved during integration; produced by interpolation
-        - atol: absolute error tolerance for the solver
-        - rtol: relative error tolerance for the solver
-        - reverse: whether solving the ODE in reverse (data to noise); default to False
-        """
-        if reverse:
-            drift = lambda x, t, model, **kwargs: self.drift(x, th.ones_like(t) * (1 - t), model, **kwargs)
-        else:
-            drift = self.drift
-
-        t0, t1 = self.transport.check_interval(
-            self.transport.train_eps,
-            self.transport.sample_eps,
-            sde=False,
-            eval=True,
-            reverse=reverse,
-            last_step_size=0.0,
-        )
-
-        _ode = ode(
-            drift=drift,
-            t0=t0,
-            t1=t1,
-            sampler_type=sampling_method,
-            num_steps=num_steps,
-            atol=atol,
-            rtol=rtol,
-        )
-        
-        return _ode.sample
-
-    def sample_ode_likelihood(
-        self,
-        *,
-        sampling_method="dopri5",
-        num_steps=50,
-        atol=1e-6,
-        rtol=1e-3,
-    ):
-        
-        """returns a sampling function for calculating likelihood with given ODE settings
-        Args:
-        - sampling_method: type of sampler used in solving the ODE; default to be Dopri5
-        - num_steps: 
-            - fixed solver (Euler, Heun): the actual number of integration steps performed
-            - adaptive solver (Dopri5): the number of datapoints saved during integration; produced by interpolation
-        - atol: absolute error tolerance for the solver
-        - rtol: relative error tolerance for the solver
-        """
-        def _likelihood_drift(x, t, model, **model_kwargs):
-            x, _ = x
-            eps = th.randint(2, x.size(), dtype=th.float, device=x.device) * 2 - 1
-            t = th.ones_like(t) * (1 - t)
-            with th.enable_grad():
-                x.requires_grad = True
-                grad = th.autograd.grad(th.sum(self.drift(x, t, model, **model_kwargs) * eps), x)[0]
-                logp_grad = th.sum(grad * eps, dim=tuple(range(1, len(x.size()))))
-                drift = self.drift(x, t, model, **model_kwargs)
-            return (-drift, logp_grad)
-        
-        t0, t1 = self.transport.check_interval(
-            self.transport.train_eps,
-            self.transport.sample_eps,
-            sde=False,
-            eval=True,
-            reverse=False,
-            last_step_size=0.0,
-        )
-
-        _ode = ode(
-            drift=_likelihood_drift,
-            t0=t0,
-            t1=t1,
-            sampler_type=sampling_method,
-            num_steps=num_steps,
-            atol=atol,
-            rtol=rtol,
-        )
-
-        def _sample_fn(x, model, **model_kwargs):
-            init_logp = th.zeros(x.size(0)).to(x)
-            input = (x, init_logp)
-            drift, delta_logp = _ode.sample(input, model, **model_kwargs)
-            drift, delta_logp = drift[-1], delta_logp[-1]
-            prior_logp = self.transport.prior_logp(drift)
-            logp = prior_logp - delta_logp
-            return logp, drift
-
-        return _sample_fn

opensora/models/diffusion/utils/transport/utils.py

@@ -1,29 +0,0 @@
-import torch as th
-
-class EasyDict:
-
-    def __init__(self, sub_dict):
-        for k, v in sub_dict.items():
-            setattr(self, k, v)
-
-    def __getitem__(self, key):
-        return getattr(self, key)
-
-def mean_flat(x):
-    """
-    Take the mean over all non-batch dimensions.
-    """
-    return th.mean(x, dim=list(range(1, len(x.size()))))
-
-def log_state(state):
-    result = []
-    
-    sorted_state = dict(sorted(state.items()))
-    for key, value in sorted_state.items():
-        # Check if the value is an instance of a class
-        if "<object" in str(value) or "object at" in str(value):
-            result.append(f"{key}: [{value.__class__.__name__}]")
-        else:
-            result.append(f"{key}: {value}")
-    
-    return '\n'.join(result)

opensora/models/text_encoder/__init__.py

@@ -1,76 +0,0 @@
-import torch
-from torch import nn
-from transformers import CLIPModel, CLIPProcessor
-
-from opensora.utils.utils import get_precision
-
-
-class T5Wrapper(nn.Module):
-    def __init__(self, args, **kwargs):
-        super(T5Wrapper, self).__init__()
-        self.model_name = args.text_encoder_name
-        if 'mt5' in self.model_name:
-            from transformers import MT5EncoderModel
-            self.text_enc = MT5EncoderModel.from_pretrained(self.model_name, cache_dir=args.cache_dir, **kwargs).eval()
-            
-        elif 't5' in self.model_name:
-            from transformers import T5EncoderModel
-            self.text_enc = T5EncoderModel.from_pretrained(self.model_name, cache_dir=args.cache_dir, **kwargs).eval()
-            # self.text_enc = T5EncoderModel.from_pretrained("/storage/ongoing/new/Open-Sora-Plan/cache_dir/models--DeepFloyd--t5-v1_1-xxl/snapshots/c9c625d2ec93667ec579ede125fd3811d1f81d37", cache_dir=args.cache_dir, **kwargs).eval()
-
-    def forward(self, input_ids, attention_mask):
-        text_encoder_embs = self.text_enc(input_ids=input_ids, attention_mask=attention_mask)['last_hidden_state']
-        return text_encoder_embs.detach()
-
-class CLIPWrapper(nn.Module):
-    def __init__(self, args):
-        super(CLIPWrapper, self).__init__()
-        self.model_name = args.text_encoder_name
-        dtype = get_precision(args)
-        model_kwargs = {'cache_dir': args.cache_dir, 'low_cpu_mem_usage': True, 'torch_dtype': dtype}
-        self.text_enc = CLIPModel.from_pretrained(self.model_name, **model_kwargs).eval()
-
-    def forward(self, prompts): # prompts = ['I am a test caption', 'Test twice']
-        # Process the prompts and move to the correct device
-        inputs = self.clip_processor(text=prompts, return_tensors="pt", padding=True, truncation=True).to(self.device)
-        
-        # Generate text features
-        with torch.no_grad():
-            text_features = self.clip_model.get_text_features(**inputs)
-        
-        return text_features.detach()
-
-
-
-text_encoder = {
-    'google/mt5-base': T5Wrapper,
-    'google/mt5-xl': T5Wrapper,
-    'google/mt5-xxl': T5Wrapper,
-    'google/umt5-xl': T5Wrapper,
-    'google/umt5-xxl': T5Wrapper,
-    'DeepFloyd/t5-v1_1-xxl': T5Wrapper,
-    'clip': CLIPWrapper
-}
-
-
-def get_text_enc(args):
-    """deprecation"""
-    encoder_key = None
-    for key in text_encoder.keys():
-        if key in args.text_encoder_name:
-            encoder_key = key
-            break
-    text_enc = text_encoder.get(encoder_key, None)
-    assert text_enc is not None
-    return text_enc(args)
-
-def get_text_warpper(text_encoder_name):
-    """deprecation"""
-    encoder_key = None
-    for key in text_encoder.keys():
-        if key in text_encoder_name:
-            encoder_key = key
-            break
-    text_enc = text_encoder.get(encoder_key, None)
-    assert text_enc is not None
-    return text_enc

opensora/models/text_encoder/clip.py

@@ -1,222 +0,0 @@
-# -*- coding: utf-8 -*-
-import os
-import re
-import ftfy
-import torch
-import html
-from PIL import Image
-from transformers import CLIPProcessor, CLIPModel, CLIPTokenizer, CLIPTextModel
-
-class CLIPEmbedder:
-    """
-    A class for embedding texts and images using a pretrained CLIP model.
-    """
-    
-    def __init__(self, device='cuda', model_name='openai/clip-vit-base-patch32', cache_dir='./cache_dir', use_text_preprocessing=True, max_length=77):
-        """
-        Initializes the CLIPEmbedder with specified model and configurations.
-        """
-        self.device = torch.device(device)
-        self.model_name = model_name
-        self.cache_dir = cache_dir
-        self.use_text_preprocessing = use_text_preprocessing
-        self.max_length = max_length
-        
-        os.makedirs(self.cache_dir, exist_ok=True)
-        
-        self.processor = CLIPProcessor.from_pretrained(model_name, cache_dir=self.cache_dir)
-        self.model = CLIPModel.from_pretrained(model_name, cache_dir=self.cache_dir).to(self.device).eval()
-        self.tokenizer = CLIPTokenizer.from_pretrained(model_name)
-        self.text_model = CLIPTextModel.from_pretrained(model_name, cache_dir=self.cache_dir).to(self.device).eval()
-        
-        for param in self.text_model.parameters():
-            param.requires_grad = False
-
-    def get_text_embeddings(self, texts):
-        """
-        Generates embeddings for a list of text prompts.
-        """
-        self._validate_input_list(texts, str)
-        
-        if self.use_text_preprocessing:
-            texts = [self._clean_text(text) for text in texts]
-        
-        inputs = self.processor(text=texts, return_tensors="pt", padding=True, truncation=True, max_length=self.max_length).to(self.device)
-        
-        with torch.no_grad():
-            embeddings = self.model.get_text_features(**inputs)
-        
-        return embeddings
-
-    def encode_text(self, texts):
-        """
-        Encodes texts into embeddings and returns the last hidden state and pooled output.
-        """
-        self._validate_input_list(texts, str)
-        
-        batch_encoding = self.tokenizer(texts, return_tensors="pt", truncation=True, max_length=self.max_length, padding="max_length").to(self.device)
-        
-        with torch.no_grad():
-            outputs = self.text_model(**batch_encoding)
-        
-        return outputs.last_hidden_state, outputs.pooler_output
-
-    def get_image_embeddings(self, image_paths):
-        """
-        Generates embeddings for a list of image file paths.
-        """
-        self._validate_input_list(image_paths, str)
-        images = [self._load_image(path) for path in image_paths]
-        
-        inputs = self.processor(images=images, return_tensors="pt").to(self.device)
-        
-        with torch.no_grad():
-            embeddings = self.model.get_image_features(**inputs)
-        
-        return embeddings
-
-    def _validate_input_list(self, input_list, expected_type):
-        """
-        Validates that the input is a list of expected type.
-        """
-        if not isinstance(input_list, list) or not all(isinstance(item, expected_type) for item in input_list):
-            raise ValueError(f"Input must be a list of {expected_type.__name__}.")
-
-    def _clean_text(self, text):
-        """
-        Applies basic cleaning and formatting to a text string.
-        """
-        text = ftfy.fix_text(text)
-        text = html.unescape(text)
-        return text.strip()
-
-    def _load_image(self, image_path):
-        """
-        Loads and preprocesses an image from a file path.
-        """
-        try:
-            image = Image.open(image_path).convert("RGB")
-        except FileNotFoundError:
-            raise FileNotFoundError(f"Image file not found: {image_path}")
-        except Exception as e:
-            raise Exception(f"Error loading image {image_path}: {e}")
-        return image
-
-    def clean_caption(self, caption):
-        caption = str(caption)
-        caption = ul.unquote_plus(caption)
-        caption = caption.strip().lower()
-        caption = re.sub('<person>', 'person', caption)
-        # urls:
-        caption = re.sub(
-            r'\b((?:https?:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))',  # noqa
-            '', caption)  # regex for urls
-        caption = re.sub(
-            r'\b((?:www:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))',  # noqa
-            '', caption)  # regex for urls
-
-        caption = BeautifulSoup(caption, features='html.parser').text
-
-
-        caption = re.sub(r'@[\w\d]+\b', '', caption)
-
-        caption = re.sub(r'[\u31c0-\u31ef]+', '', caption)
-        caption = re.sub(r'[\u31f0-\u31ff]+', '', caption)
-        caption = re.sub(r'[\u3200-\u32ff]+', '', caption)
-        caption = re.sub(r'[\u3300-\u33ff]+', '', caption)
-        caption = re.sub(r'[\u3400-\u4dbf]+', '', caption)
-        caption = re.sub(r'[\u4dc0-\u4dff]+', '', caption)
-        caption = re.sub(r'[\u4e00-\u9fff]+', '', caption)
-
-        caption = re.sub(
-            r'[\u002D\u058A\u05BE\u1400\u1806\u2010-\u2015\u2E17\u2E1A\u2E3A\u2E3B\u2E40\u301C\u3030\u30A0\uFE31\uFE32\uFE58\uFE63\uFF0D]+',  # noqa
-            '-', caption)
-
-
-        caption = re.sub(r'[`´«»“”¨]', '"', caption)
-        caption = re.sub(r'[‘’]', "'", caption)
-
-
-        caption = re.sub(r'&quot;?', '', caption)
-
-        caption = re.sub(r'&amp', '', caption)
-
-
-        caption = re.sub(r'\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}', ' ', caption)
-
-
-        caption = re.sub(r'\d:\d\d\s+$', '', caption)
-
-
-        caption = re.sub(r'\\n', ' ', caption)
-
-
-        caption = re.sub(r'#\d{1,3}\b', '', caption)
-
-        caption = re.sub(r'#\d{5,}\b', '', caption)
-        caption = re.sub(r'\b\d{6,}\b', '', caption)
-        caption = re.sub(r'[\S]+\.(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)', '', caption)
-        caption = re.sub(r'[\"\']{2,}', r'"', caption)  
-        caption = re.sub(r'[\.]{2,}', r' ', caption)  
-
-        caption = re.sub(self.bad_punct_regex, r' ', caption)  
-        caption = re.sub(r'\s+\.\s+', r' ', caption)  
-        regex2 = re.compile(r'(?:\-|\_)')
-        if len(re.findall(regex2, caption)) > 3:
-            caption = re.sub(regex2, ' ', caption)
-        caption = self.basic_clean(caption)
-        caption = re.sub(r'\b[a-zA-Z]{1,3}\d{3,15}\b', '', caption)  # jc6640
-        caption = re.sub(r'\b[a-zA-Z]+\d+[a-zA-Z]+\b', '', caption)  # jc6640vc
-        caption = re.sub(r'\b\d+[a-zA-Z]+\d+\b', '', caption)  # 6640vc231
-
-        caption = re.sub(r'(worldwide\s+)?(free\s+)?shipping', '', caption)
-        caption = re.sub(r'(free\s)?download(\sfree)?', '', caption)
-        caption = re.sub(r'\bclick\b\s(?:for|on)\s\w+', '', caption)
-        caption = re.sub(r'\b(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)(\simage[s]?)?', '', caption)
-        caption = re.sub(r'\bpage\s+\d+\b', '', caption)
-
-        caption = re.sub(r'\b\d*[a-zA-Z]+\d+[a-zA-Z]+\d+[a-zA-Z\d]*\b', r' ', caption)  # j2d1a2a...
-
-        caption = re.sub(r'\b\d+\.?\d*[xх×]\d+\.?\d*\b', '', caption)
-
-        caption = re.sub(r'\b\s+\:\s+', r': ', caption)
-        caption = re.sub(r'(\D[,\./])\b', r'\1 ', caption)
-        caption = re.sub(r'\s+', ' ', caption)
-
-        caption.strip()
-
-        caption = re.sub(r'^[\"\']([\w\W]+)[\"\']$', r'\1', caption)
-        caption = re.sub(r'^[\'\_,\-\:;]', r'', caption)
-        caption = re.sub(r'[\'\_,\-\:\-\+]$', r'', caption)
-        caption = re.sub(r'^\.\S+$', '', caption)
-
-        return caption.strip()
-
-    @staticmethod
-    def basic_clean(text):
-        text = ftfy.fix_text(text)
-        text = html.unescape(html.unescape(text))
-        return text.strip()
-
-if __name__ == '__main__':
-
-    clip_embedder = CLIPEmbedder()
-
-    # Example
-    text_prompts = [
-        'A photo of a cute puppy playing with a ball.',
-        'An image of a beautiful sunset over the ocean.',
-        'A scene depicting a busy city street.'
-    ]
-    text_embeddings = clip_embedder.get_text_embeddings(text_prompts)
-    print(f"Text embeddings shape: {text_embeddings.shape}")
-
-    image_paths = ['image1.jpg', 'image2.png']
-    try:
-        image_embeddings = clip_embedder.get_image_embeddings(image_paths)
-        print(f"Image embeddings shape: {image_embeddings.shape}")
-    except FileNotFoundError as e:
-        print(e)
-    except Exception as e:
-        print(f"An error occurred: {e}")
-

opensora/models/text_encoder/t5.py

@@ -1,197 +0,0 @@
-# -*- coding: utf-8 -*-
-import os
-import re
-import html
-import urllib.parse as ul
-
-import ftfy
-import torch
-from bs4 import BeautifulSoup
-from transformers import T5EncoderModel, AutoTokenizer
-from huggingface_hub import hf_hub_download
-
-class T5Embedder:
-
-    available_models = ['t5-v1_1-xxl']
-    bad_punct_regex = re.compile(r'['+'#®•©™&@·º½¾¿¡§~'+'\)'+'\('+'\]'+'\['+'\}'+'\{'+'\|'+'\\'+'\/'+'\*' + r']{1,}')  # noqa
-
-    def __init__(self, device, dir_or_name='t5-v1_1-xxl', *, cache_dir='./cache_dir', hf_token=None, use_text_preprocessing=True,
-                 t5_model_kwargs=None, torch_dtype=None, model_max_length=120):
-        self.device = torch.device(device)
-        self.torch_dtype = torch_dtype or torch.bfloat16
-        if t5_model_kwargs is None:
-            t5_model_kwargs = {'low_cpu_mem_usage': True, 'torch_dtype': self.torch_dtype}
-            t5_model_kwargs['device_map'] = {'shared': self.device, 'encoder': self.device}
-
-        self.use_text_preprocessing = use_text_preprocessing
-        self.hf_token = hf_token
-        self.cache_dir = cache_dir
-        self.dir_or_name = dir_or_name
-        cache_dir = os.path.join(self.cache_dir, 't5-v1_1-xxl')
-        for filename in ['config.json', 'special_tokens_map.json', 'spiece.model', 'tokenizer_config.json',
-                         'pytorch_model-00001-of-00002.bin', 'pytorch_model-00002-of-00002.bin', 'pytorch_model.bin.index.json']:
-            hf_hub_download(repo_id='DeepFloyd/t5-v1_1-xxl', filename=filename, cache_dir=cache_dir,
-                            force_filename=filename, token=self.hf_token)
-
-        print(cache_dir)
-        self.tokenizer = AutoTokenizer.from_pretrained(cache_dir)
-        self.model = T5EncoderModel.from_pretrained(cache_dir, **t5_model_kwargs).eval()
-        self.model_max_length = model_max_length
-
-    def get_text_embeddings(self, texts):
-        texts = [self.text_preprocessing(text) for text in texts]
-
-        text_tokens_and_mask = self.tokenizer(
-            texts,
-            max_length=self.model_max_length,
-            padding='max_length',
-            truncation=True,
-            return_attention_mask=True,
-            add_special_tokens=True,
-            return_tensors='pt'
-        )
-
-        text_tokens_and_mask['input_ids'] = text_tokens_and_mask['input_ids']
-        text_tokens_and_mask['attention_mask'] = text_tokens_and_mask['attention_mask']
-
-        with torch.no_grad():
-            text_encoder_embs = self.model(
-                input_ids=text_tokens_and_mask['input_ids'].to(self.device),
-                attention_mask=text_tokens_and_mask['attention_mask'].to(self.device),
-            )['last_hidden_state'].detach()
-        return text_encoder_embs, text_tokens_and_mask['attention_mask'].to(self.device)
-
-    def text_preprocessing(self, text):
-        if self.use_text_preprocessing:
-            # The exact text cleaning as was in the training stage:
-            text = self.clean_caption(text)
-            text = self.clean_caption(text)
-            return text
-        else:
-            return text.lower().strip()
-
-    @staticmethod
-    def basic_clean(text):
-        text = ftfy.fix_text(text)
-        text = html.unescape(html.unescape(text))
-        return text.strip()
-
-    def clean_caption(self, caption):
-        caption = str(caption)
-        caption = ul.unquote_plus(caption)
-        caption = caption.strip().lower()
-        caption = re.sub('<person>', 'person', caption)
-        # urls:
-        caption = re.sub(
-            r'\b((?:https?:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))',  # noqa
-            '', caption)  # regex for urls
-        caption = re.sub(
-            r'\b((?:www:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))',  # noqa
-            '', caption)  # regex for urls
-        # html:
-        caption = BeautifulSoup(caption, features='html.parser').text
-
-        # @<nickname>
-        caption = re.sub(r'@[\w\d]+\b', '', caption)
-
-        # 31C0—31EF CJK Strokes
-        # 31F0—31FF Katakana Phonetic Extensions
-        # 3200—32FF Enclosed CJK Letters and Months
-        # 3300—33FF CJK Compatibility
-        # 3400—4DBF CJK Unified Ideographs Extension A
-        # 4DC0—4DFF Yijing Hexagram Symbols
-        # 4E00—9FFF CJK Unified Ideographs
-        caption = re.sub(r'[\u31c0-\u31ef]+', '', caption)
-        caption = re.sub(r'[\u31f0-\u31ff]+', '', caption)
-        caption = re.sub(r'[\u3200-\u32ff]+', '', caption)
-        caption = re.sub(r'[\u3300-\u33ff]+', '', caption)
-        caption = re.sub(r'[\u3400-\u4dbf]+', '', caption)
-        caption = re.sub(r'[\u4dc0-\u4dff]+', '', caption)
-        caption = re.sub(r'[\u4e00-\u9fff]+', '', caption)
-        #######################################################
-
-        # все виды тире / all types of dash --> "-"
-        caption = re.sub(
-            r'[\u002D\u058A\u05BE\u1400\u1806\u2010-\u2015\u2E17\u2E1A\u2E3A\u2E3B\u2E40\u301C\u3030\u30A0\uFE31\uFE32\uFE58\uFE63\uFF0D]+',  # noqa
-            '-', caption)
-
-        # кавычки к одному стандарту
-        caption = re.sub(r'[`´«»“”¨]', '"', caption)
-        caption = re.sub(r'[‘’]', "'", caption)
-
-        # &quot;
-        caption = re.sub(r'&quot;?', '', caption)
-        # &amp
-        caption = re.sub(r'&amp', '', caption)
-
-        # ip adresses:
-        caption = re.sub(r'\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}', ' ', caption)
-
-        # article ids:
-        caption = re.sub(r'\d:\d\d\s+$', '', caption)
-
-        # \n
-        caption = re.sub(r'\\n', ' ', caption)
-
-        # "#123"
-        caption = re.sub(r'#\d{1,3}\b', '', caption)
-        # "#12345.."
-        caption = re.sub(r'#\d{5,}\b', '', caption)
-        # "123456.."
-        caption = re.sub(r'\b\d{6,}\b', '', caption)
-        # filenames:
-        caption = re.sub(r'[\S]+\.(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)', '', caption)
-
-        #
-        caption = re.sub(r'[\"\']{2,}', r'"', caption)  # """AUSVERKAUFT"""
-        caption = re.sub(r'[\.]{2,}', r' ', caption)  # """AUSVERKAUFT"""
-
-        caption = re.sub(self.bad_punct_regex, r' ', caption)  # ***AUSVERKAUFT***, #AUSVERKAUFT
-        caption = re.sub(r'\s+\.\s+', r' ', caption)  # " . "
-
-        # this-is-my-cute-cat / this_is_my_cute_cat
-        regex2 = re.compile(r'(?:\-|\_)')
-        if len(re.findall(regex2, caption)) > 3:
-            caption = re.sub(regex2, ' ', caption)
-
-        caption = self.basic_clean(caption)
-
-        caption = re.sub(r'\b[a-zA-Z]{1,3}\d{3,15}\b', '', caption)  # jc6640
-        caption = re.sub(r'\b[a-zA-Z]+\d+[a-zA-Z]+\b', '', caption)  # jc6640vc
-        caption = re.sub(r'\b\d+[a-zA-Z]+\d+\b', '', caption)  # 6640vc231
-
-        caption = re.sub(r'(worldwide\s+)?(free\s+)?shipping', '', caption)
-        caption = re.sub(r'(free\s)?download(\sfree)?', '', caption)
-        caption = re.sub(r'\bclick\b\s(?:for|on)\s\w+', '', caption)
-        caption = re.sub(r'\b(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)(\simage[s]?)?', '', caption)
-        caption = re.sub(r'\bpage\s+\d+\b', '', caption)
-
-        caption = re.sub(r'\b\d*[a-zA-Z]+\d+[a-zA-Z]+\d+[a-zA-Z\d]*\b', r' ', caption)  # j2d1a2a...
-
-        caption = re.sub(r'\b\d+\.?\d*[xх×]\d+\.?\d*\b', '', caption)
-
-        caption = re.sub(r'\b\s+\:\s+', r': ', caption)
-        caption = re.sub(r'(\D[,\./])\b', r'\1 ', caption)
-        caption = re.sub(r'\s+', ' ', caption)
-
-        caption.strip()
-
-        caption = re.sub(r'^[\"\']([\w\W]+)[\"\']$', r'\1', caption)
-        caption = re.sub(r'^[\'\_,\-\:;]', r'', caption)
-        caption = re.sub(r'[\'\_,\-\:\-\+]$', r'', caption)
-        caption = re.sub(r'^\.\S+$', '', caption)
-
-        return caption.strip()
-
-if __name__ == '__main__':
-    t5 = T5Embedder(device="cuda", cache_dir='./cache_dir', torch_dtype=torch.float)
-    device = t5.device
-    prompts = ['I am a test caption', 'Test twice']
-    with torch.no_grad():
-        caption_embs, emb_masks = t5.get_text_embeddings(prompts)
-        emb_dict = {
-            'caption_feature': caption_embs.float().cpu().data.numpy(),
-            'attention_mask': emb_masks.cpu().data.numpy(),
-        }
-    import ipdb;ipdb.set_trace()
-    print()

opensora/sample/pipeline_opensora.py

@@ -1,817 +0,0 @@
-# Copyright 2024 PixArt-Sigma Authors and 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 html
-import inspect
-import re
-import urllib.parse as ul
-from typing import Callable, List, Optional, Tuple, Union
-import math
-import torch
-from transformers import T5EncoderModel, T5Tokenizer
-
-from diffusers.models import AutoencoderKL, Transformer2DModel
-from diffusers.schedulers import DPMSolverMultistepScheduler
-from diffusers.utils import (
-    BACKENDS_MAPPING,
-    deprecate,
-    is_bs4_available,
-    is_ftfy_available,
-    logging,
-    replace_example_docstring,
-)
-from diffusers.utils.torch_utils import randn_tensor
-from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput
-
-
-
-logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
-
-if is_bs4_available():
-    from bs4 import BeautifulSoup
-
-if is_ftfy_available():
-    import ftfy
-
-
-EXAMPLE_DOC_STRING = """
-    Examples:
-        ```py
-        >>> import torch
-        >>> from diffusers import PixArtSigmaPipeline
-
-        >>> # You can replace the checkpoint id with "PixArt-alpha/PixArt-Sigma-XL-2-512-MS" too.
-        >>> pipe = PixArtSigmaPipeline.from_pretrained(
-        ...     "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS", torch_dtype=torch.float16
-        ... )
-        >>> # Enable memory optimizations.
-        >>> # pipe.enable_model_cpu_offload()
-
-        >>> prompt = "A small cactus with a happy face in the Sahara desert."
-        >>> image = pipe(prompt).images[0]
-        ```
-"""
-
-
-# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
-def retrieve_timesteps(
-    scheduler,
-    num_inference_steps: Optional[int] = None,
-    device: Optional[Union[str, torch.device]] = None,
-    timesteps: Optional[List[int]] = None,
-    **kwargs,
-):
-    """
-    Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
-    custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
-
-    Args:
-        scheduler (`SchedulerMixin`):
-            The scheduler to get timesteps from.
-        num_inference_steps (`int`):
-            The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
-            must be `None`.
-        device (`str` or `torch.device`, *optional*):
-            The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
-        timesteps (`List[int]`, *optional*):
-                Custom timesteps used to support arbitrary spacing between timesteps. If `None`, then the default
-                timestep spacing strategy of the scheduler is used. If `timesteps` is passed, `num_inference_steps`
-                must be `None`.
-
-    Returns:
-        `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
-        second element is the number of inference steps.
-    """
-    if timesteps is not None:
-        accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
-        if not accepts_timesteps:
-            raise ValueError(
-                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
-                f" timestep schedules. Please check whether you are using the correct scheduler."
-            )
-        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
-        timesteps = scheduler.timesteps
-        num_inference_steps = len(timesteps)
-    else:
-        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
-        timesteps = scheduler.timesteps
-    return timesteps, num_inference_steps
-
-
-class OpenSoraPipeline(DiffusionPipeline):
-    r"""
-    Pipeline for text-to-image generation using PixArt-Sigma.
-    """
-
-    bad_punct_regex = re.compile(
-        r"["
-        + "#®•©™&@·º½¾¿¡§~"
-        + r"\)"
-        + r"\("
-        + r"\]"
-        + r"\["
-        + r"\}"
-        + r"\{"
-        + r"\|"
-        + "\\"
-        + r"\/"
-        + r"\*"
-        + r"]{1,}"
-    )  # noqa
-
-    _optional_components = ["tokenizer", "text_encoder"]
-    model_cpu_offload_seq = "text_encoder->transformer->vae"
-
-    def __init__(
-        self,
-        tokenizer: T5Tokenizer,
-        text_encoder: T5EncoderModel,
-        vae: AutoencoderKL,
-        transformer: Transformer2DModel,
-        scheduler: DPMSolverMultistepScheduler,
-    ):
-        super().__init__()
-
-        self.register_modules(
-            tokenizer=tokenizer, text_encoder=text_encoder, vae=vae, transformer=transformer, scheduler=scheduler
-        )
-
-        # self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
-
-    # Copied from diffusers.pipelines.pixart_alpha.pipeline_pixart_alpha.PixArtAlphaPipeline.encode_prompt
-    def encode_prompt(
-        self,
-        prompt: Union[str, List[str]],
-        do_classifier_free_guidance: bool = True,
-        negative_prompt: str = "",
-        num_images_per_prompt: int = 1,
-        device: Optional[torch.device] = None,
-        prompt_embeds: Optional[torch.FloatTensor] = None,
-        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
-        prompt_attention_mask: Optional[torch.FloatTensor] = None,
-        negative_prompt_attention_mask: Optional[torch.FloatTensor] = None,
-        clean_caption: bool = False,
-        max_sequence_length: int = 120,
-        **kwargs,
-    ):
-        r"""
-        Encodes the prompt into text encoder hidden states.
-
-        Args:
-            prompt (`str` or `List[str]`, *optional*):
-                prompt to be encoded
-            negative_prompt (`str` or `List[str]`, *optional*):
-                The prompt not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds`
-                instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). For
-                PixArt-Alpha, this should be "".
-            do_classifier_free_guidance (`bool`, *optional*, defaults to `True`):
-                whether to use classifier free guidance or not
-            num_images_per_prompt (`int`, *optional*, defaults to 1):
-                number of images that should be generated per prompt
-            device: (`torch.device`, *optional*):
-                torch device to place the resulting embeddings on
-            prompt_embeds (`torch.FloatTensor`, *optional*):
-                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
-                provided, text embeddings will be generated from `prompt` input argument.
-            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
-                Pre-generated negative text embeddings. For PixArt-Alpha, it's should be the embeddings of the ""
-                string.
-            clean_caption (`bool`, defaults to `False`):
-                If `True`, the function will preprocess and clean the provided caption before encoding.
-            max_sequence_length (`int`, defaults to 120): Maximum sequence length to use for the prompt.
-        """
-
-        if "mask_feature" in kwargs:
-            deprecation_message = "The use of `mask_feature` is deprecated. It is no longer used in any computation and that doesn't affect the end results. It will be removed in a future version."
-            deprecate("mask_feature", "1.0.0", deprecation_message, standard_warn=False)
-
-        if device is None:
-            device = getattr(self, '_execution_device', None) or getattr(self, 'device', None) or torch.device('cuda')
-
-        if prompt is not None and isinstance(prompt, str):
-            batch_size = 1
-        elif prompt is not None and isinstance(prompt, list):
-            batch_size = len(prompt)
-        else:
-            batch_size = prompt_embeds.shape[0]
-
-        # See Section 3.1. of the paper.
-        max_length = max_sequence_length
-
-        if prompt_embeds is None:
-            prompt = self._text_preprocessing(prompt, clean_caption=clean_caption)
-            text_inputs = self.tokenizer(
-                prompt,
-                padding="max_length",
-                max_length=max_length,
-                truncation=True,
-                add_special_tokens=True,
-                return_tensors="pt",
-            )
-            text_input_ids = text_inputs.input_ids
-            untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
-
-            if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(
-                text_input_ids, untruncated_ids
-            ):
-                removed_text = self.tokenizer.batch_decode(untruncated_ids[:, max_length - 1 : -1])
-                logger.warning(
-                    "The following part of your input was truncated because CLIP can only handle sequences up to"
-                    f" {max_length} tokens: {removed_text}"
-                )
-
-            prompt_attention_mask = text_inputs.attention_mask
-            prompt_attention_mask = prompt_attention_mask.to(device)
-
-            prompt_embeds = self.text_encoder(text_input_ids.to(device), attention_mask=prompt_attention_mask)
-            prompt_embeds = prompt_embeds[0]
-
-        if self.text_encoder is not None:
-            dtype = self.text_encoder.dtype
-        elif self.transformer is not None:
-            dtype = self.transformer.dtype
-        else:
-            dtype = None
-
-        prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
-
-        bs_embed, seq_len, _ = prompt_embeds.shape
-        # duplicate text embeddings and attention mask for each generation per prompt, using mps friendly method
-        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
-        prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)
-        prompt_attention_mask = prompt_attention_mask.view(bs_embed, -1)
-        prompt_attention_mask = prompt_attention_mask.repeat(num_images_per_prompt, 1)
-
-        # get unconditional embeddings for classifier free guidance
-        if do_classifier_free_guidance and negative_prompt_embeds is None:
-            uncond_tokens = [negative_prompt] * batch_size
-            uncond_tokens = self._text_preprocessing(uncond_tokens, clean_caption=clean_caption)
-            max_length = prompt_embeds.shape[1]
-            uncond_input = self.tokenizer(
-                uncond_tokens,
-                padding="max_length",
-                max_length=max_length,
-                truncation=True,
-                return_attention_mask=True,
-                add_special_tokens=True,
-                return_tensors="pt",
-            )
-            negative_prompt_attention_mask = uncond_input.attention_mask
-            negative_prompt_attention_mask = negative_prompt_attention_mask.to(device)
-
-            negative_prompt_embeds = self.text_encoder(
-                uncond_input.input_ids.to(device), attention_mask=negative_prompt_attention_mask
-            )
-            negative_prompt_embeds = negative_prompt_embeds[0]
-
-        if do_classifier_free_guidance:
-            # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
-            seq_len = negative_prompt_embeds.shape[1]
-
-            negative_prompt_embeds = negative_prompt_embeds.to(dtype=dtype, device=device)
-
-            negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)
-            negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
-
-            negative_prompt_attention_mask = negative_prompt_attention_mask.view(bs_embed, -1)
-            negative_prompt_attention_mask = negative_prompt_attention_mask.repeat(num_images_per_prompt, 1)
-        else:
-            negative_prompt_embeds = None
-            negative_prompt_attention_mask = None
-
-        return prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask
-
-    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs
-    def prepare_extra_step_kwargs(self, generator, eta):
-        # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
-        # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
-        # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
-        # and should be between [0, 1]
-
-        accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
-        extra_step_kwargs = {}
-        if accepts_eta:
-            extra_step_kwargs["eta"] = eta
-
-        # check if the scheduler accepts generator
-        accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
-        if accepts_generator:
-            extra_step_kwargs["generator"] = generator
-        return extra_step_kwargs
-
-    # Copied from diffusers.pipelines.pixart_alpha.pipeline_pixart_alpha.PixArtAlphaPipeline.check_inputs
-    def check_inputs(
-        self,
-        prompt,
-        num_frames, 
-        height,
-        width,
-        negative_prompt,
-        callback_steps,
-        prompt_embeds=None,
-        negative_prompt_embeds=None,
-        prompt_attention_mask=None,
-        negative_prompt_attention_mask=None,
-    ):
-        if num_frames <= 0:
-            raise ValueError(f"`num_frames` have to be positive but is {num_frames}.")
-        if height % 8 != 0 or width % 8 != 0:
-            raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")
-
-        if (callback_steps is None) or (
-            callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0)
-        ):
-            raise ValueError(
-                f"`callback_steps` has to be a positive integer but is {callback_steps} of type"
-                f" {type(callback_steps)}."
-            )
-
-        if prompt is not None and prompt_embeds is not None:
-            raise ValueError(
-                f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
-                " only forward one of the two."
-            )
-        elif prompt is None and prompt_embeds is None:
-            raise ValueError(
-                "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
-            )
-        elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
-            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
-
-        if prompt is not None and negative_prompt_embeds is not None:
-            raise ValueError(
-                f"Cannot forward both `prompt`: {prompt} and `negative_prompt_embeds`:"
-                f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
-            )
-
-        if negative_prompt is not None and negative_prompt_embeds is not None:
-            raise ValueError(
-                f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
-                f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
-            )
-
-        if prompt_embeds is not None and prompt_attention_mask is None:
-            raise ValueError("Must provide `prompt_attention_mask` when specifying `prompt_embeds`.")
-
-        if negative_prompt_embeds is not None and negative_prompt_attention_mask is None:
-            raise ValueError("Must provide `negative_prompt_attention_mask` when specifying `negative_prompt_embeds`.")
-
-        if prompt_embeds is not None and negative_prompt_embeds is not None:
-            if prompt_embeds.shape != negative_prompt_embeds.shape:
-                raise ValueError(
-                    "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
-                    f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
-                    f" {negative_prompt_embeds.shape}."
-                )
-            if prompt_attention_mask.shape != negative_prompt_attention_mask.shape:
-                raise ValueError(
-                    "`prompt_attention_mask` and `negative_prompt_attention_mask` must have the same shape when passed directly, but"
-                    f" got: `prompt_attention_mask` {prompt_attention_mask.shape} != `negative_prompt_attention_mask`"
-                    f" {negative_prompt_attention_mask.shape}."
-                )
-
-    # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline._text_preprocessing
-    def _text_preprocessing(self, text, clean_caption=False):
-        if clean_caption and not is_bs4_available():
-            logger.warning(BACKENDS_MAPPING["bs4"][-1].format("Setting `clean_caption=True`"))
-            logger.warning("Setting `clean_caption` to False...")
-            clean_caption = False
-
-        if clean_caption and not is_ftfy_available():
-            logger.warning(BACKENDS_MAPPING["ftfy"][-1].format("Setting `clean_caption=True`"))
-            logger.warning("Setting `clean_caption` to False...")
-            clean_caption = False
-
-        if not isinstance(text, (tuple, list)):
-            text = [text]
-
-        def process(text: str):
-            if clean_caption:
-                text = self._clean_caption(text)
-                text = self._clean_caption(text)
-            else:
-                text = text.lower().strip()
-            return text
-
-        return [process(t) for t in text]
-
-    # Copied from diffusers.pipelines.deepfloyd_if.pipeline_if.IFPipeline._clean_caption
-    def _clean_caption(self, caption):
-        caption = str(caption)
-        caption = ul.unquote_plus(caption)
-        caption = caption.strip().lower()
-        caption = re.sub("<person>", "person", caption)
-        # urls:
-        caption = re.sub(
-            r"\b((?:https?:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))",  # noqa
-            "",
-            caption,
-        )  # regex for urls
-        caption = re.sub(
-            r"\b((?:www:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))",  # noqa
-            "",
-            caption,
-        )  # regex for urls
-        # html:
-        caption = BeautifulSoup(caption, features="html.parser").text
-
-        # @<nickname>
-        caption = re.sub(r"@[\w\d]+\b", "", caption)
-
-        # 31C0—31EF CJK Strokes
-        # 31F0—31FF Katakana Phonetic Extensions
-        # 3200—32FF Enclosed CJK Letters and Months
-        # 3300—33FF CJK Compatibility
-        # 3400—4DBF CJK Unified Ideographs Extension A
-        # 4DC0—4DFF Yijing Hexagram Symbols
-        # 4E00—9FFF CJK Unified Ideographs
-        caption = re.sub(r"[\u31c0-\u31ef]+", "", caption)
-        caption = re.sub(r"[\u31f0-\u31ff]+", "", caption)
-        caption = re.sub(r"[\u3200-\u32ff]+", "", caption)
-        caption = re.sub(r"[\u3300-\u33ff]+", "", caption)
-        caption = re.sub(r"[\u3400-\u4dbf]+", "", caption)
-        caption = re.sub(r"[\u4dc0-\u4dff]+", "", caption)
-        # caption = re.sub(r"[\u4e00-\u9fff]+", "", caption)
-        #######################################################
-
-        # все виды тире / all types of dash --> "-"
-        caption = re.sub(
-            r"[\u002D\u058A\u05BE\u1400\u1806\u2010-\u2015\u2E17\u2E1A\u2E3A\u2E3B\u2E40\u301C\u3030\u30A0\uFE31\uFE32\uFE58\uFE63\uFF0D]+",  # noqa
-            "-",
-            caption,
-        )
-
-        # кавычки к одному стандарту
-        caption = re.sub(r"[`´«»“”¨]", '"', caption)
-        caption = re.sub(r"[‘’]", "'", caption)
-
-        # &quot;
-        caption = re.sub(r"&quot;?", "", caption)
-        # &amp
-        caption = re.sub(r"&amp", "", caption)
-
-        # ip adresses:
-        caption = re.sub(r"\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}", " ", caption)
-
-        # article ids:
-        caption = re.sub(r"\d:\d\d\s+$", "", caption)
-
-        # \n
-        caption = re.sub(r"\\n", " ", caption)
-
-        # "#123"
-        caption = re.sub(r"#\d{1,3}\b", "", caption)
-        # "#12345.."
-        caption = re.sub(r"#\d{5,}\b", "", caption)
-        # "123456.."
-        caption = re.sub(r"\b\d{6,}\b", "", caption)
-        # filenames:
-        caption = re.sub(r"[\S]+\.(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)", "", caption)
-
-        #
-        caption = re.sub(r"[\"\']{2,}", r'"', caption)  # """AUSVERKAUFT"""
-        caption = re.sub(r"[\.]{2,}", r" ", caption)  # """AUSVERKAUFT"""
-
-        caption = re.sub(self.bad_punct_regex, r" ", caption)  # ***AUSVERKAUFT***, #AUSVERKAUFT
-        caption = re.sub(r"\s+\.\s+", r" ", caption)  # " . "
-
-        # this-is-my-cute-cat / this_is_my_cute_cat
-        regex2 = re.compile(r"(?:\-|\_)")
-        if len(re.findall(regex2, caption)) > 3:
-            caption = re.sub(regex2, " ", caption)
-
-        caption = ftfy.fix_text(caption)
-        caption = html.unescape(html.unescape(caption))
-
-        caption = re.sub(r"\b[a-zA-Z]{1,3}\d{3,15}\b", "", caption)  # jc6640
-        caption = re.sub(r"\b[a-zA-Z]+\d+[a-zA-Z]+\b", "", caption)  # jc6640vc
-        caption = re.sub(r"\b\d+[a-zA-Z]+\d+\b", "", caption)  # 6640vc231
-
-        caption = re.sub(r"(worldwide\s+)?(free\s+)?shipping", "", caption)
-        caption = re.sub(r"(free\s)?download(\sfree)?", "", caption)
-        caption = re.sub(r"\bclick\b\s(?:for|on)\s\w+", "", caption)
-        caption = re.sub(r"\b(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)(\simage[s]?)?", "", caption)
-        caption = re.sub(r"\bpage\s+\d+\b", "", caption)
-
-        caption = re.sub(r"\b\d*[a-zA-Z]+\d+[a-zA-Z]+\d+[a-zA-Z\d]*\b", r" ", caption)  # j2d1a2a...
-
-        caption = re.sub(r"\b\d+\.?\d*[xх×]\d+\.?\d*\b", "", caption)
-
-        caption = re.sub(r"\b\s+\:\s+", r": ", caption)
-        caption = re.sub(r"(\D[,\./])\b", r"\1 ", caption)
-        caption = re.sub(r"\s+", " ", caption)
-
-        caption.strip()
-
-        caption = re.sub(r"^[\"\']([\w\W]+)[\"\']$", r"\1", caption)
-        caption = re.sub(r"^[\'\_,\-\:;]", r"", caption)
-        caption = re.sub(r"[\'\_,\-\:\-\+]$", r"", caption)
-        caption = re.sub(r"^\.\S+$", "", caption)
-        return caption.strip()
-    
-    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents
-    def prepare_latents(self, batch_size, num_channels_latents, num_frames, height, width, dtype, device, generator, latents=None):
-        shape = (
-            batch_size,
-            num_channels_latents,
-            (math.ceil((int(num_frames) - 1) / self.vae.vae_scale_factor[0]) + 1) if int(num_frames) % 2 == 1 else math.ceil(int(num_frames) / self.vae.vae_scale_factor[0]), 
-            math.ceil(int(height) / self.vae.vae_scale_factor[1]),
-            math.ceil(int(width) / self.vae.vae_scale_factor[2]),
-        )
-        if isinstance(generator, list) and len(generator) != batch_size:
-            raise ValueError(
-                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
-                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
-            )
-
-        if latents is None:
-            latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
-        else:
-            latents = latents.to(device)
-
-        # scale the initial noise by the standard deviation required by the scheduler
-        latents = latents * self.scheduler.init_noise_sigma
-
-
-        return latents
-
-    @torch.no_grad()
-    def __call__(
-        self,
-        prompt: Union[str, List[str]] = None,
-        negative_prompt: str = "",
-        num_inference_steps: int = 20,
-        timesteps: List[int] = None,
-        guidance_scale: float = 4.5,
-        num_images_per_prompt: Optional[int] = 1,
-        num_frames: Optional[int] = None,
-        height: Optional[int] = None,
-        width: Optional[int] = None,
-        eta: float = 0.0,
-        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
-        latents: Optional[torch.FloatTensor] = None,
-        prompt_embeds: Optional[torch.FloatTensor] = None,
-        prompt_attention_mask: Optional[torch.FloatTensor] = None,
-        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
-        negative_prompt_attention_mask: Optional[torch.FloatTensor] = None,
-        output_type: Optional[str] = "pil",
-        return_dict: bool = True,
-        callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
-        callback_steps: int = 1,
-        clean_caption: bool = True,
-        use_resolution_binning: bool = True,
-        max_sequence_length: int = 300,
-        **kwargs,
-    ) -> Union[ImagePipelineOutput, Tuple]:
-        """
-        Function invoked when calling the pipeline for generation.
-
-        Args:
-            prompt (`str` or `List[str]`, *optional*):
-                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
-                instead.
-            negative_prompt (`str` or `List[str]`, *optional*):
-                The prompt or prompts not to guide the image generation. If not defined, one has to pass
-                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
-                less than `1`).
-            num_inference_steps (`int`, *optional*, defaults to 100):
-                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
-                expense of slower inference.
-            timesteps (`List[int]`, *optional*):
-                Custom timesteps to use for the denoising process. If not defined, equal spaced `num_inference_steps`
-                timesteps are used. Must be in descending order.
-            guidance_scale (`float`, *optional*, defaults to 4.5):
-                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
-                `guidance_scale` is defined as `w` of equation 2. of [Imagen
-                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
-                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
-                usually at the expense of lower image quality.
-            num_images_per_prompt (`int`, *optional*, defaults to 1):
-                The number of images to generate per prompt.
-            height (`int`, *optional*, defaults to self.unet.config.sample_size):
-                The height in pixels of the generated image.
-            width (`int`, *optional*, defaults to self.unet.config.sample_size):
-                The width in pixels of the generated image.
-            eta (`float`, *optional*, defaults to 0.0):
-                Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to
-                [`schedulers.DDIMScheduler`], will be ignored for others.
-            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
-                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
-                to make generation deterministic.
-            latents (`torch.FloatTensor`, *optional*):
-                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
-                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
-                tensor will ge generated by sampling using the supplied random `generator`.
-            prompt_embeds (`torch.FloatTensor`, *optional*):
-                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
-                provided, text embeddings will be generated from `prompt` input argument.
-            prompt_attention_mask (`torch.FloatTensor`, *optional*): Pre-generated attention mask for text embeddings.
-            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
-                Pre-generated negative text embeddings. For PixArt-Sigma this negative prompt should be "". If not
-                provided, negative_prompt_embeds will be generated from `negative_prompt` input argument.
-            negative_prompt_attention_mask (`torch.FloatTensor`, *optional*):
-                Pre-generated attention mask for negative text embeddings.
-            output_type (`str`, *optional*, defaults to `"pil"`):
-                The output format of the generate image. Choose between
-                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
-            return_dict (`bool`, *optional*, defaults to `True`):
-                Whether or not to return a [`~pipelines.stable_diffusion.IFPipelineOutput`] instead of a plain tuple.
-            callback (`Callable`, *optional*):
-                A function that will be called every `callback_steps` steps during inference. The function will be
-                called with the following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`.
-            callback_steps (`int`, *optional*, defaults to 1):
-                The frequency at which the `callback` function will be called. If not specified, the callback will be
-                called at every step.
-            clean_caption (`bool`, *optional*, defaults to `True`):
-                Whether or not to clean the caption before creating embeddings. Requires `beautifulsoup4` and `ftfy` to
-                be installed. If the dependencies are not installed, the embeddings will be created from the raw
-                prompt.
-            use_resolution_binning (`bool` defaults to `True`):
-                If set to `True`, the requested height and width are first mapped to the closest resolutions using
-                `ASPECT_RATIO_1024_BIN`. After the produced latents are decoded into images, they are resized back to
-                the requested resolution. Useful for generating non-square images.
-            max_sequence_length (`int` defaults to 120): Maximum sequence length to use with the `prompt`.
-
-        Examples:
-
-        Returns:
-            [`~pipelines.ImagePipelineOutput`] or `tuple`:
-                If `return_dict` is `True`, [`~pipelines.ImagePipelineOutput`] is returned, otherwise a `tuple` is
-                returned where the first element is a list with the generated images
-        """
-        # 1. Check inputs. Raise error if not correct
-        num_frames = num_frames or self.transformer.config.sample_size_t * self.vae.vae_scale_factor[0]
-        height = height or self.transformer.config.sample_size[0] * self.vae.vae_scale_factor[1]
-        width = width or self.transformer.config.sample_size[1] * self.vae.vae_scale_factor[2]
-        self.check_inputs(
-            prompt,
-            num_frames, 
-            height,
-            width,
-            negative_prompt,
-            callback_steps,
-            prompt_embeds,
-            negative_prompt_embeds,
-            prompt_attention_mask,
-            negative_prompt_attention_mask,
-        )
-
-        # 2. Default height and width to transformer
-        if prompt is not None and isinstance(prompt, str):
-            batch_size = 1
-        elif prompt is not None and isinstance(prompt, list):
-            batch_size = len(prompt)
-        else:
-            batch_size = prompt_embeds.shape[0]
-        # import ipdb;ipdb.set_trace()
-        device = getattr(self, '_execution_device', None) or getattr(self, 'device', None) or torch.device('cuda')
-
-        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
-        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
-        # corresponds to doing no classifier free guidance.
-        do_classifier_free_guidance = guidance_scale > 1.0
-
-        # 3. Encode input prompt
-        (
-            prompt_embeds,
-            prompt_attention_mask,
-            negative_prompt_embeds,
-            negative_prompt_attention_mask,
-        ) = self.encode_prompt(
-            prompt,
-            do_classifier_free_guidance,
-            negative_prompt=negative_prompt,
-            num_images_per_prompt=num_images_per_prompt,
-            device=device,
-            prompt_embeds=prompt_embeds,
-            negative_prompt_embeds=negative_prompt_embeds,
-            prompt_attention_mask=prompt_attention_mask,
-            negative_prompt_attention_mask=negative_prompt_attention_mask,
-            clean_caption=clean_caption,
-            max_sequence_length=max_sequence_length,
-        )
-        if do_classifier_free_guidance:
-            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
-            prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], dim=0)
-
-        # 4. Prepare timesteps
-        timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps)
-
-        # 5. Prepare latents.
-        latent_channels = self.transformer.config.in_channels
-        latents = self.prepare_latents(
-            batch_size * num_images_per_prompt,
-            latent_channels,
-            num_frames, 
-            height,
-            width,
-            prompt_embeds.dtype,
-            device,
-            generator,
-            latents,
-        )
-
-        # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
-        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
-
-        # 6.1 Prepare micro-conditions.
-        added_cond_kwargs = {"resolution": None, "aspect_ratio": None}
-
-        # 7. Denoising loop
-        num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
-
-        with self.progress_bar(total=num_inference_steps) as progress_bar:
-            for i, t in enumerate(timesteps):
-                latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
-                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
-
-                current_timestep = t
-                if not torch.is_tensor(current_timestep):
-                    # 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 = latent_model_input.device.type == "mps"
-                    if isinstance(current_timestep, float):
-                        dtype = torch.float32 if is_mps else torch.float64
-                    else:
-                        dtype = torch.int32 if is_mps else torch.int64
-                    current_timestep = torch.tensor([current_timestep], dtype=dtype, device=latent_model_input.device)
-                elif len(current_timestep.shape) == 0:
-                    current_timestep = current_timestep[None].to(latent_model_input.device)
-                # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
-                current_timestep = current_timestep.expand(latent_model_input.shape[0])
-
-                # import ipdb;ipdb.set_trace()
-                if prompt_embeds.ndim == 3:
-                    prompt_embeds = prompt_embeds.unsqueeze(1)  # b l d -> b 1 l d
-                if prompt_attention_mask.ndim == 2:
-                    prompt_attention_mask = prompt_attention_mask.unsqueeze(1)  # b l -> b 1 l
-                # prepare attention_mask.
-                # b c t h w -> b t h w
-                attention_mask = torch.ones_like(latent_model_input)[:, 0]
-                # predict noise model_output
-                noise_pred = self.transformer(
-                    latent_model_input,
-                    attention_mask=attention_mask, 
-                    encoder_hidden_states=prompt_embeds,
-                    encoder_attention_mask=prompt_attention_mask,
-                    timestep=current_timestep,
-                    added_cond_kwargs=added_cond_kwargs,
-                    return_dict=False,
-                )[0]
-
-                # perform guidance
-                if do_classifier_free_guidance:
-                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
-                    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
-
-                # learned sigma
-                if self.transformer.config.out_channels // 2 == latent_channels:
-                    noise_pred = noise_pred.chunk(2, dim=1)[0]
-                else:
-                    noise_pred = noise_pred
-
-                # compute previous image: x_t -> x_t-1
-                latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
-                # print(f'latents_{i}_{t}', torch.max(latents), torch.min(latents), torch.mean(latents), torch.std(latents))
-                # call the callback, if provided
-                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
-                    progress_bar.update()
-                    if callback is not None and i % callback_steps == 0:
-                        step_idx = i // getattr(self.scheduler, "order", 1)
-                        callback(step_idx, t, latents)
-        # import ipdb;ipdb.set_trace()
-        # latents = latents.squeeze(2)
-        if not output_type == "latent":
-            # b t h w c
-            image = self.decode_latents(latents)
-            image = image[:, :num_frames, :height, :width]
-        else:
-            image = latents
-
-        # Offload all models
-        self.maybe_free_model_hooks()
-
-        if not return_dict:
-            return (image,)
-
-        return ImagePipelineOutput(images=image)
-
-
-    def decode_latents(self, latents):
-        # print(f'before vae decode', torch.max(latents).item(), torch.min(latents).item(), torch.mean(latents).item(), torch.std(latents).item())
-        video = self.vae.decode(latents.to(self.vae.vae.dtype))
-        # print(f'after vae decode', torch.max(video).item(), torch.min(video).item(), torch.mean(video).item(), torch.std(video).item())
-        video = ((video / 2.0 + 0.5).clamp(0, 1) * 255).to(dtype=torch.uint8).cpu().permute(0, 1, 3, 4, 2).contiguous() # b t h w c
-        # we always cast to float32 as this does not cause significant overhead and is compatible with bfloa16
-        return video

opensora/sample/sample_t2v.py

@@ -1,224 +0,0 @@
-import math
-import os
-import torch
-import argparse
-import torchvision
-
-from diffusers.schedulers import (DDIMScheduler, DDPMScheduler, PNDMScheduler,
-                                  EulerDiscreteScheduler, DPMSolverMultistepScheduler,
-                                  HeunDiscreteScheduler, EulerAncestralDiscreteScheduler,
-                                  DEISMultistepScheduler, KDPM2AncestralDiscreteScheduler)
-from diffusers.schedulers.scheduling_dpmsolver_singlestep import DPMSolverSinglestepScheduler
-from torchvision.utils import save_image
-from transformers import  MT5EncoderModel,  AutoTokenizer
-
-
-from opensora.models.causalvideovae import ae_stride_config, CausalVAEModelWrapper
-from opensora.models.diffusion.opensora.modeling_opensora import OpenSoraT2V
-from opensora.utils.utils import save_video_grid
-
-from opensora.sample.pipeline_opensora import OpenSoraPipeline
-
-import imageio
-
-
-def main(args):
-    # torch.manual_seed(args.seed)
-    # torch.backends.cuda.matmul.allow_tf32 = False
-    weight_dtype = torch.bfloat16
-    device = torch.device(args.device)
-
-    # vae = getae_wrapper(args.ae)(args.model_path, subfolder="vae", cache_dir=args.cache_dir)
-    vae = CausalVAEModelWrapper(args.ae_path)
-    vae.vae = vae.vae.to(device=device, dtype=weight_dtype)
-    if args.enable_tiling:
-        vae.vae.enable_tiling()
-        vae.vae.tile_overlap_factor = args.tile_overlap_factor
-        vae.vae.tile_sample_min_size = 512
-        vae.vae.tile_latent_min_size = 64
-        vae.vae.tile_sample_min_size_t = 29
-        vae.vae.tile_latent_min_size_t = 8
-        if args.save_memory:
-            vae.vae.tile_sample_min_size = 256
-            vae.vae.tile_latent_min_size = 32
-            vae.vae.tile_sample_min_size_t = 29
-            vae.vae.tile_latent_min_size_t = 8
-    vae.vae_scale_factor = ae_stride_config[args.ae]
-
-    # if args.model_3d:
-    #     transformer_model = OpenSoraT2V.from_pretrained(args.model_path, subfolder=args.version, cache_dir=args.cache_dir, low_cpu_mem_usage=False, device_map=None, torch_dtype=weight_dtype)
-    # else:
-    #     transformer_model = LatteT2V.from_pretrained(args.model_path, subfolder=args.version, cache_dir=args.cache_dir, low_cpu_mem_usage=False, device_map=None, torch_dtype=weight_dtype)
-    
-    if args.model_type == 'dit':
-        transformer_model = OpenSoraT2V.from_pretrained(args.model_path, cache_dir=args.cache_dir, 
-                                                        low_cpu_mem_usage=False, device_map=None, torch_dtype=weight_dtype)
-        
-    text_encoder = MT5EncoderModel.from_pretrained(args.text_encoder_name, cache_dir=args.cache_dir, low_cpu_mem_usage=True, torch_dtype=weight_dtype)
-    tokenizer = AutoTokenizer.from_pretrained(args.text_encoder_name, cache_dir=args.cache_dir)
-    # text_encoder = MT5EncoderModel.from_pretrained("/storage/ongoing/new/Open-Sora-Plan/cache_dir/mt5-xxl", cache_dir=args.cache_dir, low_cpu_mem_usage=True, torch_dtype=weight_dtype)
-    # tokenizer = AutoTokenizer.from_pretrained("/storage/ongoing/new/Open-Sora-Plan/cache_dir/mt5-xxl", cache_dir=args.cache_dir)
-    
-    
-    # set eval mode
-    transformer_model.eval()
-    vae.eval()
-    text_encoder.eval()
-
-    if args.sample_method == 'DDIM':  #########
-        scheduler = DDIMScheduler(clip_sample=False)
-    elif args.sample_method == 'EulerDiscrete':
-        scheduler = EulerDiscreteScheduler()
-    elif args.sample_method == 'DDPM':  #############
-        scheduler = DDPMScheduler(clip_sample=False)
-    elif args.sample_method == 'DPMSolverMultistep':
-        '''
-        DPM++ 2M	        DPMSolverMultistepScheduler	
-        DPM++ 2M Karras	    DPMSolverMultistepScheduler	init with use_karras_sigmas=True
-        DPM++ 2M SDE	    DPMSolverMultistepScheduler	init with algorithm_type="sde-dpmsolver++"
-        DPM++ 2M SDE Karras	DPMSolverMultistepScheduler	init with use_karras_sigmas=True and algorithm_type="sde-dpmsolver++"
-        
-        DPM++ SDE	        DPMSolverSinglestepScheduler	
-        DPM++ SDE Karras	DPMSolverSinglestepScheduler	init with use_karras_sigmas=True
-        DPM2	            KDPM2DiscreteScheduler	
-        DPM2 Karras	        KDPM2DiscreteScheduler	init with use_karras_sigmas=True
-        DPM2 a	            KDPM2AncestralDiscreteScheduler	
-        DPM2 a Karras	    KDPM2AncestralDiscreteScheduler	init with use_karras_sigmas=True
-        '''
-        # scheduler = DPMSolverMultistepScheduler(use_karras_sigmas=True)
-        scheduler = DPMSolverMultistepScheduler()
-    elif args.sample_method == 'DPMSolverSinglestep':
-        scheduler = DPMSolverSinglestepScheduler()
-    elif args.sample_method == 'PNDM':
-        scheduler = PNDMScheduler()
-    elif args.sample_method == 'HeunDiscrete':  ########
-        scheduler = HeunDiscreteScheduler()
-    elif args.sample_method == 'EulerAncestralDiscrete':
-        scheduler = EulerAncestralDiscreteScheduler()
-    elif args.sample_method == 'DEISMultistep':
-        scheduler = DEISMultistepScheduler()
-    elif args.sample_method == 'KDPM2AncestralDiscrete':  #########
-        scheduler = KDPM2AncestralDiscreteScheduler()
-    elif args.sample_method == 'EulerDiscreteSVD':
-        scheduler = EulerDiscreteScheduler.from_pretrained("stabilityai/stable-video-diffusion-img2vid", 
-                                                        subfolder="scheduler", cache_dir=args.cache_dir)
-    pipeline = OpenSoraPipeline(vae=vae,
-                                text_encoder=text_encoder,
-                                tokenizer=tokenizer,
-                                scheduler=scheduler,
-                                transformer=transformer_model)
-    pipeline.to(device)
-    if args.compile:
-        # 5%  https://github.com/siliconflow/onediff/tree/main/src/onediff/infer_compiler/backends/nexfort
-        options = '{"mode": "max-optimize:max-autotune:freezing:benchmark:low-precision",             \
-                    "memory_format": "channels_last", "options": {"inductor.optimize_linear_epilogue": false, \
-                    "triton.fuse_attention_allow_fp16_reduction": false}}'
-        # options = '{"mode": "max-autotune", "memory_format": "channels_last",              \
-        #             "options": {"inductor.optimize_linear_epilogue": false, "triton.fuse_attention_allow_fp16_reduction": false}}'
-        from onediffx import compile_pipe
-        pipeline = compile_pipe(
-                pipeline, backend="nexfort", options=options, fuse_qkv_projections=True
-            )
-
-        # 4%
-        # pipeline.transformer = torch.compile(pipeline.transformer)
-    
-    if not os.path.exists(args.save_img_path):
-        os.makedirs(args.save_img_path)
-        
-    if not isinstance(args.text_prompt, list):
-        args.text_prompt = [args.text_prompt]
-    if len(args.text_prompt) == 1 and args.text_prompt[0].endswith('txt'):
-        text_prompt = open(args.text_prompt[0], 'r').readlines()
-        text_prompt = [i.strip() for i in text_prompt]
-
-    positive_prompt = """
-    (masterpiece), (best quality), (ultra-detailed), 
-    {}. 
-    """
-    
-    negative_prompt = """
-    nsfw, lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, 
-    low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry.
-    """
-
-    # positive_prompt = """
-    # (masterpiece), (best quality), (ultra-detailed), 
-    # {}. 
-    # emotional, harmonious, vignette, 4k epic detailed, shot on kodak, 35mm photo, 
-    # sharp focus, high budget, cinemascope, moody, epic, gorgeous
-    # """
-    
-    # negative_prompt = """
-    # disfigured, poorly drawn face, longbody, lowres, bad anatomy, bad hands, missing fingers, cropped, worst quality, low quality
-    # """
-
-    video_grids = []
-    for idx, prompt in enumerate(text_prompt):
-        videos = pipeline(positive_prompt.format(prompt),
-                          negative_prompt=negative_prompt, 
-                          num_frames=args.num_frames,
-                          height=args.height,
-                          width=args.width,
-                          num_inference_steps=args.num_sampling_steps,
-                          guidance_scale=args.guidance_scale,
-                          num_images_per_prompt=1,
-                          mask_feature=True,
-                          device=args.device, 
-                          max_sequence_length=args.max_sequence_length, 
-                          ).images
-        try:
-            if args.num_frames == 1:
-                ext = 'jpg'
-                videos = videos[:, 0].permute(0, 3, 1, 2)  # b t h w c -> b c h w
-                save_image(videos / 255.0, os.path.join(args.save_img_path, f'{idx}.{ext}'), nrow=1, normalize=True, value_range=(0, 1))  # t c h w
-
-            else:
-                ext = 'mp4'
-                imageio.mimwrite(
-                    os.path.join(args.save_img_path, f'{idx}.{ext}'), videos[0], fps=args.fps, quality=6)  # highest quality is 10, lowest is 0
-        except:
-            print('Error when saving {}'.format(prompt))
-        video_grids.append(videos)
-    video_grids = torch.cat(video_grids, dim=0)
-
-
-    # torchvision.io.write_video(args.save_img_path + '_%04d' % args.run_time + '-.mp4', video_grids, fps=6)
-    if args.num_frames == 1:
-        save_image(video_grids / 255.0, os.path.join(args.save_img_path, f'{args.sample_method}_gs{args.guidance_scale}_s{args.num_sampling_steps}.{ext}'),
-                   nrow=math.ceil(math.sqrt(len(video_grids))), normalize=True, value_range=(0, 1))
-    else:
-        video_grids = save_video_grid(video_grids)
-        imageio.mimwrite(os.path.join(args.save_img_path, f'{args.sample_method}_gs{args.guidance_scale}_s{args.num_sampling_steps}.{ext}'), video_grids, fps=args.fps, quality=6)
-
-    print('save path {}'.format(args.save_img_path))
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--model_path", type=str, default='LanguageBind/Open-Sora-Plan-v1.0.0')
-    parser.add_argument("--version", type=str, default=None, choices=[None, '65x512x512', '65x256x256', '17x256x256'])
-    parser.add_argument("--num_frames", type=int, default=1)
-    parser.add_argument("--height", type=int, default=512)
-    parser.add_argument("--width", type=int, default=512)
-    parser.add_argument("--device", type=str, default='cuda:0')
-    parser.add_argument("--cache_dir", type=str, default='./cache_dir')
-    parser.add_argument("--ae", type=str, default='CausalVAEModel_4x8x8')
-    parser.add_argument("--ae_path", type=str, default='CausalVAEModel_4x8x8')
-    parser.add_argument("--text_encoder_name", type=str, default='DeepFloyd/t5-v1_1-xxl')
-    parser.add_argument("--save_img_path", type=str, default="./sample_videos/t2v")
-    parser.add_argument("--guidance_scale", type=float, default=2.5)
-    parser.add_argument("--sample_method", type=str, default="PNDM")
-    parser.add_argument("--max_sequence_length", type=int, default=300)
-    parser.add_argument("--num_sampling_steps", type=int, default=50)
-    parser.add_argument("--fps", type=int, default=24)
-    parser.add_argument("--run_time", type=int, default=0)
-    parser.add_argument("--text_prompt", nargs='+')
-    parser.add_argument('--tile_overlap_factor', type=float, default=0.125)
-    parser.add_argument('--enable_tiling', action='store_true')
-    parser.add_argument('--compile', action='store_true')
-    parser.add_argument('--model_type', type=str, default="dit", choices=['dit'])
-    parser.add_argument('--save_memory', action='store_true')
-    args = parser.parse_args()
-
-    main(args)

opensora/sample/sample_t2v_mochi.py

@@ -1,37 +0,0 @@
-import torch
-from diffusers import MochiPipeline
-
-from diffusers.utils import export_to_video, load_image, load_video
-
-def main():
-    # Set the random seed for reproducibility
-    generator = torch.Generator("cpu").manual_seed(12345)
-
-    # Load the Mochi pipeline
-    model_path = "data/mochi"
-    prompt = "A hand with delicate fingers picks up a bright yellow lemon from a wooden bowl filled with lemons and sprigs of mint against a peach-colored background. The hand gently tosses the lemon up and catches it, showcasing its smooth texture. A beige string bag sits beside the bowl, adding a rustic touch to the scene. Additional lemons, one halved, are scattered around the base of the bowl. The even lighting enhances the vibrant colors and creates a fresh, inviting atmosphere."
-    num_frames = 163
-    height = 480
-    width = 848
-    num_inference_steps = 64
-
-    pipe = MochiPipeline.from_pretrained(model_path, torch_dtype=torch.float16)
-    pipe.enable_vae_tiling()
-    pipe.enable_model_cpu_offload()
-
-    # Generate videos from the input prompt
-    video = pipe(
-        prompt=prompt,
-        height=height,
-        width=width,
-        num_frames=num_frames,
-        generator=generator,
-        num_inference_steps=num_inference_steps,
-        guidance_scale=4.5
-    ).frames[0]
-
-
-    export_to_video(video, './outputs.mp4', fps=30)
-
-if __name__ == "__main__":
-    main()

opensora/sample/sample_t2v_sp.py

@@ -1,274 +0,0 @@
-import math
-import os
-import torch
-import argparse
-import torchvision
-import torch.distributed as dist
-
-from diffusers.schedulers import (DDIMScheduler, DDPMScheduler, PNDMScheduler,
-                                  EulerDiscreteScheduler, DPMSolverMultistepScheduler,
-                                  HeunDiscreteScheduler, EulerAncestralDiscreteScheduler,
-                                  DEISMultistepScheduler, KDPM2AncestralDiscreteScheduler)
-from diffusers.schedulers.scheduling_dpmsolver_singlestep import DPMSolverSinglestepScheduler
-from diffusers.models import AutoencoderKL, AutoencoderKLTemporalDecoder, Transformer2DModel
-from omegaconf import OmegaConf
-from torchvision.utils import save_image
-from transformers import T5EncoderModel, T5Tokenizer, AutoTokenizer, MT5EncoderModel
-
-import os, sys
-
-from opensora.models.causalvideovae import ae_stride_config, ae_channel_config, ae_norm, ae_denorm, CausalVAEModelWrapper
-from opensora.models.diffusion.udit.modeling_udit import UDiTT2V
-from opensora.models.diffusion.opensora.modeling_opensora import OpenSoraT2V
-
-from opensora.models.text_encoder import get_text_enc
-from opensora.utils.utils import save_video_grid
-
-from opensora.sample.pipeline_opensora_sp import OpenSoraPipeline
-
-import imageio
-from opensora.utils.parallel_states import initialize_sequence_parallel_state, nccl_info
-import time
-
-
-
-def load_t2v_checkpoint(model_path):
-    print('load_t2v_checkpoint, ', model_path)
-    if args.model_type == 'udit':
-        transformer_model = UDiTT2V.from_pretrained(model_path, cache_dir=args.cache_dir,
-                                                        low_cpu_mem_usage=False, device_map=None,
-                                                        torch_dtype=weight_dtype)
-    elif args.model_type == 'dit':
-        transformer_model = OpenSoraT2V.from_pretrained(model_path, cache_dir=args.cache_dir,
-                                                        low_cpu_mem_usage=False, device_map=None,
-                                                        torch_dtype=weight_dtype)
-    else:
-        transformer_model = LatteT2V.from_pretrained(model_path, cache_dir=args.cache_dir, low_cpu_mem_usage=False,
-                                                     device_map=None, torch_dtype=weight_dtype)
-    # print(transformer_model.config)
-
-    # set eval mode
-    transformer_model.eval()
-    pipeline = OpenSoraPipeline(vae=vae,
-                                text_encoder=text_encoder,
-                                tokenizer=tokenizer,
-                                scheduler=scheduler,
-                                transformer=transformer_model).to(device)
-
-
-    if args.compile:
-        # 5%  https://github.com/siliconflow/onediff/tree/main/src/onediff/infer_compiler/backends/nexfort
-        # options = '{"mode": "max-optimize:max-autotune:freezing:benchmark:low-precision",             \
-        #             "memory_format": "channels_last", "options": {"inductor.optimize_linear_epilogue": false, \
-        #             "triton.fuse_attention_allow_fp16_reduction": false}}'
-        # # options = '{"mode": "max-autotune", "memory_format": "channels_last",              \
-        # #             "options": {"inductor.optimize_linear_epilogue": false, "triton.fuse_attention_allow_fp16_reduction": false}}'
-        # from onediffx import compile_pipe
-        # pipeline = compile_pipe(
-        #         pipeline, backend="nexfort", options=options, fuse_qkv_projections=True
-        #     )
-
-        # 4%
-        pipeline.transformer = torch.compile(pipeline.transformer)
-    return pipeline
-
-
-def get_latest_path():
-    # Get the most recent checkpoint
-    dirs = os.listdir(args.model_path)
-    dirs = [d for d in dirs if d.startswith("checkpoint")]
-    dirs = sorted(dirs, key=lambda x: int(x.split("-")[1]))
-    path = dirs[-1] if len(dirs) > 0 else None
-
-    return path
-
-
-def run_model_and_save_images(pipeline, model_path):
-    video_grids = []
-    if not isinstance(args.text_prompt, list):
-        args.text_prompt = [args.text_prompt]
-    if len(args.text_prompt) == 1 and args.text_prompt[0].endswith('txt'):
-        text_prompt = open(args.text_prompt[0], 'r').readlines()
-        args.text_prompt = [i.strip() for i in text_prompt]
-
-    checkpoint_name = f"{os.path.basename(model_path)}"
-    positive_prompt = """
-    (masterpiece), (best quality), (ultra-detailed), (unwatermarked), 
-    {}. 
-    emotional, harmonious, vignette, 4k epic detailed, shot on kodak, 35mm photo, 
-    sharp focus, high budget, cinemascope, moody, epic, gorgeous
-    """
-    
-    negative_prompt = """
-    nsfw, lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, 
-    low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry.
-    """
-
-    for index, prompt in enumerate(args.text_prompt):
-        # if index % world_size != local_rank:
-        #     continue
-        # print('Processing the ({}) prompt, device ({})'.format(prompt, device))
-        videos = pipeline(positive_prompt.format(prompt),
-                        negative_prompt=negative_prompt, 
-                        num_frames=args.num_frames,
-                        height=args.height,
-                        width=args.width,
-                        num_inference_steps=args.num_sampling_steps,
-                        guidance_scale=args.guidance_scale,
-                        num_images_per_prompt=1,
-                        mask_feature=True,
-                        device=args.device,
-                        max_sequence_length=args.max_sequence_length,
-                        ).images
-        print(videos.shape)
-        
-        if nccl_info.rank <= 0:
-            try:
-                if args.num_frames == 1:
-                    videos = videos[:, 0].permute(0, 3, 1, 2)  # b t h w c -> b c h w
-                    save_image(videos / 255.0, os.path.join(args.save_img_path,
-                                                            f'{args.sample_method}_{index}_{checkpoint_name}_gs{args.guidance_scale}_s{args.num_sampling_steps}.{ext}'),
-                            nrow=1, normalize=True, value_range=(0, 1))  # t c h w
-
-                else:
-                    imageio.mimwrite(
-                        os.path.join(
-                            args.save_img_path,
-                            f'{args.sample_method}_{index}_{checkpoint_name}__gs{args.guidance_scale}_s{args.num_sampling_steps}.{ext}'
-                        ), videos[0],
-                        fps=args.fps, quality=6, codec='libx264',
-                        output_params=['-threads', '20'])  # highest quality is 10, lowest is 0
-            except:
-                print('Error when saving {}'.format(prompt))
-            video_grids.append(videos)
-    if nccl_info.rank <= 0:
-        video_grids = torch.cat(video_grids, dim=0)
-
-        def get_file_name():
-            return os.path.join(args.save_img_path,
-                                f'{args.sample_method}_gs{args.guidance_scale}_s{args.num_sampling_steps}_{checkpoint_name}.{ext}')
-
-        if args.num_frames == 1:
-            save_image(video_grids / 255.0, get_file_name(),
-                    nrow=math.ceil(math.sqrt(len(video_grids))), normalize=True, value_range=(0, 1))
-        else:
-            video_grids = save_video_grid(video_grids)
-            imageio.mimwrite(get_file_name(), video_grids, fps=args.fps, quality=6)
-
-        print('save path {}'.format(args.save_img_path))
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--model_path", type=str, default='LanguageBind/Open-Sora-Plan-v1.0.0')
-    parser.add_argument("--version", type=str, default=None, choices=[None, '65x512x512', '65x256x256', '17x256x256'])
-    parser.add_argument("--num_frames", type=int, default=1)
-    parser.add_argument("--height", type=int, default=512)
-    parser.add_argument("--width", type=int, default=512)
-    parser.add_argument("--device", type=str, default='cuda:0')
-    parser.add_argument("--cache_dir", type=str, default='./cache_dir')
-    parser.add_argument("--ae", type=str, default='CausalVAEModel_4x8x8')
-    parser.add_argument("--ae_path", type=str, default='CausalVAEModel_4x8x8')
-    parser.add_argument("--text_encoder_name", type=str, default='DeepFloyd/t5-v1_1-xxl')
-    parser.add_argument("--save_img_path", type=str, default="./sample_videos/t2v")
-    parser.add_argument("--guidance_scale", type=float, default=7.5)
-    parser.add_argument("--sample_method", type=str, default="PNDM")
-    parser.add_argument("--num_sampling_steps", type=int, default=50)
-    parser.add_argument("--fps", type=int, default=24)
-    parser.add_argument("--max_sequence_length", type=int, default=512)
-    parser.add_argument("--text_prompt", nargs='+')
-    parser.add_argument('--tile_overlap_factor', type=float, default=0.25)
-    parser.add_argument('--enable_tiling', action='store_true')
-    parser.add_argument('--model_type', type=str, default="dit", choices=['dit', 'udit', 'latte'])
-    parser.add_argument('--compile', action='store_true')
-    parser.add_argument('--save_memory', action='store_true')
-    args = parser.parse_args()
-
-
-    # 初始化分布式环境
-    local_rank = int(os.getenv('RANK', 0))
-    world_size = int(os.getenv('WORLD_SIZE', 1))
-    print('world_size', world_size)
-
-    torch.cuda.set_device(local_rank)
-    dist.init_process_group(backend='nccl', init_method='env://', world_size=world_size, rank=local_rank)
-    initialize_sequence_parallel_state(world_size)
-    # torch.manual_seed(args.seed)
-    weight_dtype = torch.bfloat16
-    device = torch.cuda.current_device()
-    # print(11111111111111111111, local_rank, device)
-    # vae = getae_wrapper(args.ae)(args.model_path, subfolder="vae", cache_dir=args.cache_dir)
-    vae = CausalVAEModelWrapper(args.ae_path)
-    vae.vae = vae.vae.to(device=device, dtype=weight_dtype)
-    if args.enable_tiling:
-        vae.vae.enable_tiling()
-        vae.vae.tile_overlap_factor = args.tile_overlap_factor
-        vae.vae.tile_sample_min_size = 512
-        vae.vae.tile_latent_min_size = 64
-        vae.vae.tile_sample_min_size_t = 29
-        vae.vae.tile_latent_min_size_t = 8
-        if args.save_memory:
-            vae.vae.tile_sample_min_size = 256
-            vae.vae.tile_latent_min_size = 32
-            vae.vae.tile_sample_min_size_t = 29
-            vae.vae.tile_latent_min_size_t = 8
-
-    vae.vae_scale_factor = ae_stride_config[args.ae]
-
-    # text_encoder = MT5EncoderModel.from_pretrained("/storage/ongoing/new/Open-Sora-Plan/cache_dir/mt5-xxl", 
-    #                                                cache_dir=args.cache_dir, low_cpu_mem_usage=True, 
-    #                                                torch_dtype=weight_dtype).to(device)
-    # tokenizer = AutoTokenizer.from_pretrained("/storage/ongoing/new/Open-Sora-Plan/cache_dir/mt5-xxl", 
-    #                                           cache_dir=args.cache_dir)
-    
-    text_encoder = MT5EncoderModel.from_pretrained(args.text_encoder_name, cache_dir=args.cache_dir,
-                                                  low_cpu_mem_usage=True, torch_dtype=weight_dtype).to(device)
-    tokenizer = T5Tokenizer.from_pretrained(args.text_encoder_name, cache_dir=args.cache_dir)
-
-    # set eval mode
-    vae.eval()
-    text_encoder.eval()
-
-    if args.sample_method == 'DDIM':  #########
-        scheduler = DDIMScheduler(clip_sample=False)
-    elif args.sample_method == 'EulerDiscrete':
-        scheduler = EulerDiscreteScheduler()
-    elif args.sample_method == 'DDPM':  #############
-        scheduler = DDPMScheduler(clip_sample=False)
-    elif args.sample_method == 'DPMSolverMultistep':
-        scheduler = DPMSolverMultistepScheduler()
-    elif args.sample_method == 'DPMSolverSinglestep':
-        scheduler = DPMSolverSinglestepScheduler()
-    elif args.sample_method == 'PNDM':
-        scheduler = PNDMScheduler()
-    elif args.sample_method == 'HeunDiscrete':  ########
-        scheduler = HeunDiscreteScheduler()
-    elif args.sample_method == 'EulerAncestralDiscrete':
-        scheduler = EulerAncestralDiscreteScheduler()
-    elif args.sample_method == 'DEISMultistep':
-        scheduler = DEISMultistepScheduler()
-    elif args.sample_method == 'KDPM2AncestralDiscrete':  #########
-        scheduler = KDPM2AncestralDiscreteScheduler()
-
-    if not os.path.exists(args.save_img_path):
-        os.makedirs(args.save_img_path, exist_ok=True)
-
-    if args.num_frames == 1:
-        video_length = 1
-        ext = 'jpg'
-    else:
-        ext = 'mp4'
-
-    latest_path = None
-    save_img_path = args.save_img_path
-
-
-    if latest_path is None:
-        latest_path = ''
-    # full_path = f"{args.model_path}/{latest_path}/model_ema"
-    full_path = f"{args.model_path}"
-    # full_path = f"{args.model_path}/{latest_path}/model"
-    pipeline = load_t2v_checkpoint(full_path)
-    print('load model')
-    
-    run_model_and_save_images(pipeline, latest_path)

opensora/sample/vae_reconstruct_video.py

@@ -1,120 +0,0 @@
-import math
-import random
-import argparse
-from typing import Optional
-
-import cv2
-import numpy as np
-import numpy.typing as npt
-import torch
-from decord import VideoReader, cpu
-from torch.nn import functional as F
-from torchvision.transforms import Lambda, Compose
-from opensora.models import CausalVAEModelWrapper
-from opensora.dataset.transform import Normalize255, CenterCropResizeVideo
-
-
-def array_to_video(image_array: npt.NDArray, fps: float = 30.0, output_file: str = 'output_video.mp4') -> None:
-    height, width, channels = image_array[0].shape
-    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-    video_writer = cv2.VideoWriter(output_file, fourcc, float(fps), (width, height))
-
-    for image in image_array:
-        image_rgb = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
-        video_writer.write(image_rgb)
-
-    video_writer.release()
-
-
-def custom_to_video(x: torch.Tensor, fps: float = 2.0, output_file: str = 'output_video.mp4') -> None:
-    x = x.detach().cpu()
-    x = torch.clamp(x, -1, 1)
-    x = (x + 1) / 2
-    x = x.permute(0, 2, 3, 1).numpy()
-    x = (255 * x).astype(np.uint8)
-    array_to_video(x, fps=fps, output_file=output_file)
-    return
-
-
-def read_video(video_path: str, decode_frames: int) -> torch.Tensor:
-    decord_vr = VideoReader(video_path, ctx=cpu(0))
-    total_frames = len(decord_vr)
-    if total_frames < decode_frames:
-        decode_frames = total_frames
-        print(f"Video has less frames than decode_frames. Setting decode_frames to {total_frames}")
-    s = random.randint(0, total_frames - decode_frames)
-    e = s + decode_frames
-    
-    frame_id_list = np.arange(s, e)
-    video_data = decord_vr.get_batch(frame_id_list).asnumpy()
-    video_data = torch.from_numpy(video_data)
-    video_data = video_data.permute(3, 0, 1, 2)  # (T, H, W, C) -> (C, T, H, W)
-    return video_data
-
-
-def preprocess(video_data: torch.Tensor, height: int = 128, width: int = 128) -> torch.Tensor:
-    transform = Compose(
-        [
-            Normalize255(),
-            CenterCropResizeVideo((height, width)),
-            Lambda(lambda x: 2. * x - 1.)
-        ]
-    )
-
-    video_outputs = transform(video_data)
-    video_outputs = torch.unsqueeze(video_outputs, 0)
-
-    return video_outputs
-
-
-def main(args: argparse.Namespace):
-    device = args.device
-    kwarg = {}
-    # vae = getae_wrapper(args.ae)(args.model_path, subfolder="vae", cache_dir='cache_dir', **kwarg).to(device)
-    vae = CausalVAEModelWrapper(args.ae_path, **kwarg).to(device)
-    if args.enable_tiling:
-        vae.vae.enable_tiling()
-        vae.vae.tile_overlap_factor = args.tile_overlap_factor
-        vae.vae.tile_overlap_factor = args.tile_overlap_factor
-        if args.save_memory:
-            vae.vae.tile_sample_min_size = 256
-            vae.vae.tile_latent_min_size = 32
-            vae.vae.tile_sample_min_size_t = 29
-            vae.vae.tile_latent_min_size_t = 8
-    vae.eval()
-    vae = vae.to(device)
-    vae = vae.half()
-
-    with torch.no_grad():
-        x_vae = preprocess(read_video(args.video_path, args.decode_frames), args.height,
-                           args.width)
-        print(x_vae.shape)
-        x_vae = x_vae.to(device, dtype=torch.float16)  # b c t h w
-        latents = vae.encode(x_vae)
-        print(latents.shape)
-        latents = latents.to(torch.float16)
-        video_recon = vae.decode(latents)  # b t c h w
-        print(video_recon.shape)
-
-    custom_to_video(video_recon[0], fps=args.fps, output_file=args.rec_path)
-
-
-if __name__ == '__main__':
-    parser = argparse.ArgumentParser()
-    parser.add_argument('--video_path', type=str, default='')
-    parser.add_argument('--rec_path', type=str, default='')
-    parser.add_argument('--ae', type=str, default='')
-    parser.add_argument('--ae_path', type=str, default='')
-    parser.add_argument('--model_path', type=str, default='results/pretrained')
-    parser.add_argument('--fps', type=int, default=30)
-    parser.add_argument('--height', type=int, default=336)
-    parser.add_argument('--width', type=int, default=336)
-    parser.add_argument('--decode_frames', type=int, default=100)
-    parser.add_argument('--device', type=str, default="cuda")
-    parser.add_argument('--tile_overlap_factor', type=float, default=0.25)
-    parser.add_argument('--tile_sample_min_size', type=int, default=256)
-    parser.add_argument('--enable_tiling', action='store_true')
-    parser.add_argument('--save_memory', action='store_true')
-
-    args = parser.parse_args()
-    main(args)