betterments.

training/mochi-1/dataset.py

@@ -1,243 +0,0 @@
-from pathlib import Path
-from typing import Any, Dict, Tuple
-
-import numpy as np
-import torch
-import torchvision.transforms as TT
-from accelerate.logging import get_logger
-from torchvision import transforms
-from torchvision.transforms import InterpolationMode
-from torchvision.transforms.functional import resize
-
-
-# Must import after torch because this can sometimes lead to a nasty segmentation fault, or stack smashing error
-# Very few bug reports but it happens. Look in decord Github issues for more relevant information.
-import decord  # isort:skip
-
-decord.bridge.set_bridge("torch")
-
-import sys 
-sys.path.append("..")
-
-from dataset import VideoDataset as VDS
-from dataset import BucketSampler
-
-logger = get_logger(__name__)
-
-# TODO (sayakpaul): probably not all buckets are needed for Mochi-1? 
-HEIGHT_BUCKETS = [256, 320, 384, 480, 512, 576, 720, 768, 960, 1024, 1280, 1536]
-WIDTH_BUCKETS = [256, 320, 384, 480, 512, 576, 720, 768, 848, 960, 1024, 1280, 1536]
-FRAME_BUCKETS = [16, 24, 32, 48, 64, 80, 84]
-
-VAE_SPATIAL_SCALE_FACTOR = 8
-VAE_TEMPORAL_SCALE_FACTOR = 6
-
-class VideoDataset(VDS):
-    def __init__(self, *args, **kwargs) -> None:
-        super().__init__(*args, **kwargs)
-
-    # Overriding this because we calculate `num_frames` differently.
-    def __getitem__(self, index: int) -> Dict[str, Any]:
-        if isinstance(index, list):
-            # Here, index is actually a list of data objects that we need to return.
-            # The BucketSampler should ideally return indices. But, in the sampler, we'd like
-            # to have information about num_frames, height and width. Since this is not stored
-            # as metadata, we need to read the video to get this information. You could read this
-            # information without loading the full video in memory, but we do it anyway. In order
-            # to not load the video twice (once to get the metadata, and once to return the loaded video
-            # based on sampled indices), we cache it in the BucketSampler. When the sampler is
-            # to yield, we yield the cache data instead of indices. So, this special check ensures
-            # that data is not loaded a second time. PRs are welcome for improvements.
-            return index
-
-        if self.load_tensors:
-            image_latents, video_latents, prompt_embeds, prompt_attention_mask = self._preprocess_video(self.video_paths[index])
-
-            # This is hardcoded for now.
-            # Output of the VAE encoding is 2 * output_channels and then it's
-            # temporal compression factor is 6. Initially, the VAE encodings will have
-            # 24 latent number of frames. So, if we were to train with a 
-            # max frame size of 84 and frame bucket of [84], we need to have the following logic.
-            latent_num_frames = video_latents.size(0)
-            num_frames = (latent_num_frames // 2) * (VAE_TEMPORAL_SCALE_FACTOR + 1)
-
-            height = video_latents.size(2) * VAE_SPATIAL_SCALE_FACTOR
-            width = video_latents.size(3) * VAE_SPATIAL_SCALE_FACTOR
-
-            return {
-                "prompt": prompt_embeds,
-                "prompt_attention_mask": prompt_attention_mask,
-                "image": image_latents,
-                "video": video_latents,
-                "video_metadata": {
-                    "num_frames": num_frames,
-                    "height": height,
-                    "width": width,
-                },
-            }
-        else:
-            image, video, _ = self._preprocess_video(self.video_paths[index])
-            if video is not None:
-                return {
-                    "prompt": self.id_token + self.prompts[index],
-                    "image": image,
-                    "video": video,
-                    "video_metadata": {
-                        "num_frames": video.shape[0],
-                        "height": video.shape[2],
-                        "width": video.shape[3],
-                    },
-                }
-
-    # Overriding this because we need `prompt_attention_mask`.
-    def _load_preprocessed_latents_and_embeds(self, path: Path) -> Tuple[torch.Tensor, torch.Tensor]:
-        filename_without_ext = path.name.split(".")[0]
-        pt_filename = f"{filename_without_ext}.pt"
-
-        # The current path is something like: /a/b/c/d/videos/00001.mp4
-        # We need to reach: /a/b/c/d/video_latents/00001.pt
-        image_latents_path = path.parent.parent.joinpath("image_latents")
-        video_latents_path = path.parent.parent.joinpath("video_latents")
-        embeds_path = path.parent.parent.joinpath("prompt_embeds")
-        attention_mask_path = path.parent.parent.joinpath("prompt_attention_mask")
-
-        if (
-            not video_latents_path.exists()
-            or not embeds_path.exists()
-            or not attention_mask_path.exists()
-            or (self.image_to_video and not image_latents_path.exists())
-        ):
-            raise ValueError(
-                f"When setting the load_tensors parameter to `True`, it is expected that the `{self.data_root=}` contains three folders named `video_latents`, `prompt_embeds`, and `prompt_attention_mask`. However, these folders were not found. Please make sure to have prepared your data correctly using `prepare_data.py`. Additionally, if you're training image-to-video, it is expected that an `image_latents` folder is also present."
-            )
-
-        if self.image_to_video:
-            image_latent_filepath = image_latents_path.joinpath(pt_filename)
-        video_latent_filepath = video_latents_path.joinpath(pt_filename)
-        embeds_filepath = embeds_path.joinpath(pt_filename)
-        attention_mask_filepath = attention_mask_path.joinpath(pt_filename)
-
-        if not video_latent_filepath.is_file() or not embeds_filepath.is_file() or not attention_mask_filepath.is_file():
-            if self.image_to_video:
-                image_latent_filepath = image_latent_filepath.as_posix()
-            video_latent_filepath = video_latent_filepath.as_posix()
-            embeds_filepath = embeds_filepath.as_posix()
-            attention_mask_filepath = attention_mask_filepath.as_posix()
-            raise ValueError(
-                f"The file {video_latent_filepath=} or {embeds_filepath=} or {attention_mask_filepath=} could not be found. Please ensure that you've correctly executed `prepare_dataset.py`."
-            )
-
-        images = (
-            torch.load(image_latent_filepath, map_location="cpu", weights_only=True) if self.image_to_video else None
-        )
-        latents = torch.load(video_latent_filepath, map_location="cpu", weights_only=True)
-        embeds = torch.load(embeds_filepath, map_location="cpu", weights_only=True)
-        attention_masks = torch.load(attention_mask_filepath,  map_location="cpu", weights_only=True)
-
-        return images, latents, embeds, attention_masks
-
-
-# We need the `VideoDatasetWithResizing` and `VideoDatasetWithResizeAndRectangleCrop` classes to subclass from
-# the new `VideoDataset` class defined in this file. And also because of the changes in 
-# `_preprocess_video()` (how we handle `nearest_frame_bucket`). 
-
-class VideoDatasetWithResizing(VideoDataset):
-    def __init__(self, *args, **kwargs) -> None:
-        super().__init__(*args, **kwargs)
-
-    def _preprocess_video(self, path: Path) -> torch.Tensor:
-        if self.load_tensors:
-            return self._load_preprocessed_latents_and_embeds(path)
-        else:
-            video_reader = decord.VideoReader(uri=path.as_posix())
-            video_num_frames = len(video_reader)
-
-            nearest_frame_bucket = min(
-                [bucket for bucket in self.frame_buckets if bucket <= video_num_frames],
-                key=lambda x: abs(x - min(video_num_frames, self.max_num_frames)),
-                default=1,
-            )
-            
-            frame_indices = list(range(0, video_num_frames, video_num_frames // nearest_frame_bucket))
-            frames = video_reader.get_batch(frame_indices)
-            frames = frames[:nearest_frame_bucket].float()
-            frames = frames.permute(0, 3, 1, 2).contiguous()
-
-            nearest_res = self._find_nearest_resolution(frames.shape[2], frames.shape[3])
-            frames_resized = torch.stack([resize(frame, nearest_res) for frame in frames], dim=0)
-            frames = torch.stack([self.video_transforms(frame) for frame in frames_resized], dim=0)
-
-            image = frames[:1].clone() if self.image_to_video else None
-
-            return image, frames, None
-
-    def _find_nearest_resolution(self, height, width):
-        nearest_res = min(self.resolutions, key=lambda x: abs(x[1] - height) + abs(x[2] - width))
-        return nearest_res[1], nearest_res[2]
-
-
-class VideoDatasetWithResizeAndRectangleCrop(VideoDataset):
-    def __init__(self, video_reshape_mode: str = "center", *args, **kwargs) -> None:
-        super().__init__(*args, **kwargs)
-        self.video_reshape_mode = video_reshape_mode
-
-    def _resize_for_rectangle_crop(self, arr, image_size):
-        reshape_mode = self.video_reshape_mode
-        if arr.shape[3] / arr.shape[2] > image_size[1] / image_size[0]:
-            arr = resize(
-                arr,
-                size=[image_size[0], int(arr.shape[3] * image_size[0] / arr.shape[2])],
-                interpolation=InterpolationMode.BICUBIC,
-            )
-        else:
-            arr = resize(
-                arr,
-                size=[int(arr.shape[2] * image_size[1] / arr.shape[3]), image_size[1]],
-                interpolation=InterpolationMode.BICUBIC,
-            )
-
-        h, w = arr.shape[2], arr.shape[3]
-        arr = arr.squeeze(0)
-
-        delta_h = h - image_size[0]
-        delta_w = w - image_size[1]
-
-        if reshape_mode == "random" or reshape_mode == "none":
-            top = np.random.randint(0, delta_h + 1)
-            left = np.random.randint(0, delta_w + 1)
-        elif reshape_mode == "center":
-            top, left = delta_h // 2, delta_w // 2
-        else:
-            raise NotImplementedError
-        arr = TT.functional.crop(arr, top=top, left=left, height=image_size[0], width=image_size[1])
-        return arr
-
-    def _preprocess_video(self, path: Path) -> torch.Tensor:
-        if self.load_tensors:
-            return self._load_preprocessed_latents_and_embeds(path)
-        else:
-            video_reader = decord.VideoReader(uri=path.as_posix())
-            video_num_frames = len(video_reader)
-            nearest_frame_bucket = min(
-                [bucket for bucket in self.frame_buckets if bucket <= video_num_frames],
-                key=lambda x: abs(x - min(video_num_frames, self.max_num_frames)),
-                default=1,
-            )
-
-            frame_indices = list(range(0, video_num_frames, video_num_frames // nearest_frame_bucket))
-
-            frames = video_reader.get_batch(frame_indices)
-            frames = frames[:nearest_frame_bucket].float()
-            frames = frames.permute(0, 3, 1, 2).contiguous()
-
-            nearest_res = self._find_nearest_resolution(frames.shape[2], frames.shape[3])
-            frames_resized = self._resize_for_rectangle_crop(frames, nearest_res)
-            frames = torch.stack([self.video_transforms(frame) for frame in frames_resized], dim=0)
-
-            image = frames[:1].clone() if self.image_to_video else None
-
-            return image, frames, None
-
-    def _find_nearest_resolution(self, height, width):
-        nearest_res = min(self.resolutions, key=lambda x: abs(x[1] - height) + abs(x[2] - width))
-        return nearest_res[1], nearest_res[2]

training/mochi-1/prepare_dataset.py

@@ -21,17 +21,19 @@ from torch.utils.data import DataLoader
 from torchvision import transforms
 from tqdm import tqdm
 from transformers import T5EncoderModel, T5Tokenizer
+from dataset_mochi import VideoDatasetWithResizing, VideoDatasetWithResizeAndRectangleCrop
 
 
 import decord  # isort:skip
 
-import sys
-sys.path.append(".")
-from dataset import BucketSampler, VideoDatasetWithResizing, VideoDatasetWithResizeAndRectangleCrop  # isort:skip
-
-
 decord.bridge.set_bridge("torch")
 
+import sys
+sys.path.append("..")
+
+from dataset import BucketSampler
+
+
 logger = get_logger(__name__)
 
 DTYPE_MAPPING = {

training/mochi-1/text_to_video_lora.py

@@ -54,11 +54,12 @@ from tqdm.auto import tqdm
 from transformers import AutoTokenizer, T5EncoderModel
 
 from args import get_args  # isort:skip
+from dataset_mochi import VideoDatasetWithResizing, VideoDatasetWithResizeAndRectangleCrop  # isort:skip
 
 import sys
-sys.path.append(".")
+sys.path.append("..")
 
-from dataset import BucketSampler, VideoDatasetWithResizing, VideoDatasetWithResizeAndRectangleCrop  # isort:skip
+from dataset import BucketSampler # isort:skip
 from text_encoder import compute_prompt_embeddings  # isort:skip
 from utils import get_gradient_norm, get_optimizer, print_memory, reset_memory  # isort:skip