From ec81e2427c6d0c616ed45693e171a7b860feabf5 Mon Sep 17 00:00:00 2001 From: qqingzheng <2533221180@qq.com> Date: Wed, 27 Mar 2024 12:49:30 +0000 Subject: [PATCH] add causalvae Former-commit-id: ac936276fd3213ee1815dfe096cd64cecb601428 --- .gitignore | 2 + examples/get_latents_std.py | 38 ++ examples/rec_image.py | 57 ++ .../rec_video.py => examples/rec_video_ae.py | 49 +- examples/rec_video_vae.py | 193 ++++++ opensora/models/ae/videobase/__init__.py | 7 +- .../ae/videobase/causal_vae/__init__.py | 4 + .../causal_vae/configuration_causalvae.py | 50 ++ .../videobase/causal_vae/dataset_causalvae.py | 96 +++ .../causal_vae/modeling_causalvae.py | 407 ++++++++++++ .../videobase/causal_vae/trainer_causalvae.py | 20 + .../causal_vqvae/trainer_causalvqvae.py | 14 +- .../models/ae/videobase/modeling_videobase.py | 10 +- .../videogpt/scripts => modules}/__init__.py | 0 .../models/ae/videobase/modules/attention.py | 132 ++++ opensora/models/ae/videobase/modules/conv.py | 164 +++++ .../ae/videobase/modules/lpips_video.py | 120 ++++ .../models/ae/videobase/modules/normalize.py | 6 + opensora/models/ae/videobase/modules/ops.py | 7 + .../ae/videobase/modules/perceptual_loss.py | 91 +++ opensora/models/ae/videobase/modules/quant.py | 100 +++ .../ae/videobase/modules/resnet_block.py | 134 ++++ .../ae/videobase/modules/updownsample.py | 171 +++++ .../models/ae/videobase/trainer_videobase.py | 26 + .../ae/videobase/vqvae/trainer_vqvae.py | 13 +- .../videobase/vqvae/videogpt/requirements.txt | 9 - .../vqvae/videogpt/scripts/compute_fvd.py | 131 ---- .../videogpt/scripts/preprocess/__init__.py | 0 .../scripts/preprocess/bair/__init__.py | 0 .../preprocess/bair/bair_extract_images.py | 66 -- .../preprocess/bair/bair_image_to_hdf5.py | 50 -- .../preprocess/bair/create_bair_dataset.sh | 11 - .../scripts/preprocess/ucf101/__init__.py | 0 .../preprocess/ucf101/create_ucf_dataset.sh | 19 - .../preprocess/ucf101/ucf_split_train_test.py | 24 - .../vqvae/videogpt/scripts/sample_videogpt.py | 30 - .../vqvae/videogpt/scripts/train_videogpt.py | 44 -- .../vqvae/videogpt/scripts/train_vqvae.py | 40 -- .../ae/videobase/vqvae/videogpt/setup.py | 16 - .../vqvae/videogpt/videogpt/__init__.py | 6 - .../vqvae/videogpt/videogpt/attention.py | 585 ------------------ .../videobase/vqvae/videogpt/videogpt/data.py | 213 ------- .../vqvae/videogpt/videogpt/download.py | 60 -- .../vqvae/videogpt/videogpt/fvd/__init__.py | 0 .../videogpt/fvd/convert_tf_pretrained.py | 127 ---- .../vqvae/videogpt/videogpt/fvd/fvd.py | 85 --- .../videogpt/videogpt/fvd/pytorch_i3d.py | 322 ---------- .../videobase/vqvae/videogpt/videogpt/gpt.py | 185 ------ .../vqvae/videogpt/videogpt/resnet.py | 149 ----- .../vqvae/videogpt/videogpt/utils.py | 87 --- .../vqvae/videogpt/videogpt/vqvae.py | 324 ---------- opensora/train/train_causalvae.py | 104 ++++ opensora/utils/taming_download.py | 145 +++++ 53 files changed, 2115 insertions(+), 2628 deletions(-) create mode 100644 examples/get_latents_std.py create mode 100644 examples/rec_image.py rename opensora/models/ae/videobase/vqvae/videogpt/rec_video.py => examples/rec_video_ae.py (73%) create mode 100644 examples/rec_video_vae.py create mode 100644 opensora/models/ae/videobase/causal_vae/__init__.py create mode 100644 opensora/models/ae/videobase/causal_vae/configuration_causalvae.py create mode 100644 opensora/models/ae/videobase/causal_vae/dataset_causalvae.py create mode 100644 opensora/models/ae/videobase/causal_vae/modeling_causalvae.py create mode 100644 opensora/models/ae/videobase/causal_vae/trainer_causalvae.py rename opensora/models/ae/videobase/{vqvae/videogpt/scripts => modules}/__init__.py (100%) create mode 100644 opensora/models/ae/videobase/modules/attention.py create mode 100644 opensora/models/ae/videobase/modules/conv.py create mode 100644 opensora/models/ae/videobase/modules/lpips_video.py create mode 100644 opensora/models/ae/videobase/modules/normalize.py create mode 100644 opensora/models/ae/videobase/modules/ops.py create mode 100644 opensora/models/ae/videobase/modules/perceptual_loss.py create mode 100644 opensora/models/ae/videobase/modules/quant.py create mode 100644 opensora/models/ae/videobase/modules/resnet_block.py create mode 100644 opensora/models/ae/videobase/modules/updownsample.py create mode 100644 opensora/models/ae/videobase/trainer_videobase.py delete mode 100644 opensora/models/ae/videobase/vqvae/videogpt/requirements.txt delete mode 100644 opensora/models/ae/videobase/vqvae/videogpt/scripts/compute_fvd.py delete mode 100644 opensora/models/ae/videobase/vqvae/videogpt/scripts/preprocess/__init__.py delete mode 100644 opensora/models/ae/videobase/vqvae/videogpt/scripts/preprocess/bair/__init__.py delete mode 100644 opensora/models/ae/videobase/vqvae/videogpt/scripts/preprocess/bair/bair_extract_images.py delete mode 100644 opensora/models/ae/videobase/vqvae/videogpt/scripts/preprocess/bair/bair_image_to_hdf5.py delete mode 100644 opensora/models/ae/videobase/vqvae/videogpt/scripts/preprocess/bair/create_bair_dataset.sh delete mode 100644 opensora/models/ae/videobase/vqvae/videogpt/scripts/preprocess/ucf101/__init__.py delete mode 100644 opensora/models/ae/videobase/vqvae/videogpt/scripts/preprocess/ucf101/create_ucf_dataset.sh delete mode 100644 opensora/models/ae/videobase/vqvae/videogpt/scripts/preprocess/ucf101/ucf_split_train_test.py delete mode 100644 opensora/models/ae/videobase/vqvae/videogpt/scripts/sample_videogpt.py delete mode 100644 opensora/models/ae/videobase/vqvae/videogpt/scripts/train_videogpt.py delete mode 100644 opensora/models/ae/videobase/vqvae/videogpt/scripts/train_vqvae.py delete mode 100644 opensora/models/ae/videobase/vqvae/videogpt/setup.py delete mode 100644 opensora/models/ae/videobase/vqvae/videogpt/videogpt/__init__.py delete mode 100644 opensora/models/ae/videobase/vqvae/videogpt/videogpt/attention.py delete mode 100644 opensora/models/ae/videobase/vqvae/videogpt/videogpt/data.py delete mode 100644 opensora/models/ae/videobase/vqvae/videogpt/videogpt/download.py delete mode 100644 opensora/models/ae/videobase/vqvae/videogpt/videogpt/fvd/__init__.py delete mode 100644 opensora/models/ae/videobase/vqvae/videogpt/videogpt/fvd/convert_tf_pretrained.py delete mode 100644 opensora/models/ae/videobase/vqvae/videogpt/videogpt/fvd/fvd.py delete mode 100644 opensora/models/ae/videobase/vqvae/videogpt/videogpt/fvd/pytorch_i3d.py delete mode 100644 opensora/models/ae/videobase/vqvae/videogpt/videogpt/gpt.py delete mode 100644 opensora/models/ae/videobase/vqvae/videogpt/videogpt/resnet.py delete mode 100644 opensora/models/ae/videobase/vqvae/videogpt/videogpt/utils.py delete mode 100644 opensora/models/ae/videobase/vqvae/videogpt/videogpt/vqvae.py create mode 100644 opensora/train/train_causalvae.py create mode 100644 opensora/utils/taming_download.py diff --git a/.gitignore b/.gitignore index 27b4cead2..3cd8b087b 100644 --- a/.gitignore +++ b/.gitignore @@ -10,3 +10,5 @@ build/ opensora.egg-info/ wandb/ .idea +*.ipynb +*.jpg diff --git a/examples/get_latents_std.py b/examples/get_latents_std.py new file mode 100644 index 000000000..5f9937024 --- /dev/null +++ b/examples/get_latents_std.py @@ -0,0 +1,38 @@ +import torch +from torch.utils.data import DataLoader, Subset +import sys +sys.path.append(".") +from opensora.models.ae.videobase import CausalVAEModel, CausalVAEDataset + +num_workers = 4 +batch_size = 12 + +torch.manual_seed(0) +torch.set_grad_enabled(False) + +device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') + +pretrained_model_name_or_path = 'results/causalvae/checkpoint-26000' +data_path = '/remote-home1/dataset/UCF-101' +video_num_frames = 17 +resolution = 128 +sample_rate = 10 + +vae = CausalVAEModel.load_from_checkpoint(pretrained_model_name_or_path) +vae.to(device) + +dataset = CausalVAEDataset(data_path, sequence_length=video_num_frames, resolution=resolution, sample_rate=sample_rate) +subset_indices = list(range(1000)) +subset_dataset = Subset(dataset, subset_indices) +loader = DataLoader(subset_dataset, batch_size=8, pin_memory=True) + +all_latents = [] +for video_data in loader: + video_data = video_data['video'].to(device) + latents = vae.encode(video_data).sample() + all_latents.append(video_data.cpu()) + +all_latents_tensor = torch.cat(all_latents) +std = all_latents_tensor.std().item() +normalizer = 1 / std +print(f'{normalizer = }') \ No newline at end of file diff --git a/examples/rec_image.py b/examples/rec_image.py new file mode 100644 index 000000000..a8e95cb30 --- /dev/null +++ b/examples/rec_image.py @@ -0,0 +1,57 @@ +import sys +sys.path.append(".") +from PIL import Image +import torch +from torchvision.transforms import ToTensor, Compose, Resize, Normalize +from torch.nn import functional as F +from opensora.models.ae.videobase import CausalVAEModel +import argparse +import numpy as np + +def preprocess(video_data: torch.Tensor, short_size: int = 128) -> torch.Tensor: + transform = Compose( + [ + ToTensor(), + Normalize((0.5), (0.5)), + Resize(size=short_size), + ] + ) + outputs = transform(video_data) + outputs = outputs.unsqueeze(0).unsqueeze(2) + return outputs + +def main(args: argparse.Namespace): + image_path = args.image_path + resolution = args.resolution + device = args.device + + vqvae = CausalVAEModel.load_from_checkpoint(args.ckpt) + vqvae.eval() + vqvae = vqvae.to(device) + + with torch.no_grad(): + x_vae = preprocess(Image.open(image_path), resolution) + x_vae = x_vae.to(device) + latents = vqvae.encode(x_vae) + recon = vqvae.decode(latents.sample()) + x = recon[0, :, 0, :, :] + x = x.squeeze() + x = x.detach().cpu().numpy() + x = np.clip(x, -1, 1) + x = (x + 1) / 2 + x = (255*x).astype(np.uint8) + x = x.transpose(1,2,0) + image = Image.fromarray(x) + image.save(args.rec_path) + + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument('--image-path', type=str, default='') + parser.add_argument('--rec-path', type=str, default='') + parser.add_argument('--ckpt', type=str, default='') + parser.add_argument('--resolution', type=int, default=336) + parser.add_argument('--device', type=str, default='cuda') + + args = parser.parse_args() + main(args) diff --git a/opensora/models/ae/videobase/vqvae/videogpt/rec_video.py b/examples/rec_video_ae.py similarity index 73% rename from opensora/models/ae/videobase/vqvae/videogpt/rec_video.py rename to examples/rec_video_ae.py index 5bc2d2f19..5acbeb6fb 100644 --- a/opensora/models/ae/videobase/vqvae/videogpt/rec_video.py +++ b/examples/rec_video_ae.py @@ -1,20 +1,26 @@ import random +import argparse +from typing import Optional + import cv2 +import imageio 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.io import read_video from torchvision.transforms import Lambda, Compose from torchvision.transforms._transforms_video import RandomCropVideo -from torch.nn import functional as F -from videogpt import load_vqvae -from videogpt.data import preprocess -import argparse -def array_to_video(image_array, fps=30.0, output_file='output_video.mp4'): +import sys +sys.path.append(".") +from opensora.models.ae import VQVAEModel + + +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') + fourcc = cv2.VideoWriter_fourcc(*'mp4v') # type: ignore video_writer = cv2.VideoWriter(output_file, fourcc, float(fps), (width, height)) for image in image_array: @@ -23,16 +29,17 @@ def array_to_video(image_array, fps=30.0, output_file='output_video.mp4'): video_writer.release() -def custom_to_video(x, fps=2.0, output_file='output_video.mp4'): +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, -0.5, 0.5) x = (x + 0.5) x = x.permute(1, 2, 3, 0).numpy() # (C, T, H, W) -> (T, H, W, C) x = (255*x).astype(np.uint8) - array_to_video(x, fps=fps, output_file=output_file) + # array_to_video(x, fps=fps, output_file=output_file) + imageio.mimwrite(output_file, x, fps=fps, quality=9) return -def read_video(video_path, num_frames, sample_rate): +def read_video(video_path: str, num_frames: int, sample_rate: int) -> torch.Tensor: decord_vr = VideoReader(video_path, ctx=cpu(0)) total_frames = len(decord_vr) sample_frames_len = sample_rate * num_frames @@ -55,7 +62,7 @@ def read_video(video_path, num_frames, sample_rate): video_data = video_data.permute(3, 0, 1, 2) # (T, H, W, C) -> (C, T, H, W) return video_data -def preprocess(video_data, short_size=128, crop_size=None): +def preprocess(video_data: torch.Tensor, short_size: int = 128, crop_size: Optional[int] = None) -> torch.Tensor: transform = Compose( [ @@ -74,7 +81,7 @@ def preprocess(video_data, short_size=128, crop_size=None): return video_outputs -def main(args): +def main(args: argparse.Namespace): video_path = args.video_path num_frames = args.num_frames resolution = args.resolution @@ -82,26 +89,28 @@ def main(args): sample_fps = args.sample_fps sample_rate = args.sample_rate device = torch.device('cuda') - - vqvae = load_vqvae(args.ckpt, root='./').to(device) - - - x_vae = preprocess(read_video(video_path, num_frames, sample_rate), resolution, crop_size) - x_vae = x_vae.to(device) + if args.ckpt in ['bair_stride4x2x2', 'ucf101_stride4x4x4', 'kinetics_stride4x4x4', 'kinetics_stride2x4x4']: + vqvae = VQVAEModel.download_and_load_model(args.ckpt) + else: + vqvae = VQVAEModel.load_from_checkpoint(args.ckpt) + vqvae.eval() + vqvae = vqvae.to(device) with torch.no_grad(): + x_vae = preprocess(read_video(video_path, num_frames, sample_rate), resolution, crop_size) + x_vae = x_vae.to(device) encodings, embeddings = vqvae.encode(x_vae, include_embeddings=True) video_recon = vqvae.decode(encodings) # custom_to_video(x_vae[0], fps=sample_fps/sample_rate, output_file='origin_input.mp4') custom_to_video(video_recon[0], fps=sample_fps/sample_rate, output_file=args.rec_path) + if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument('--video-path', type=str, default='') parser.add_argument('--rec-path', type=str, default='') - parser.add_argument('--ckpt', type=str, default='ucf101_stride4x4x4', - choices=['bair_stride4x2x2', 'ucf101_stride4x4x4', 'kinetics_stride4x4x4', 'kinetics_stride2x4x4']) + parser.add_argument('--ckpt', type=str, default='ucf101_stride4x4x4') 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) diff --git a/examples/rec_video_vae.py b/examples/rec_video_vae.py new file mode 100644 index 000000000..d528f4136 --- /dev/null +++ b/examples/rec_video_vae.py @@ -0,0 +1,193 @@ +import random +import argparse +from typing import Optional + +import cv2 +import numpy as np +import numpy.typing as npt +import torch +from decord import VideoReader, cpu +from torch.nn import functional as F +from pytorchvideo.transforms import ShortSideScale +from torchvision.transforms import Lambda, Compose +from torchvision.transforms._transforms_video import RandomCropVideo +import sys +from tqdm import tqdm + +sys.path.append(".") +from opensora.models.ae.videobase import CausalVAEModel +from opensora.models.ae.videobase.modules.conv import BlockwiseConv3d +import torch.nn as nn +import math + +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).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 + + +def preprocess( + video_data: torch.Tensor, short_size: int = 128, crop_size: Optional[int] = None +) -> torch.Tensor: + + transform = Compose( + [ + Lambda(lambda x: ((x / 255.0) * 2 - 1)), + ShortSideScale(size=short_size), + ( + RandomCropVideo(size=crop_size) + if crop_size is not None + else Lambda(lambda x: x) + ), + ] + ) + + video_outputs = transform(video_data) + video_outputs = torch.unsqueeze(video_outputs, 0) + + return video_outputs + + +def replace_conv3d_with_blockwise(model): + for name, module in model.named_children(): + if isinstance(module, nn.Conv3d): + new_module = BlockwiseConv3d( + module.in_channels, + module.out_channels, + module.kernel_size, + module.stride, + module.padding, + (4, 64, 64), + ) + new_module.set_weights( + module.weight.data, + module.bias.data if module.bias is not None else None, + ) + setattr(model, name, new_module) + else: + replace_conv3d_with_blockwise(module) + +def process_in_chunks(video_data: torch.Tensor, model: nn.Module, chunk_size: int, overlap: int, device: str): + assert (chunk_size + overlap - 1) % 4 == 0 + num_frames = video_data.size(2) + output_chunks = [] + + start = 0 + while start < num_frames: + end = min(start + chunk_size, num_frames) + if start + chunk_size + overlap < num_frames: + end += overlap + chunk = video_data[:, :, start:end, :, :] + + with torch.no_grad(): + chunk = chunk.half().to(device) + latents = model.encode(chunk) + recon_chunk = model.decode(latents.sample().half()).cpu().float() + + if output_chunks: + overlap_step = min(overlap, recon_chunk.shape[2]) + overlap_tensor = output_chunks[-1][:, :, -overlap_step:] * 1 / 4 + recon_chunk[:, :, :overlap_step] * 3 / 4 + output_chunks[-1] = torch.cat((output_chunks[-1][:, :, :-overlap], overlap_tensor), dim=2) + if end < num_frames: + output_chunks.append(recon_chunk[:, :, overlap:]) + else: + output_chunks.append(recon_chunk[:, :, :, :, :]) + else: + output_chunks.append(recon_chunk) + start += chunk_size + return torch.cat(output_chunks, dim=2) + +def main(args: argparse.Namespace): + video_path = args.video_path + num_frames = args.num_frames + resolution = args.resolution + crop_size = args.crop_size + sample_fps = args.sample_fps + sample_rate = args.sample_rate + + device = args.device + torch.cuda.reset_max_memory_allocated() + start_memory = torch.cuda.memory_allocated() / 2**30 + vqvae = CausalVAEModel.load_from_checkpoint(args.ckpt) + # replace_conv3d_with_blockwise(vqvae) + vqvae = vqvae.half().to(device) + with torch.no_grad(): + x_vae = preprocess( + read_video(video_path, num_frames, sample_rate), resolution, crop_size + ) + print(x_vae.shape) + if x_vae.size(2) > 32: + video_recon = process_in_chunks(x_vae, vqvae, chunk_size=4, overlap=1, device=device) + else: + x_vae = x_vae.half().to(device) + latents = vqvae.encode(x_vae) + video_recon = vqvae.decode(latents.sample().half()) + + peak_memory_usage = torch.cuda.max_memory_allocated() / 2**30 + print(peak_memory_usage - start_memory) + print(video_recon[0].shape) + custom_to_video( + video_recon[0], fps=sample_fps / sample_rate, output_file=args.rec_path + ) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--video-path", type=str, default="") + parser.add_argument("--rec-path", type=str, default="") + parser.add_argument("--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=100) + parser.add_argument("--sample-rate", type=int, default=1) + parser.add_argument("--device", type=str, default="cuda") + + args = parser.parse_args() + main(args) diff --git a/opensora/models/ae/videobase/__init__.py b/opensora/models/ae/videobase/__init__.py index 9cceadd62..4ec7b431c 100644 --- a/opensora/models/ae/videobase/__init__.py +++ b/opensora/models/ae/videobase/__init__.py @@ -13,7 +13,12 @@ from .causal_vqvae import ( CausalVQVAETrainer, CausalVQVAEModel, CausalVQVAEModelWrapper ) - +from .causal_vae import ( + CausalVAETrainer, + CausalVAEConfiguration, + CausalVAEDataset, + CausalVAEModel +) """ old code """ diff --git a/opensora/models/ae/videobase/causal_vae/__init__.py b/opensora/models/ae/videobase/causal_vae/__init__.py new file mode 100644 index 000000000..dedb33555 --- /dev/null +++ b/opensora/models/ae/videobase/causal_vae/__init__.py @@ -0,0 +1,4 @@ +from .configuration_causalvae import CausalVAEConfiguration +from .dataset_causalvae import CausalVAEDataset +from .modeling_causalvae import CausalVAEModel +from .trainer_causalvae import CausalVAETrainer diff --git a/opensora/models/ae/videobase/causal_vae/configuration_causalvae.py b/opensora/models/ae/videobase/causal_vae/configuration_causalvae.py new file mode 100644 index 000000000..2937728a7 --- /dev/null +++ b/opensora/models/ae/videobase/causal_vae/configuration_causalvae.py @@ -0,0 +1,50 @@ +from ..configuration_videobase import VideoBaseConfiguration +from typing import Union, Tuple + +class CausalVAEConfiguration(VideoBaseConfiguration): + + def __init__( + self, + in_channels: int = 3, + out_channels: int = 3, + hidden_size: int = 128, + z_channels: int = 16, + ch_mult: Tuple[int] = (1,1,2,2,4), + num_res_block: int = 2, + attn_resolutions: Tuple[int] = [16], + dropout: float = 0.0, + resolution: int = 256, + attn_type: str = "vanilla3D", + use_linear_attn: bool = False, + embed_dim: int = 16, + time_compress: int = 2, + # ---- for loss ---- + logvar_init: float = 0., + kl_weight: float = 1e-6, + pixelloss_weight: int = 1, + perceptual_weight: int = 1, + disc_loss: str = "hinge", + **kwargs + ) -> None: + + super().__init__(**kwargs) + self.hidden_size = hidden_size + self.in_channels = in_channels + self.out_channels = out_channels + self.z_channels = z_channels + self.ch_mult = ch_mult + self.num_res_block = num_res_block + self.attn_resolutions = attn_resolutions + self.dropout = dropout + self.resolution = resolution + self.attn_type = attn_type + self.use_linear_attn = use_linear_attn + self.embed_dim = embed_dim + self.time_compress = time_compress + + # ---- for loss ---- + self.logvar_init = logvar_init + self.kl_weight = kl_weight + self.pixelloss_weight = pixelloss_weight + self.perceptual_weight = perceptual_weight + self.disc_loss = disc_loss \ No newline at end of file diff --git a/opensora/models/ae/videobase/causal_vae/dataset_causalvae.py b/opensora/models/ae/videobase/causal_vae/dataset_causalvae.py new file mode 100644 index 000000000..71e2fb147 --- /dev/null +++ b/opensora/models/ae/videobase/causal_vae/dataset_causalvae.py @@ -0,0 +1,96 @@ +import os.path as osp +import math +import pickle +import warnings +import glob +from PIL import Image +import torch.utils.data as data +import torch.distributed as dist +import torch.nn.functional as F +from torchvision.datasets.video_utils import VideoClips + +class CausalVAEDataset(data.Dataset): + video_exts = ['avi', 'mp4', 'webm'] + def __init__(self, video_folder, sequence_length, sample_rate=1, train=True, resolution=64): + super().__init__() + self.train = train + self.sequence_length = sequence_length + self.resolution = resolution + + video_files = [] + for data_folder in [video_folder]: + if data_folder is None: + continue + folder = data_folder + video_files += sum([glob.glob(osp.join(folder, '**', f'*.{ext}'), recursive=True) + for ext in self.video_exts], []) + + warnings.filterwarnings('ignore') + cache_file = osp.join(folder, f"metadata_{sequence_length}_{sample_rate}.pkl") + if dist.is_initialized() and dist.get_rank() != 0: + dist.barrier() + if not osp.exists(cache_file): + clips = VideoClips(video_files, sequence_length, frame_rate=sample_rate, num_workers=32) + pickle.dump(clips.metadata, open(cache_file, 'wb')) + else: + metadata = pickle.load(open(cache_file, 'rb')) + clips = VideoClips(video_files, sequence_length, frame_rate=sample_rate, + _precomputed_metadata=metadata) + if dist.is_initialized() and dist.get_rank() == 0: + dist.barrier() + self._clips = clips + self._clips_num = self._clips.num_clips() + + @property + def n_classes(self): + return len(self.classes) + + def __len__(self): + return self._clips_num + + def __getitem__(self, idx): + resolution = self.resolution + if idx < self._clips_num: + video, _, _, idx = self._clips.get_clip(idx) + video = preprocess(video, resolution) + class_name = get_parent_dir(self._clips.video_paths[idx]) + else: + idx -= self._clips_num + image = Image.open(self.image_files[idx]) + video = preprocess_image(image, resolution, self.sequence_length) + # label = self.class_to_label[class_name] + return dict(video=video, label="") + +def get_parent_dir(path): + return osp.basename(osp.dirname(path)) + +def preprocess(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 + video = 2.0 * video - 1 + + return video diff --git a/opensora/models/ae/videobase/causal_vae/modeling_causalvae.py b/opensora/models/ae/videobase/causal_vae/modeling_causalvae.py new file mode 100644 index 000000000..b15921aac --- /dev/null +++ b/opensora/models/ae/videobase/causal_vae/modeling_causalvae.py @@ -0,0 +1,407 @@ +import torch +from torch import nn +from .configuration_causalvae import CausalVAEConfiguration +from ..modules.attention import make_attn +from ..modules.resnet_block import ResnetBlock3D +from ..modules.conv import CausalConv3d +from ..modules.normalize import Normalize +from ..modules.perceptual_loss import LPIPSWithDiscriminator +from ..modules.updownsample import ( + SpatialDownsample2x, + SpatialUpsample2x, + TimeDownsample2x, + TimeUpsample2x, +) +import numpy as np +from ..modules.ops import nonlinearity +from ..modeling_videobase import VideoBaseAE + +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.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.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 + + +class Encoder(nn.Module): + def __init__( + self, + *, + ch, + out_ch, + ch_mult=(1, 2, 4, 8), + num_res_blocks, + attn_resolutions, + dropout=0.0, + resamp_with_conv=True, + in_channels, + resolution, + z_channels, + double_z=True, + use_linear_attn=False, + attn_type="vanilla", + time_compress=2, + **ignore_kwargs, + ): + super().__init__() + if use_linear_attn: + attn_type = "linear" + self.ch = ch + self.temb_ch = 0 + self.time_compress = time_compress + self.num_resolutions = len(ch_mult) + self.num_res_blocks = num_res_blocks + self.resolution = resolution + self.in_channels = in_channels + + # downsampling + self.conv_in = CausalConv3d( + in_channels, self.ch, kernel_size=3, stride=1, padding=1 + ) + + curr_res = resolution + in_ch_mult = (1,) + tuple(ch_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 = ch * in_ch_mult[i_level] + block_out = ch * ch_mult[i_level] + for i_block in range(self.num_res_blocks): + block.append( + ResnetBlock3D( + in_channels=block_in, + out_channels=block_out, + dropout=dropout, + ) + ) + block_in = block_out + if curr_res in attn_resolutions: + attn.append(make_attn(block_in, attn_type=attn_type)) + down = nn.Module() + down.block = block + down.attn = attn + if i_level != self.num_resolutions - 1: + down.downsample = SpatialDownsample2x(block_in, block_in) + curr_res = curr_res // 2 + # if i_level < self.time_compress and i_level != self.num_resolutions - 1: + if i_level < self.time_compress: + down.time_downsample = TimeDownsample2x() + self.down.append(down) + + # middle + self.mid = nn.Module() + self.mid.block_1 = ResnetBlock3D( + in_channels=block_in, + out_channels=block_in, + dropout=dropout, + ) + self.mid.attn_1 = make_attn(block_in, attn_type=attn_type) + self.mid.block_2 = ResnetBlock3D( + in_channels=block_in, + out_channels=block_in, + dropout=dropout, + ) + + # end + self.norm_out = Normalize(block_in) + self.conv_out = CausalConv3d( + block_in, + 2 * z_channels if double_z else z_channels, + kernel_size=3, + stride=1, + padding=1, + ) + + def forward(self, x): + # downsampling + 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 i_level != self.num_resolutions - 1: + hs.append(self.down[i_level].downsample(hs[-1])) + # if i_level < self.time_compress and i_level != self.num_resolutions - 1: + if i_level < self.time_compress: + hs.append(self.down[i_level].time_downsample(hs[-1])) + + # middle + h = hs[-1] + h = self.mid.block_1(h) + h = self.mid.attn_1(h) + h = self.mid.block_2(h) + # end + h = self.norm_out(h) + h = nonlinearity(h) + h = self.conv_out(h) + return h + + +class Decoder(nn.Module): + def __init__( + self, + *, + ch, + out_ch, + ch_mult=(1, 2, 4, 8), + num_res_blocks, + attn_resolutions, + dropout=0.0, + resamp_with_conv=True, + in_channels, + resolution, + z_channels, + give_pre_end=False, + tanh_out=False, + use_linear_attn=False, + attn_type="vanilla", + time_compress=2, + **ignorekwargs, + ): + super().__init__() + if use_linear_attn: + attn_type = "linear" + self.ch = ch + self.time_compress = time_compress + self.temb_ch = 0 + self.num_resolutions = len(ch_mult) + self.num_res_blocks = num_res_blocks + self.resolution = resolution + self.in_channels = in_channels + self.give_pre_end = give_pre_end + self.tanh_out = tanh_out + + # compute in_ch_mult, block_in and curr_res at lowest res + in_ch_mult = (1,) + tuple(ch_mult) + block_in = ch * ch_mult[self.num_resolutions - 1] + curr_res = resolution // 2 ** (self.num_resolutions - 1) + self.z_shape = (1, z_channels, curr_res, curr_res) + print( + "Working with z of shape {} = {} dimensions.".format( + self.z_shape, np.prod(self.z_shape) + ) + ) + + # z to block_in + self.conv_in = CausalConv3d(z_channels, block_in, kernel_size=3, padding=1) + + # middle + self.mid = nn.Module() + self.mid.block_1 = ResnetBlock3D( + in_channels=block_in, + out_channels=block_in, + dropout=dropout, + ) + self.mid.attn_1 = make_attn(block_in, attn_type=attn_type) + self.mid.block_2 = ResnetBlock3D( + in_channels=block_in, + out_channels=block_in, + dropout=dropout, + ) + + # upsampling + self.up = nn.ModuleList() + for i_level in reversed(range(self.num_resolutions)): + block = nn.ModuleList() + attn = nn.ModuleList() + block_out = ch * ch_mult[i_level] + for i_block in range(self.num_res_blocks + 1): + block.append( + ResnetBlock3D( + in_channels=block_in, + out_channels=block_out, + dropout=dropout, + ) + ) + block_in = block_out + if curr_res in attn_resolutions: + attn.append(make_attn(block_in, attn_type=attn_type)) + up = nn.Module() + up.block = block + up.attn = attn + if i_level != 0: + up.upsample = SpatialUpsample2x(block_in, block_in) + curr_res = curr_res * 2 + if i_level > self.num_resolutions - 1 - self.time_compress and i_level != 0: + # if i_level <= self.time_compress and i_level != 0: + up.time_upsample = TimeUpsample2x() + + self.up.insert(0, up) # prepend to get consistent order + + # end + self.norm_out = Normalize(block_in) + self.conv_out = CausalConv3d(block_in, out_ch, kernel_size=3, padding=1) + + def forward(self, z): + self.last_z_shape = z.shape + # z to block_in + h = self.conv_in(z) + # middle + h = self.mid.block_1(h) + h = self.mid.attn_1(h) + h = self.mid.block_2(h) + # upsampling + 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 i_level != 0: + h = self.up[i_level].upsample(h) + + if i_level > self.num_resolutions - 1 - self.time_compress and i_level != 0: + # if i_level <= self.time_compress and i_level != 0: + h = self.up[i_level].time_upsample(h) + + # end + if self.give_pre_end: + return h + + h = self.norm_out(h) + h = nonlinearity(h) + h = self.conv_out(h) + if self.tanh_out: + h = torch.tanh(h) + return h + + +class CausalVAEModel(VideoBaseAE): + CONFIGURATION_CLS = CausalVAEConfiguration + + def __init__(self, config: CausalVAEConfiguration): + super().__init__() + self.config = config + self.loss = LPIPSWithDiscriminator( + logvar_init=config.logvar_init, + kl_weight=config.kl_weight, + pixelloss_weight=config.pixelloss_weight, + perceptual_weight=config.perceptual_weight, + disc_loss=config.disc_loss, + ) + self.loss.logvar.requires_grad = False + self.encoder = Encoder( + ch=config.hidden_size, + out_ch=config.out_channels, + ch_mult=config.ch_mult, + z_channels=config.z_channels, + num_res_blocks=config.num_res_block, + attn_resolutions=config.attn_resolutions, + dropout=config.dropout, + in_channels=config.in_channels, + resolution=config.resolution, + resamp_with_conv=True, + attn_type=config.attn_type, + use_linear_attn=config.use_linear_attn, + time_compress=config.time_compress, + ) + self.decoder = Decoder( + ch=config.hidden_size, + out_ch=config.out_channels, + ch_mult=config.ch_mult, + z_channels=config.z_channels, + num_res_blocks=config.num_res_block, + attn_resolutions=config.attn_resolutions, + dropout=config.dropout, + in_channels=config.in_channels, + resolution=config.resolution, + resamp_with_conv=True, + attn_type=config.attn_type, + use_linear_attn=config.use_linear_attn, + time_compress=config.time_compress, + ) + self.quant_conv = CausalConv3d( + 2 * config.z_channels, 2 * config.embed_dim, 1 + ) + self.post_quant_conv = CausalConv3d( + config.embed_dim, config.z_channels, 1 + ) + self.embed_dim = config.embed_dim + + def encode(self, x): + h = self.encoder(x) + moments = self.quant_conv(h) + posterior = DiagonalGaussianDistribution(moments) + return posterior + + def decode(self, z): + 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 get_input(self, batch, k): + x = batch[k] + if len(x.shape) == 3: + x = x[..., None] + x = x.to(memory_format=torch.contiguous_format).float() + return x + + def get_last_layer(self): + return self.decoder.conv_out.weight + + # def init_from_ckpt(self, path, ignore_keys=list()): + # sd = torch.load(path, map_location="cpu") + # 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] + # self.load_state_dict(sd, strict=False) + # print(f"Restored from {path}") \ No newline at end of file diff --git a/opensora/models/ae/videobase/causal_vae/trainer_causalvae.py b/opensora/models/ae/videobase/causal_vae/trainer_causalvae.py new file mode 100644 index 000000000..3530085be --- /dev/null +++ b/opensora/models/ae/videobase/causal_vae/trainer_causalvae.py @@ -0,0 +1,20 @@ +from ..trainer_videobase import VideoBaseTrainer +from typing import Optional +import os +import torch +from transformers.utils import WEIGHTS_NAME +import json + +class CausalVAETrainer(VideoBaseTrainer): + + def compute_loss(self, model, inputs, return_outputs=False): + model = model.module + x = inputs.get("video") + reconstructions, posterior = model(x) + aeloss, _ = model.loss( + x, + reconstructions, + posterior, + split="train", + ) + return aeloss \ No newline at end of file diff --git a/opensora/models/ae/videobase/causal_vqvae/trainer_causalvqvae.py b/opensora/models/ae/videobase/causal_vqvae/trainer_causalvqvae.py index b60cc9a4b..53519ee76 100644 --- a/opensora/models/ae/videobase/causal_vqvae/trainer_causalvqvae.py +++ b/opensora/models/ae/videobase/causal_vqvae/trainer_causalvqvae.py @@ -1,4 +1,4 @@ -from transformers import Trainer +from ..trainer_videobase import VideoBaseTrainer import torch.nn.functional as F from typing import Optional import os @@ -6,7 +6,7 @@ import torch from transformers.utils import WEIGHTS_NAME import json -class CausalVQVAETrainer(Trainer): +class CausalVQVAETrainer(VideoBaseTrainer): def compute_loss(self, model, inputs, return_outputs=False): model = model.module @@ -18,13 +18,3 @@ class CausalVQVAETrainer(Trainer): commitment_loss = vq_output['commitment_loss'] loss = recon_loss + commitment_loss return loss - - 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() - 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/ae/videobase/modeling_videobase.py b/opensora/models/ae/videobase/modeling_videobase.py index bdcc85c23..9880ac44b 100644 --- a/opensora/models/ae/videobase/modeling_videobase.py +++ b/opensora/models/ae/videobase/modeling_videobase.py @@ -1,5 +1,7 @@ import torch from torch import nn +import os +import json """ We are currently working on the abstraction of the code structure, @@ -13,7 +15,13 @@ class VideoBaseAE(nn.Module): @classmethod def load_from_checkpoint(cls, model_path): - pass + with open(os.path.join(model_path, "config.json"), "r") as file: + config = json.load(file) + state_dict = torch.load(os.path.join(model_path, "pytorch_model.bin"), map_location="cpu") + model = cls(config=cls.CONFIGURATION_CLS(**config)) + model.load_state_dict(state_dict, strict=False) + return model + @classmethod def download_and_load_model(cls, model_name, cache_dir=None): pass diff --git a/opensora/models/ae/videobase/vqvae/videogpt/scripts/__init__.py b/opensora/models/ae/videobase/modules/__init__.py similarity index 100% rename from opensora/models/ae/videobase/vqvae/videogpt/scripts/__init__.py rename to opensora/models/ae/videobase/modules/__init__.py diff --git a/opensora/models/ae/videobase/modules/attention.py b/opensora/models/ae/videobase/modules/attention.py new file mode 100644 index 000000000..b5421e4ba --- /dev/null +++ b/opensora/models/ae/videobase/modules/attention.py @@ -0,0 +1,132 @@ +import torch.nn as nn +from .normalize import Normalize +from ..modules.conv import CausalConv3d +import torch +import numpy as np +from einops import rearrange + + +class LinearAttention(nn.Module): + 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(nn.Module): + 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 AttnBlock(nn.Module): + 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 + ) + + 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_ + + +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/ae/videobase/modules/conv.py b/opensora/models/ae/videobase/modules/conv.py new file mode 100644 index 000000000..e4c240a2d --- /dev/null +++ b/opensora/models/ae/videobase/modules/conv.py @@ -0,0 +1,164 @@ +import torch.nn as nn +from typing import Union, Tuple +import torch.nn.functional as F +import torch +from einops import rearrange, pack, unpack +from .ops import cast_tuple + + +class CausalConv3d(nn.Module): + def __init__( + self, chan_in, chan_out, kernel_size: Union[int, Tuple[int, int, int]], **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._init_weights() + + def _init_weights(self): + ks = torch.tensor(self.kernel_size) + mean_weight = torch.full( + (self.chan_out, self.chan_in, *self.kernel_size), + fill_value=1 / (torch.prod(ks) * self.chan_in), + ) + self.conv.weight = nn.Parameter(mean_weight, requires_grad=True) + if self.conv.bias is not None: + nn.init.constant_(self.conv.bias, 0) + + def forward(self, x): + first_frame_pad = x[:, :, :1, :, :].repeat( + (1, 1, self.time_kernel_size - 1, 1, 1) + ) + x = torch.concatenate((first_frame_pad, x), dim=2) + return self.conv(x) + + +def cal_idx_by_step_begin(kernel_size, stride, step): + return stride * step + +def cal_idx_by_step_end(kernel_size, stride, step): + return kernel_size + (step - 1) * stride + +def get_max_step(kernel_size, stride, length): + return (length - kernel_size) // stride + 1 + +def get_new_size(kernel_size, stride, length): + return (length - kernel_size) // stride + 1 + + +class BlockwiseConv3d(nn.Module): + def __init__( + self, + chan_in, + chan_out, + kernel_size, + stride=(1, 1, 1), + padding=(0, 0, 0), + block_step=(15, 15, 15), + ): + super().__init__() + self.chan_in = chan_in + self.chan_out = chan_out + self.kernel_size = kernel_size + self.padding = padding + self.tile_padding = torch.repeat_interleave(torch.tensor(padding), 2).tolist() + self.tile_padding.reverse() + self.stride = stride + self.block_step = block_step + self.conv = nn.Conv3d(chan_in, chan_out, kernel_size, stride, padding=0) + + def forward(self, x): + device = x.device + batch_size = x.shape[0] + dtype = x.dtype + padded_x = nn.functional.pad(x, self.tile_padding) + del x + time_length = get_new_size( + self.kernel_size[0], self.stride[0], padded_x.shape[2] + ) + height_length = get_new_size( + self.kernel_size[1], self.stride[1], padded_x.shape[3] + ) + width_length = get_new_size( + self.kernel_size[2], self.stride[2], padded_x.shape[4] + ) + + result = torch.zeros( + (batch_size, self.chan_out, time_length, height_length, width_length), device=device, dtype=dtype + ) + + origin_time_length = padded_x.shape[2] + origin_height_length = padded_x.shape[3] + origin_width_length = padded_x.shape[4] + + for time in range( + 0, + get_max_step(self.kernel_size[0], self.stride[0], origin_time_length), + self.block_step[0], + ): + for height in range( + 0, + get_max_step(self.kernel_size[1], self.stride[1], origin_height_length), + self.block_step[1], + ): + for width in range( + 0, + get_max_step( + self.kernel_size[2], self.stride[2], origin_width_length + ), + self.block_step[2], + ): + time_start = cal_idx_by_step_begin( + self.kernel_size[0], self.stride[0], time + ) + height_start = cal_idx_by_step_begin( + self.kernel_size[1], self.stride[1], height + ) + width_start = cal_idx_by_step_begin( + self.kernel_size[2], self.stride[2], width + ) + if ( + time_start >= origin_time_length + or height_start >= origin_height_length + or width_start >= origin_width_length + ): + continue + time_end = cal_idx_by_step_end( + self.kernel_size[0], self.stride[0], time + self.block_step[0] + ) + height_end = cal_idx_by_step_end( + self.kernel_size[1], self.stride[1], height + self.block_step[1] + ) + width_end = cal_idx_by_step_end( + self.kernel_size[2], self.stride[2], width + self.block_step[2] + ) + block = padded_x[ + :, + :, + time_start : min(time_end, origin_time_length), + height_start : min(height_end, origin_height_length), + width_start : min(width_end, origin_width_length), + ] + result_block = self.conv(block) + result[ + :, + :, + time : time + self.block_step[0], + height : height + self.block_step[1], + width : width + self.block_step[2], + ] = result_block + return result + + def set_weights(self, weights, bias): + self.conv.weight = nn.Parameter(weights) + self.conv.bias = nn.Parameter(bias) \ No newline at end of file diff --git a/opensora/models/ae/videobase/modules/lpips_video.py b/opensora/models/ae/videobase/modules/lpips_video.py new file mode 100644 index 000000000..8b7062cdd --- /dev/null +++ b/opensora/models/ae/videobase/modules/lpips_video.py @@ -0,0 +1,120 @@ +"""Stripped version of https://github.com/richzhang/PerceptualSimilarity/tree/master/models""" + +import torch +import torch.nn as nn +from torchvision import models +from collections import namedtuple +from .....utils.taming_download import get_ckpt_path + +class LPIPS(nn.Module): + # Learned perceptual metric + def __init__(self, use_dropout=True): + super().__init__() + self.scaling_layer = ScalingLayer() + self.chns = [64, 128, 256, 512, 512] # vg16 features + self.net = vgg16(pretrained=True, requires_grad=False) + 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.load_from_pretrained() + for param in self.parameters(): + param.requires_grad = False + + def load_from_pretrained(self, name="vgg_lpips"): + ckpt = get_ckpt_path(name, "taming/modules/autoencoder/lpips") + self.load_state_dict(torch.load(ckpt, map_location=torch.device("cpu")), strict=False) + print("loaded pretrained LPIPS loss from {}".format(ckpt)) + + @classmethod + def from_pretrained(cls, name="vgg_lpips"): + if name != "vgg_lpips": + raise NotImplementedError + model = cls() + ckpt = get_ckpt_path(name) + model.load_state_dict(torch.load(ckpt, map_location=torch.device("cpu")), strict=False) + return model + + def forward(self, input, target): + in0_input, in1_input = (self.scaling_layer(input), self.scaling_layer(target)) + outs0, outs1 = self.net(in0_input), self.net(in1_input) + feats0, feats1, diffs = {}, {}, {} + lins = [self.lin0, self.lin1, self.lin2, self.lin3, self.lin4] + for kk in range(len(self.chns)): + feats0[kk], feats1[kk] = normalize_tensor(outs0[kk]), normalize_tensor(outs1[kk]) + diffs[kk] = (feats0[kk] - feats1[kk]) ** 2 + + res = [spatial_average(lins[kk].model(diffs[kk]), keepdim=True) for kk in range(len(self.chns))] + val = res[0] + for l in range(1, len(self.chns)): + val += res[l] + return val + + +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) + + +class vgg16(torch.nn.Module): + def __init__(self, requires_grad=False, pretrained=True): + super(vgg16, self).__init__() + vgg_pretrained_features = models.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 + + +def normalize_tensor(x,eps=1e-10): + norm_factor = torch.sqrt(torch.sum(x**2,dim=1,keepdim=True)) + return x/(norm_factor+eps) + + +def spatial_average(x, keepdim=True): + return x.mean([2,3],keepdim=keepdim) diff --git a/opensora/models/ae/videobase/modules/normalize.py b/opensora/models/ae/videobase/modules/normalize.py new file mode 100644 index 000000000..f9019c0f6 --- /dev/null +++ b/opensora/models/ae/videobase/modules/normalize.py @@ -0,0 +1,6 @@ +import torch + +def Normalize(in_channels, num_groups=32): + return torch.nn.GroupNorm( + num_groups=num_groups, num_channels=in_channels, eps=1e-6, affine=True + ) diff --git a/opensora/models/ae/videobase/modules/ops.py b/opensora/models/ae/videobase/modules/ops.py new file mode 100644 index 000000000..51d18a9b6 --- /dev/null +++ b/opensora/models/ae/videobase/modules/ops.py @@ -0,0 +1,7 @@ +import torch + +def nonlinearity(x): + return x * torch.sigmoid(x) + +def cast_tuple(t, length=1): + return t if isinstance(t, tuple) else ((t,) * length) diff --git a/opensora/models/ae/videobase/modules/perceptual_loss.py b/opensora/models/ae/videobase/modules/perceptual_loss.py new file mode 100644 index 000000000..a1ad760c7 --- /dev/null +++ b/opensora/models/ae/videobase/modules/perceptual_loss.py @@ -0,0 +1,91 @@ +import torch +from torch import nn +import torch.nn.functional as F +from .lpips_video import LPIPS + +from einops import rearrange + +def hinge_d_loss_with_exemplar_weights(logits_real, logits_fake, weights): + assert weights.shape[0] == logits_real.shape[0] == logits_fake.shape[0] + loss_real = torch.mean(F.relu(1.0 - logits_real), dim=[1, 2, 3]) + loss_fake = torch.mean(F.relu(1.0 + logits_fake), dim=[1, 2, 3]) + loss_real = (weights * loss_real).sum() / weights.sum() + loss_fake = (weights * loss_fake).sum() / weights.sum() + d_loss = 0.5 * (loss_real + loss_fake) + return d_loss + + +def adopt_weight(weight, global_step, threshold=0, value=0.0): + if global_step < threshold: + weight = value + return weight + + +def measure_perplexity(predicted_indices, n_embed): + # src: https://github.com/karpathy/deep-vector-quantization/blob/main/model.py + # eval cluster perplexity. when perplexity == num_embeddings then all clusters are used exactly equally + encodings = F.one_hot(predicted_indices, n_embed).float().reshape(-1, n_embed) + avg_probs = encodings.mean(0) + perplexity = (-(avg_probs * torch.log(avg_probs + 1e-10)).sum()).exp() + cluster_use = torch.sum(avg_probs > 0) + return perplexity, cluster_use + + +def l1(x, y): + return torch.abs(x - y) + + +def l2(x, y): + return torch.pow((x - y), 2) + +class LPIPSWithDiscriminator(nn.Module): + def __init__( + self, + logvar_init=0.0, + kl_weight=1.0, + pixelloss_weight=1.0, + perceptual_weight=1.0, + disc_loss="hinge", + ): + + super().__init__() + assert disc_loss in ["hinge", "vanilla"] + self.kl_weight = kl_weight + self.pixel_weight = pixelloss_weight + self.perceptual_loss = LPIPS().eval() + self.perceptual_weight = perceptual_weight + self.logvar = nn.Parameter(torch.ones(size=()) * logvar_init) + + def forward( + self, + inputs, + reconstructions, + posteriors, + split="train", + weights=None, + ): + inputs = rearrange(inputs, "b c t h w -> (b t) c h w").contiguous() + reconstructions = rearrange(reconstructions, "b c t h w -> (b t) c h w").contiguous() + rec_loss = torch.abs(inputs - reconstructions) + if self.perceptual_weight > 0: + p_loss = self.perceptual_loss(inputs, reconstructions) + rec_loss = rec_loss + self.perceptual_weight * p_loss + nll_loss = rec_loss / torch.exp(self.logvar) + self.logvar + weighted_nll_loss = nll_loss + if weights is not None: + weighted_nll_loss = weights * nll_loss + weighted_nll_loss = torch.sum(weighted_nll_loss) / weighted_nll_loss.shape[0] + nll_loss = torch.sum(nll_loss) / nll_loss.shape[0] + kl_loss = posteriors.kl() + kl_loss = torch.sum(kl_loss) / kl_loss.shape[0] + loss = weighted_nll_loss + self.kl_weight * kl_loss + log = { + "{}/total_loss".format(split): loss.clone().detach().mean(), + "{}/logvar".format(split): self.logvar.detach(), + "{}/kl_loss".format(split): kl_loss.detach().mean(), + "{}/nll_loss".format(split): nll_loss.detach().mean(), + "{}/rec_loss".format(split): rec_loss.detach().mean(), + } + if self.perceptual_weight > 0: + log.update({"{}/p_loss".format(split): p_loss.detach().mean()}) + return loss, log diff --git a/opensora/models/ae/videobase/modules/quant.py b/opensora/models/ae/videobase/modules/quant.py new file mode 100644 index 000000000..eebf451e6 --- /dev/null +++ b/opensora/models/ae/videobase/modules/quant.py @@ -0,0 +1,100 @@ +import torch +import torch.nn as nn +import torch.distributed as dist +import numpy as np +import torch.nn.functional as F +from ..utils.op_utils 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/ae/videobase/modules/resnet_block.py b/opensora/models/ae/videobase/modules/resnet_block.py new file mode 100644 index 000000000..9be406276 --- /dev/null +++ b/opensora/models/ae/videobase/modules/resnet_block.py @@ -0,0 +1,134 @@ +import torch +import torch.nn as nn +from einops import rearrange, pack, unpack +from .normalize import Normalize +from .ops import nonlinearity +from .conv import CausalConv3d + + +class ResnetBlock(nn.Module): + 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 + ) + + 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 ResnetBlock2D(nn.Module): + 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 + ) + + def forward(self, x): + x = rearrange(x, "b c f h w -> b f c h w") + x, ps = pack([x], "* c h w") + 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 + x = unpack(x, ps, "* c h w")[0] + x = rearrange(x, "b f c h w -> b c f h w") + return x + + +class ResnetBlock3D(nn.Module): + 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 \ No newline at end of file diff --git a/opensora/models/ae/videobase/modules/updownsample.py b/opensora/models/ae/videobase/modules/updownsample.py new file mode 100644 index 000000000..31905eb79 --- /dev/null +++ b/opensora/models/ae/videobase/modules/updownsample.py @@ -0,0 +1,171 @@ +from typing import Union, Tuple +import torch +import torch.nn as nn +import torch.nn.functional as F +from .ops import cast_tuple +from .conv import CausalConv3d +from einops import rearrange + +class Upsample(nn.Module): + 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) + + 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(nn.Module): + def __init__(self, in_channels, out_channels): + super().__init__() + self.with_conv = True + if self.with_conv: + # no asymmetric padding in torch conv, must do it ourselves + self.conv = torch.nn.Conv2d(in_channels, + out_channels, + kernel_size=3, + stride=2, + padding=0) + + def forward(self, x): + if self.with_conv: + 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(nn.Module): + def __init__( + self, + chan_in, + chan_out, + kernel_size: Union[int, Tuple[int]] = (3, 3), + stride: Union[int, Tuple[int]] = (2, 2), + ): + 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(nn.Module): + def __init__( + self, + chan_in, + chan_out, + kernel_size: Union[int, Tuple[int]] = (3, 3), + stride: Union[int, Tuple[int]] = (1, 1), + ): + super().__init__() + 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=1 + ) + + def forward(self, x): + 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(nn.Module): + def __init__( + self, + 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(nn.Module): + def __init__( + self, + ): + 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 TimeDownsample2x(nn.Module): +# def __init__( +# self, +# chan_in, +# chan_out, +# kernel_size: int = 3, +# padding: int = 1 +# ): +# super().__init__() +# self.chan_in = chan_in +# self.chan_out = chan_out +# self.kernel_size = kernel_size +# self.conv = CausalConv3d(chan_in, chan_out, kernel_size, stride=(2, 1, 1), padding=padding) + +# def forward(self, x): +# return self.conv(x) + +# class TimeUpsample2x(nn.Module): +# def __init__( +# self, +# chan_in, +# chan_out, +# kernel_size: int = 3, +# ): +# super().__init__() +# self.chan_in = chan_in +# self.chan_out = chan_out +# self.kernel_size = kernel_size +# self.conv = CausalConv3d(chan_in, chan_out, kernel_size, stride=1, padding=(0,1,1)) + +# def forward(self, x): +# if x.size(2) > 1: +# x,x_= x[:,:,:1],x[:,:,1:] +# x = F.interpolate(x, scale_factor=(1,1,1), mode='nearest') +# x_= F.interpolate(x_, scale_factor=(2,1,1), mode='nearest') +# x = torch.concat([x, x_], dim=2) +# x = self.conv(x) +# return x + \ No newline at end of file diff --git a/opensora/models/ae/videobase/trainer_videobase.py b/opensora/models/ae/videobase/trainer_videobase.py new file mode 100644 index 000000000..b6c4b0c7d --- /dev/null +++ b/opensora/models/ae/videobase/trainer_videobase.py @@ -0,0 +1,26 @@ +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/ae/videobase/vqvae/trainer_vqvae.py b/opensora/models/ae/videobase/vqvae/trainer_vqvae.py index 419c18244..a38c0c012 100644 --- a/opensora/models/ae/videobase/vqvae/trainer_vqvae.py +++ b/opensora/models/ae/videobase/vqvae/trainer_vqvae.py @@ -1,4 +1,4 @@ -from transformers import Trainer +from ..trainer_videobase import VideoBaseTrainer import torch.nn.functional as F from typing import Optional import os @@ -6,7 +6,7 @@ import torch from transformers.utils import WEIGHTS_NAME import json -class VQVAETrainer(Trainer): +class VQVAETrainer(VideoBaseTrainer): def compute_loss(self, model, inputs, return_outputs=False): model = model.module @@ -19,12 +19,3 @@ class VQVAETrainer(Trainer): loss = recon_loss + commitment_loss return loss - 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() - 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/ae/videobase/vqvae/videogpt/requirements.txt b/opensora/models/ae/videobase/vqvae/videogpt/requirements.txt deleted file mode 100644 index fde2a9f3d..000000000 --- a/opensora/models/ae/videobase/vqvae/videogpt/requirements.txt +++ /dev/null @@ -1,9 +0,0 @@ -tqdm -requests -h5py -scikit-video -av -gradio -moviepy -imageio -gdown diff --git a/opensora/models/ae/videobase/vqvae/videogpt/scripts/compute_fvd.py b/opensora/models/ae/videobase/vqvae/videogpt/scripts/compute_fvd.py deleted file mode 100644 index 1aca1551e..000000000 --- a/opensora/models/ae/videobase/vqvae/videogpt/scripts/compute_fvd.py +++ /dev/null @@ -1,131 +0,0 @@ -import os -import functools -import argparse -from videogpt.download import load_i3d_pretrained -from tqdm import tqdm -import numpy as np - -import torch -import torch.multiprocessing as mp -import torch.distributed as dist - -from videogpt.fvd.fvd import get_fvd_logits, frechet_distance -from videogpt import VideoData, VideoGPT, load_videogpt - - -MAX_BATCH = 32 - - -def main(): - assert torch.cuda.is_available() - ngpus = torch.cuda.device_count() - assert 256 % ngpus == 0, f"Must have 256 % n_gpus == 0" - - mp.spawn(main_worker, nprocs=ngpus, args=(ngpus, args), join=True) - - -def main_worker(rank, size, args_in): - global args - args = args_in - is_root = rank == 0 - dist.init_process_group(backend='nccl', init_method=f'tcp://localhost:{args.port}', - world_size=size, rank=rank) - device = torch.device(f"cuda:{rank}") - torch.cuda.set_device(device) - torch.set_grad_enabled(False) - - n_trials = args.n_trials - - #################### Load VideoGPT ######################################## - if not os.path.exists(args.ckpt): - gpt = load_videogpt(args.ckpt, device=device) - else: - gpt = VideoGPT.load_from_checkpoint(args.ckpt).to(device) - gpt.eval() - args = gpt.hparams['args'] - - args.batch_size = 256 // dist.get_world_size() - loader = VideoData(args).test_dataloader() - - #################### Load I3D ######################################## - i3d = load_i3d_pretrained(device) - - #################### Compute FVD ############################### - fvds = [] - fvds_star = [] - if is_root: - pbar = tqdm(total=n_trials) - for _ in range(n_trials): - fvd, fvd_star = eval_fvd(i3d, gpt, loader, device) - fvds.append(fvd) - fvds_star.append(fvd_star) - - if is_root: - pbar.update(1) - fvd_mean = np.mean(fvds) - fvd_std = np.std(fvds) - - fvd_star_mean = np.mean(fvds_star) - fvd_star_std = np.std(fvds_star) - - pbar.set_description(f"FVD {fvd_mean:.2f} +/- {fvd_std:.2f}, FVD* {fvd_star_mean:.2f} +/0 {fvd_star_std:.2f}") - if is_root: - pbar.close() - print(f"Final FVD {fvd_mean:.2f} +/- {fvd_std:.2f}, FVD* {fvd_star_mean:.2f} +/- {fvd_star_std:.2f}") - - -def all_gather(tensor): - rank, size = dist.get_rank(), dist.get_world_size() - tensor_list = [torch.zeros_like(tensor) for _ in range(size)] - dist.all_gather(tensor_list, tensor) - return torch.cat(tensor_list) - - -def eval_fvd(i3d, videogpt, loader, device): - rank, size = dist.get_rank(), dist.get_world_size() - is_root = rank == 0 - - batch = next(iter(loader)) - batch = {k: v.to(device) for k, v in batch.items()} - - fake_embeddings = [] - for i in range(0, batch['video'].shape[0], MAX_BATCH): - fake = videogpt.sample(MAX_BATCH, {k: v[i:i+MAX_BATCH] for k, v in batch.items()}) - fake = fake.permute(0, 2, 3, 4, 1).cpu().numpy() # BCTHW -> BTHWC - fake = (fake * 255).astype('uint8') - fake_embeddings.append(get_fvd_logits(fake, i3d=i3d, device=device)) - fake_embeddings = torch.cat(fake_embeddings) - - real = batch['video'].to(device) - real_recon_embeddings = [] - for i in range(0, batch['video'].shape[0], MAX_BATCH): - real_recon = (videogpt.get_reconstruction(batch['video'][i:i+MAX_BATCH]) + 0.5).clamp(0, 1) - real_recon = real_recon.permute(0, 2, 3, 4, 1).cpu().numpy() - real_recon = (real_recon * 255).astype('uint8') - real_recon_embeddings.append(get_fvd_logits(real_recon, i3d=i3d, device=device)) - real_recon_embeddings = torch.cat(real_recon_embeddings) - - real = real + 0.5 - real = real.permute(0, 2, 3, 4, 1).cpu().numpy() # BCTHW -> BTHWC - real = (real * 255).astype('uint8') - real_embeddings = get_fvd_logits(real, i3d=i3d, device=device) - - fake_embeddings = all_gather(fake_embeddings) - real_recon_embeddings = all_gather(real_recon_embeddings) - real_embeddings = all_gather(real_embeddings) - - assert fake_embeddings.shape[0] == real_recon_embeddings.shape[0] == real_embeddings.shape[0] == 256 - - fvd = frechet_distance(fake_embeddings.clone(), real_embeddings) - fvd_star = frechet_distance(fake_embeddings.clone(), real_recon_embeddings) - return fvd.item(), fvd_star.item() - - -if __name__ == '__main__': - parser = argparse.ArgumentParser() - parser.add_argument('--ckpt', type=str, default='bair_gpt') - parser.add_argument('--n_trials', type=int, default=1, help="Number of trials to compute mean/std") - parser.add_argument('--port', type=int, default=23452) - args = parser.parse_args() - - main() diff --git a/opensora/models/ae/videobase/vqvae/videogpt/scripts/preprocess/__init__.py b/opensora/models/ae/videobase/vqvae/videogpt/scripts/preprocess/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/opensora/models/ae/videobase/vqvae/videogpt/scripts/preprocess/bair/__init__.py b/opensora/models/ae/videobase/vqvae/videogpt/scripts/preprocess/bair/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/opensora/models/ae/videobase/vqvae/videogpt/scripts/preprocess/bair/bair_extract_images.py b/opensora/models/ae/videobase/vqvae/videogpt/scripts/preprocess/bair/bair_extract_images.py deleted file mode 100644 index 3f3d0d958..000000000 --- a/opensora/models/ae/videobase/vqvae/videogpt/scripts/preprocess/bair/bair_extract_images.py +++ /dev/null @@ -1,66 +0,0 @@ -import os -import io - -import numpy as np -from PIL import Image -import tensorflow as tf - -from tensorflow.python.platform import flags -from tensorflow.python.platform import gfile - -import imageio - -import argparse -parser = argparse.ArgumentParser() -parser.add_argument('--data_dir', default='', help='base directory to save processed data') -opt = parser.parse_args() - -def get_seq(dname): - data_dir = '%s/softmotion30_44k/%s' % (opt.data_dir, dname) - - filenames = gfile.Glob(os.path.join(data_dir, '*')) - if not filenames: - raise RuntimeError('No data files found.') - - for f in filenames: - k=0 - for serialized_example in tf.python_io.tf_record_iterator(f): - example = tf.train.Example() - example.ParseFromString(serialized_example) - image_seq, action_seq = [], [] - for i in range(30): - image_name = str(i) + '/image_aux1/encoded' - byte_str = example.features.feature[image_name].bytes_list.value[0] - img = Image.frombytes('RGB', (64, 64), byte_str) - arr = np.array(img.getdata()).reshape(img.size[1], img.size[0], 3) - image_seq.append(arr.reshape(1, 64, 64, 3)) - - action_name = str(i) + '/action' - action = example.features.feature[action_name].float_list.value - action = np.array(action).astype('float32') - action_seq.append(action) - image_seq = np.concatenate(image_seq, axis=0) - action_seq = np.stack(action_seq, axis=0) - k=k+1 - yield f, k, image_seq, action_seq - -def convert_data(dname): - seq_generator = get_seq(dname) - n = 0 - while True: - n+=1 - try: - f, k, seq, actions = next(seq_generator) - seq = seq.astype('uint8') - except StopIteration: - break - f = f.split('/')[-1] - os.makedirs('%s/processed_data/%s/%s/%d/' % (opt.data_dir, dname, f[:-10], k), exist_ok=True) - for i in range(len(seq)): - imageio.imwrite('%s/processed_data/%s/%s/%d/%d.png' % (opt.data_dir, dname, f[:-10], k, i), seq[i]) - np.save('%s/processed_data/%s/%s/%d/actions.npy' % (opt.data_dir, dname, f[:-10], k), actions) - - print('%s data: %s (%d) (%d)' % (dname, f, k, n)) - -convert_data('test') -convert_data('train') diff --git a/opensora/models/ae/videobase/vqvae/videogpt/scripts/preprocess/bair/bair_image_to_hdf5.py b/opensora/models/ae/videobase/vqvae/videogpt/scripts/preprocess/bair/bair_image_to_hdf5.py deleted file mode 100644 index b1c97ac71..000000000 --- a/opensora/models/ae/videobase/vqvae/videogpt/scripts/preprocess/bair/bair_image_to_hdf5.py +++ /dev/null @@ -1,50 +0,0 @@ -import glob -import argparse -import h5py -import numpy as np -from PIL import Image -import os -import os.path as osp -from tqdm import tqdm -import sys - -def convert_data(f, split): - root_dir = args.data_dir - path = osp.join(root_dir, 'processed_data', split) - traj_paths = glob.glob(osp.join(path, '*', '*')) - trajs, actions = [], [] - for traj_path in tqdm(traj_paths): - image_paths = glob.glob(osp.join(traj_path, '*.png')) - image_paths.sort(key=lambda x: int(osp.splitext(osp.basename(x))[0])) - traj = [] - for img_path in image_paths: - img = Image.open(img_path) - arr = np.array(img) # HWC - traj.append(arr) - traj = np.stack(traj, axis=0) # THWC - trajs.append(traj) - - actions.append(np.load(osp.join(traj_path, 'actions.npy'))) - - idxs = np.arange(len(trajs)) * 30 - trajs = np.concatenate(trajs, axis=0) # (NT)HWC - actions = np.concatenate(actions, axis=0) # (NT)(act_dim) - - f.create_dataset(f'{split}_data', data=trajs) - f.create_dataset(f'{split}_actions', data=actions) - f.create_dataset(f'{split}_idx', data=idxs) - - print(split) - print(f'\timages: {f[f"{split}_data"].shape}, {f[f"{split}_data"].dtype}') - print(f'\timages: {f[f"{split}_idx"].shape}, {f[f"{split}_idx"].dtype}') - -parser = argparse.ArgumentParser() -parser.add_argument('--data_dir', type=str, required=True) -parser.add_argument('--output_dir', type=str, required=True) -args = parser.parse_args() - -os.makedirs(args.output_dir, exist_ok=True) -f = h5py.File(osp.join(args.output_dir, 'bair.hdf5'), 'a') -convert_data(f, 'train') -convert_data(f, 'test') -f.close() diff --git a/opensora/models/ae/videobase/vqvae/videogpt/scripts/preprocess/bair/create_bair_dataset.sh b/opensora/models/ae/videobase/vqvae/videogpt/scripts/preprocess/bair/create_bair_dataset.sh deleted file mode 100644 index ab1c7e52d..000000000 --- a/opensora/models/ae/videobase/vqvae/videogpt/scripts/preprocess/bair/create_bair_dataset.sh +++ /dev/null @@ -1,11 +0,0 @@ -#! /bin/sh - -mkdir -p ${1} -mkdir -p ~/.cache/bair -wget -P ~/.cache/bair/ http://rail.eecs.berkeley.edu/datasets/bair_robot_pushing_dataset_v0.tar -tar -xf ~/.cache/bair/bair_robot_pushing_dataset_v0.tar -C ~/.cache/bair/ - -python scripts/preprocess/bair/bair_extract_images.py --data_dir ~/.cache/bair -python scripts/preprocess/bair/bair_image_to_hdf5.py --data_dir ~/.cache/bair --output_dir ${1} - -rm -r ~/.cache/bair diff --git a/opensora/models/ae/videobase/vqvae/videogpt/scripts/preprocess/ucf101/__init__.py b/opensora/models/ae/videobase/vqvae/videogpt/scripts/preprocess/ucf101/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/opensora/models/ae/videobase/vqvae/videogpt/scripts/preprocess/ucf101/create_ucf_dataset.sh b/opensora/models/ae/videobase/vqvae/videogpt/scripts/preprocess/ucf101/create_ucf_dataset.sh deleted file mode 100644 index 5754671a5..000000000 --- a/opensora/models/ae/videobase/vqvae/videogpt/scripts/preprocess/ucf101/create_ucf_dataset.sh +++ /dev/null @@ -1,19 +0,0 @@ -#! /bin/sh - -mkdir -p ${1} - -# Download UCF-101 video files -wget --no-check-certificate -P ${1} https://www.crcv.ucf.edu/data/UCF101/UCF101.rar -unrar x ${1}/UCF101.rar ${1} -rm ${1}/UCF101.rar - -# Download UCF-101 train/test splits -wget --no-check-certificate -P ${1} https://www.crcv.ucf.edu/data/UCF101/UCF101TrainTestSplits-RecognitionTask.zip -unzip ${1}/UCF101TrainTestSplits-RecognitionTask.zip -d ${1} -rm ${1}/UCF101TrainTestSplits-RecognitionTask.zip - -# Move video files into train / test directories based on train/test split -python scripts/preprocess/ucf101/ucf_split_train_test.py ${1} 1 - -# Delete leftover files -rm -r ${1}/UCF-101 diff --git a/opensora/models/ae/videobase/vqvae/videogpt/scripts/preprocess/ucf101/ucf_split_train_test.py b/opensora/models/ae/videobase/vqvae/videogpt/scripts/preprocess/ucf101/ucf_split_train_test.py deleted file mode 100644 index 797a43127..000000000 --- a/opensora/models/ae/videobase/vqvae/videogpt/scripts/preprocess/ucf101/ucf_split_train_test.py +++ /dev/null @@ -1,24 +0,0 @@ -import sys -import os -import os.path as osp -import shutil - -root = sys.argv[1] -fold = int(sys.argv[2]) -assert fold in [1, 2, 3] - -def move_files(files, split): - split_dir = osp.join(root, split) - os.makedirs(split_dir, exist_ok=True) - for filename in files: - folder = osp.join(split_dir, osp.dirname(filename)) - os.makedirs(folder, exist_ok=True) - shutil.move(osp.join(root, 'UCF-101', filename), osp.join(split_dir, filename)) - -with open(osp.join(root, 'ucfTrainTestlist', f'trainlist0{fold}.txt'), 'r') as f: - train_files = [p.strip().split(' ')[0] for p in f.readlines()] -move_files(train_files, 'train') - -with open(osp.join(root, 'ucfTrainTestlist', f'testlist0{fold}.txt'), 'r') as f: - test_files = [p.strip() for p in f.readlines()] -move_files(test_files, 'test') diff --git a/opensora/models/ae/videobase/vqvae/videogpt/scripts/sample_videogpt.py b/opensora/models/ae/videobase/vqvae/videogpt/scripts/sample_videogpt.py deleted file mode 100644 index 2ead14a3f..000000000 --- a/opensora/models/ae/videobase/vqvae/videogpt/scripts/sample_videogpt.py +++ /dev/null @@ -1,30 +0,0 @@ -import os -import argparse -import torch - -from videogpt import VideoData, VideoGPT, load_videogpt -from videogpt.utils import save_video_grid - - -parser = argparse.ArgumentParser() -parser.add_argument('--ckpt', type=str, default='ucf101_uncond_gpt') -parser.add_argument('--n', type=int, default=8) -args = parser.parse_args() -n = args.n - -if not os.path.exists(args.ckpt): - gpt = load_videogpt(args.ckpt) -else: - gpt = VideoGPT.load_from_checkpoint(args.ckpt) -gpt = gpt.cuda() -gpt.eval() -args = gpt.hparams['args'] - -args.batch_size = n -data = VideoData(args) -loader = data.test_dataloader() -batch = next(iter(loader)) -batch = {k: v.cuda() for k, v in batch.items()} - -samples = gpt.sample(n, batch) -save_video_grid(samples, 'samples.mp4') diff --git a/opensora/models/ae/videobase/vqvae/videogpt/scripts/train_videogpt.py b/opensora/models/ae/videobase/vqvae/videogpt/scripts/train_videogpt.py deleted file mode 100644 index 38784713f..000000000 --- a/opensora/models/ae/videobase/vqvae/videogpt/scripts/train_videogpt.py +++ /dev/null @@ -1,44 +0,0 @@ -import argparse -import pytorch_lightning as pl -from pytorch_lightning.callbacks import ModelCheckpoint -from videogpt import VideoGPT, VideoData - - -def main(): - pl.seed_everything(1234) - - parser = argparse.ArgumentParser() - parser = pl.Trainer.add_argparse_args(parser) - parser = VideoGPT.add_model_specific_args(parser) - parser.add_argument('--data_path', type=str, required=True) - parser.add_argument('--resolution', type=int, default=128) - parser.add_argument('--sequence_length', type=int, default=16) - parser.add_argument('--batch_size', type=int, default=8) - parser.add_argument('--num_workers', type=int, default=8) - args = parser.parse_args() - - data = VideoData(args) - # pre-make relevant cached files if necessary - data.train_dataloader() - data.test_dataloader() - - args.class_cond_dim = data.n_classes if args.class_cond else None - model = VideoGPT(args) - - callbacks = [] - callbacks.append(ModelCheckpoint(monitor='val/loss', mode='min', save_top_k=-1)) - - kwargs = dict() - if args.gpus > 1: - # find_unused_parameters = False to support gradient checkpointing - kwargs = dict(distributed_backend='ddp', gpus=args.gpus, - plugins=[pl.plugins.DDPPlugin(find_unused_parameters=False)]) - trainer = pl.Trainer.from_argparse_args(args, callbacks=callbacks, - max_steps=args.max_steps, **kwargs) - - trainer.fit(model, data) - - -if __name__ == '__main__': - main() - diff --git a/opensora/models/ae/videobase/vqvae/videogpt/scripts/train_vqvae.py b/opensora/models/ae/videobase/vqvae/videogpt/scripts/train_vqvae.py deleted file mode 100644 index 650fd3691..000000000 --- a/opensora/models/ae/videobase/vqvae/videogpt/scripts/train_vqvae.py +++ /dev/null @@ -1,40 +0,0 @@ -import argparse -import pytorch_lightning as pl -from pytorch_lightning.callbacks import ModelCheckpoint -from videogpt import VQVAE, VideoData -from transformers import TrainingArguments, HfArgumentParser - -def main(): - pl.seed_everything(1234) - - parser = argparse.ArgumentParser() - parser = pl.Trainer.add_argparse_args(parser) - parser = VQVAE.add_model_specific_args(parser) - parser.add_argument('--data_path', type=str, default='/home/wilson/data/datasets/bair.hdf5') - parser.add_argument('--sequence_length', type=int, default=16) - parser.add_argument('--resolution', type=int, default=64) - parser.add_argument('--batch_size', type=int, default=32) - parser.add_argument('--num_workers', type=int, default=8) - args = parser.parse_args() - - data = VideoData(args) - # pre-make relevant cached files if necessary - data.train_dataloader() - data.test_dataloader() - model = VQVAE(args) - - callbacks = [] - callbacks.append(ModelCheckpoint(monitor='val/recon_loss', mode='min')) - - kwargs = dict() - if args.gpus > 1: - kwargs = dict(distributed_backend='ddp', gpus=args.gpus) - trainer = pl.Trainer.from_argparse_args(args, callbacks=callbacks, - max_steps=200000, **kwargs) - - trainer.fit(model, data) - - -if __name__ == '__main__': - main() - diff --git a/opensora/models/ae/videobase/vqvae/videogpt/setup.py b/opensora/models/ae/videobase/vqvae/videogpt/setup.py deleted file mode 100644 index 364b0913f..000000000 --- a/opensora/models/ae/videobase/vqvae/videogpt/setup.py +++ /dev/null @@ -1,16 +0,0 @@ -from setuptools import setup - -def parse_requirements(filename): - lines = (line.strip() for line in open(filename)) - return [line for line in lines if line and not line.startswith('#')] - - -setup(name='VideoGPT', version='1.0', - description='PyTorch package for VideoGPT', - url='http://github.com/wilson1yan/VideoGPT', - author='Wilson Yan', - author_email='wilson1.yan@berkeley.edu', - license='BSD', - packages=['videogpt'], - install_requires=parse_requirements('requirements.txt'), - zip_safe=True) diff --git a/opensora/models/ae/videobase/vqvae/videogpt/videogpt/__init__.py b/opensora/models/ae/videobase/vqvae/videogpt/videogpt/__init__.py deleted file mode 100644 index d41e91cd4..000000000 --- a/opensora/models/ae/videobase/vqvae/videogpt/videogpt/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ - -from .vqvae import VQVAE -from .gpt import VideoGPT -from .data import VideoData -from .download import load_vqvae, load_videogpt, load_i3d_pretrained, download - diff --git a/opensora/models/ae/videobase/vqvae/videogpt/videogpt/attention.py b/opensora/models/ae/videobase/vqvae/videogpt/videogpt/attention.py deleted file mode 100644 index d35322ce0..000000000 --- a/opensora/models/ae/videobase/vqvae/videogpt/videogpt/attention.py +++ /dev/null @@ -1,585 +0,0 @@ -import numpy as np - -import torch -import torch.nn as nn -import torch.nn.functional as F -from torch.utils.checkpoint import checkpoint - -from .utils import shift_dim, view_range, tensor_slice - - -class AttentionStack(nn.Module): - def __init__( - self, shape, embd_dim, n_head, n_layer, dropout, - attn_type, attn_dropout, class_cond_dim, frame_cond_shape, - ): - super().__init__() - self.shape = shape - self.embd_dim = embd_dim - self.use_frame_cond = frame_cond_shape is not None - - self.right_shift = RightShift(embd_dim) - self.pos_embd = AddBroadcastPosEmbed( - shape=shape, embd_dim=embd_dim - ) - - self.attn_nets = nn.ModuleList( - [ - AttentionBlock( - shape=shape, - embd_dim=embd_dim, - n_head=n_head, - n_layer=n_layer, - dropout=dropout, - attn_type=attn_type, - attn_dropout=attn_dropout, - class_cond_dim=class_cond_dim, - frame_cond_shape=frame_cond_shape - ) - for i in range(n_layer) - ] - ) - - def forward(self, x, cond, decode_step, decode_idx): - """ - Args - ------ - x: (b, d1, d2, ..., dn, embd_dim) - cond: a dictionary of conditioning tensors - - (below is used only when sampling for fast decoding) - decode: the enumerated rasterscan order of the current idx being sampled - decode_step: a tuple representing the current idx being sampled - """ - x = self.right_shift(x, decode_step) - x = self.pos_embd(x, decode_step, decode_idx) - for net in self.attn_nets: - x = net(x, cond, decode_step, decode_idx) - - return x - - -class AttentionBlock(nn.Module): - def __init__(self, shape, embd_dim, n_head, n_layer, dropout, - attn_type, attn_dropout, class_cond_dim, frame_cond_shape): - super().__init__() - self.use_frame_cond = frame_cond_shape is not None - - self.pre_attn_norm = LayerNorm(embd_dim, class_cond_dim) - self.post_attn_dp = nn.Dropout(dropout) - self.attn = MultiHeadAttention(shape, embd_dim, embd_dim, n_head, - n_layer, causal=True, attn_type=attn_type, - attn_kwargs=dict(attn_dropout=attn_dropout)) - - if frame_cond_shape is not None: - enc_len = np.prod(frame_cond_shape[:-1]) - self.pre_enc_norm = LayerNorm(embd_dim, class_cond_dim) - self.post_enc_dp = nn.Dropout(dropout) - self.enc_attn = MultiHeadAttention(shape, embd_dim, frame_cond_shape[-1], - n_head, n_layer, attn_type='full', - attn_kwargs=dict(attn_dropout=0.), causal=False) - - self.pre_fc_norm = LayerNorm(embd_dim, class_cond_dim) - self.post_fc_dp = nn.Dropout(dropout) - self.fc_block = nn.Sequential( - nn.Linear(in_features=embd_dim, out_features=embd_dim * 4), - GeLU2(), - nn.Linear(in_features=embd_dim * 4, out_features=embd_dim), - ) - - def forward(self, x, cond, decode_step, decode_idx): - h = self.pre_attn_norm(x, cond) - if self.training: - h = checkpoint(self.attn, h, h, h, decode_step, decode_idx) - else: - h = self.attn(h, h, h, decode_step, decode_idx) - h = self.post_attn_dp(h) - x = x + h - - if self.use_frame_cond: - h = self.pre_enc_norm(x, cond) - if self.training: - h = checkpoint(self.enc_attn, h, cond['frame_cond'], cond['frame_cond'], - decode_step, decode_idx) - else: - h = self.enc_attn(h, cond['frame_cond'], cond['frame_cond'], - decode_step, decode_idx) - h = self.post_enc_dp(h) - x = x + h - - h = self.pre_fc_norm(x, cond) - if self.training: - h = checkpoint(self.fc_block, h) - else: - h = self.fc_block(h) - h = self.post_fc_dp(h) - x = x + h - - return x - - -class MultiHeadAttention(nn.Module): - def __init__(self, shape, dim_q, dim_kv, n_head, n_layer, - causal, attn_type, attn_kwargs): - super().__init__() - self.causal = causal - self.shape = shape - - self.d_k = dim_q // n_head - self.d_v = dim_kv // n_head - self.n_head = n_head - - self.w_qs = nn.Linear(dim_q, n_head * self.d_k, bias=False) # q - self.w_qs.weight.data.normal_(std=1.0 / np.sqrt(dim_q)) - - self.w_ks = nn.Linear(dim_kv, n_head * self.d_k, bias=False) # k - self.w_ks.weight.data.normal_(std=1.0 / np.sqrt(dim_kv)) - - self.w_vs = nn.Linear(dim_kv, n_head * self.d_v, bias=False) # v - self.w_vs.weight.data.normal_(std=1.0 / np.sqrt(dim_kv)) - - self.fc = nn.Linear(n_head * self.d_v, dim_q, bias=True) # c - self.fc.weight.data.normal_(std=1.0 / np.sqrt(dim_q * n_layer)) - - if attn_type == 'full': - self.attn = FullAttention(shape, causal, **attn_kwargs) - elif attn_type == 'axial': - assert not causal, 'causal axial attention is not supported' - self.attn = AxialAttention(len(shape), **attn_kwargs) - elif attn_type == 'sparse': - self.attn = SparseAttention(shape, n_head, causal, **attn_kwargs) - - self.cache = None - - def forward(self, q, k, v, decode_step=None, decode_idx=None): - """ Compute multi-head attention - Args - q, k, v: a [b, d1, ..., dn, c] tensor or - a [b, 1, ..., 1, c] tensor if decode_step is not None - - Returns - The output after performing attention - """ - - # compute k, q, v - d_k, d_v, n_head = self.d_k, self.d_v, self.n_head - q = view_range(self.w_qs(q), -1, None, (n_head, d_k)) - k = view_range(self.w_ks(k), -1, None, (n_head, d_k)) - v = view_range(self.w_vs(v), -1, None, (n_head, d_v)) - - # b x n_head x seq_len x d - # (b, *d_shape, n_head, d) -> (b, n_head, *d_shape, d) - q = shift_dim(q, -2, 1) - k = shift_dim(k, -2, 1) - v = shift_dim(v, -2, 1) - - # fast decoding - if decode_step is not None: - if decode_step == 0: - if self.causal: - k_shape = (q.shape[0], n_head, *self.shape, self.d_k) - v_shape = (q.shape[0], n_head, *self.shape, self.d_v) - self.cache = dict(k=torch.zeros(k_shape, dtype=k.dtype, device=q.device), - v=torch.zeros(v_shape, dtype=v.dtype, device=q.device)) - else: - # cache only once in the non-causal case - self.cache = dict(k=k.clone(), v=v.clone()) - if self.causal: - idx = (slice(None, None), slice(None, None), *[slice(i, i+ 1) for i in decode_idx]) - self.cache['k'][idx] = k - self.cache['v'][idx] = v - k, v = self.cache['k'], self.cache['v'] - - a = self.attn(q, k, v, decode_step, decode_idx) - - # (b, *d_shape, n_head, d) -> (b, *d_shape, n_head * d) - a = shift_dim(a, 1, -2).flatten(start_dim=-2) - a = self.fc(a) # (b x seq_len x embd_dim) - - return a - -############## Attention ####################### -class FullAttention(nn.Module): - def __init__(self, shape, causal, attn_dropout): - super().__init__() - self.causal = causal - self.attn_dropout = attn_dropout - - seq_len = np.prod(shape) - if self.causal: - self.register_buffer('mask', torch.tril(torch.ones(seq_len, seq_len))) - - def forward(self, q, k, v, decode_step, decode_idx): - mask = self.mask if self.causal else None - if decode_step is not None and mask is not None: - mask = mask[[decode_step]] - - old_shape = q.shape[2:-1] - q = q.flatten(start_dim=2, end_dim=-2) - k = k.flatten(start_dim=2, end_dim=-2) - v = v.flatten(start_dim=2, end_dim=-2) - - out = scaled_dot_product_attention(q, k, v, mask=mask, - attn_dropout=self.attn_dropout, - training=self.training) - - return view_range(out, 2, 3, old_shape) - -class AxialAttention(nn.Module): - def __init__(self, n_dim, axial_dim): - super().__init__() - if axial_dim < 0: - axial_dim = 2 + n_dim + 1 + axial_dim - else: - axial_dim += 2 # account for batch, head, dim - self.axial_dim = axial_dim - - def forward(self, q, k, v, decode_step, decode_idx): - q = shift_dim(q, self.axial_dim, -2).flatten(end_dim=-3) - k = shift_dim(k, self.axial_dim, -2).flatten(end_dim=-3) - v = shift_dim(v, self.axial_dim, -2) - old_shape = list(v.shape) - v = v.flatten(end_dim=-3) - - out = scaled_dot_product_attention(q, k, v, training=self.training) - out = out.view(*old_shape) - out = shift_dim(out, -2, self.axial_dim) - return out - - -class SparseAttention(nn.Module): - ops = dict() - attn_mask = dict() - block_layout = dict() - - def __init__(self, shape, n_head, causal, num_local_blocks=4, block=32, - attn_dropout=0.): # does not use attn_dropout - super().__init__() - self.causal = causal - self.shape = shape - - self.sparsity_config = StridedSparsityConfig(shape=shape, n_head=n_head, - causal=causal, block=block, - num_local_blocks=num_local_blocks) - - if self.shape not in SparseAttention.block_layout: - SparseAttention.block_layout[self.shape] = self.sparsity_config.make_layout() - if causal and self.shape not in SparseAttention.attn_mask: - SparseAttention.attn_mask[self.shape] = self.sparsity_config.make_sparse_attn_mask() - - def get_ops(self): - try: - from deepspeed.ops.sparse_attention import MatMul, Softmax - except: - raise Exception('Error importing deepspeed. Please install using `DS_BUILD_SPARSE_ATTN=1 pip install deepspeed`') - if self.shape not in SparseAttention.ops: - sparsity_layout = self.sparsity_config.make_layout() - sparse_dot_sdd_nt = MatMul(sparsity_layout, - self.sparsity_config.block, - 'sdd', - trans_a=False, - trans_b=True) - - sparse_dot_dsd_nn = MatMul(sparsity_layout, - self.sparsity_config.block, - 'dsd', - trans_a=False, - trans_b=False) - - sparse_softmax = Softmax(sparsity_layout, self.sparsity_config.block) - - SparseAttention.ops[self.shape] = (sparse_dot_sdd_nt, - sparse_dot_dsd_nn, - sparse_softmax) - return SparseAttention.ops[self.shape] - - def forward(self, q, k, v, decode_step, decode_idx): - if self.training and self.shape not in SparseAttention.ops: - self.get_ops() - - SparseAttention.block_layout[self.shape] = SparseAttention.block_layout[self.shape].to(q) - if self.causal: - SparseAttention.attn_mask[self.shape] = SparseAttention.attn_mask[self.shape].to(q).type_as(q) - attn_mask = SparseAttention.attn_mask[self.shape] if self.causal else None - - old_shape = q.shape[2:-1] - q = q.flatten(start_dim=2, end_dim=-2) - k = k.flatten(start_dim=2, end_dim=-2) - v = v.flatten(start_dim=2, end_dim=-2) - - if decode_step is not None: - mask = self.sparsity_config.get_non_block_layout_row(SparseAttention.block_layout[self.shape], decode_step) - out = scaled_dot_product_attention(q, k, v, mask=mask, training=self.training) - else: - if q.shape != k.shape or k.shape != v.shape: - raise Exception('SparseAttention only support self-attention') - sparse_dot_sdd_nt, sparse_dot_dsd_nn, sparse_softmax = self.get_ops() - scaling = float(q.shape[-1]) ** -0.5 - - attn_output_weights = sparse_dot_sdd_nt(q, k) - if attn_mask is not None: - attn_output_weights = attn_output_weights.masked_fill(attn_mask == 0, - float('-inf')) - attn_output_weights = sparse_softmax( - attn_output_weights, - scale=scaling - ) - - out = sparse_dot_dsd_nn(attn_output_weights, v) - - return view_range(out, 2, 3, old_shape) - - -class StridedSparsityConfig(object): - """ - Strided Sparse configuration specified in https://arxiv.org/abs/1904.10509 that - generalizes to arbitrary dimensions - """ - def __init__(self, shape, n_head, causal, block, num_local_blocks): - self.n_head = n_head - self.shape = shape - self.causal = causal - self.block = block - self.num_local_blocks = num_local_blocks - - assert self.num_local_blocks >= 1, 'Must have at least 1 local block' - assert self.seq_len % self.block == 0, 'seq len must be divisible by block size' - - self._block_shape = self._compute_block_shape() - self._block_shape_cum = self._block_shape_cum_sizes() - - @property - def seq_len(self): - return np.prod(self.shape) - - @property - def num_blocks(self): - return self.seq_len // self.block - - def set_local_layout(self, layout): - num_blocks = self.num_blocks - for row in range(0, num_blocks): - end = min(row + self.num_local_blocks, num_blocks) - for col in range( - max(0, row - self.num_local_blocks), - (row + 1 if self.causal else end)): - layout[:, row, col] = 1 - return layout - - def set_global_layout(self, layout): - num_blocks = self.num_blocks - n_dim = len(self._block_shape) - for row in range(num_blocks): - assert self._to_flattened_idx(self._to_unflattened_idx(row)) == row - cur_idx = self._to_unflattened_idx(row) - # no strided attention over last dim - for d in range(n_dim - 1): - end = self._block_shape[d] - for i in range(0, (cur_idx[d] + 1 if self.causal else end)): - new_idx = list(cur_idx) - new_idx[d] = i - new_idx = tuple(new_idx) - - col = self._to_flattened_idx(new_idx) - layout[:, row, col] = 1 - - return layout - - def make_layout(self): - layout = torch.zeros((self.n_head, self.num_blocks, self.num_blocks), dtype=torch.int64) - layout = self.set_local_layout(layout) - layout = self.set_global_layout(layout) - return layout - - def make_sparse_attn_mask(self): - block_layout = self.make_layout() - assert block_layout.shape[1] == block_layout.shape[2] == self.num_blocks - - num_dense_blocks = block_layout.sum().item() - attn_mask = torch.ones(num_dense_blocks, self.block, self.block) - counter = 0 - for h in range(self.n_head): - for i in range(self.num_blocks): - for j in range(self.num_blocks): - elem = block_layout[h, i, j].item() - if elem == 1: - assert i >= j - if i == j: # need to mask within block on diagonals - attn_mask[counter] = torch.tril(attn_mask[counter]) - counter += 1 - assert counter == num_dense_blocks - - return attn_mask.unsqueeze(0) - - def get_non_block_layout_row(self, block_layout, row): - block_row = row // self.block - block_row = block_layout[:, [block_row]] # n_head x 1 x n_blocks - block_row = block_row.repeat_interleave(self.block, dim=-1) - block_row[:, :, row + 1:] = 0. - return block_row - - ############# Helper functions ########################## - - def _compute_block_shape(self): - n_dim = len(self.shape) - cum_prod = 1 - for i in range(n_dim - 1, -1, -1): - cum_prod *= self.shape[i] - if cum_prod > self.block: - break - assert cum_prod % self.block == 0 - new_shape = (*self.shape[:i], cum_prod // self.block) - - assert np.prod(new_shape) == np.prod(self.shape) // self.block - - return new_shape - - def _block_shape_cum_sizes(self): - bs = np.flip(np.array(self._block_shape)) - return tuple(np.flip(np.cumprod(bs)[:-1])) + (1,) - - def _to_flattened_idx(self, idx): - assert len(idx) == len(self._block_shape), f"{len(idx)} != {len(self._block_shape)}" - flat_idx = 0 - for i in range(len(self._block_shape)): - flat_idx += idx[i] * self._block_shape_cum[i] - return flat_idx - - def _to_unflattened_idx(self, flat_idx): - assert flat_idx < np.prod(self._block_shape) - idx = [] - for i in range(len(self._block_shape)): - idx.append(flat_idx // self._block_shape_cum[i]) - flat_idx %= self._block_shape_cum[i] - return tuple(idx) - - -################ Spatiotemporal broadcasted positional embeddings ############### -class AddBroadcastPosEmbed(nn.Module): - def __init__(self, shape, embd_dim, dim=-1): - super().__init__() - assert dim in [-1, 1] # only first or last dim supported - self.shape = shape - self.n_dim = n_dim = len(shape) - self.embd_dim = embd_dim - self.dim = dim - - assert embd_dim % n_dim == 0, f"{embd_dim} % {n_dim} != 0" - self.emb = nn.ParameterDict({ - f'd_{i}': nn.Parameter(torch.randn(shape[i], embd_dim // n_dim) * 0.01 - if dim == -1 else - torch.randn(embd_dim // n_dim, shape[i]) * 0.01) - for i in range(n_dim) - }) - - def forward(self, x, decode_step=None, decode_idx=None): - embs = [] - for i in range(self.n_dim): - e = self.emb[f'd_{i}'] - if self.dim == -1: - # (1, 1, ..., 1, self.shape[i], 1, ..., -1) - e = e.view(1, *((1,) * i), self.shape[i], *((1,) * (self.n_dim - i - 1)), -1) - e = e.expand(1, *self.shape, -1) - else: - e = e.view(1, -1, *((1,) * i), self.shape[i], *((1,) * (self.n_dim - i - 1))) - e = e.expand(1, -1, *self.shape) - embs.append(e) - - embs = torch.cat(embs, dim=self.dim) - if decode_step is not None: - embs = tensor_slice(embs, [0, *decode_idx, 0], - [x.shape[0], *(1,) * self.n_dim, x.shape[-1]]) - - return x + embs - -################# Helper Functions ################################### -def scaled_dot_product_attention(q, k, v, mask=None, attn_dropout=0., training=True): - # Performs scaled dot-product attention over the second to last dimension dn - - # (b, n_head, d1, ..., dn, d) - attn = torch.matmul(q, k.transpose(-1, -2)) - attn = attn / np.sqrt(q.shape[-1]) - if mask is not None: - attn = attn.masked_fill(mask == 0, float('-inf')) - attn_float = F.softmax(attn, dim=-1) - attn = attn_float.type_as(attn) # b x n_head x d1 x ... x dn x d - attn = F.dropout(attn, p=attn_dropout, training=training) - - a = torch.matmul(attn, v) # b x n_head x d1 x ... x dn x d - - return a - - -class RightShift(nn.Module): - def __init__(self, embd_dim): - super().__init__() - self.embd_dim = embd_dim - self.sos = nn.Parameter(torch.FloatTensor(embd_dim).normal_(std=0.02), requires_grad=True) - - def forward(self, x, decode_step): - if decode_step is not None and decode_step > 0: - return x - - x_shape = list(x.shape) - x = x.flatten(start_dim=1, end_dim=-2) # (b, seq_len, embd_dim) - sos = torch.ones(x_shape[0], 1, self.embd_dim, dtype=torch.float32).to(self.sos) * self.sos - sos = sos.type_as(x) - x = torch.cat([sos, x[:, :-1, :]], axis=1) - x = x.view(*x_shape) - - return x - - -class GeLU2(nn.Module): - def forward(self, x): - return (1.702 * x).sigmoid() * x - - -class LayerNorm(nn.Module): - def __init__(self, embd_dim, class_cond_dim): - super().__init__() - self.conditional = class_cond_dim is not None - - if self.conditional: - self.w = nn.Linear(class_cond_dim, embd_dim, bias=False) - nn.init.constant_(self.w.weight.data, 1. / np.sqrt(class_cond_dim)) - self.wb = nn.Linear(class_cond_dim, embd_dim, bias=False) - else: - self.g = nn.Parameter(torch.ones(embd_dim, dtype=torch.float32), requires_grad=True) - self.b = nn.Parameter(torch.zeros(embd_dim, dtype=torch.float32), requires_grad=True) - - def forward(self, x, cond): - if self.conditional: # (b, cond_dim) - g = 1 + self.w(cond['class_cond']).view(x.shape[0], *(1,)*(len(x.shape)-2), x.shape[-1]) # (b, ..., embd_dim) - b = self.wb(cond['class_cond']).view(x.shape[0], *(1,)*(len(x.shape)-2), x.shape[-1]) - else: - g = self.g # (embd_dim,) - b = self.b - - x_float = x.float() - - mu = x_float.mean(dim=-1, keepdims=True) - s = (x_float - mu).square().mean(dim=-1, keepdims=True) - x_float = (x_float - mu) * (1e-5 + s.rsqrt()) # (b, ..., embd_dim) - x_float = x_float * g + b - - x = x_float.type_as(x) - return x - -def view_range(x, i, j, shape): - shape = tuple(shape) - - n_dims = len(x.shape) - if i < 0: - i = n_dims + i - - if j is None: - j = n_dims - elif j < 0: - j = n_dims + j - - assert 0 <= i < j <= n_dims - - x_shape = x.shape - target_shape = x_shape[:i] + shape + x_shape[j:] - return x.view(target_shape) \ No newline at end of file diff --git a/opensora/models/ae/videobase/vqvae/videogpt/videogpt/data.py b/opensora/models/ae/videobase/vqvae/videogpt/videogpt/data.py deleted file mode 100644 index bd0c23b0a..000000000 --- a/opensora/models/ae/videobase/vqvae/videogpt/videogpt/data.py +++ /dev/null @@ -1,213 +0,0 @@ -import os -import os.path as osp -import math -import random -import pickle -import warnings - -import glob -import h5py -import numpy as np - -import torch -import torch.utils.data as data -import torch.nn.functional as F -import torch.distributed as dist -from torchvision.datasets.video_utils import VideoClips -import pytorch_lightning as pl - - -class VideoDataset(data.Dataset): - """ Generic dataset for videos files stored in folders - Returns BCTHW videos in the range [-0.5, 0.5] """ - exts = ['avi', 'mp4', 'webm'] - - def __init__(self, data_folder, sequence_length, train=True, resolution=64): - """ - Args: - data_folder: path to the folder with videos. The folder - should contain a 'train' and a 'test' directory, - each with corresponding videos stored - sequence_length: length of extracted video sequences - """ - super().__init__() - self.train = train - self.sequence_length = sequence_length - self.resolution = resolution - - folder = osp.join(data_folder, 'train' if train else 'test') - files = sum([glob.glob(osp.join(folder, '**', f'*.{ext}'), recursive=True) - for ext in self.exts], []) - - # hacky way to compute # of classes (count # of unique parent directories) - self.classes = list(set([get_parent_dir(f) for f in files])) - self.classes.sort() - self.class_to_label = {c: i for i, c in enumerate(self.classes)} - - warnings.filterwarnings('ignore') - cache_file = osp.join(folder, f"metadata_{sequence_length}.pkl") - if not osp.exists(cache_file): - clips = VideoClips(files, sequence_length, num_workers=32) - pickle.dump(clips.metadata, open(cache_file, 'wb')) - else: - metadata = pickle.load(open(cache_file, 'rb')) - clips = VideoClips(files, sequence_length, - _precomputed_metadata=metadata) - self._clips = clips - - @property - def n_classes(self): - return len(self.classes) - - def __len__(self): - return self._clips.num_clips() - - def __getitem__(self, idx): - resolution = self.resolution - video, _, _, idx = self._clips.get_clip(idx) - - class_name = get_parent_dir(self._clips.video_paths[idx]) - label = self.class_to_label[class_name] - return dict(video=preprocess(video, resolution), label=label) - - -def get_parent_dir(path): - return osp.basename(osp.dirname(path)) - - -def preprocess(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 - - -class HDF5Dataset(data.Dataset): - """ Generic dataset for data stored in h5py as uint8 numpy arrays. - Reads videos in {0, ..., 255} and returns in range [-0.5, 0.5] """ - def __init__(self, data_file, sequence_length, train=True, resolution=64): - """ - Args: - data_file: path to the pickled data file with the - following format: - { - 'train_data': [B, H, W, 3] np.uint8, - 'train_idx': [B], np.int64 (start indexes for each video) - 'test_data': [B', H, W, 3] np.uint8, - 'test_idx': [B'], np.int64 - } - sequence_length: length of extracted video sequences - """ - super().__init__() - self.train = train - self.sequence_length = sequence_length - self.resolution = resolution - - # read in data - self.data_file = data_file - self.data = h5py.File(data_file, 'r') - self.prefix = 'train' if train else 'test' - self._images = self.data[f'{self.prefix}_data'] - self._idx = self.data[f'{self.prefix}_idx'] - self.size = len(self._idx) - - @property - def n_classes(self): - raise Exception('class conditioning not support for HDF5Dataset') - - def __getstate__(self): - state = self.__dict__ - state['data'] = None - state['_images'] = None - state['_idx'] = None - return state - - def __setstate__(self, state): - self.__dict__ = state - self.data = h5py.File(self.data_file, 'r') - self._images = self.data[f'{self.prefix}_data'] - self._idx = self.data[f'{self.prefix}_idx'] - - def __len__(self): - return self.size - - def __getitem__(self, idx): - start = self._idx[idx] - end = self._idx[idx + 1] if idx < len(self._idx) - 1 else len(self._images) - assert end - start >= 0 - - start = start + np.random.randint(low=0, high=end - start - self.sequence_length) - assert start < start + self.sequence_length <= end - video = torch.tensor(self._images[start:start + self.sequence_length]) - return dict(video=preprocess(video, self.resolution)) - - -class VideoData(pl.LightningDataModule): - - def __init__(self, args): - super().__init__() - self.args = args - - @property - def n_classes(self): - dataset = self._dataset(True) - return dataset.n_classes - - - def _dataset(self, train): - Dataset = VideoDataset if osp.isdir(self.args.data_path) else HDF5Dataset - dataset = Dataset(self.args.data_path, self.args.sequence_length, - train=train, resolution=self.args.resolution) - return dataset - - - def _dataloader(self, train): - dataset = self._dataset(train) - if dist.is_initialized(): - sampler = data.distributed.DistributedSampler( - dataset, num_replicas=dist.get_world_size(), rank=dist.get_rank() - ) - else: - sampler = None - dataloader = data.DataLoader( - dataset, - batch_size=self.args.batch_size, - num_workers=self.args.num_workers, - pin_memory=True, - sampler=sampler, - shuffle=sampler is None - ) - return dataloader - - def train_dataloader(self): - return self._dataloader(True) - - def val_dataloader(self): - return self._dataloader(False) - - def test_dataloader(self): - return self.val_dataloader() diff --git a/opensora/models/ae/videobase/vqvae/videogpt/videogpt/download.py b/opensora/models/ae/videobase/vqvae/videogpt/videogpt/download.py deleted file mode 100644 index 9879d8e39..000000000 --- a/opensora/models/ae/videobase/vqvae/videogpt/videogpt/download.py +++ /dev/null @@ -1,60 +0,0 @@ -import requests -from tqdm import tqdm -import os -import gdown -import torch - -from .vqvae import VQVAE -from .gpt import VideoGPT - - -def download(id, fname, root=os.path.expanduser('~/.cache/videogpt')): - os.makedirs(root, exist_ok=True) - destination = os.path.join(root, fname) - - if os.path.exists(destination): - return destination - - gdown.download(id=id, output=destination, quiet=False) - return destination - - -_VQVAE = { - 'bair_stride4x2x2': '1iIAYJ2Qqrx5Q94s5eIXQYJgAydzvT_8L', # trained on 16 frames of 64 x 64 images - 'ucf101_stride4x4x4': '1uuB_8WzHP_bbBmfuaIV7PK_Itl3DyHY5', # trained on 16 frames of 128 x 128 images - 'kinetics_stride4x4x4': '1DOvOZnFAIQmux6hG7pN_HkyJZy3lXbCB', # trained on 16 frames of 128 x 128 images - 'kinetics_stride2x4x4': '1jvtjjtrtE4cy6pl7DK_zWFEPY3RZt2pB' # trained on 16 frames of 128 x 128 images -} - -def load_vqvae(model_name, device=torch.device('cpu'), root=os.path.expanduser('~/.cache/videogpt')): - assert model_name in _VQVAE, f"Invalid model_name: {model_name}" - filepath = download(_VQVAE[model_name], model_name, root=root) - vqvae = VQVAE.load_from_checkpoint(filepath).to(device) - vqvae.eval() - - return vqvae - - -_VIDEOGPT = { - 'bair_gpt': '1fNTtJAgO6grEtPNrufkpbee1CfGztW-1', # 1-frame conditional, 16 frames of 64 x 64 images - 'ucf101_uncond_gpt': '1QkF_Sb2XVRgSbFT_SxQ6aZUeDFoliPQq', # unconditional, 16 frames of 128 x 128 images -} - -def load_videogpt(model_name, device=torch.device('cpu')): - assert model_name in _VIDEOGPT, f"Invalid model_name: {model_name}" - filepath = download(_VIDEOGPT[model_name], model_name) - gpt = VideoGPT.load_from_checkpoint(filepath).to(device) - gpt.eval() - - return gpt - - -_I3D_PRETRAINED_ID = '1mQK8KD8G6UWRa5t87SRMm5PVXtlpneJT' - -def load_i3d_pretrained(device=torch.device('cpu')): - from .fvd.pytorch_i3d import InceptionI3d - i3d = InceptionI3d(400, in_channels=3).to(device) - filepath = download(_I3D_PRETRAINED_ID, 'i3d_pretrained_400.pt') - i3d.load_state_dict(torch.load(filepath, map_location=device)) - i3d.eval() - return i3d diff --git a/opensora/models/ae/videobase/vqvae/videogpt/videogpt/fvd/__init__.py b/opensora/models/ae/videobase/vqvae/videogpt/videogpt/fvd/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/opensora/models/ae/videobase/vqvae/videogpt/videogpt/fvd/convert_tf_pretrained.py b/opensora/models/ae/videobase/vqvae/videogpt/videogpt/fvd/convert_tf_pretrained.py deleted file mode 100644 index 40690a959..000000000 --- a/opensora/models/ae/videobase/vqvae/videogpt/videogpt/fvd/convert_tf_pretrained.py +++ /dev/null @@ -1,127 +0,0 @@ -import sys -from collections import OrderedDict -import tensorflow_hub as hub -import torch - -from src_pytorch.fvd.pytorch_i3d import InceptionI3d - - -def convert_name(name): - mapping = { - 'conv_3d': 'conv3d', - 'batch_norm': 'bn', - 'w:0': 'weight', - 'b:0': 'bias', - 'moving_mean:0': 'running_mean', - 'moving_variance:0': 'running_var', - 'beta:0': 'bias' - } - - segs = name.split('/') - new_segs = [] - i = 0 - while i < len(segs): - seg = segs[i] - if 'Mixed' in seg: - new_segs.append(seg) - elif 'Conv' in seg and 'Mixed' not in name: - new_segs.append(seg) - elif 'Branch' in seg: - branch_i = int(seg.split('_')[-1]) - i += 1 - seg = segs[i] - - # special case due to typo in original code - if 'Mixed_5b' in name and branch_i == 2: - if '1x1' in seg: - new_segs.append(f'b{branch_i}a') - elif '3x3' in seg: - new_segs.append(f'b{branch_i}b') - else: - raise Exception() - # Either Conv3d_{i}a_... or Conv3d_{i}b_... - elif 'a' in seg: - if branch_i == 0: - new_segs.append('b0') - else: - new_segs.append(f'b{branch_i}a') - elif 'b' in seg: - new_segs.append(f'b{branch_i}b') - else: - raise Exception - elif seg == 'Logits': - new_segs.append('logits') - i += 1 - elif seg in mapping: - new_segs.append(mapping[seg]) - else: - raise Exception(f"No match found for seg {seg} in name {name}") - - i += 1 - return '.'.join(new_segs) - -def convert_tensor(tensor): - tensor_dim = len(tensor.shape) - if tensor_dim == 5: # conv or bn - if all([t == 1 for t in tensor.shape[:-1]]): - tensor = tensor.squeeze() - else: - tensor = tensor.permute(4, 3, 0, 1, 2).contiguous() - elif tensor_dim == 1: # conv bias - pass - else: - raise Exception(f"Invalid shape {tensor.shape}") - return tensor - -n_class = int(sys.argv[1]) # 600 or 400 -assert n_class in [400, 600] - -# Converts model from https://github.com/google-research/google-research/tree/master/frechet_video_distance -# to pytorch version for loading -model_url = f"https://tfhub.dev/deepmind/i3d-kinetics-{n_class}/1" -i3d = hub.load(model_url) -name_prefix = 'RGB/inception_i3d/' - -print('Creating state_dict...') -all_names = [] -state_dict = OrderedDict() -for var in i3d.variables: - name = var.name[len(name_prefix):] - new_name = convert_name(name) - all_names.append(new_name) - - tensor = torch.FloatTensor(var.value().numpy()) - new_tensor = convert_tensor(tensor) - - state_dict[new_name] = new_tensor - - if 'bn.bias' in new_name: - new_name = new_name[:-4] + 'weight' # bn.weight - new_tensor = torch.ones_like(new_tensor).float() - state_dict[new_name] = new_tensor - -print(f'Complete state_dict with {len(state_dict)} entries') - -s = dict() -for i, n in enumerate(all_names): - s[n] = s.get(n, []) + [i] - -for k, v in s.items(): - if len(v) > 1: - print('dup', k) - for i in v: - print('\t', i3d.variables[i].name) - -print('Testing load_state_dict...') -print('Creating model...') - -i3d = InceptionI3d(n_class, in_channels=3) - -print('Loading state_dict...') -i3d.load_state_dict(state_dict) - -print(f'Saving state_dict as fvd/i3d_pretrained_{n_class}.pt') -torch.save(state_dict, f'fvd/i3d_pretrained_{n_class}.pt') - -print('Done') - diff --git a/opensora/models/ae/videobase/vqvae/videogpt/videogpt/fvd/fvd.py b/opensora/models/ae/videobase/vqvae/videogpt/videogpt/fvd/fvd.py deleted file mode 100644 index 7542f4c7a..000000000 --- a/opensora/models/ae/videobase/vqvae/videogpt/videogpt/fvd/fvd.py +++ /dev/null @@ -1,85 +0,0 @@ -import torch -from ..data import preprocess as preprocess_single - - -def preprocess(videos, target_resolution=224): - # videos in {0, ..., 255} as np.uint8 array - 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): - videos = preprocess(videos) - embeddings = get_logits(i3d, videos, device) - 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) - - sqrt_trace_component = trace_sqrt_product(sigma, sigma_w) - trace = torch.trace(sigma + sigma_w) - 2.0 * sqrt_trace_component - - mean = torch.sum((m - m_w) ** 2) - fd = trace + mean - return fd - - -def get_logits(i3d, videos, device): - assert videos.shape[0] % 16 == 0 - with torch.no_grad(): - logits = [] - for i in range(0, videos.shape[0], 16): - batch = videos[i:i + 16].to(device) - logits.append(i3d(batch)) - logits = torch.cat(logits, dim=0) - return logits diff --git a/opensora/models/ae/videobase/vqvae/videogpt/videogpt/fvd/pytorch_i3d.py b/opensora/models/ae/videobase/vqvae/videogpt/videogpt/fvd/pytorch_i3d.py deleted file mode 100644 index 77d4830f5..000000000 --- a/opensora/models/ae/videobase/vqvae/videogpt/videogpt/fvd/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) diff --git a/opensora/models/ae/videobase/vqvae/videogpt/videogpt/gpt.py b/opensora/models/ae/videobase/vqvae/videogpt/videogpt/gpt.py deleted file mode 100644 index d6cf35bb8..000000000 --- a/opensora/models/ae/videobase/vqvae/videogpt/videogpt/gpt.py +++ /dev/null @@ -1,185 +0,0 @@ -import os -import itertools -import numpy as np -from tqdm import tqdm -import argparse - -import torch -import torch.nn as nn -import torch.nn.functional as F -import torch.optim.lr_scheduler as lr_scheduler -import pytorch_lightning as pl - -from .resnet import resnet34 -from .attention import AttentionStack, LayerNorm, AddBroadcastPosEmbed -from .utils import shift_dim - - -class VideoGPT(pl.LightningModule): - def __init__(self, args): - super().__init__() - self.args = args - - # Load VQ-VAE and set all parameters to no grad - from .vqvae import VQVAE - from .download import load_vqvae - if not os.path.exists(args.vqvae): - self.vqvae = load_vqvae(args.vqvae) - else: - self.vqvae = VQVAE.load_from_checkpoint(args.vqvae) - for p in self.vqvae.parameters(): - p.requires_grad = False - self.vqvae.codebook._need_init = False - self.vqvae.eval() - - # ResNet34 for frame conditioning - self.use_frame_cond = args.n_cond_frames > 0 - if self.use_frame_cond: - frame_cond_shape = (args.n_cond_frames, - args.resolution // 4, - args.resolution // 4, - 240) - self.resnet = resnet34(1, (1, 4, 4), resnet_dim=240) - self.cond_pos_embd = AddBroadcastPosEmbed( - shape=frame_cond_shape[:-1], embd_dim=frame_cond_shape[-1] - ) - else: - frame_cond_shape = None - - # VideoGPT transformer - self.shape = self.vqvae.latent_shape - - self.fc_in = nn.Linear(self.vqvae.embedding_dim, args.hidden_dim, bias=False) - self.fc_in.weight.data.normal_(std=0.02) - - self.attn_stack = AttentionStack( - self.shape, args.hidden_dim, args.heads, args.layers, args.dropout, - args.attn_type, args.attn_dropout, args.class_cond_dim, frame_cond_shape - ) - - self.norm = LayerNorm(args.hidden_dim, args.class_cond_dim) - - self.fc_out = nn.Linear(args.hidden_dim, self.vqvae.n_codes, bias=False) - self.fc_out.weight.data.copy_(torch.zeros(self.vqvae.n_codes, args.hidden_dim)) - - # caches for faster decoding (if necessary) - self.frame_cond_cache = None - - self.save_hyperparameters() - - def get_reconstruction(self, videos): - return self.vqvae.decode(self.vqvae.encode(videos)) - - def sample(self, n, batch=None): - device = self.fc_in.weight.device - - cond = dict() - if self.use_frame_cond or self.args.class_cond: - assert batch is not None - video = batch['video'] - - if self.args.class_cond: - label = batch['label'] - cond['class_cond'] = F.one_hot(label, self.args.class_cond_dim).type_as(video) - if self.use_frame_cond: - cond['frame_cond'] = video[:, :, :self.args.n_cond_frames] - - samples = torch.zeros((n,) + self.shape).long().to(device) - idxs = list(itertools.product(*[range(s) for s in self.shape])) - - with torch.no_grad(): - prev_idx = None - for i, idx in enumerate(tqdm(idxs)): - batch_idx_slice = (slice(None, None), *[slice(i, i + 1) for i in idx]) - batch_idx = (slice(None, None), *idx) - embeddings = self.vqvae.codebook.dictionary_lookup(samples) - - if prev_idx is None: - # set arbitrary input values for the first token - # does not matter what value since it will be shifted anyways - embeddings_slice = embeddings[batch_idx_slice] - samples_slice = samples[batch_idx_slice] - else: - embeddings_slice = embeddings[prev_idx] - samples_slice = samples[prev_idx] - - logits = self(embeddings_slice, samples_slice, cond, - decode_step=i, decode_idx=idx)[1] - # squeeze all possible dim except batch dimension - logits = logits.squeeze().unsqueeze(0) if logits.shape[0] == 1 else logits.squeeze() - probs = F.softmax(logits, dim=-1) - samples[batch_idx] = torch.multinomial(probs, 1).squeeze(-1) - - prev_idx = batch_idx_slice - samples = self.vqvae.decode(samples) - samples = torch.clamp(samples, -0.5, 0.5) + 0.5 - - return samples # BCTHW in [0, 1] - - - def forward(self, x, targets, cond, decode_step=None, decode_idx=None): - if self.use_frame_cond: - if decode_step is None: - cond['frame_cond'] = self.cond_pos_embd(self.resnet(cond['frame_cond'])) - elif decode_step == 0: - self.frame_cond_cache = self.cond_pos_embd(self.resnet(cond['frame_cond'])) - cond['frame_cond'] = self.frame_cond_cache - else: - cond['frame_cond'] = self.frame_cond_cache - - h = self.fc_in(x) - h = self.attn_stack(h, cond, decode_step, decode_idx) - h = self.norm(h, cond) - logits = self.fc_out(h) - - loss = F.cross_entropy(shift_dim(logits, -1, 1), targets) - - return loss, logits - - def training_step(self, batch, batch_idx): - self.vqvae.eval() - x = batch['video'] - - cond = dict() - if self.args.class_cond: - label = batch['label'] - cond['class_cond'] = F.one_hot(label, self.args.class_cond_dim).type_as(x) - if self.use_frame_cond: - cond['frame_cond'] = x[:, :, :self.args.n_cond_frames] - - with torch.no_grad(): - targets, x = self.vqvae.encode(x, include_embeddings=True) - x = shift_dim(x, 1, -1) - - loss, _ = self(x, targets, cond) - return loss - - def validation_step(self, batch, batch_idx): - loss = self.training_step(batch, batch_idx) - self.log('val/loss', loss, prog_bar=True) - - def configure_optimizers(self): - optimizer = torch.optim.Adam(self.parameters(), lr=3e-4, betas=(0.9, 0.999)) - assert hasattr(self.args, 'max_steps') and self.args.max_steps is not None, f"Must set max_steps argument" - scheduler = lr_scheduler.CosineAnnealingLR(optimizer, self.args.max_steps) - return [optimizer], [dict(scheduler=scheduler, interval='step', frequency=1)] - - - @staticmethod - def add_model_specific_args(parent_parser): - parser = argparse.ArgumentParser(parents=[parent_parser], add_help=False) - parser.add_argument('--vqvae', type=str, default='kinetics_stride4x4x4', - help='path to vqvae ckpt, or model name to download pretrained') - parser.add_argument('--n_cond_frames', type=int, default=0) - parser.add_argument('--class_cond', action='store_true') - - # VideoGPT hyperparmeters - parser.add_argument('--hidden_dim', type=int, default=576) - parser.add_argument('--heads', type=int, default=4) - parser.add_argument('--layers', type=int, default=8) - parser.add_argument('--dropout', type=float, default=0.2) - parser.add_argument('--attn_type', type=str, default='full', - choices=['full', 'sparse']) - parser.add_argument('--attn_dropout', type=float, default=0.3) - - return parser diff --git a/opensora/models/ae/videobase/vqvae/videogpt/videogpt/resnet.py b/opensora/models/ae/videobase/vqvae/videogpt/videogpt/resnet.py deleted file mode 100644 index d6882f81d..000000000 --- a/opensora/models/ae/videobase/vqvae/videogpt/videogpt/resnet.py +++ /dev/null @@ -1,149 +0,0 @@ -import math -import numpy as np - -import torch -import torch.nn as nn -import torch.nn.functional as F - -from .utils import shift_dim - -class ChannelLayerNorm(nn.Module): - # layer norm on channels - def __init__(self, in_features): - super().__init__() - self.norm = nn.LayerNorm(in_features) - - def forward(self, x): - x = shift_dim(x, 1, -1) - x = self.norm(x) - x = shift_dim(x, -1, 1) - return x - - -class NormReLU(nn.Module): - - def __init__(self, channels, relu=True, affine=True): - super().__init__() - - self.relu = relu - self.norm = ChannelLayerNorm(channels) - - def forward(self, x): - x_float = x.float() - x_float = self.norm(x_float) - x = x_float.type_as(x) - if self.relu: - x = F.relu(x, inplace=True) - return x - - -class ResidualBlock(nn.Module): - - def __init__(self, in_channels, filters, stride, use_projection=False): - super().__init__() - - if use_projection: - self.proj_conv = nn.Conv3d(in_channels, filters, kernel_size=1, - stride=stride, bias=False) - self.proj_bnr = NormReLU(filters, relu=False) - - self.conv1 = nn.Conv3d(in_channels, filters, kernel_size=3, - stride=stride, bias=False, padding=1) - self.bnr1 = NormReLU(filters) - - self.conv2 = nn.Conv3d(filters, filters, kernel_size=3, - stride=1, bias=False, padding=1) - self.bnr2 = NormReLU(filters) - - self.use_projection = use_projection - - def forward(self, x): - shortcut = x - if self.use_projection: - shortcut = self.proj_bnr(self.proj_conv(x)) - x = self.bnr1(self.conv1(x)) - x = self.bnr2(self.conv2(x)) - - return F.relu(x + shortcut, inplace=True) - -class BlockGroup(nn.Module): - - def __init__(self, in_channels, filters, blocks, stride): - super().__init__() - - self.start_block = ResidualBlock(in_channels, filters, stride, use_projection=True) - in_channels = filters - - self.blocks = [] - for _ in range(1, blocks): - self.blocks.append(ResidualBlock(in_channels, filters, 1)) - self.blocks = nn.Sequential(*self.blocks) - - def forward(self, x): - x = self.start_block(x) - x = self.blocks(x) - return x - - -class ResNet(nn.Module): - - def __init__(self, in_channels, layers, width_multiplier, - stride, resnet_dim=240, cifar_stem=True): - super().__init__() - self.width_multiplier = width_multiplier - self.resnet_dim = resnet_dim - - assert all([int(math.log2(d)) == math.log2(d) for d in stride]), stride - n_times_downsample = np.array([int(math.log2(d)) for d in stride]) - - if cifar_stem: - self.stem = nn.Sequential( - nn.Conv3d(in_channels, 64 * width_multiplier, - kernel_size=3, padding=1, bias=False), - NormReLU(64 * width_multiplier) - ) - else: - stride = tuple([2 if d > 0 else 1 for d in n_times_downsample]) - n_times_downsample -= 1 # conv - n_times_downsample[-2:] = n_times_downsample[-2:] - 1 # pooling - self.stem = nn.Sequential( - nn.Conv3d(in_channels, 64 * width_multiplier, - kernel_size=7, stride=stride, bias=False, - padding=3), - NormReLU(64 * width_multiplier), - nn.MaxPool3d(kernel_size=3, stride=(1, 2, 2), padding=1) - ) - - self.group1 = BlockGroup(64 * width_multiplier, 64 * width_multiplier, - blocks=layers[0], stride=1) - - stride = tuple([2 if d > 0 else 1 for d in n_times_downsample]) - n_times_downsample -= 1 - self.group2 = BlockGroup(64 * width_multiplier, 128 * width_multiplier, - blocks=layers[1], stride=stride) - - stride = tuple([2 if d > 0 else 1 for d in n_times_downsample]) - n_times_downsample -= 1 - self.group3 = BlockGroup(128 * width_multiplier, 256 * width_multiplier, - blocks=layers[2], stride=stride) - - stride = tuple([2 if d > 0 else 1 for d in n_times_downsample]) - n_times_downsample -= 1 - self.group4 = BlockGroup(256 * width_multiplier, resnet_dim, - blocks=layers[3], stride=stride) - assert all([d <= 0 for d in n_times_downsample]), f'final downsample {n_times_downsample}' - - def forward(self, x): - x = self.stem(x) - x = self.group1(x) - x = self.group2(x) - x = self.group3(x) - x = self.group4(x) - x = shift_dim(x, 1, -1) - - return x - - -def resnet34(width_multiplier, stride, cifar_stem=True, resnet_dim=240): - return ResNet(3, [3, 4, 6, 3], width_multiplier, - stride, cifar_stem=cifar_stem, resnet_dim=resnet_dim) diff --git a/opensora/models/ae/videobase/vqvae/videogpt/videogpt/utils.py b/opensora/models/ae/videobase/vqvae/videogpt/videogpt/utils.py deleted file mode 100644 index d35f44d21..000000000 --- a/opensora/models/ae/videobase/vqvae/videogpt/videogpt/utils.py +++ /dev/null @@ -1,87 +0,0 @@ -# Shifts src_tf dim to dest dim -# i.e. shift_dim(x, 1, -1) would be (b, c, t, h, w) -> (b, t, h, w, c) -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 - -# reshapes tensor start from dim i (inclusive) -# to dim j (exclusive) to the desired shape -# e.g. if x.shape = (b, thw, c) then -# view_range(x, 1, 2, (t, h, w)) returns -# x of shape (b, t, h, w, c) -def view_range(x, i, j, shape): - shape = tuple(shape) - - n_dims = len(x.shape) - if i < 0: - i = n_dims + i - - if j is None: - j = n_dims - elif j < 0: - j = n_dims + j - - assert 0 <= i < j <= n_dims - - x_shape = x.shape - target_shape = x_shape[:i] + shape + x_shape[j:] - return x.view(target_shape) - - -def tensor_slice(x, begin, size): - assert all([b >= 0 for b in begin]) - size = [l - b if s == -1 else s - for s, b, l in zip(size, begin, x.shape)] - assert all([s >= 0 for s in size]) - - slices = [slice(b, b + s) for b, s in zip(begin, size)] - return x[slices] - - -import math -import numpy as np -import skvideo.io -def save_video_grid(video, fname, nrow=None): - b, c, t, h, w = video.shape - video = video.permute(0, 2, 3, 4, 1) - video = (video.cpu().numpy() * 255).astype('uint8') - - if nrow is None: - nrow = math.ceil(math.sqrt(b)) - ncol = math.ceil(b / nrow) - padding = 1 - video_grid = np.zeros((t, (padding + h) * nrow + padding, - (padding + w) * ncol + padding, c), dtype='uint8') - 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] - - skvideo.io.vwrite(fname, video_grid, inputdict={'-r': '5'}) - print('saved videos to', fname) - - diff --git a/opensora/models/ae/videobase/vqvae/videogpt/videogpt/vqvae.py b/opensora/models/ae/videobase/vqvae/videogpt/videogpt/vqvae.py deleted file mode 100644 index fb7fa82a4..000000000 --- a/opensora/models/ae/videobase/vqvae/videogpt/videogpt/vqvae.py +++ /dev/null @@ -1,324 +0,0 @@ -import math -import argparse -import numpy as np - -import pytorch_lightning as pl -import torch -import torch.nn as nn -import torch.nn.functional as F -import torch.distributed as dist - - - -from .attention import MultiHeadAttention -from .utils import shift_dim - - -class VQVAE(pl.LightningModule): - def __init__(self, args): - super().__init__() - self.args = args - self.embedding_dim = args.embedding_dim - self.n_codes = args.n_codes - - self.encoder = Encoder(args.n_hiddens, args.n_res_layers, args.downsample) - self.decoder = Decoder(args.n_hiddens, args.n_res_layers, args.downsample) - - self.pre_vq_conv = SamePadConv3d(args.n_hiddens, args.embedding_dim, 1) - self.post_vq_conv = SamePadConv3d(args.embedding_dim, args.n_hiddens, 1) - - self.codebook = Codebook(args.n_codes, args.embedding_dim) - self.save_hyperparameters() - - @property - def latent_shape(self): - input_shape = (self.args.sequence_length, self.args.resolution, - self.args.resolution) - return tuple([s // d for s, d in zip(input_shape, - self.args.downsample)]) - - def encode(self, x, include_embeddings=False): - h = self.pre_vq_conv(self.encoder(x)) - print('before vq', h.max(), h.min()) - vq_output = self.codebook(h) - if include_embeddings: - return vq_output['encodings'], vq_output['embeddings'] - else: - return vq_output['encodings'] - - def decode(self, encodings): - - h = F.embedding(encodings, self.codebook.embeddings) - # print('after vq', h.max(), h.min()) - h = self.post_vq_conv(shift_dim(h, -1, 1)) - return self.decoder(h) - - def forward(self, x): - z = self.pre_vq_conv(self.encoder(x)) - vq_output = self.codebook(z) - x_recon = self.decoder(self.post_vq_conv(vq_output['embeddings'])) - recon_loss = F.mse_loss(x_recon, x) / 0.06 - - return recon_loss, x_recon, vq_output - - def training_step(self, batch, batch_idx): - x = batch['video'] - recon_loss, _, vq_output = self.forward(x) - commitment_loss = vq_output['commitment_loss'] - loss = recon_loss + commitment_loss - return loss - - def validation_step(self, batch, batch_idx): - x = batch['video'] - recon_loss, _, vq_output = self.forward(x) - self.log('val/recon_loss', recon_loss, prog_bar=True) - self.log('val/perplexity', vq_output['perplexity'], prog_bar=True) - self.log('val/commitment_loss', vq_output['commitment_loss'], prog_bar=True) - - def configure_optimizers(self): - return torch.optim.Adam(self.parameters(), lr=3e-4, betas=(0.9, 0.999)) - - @staticmethod - def add_model_specific_args(parent_parser): - parser = argparse.ArgumentParser(parents=[parent_parser], add_help=False) - parser.add_argument('--embedding_dim', type=int, default=256) - parser.add_argument('--n_codes', type=int, default=2048) - parser.add_argument('--n_hiddens', type=int, default=240) - parser.add_argument('--n_res_layers', type=int, default=4) - parser.add_argument('--downsample', nargs='+', type=int, default=(4, 4, 4)) - return parser - - -class AxialBlock(nn.Module): - def __init__(self, n_hiddens, n_head): - super().__init__() - kwargs = dict(shape=(0,) * 3, dim_q=n_hiddens, - dim_kv=n_hiddens, n_head=n_head, - n_layer=1, causal=False, attn_type='axial') - self.attn_w = MultiHeadAttention(attn_kwargs=dict(axial_dim=-2), - **kwargs) - self.attn_h = MultiHeadAttention(attn_kwargs=dict(axial_dim=-3), - **kwargs) - self.attn_t = MultiHeadAttention(attn_kwargs=dict(axial_dim=-4), - **kwargs) - - def forward(self, x): - x = shift_dim(x, 1, -1) - x = self.attn_w(x, x, x) + self.attn_h(x, x, x) + self.attn_t(x, x, x) - x = shift_dim(x, -1, 1) - return x - - -class AttentionResidualBlock(nn.Module): - def __init__(self, n_hiddens): - super().__init__() - self.block = nn.Sequential( - nn.BatchNorm3d(n_hiddens), - nn.ReLU(), - SamePadConv3d(n_hiddens, n_hiddens // 2, 3, bias=False), - nn.BatchNorm3d(n_hiddens // 2), - nn.ReLU(), - SamePadConv3d(n_hiddens // 2, n_hiddens, 1, bias=False), - nn.BatchNorm3d(n_hiddens), - nn.ReLU(), - AxialBlock(n_hiddens, 2) - ) - - def forward(self, x): - return x + self.block(x) - -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 - -class Encoder(nn.Module): - def __init__(self, n_hiddens, n_res_layers, downsample): - super().__init__() - n_times_downsample = np.array([int(math.log2(d)) for d in downsample]) - self.convs = nn.ModuleList() - max_ds = n_times_downsample.max() - for i in range(max_ds): - in_channels = 3 if i == 0 else n_hiddens - stride = tuple([2 if d > 0 else 1 for d in n_times_downsample]) - conv = SamePadConv3d(in_channels, n_hiddens, 4, stride=stride) - self.convs.append(conv) - n_times_downsample -= 1 - self.conv_last = SamePadConv3d(in_channels, n_hiddens, kernel_size=3) - - self.res_stack = nn.Sequential( - *[AttentionResidualBlock(n_hiddens) - for _ in range(n_res_layers)], - nn.BatchNorm3d(n_hiddens), - nn.ReLU() - ) - - def forward(self, x): - h = x - for conv in self.convs: - h = F.relu(conv(h)) - h = self.conv_last(h) - h = self.res_stack(h) - return h - - -class Decoder(nn.Module): - def __init__(self, n_hiddens, n_res_layers, upsample): - super().__init__() - self.res_stack = nn.Sequential( - *[AttentionResidualBlock(n_hiddens) - for _ in range(n_res_layers)], - nn.BatchNorm3d(n_hiddens), - nn.ReLU() - ) - - n_times_upsample = np.array([int(math.log2(d)) for d in upsample]) - max_us = n_times_upsample.max() - self.convts = nn.ModuleList() - for i in range(max_us): - out_channels = 3 if i == max_us - 1 else n_hiddens - us = tuple([2 if d > 0 else 1 for d in n_times_upsample]) - convt = SamePadConvTranspose3d(n_hiddens, out_channels, 4, - stride=us) - self.convts.append(convt) - n_times_upsample -= 1 - - def forward(self, x): - h = self.res_stack(x) - for i, convt in enumerate(self.convts): - h = convt(h) - if i < len(self.convts) - 1: - h = F.relu(h) - return h - - -# Does not support dilation -class SamePadConv3d(nn.Module): - def __init__(self, in_channels, out_channels, kernel_size, stride=1, bias=True): - super().__init__() - if isinstance(kernel_size, int): - kernel_size = (kernel_size,) * 3 - if isinstance(stride, int): - stride = (stride,) * 3 - - # assumes that the input shape is divisible by stride - total_pad = tuple([k - s for k, s in zip(kernel_size, stride)]) - pad_input = [] - for p in total_pad[::-1]: # reverse since F.pad starts from last dim - pad_input.append((p // 2 + p % 2, p // 2)) - pad_input = sum(pad_input, tuple()) - self.pad_input = pad_input - - self.conv = nn.Conv3d(in_channels, out_channels, kernel_size, - stride=stride, padding=0, bias=bias) - - def forward(self, x): - return self.conv(F.pad(x, self.pad_input)) - - -class SamePadConvTranspose3d(nn.Module): - def __init__(self, in_channels, out_channels, kernel_size, stride=1, bias=True): - super().__init__() - if isinstance(kernel_size, int): - kernel_size = (kernel_size,) * 3 - if isinstance(stride, int): - stride = (stride,) * 3 - - total_pad = tuple([k - s for k, s in zip(kernel_size, stride)]) - pad_input = [] - for p in total_pad[::-1]: # reverse since F.pad starts from last dim - pad_input.append((p // 2 + p % 2, p // 2)) - pad_input = sum(pad_input, tuple()) - self.pad_input = pad_input - - self.convt = nn.ConvTranspose3d(in_channels, out_channels, kernel_size, - stride=stride, bias=bias, - padding=tuple([k - 1 for k in kernel_size])) - - def forward(self, x): - return self.convt(F.pad(x, self.pad_input)) - - diff --git a/opensora/train/train_causalvae.py b/opensora/train/train_causalvae.py new file mode 100644 index 000000000..552982dd4 --- /dev/null +++ b/opensora/train/train_causalvae.py @@ -0,0 +1,104 @@ +import sys +sys.path.append(".") + +from opensora.models.ae.videobase import ( + CausalVAEConfiguration, + CausalVAEDataset, + CausalVAEModel, + CausalVAETrainer, +) +import argparse +from typing import Optional +from accelerate.utils import set_seed +from transformers import HfArgumentParser, TrainingArguments +from dataclasses import dataclass, field, asdict +from typing import Tuple, List +import torch.distributed as dist +import os +import cv2 +import torch +import numpy as np + +def array_to_video(image_array, fps: float = 30.0, output_file: str = 'output_video.mp4') -> None: + height, width, channels = image_array[0].shape + fourcc = cv2.VideoWriter_fourcc(*'mp4v') # type: ignore + 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, -0.5, 0.5) + x = (x + 0.5) + x = x.permute(1, 2, 3, 0).numpy() # (C, T, H, W) -> (T, H, W, C) + x = (255*x).astype(np.uint8) + array_to_video(x, fps=fps, output_file=output_file) + # imageio.mimwrite(output_file, x, fps=fps, quality=9) + return + + +@dataclass +class DataArguments: + data_path: str = field(default="./UCF101") + video_num_frames: int = field(default=17) + sample_rate: int = field(default=1) + +@dataclass +class CausalVAEArguments: + in_channels: int = field(default=3) + out_channels: int = field(default=3) + hidden_size: int = field(default=128) + z_channels: int = field(default=16) + ch_mult: Tuple[int] = field(default=(1, 1, 2, 2, 4)) + num_res_block: int = field(default=2) + attn_resolutions: List[int] = field(default_factory=lambda: [16]) + dropout: float = field(default=0.0) + resolution: int = field(default=256) + attn_type: str = field(default="vanilla3D") + use_linear_attn: bool = field(default=False) + embed_dim: int = field(default=16) + time_compress: int = field(default=2) + logvar_init: float = field(default=0.0) + kl_weight: float = field(default=1e-6) + pixelloss_weight: int = field(default=1) + perceptual_weight: int = field(default=1) + disc_loss: str = field(default="hinge") + +@dataclass +class CausalVAETrainingArguments(TrainingArguments): + remove_unused_columns: Optional[bool] = field( + default=False, metadata={"help": "Remove columns not required by the model when using an nlp.Dataset."} + ) + +def train(args, vqvae_args, training_args): + # Load Config + config = CausalVAEConfiguration(**asdict(vqvae_args)) + # Load Model + if training_args.resume_from_checkpoint is not None: + model = CausalVAEModel.load_from_checkpoint(training_args.resume_from_checkpoint) + else: + model = CausalVAEModel(config) + if (dist.is_initialized() and dist.get_rank() == 0) or not dist.is_initialized(): + print(model) + # Load Dataset + dataset = CausalVAEDataset(args.data_path, sequence_length=args.video_num_frames, resolution=config.resolution, sample_rate=args.sample_rate) + custom_to_video(dataset[0]['video'], fps=10, output_file="dataset.mp4") + # Load Trainer + trainer = CausalVAETrainer(model, training_args, train_dataset=dataset) + trainer.train(resume_from_checkpoint=training_args.resume_from_checkpoint) + + +if __name__ == "__main__": + parser = HfArgumentParser((CausalVAEArguments, CausalVAETrainingArguments, DataArguments)) + + if not sys.argv[-1].endswith(".yaml"): + raise Exception("Please use yaml config.") + + vqvae_args, training_args, data_args = parser.parse_yaml_file(yaml_file=os.path.abspath(sys.argv[-1])) + args = argparse.Namespace(**vars(vqvae_args), **vars(training_args), **vars(data_args)) + set_seed(args.seed) + train(args, vqvae_args, training_args) diff --git a/opensora/utils/taming_download.py b/opensora/utils/taming_download.py new file mode 100644 index 000000000..5a62be778 --- /dev/null +++ b/opensora/utils/taming_download.py @@ -0,0 +1,145 @@ +"""Modified from https://github.com/CompVis/taming-transformers.git""" + +import os, hashlib +import requests +from tqdm import tqdm + +URL_MAP = { + "vgg_lpips": "https://heibox.uni-heidelberg.de/f/607503859c864bc1b30b/?dl=1" +} + +CKPT_MAP = { + "vgg_lpips": "vgg.pth" +} + +MD5_MAP = { + "vgg_lpips": "d507d7349b931f0638a25a48a722f98a" +} + + +def download(url, local_path, chunk_size=1024): + os.makedirs(os.path.split(local_path)[0], exist_ok=True) + with requests.get(url, stream=True) as r: + total_size = int(r.headers.get("content-length", 0)) + with tqdm(total=total_size, unit="B", unit_scale=True) as pbar: + with open(local_path, "wb") as f: + for data in r.iter_content(chunk_size=chunk_size): + if data: + f.write(data) + pbar.update(chunk_size) + + +def md5_hash(path): + with open(path, "rb") as f: + content = f.read() + return hashlib.md5(content).hexdigest() + + +def get_ckpt_path(name, root, check=False): + assert name in URL_MAP + path = os.path.join(root, CKPT_MAP[name]) + if not os.path.exists(path) or (check and not md5_hash(path) == MD5_MAP[name]): + print("Downloading {} model from {} to {}".format(name, URL_MAP[name], path)) + download(URL_MAP[name], path) + md5 = md5_hash(path) + assert md5 == MD5_MAP[name], md5 + return path + + +class KeyNotFoundError(Exception): + def __init__(self, cause, keys=None, visited=None): + self.cause = cause + self.keys = keys + self.visited = visited + messages = list() + if keys is not None: + messages.append("Key not found: {}".format(keys)) + if visited is not None: + messages.append("Visited: {}".format(visited)) + messages.append("Cause:\n{}".format(cause)) + message = "\n".join(messages) + super().__init__(message) + + +def retrieve( + list_or_dict, key, splitval="/", default=None, expand=True, pass_success=False +): + """Given a nested list or dict return the desired value at key expanding + callable nodes if necessary and :attr:`expand` is ``True``. The expansion + is done in-place. + + Parameters + ---------- + list_or_dict : list or dict + Possibly nested list or dictionary. + key : str + key/to/value, path like string describing all keys necessary to + consider to get to the desired value. List indices can also be + passed here. + splitval : str + String that defines the delimiter between keys of the + different depth levels in `key`. + default : obj + Value returned if :attr:`key` is not found. + expand : bool + Whether to expand callable nodes on the path or not. + + Returns + ------- + The desired value or if :attr:`default` is not ``None`` and the + :attr:`key` is not found returns ``default``. + + Raises + ------ + Exception if ``key`` not in ``list_or_dict`` and :attr:`default` is + ``None``. + """ + + keys = key.split(splitval) + + success = True + try: + visited = [] + parent = None + last_key = None + for key in keys: + if callable(list_or_dict): + if not expand: + raise KeyNotFoundError( + ValueError( + "Trying to get past callable node with expand=False." + ), + keys=keys, + visited=visited, + ) + list_or_dict = list_or_dict() + parent[last_key] = list_or_dict + + last_key = key + parent = list_or_dict + + try: + if isinstance(list_or_dict, dict): + list_or_dict = list_or_dict[key] + else: + list_or_dict = list_or_dict[int(key)] + except (KeyError, IndexError, ValueError) as e: + raise KeyNotFoundError(e, keys=keys, visited=visited) + + visited += [key] + # final expansion of retrieved value + if expand and callable(list_or_dict): + list_or_dict = list_or_dict() + parent[last_key] = list_or_dict + except KeyNotFoundError as e: + if default is None: + raise e + else: + list_or_dict = default + success = False + + if not pass_success: + return list_or_dict + else: + return list_or_dict, success +