add sample script

Former-commit-id: 9ce847651bb6dba28c51ce5807e222ae2902a78a

README.md

@@ -47,12 +47,13 @@ Project stages:
 - [x] Add VAE model, which is borrowed from [Stable Diffusion](https://github.com/CompVis/latent-diffusion).
 - [x] Joint dynamic mask input with VAE.
 - [ ] Make the codebase ready for the cluster training. Add SLURM scripts.
-- [ ] Add sampling script.
+- [x] Add sampling script.
 - [ ] Incorporating [SiT](https://github.com/willisma/SiT).
 
 #### Train models that boost resolution and duration
 - [ ] Add [PI](https://arxiv.org/abs/2306.15595) to support out-of-domain size.
 - [x] Add frame interpolation model.
+- [x] Add super resolution model.
 - [ ] Add accelerate to automatically manage training.
 - [ ] Joint training with images.
 

opensora/dataset/__init__.py

@@ -8,6 +8,7 @@ from torchvision.transforms import Lambda
 
 from .transform import ToTensorVideo, TemporalRandomCrop, RandomHorizontalFlipVideo, CenterCropResizeVideo
 from .ucf101 import UCF101
+from .sky_datasets import Sky
 
 
 def getdataset(args):
@@ -35,4 +36,14 @@ def getdataset(args):
                 norm_fun,
             ]
         )
-        return UCF101(args, transform=transform, temporal_sample=temporal_sample)
+        return UCF101(args, transform=transform, temporal_sample=temporal_sample)
+    elif args.dataset == 'sky':
+        transform_sky = transforms.Compose([
+            ToTensorVideo(),
+            CenterCropResizeVideo(args.max_image_size),
+            # video_transforms.RandomHorizontalFlipVideo(),
+            norm_fun
+        ])
+        return Sky(args, transform=transform_sky, temporal_sample=temporal_sample)
+    else:
+        raise NotImplementedError(args.dataset)

opensora/dataset/sky_datasets.py

@@ -0,0 +1,112 @@
+import os
+import torch
+import random
+import torch.utils.data as data
+
+import numpy as np
+
+from PIL import Image
+
+IMG_EXTENSIONS = ['.jpg', '.JPG', '.jpeg', '.JPEG', '.png', '.PNG']
+
+def is_image_file(filename):
+    return any(filename.endswith(extension) for extension in IMG_EXTENSIONS)
+
+class Sky(data.Dataset):
+    def __init__(self, configs, transform, temporal_sample=None, train=True):
+
+        self.configs = configs
+        self.data_path = configs.data_path
+        self.transform = transform
+        self.temporal_sample = temporal_sample
+        self.target_video_len = self.configs.num_frames
+        self.frame_interval = self.configs.sample_rate
+        self.data_all = self.load_video_frames(self.data_path)
+
+    def __getitem__(self, index):
+
+        vframes = self.data_all[index]
+        total_frames = len(vframes)
+
+        # Sampling video frames
+        start_frame_ind, end_frame_ind = self.temporal_sample(total_frames)
+        assert end_frame_ind - start_frame_ind >= self.target_video_len
+        frame_indice = np.linspace(start_frame_ind, end_frame_ind-1, num=self.target_video_len, dtype=int) # start, stop, num=50
+
+        select_video_frames = vframes[frame_indice[0]: frame_indice[-1]+1: self.frame_interval] 
+
+        video_frames = []
+        for path in select_video_frames:
+            video_frame = torch.as_tensor(np.array(Image.open(path), dtype=np.uint8, copy=True)).unsqueeze(0)
+            video_frames.append(video_frame)
+        video_clip = torch.cat(video_frames, dim=0).permute(0, 3, 1, 2)
+        video_clip = self.transform(video_clip)
+        video_clip = video_clip.transpose(0, 1)  # T C H W -> C T H W
+
+        return video_clip, 1
+
+    def __len__(self):
+        return self.video_num
+    
+    def load_video_frames(self, dataroot):
+        data_all = []
+        frame_list = os.walk(dataroot)
+        for _, meta in enumerate(frame_list):
+            root = meta[0]
+            try:
+                frames = sorted(meta[2], key=lambda item: int(item.split('.')[0].split('_')[-1]))
+            except:
+                # pass
+                print(meta[0]) # root
+                print(meta[2]) # files
+            frames = [os.path.join(root, item) for item in frames if is_image_file(item)]
+            if len(frames) > max(0, self.target_video_len * self.frame_interval): # need all > (16 * frame-interval) videos
+            # if len(frames) >= max(0, self.target_video_len): # need all > 16 frames videos
+                data_all.append(frames)
+        self.video_num = len(data_all)
+        return data_all
+    
+
+if __name__ == '__main__':
+
+    import argparse
+    import torchvision
+    import video_transforms
+    import torch.utils.data as data
+
+    from torchvision import transforms
+    from torchvision.utils import save_image
+
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--num_frames", type=int, default=16)
+    parser.add_argument("--frame_interval", type=int, default=4)
+    parser.add_argument("--data-path", type=str, default="/path/to/datasets/sky_timelapse/sky_train/")
+    config = parser.parse_args()
+
+
+    target_video_len = config.num_frames
+
+    temporal_sample = video_transforms.TemporalRandomCrop(target_video_len * config.frame_interval)
+    trans = transforms.Compose([
+        video_transforms.ToTensorVideo(),
+        # video_transforms.CenterCropVideo(256),
+        video_transforms.CenterCropResizeVideo(256),
+        # video_transforms.RandomHorizontalFlipVideo(),
+        transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True)
+    ])
+
+    taichi_dataset = Sky(config, transform=trans, temporal_sample=temporal_sample)
+    print(len(taichi_dataset))
+    taichi_dataloader = data.DataLoader(dataset=taichi_dataset, batch_size=1, shuffle=False, num_workers=1)
+
+    for i, video_data in enumerate(taichi_dataloader):
+        print(video_data['video'].shape)
+        
+        # print(video_data.dtype)
+        # for i in range(target_video_len):
+        #     save_image(video_data[0][i], os.path.join('./test_data', '%04d.png' % i), normalize=True, value_range=(-1, 1))
+
+        # video_ = ((video_data[0] * 0.5 + 0.5) * 255).add_(0.5).clamp_(0, 255).to(dtype=torch.uint8).cpu().permute(0, 2, 3, 1)
+        # torchvision.io.write_video('./test_data' + 'test.mp4', video_, fps=8)
+        # exit()

