From cb16cbae5fffe1f966255b0830742d448c7c07c8 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Tue, 19 Nov 2024 11:35:00 +0530 Subject: [PATCH] dataset_mochi --- dataset_mochi.py | 242 +++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 242 insertions(+) create mode 100644 dataset_mochi.py diff --git a/dataset_mochi.py b/dataset_mochi.py new file mode 100644 index 0000000..c40a747 --- /dev/null +++ b/dataset_mochi.py @@ -0,0 +1,242 @@ +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 + +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