diff --git a/training/mochi-1/dataset.py b/training/mochi-1/dataset.py deleted file mode 100644 index fa109c4..0000000 --- a/training/mochi-1/dataset.py +++ /dev/null @@ -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] \ No newline at end of file diff --git a/training/mochi-1/prepare_dataset.py b/training/mochi-1/prepare_dataset.py index d9da99a..0826c27 100644 --- a/training/mochi-1/prepare_dataset.py +++ b/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 = { diff --git a/training/mochi-1/text_to_video_lora.py b/training/mochi-1/text_to_video_lora.py index d0d1736..2b9ee42 100644 --- a/training/mochi-1/text_to_video_lora.py +++ b/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