opensora/dataset/ucf101.py

@@ -13,30 +13,8 @@ from pytorchvideo.transforms import ApplyTransformToKey, ShortSideScale, Uniform
 from torch.nn import functional as F
 import random
 
+from opensora.utils.dataset_utils import DecordInit
 
-class DecordInit(object):
-    """Using Decord(https://github.com/dmlc/decord) to initialize the video_reader."""
-
-    def __init__(self, num_threads=1):
-        self.num_threads = num_threads
-        self.ctx = decord.cpu(0)
-
-    def __call__(self, filename):
-        """Perform the Decord initialization.
-        Args:
-            results (dict): The resulting dict to be modified and passed
-                to the next transform in pipeline.
-        """
-        reader = decord.VideoReader(filename,
-                                    ctx=self.ctx,
-                                    num_threads=self.num_threads)
-        return reader
-
-    def __repr__(self):
-        repr_str = (f'{self.__class__.__name__}('
-                    f'sr={self.sr},'
-                    f'num_threads={self.num_threads})')
-        return repr_str
 
 class UCF101(Dataset):
     def __init__(self, args, transform, temporal_sample):

opensora/models/diffusion/latte/latte.py

@@ -371,11 +371,13 @@ class Latte(nn.Module):
         """
         if use_fp16:
             x = x.to(dtype=torch.float16)
-        attention_mask_spatial = rearrange(attention_mask, 'b t h w -> (b t) h w')
-        attention_mask_spatial = self.make_mask(attention_mask_spatial, x.dtype)
+        attention_mask_temproal, attention_mask_spatial = None, None
+        if attention_mask is not None:
+            attention_mask_spatial = rearrange(attention_mask, 'b t h w -> (b t) h w')
+            attention_mask_spatial = self.make_mask(attention_mask_spatial, x.dtype)
 
-        attention_mask_temproal = rearrange(attention_mask, 'b t h w -> (b h w) t')
-        attention_mask_temproal = self.make_mask(attention_mask_temproal, x.dtype)
+            attention_mask_temproal = rearrange(attention_mask, 'b t h w -> (b h w) t')
+            attention_mask_temproal = self.make_mask(attention_mask_temproal, x.dtype)
 
         batches, frames, channels, high, weight = x.shape
         num_patches_height = high // self.patch_size

opensora/sample/sample.py

@@ -0,0 +1,127 @@
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""
+Sample new images from a pre-trained Latte.
+"""
+import os
+import sys
+
+from opensora.models.ae import ae_channel_config, getae, ae_stride_config
+from opensora.models.diffusion import Diffusion_models
+from opensora.models.diffusion.diffusion import create_diffusion
+from opensora.utils.utils import find_model
+
+import torch
+import argparse
+
+from einops import rearrange
+import imageio
+
+torch.backends.cuda.matmul.allow_tf32 = True
+torch.backends.cudnn.allow_tf32 = True
+
+def main(args):
+    # Setup PyTorch:
+    # torch.manual_seed(args.seed)
+    torch.set_grad_enabled(False)
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    using_cfg = args.cfg_scale > 1.0
+
+    # Load model:
+    latent_size = args.image_size // ae_stride_config[args.ae][1]
+    args.latent_size = latent_size
+    model = Diffusion_models[args.model](
+        input_size=latent_size,
+        num_classes=args.num_classes,
+        in_channels=ae_channel_config[args.ae],
+        extras=args.extras
+    ).to(device)
+
+    if args.use_compile:
+        model = torch.compile(model)
+
+    # a pre-trained model or load a custom Latte checkpoint from train.py:
+    ckpt_path = args.ckpt
+    state_dict = find_model(ckpt_path)
+    model.load_state_dict(state_dict)
+
+    model.eval()  # important!
+    diffusion = create_diffusion(str(args.num_sampling_steps))
+    ae = getae(args).to(device)
+
+    if args.use_fp16:
+        print('WARNING: using half percision for inferencing!')
+        ae.to(dtype=torch.float16)
+        model.to(dtype=torch.float16)
+
+    # Labels to condition the model with (feel free to change):
+
+    # Create sampling noise:
+    if args.use_fp16:
+        z = torch.randn(1, args.num_frames, model.in_channels, latent_size, latent_size, dtype=torch.float16, device=device) # b c f h w
+    else:
+        z = torch.randn(1, args.num_frames, model.in_channels, latent_size, latent_size, device=device)
+
+    # Setup classifier-free guidance:
+    if using_cfg:
+        z = torch.cat([z, z], 0)
+        y = torch.randint(0, args.num_classes, (1,), device=device)
+        y_null = torch.tensor([args.num_classes] * 1, device=device)
+        y = torch.cat([y, y_null], dim=0)
+        model_kwargs = dict(y=y, cfg_scale=args.cfg_scale, use_fp16=args.use_fp16)
+        sample_fn = model.forward_with_cfg
+    else:
+        sample_fn = model.forward
+        model_kwargs = dict(y=None, use_fp16=args.use_fp16)
+
+    # Sample images:
+    if args.sample_method == 'ddim':
+        samples = diffusion.ddim_sample_loop(
+            sample_fn, z.shape, z, clip_denoised=False, model_kwargs=model_kwargs, progress=True, device=device
+        )
+    elif args.sample_method == 'ddpm':
+        samples = diffusion.p_sample_loop(
+            sample_fn, z.shape, z, clip_denoised=False, model_kwargs=model_kwargs, progress=True, device=device
+        )
+
+    if args.use_fp16:
+        samples = samples.to(dtype=torch.float16)
+    b, f, c, h, w = samples.shape
+    samples = rearrange(samples, 'b f c h w -> (b f) c h w')
+    samples = ae.decode(samples)
+    samples = rearrange(samples, '(b f) c h w -> b f c h w', b=b)
+    # Save and display images:
+
+    if not os.path.exists(args.save_video_path):
+        os.makedirs(args.save_video_path)
+
+
+    video_ = ((samples[0] * 0.5 + 0.5) * 255).add_(0.5).clamp_(0, 255).to(dtype=torch.uint8).cpu().permute(0, 2, 3, 1).contiguous()
+    video_save_path = os.path.join(args.save_video_path, 'sample' + '.mp4')
+    print(video_save_path)
+    imageio.mimwrite(video_save_path, video_, fps=args.fps, quality=9)
+    print('save path {}'.format(args.save_video_path))
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--ckpt", type=str, default="")
+    parser.add_argument("--model", type=str, default='Latte-XL/122')
+    parser.add_argument("--ae", type=str, default='stabilityai/sd-vae-ft-mse')
+    parser.add_argument("--save-video-path", type=str, default="./sample_videos/")
+    parser.add_argument("--fps", type=int, default=10)
+    parser.add_argument("--num-classes", type=int, default=101)
+    parser.add_argument("--num-frames", type=int, default=16)
+    parser.add_argument("--image-size", type=int, default=256, choices=[256, 512])
+    parser.add_argument("--extras", type=int, default=1)
+    parser.add_argument("--num-sampling-steps", type=int, default=250)
+    parser.add_argument("--cfg-scale", type=float, default=1.0)
+    parser.add_argument("--use-fp16", action="store_true")
+    parser.add_argument("--use-compile", action="store_true")
+    parser.add_argument("--sample-method", type=str, default='ddpm')
+    args = parser.parse_args()
+    main(args)

