diff --git a/.gitignore b/.gitignore index 561f7c7c7..47053098a 100644 --- a/.gitignore +++ b/.gitignore @@ -7,7 +7,7 @@ UCF-101/ results/ vae build/ -opensora.egg-info/ +fastvideo.egg-info/ wandb/ .idea *.ipynb diff --git a/README.md b/README.md index 4f46973fd..f3b5f6d2c 100644 --- a/README.md +++ b/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 diff --git a/opensora/dataset/__init__.py b/fastvideo/dataset/__init__.py similarity index 79% rename from opensora/dataset/__init__.py rename to fastvideo/dataset/__init__.py index bb6654787..aaeda7f94 100644 --- a/opensora/dataset/__init__.py +++ b/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, diff --git a/fastvideo/dataset/latent_datasets.py b/fastvideo/dataset/latent_datasets.py new file mode 100644 index 000000000..893e3261f --- /dev/null +++ b/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() \ No newline at end of file diff --git a/opensora/dataset/t2v_datasets.py b/fastvideo/dataset/t2v_datasets.py similarity index 97% rename from opensora/dataset/t2v_datasets.py rename to fastvideo/dataset/t2v_datasets.py index da28fcbd9..5d29517b5 100644 --- a/opensora/dataset/t2v_datasets.py +++ b/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, diff --git a/opensora/dataset/transform.py b/fastvideo/dataset/transform.py similarity index 99% rename from opensora/dataset/transform.py rename to fastvideo/dataset/transform.py index ca68d085d..df19d603a 100644 --- a/opensora/dataset/transform.py +++ b/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)) \ No newline at end of file diff --git a/fastvideo/sample/generate_synthetic.py b/fastvideo/sample/generate_synthetic.py new file mode 100644 index 000000000..076e0efb5 --- /dev/null +++ b/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) + + \ No newline at end of file diff --git a/fastvideo/sample/monkey_patch/sp_attention.py b/fastvideo/sample/monkey_patch/sp_attention.py new file mode 100644 index 000000000..ed6e6c4ed --- /dev/null +++ b/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 + diff --git a/fastvideo/sample/monkey_patch/sp_rope.py b/fastvideo/sample/monkey_patch/sp_rope.py new file mode 100644 index 000000000..9c94558e4 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/synthetic_generation/mochi_pipe.py b/fastvideo/sample/pipeline_mochi.py similarity index 97% rename from synthetic_generation/mochi_pipe.py rename to fastvideo/sample/pipeline_mochi.py index 3556026c5..1c9ce0299 100644 --- a/synthetic_generation/mochi_pipe.py +++ b/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) \ No newline at end of file + return MochiPipelineOutput(frames=video) diff --git a/fastvideo/sample/sample_t2v_mochi.py b/fastvideo/sample/sample_t2v_mochi.py new file mode 100644 index 000000000..48283ce6e --- /dev/null +++ b/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) diff --git a/fastvideo/sample/sample_t2v_mochi_sp.py b/fastvideo/sample/sample_t2v_mochi_sp.py new file mode 100644 index 000000000..d894b170e --- /dev/null +++ b/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) diff --git a/fastvideo/train.py b/fastvideo/train.py new file mode 100644 index 000000000..6cbf3038d --- /dev/null +++ b/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) diff --git a/opensora/utils/communications.py b/fastvideo/utils/communications.py similarity index 69% rename from opensora/utils/communications.py rename to fastvideo/utils/communications.py index c7a79e3a6..ce29b52eb 100644 --- a/opensora/utils/communications.py +++ b/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 \ No newline at end of file + return hidden_states, encoder_hidden_states, attention_mask, encoder_attention_mask \ No newline at end of file diff --git a/opensora/utils/dataset_utils.py b/fastvideo/utils/dataset_utils.py similarity index 100% rename from opensora/utils/dataset_utils.py rename to fastvideo/utils/dataset_utils.py diff --git a/opensora/utils/ema.py b/fastvideo/utils/ema.py similarity index 100% rename from opensora/utils/ema.py rename to fastvideo/utils/ema.py diff --git a/opensora/train/train_t2v_diffusers.py b/fastvideo/utils/encode_vae_t5.py similarity index 87% rename from opensora/train/train_t2v_diffusers.py rename to fastvideo/utils/encode_vae_t5.py index 967642b25..527564262 100644 --- a/opensora/train/train_t2v_diffusers.py +++ b/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.") diff --git a/fastvideo/utils/optimizer.py b/fastvideo/utils/optimizer.py new file mode 100644 index 000000000..e5aad89f8 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/opensora/utils/parallel_states.py b/fastvideo/utils/parallel_states.py similarity index 96% rename from opensora/utils/parallel_states.py rename to fastvideo/utils/parallel_states.py index 7b7465817..6768f5f3c 100644 --- a/opensora/utils/parallel_states.py +++ b/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 diff --git a/opensora/utils/utils.py b/fastvideo/utils/utils.py similarity index 100% rename from opensora/utils/utils.py rename to fastvideo/utils/utils.py diff --git a/generate_synthetic.sh b/generate_synthetic.sh new file mode 100644 index 000000000..3e7c25ec7 --- /dev/null +++ b/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 + + \ No newline at end of file diff --git a/infenrece_mochi.sh b/infenrece_mochi.sh new file mode 100644 index 000000000..98b7cd67b --- /dev/null +++ b/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 diff --git a/infenrece_mochi_sp.sh b/infenrece_mochi_sp.sh new file mode 100644 index 000000000..23e7777d3 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/opensora/__init__.py b/opensora/__init__.py deleted file mode 100644 index 4e768b56d..000000000 --- a/opensora/__init__.py +++ /dev/null @@ -1 +0,0 @@ -# \ No newline at end of file diff --git a/opensora/adaptor/__init__.py b/opensora/adaptor/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/opensora/adaptor/modules.py b/opensora/adaptor/modules.py deleted file mode 100644 index e53599d55..000000000 --- a/opensora/adaptor/modules.py +++ /dev/null @@ -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 diff --git a/opensora/dataset/inpaint_datasets.py b/opensora/dataset/inpaint_datasets.py deleted file mode 100644 index 7f22baa23..000000000 --- a/opensora/dataset/inpaint_datasets.py +++ /dev/null @@ -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) - \ No newline at end of file diff --git a/opensora/eval/cal_flolpips.py b/opensora/eval/cal_flolpips.py deleted file mode 100644 index eaf7ca9e1..000000000 --- a/opensora/eval/cal_flolpips.py +++ /dev/null @@ -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() \ No newline at end of file diff --git a/opensora/eval/cal_fvd.py b/opensora/eval/cal_fvd.py deleted file mode 100644 index 1f1a9806e..000000000 --- a/opensora/eval/cal_fvd.py +++ /dev/null @@ -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() diff --git a/opensora/eval/cal_lpips.py b/opensora/eval/cal_lpips.py deleted file mode 100644 index 10edc610b..000000000 --- a/opensora/eval/cal_lpips.py +++ /dev/null @@ -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() \ No newline at end of file diff --git a/opensora/eval/cal_psnr.py b/opensora/eval/cal_psnr.py deleted file mode 100644 index b325106c6..000000000 --- a/opensora/eval/cal_psnr.py +++ /dev/null @@ -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() \ No newline at end of file diff --git a/opensora/eval/cal_ssim.py b/opensora/eval/cal_ssim.py deleted file mode 100644 index b2de811b6..000000000 --- a/opensora/eval/cal_ssim.py +++ /dev/null @@ -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() \ No newline at end of file diff --git a/opensora/eval/eval_clip_score.py b/opensora/eval/eval_clip_score.py deleted file mode 100755 index ed2f51082..000000000 --- a/opensora/eval/eval_clip_score.py +++ /dev/null @@ -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() diff --git a/opensora/eval/eval_common_metric.py b/opensora/eval/eval_common_metric.py deleted file mode 100644 index 37c24f366..000000000 --- a/opensora/eval/eval_common_metric.py +++ /dev/null @@ -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() diff --git a/opensora/eval/flolpips/correlation/correlation.py b/opensora/eval/flolpips/correlation/correlation.py deleted file mode 100644 index 7c91055ba..000000000 --- a/opensora/eval/flolpips/correlation/correlation.py +++ /dev/null @@ -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 \ No newline at end of file diff --git a/opensora/eval/flolpips/flolpips.py b/opensora/eval/flolpips/flolpips.py deleted file mode 100644 index 66f473149..000000000 --- a/opensora/eval/flolpips/flolpips.py +++ /dev/null @@ -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 \ No newline at end of file diff --git a/opensora/eval/flolpips/pretrained_networks.py b/opensora/eval/flolpips/pretrained_networks.py deleted file mode 100644 index 82f14e206..000000000 --- a/opensora/eval/flolpips/pretrained_networks.py +++ /dev/null @@ -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 diff --git a/opensora/eval/flolpips/pwcnet.py b/opensora/eval/flolpips/pwcnet.py deleted file mode 100644 index 374f18e20..000000000 --- a/opensora/eval/flolpips/pwcnet.py +++ /dev/null @@ -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 \ No newline at end of file diff --git a/opensora/eval/flolpips/utils.py b/opensora/eval/flolpips/utils.py deleted file mode 100644 index 98d4502d5..000000000 --- a/opensora/eval/flolpips/utils.py +++ /dev/null @@ -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 diff --git a/opensora/eval/fvd/styleganv/fvd.py b/opensora/eval/fvd/styleganv/fvd.py deleted file mode 100644 index 3043a2a4a..000000000 --- a/opensora/eval/fvd/styleganv/fvd.py +++ /dev/null @@ -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) \ No newline at end of file diff --git a/opensora/eval/fvd/videogpt/fvd.py b/opensora/eval/fvd/videogpt/fvd.py deleted file mode 100644 index e81c82904..000000000 --- a/opensora/eval/fvd/videogpt/fvd.py +++ /dev/null @@ -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 diff --git a/opensora/eval/fvd/videogpt/pytorch_i3d.py b/opensora/eval/fvd/videogpt/pytorch_i3d.py deleted file mode 100644 index 58a16cdd7..000000000 --- a/opensora/eval/fvd/videogpt/pytorch_i3d.py +++ /dev/null @@ -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) \ No newline at end of file diff --git a/opensora/eval/script/cal_clip_score.sh b/opensora/eval/script/cal_clip_score.sh deleted file mode 100644 index 8bc241c15..000000000 --- a/opensora/eval/script/cal_clip_score.sh +++ /dev/null @@ -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" diff --git a/opensora/eval/script/cal_fvd.sh b/opensora/eval/script/cal_fvd.sh deleted file mode 100644 index e344242a0..000000000 --- a/opensora/eval/script/cal_fvd.sh +++ /dev/null @@ -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' diff --git a/opensora/eval/script/cal_lpips.sh b/opensora/eval/script/cal_lpips.sh deleted file mode 100644 index 9ebaf26e8..000000000 --- a/opensora/eval/script/cal_lpips.sh +++ /dev/null @@ -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' \ No newline at end of file diff --git a/opensora/eval/script/cal_psnr.sh b/opensora/eval/script/cal_psnr.sh deleted file mode 100644 index a60a5d1c1..000000000 --- a/opensora/eval/script/cal_psnr.sh +++ /dev/null @@ -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' \ No newline at end of file diff --git a/opensora/eval/script/cal_ssim.sh b/opensora/eval/script/cal_ssim.sh deleted file mode 100644 index 404d8da9a..000000000 --- a/opensora/eval/script/cal_ssim.sh +++ /dev/null @@ -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' \ No newline at end of file diff --git a/opensora/models/__init__.py b/opensora/models/__init__.py deleted file mode 100644 index 3eca32c73..000000000 --- a/opensora/models/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from .causalvideovae.model import CausalVAEModelWrapper \ No newline at end of file diff --git a/opensora/models/causalvideovae/__init__.py b/opensora/models/causalvideovae/__init__.py deleted file mode 100644 index bc1357393..000000000 --- a/opensora/models/causalvideovae/__init__.py +++ /dev/null @@ -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.), -} \ No newline at end of file diff --git a/opensora/models/causalvideovae/model/__init__.py b/opensora/models/causalvideovae/model/__init__.py deleted file mode 100755 index 7ca9d1546..000000000 --- a/opensora/models/causalvideovae/model/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -from .causal_vae import ( - CausalVAEModel, CausalVAEModelWrapper -) -from .ema_model import EMA \ No newline at end of file diff --git a/opensora/models/causalvideovae/model/causal_vae/__init__.py b/opensora/models/causalvideovae/model/causal_vae/__init__.py deleted file mode 100755 index fe97ed155..000000000 --- a/opensora/models/causalvideovae/model/causal_vae/__init__.py +++ /dev/null @@ -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 \ No newline at end of file diff --git a/opensora/models/causalvideovae/model/causal_vae/modeling_causalvae.py b/opensora/models/causalvideovae/model/causal_vae/modeling_causalvae.py deleted file mode 100755 index 7c38d2748..000000000 --- a/opensora/models/causalvideovae/model/causal_vae/modeling_causalvae.py +++ /dev/null @@ -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}" - - - - diff --git a/opensora/models/causalvideovae/model/configuration_videobase.py b/opensora/models/causalvideovae/model/configuration_videobase.py deleted file mode 100755 index f25a25c5b..000000000 --- a/opensora/models/causalvideovae/model/configuration_videobase.py +++ /dev/null @@ -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) \ No newline at end of file diff --git a/opensora/models/causalvideovae/model/dataset_videobase.py b/opensora/models/causalvideovae/model/dataset_videobase.py deleted file mode 100755 index 564bc1f72..000000000 --- a/opensora/models/causalvideovae/model/dataset_videobase.py +++ /dev/null @@ -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 \ No newline at end of file diff --git a/opensora/models/causalvideovae/model/ema_model.py b/opensora/models/causalvideovae/model/ema_model.py deleted file mode 100755 index c4cf50172..000000000 --- a/opensora/models/causalvideovae/model/ema_model.py +++ /dev/null @@ -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 = {} - - \ No newline at end of file diff --git a/opensora/models/causalvideovae/model/modeling_videobase.py b/opensora/models/causalvideovae/model/modeling_videobase.py deleted file mode 100755 index 260575675..000000000 --- a/opensora/models/causalvideovae/model/modeling_videobase.py +++ /dev/null @@ -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) \ No newline at end of file diff --git a/opensora/models/causalvideovae/model/modules/__init__.py b/opensora/models/causalvideovae/model/modules/__init__.py deleted file mode 100755 index 7ef41edd2..000000000 --- a/opensora/models/causalvideovae/model/modules/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -from .block import Block -from .attention import * -from .conv import * -from .normalize import * -from .resnet_block import * -from .updownsample import * diff --git a/opensora/models/causalvideovae/model/modules/attention.py b/opensora/models/causalvideovae/model/modules/attention.py deleted file mode 100755 index c3b9bb19b..000000000 --- a/opensora/models/causalvideovae/model/modules/attention.py +++ /dev/null @@ -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) \ No newline at end of file diff --git a/opensora/models/causalvideovae/model/modules/block.py b/opensora/models/causalvideovae/model/modules/block.py deleted file mode 100755 index e93672d20..000000000 --- a/opensora/models/causalvideovae/model/modules/block.py +++ /dev/null @@ -1,5 +0,0 @@ -import torch.nn as nn - -class Block(nn.Module): - def __init__(self, *args, **kwargs) -> None: - super().__init__(*args, **kwargs) \ No newline at end of file diff --git a/opensora/models/causalvideovae/model/modules/conv.py b/opensora/models/causalvideovae/model/modules/conv.py deleted file mode 100755 index a4c6b78c0..000000000 --- a/opensora/models/causalvideovae/model/modules/conv.py +++ /dev/null @@ -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) \ No newline at end of file diff --git a/opensora/models/causalvideovae/model/modules/normalize.py b/opensora/models/causalvideovae/model/modules/normalize.py deleted file mode 100755 index 7c8c05f0f..000000000 --- a/opensora/models/causalvideovae/model/modules/normalize.py +++ /dev/null @@ -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 diff --git a/opensora/models/causalvideovae/model/modules/ops.py b/opensora/models/causalvideovae/model/modules/ops.py deleted file mode 100755 index e5987e740..000000000 --- a/opensora/models/causalvideovae/model/modules/ops.py +++ /dev/null @@ -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 \ No newline at end of file diff --git a/opensora/models/causalvideovae/model/modules/quant.py b/opensora/models/causalvideovae/model/modules/quant.py deleted file mode 100755 index bb702cee5..000000000 --- a/opensora/models/causalvideovae/model/modules/quant.py +++ /dev/null @@ -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 \ No newline at end of file diff --git a/opensora/models/causalvideovae/model/modules/resnet_block.py b/opensora/models/causalvideovae/model/modules/resnet_block.py deleted file mode 100755 index d0e2551b5..000000000 --- a/opensora/models/causalvideovae/model/modules/resnet_block.py +++ /dev/null @@ -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 \ No newline at end of file diff --git a/opensora/models/causalvideovae/model/modules/updownsample.py b/opensora/models/causalvideovae/model/modules/updownsample.py deleted file mode 100755 index b2592851e..000000000 --- a/opensora/models/causalvideovae/model/modules/updownsample.py +++ /dev/null @@ -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) \ No newline at end of file diff --git a/opensora/models/causalvideovae/model/trainer_videobase.py b/opensora/models/causalvideovae/model/trainer_videobase.py deleted file mode 100755 index b6c4b0c7d..000000000 --- a/opensora/models/causalvideovae/model/trainer_videobase.py +++ /dev/null @@ -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")) diff --git a/opensora/models/causalvideovae/model/utils/__init__.py b/opensora/models/causalvideovae/model/utils/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/opensora/models/causalvideovae/model/utils/distrib_utils.py b/opensora/models/causalvideovae/model/utils/distrib_utils.py deleted file mode 100755 index 760c0673f..000000000 --- a/opensora/models/causalvideovae/model/utils/distrib_utils.py +++ /dev/null @@ -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 diff --git a/opensora/models/causalvideovae/model/utils/module_utils.py b/opensora/models/causalvideovae/model/utils/module_utils.py deleted file mode 100755 index 3f843e4f3..000000000 --- a/opensora/models/causalvideovae/model/utils/module_utils.py +++ /dev/null @@ -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) \ No newline at end of file diff --git a/opensora/models/causalvideovae/model/utils/scheduler_utils.py b/opensora/models/causalvideovae/model/utils/scheduler_utils.py deleted file mode 100755 index 0650a3f79..000000000 --- a/opensora/models/causalvideovae/model/utils/scheduler_utils.py +++ /dev/null @@ -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 \ No newline at end of file diff --git a/opensora/models/causalvideovae/model/utils/video_utils.py b/opensora/models/causalvideovae/model/utils/video_utils.py deleted file mode 100755 index 4fbd3f04c..000000000 --- a/opensora/models/causalvideovae/model/utils/video_utils.py +++ /dev/null @@ -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 \ No newline at end of file diff --git a/opensora/models/causalvideovae/sample/rec_video_vae.py b/opensora/models/causalvideovae/sample/rec_video_vae.py deleted file mode 100644 index e49629d0e..000000000 --- a/opensora/models/causalvideovae/sample/rec_video_vae.py +++ /dev/null @@ -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) - \ No newline at end of file diff --git a/opensora/models/causalvideovae/utils/__init__.py b/opensora/models/causalvideovae/utils/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/opensora/models/causalvideovae/utils/dataset_utils.py b/opensora/models/causalvideovae/utils/dataset_utils.py deleted file mode 100755 index 1309989b1..000000000 --- a/opensora/models/causalvideovae/utils/dataset_utils.py +++ /dev/null @@ -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 diff --git a/opensora/models/causalvideovae/utils/downloader.py b/opensora/models/causalvideovae/utils/downloader.py deleted file mode 100755 index f2ac4017b..000000000 --- a/opensora/models/causalvideovae/utils/downloader.py +++ /dev/null @@ -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 diff --git a/opensora/models/causalvideovae/utils/utils.py b/opensora/models/causalvideovae/utils/utils.py deleted file mode 100755 index 66a51417f..000000000 --- a/opensora/models/causalvideovae/utils/utils.py +++ /dev/null @@ -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', 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 - - # @ - 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) - - # " - caption = re.sub(r'"?', '', caption) - # & - caption = re.sub(r'&', '', 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 \ No newline at end of file diff --git a/opensora/models/diffusion/__init__.py b/opensora/models/diffusion/__init__.py deleted file mode 100644 index 2fdcec59c..000000000 --- a/opensora/models/diffusion/__init__.py +++ /dev/null @@ -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) \ No newline at end of file diff --git a/opensora/models/diffusion/opensora/__init__.py b/opensora/models/diffusion/opensora/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/opensora/models/diffusion/opensora/modeling_opensora.py b/opensora/models/diffusion/opensora/modeling_opensora.py deleted file mode 100644 index 6fc1692f5..000000000 --- a/opensora/models/diffusion/opensora/modeling_opensora.py +++ /dev/null @@ -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) - - - - diff --git a/opensora/models/diffusion/opensora/modules.py b/opensora/models/diffusion/opensora/modules.py deleted file mode 100644 index 5c0e87803..000000000 --- a/opensora/models/diffusion/opensora/modules.py +++ /dev/null @@ -1,1277 +0,0 @@ -from einops import rearrange -from torch import nn -import torch -import numpy as np - -from einops import rearrange, repeat -from typing import Any, Dict, Optional, Tuple -from diffusers.utils.torch_utils import maybe_allow_in_graph -from typing import Any, Dict, Optional -import re -import torch -import torch.nn.functional as F - -from torch import nn -from diffusers.utils import deprecate, logging -from diffusers.utils.torch_utils import maybe_allow_in_graph -from diffusers.models.attention import FeedForward, GatedSelfAttentionDense -from diffusers.models.attention_processor import Attention as Attention_ -from diffusers.models.embeddings import SinusoidalPositionalEmbedding -from diffusers.models.normalization import AdaLayerNorm, AdaLayerNormContinuous, AdaLayerNormZero, RMSNorm -from .rope import PositionGetter3D, RoPE3D -from opensora.utils.parallel_states import get_sequence_parallel_state, nccl_info -from opensora.utils.communications import all_to_all_SBH -logger = logging.get_logger(__name__) - -def get_3d_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_t = np.arange(grid_size[0], dtype=np.float32) / (grid_size[0] / base_size[0]) / interpolation_scale[0] - grid_h = np.arange(grid_size[1], dtype=np.float32) / (grid_size[1] / base_size[1]) / interpolation_scale[1] - grid_w = np.arange(grid_size[2], dtype=np.float32) / (grid_size[2] / base_size[2]) / interpolation_scale[2] - grid = np.meshgrid(grid_w, grid_h, grid_t) # here w goes first - grid = np.stack(grid, axis=0) - - grid = grid.reshape([3, 1, grid_size[2], grid_size[1], grid_size[0]]) - pos_embed = get_3d_sincos_pos_embed_from_grid(embed_dim, grid) - # import ipdb;ipdb.set_trace() - 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_3d_sincos_pos_embed_from_grid(embed_dim, grid): - if embed_dim % 3 != 0: - raise ValueError("embed_dim must be divisible by 3") - - # import ipdb;ipdb.set_trace() - # use 1/3 of dimensions to encode grid_t/h/w - emb_t = get_1d_sincos_pos_embed_from_grid(embed_dim // 3, grid[0]) # (T*H*W, D/3) - emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 3, grid[1]) # (T*H*W, D/3) - emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 3, grid[2]) # (T*H*W, D/3) - - emb = np.concatenate([emb_t, emb_h, emb_w], axis=1) # (T*H*W, D) - return emb - - -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[0]) / interpolation_scale[0] - grid_w = np.arange(grid_size[1], dtype=np.float32) / (grid_size[1] / base_size[1]) / interpolation_scale[1] - 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 1/3 of dimensions to encode grid_t/h/w - 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, grid_size, cls_token=False, extra_tokens=0, interpolation_scale=1.0, base_size=16, -): - """ - grid_size: int of the grid return: pos_embed: [grid_size, embed_dim] or - [1+grid_size, embed_dim] (w/ or w/o cls_token) - """ - # if isinstance(grid_size, int): - # grid_size = (grid_size, grid_size) - - grid = np.arange(grid_size, dtype=np.float32) / (grid_size / base_size) / interpolation_scale - pos_embed = get_1d_sincos_pos_embed_from_grid(embed_dim, grid) # (H*W, D/2) - 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_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 - - -class PatchEmbed2D(nn.Module): - """2D Image to Patch Embedding but with 3D position embedding""" - - def __init__( - self, - num_frames=1, - height=224, - width=224, - patch_size_t=1, - patch_size=16, - in_channels=3, - embed_dim=768, - layer_norm=False, - flatten=True, - bias=True, - interpolation_scale=(1, 1), - interpolation_scale_t=1, - use_abs_pos=True, - ): - super().__init__() - # assert num_frames == 1 - self.use_abs_pos = use_abs_pos - self.flatten = flatten - self.layer_norm = layer_norm - - self.proj = nn.Conv2d( - in_channels, embed_dim, kernel_size=(patch_size, patch_size), stride=(patch_size, patch_size), bias=bias - ) - if layer_norm: - self.norm = nn.LayerNorm(embed_dim, elementwise_affine=False, eps=1e-6) - else: - self.norm = None - - self.patch_size_t = patch_size_t - self.patch_size = patch_size - # See: - # https://github.com/PixArt-alpha/PixArt-alpha/blob/0f55e922376d8b797edd44d25d0e7464b260dcab/diffusion/model/nets/PixArtMS.py#L161 - - self.height, self.width = height // patch_size, width // patch_size - self.base_size = (height // patch_size, width // patch_size) - self.interpolation_scale = (interpolation_scale[0], interpolation_scale[1]) - pos_embed = get_2d_sincos_pos_embed( - embed_dim, (self.height, self.width), base_size=self.base_size, interpolation_scale=self.interpolation_scale - ) - self.register_buffer("pos_embed", torch.from_numpy(pos_embed).float().unsqueeze(0), persistent=False) - - self.num_frames = (num_frames - 1) // patch_size_t + 1 if num_frames % 2 == 1 else num_frames // patch_size_t - self.base_size_t = (num_frames - 1) // patch_size_t + 1 if num_frames % 2 == 1 else num_frames // patch_size_t - self.interpolation_scale_t = interpolation_scale_t - temp_pos_embed = get_1d_sincos_pos_embed(embed_dim, self.num_frames, base_size=self.base_size_t, interpolation_scale=self.interpolation_scale_t) - self.register_buffer("temp_pos_embed", torch.from_numpy(temp_pos_embed).float().unsqueeze(0), persistent=False) - # self.temp_embed_gate = nn.Parameter(torch.tensor([0.0])) - - def forward(self, latent, num_frames): - b, _, _, _, _ = latent.shape - video_latent, image_latent = None, None - # b c 1 h w - # assert latent.shape[-3] == 1 and num_frames == 1 - height, width = latent.shape[-2] // self.patch_size, latent.shape[-1] // self.patch_size - latent = rearrange(latent, 'b c t h w -> (b t) c h w') - latent = self.proj(latent) - - if self.flatten: - latent = latent.flatten(2).transpose(1, 2) # BT C H W -> BT N C - if self.layer_norm: - latent = self.norm(latent) - - if self.use_abs_pos: - # Interpolate positional embeddings if needed. - # (For PixArt-Alpha: https://github.com/PixArt-alpha/PixArt-alpha/blob/0f55e922376d8b797edd44d25d0e7464b260dcab/diffusion/model/nets/PixArtMS.py#L162C151-L162C160) - if self.height != height or self.width != width: - # raise NotImplementedError - pos_embed = get_2d_sincos_pos_embed( - embed_dim=self.pos_embed.shape[-1], - grid_size=(height, width), - base_size=self.base_size, - interpolation_scale=self.interpolation_scale, - ) - pos_embed = torch.from_numpy(pos_embed) - pos_embed = pos_embed.float().unsqueeze(0).to(latent.device) - else: - pos_embed = self.pos_embed - - - if self.num_frames != num_frames: - # import ipdb;ipdb.set_trace() - # raise NotImplementedError - if get_sequence_parallel_state(): - - sp_size = nccl_info.world_size - temp_pos_embed = get_1d_sincos_pos_embed( - embed_dim=self.temp_pos_embed.shape[-1], - grid_size=num_frames * sp_size, - base_size=self.base_size_t, - interpolation_scale=self.interpolation_scale_t, - ) - rank = nccl_info.rank % sp_size - st_frame = rank * num_frames - ed_frame = st_frame + num_frames - temp_pos_embed = temp_pos_embed[st_frame: ed_frame] - - else: - temp_pos_embed = get_1d_sincos_pos_embed( - embed_dim=self.temp_pos_embed.shape[-1], - grid_size=num_frames, - base_size=self.base_size_t, - interpolation_scale=self.interpolation_scale_t, - ) - temp_pos_embed = torch.from_numpy(temp_pos_embed) - temp_pos_embed = temp_pos_embed.float().unsqueeze(0).to(latent.device) - else: - temp_pos_embed = self.temp_pos_embed - - latent = (latent + pos_embed).to(latent.dtype) - - latent = rearrange(latent, '(b t) n c -> b t n c', b=b) - video_latent, image_latent = latent[:, :num_frames], latent[:, num_frames:] - - if self.use_abs_pos: - # temp_pos_embed = temp_pos_embed.unsqueeze(2) * self.temp_embed_gate.tanh() - temp_pos_embed = temp_pos_embed.unsqueeze(2) - video_latent = (video_latent + temp_pos_embed).to(video_latent.dtype) if video_latent is not None and video_latent.numel() > 0 else None - image_latent = (image_latent + temp_pos_embed[:, :1]).to(image_latent.dtype) if image_latent is not None and image_latent.numel() > 0 else None - - video_latent = rearrange(video_latent, 'b t n c -> b (t n) c') if video_latent is not None and video_latent.numel() > 0 else None - image_latent = rearrange(image_latent, 'b t n c -> (b t) n c') if image_latent is not None and image_latent.numel() > 0 else None - - if num_frames == 1 and image_latent is None and not get_sequence_parallel_state(): - image_latent = video_latent - video_latent = None - # print('video_latent is None, image_latent is None', video_latent is None, image_latent is None) - return video_latent, image_latent - - - -class OverlapPatchEmbed3D(nn.Module): - """2D Image to Patch Embedding but with 3D position embedding""" - - def __init__( - self, - num_frames=1, - height=224, - width=224, - patch_size_t=1, - patch_size=16, - in_channels=3, - embed_dim=768, - layer_norm=False, - flatten=True, - bias=True, - interpolation_scale=(1, 1), - interpolation_scale_t=1, - use_abs_pos=True, - ): - super().__init__() - # assert patch_size_t == 1 and patch_size == 1 - self.use_abs_pos = use_abs_pos - self.flatten = flatten - self.layer_norm = layer_norm - - self.proj = nn.Conv3d( - in_channels, embed_dim, kernel_size=(patch_size_t, patch_size, patch_size), stride=(patch_size_t, patch_size, patch_size), bias=bias - ) - if layer_norm: - self.norm = nn.LayerNorm(embed_dim, elementwise_affine=False, eps=1e-6) - else: - self.norm = None - - self.patch_size_t = patch_size_t - self.patch_size = patch_size - # See: - # https://github.com/PixArt-alpha/PixArt-alpha/blob/0f55e922376d8b797edd44d25d0e7464b260dcab/diffusion/model/nets/PixArtMS.py#L161 - - self.height, self.width = height // patch_size, width // patch_size - self.base_size = (height // patch_size, width // patch_size) - self.interpolation_scale = (interpolation_scale[0], interpolation_scale[1]) - pos_embed = get_2d_sincos_pos_embed( - embed_dim, (self.height, self.width), base_size=self.base_size, interpolation_scale=self.interpolation_scale - ) - self.register_buffer("pos_embed", torch.from_numpy(pos_embed).float().unsqueeze(0), persistent=False) - - self.num_frames = (num_frames - 1) // patch_size_t + 1 if num_frames % 2 == 1 else num_frames // patch_size_t - self.base_size_t = (num_frames - 1) // patch_size_t + 1 if num_frames % 2 == 1 else num_frames // patch_size_t - self.interpolation_scale_t = interpolation_scale_t - temp_pos_embed = get_1d_sincos_pos_embed(embed_dim, self.num_frames, base_size=self.base_size_t, interpolation_scale=self.interpolation_scale_t) - self.register_buffer("temp_pos_embed", torch.from_numpy(temp_pos_embed).float().unsqueeze(0), persistent=False) - # self.temp_embed_gate = nn.Parameter(torch.tensor([0.0])) - - def forward(self, latent, num_frames): - b, _, _, _, _ = latent.shape - video_latent, image_latent = None, None - # b c 1 h w - # assert latent.shape[-3] == 1 and num_frames == 1 - height, width = latent.shape[-2] // self.patch_size, latent.shape[-1] // self.patch_size - # latent = rearrange(latent, 'b c t h w -> (b t) c h w') - latent = self.proj(latent) - - if self.flatten: - # latent = latent.flatten(2).transpose(1, 2) # BT C H W -> BT N C - latent = rearrange(latent, 'b c t h w -> (b t) (h w) c ') - if self.layer_norm: - latent = self.norm(latent) - - if self.use_abs_pos: - # Interpolate positional embeddings if needed. - # (For PixArt-Alpha: https://github.com/PixArt-alpha/PixArt-alpha/blob/0f55e922376d8b797edd44d25d0e7464b260dcab/diffusion/model/nets/PixArtMS.py#L162C151-L162C160) - if self.height != height or self.width != width: - # raise NotImplementedError - pos_embed = get_2d_sincos_pos_embed( - embed_dim=self.pos_embed.shape[-1], - grid_size=(height, width), - base_size=self.base_size, - interpolation_scale=self.interpolation_scale, - ) - pos_embed = torch.from_numpy(pos_embed) - pos_embed = pos_embed.float().unsqueeze(0).to(latent.device) - else: - pos_embed = self.pos_embed - - - if self.num_frames != num_frames: - # import ipdb;ipdb.set_trace() - # raise NotImplementedError - temp_pos_embed = get_1d_sincos_pos_embed( - embed_dim=self.temp_pos_embed.shape[-1], - grid_size=num_frames, - base_size=self.base_size_t, - interpolation_scale=self.interpolation_scale_t, - ) - temp_pos_embed = torch.from_numpy(temp_pos_embed) - temp_pos_embed = temp_pos_embed.float().unsqueeze(0).to(latent.device) - else: - temp_pos_embed = self.temp_pos_embed - - latent = (latent + pos_embed).to(latent.dtype) - - latent = rearrange(latent, '(b t) n c -> b t n c', b=b) - video_latent, image_latent = latent[:, :num_frames], latent[:, num_frames:] - - if self.use_abs_pos: - # temp_pos_embed = temp_pos_embed.unsqueeze(2) * self.temp_embed_gate.tanh() - temp_pos_embed = temp_pos_embed.unsqueeze(2) - video_latent = (video_latent + temp_pos_embed).to(video_latent.dtype) if video_latent is not None and video_latent.numel() > 0 else None - image_latent = (image_latent + temp_pos_embed[:, :1]).to(image_latent.dtype) if image_latent is not None and image_latent.numel() > 0 else None - - - video_latent = rearrange(video_latent, 'b t n c -> b (t n) c') if video_latent is not None and video_latent.numel() > 0 else None - image_latent = rearrange(image_latent, 'b t n c -> (b t) n c') if image_latent is not None and image_latent.numel() > 0 else None - - if num_frames == 1 and image_latent is None: - image_latent = video_latent - video_latent = None - return video_latent, image_latent - - - -class OverlapPatchEmbed2D(nn.Module): - """2D Image to Patch Embedding but with 3D position embedding""" - - def __init__( - self, - num_frames=1, - height=224, - width=224, - patch_size_t=1, - patch_size=16, - in_channels=3, - embed_dim=768, - layer_norm=False, - flatten=True, - bias=True, - interpolation_scale=(1, 1), - interpolation_scale_t=1, - use_abs_pos=True, - ): - super().__init__() - assert patch_size_t == 1 - self.use_abs_pos = use_abs_pos - self.flatten = flatten - self.layer_norm = layer_norm - - self.proj = nn.Conv2d( - in_channels, embed_dim, kernel_size=(patch_size, patch_size), stride=(patch_size, patch_size), bias=bias - ) - if layer_norm: - self.norm = nn.LayerNorm(embed_dim, elementwise_affine=False, eps=1e-6) - else: - self.norm = None - - self.patch_size_t = patch_size_t - self.patch_size = patch_size - # See: - # https://github.com/PixArt-alpha/PixArt-alpha/blob/0f55e922376d8b797edd44d25d0e7464b260dcab/diffusion/model/nets/PixArtMS.py#L161 - - self.height, self.width = height // patch_size, width // patch_size - self.base_size = (height // patch_size, width // patch_size) - self.interpolation_scale = (interpolation_scale[0], interpolation_scale[1]) - pos_embed = get_2d_sincos_pos_embed( - embed_dim, (self.height, self.width), base_size=self.base_size, interpolation_scale=self.interpolation_scale - ) - self.register_buffer("pos_embed", torch.from_numpy(pos_embed).float().unsqueeze(0), persistent=False) - - self.num_frames = (num_frames - 1) // patch_size_t + 1 if num_frames % 2 == 1 else num_frames // patch_size_t - self.base_size_t = (num_frames - 1) // patch_size_t + 1 if num_frames % 2 == 1 else num_frames // patch_size_t - self.interpolation_scale_t = interpolation_scale_t - temp_pos_embed = get_1d_sincos_pos_embed(embed_dim, self.num_frames, base_size=self.base_size_t, interpolation_scale=self.interpolation_scale_t) - self.register_buffer("temp_pos_embed", torch.from_numpy(temp_pos_embed).float().unsqueeze(0), persistent=False) - # self.temp_embed_gate = nn.Parameter(torch.tensor([0.0])) - - def forward(self, latent, num_frames): - b, _, _, _, _ = latent.shape - video_latent, image_latent = None, None - # b c 1 h w - # assert latent.shape[-3] == 1 and num_frames == 1 - height, width = latent.shape[-2] // self.patch_size, latent.shape[-1] // self.patch_size - latent = rearrange(latent, 'b c t h w -> (b t) c h w') - latent = self.proj(latent) - - if self.flatten: - latent = latent.flatten(2).transpose(1, 2) # BT C H W -> BT N C - if self.layer_norm: - latent = self.norm(latent) - - if self.use_abs_pos: - # Interpolate positional embeddings if needed. - # (For PixArt-Alpha: https://github.com/PixArt-alpha/PixArt-alpha/blob/0f55e922376d8b797edd44d25d0e7464b260dcab/diffusion/model/nets/PixArtMS.py#L162C151-L162C160) - if self.height != height or self.width != width: - # raise NotImplementedError - pos_embed = get_2d_sincos_pos_embed( - embed_dim=self.pos_embed.shape[-1], - grid_size=(height, width), - base_size=self.base_size, - interpolation_scale=self.interpolation_scale, - ) - pos_embed = torch.from_numpy(pos_embed) - pos_embed = pos_embed.float().unsqueeze(0).to(latent.device) - else: - pos_embed = self.pos_embed - - - if self.num_frames != num_frames: - # import ipdb;ipdb.set_trace() - # raise NotImplementedError - temp_pos_embed = get_1d_sincos_pos_embed( - embed_dim=self.temp_pos_embed.shape[-1], - grid_size=num_frames, - base_size=self.base_size_t, - interpolation_scale=self.interpolation_scale_t, - ) - temp_pos_embed = torch.from_numpy(temp_pos_embed) - temp_pos_embed = temp_pos_embed.float().unsqueeze(0).to(latent.device) - else: - temp_pos_embed = self.temp_pos_embed - - latent = (latent + pos_embed).to(latent.dtype) - - latent = rearrange(latent, '(b t) n c -> b t n c', b=b) - video_latent, image_latent = latent[:, :num_frames], latent[:, num_frames:] - - if self.use_abs_pos: - # temp_pos_embed = temp_pos_embed.unsqueeze(2) * self.temp_embed_gate.tanh() - temp_pos_embed = temp_pos_embed.unsqueeze(2) - video_latent = (video_latent + temp_pos_embed).to(video_latent.dtype) if video_latent is not None and video_latent.numel() > 0 else None - image_latent = (image_latent + temp_pos_embed[:, :1]).to(image_latent.dtype) if image_latent is not None and image_latent.numel() > 0 else None - - - video_latent = rearrange(video_latent, 'b t n c -> b (t n) c') if video_latent is not None and video_latent.numel() > 0 else None - image_latent = rearrange(image_latent, 'b t n c -> (b t) n c') if image_latent is not None and image_latent.numel() > 0 else None - - if num_frames == 1 and image_latent is None: - image_latent = video_latent - video_latent = None - return video_latent, image_latent - -class Attention(Attention_): - def __init__(self, downsampler, attention_mode, use_rope, interpolation_scale_thw, **kwags): - processor = AttnProcessor2_0(attention_mode=attention_mode, use_rope=use_rope, interpolation_scale_thw=interpolation_scale_thw) - super().__init__(processor=processor, **kwags) - self.downsampler = None - if downsampler: # downsampler k155_s122 - downsampler_ker_size = list(re.search(r'k(\d{2,3})', downsampler).group(1)) # 122 - down_factor = list(re.search(r's(\d{2,3})', downsampler).group(1)) - downsampler_ker_size = [int(i) for i in downsampler_ker_size] - downsampler_padding = [(i - 1) // 2 for i in downsampler_ker_size] - down_factor = [int(i) for i in down_factor] - - if len(downsampler_ker_size) == 2: - self.downsampler = DownSampler2d(kwags['query_dim'], kwags['query_dim'], kernel_size=downsampler_ker_size, stride=1, - padding=downsampler_padding, groups=kwags['query_dim'], down_factor=down_factor, - down_shortcut=True) - elif len(downsampler_ker_size) == 3: - self.downsampler = DownSampler3d(kwags['query_dim'], kwags['query_dim'], kernel_size=downsampler_ker_size, stride=1, - padding=downsampler_padding, groups=kwags['query_dim'], down_factor=down_factor, - down_shortcut=True) - - # self.q_norm = nn.LayerNorm(kwags['dim_head'], elementwise_affine=True, eps=1e-6) - # self.k_norm = nn.LayerNorm(kwags['dim_head'], elementwise_affine=True, eps=1e-6) - - def prepare_attention_mask( - self, attention_mask: torch.Tensor, target_length: int, batch_size: int, out_dim: int = 3 - ) -> torch.Tensor: - r""" - Prepare the attention mask for the attention computation. - - Args: - attention_mask (`torch.Tensor`): - The attention mask to prepare. - target_length (`int`): - The target length of the attention mask. This is the length of the attention mask after padding. - batch_size (`int`): - The batch size, which is used to repeat the attention mask. - out_dim (`int`, *optional*, defaults to `3`): - The output dimension of the attention mask. Can be either `3` or `4`. - - Returns: - `torch.Tensor`: The prepared attention mask. - """ - head_size = self.heads - if get_sequence_parallel_state(): - head_size = head_size // nccl_info.world_size - if attention_mask is None: - return attention_mask - - current_length: int = attention_mask.shape[-1] - if current_length != target_length: - if attention_mask.device.type == "mps": - # HACK: MPS: Does not support padding by greater than dimension of input tensor. - # Instead, we can manually construct the padding tensor. - padding_shape = (attention_mask.shape[0], attention_mask.shape[1], target_length) - padding = torch.zeros(padding_shape, dtype=attention_mask.dtype, device=attention_mask.device) - attention_mask = torch.cat([attention_mask, padding], dim=2) - else: - # TODO: for pipelines such as stable-diffusion, padding cross-attn mask: - # we want to instead pad by (0, remaining_length), where remaining_length is: - # remaining_length: int = target_length - current_length - # TODO: re-enable tests/models/test_models_unet_2d_condition.py#test_model_xattn_padding - attention_mask = F.pad(attention_mask, (0, target_length), value=0.0) - - if out_dim == 3: - if attention_mask.shape[0] < batch_size * head_size: - attention_mask = attention_mask.repeat_interleave(head_size, dim=0) - elif out_dim == 4: - attention_mask = attention_mask.unsqueeze(1) - attention_mask = attention_mask.repeat_interleave(head_size, dim=1) - - return attention_mask - -class DownSampler3d(nn.Module): - def __init__(self, *args, **kwargs): - ''' Required kwargs: down_factor, downsampler''' - super().__init__() - self.down_factor = kwargs.pop('down_factor') - self.down_shortcut = kwargs.pop('down_shortcut') - self.layer = nn.Conv3d(*args, **kwargs) - - def forward(self, x, attention_mask, t, h, w): - b = x.shape[0] - x = rearrange(x, 'b (t h w) d -> b d t h w', t=t, h=h, w=w) - x = self.layer(x) + (x if self.down_shortcut else 0) - self.t = t//self.down_factor[0] - self.h = h//self.down_factor[1] - self.w = w//self.down_factor[2] - x = rearrange(x, 'b d (t dt) (h dh) (w dw) -> (b dt dh dw) (t h w) d', - t=t//self.down_factor[0], h=h//self.down_factor[1], w=w//self.down_factor[2], - dt=self.down_factor[0], dh=self.down_factor[1], dw=self.down_factor[2]) - - - attention_mask = rearrange(attention_mask, 'b 1 (t h w) -> b 1 t h w', t=t, h=h, w=w) - attention_mask = rearrange(attention_mask, 'b 1 (t dt) (h dh) (w dw) -> (b dt dh dw) 1 (t h w)', - t=t//self.down_factor[0], h=h//self.down_factor[1], w=w//self.down_factor[2], - dt=self.down_factor[0], dh=self.down_factor[1], dw=self.down_factor[2]) - return x, attention_mask - - def reverse(self, x, t, h, w): - x = rearrange(x, '(b dt dh dw) (t h w) d -> b (t dt h dh w dw) d', - t=t, h=h, w=w, - dt=self.down_factor[0], dh=self.down_factor[1], dw=self.down_factor[2]) - return x - - -class DownSampler2d(nn.Module): - def __init__(self, *args, **kwargs): - ''' Required kwargs: down_factor, downsampler''' - super().__init__() - self.down_factor = kwargs.pop('down_factor') - self.down_shortcut = kwargs.pop('down_shortcut') - self.layer = nn.Conv2d(*args, **kwargs) - - def forward(self, x, attention_mask, t, h, w): - b = x.shape[0] - x = rearrange(x, 'b (t h w) d -> (b t) d h w', t=t, h=h, w=w) - x = self.layer(x) + (x if self.down_shortcut else 0) - - self.t = 1 - self.h = h//self.down_factor[0] - self.w = w//self.down_factor[1] - - x = rearrange(x, 'b d (h dh) (w dw) -> (b dh dw) (h w) d', - h=h//self.down_factor[0], w=w//self.down_factor[1], - dh=self.down_factor[0], dw=self.down_factor[1]) - - attention_mask = rearrange(attention_mask, 'b 1 (t h w) -> (b t) 1 h w', h=h, w=w) - attention_mask = rearrange(attention_mask, 'b 1 (h dh) (w dw) -> (b dh dw) 1 (h w)', - h=h//self.down_factor[0], w=w//self.down_factor[1], - dh=self.down_factor[0], dw=self.down_factor[1]) - return x, attention_mask - - def reverse(self, x, t, h, w): - x = rearrange(x, '(b t dh dw) (h w) d -> b (t h dh w dw) d', - t=t, h=h, w=w, - dh=self.down_factor[0], dw=self.down_factor[1]) - return x - -class AttnProcessor2_0: - r""" - Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0). - """ - - def __init__(self, attention_mode='xformers', use_rope=False, interpolation_scale_thw=(1, 1, 1)): - self.use_rope = use_rope - self.interpolation_scale_thw = interpolation_scale_thw - if self.use_rope: - self._init_rope(interpolation_scale_thw) - self.attention_mode = attention_mode - if not hasattr(F, "scaled_dot_product_attention"): - raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.") - - - def _init_rope(self, interpolation_scale_thw): - self.rope = RoPE3D(interpolation_scale_thw=interpolation_scale_thw) - self.position_getter = PositionGetter3D() - - def __call__( - self, - attn: Attention, - hidden_states: torch.FloatTensor, - encoder_hidden_states: Optional[torch.FloatTensor] = None, - attention_mask: Optional[torch.FloatTensor] = None, - temb: Optional[torch.FloatTensor] = None, - frame: int = 8, - height: int = 16, - width: int = 16, - *args, - **kwargs, - ) -> torch.FloatTensor: - if len(args) > 0 or kwargs.get("scale", None) is not None: - deprecation_message = "The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`." - deprecate("scale", "1.0.0", deprecation_message) - - - if attn.downsampler is not None: - hidden_states, attention_mask = attn.downsampler(hidden_states, attention_mask, t=frame, h=height, w=width) - frame, height, width = attn.downsampler.t, attn.downsampler.h, attn.downsampler.w - - residual = hidden_states - - if attn.spatial_norm is not None: - hidden_states = attn.spatial_norm(hidden_states, temb) - - input_ndim = hidden_states.ndim - - if input_ndim == 4: - batch_size, channel, height, width = hidden_states.shape - hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2) - - if get_sequence_parallel_state(): - - sequence_length, batch_size, _ = ( - hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape - ) - else: - batch_size, sequence_length, _ = ( - hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape - ) - - if attention_mask is not None: - attention_mask = attention_mask.view(batch_size, 1, -1, attention_mask.shape[-1]) - - if attn.group_norm is not None: - hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2) - - query = attn.to_q(hidden_states) - - if encoder_hidden_states is None: - encoder_hidden_states = hidden_states - elif attn.norm_cross: - encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states) - key = attn.to_k(encoder_hidden_states) - value = attn.to_v(encoder_hidden_states) - - inner_dim = key.shape[-1] - head_dim = inner_dim // attn.heads - - if get_sequence_parallel_state(): - query = query.reshape(-1, attn.heads, head_dim) # [s // sp, b, h * d] -> [s // sp * b, h, d] - key = key.reshape(-1, attn.heads, head_dim) - value = value.reshape(-1, attn.heads, head_dim) - # query = attn.q_norm(query) - # key = attn.k_norm(key) - h_size = attn.heads * head_dim - sp_size = nccl_info.world_size - h_size_sp = h_size // sp_size - # apply all_to_all to gather sequence and split attention heads [s // sp * b, h, d] -> [s * b, h // sp, d] - query = all_to_all_SBH(query, scatter_dim=1, gather_dim=0).reshape(-1, batch_size, h_size_sp) - key = all_to_all_SBH(key, scatter_dim=1, gather_dim=0).reshape(-1, batch_size, h_size_sp) - value = all_to_all_SBH(value, scatter_dim=1, gather_dim=0).reshape(-1, batch_size, h_size_sp) - query = query.reshape(-1, batch_size, attn.heads // sp_size, head_dim) - key = key.reshape(-1, batch_size, attn.heads // sp_size, head_dim) - value = value.reshape(-1, batch_size, attn.heads // sp_size, head_dim) - # print('query', query.shape, 'key', key.shape, 'value', value.shape) - if self.use_rope: - # require the shape of (batch_size x nheads x ntokens x dim) - pos_thw = self.position_getter(batch_size, t=frame * sp_size, h=height, w=width, device=query.device) - query = self.rope(query, pos_thw) - key = self.rope(key, pos_thw) - - # print('after rope query', query.shape, 'key', key.shape, 'value', value.shape) - query = rearrange(query, 's b h d -> b h s d') - key = rearrange(key, 's b h d -> b h s d') - value = rearrange(value, 's b h d -> b h s d') - # print('rearrange query', query.shape, 'key', key.shape, 'value', value.shape) - - # 0, -10000 ->(bool) False, True ->(any) True ->(not) False - # 0, 0 ->(bool) False, False ->(any) False ->(not) True - if attention_mask is None or not torch.any(attention_mask.bool()): # 0 mean visible - attention_mask = None - # the output of sdp = (batch, num_heads, seq_len, head_dim) - # TODO: add support for attn.scale when we move to Torch 2.1 - # import ipdb;ipdb.set_trace() - # print(attention_mask) - if self.attention_mode == 'flash': - assert attention_mask is None, 'flash-attn do not support attention_mask' - with torch.nn.attention.sdpa_kernel(torch.nn.attention.SDPBackend.FLASH_ATTENTION): - hidden_states = F.scaled_dot_product_attention( - query, key, value, dropout_p=0.0, is_causal=False - ) - elif self.attention_mode == 'xformers': - with torch.nn.attention.sdpa_kernel(torch.nn.attention.SDPBackend.EFFICIENT_ATTENTION): - hidden_states = F.scaled_dot_product_attention( - query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False - ) - elif self.attention_mode == 'math': - hidden_states = F.scaled_dot_product_attention( - query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False - ) - else: - raise NotImplementedError(f'Found attention_mode: {self.attention_mode}') - - hidden_states = rearrange(hidden_states, 'b h s d -> s b h d') - - hidden_states = hidden_states.reshape(-1, attn.heads // sp_size, head_dim) - - # [s * b, h // sp, d] -> [s // sp * b, h, d] -> [s // sp, b, h * d] - hidden_states = all_to_all_SBH(hidden_states, scatter_dim=0, gather_dim=1).reshape(-1, batch_size, h_size) - else: - query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - - # qk norm - # query = attn.q_norm(query) - # key = attn.k_norm(key) - - if self.use_rope: - # require the shape of (batch_size x nheads x ntokens x dim) - pos_thw = self.position_getter(batch_size, t=frame, h=height, w=width, device=query.device) - query = self.rope(query, pos_thw) - key = self.rope(key, pos_thw) - - value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) - - # 0, -10000 ->(bool) False, True ->(any) True ->(not) False - # 0, 0 ->(bool) False, False ->(any) False ->(not) True - if attention_mask is None or not torch.any(attention_mask.bool()): # 0 mean visible - attention_mask = None - # the output of sdp = (batch, num_heads, seq_len, head_dim) - # TODO: add support for attn.scale when we move to Torch 2.1 - # import ipdb;ipdb.set_trace() - # print(attention_mask) - if self.attention_mode == 'flash': - assert attention_mask is None, 'flash-attn do not support attention_mask' - with torch.nn.attention.sdpa_kernel(torch.nn.attention.SDPBackend.FLASH_ATTENTION): - hidden_states = F.scaled_dot_product_attention( - query, key, value, dropout_p=0.0, is_causal=False - ) - elif self.attention_mode == 'xformers': - with torch.nn.attention.sdpa_kernel(torch.nn.attention.SDPBackend.EFFICIENT_ATTENTION): - hidden_states = F.scaled_dot_product_attention( - query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False - ) - elif self.attention_mode == 'math': - hidden_states = F.scaled_dot_product_attention( - query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False - ) - else: - raise NotImplementedError(f'Found attention_mode: {self.attention_mode}') - hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) - hidden_states = 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 input_ndim == 4: - hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width) - - if attn.residual_connection: - hidden_states = hidden_states + residual - - hidden_states = hidden_states / attn.rescale_output_factor - - if attn.downsampler is not None: - hidden_states = attn.downsampler.reverse(hidden_states, t=frame, h=height, w=width) - return hidden_states - - - -class FeedForward_Conv3d(nn.Module): - def __init__(self, downsampler, dim, hidden_features, bias=True): - super(FeedForward_Conv3d, self).__init__() - - self.bias = bias - - self.project_in = nn.Linear(dim, hidden_features, bias=bias) - - self.dwconv = nn.ModuleList([ - nn.Conv3d(hidden_features, hidden_features, kernel_size=(5, 5, 5), stride=1, padding=(2, 2, 2), dilation=1, - groups=hidden_features, bias=bias), - nn.Conv3d(hidden_features, hidden_features, kernel_size=(3, 3, 3), stride=1, padding=(1, 1, 1), dilation=1, - groups=hidden_features, bias=bias), - nn.Conv3d(hidden_features, hidden_features, kernel_size=(1, 1, 1), stride=1, padding=(0, 0, 0), dilation=1, - groups=hidden_features, bias=bias) - ]) - - self.project_out = nn.Linear(hidden_features, dim, bias=bias) - - - def forward(self, x, t, h, w): - # import ipdb;ipdb.set_trace() - x = self.project_in(x) - x = rearrange(x, 'b (t h w) d -> b d t h w', t=t, h=h, w=w) - x = F.gelu(x) - out = x - for module in self.dwconv: - out = out + module(x) - out = rearrange(out, 'b d t h w -> b (t h w) d', t=t, h=h, w=w) - x = self.project_out(out) - return x - - -class FeedForward_Conv2d(nn.Module): - def __init__(self, downsampler, dim, hidden_features, bias=True): - super(FeedForward_Conv2d, self).__init__() - - self.bias = bias - - self.project_in = nn.Linear(dim, hidden_features, bias=bias) - - self.dwconv = nn.ModuleList([ - nn.Conv2d(hidden_features, hidden_features, kernel_size=(5, 5), stride=1, padding=(2, 2), dilation=1, - groups=hidden_features, bias=bias), - nn.Conv2d(hidden_features, hidden_features, kernel_size=(3, 3), stride=1, padding=(1, 1), dilation=1, - groups=hidden_features, bias=bias), - nn.Conv2d(hidden_features, hidden_features, kernel_size=(1, 1), stride=1, padding=(0, 0), dilation=1, - groups=hidden_features, bias=bias) - ]) - - self.project_out = nn.Linear(hidden_features, dim, bias=bias) - - - def forward(self, x, t, h, w): - # import ipdb;ipdb.set_trace() - x = self.project_in(x) - x = rearrange(x, 'b (t h w) d -> (b t) d h w', t=t, h=h, w=w) - x = F.gelu(x) - out = x - for module in self.dwconv: - out = out + module(x) - out = rearrange(out, '(b t) d h w -> b (t h w) d', t=t, h=h, w=w) - x = self.project_out(out) - return x - -@maybe_allow_in_graph -class BasicTransformerBlock(nn.Module): - r""" - A basic Transformer block. - - Parameters: - dim (`int`): The number of channels in the input and output. - num_attention_heads (`int`): The number of heads to use for multi-head attention. - attention_head_dim (`int`): The number of channels in each head. - dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use. - cross_attention_dim (`int`, *optional*): The size of the encoder_hidden_states vector for cross attention. - activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward. - num_embeds_ada_norm (: - obj: `int`, *optional*): The number of diffusion steps used during training. See `Transformer2DModel`. - attention_bias (: - obj: `bool`, *optional*, defaults to `False`): Configure if the attentions should contain a bias parameter. - only_cross_attention (`bool`, *optional*): - Whether to use only cross-attention layers. In this case two cross attention layers are used. - double_self_attention (`bool`, *optional*): - Whether to use two self-attention layers. In this case no cross attention layers are used. - upcast_attention (`bool`, *optional*): - Whether to upcast the attention computation to float32. This is useful for mixed precision training. - norm_elementwise_affine (`bool`, *optional*, defaults to `True`): - Whether to use learnable elementwise affine parameters for normalization. - norm_type (`str`, *optional*, defaults to `"layer_norm"`): - The normalization layer to use. Can be `"layer_norm"`, `"ada_norm"` or `"ada_norm_zero"`. - final_dropout (`bool` *optional*, defaults to False): - Whether to apply a final dropout after the last feed-forward layer. - attention_type (`str`, *optional*, defaults to `"default"`): - The type of attention to use. Can be `"default"` or `"gated"` or `"gated-text-image"`. - positional_embeddings (`str`, *optional*, defaults to `None`): - The type of positional embeddings to apply to. - num_positional_embeddings (`int`, *optional*, defaults to `None`): - The maximum number of positional embeddings to apply. - """ - - def __init__( - self, - dim: int, - num_attention_heads: int, - attention_head_dim: int, - dropout=0.0, - cross_attention_dim: Optional[int] = None, - activation_fn: str = "geglu", - num_embeds_ada_norm: Optional[int] = None, - attention_bias: bool = False, - only_cross_attention: bool = False, - double_self_attention: bool = False, - upcast_attention: bool = False, - norm_elementwise_affine: bool = True, - norm_type: str = "layer_norm", # 'layer_norm', 'ada_norm', 'ada_norm_zero', 'ada_norm_single', 'ada_norm_continuous', 'layer_norm_i2vgen' - norm_eps: float = 1e-5, - final_dropout: bool = False, - attention_type: str = "default", - positional_embeddings: Optional[str] = None, - num_positional_embeddings: Optional[int] = None, - ada_norm_continous_conditioning_embedding_dim: Optional[int] = None, - ada_norm_bias: Optional[int] = None, - ff_inner_dim: Optional[int] = None, - ff_bias: bool = True, - attention_out_bias: bool = True, - attention_mode: str = "xformers", - downsampler: str = None, - use_rope: bool = False, - interpolation_scale_thw: Tuple[int] = (1, 1, 1), - ): - super().__init__() - self.only_cross_attention = only_cross_attention - self.downsampler = downsampler - - # We keep these boolean flags for backward-compatibility. - self.use_ada_layer_norm_zero = (num_embeds_ada_norm is not None) and norm_type == "ada_norm_zero" - self.use_ada_layer_norm = (num_embeds_ada_norm is not None) and norm_type == "ada_norm" - self.use_ada_layer_norm_single = norm_type == "ada_norm_single" - self.use_layer_norm = norm_type == "layer_norm" - self.use_ada_layer_norm_continuous = norm_type == "ada_norm_continuous" - - if norm_type in ("ada_norm", "ada_norm_zero") and num_embeds_ada_norm is None: - raise ValueError( - f"`norm_type` is set to {norm_type}, but `num_embeds_ada_norm` is not defined. Please make sure to" - f" define `num_embeds_ada_norm` if setting `norm_type` to {norm_type}." - ) - - self.norm_type = norm_type - self.num_embeds_ada_norm = num_embeds_ada_norm - - if positional_embeddings and (num_positional_embeddings is None): - raise ValueError( - "If `positional_embedding` type is defined, `num_positition_embeddings` must also be defined." - ) - - if positional_embeddings == "sinusoidal": - self.pos_embed = SinusoidalPositionalEmbedding(dim, max_seq_length=num_positional_embeddings) - else: - self.pos_embed = None - - # Define 3 blocks. Each block has its own normalization layer. - # 1. Self-Attn - if norm_type == "ada_norm": - self.norm1 = AdaLayerNorm(dim, num_embeds_ada_norm) - elif norm_type == "ada_norm_zero": - self.norm1 = AdaLayerNormZero(dim, num_embeds_ada_norm) - elif norm_type == "ada_norm_continuous": - self.norm1 = AdaLayerNormContinuous( - dim, - ada_norm_continous_conditioning_embedding_dim, - norm_elementwise_affine, - norm_eps, - ada_norm_bias, - "rms_norm", - ) - else: - self.norm1 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps) - - self.attn1 = Attention( - query_dim=dim, - heads=num_attention_heads, - dim_head=attention_head_dim, - dropout=dropout, - bias=attention_bias, - cross_attention_dim=cross_attention_dim if only_cross_attention else None, - upcast_attention=upcast_attention, - out_bias=attention_out_bias, - attention_mode=attention_mode, - downsampler=downsampler, - use_rope=use_rope, - interpolation_scale_thw=interpolation_scale_thw, - ) - - # 2. Cross-Attn - if cross_attention_dim is not None or double_self_attention: - # We currently only use AdaLayerNormZero for self attention where there will only be one attention block. - # I.e. the number of returned modulation chunks from AdaLayerZero would not make sense if returned during - # the second cross attention block. - if norm_type == "ada_norm": - self.norm2 = AdaLayerNorm(dim, num_embeds_ada_norm) - elif norm_type == "ada_norm_continuous": - self.norm2 = AdaLayerNormContinuous( - dim, - ada_norm_continous_conditioning_embedding_dim, - norm_elementwise_affine, - norm_eps, - ada_norm_bias, - "rms_norm", - ) - else: - self.norm2 = nn.LayerNorm(dim, norm_eps, norm_elementwise_affine) - - self.attn2 = Attention( - query_dim=dim, - cross_attention_dim=cross_attention_dim if not double_self_attention else None, - heads=num_attention_heads, - dim_head=attention_head_dim, - dropout=dropout, - bias=attention_bias, - upcast_attention=upcast_attention, - out_bias=attention_out_bias, - attention_mode=attention_mode, - downsampler=False, - use_rope=False, - interpolation_scale_thw=interpolation_scale_thw, - ) # is self-attn if encoder_hidden_states is none - else: - self.norm2 = None - self.attn2 = None - - # 3. Feed-forward - if norm_type == "ada_norm_continuous": - self.norm3 = AdaLayerNormContinuous( - dim, - ada_norm_continous_conditioning_embedding_dim, - norm_elementwise_affine, - norm_eps, - ada_norm_bias, - "layer_norm", - ) - - elif norm_type in ["ada_norm_zero", "ada_norm", "layer_norm", "ada_norm_continuous"]: - self.norm3 = nn.LayerNorm(dim, norm_eps, norm_elementwise_affine) - elif norm_type == "layer_norm_i2vgen": - self.norm3 = None - - if downsampler: - downsampler_ker_size = list(re.search(r'k(\d{2,3})', downsampler).group(1)) # 122 - # if len(downsampler_ker_size) == 3: - # self.ff = FeedForward_Conv3d( - # downsampler, - # dim, - # 2 * dim, - # bias=ff_bias, - # ) - # elif len(downsampler_ker_size) == 2: - self.ff = FeedForward_Conv2d( - downsampler, - dim, - 2 * dim, - bias=ff_bias, - ) - else: - self.ff = FeedForward( - dim, - dropout=dropout, - activation_fn=activation_fn, - final_dropout=final_dropout, - inner_dim=ff_inner_dim, - bias=ff_bias, - ) - - # 4. Fuser - if attention_type == "gated" or attention_type == "gated-text-image": - self.fuser = GatedSelfAttentionDense(dim, cross_attention_dim, num_attention_heads, attention_head_dim) - - # 5. Scale-shift for PixArt-Alpha. - if norm_type == "ada_norm_single": - self.scale_shift_table = nn.Parameter(torch.randn(6, dim) / dim**0.5) - - # let chunk size default to None - self._chunk_size = None - self._chunk_dim = 0 - - def set_chunk_feed_forward(self, chunk_size: Optional[int], dim: int = 0): - # Sets chunk feed-forward - self._chunk_size = chunk_size - self._chunk_dim = dim - - def forward( - self, - hidden_states: torch.FloatTensor, - attention_mask: Optional[torch.FloatTensor] = None, - encoder_hidden_states: Optional[torch.FloatTensor] = None, - encoder_attention_mask: Optional[torch.FloatTensor] = None, - timestep: Optional[torch.LongTensor] = None, - cross_attention_kwargs: Dict[str, Any] = None, - class_labels: Optional[torch.LongTensor] = None, - frame: int = None, - height: int = None, - width: int = None, - added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None, - ) -> torch.FloatTensor: - if cross_attention_kwargs is not None: - if cross_attention_kwargs.get("scale", None) is not None: - logger.warning("Passing `scale` to `cross_attention_kwargs` is deprecated. `scale` will be ignored.") - - # Notice that normalization is always applied before the real computation in the following blocks. - # 0. Self-Attention - batch_size = hidden_states.shape[0] - - # import ipdb;ipdb.set_trace() - if self.norm_type == "ada_norm": - norm_hidden_states = self.norm1(hidden_states, timestep) - elif self.norm_type == "ada_norm_zero": - norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1( - hidden_states, timestep, class_labels, hidden_dtype=hidden_states.dtype - ) - elif self.norm_type in ["layer_norm", "layer_norm_i2vgen"]: - norm_hidden_states = self.norm1(hidden_states) - elif self.norm_type == "ada_norm_continuous": - norm_hidden_states = self.norm1(hidden_states, added_cond_kwargs["pooled_text_emb"]) - elif self.norm_type == "ada_norm_single": - # import ipdb;ipdb.set_trace() - if get_sequence_parallel_state(): - batch_size = hidden_states.shape[1] - # print('hidden_states', hidden_states.shape) - # print('timestep', timestep.shape) - shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( - self.scale_shift_table[:, None] + timestep.reshape(6, batch_size, -1) - ).chunk(6, dim=0) - else: - shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( - self.scale_shift_table[None] + timestep.reshape(batch_size, 6, -1) - ).chunk(6, dim=1) - norm_hidden_states = self.norm1(hidden_states) - norm_hidden_states = norm_hidden_states * (1 + scale_msa) + shift_msa - # norm_hidden_states = norm_hidden_states.squeeze(1) - else: - raise ValueError("Incorrect norm used") - - if self.pos_embed is not None: - norm_hidden_states = self.pos_embed(norm_hidden_states) - - # 1. Prepare GLIGEN inputs - cross_attention_kwargs = cross_attention_kwargs.copy() if cross_attention_kwargs is not None else {} - gligen_kwargs = cross_attention_kwargs.pop("gligen", None) - - attn_output = self.attn1( - norm_hidden_states, - encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None, - attention_mask=attention_mask, frame=frame, height=height, width=width, - **cross_attention_kwargs, - ) - if self.norm_type == "ada_norm_zero": - attn_output = gate_msa.unsqueeze(1) * attn_output - elif self.norm_type == "ada_norm_single": - attn_output = gate_msa * attn_output - - hidden_states = attn_output + hidden_states - if hidden_states.ndim == 4: - hidden_states = hidden_states.squeeze(1) - - # 1.2 GLIGEN Control - if gligen_kwargs is not None: - hidden_states = self.fuser(hidden_states, gligen_kwargs["objs"]) - - # 3. Cross-Attention - if self.attn2 is not None: - if self.norm_type == "ada_norm": - norm_hidden_states = self.norm2(hidden_states, timestep) - elif self.norm_type in ["ada_norm_zero", "layer_norm", "layer_norm_i2vgen"]: - norm_hidden_states = self.norm2(hidden_states) - elif self.norm_type == "ada_norm_single": - # For PixArt norm2 isn't applied here: - # https://github.com/PixArt-alpha/PixArt-alpha/blob/0f55e922376d8b797edd44d25d0e7464b260dcab/diffusion/model/nets/PixArtMS.py#L70C1-L76C103 - norm_hidden_states = hidden_states - elif self.norm_type == "ada_norm_continuous": - norm_hidden_states = self.norm2(hidden_states, added_cond_kwargs["pooled_text_emb"]) - else: - raise ValueError("Incorrect norm") - - if self.pos_embed is not None and self.norm_type != "ada_norm_single": - norm_hidden_states = self.pos_embed(norm_hidden_states) - - attn_output = self.attn2( - norm_hidden_states, - encoder_hidden_states=encoder_hidden_states, - attention_mask=encoder_attention_mask, - **cross_attention_kwargs, - ) - hidden_states = attn_output + hidden_states - - # 4. Feed-forward - # i2vgen doesn't have this norm 🤷‍♂️ - if self.norm_type == "ada_norm_continuous": - norm_hidden_states = self.norm3(hidden_states, added_cond_kwargs["pooled_text_emb"]) - elif not self.norm_type == "ada_norm_single": - norm_hidden_states = self.norm3(hidden_states) - - if self.norm_type == "ada_norm_zero": - norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None] - - if self.norm_type == "ada_norm_single": - norm_hidden_states = self.norm2(hidden_states) - norm_hidden_states = norm_hidden_states * (1 + scale_mlp) + shift_mlp - - # if self._chunk_size is not None: - # # "feed_forward_chunk_size" can be used to save memory - # ff_output = _chunked_feed_forward(self.ff, norm_hidden_states, self._chunk_dim, self._chunk_size) - # else: - - if self.downsampler: - ff_output = self.ff(norm_hidden_states, t=frame, h=height, w=width) - else: - ff_output = self.ff(norm_hidden_states) - - if self.norm_type == "ada_norm_zero": - ff_output = gate_mlp.unsqueeze(1) * ff_output - elif self.norm_type == "ada_norm_single": - ff_output = gate_mlp * ff_output - - hidden_states = ff_output + hidden_states - if hidden_states.ndim == 4: - hidden_states = hidden_states.squeeze(1) - - return hidden_states diff --git a/opensora/models/diffusion/opensora/rope.py b/opensora/models/diffusion/opensora/rope.py deleted file mode 100644 index f43dcae04..000000000 --- a/opensora/models/diffusion/opensora/rope.py +++ /dev/null @@ -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 \ No newline at end of file diff --git a/opensora/models/diffusion/utils/curope/__init__.py b/opensora/models/diffusion/utils/curope/__init__.py deleted file mode 100644 index 25e3d48a1..000000000 --- a/opensora/models/diffusion/utils/curope/__init__.py +++ /dev/null @@ -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 diff --git a/opensora/models/diffusion/utils/curope/curope.cpp b/opensora/models/diffusion/utils/curope/curope.cpp deleted file mode 100644 index 8fe9058e0..000000000 --- a/opensora/models/diffusion/utils/curope/curope.cpp +++ /dev/null @@ -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 - -// 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(); - auto pos = positions.accessor(); - - 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"); -} diff --git a/opensora/models/diffusion/utils/curope/curope2d.py b/opensora/models/diffusion/utils/curope/curope2d.py deleted file mode 100644 index a49c12f8c..000000000 --- a/opensora/models/diffusion/utils/curope/curope2d.py +++ /dev/null @@ -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 \ No newline at end of file diff --git a/opensora/models/diffusion/utils/curope/kernels.cu b/opensora/models/diffusion/utils/curope/kernels.cu deleted file mode 100644 index 7156cd1bb..000000000 --- a/opensora/models/diffusion/utils/curope/kernels.cu +++ /dev/null @@ -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 -#include -#include -#include - -#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 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 <<>> ( - //tokens.data_ptr(), - tokens.packed_accessor32(), - pos.data_ptr(), - base, fwd); //, N, H, D ); - })); -} diff --git a/opensora/models/diffusion/utils/curope/setup.py b/opensora/models/diffusion/utils/curope/setup.py deleted file mode 100644 index 230632ed0..000000000 --- a/opensora/models/diffusion/utils/curope/setup.py +++ /dev/null @@ -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 - }) diff --git a/opensora/models/diffusion/utils/pos_embed.py b/opensora/models/diffusion/utils/pos_embed.py deleted file mode 100644 index 61872bc97..000000000 --- a/opensora/models/diffusion/utils/pos_embed.py +++ /dev/null @@ -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 \ No newline at end of file diff --git a/opensora/models/diffusion/utils/transport/__init__.py b/opensora/models/diffusion/utils/transport/__init__.py deleted file mode 100644 index db68edd20..000000000 --- a/opensora/models/diffusion/utils/transport/__init__.py +++ /dev/null @@ -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 \ No newline at end of file diff --git a/opensora/models/diffusion/utils/transport/integrators.py b/opensora/models/diffusion/utils/transport/integrators.py deleted file mode 100644 index adf7c7b4c..000000000 --- a/opensora/models/diffusion/utils/transport/integrators.py +++ /dev/null @@ -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 \ No newline at end of file diff --git a/opensora/models/diffusion/utils/transport/path.py b/opensora/models/diffusion/utils/transport/path.py deleted file mode 100644 index 156a7b0de..000000000 --- a/opensora/models/diffusion/utils/transport/path.py +++ /dev/null @@ -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)) \ No newline at end of file diff --git a/opensora/models/diffusion/utils/transport/transport.py b/opensora/models/diffusion/utils/transport/transport.py deleted file mode 100644 index 396c516cf..000000000 --- a/opensora/models/diffusion/utils/transport/transport.py +++ /dev/null @@ -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 \ No newline at end of file diff --git a/opensora/models/diffusion/utils/transport/utils.py b/opensora/models/diffusion/utils/transport/utils.py deleted file mode 100644 index 446460355..000000000 --- a/opensora/models/diffusion/utils/transport/utils.py +++ /dev/null @@ -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 "