Merge pull request #162 from qqingzheng/add_causal_vae

[feat] add causalvae ✨

Former-commit-id: 4c1b27c277f6dee020d7aa9d6f719415f0eed898

.gitignore

@@ -10,3 +10,5 @@ build/
 opensora.egg-info/
 wandb/
 .idea
+*.ipynb
+*.jpg

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 = }')

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)

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)

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)

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
 """

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

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

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

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}")

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

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"))

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

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)

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)

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)

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
+    )

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)

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

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

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

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
+    

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"))

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"))

opensora/models/ae/videobase/vqvae/videogpt/requirements.txt

@@ -1,9 +0,0 @@
-tqdm
-requests
-h5py
-scikit-video
-av
-gradio
-moviepy
-imageio
-gdown

opensora/models/ae/videobase/vqvae/videogpt/scripts/compute_fvd.py

@@ -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()

opensora/models/ae/videobase/vqvae/videogpt/scripts/preprocess/bair/bair_extract_images.py

@@ -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')

opensora/models/ae/videobase/vqvae/videogpt/scripts/preprocess/bair/bair_image_to_hdf5.py

@@ -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()

opensora/models/ae/videobase/vqvae/videogpt/scripts/preprocess/bair/create_bair_dataset.sh

@@ -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

opensora/models/ae/videobase/vqvae/videogpt/scripts/preprocess/ucf101/create_ucf_dataset.sh

@@ -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

opensora/models/ae/videobase/vqvae/videogpt/scripts/preprocess/ucf101/ucf_split_train_test.py

@@ -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')

opensora/models/ae/videobase/vqvae/videogpt/scripts/sample_videogpt.py

@@ -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')

opensora/models/ae/videobase/vqvae/videogpt/scripts/train_videogpt.py

@@ -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()
-

opensora/models/ae/videobase/vqvae/videogpt/scripts/train_vqvae.py

@@ -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()
-

opensora/models/ae/videobase/vqvae/videogpt/setup.py

@@ -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)

opensora/models/ae/videobase/vqvae/videogpt/videogpt/__init__.py

@@ -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
-

opensora/models/ae/videobase/vqvae/videogpt/videogpt/attention.py

@@ -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)

opensora/models/ae/videobase/vqvae/videogpt/videogpt/data.py

@@ -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()

opensora/models/ae/videobase/vqvae/videogpt/videogpt/download.py

@@ -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

opensora/models/ae/videobase/vqvae/videogpt/videogpt/fvd/convert_tf_pretrained.py

@@ -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')
-

opensora/models/ae/videobase/vqvae/videogpt/videogpt/fvd/fvd.py

@@ -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

opensora/models/ae/videobase/vqvae/videogpt/videogpt/fvd/pytorch_i3d.py

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

opensora/models/ae/videobase/vqvae/videogpt/videogpt/gpt.py

@@ -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

opensora/models/ae/videobase/vqvae/videogpt/videogpt/resnet.py

@@ -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)

opensora/models/ae/videobase/vqvae/videogpt/videogpt/utils.py

@@ -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)
-
-

opensora/models/ae/videobase/vqvae/videogpt/videogpt/vqvae.py

@@ -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))
-
-

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)

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
+