opensora/utils/convert_video_offline.py

@@ -13,14 +13,20 @@ import imageio
 import torch
 
 # the first flag below was False when we tested this script but True makes A100 training a lot faster:
+import torchvision
 from einops import rearrange
+from torchvision.transforms import Compose, Lambda
 from tqdm import tqdm
 
+from opensora.dataset import ToTensorVideo, CenterCropResizeVideo, TemporalRandomCrop, UCF101
+from opensora.dataset.ucf101 import DecordInit
+from opensora.models.ae import videovqvae, videovae, vqvae, getae, vae
+
 torch.backends.cuda.matmul.allow_tf32 = True
 torch.backends.cudnn.allow_tf32 = True
 import torch.distributed as dist
 from torch.nn.parallel import DistributedDataParallel as DDP
-from torch.utils.data import DataLoader
+from torch.utils.data import DataLoader, Dataset
 from torch.utils.data.distributed import DistributedSampler
 from torchvision.datasets import ImageFolder
 from torchvision import transforms
@@ -34,40 +40,42 @@ import argparse
 import logging
 import os
 
-from datasets import video_transforms, UCF101
-from diffusers.models import AutoencoderKL
 
 
-#################################################################################
-#                             Training Helper Functions                         #
-#################################################################################
+class FolderExtract(Dataset):
+    def __init__(self, args, transform):
+        self.data_path = args.data_path
+        self.sample_rate = args.sample_rate
+        self.transform = transform
+        self.v_decoder = DecordInit()
+        self.samples = list(glob(f'{self.data_path}/*mp4'))
 
+    def __len__(self):
+        return len(self.samples)
 
-def center_crop_arr(pil_image, image_size):
-    """
-    Center cropping implementation from ADM.
-    https://github.com/openai/guided-diffusion/blob/8fb3ad9197f16bbc40620447b2742e13458d2831/guided_diffusion/image_datasets.py#L126
-    """
-    while min(*pil_image.size) >= 2 * image_size:
-        pil_image = pil_image.resize(
-            tuple(x // 2 for x in pil_image.size), resample=Image.BOX
-        )
+    def __getitem__(self, idx):
+        video_path = self.samples[idx]
+        video = self.decord_read(video_path)
+        video = self.transform(video)  # T C H W -> T C H W
+        return video
 
-    scale = image_size / min(*pil_image.size)
-    pil_image = pil_image.resize(
-        tuple(round(x * scale) for x in pil_image.size), resample=Image.BICUBIC
-    )
+    def tv_read(self, path):
+        vframes, aframes, info = torchvision.io.read_video(filename=path, pts_unit='sec', output_format='TCHW')
+        total_frames = len(vframes)
+        frame_indice = list(range(total_frames))[::self.sample_rate]
+        video = vframes[frame_indice]
+        return video
 
-    arr = np.array(pil_image)
-    crop_y = (arr.shape[0] - image_size) // 2
-    crop_x = (arr.shape[1] - image_size) // 2
-    return Image.fromarray(arr[crop_y: crop_y + image_size, crop_x: crop_x + image_size])
+    def decord_read(self, path):
+        decord_vr = self.v_decoder(path)
+        total_frames = len(decord_vr)
+        frame_indice = list(range(total_frames))[::self.sample_rate]
+        video_data = decord_vr.get_batch(frame_indice).asnumpy()
+        video_data = torch.from_numpy(video_data)
+        video_data = video_data.permute(0, 3, 1, 2)  # (T, H, W, C) -> (T C H W)
+        return video_data
 
 
-#################################################################################
-#                                  Training Loop                                #
-#################################################################################
-
 def main(args):
     """
     Trains a new DiT model.
@@ -78,32 +86,23 @@ def main(args):
     seed = args.global_seed
     torch.manual_seed(seed)
     torch.cuda.set_device(device)
-    # args.data_path = "E:/UCF-101"
-    args.frame_interval = 3
-    args.num_frames = 16
 
     # Setup a feature folder:
     os.makedirs(args.features_path, exist_ok=True)
     os.makedirs(os.path.join(args.features_path, args.features_name), exist_ok=True)
-    # os.makedirs(os.path.join(args.features_path, 'imagenet256_labels'), exist_ok=True)
 
     # Create model:
-    assert args.image_size % 8 == 0, "Image size must be divisible by 8 (for the VAE encoder)."
-    latent_size = args.image_size // 8
-    vae = AutoencoderKL.from_pretrained(f"stabilityai/sd-vae-ft-{args.vae}", cache_dir='cache_dir').to(device)
+    ae = getae(args).to(device)
 
     # Setup data:
-
-    temporal_sample = video_transforms.TemporalRandomCrop(args.num_frames * args.frame_interval) # 16 1
-    transform_ucf101 = transforms.Compose([
-        video_transforms.ToTensorVideo(),  # TCHW
-        # video_transforms.RandomHorizontalFlipVideo(),
-        video_transforms.UCFCenterCropVideo(args.image_size),
-        transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True)
-    ])
-
-    dataset = UCF101(args, transform=transform_ucf101, temporal_sample=temporal_sample)
-
+    norm_fun = transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True)
+    transform = Compose(
+        [
+            ToTensorVideo(),  # TCHW
+            norm_fun,
+        ]
+    )
+    dataset = FolderExtract(args, transform=transform)
 
     loader = DataLoader(
         dataset,
@@ -116,93 +115,59 @@ def main(args):
 
     train_steps = 0
     for x in tqdm(loader):
-        # x = x.to(device)
-        x = x['video'].to(device, non_blocking=True)
-        # y = y.to(device)
+        x = x.to(device, non_blocking=True)
         with torch.no_grad():
-            # Map input images to latent space + normalize latents:
-            # x = vae.encode(x).latent_dist.sample().mul_(0.18215)
-            b, _, _, _, _ = x.shape
-            x = rearrange(x, 'b f c h w -> (b f) c h w').contiguous()
+            b, _, _, _, _ = x.shape  # b t c h w
+            assert b == 1
+            x = x[0]
             B = x.shape[0]
-            n_sample = 64
+            n_sample = 32
             n_round = B // n_sample
             x_enc = []
-            print(x.shape)
+            # print(x.shape)
             for i in range(n_round+1):
                 if i*n_sample == x.shape[0]:
                     break
-                x_ = vae.encode(x[i*n_sample:(i+1)*n_sample]).latent_dist.sample().mul_(0.18215)
+                x_ = ae.encode(x[i*n_sample:(i+1)*n_sample]) # B T C H W
                 x_enc.append(x_)
             x = torch.cat(x_enc, dim=0)
-            x = rearrange(x, '(b f) c h w -> b f c h w', b=b).contiguous()
+            x = rearrange(x, '(b t) c h w -> b t c h w', b=b).contiguous()
 
 
+            if train_steps == 0:
+                b, t, c, h, w = x.shape
+                samples = rearrange(x, 'b t c h w -> (b t) c h w')
+                samples_dec = []
+                for i in range(n_round+1):
+                    samples_ = ae.decode(samples[i*n_sample:(i+1)*n_sample])
+                    samples_dec.append(samples_)
+                samples = torch.cat(samples_dec, dim=0)
 
-        #     b, f, c, h, w = x.shape
-        #     samples = rearrange(x, 'b f c h w -> (b f) c h w')
-        #     samples_dec = []
-        #     for i in tqdm(range(n_round+1)):
-        #         samples_ = vae.decode(samples[i*n_sample:(i+1)*n_sample] / 0.18215).sample
-        #         samples_dec.append(samples_)
-        #     samples = torch.cat(samples_dec, dim=0)
-        #
-        #     samples = rearrange(samples, '(b f) c h w -> b f c h w', b=b)
-        #
-        # video_ = ((samples[0] * 0.5 + 0.5) * 255).add_(0.5).clamp_(0, 255).to(dtype=torch.uint8).cpu().permute(0, 2, 3, 1).contiguous()
-        #
-        # imageio.mimwrite('1.mp4', video_, fps=8, quality=9)
-        # sys.exit()
+                samples = rearrange(samples, '(b t) c h w -> b t c h w', b=b)
+
+                video_ = ((samples[0] * 0.5 + 0.5) * 255).add_(0.5).clamp_(0, 255).to(dtype=torch.uint8).cpu().permute(0, 2, 3, 1).contiguous()
+
+                imageio.mimwrite('test.mp4', video_, fps=8, quality=9)
 
         x = x.detach().cpu().numpy()  # (1, 4, 32, 32)
         np.save(f'{args.features_path}/{args.features_name}/{train_steps}.npy', x)
 
         # y = y.detach().cpu().numpy()  # (1,)
-        # np.save(f'{args.features_path}/imagenet256_labels/{train_steps}.npy', y)
+        # np.save(f'{args.features_path}/{args.features_name}/{train_steps}.npy', y)
 
         train_steps += 1
         # print(train_steps)
 
 
 if __name__ == "__main__":
-    # Default args here will train DiT-XL/2 with the hyperparameters we used in our paper (except training iters).
-    # parser = argparse.ArgumentParser()
-    # parser.add_argument("--data-path", type=str)
-    # parser.add_argument("--features-path", type=str, default="features")
-    # parser.add_argument("--features-name", type=str, default="512video")
-    # parser.add_argument("--results-dir", type=str, default="results")
-    # parser.add_argument("--image-size", type=int, choices=[128, 256, 512], default=512)
-    # parser.add_argument("--num-classes", type=int, default=1000)
-    # parser.add_argument("--epochs", type=int, default=1400)
-    # parser.add_argument("--global-batch-size", type=int, default=256)
-    # parser.add_argument("--global-seed", type=int, default=0)
-    # parser.add_argument("--vae", type=str, choices=["ema", "mse"], default="mse")  # Choice doesn't affect training
-    # parser.add_argument("--num-workers", type=int, default=0)
-    # parser.add_argument("--log-every", type=int, default=100)
-    # parser.add_argument("--ckpt-every", type=int, default=50_000)
-    # args = parser.parse_args()
-    # main(args)
-    device = torch.device('cuda:0')
-    x = torch.from_numpy(np.load('1.npy')).to(device)
-    vae = AutoencoderKL.from_pretrained(f"stabilityai/sd-vae-ft-mse", cache_dir='cache_dir').to(device)
-    vae.eval()
-    b, f, c, h, w = x.shape
-
-    n_sample = 1
-    n_round = b*f // n_sample
-    samples = rearrange(x, 'b f c h w -> (b f) c h w')
-    samples_dec = []
-
-    for i in tqdm(range(n_round+1)):
-        with torch.no_grad():
-            samples_ = vae.decode(samples[i*n_sample:(i+1)*n_sample] / 0.18215).sample
-        samples_dec.append(samples_.cpu())
-    samples = torch.cat(samples_dec, dim=0)
-
-    samples = rearrange(samples, '(b f) c h w -> b f c h w', b=b)
-
-    video_ = ((samples[0] * 0.5 + 0.5) * 255).add_(0.5).clamp_(0, 255).to(dtype=torch.uint8).permute(0, 2, 3, 1).contiguous()
-
-    imageio.mimwrite('1.mp4', video_, fps=10, quality=9)
-
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--data-path", type=str)
+    parser.add_argument("--ae", type=str, default="stabilityai/sd-vae-ft-mse")
+    parser.add_argument("--features-path", type=str, default="features")
+    parser.add_argument("--features-name", type=str, default="movie")
+    parser.add_argument("--global-seed", type=int, default=0)
+    parser.add_argument("--num-workers", type=int, default=0)
+    parser.add_argument("--sample-rate", type=int, default=3)
+    args = parser.parse_args()
+    main(args)
 

opensora/utils/dataset_utils.py

@@ -1,8 +1,34 @@
 import math
+
+import decord
 from torch.nn import functional as F
 import torch
 
 
+class DecordInit(object):
+    """Using Decord(https://github.com/dmlc/decord) to initialize the video_reader."""
+
+    def __init__(self, num_threads=1):
+        self.num_threads = num_threads
+        self.ctx = decord.cpu(0)
+
+    def __call__(self, filename):
+        """Perform the Decord initialization.
+        Args:
+            results (dict): The resulting dict to be modified and passed
+                to the next transform in pipeline.
+        """
+        reader = decord.VideoReader(filename,
+                                    ctx=self.ctx,
+                                    num_threads=self.num_threads)
+        return reader
+
+    def __repr__(self):
+        repr_str = (f'{self.__class__.__name__}('
+                    f'sr={self.sr},'
+                    f'num_threads={self.num_threads})')
+        return repr_str
+
 def pad_to_multiple(number, ds_stride):
     remainder = number % ds_stride
     if remainder == 0:

opensora/utils/utils.py

@@ -0,0 +1,19 @@
+import os
+
+import torch
+
+
+def find_model(model_name):
+    """
+    Finds a pre-trained Latte model, downloading it if necessary. Alternatively, loads a model from a local path.
+    """
+    assert os.path.isfile(model_name), f'Could not find Latte checkpoint at {model_name}'
+    checkpoint = torch.load(model_name, map_location=lambda storage, loc: storage)
+
+    # if "ema" in checkpoint:  # supports checkpoints from train.py
+    #     print('Using Ema!')
+    #     checkpoint = checkpoint["ema"]
+    # else:
+    print('Using model!')
+    checkpoint = checkpoint['model']
+    return checkpoint

scripts/sample.sh

@@ -0,0 +1,4 @@
+python opensora/sample/sample.py \
+  --model Latte-XL/122 --ae stabilityai/sd-vae-ft-mse \
+  --ckpt results/047-Latte-XL-122/checkpoints/0005000.pt --extras 1 \
+  --fps 10 --num-frames 16 --image-size 256

scripts/train.sh

@@ -1,7 +1,7 @@
-torchrun  --nproc_per_node=1 --master_port=29501 opensora/train/train.py \
-  --model DiT-XL/122 --dataset ucf101\
-  --ae ucf101_stride4x4x4 \
-  --data-path /remote-home/yeyang/UCF-101 --num-classes 101 \
-  --sample-rate 2 --num-frames 8 --max-image-size 128 --clip-grad-norm 1 \
-  --epochs 14000 --global-batch-size 2 --lr 1e-4 \
+torchrun  --nproc_per_node=8 --master_port=29501 opensora/train/train.py \
+  --model Latte-XL/122 --dataset sky \
+  --ae stabilityai/sd-vae-ft-mse \
+  --data-path /remote-home/yeyang/sky_timelapse/sky_train/ --extras 1 \
+  --sample-rate 3 --num-frames 16 --max-image-size 256 \
+  --epochs 14000 --global-batch-size 40 --lr 1e-4 \
   --ckpt-every 1000 --log-every 50