add VideoDatasetWithResizeAndRectangleCrop dataset resize crop (#13)

* fix issue 12:  train dataset _preprocess_data crop video

* add_argument  video_reshape_mode input videos are reshaped to this mode

* arg enable_model_cpu_offloading

* Add imageio-ffmpeg and imageio to requirements.txt

* apply changes to i2v lora script

* make style

---------

Co-authored-by: --unset <--unset>
Co-authored-by: Aryan <[email protected]>

Author: glide-the

requirements.txt

@@ -11,3 +11,7 @@ pandas
 torch
 torchvision
 torchao
+sentencepiece
+imageio-ffmpeg
+imageio
+numpy==2.1.1

training/args.py

@@ -192,6 +192,12 @@ def _get_training_args(parser: argparse.ArgumentParser) -> None:
         default=720,
         help="All input videos are resized to this width.",
     )
+    parser.add_argument(
+        "--video_reshape_mode",
+        type=str,
+        default=None,
+        help="All input videos are reshaped to this mode. Choose between ['center', 'random', 'none']",
+    )
     parser.add_argument("--fps", type=int, default=8, help="All input videos will be used at this FPS.")
     parser.add_argument(
         "--max_num_frames",

training/cogvideox_image_to_video_lora.py

@@ -55,7 +55,7 @@
 
 
 from args import get_args  # isort:skip
-from dataset import BucketSampler, VideoDatasetWithResizing  # isort:skip
+from dataset import BucketSampler, VideoDatasetWithResizing, VideoDatasetWithResizeAndRectangleCrop  # isort:skip
 from text_encoder import compute_prompt_embeddings  # isort:skip
 from utils import get_gradient_norm, get_optimizer, prepare_rotary_positional_embeddings, print_memory, reset_memory  # isort:skip
 
@@ -465,20 +465,37 @@ def load_model_hook(models, input_dir):
     )
 
     # Dataset and DataLoader
-    train_dataset = VideoDatasetWithResizing(
-        data_root=args.data_root,
-        dataset_file=args.dataset_file,
-        caption_column=args.caption_column,
-        video_column=args.video_column,
-        max_num_frames=args.max_num_frames,
-        id_token=args.id_token,
-        height_buckets=args.height_buckets,
-        width_buckets=args.width_buckets,
-        frame_buckets=args.frame_buckets,
-        load_tensors=args.load_tensors,
-        random_flip=args.random_flip,
-        image_to_video=True,
-    )
+    if not args.video_reshape_mode:
+        train_dataset = VideoDatasetWithResizing(
+            data_root=args.data_root,
+            dataset_file=args.dataset_file,
+            caption_column=args.caption_column,
+            video_column=args.video_column,
+            max_num_frames=args.max_num_frames,
+            id_token=args.id_token,
+            height_buckets=args.height_buckets,
+            width_buckets=args.width_buckets,
+            frame_buckets=args.frame_buckets,
+            load_tensors=args.load_tensors,
+            random_flip=args.random_flip,
+            image_to_video=True,
+        )
+    else:
+        train_dataset = VideoDatasetWithResizeAndRectangleCrop(
+            video_reshape_mode=args.video_reshape_mode,
+            data_root=args.data_root,
+            dataset_file=args.dataset_file,
+            caption_column=args.caption_column,
+            video_column=args.video_column,
+            max_num_frames=args.max_num_frames,
+            id_token=args.id_token,
+            height_buckets=args.height_buckets,
+            width_buckets=args.width_buckets,
+            frame_buckets=args.frame_buckets,
+            load_tensors=args.load_tensors,
+            random_flip=args.random_flip,
+            image_to_video=True,
+        )
 
     def collate_fn(data):
         prompts = [x["prompt"] for x in data[0]]

training/cogvideox_text_to_video_lora.py

@@ -54,7 +54,7 @@
 
 
 from args import get_args  # isort:skip
-from dataset import BucketSampler, VideoDatasetWithResizing  # isort:skip
+from dataset import BucketSampler, VideoDatasetWithResizing, VideoDatasetWithResizeAndRectangleCrop  # isort:skip
 from text_encoder import compute_prompt_embeddings  # isort:skip
 from utils import get_gradient_norm, get_optimizer, prepare_rotary_positional_embeddings, print_memory, reset_memory  # isort:skip
 
@@ -462,19 +462,35 @@ def load_model_hook(models, input_dir):
     )
 
     # Dataset and DataLoader
-    train_dataset = VideoDatasetWithResizing(
-        data_root=args.data_root,
-        dataset_file=args.dataset_file,
-        caption_column=args.caption_column,
-        video_column=args.video_column,
-        max_num_frames=args.max_num_frames,
-        id_token=args.id_token,
-        height_buckets=args.height_buckets,
-        width_buckets=args.width_buckets,
-        frame_buckets=args.frame_buckets,
-        load_tensors=args.load_tensors,
-        random_flip=args.random_flip,
-    )
+    if not args.video_reshape_mode:
+        train_dataset = VideoDatasetWithResizing(
+            data_root=args.data_root,
+            dataset_file=args.dataset_file,
+            caption_column=args.caption_column,
+            video_column=args.video_column,
+            max_num_frames=args.max_num_frames,
+            id_token=args.id_token,
+            height_buckets=args.height_buckets,
+            width_buckets=args.width_buckets,
+            frame_buckets=args.frame_buckets,
+            load_tensors=args.load_tensors,
+            random_flip=args.random_flip,
+        )
+    else:
+        train_dataset = VideoDatasetWithResizeAndRectangleCrop(
+            video_reshape_mode=args.video_reshape_mode,
+            data_root=args.data_root,
+            dataset_file=args.dataset_file,
+            caption_column=args.caption_column,
+            video_column=args.video_column,
+            max_num_frames=args.max_num_frames,
+            id_token=args.id_token,
+            height_buckets=args.height_buckets,
+            width_buckets=args.width_buckets,
+            frame_buckets=args.frame_buckets,
+            load_tensors=args.load_tensors,
+            random_flip=args.random_flip,
+        )
 
     def collate_fn(data):
         prompts = [x["prompt"] for x in data[0]]

training/cogvideox_text_to_video_sft.py

@@ -53,7 +53,7 @@
 
 
 from args import get_args  # isort:skip
-from dataset import BucketSampler, VideoDatasetWithResizing  # isort:skip
+from dataset import BucketSampler, VideoDatasetWithResizing, VideoDatasetWithResizeAndRectangleCrop  # isort:skip
 from text_encoder import compute_prompt_embeddings  # isort:skip
 from utils import get_gradient_norm, get_optimizer, prepare_rotary_positional_embeddings, print_memory, reset_memory  # isort:skip
 
@@ -426,19 +426,35 @@ def load_model_hook(models, input_dir):
     )
 
     # Dataset and DataLoader
-    train_dataset = VideoDatasetWithResizing(
-        data_root=args.data_root,
-        dataset_file=args.dataset_file,
-        caption_column=args.caption_column,
-        video_column=args.video_column,
-        max_num_frames=args.max_num_frames,
-        id_token=args.id_token,
-        height_buckets=args.height_buckets,
-        width_buckets=args.width_buckets,
-        frame_buckets=args.frame_buckets,
-        load_tensors=args.load_tensors,
-        random_flip=args.random_flip,
-    )
+    if not args.video_reshape_mode:
+        train_dataset = VideoDatasetWithResizing(
+            data_root=args.data_root,
+            dataset_file=args.dataset_file,
+            caption_column=args.caption_column,
+            video_column=args.video_column,
+            max_num_frames=args.max_num_frames,
+            id_token=args.id_token,
+            height_buckets=args.height_buckets,
+            width_buckets=args.width_buckets,
+            frame_buckets=args.frame_buckets,
+            load_tensors=args.load_tensors,
+            random_flip=args.random_flip,
+        )
+    else:
+        train_dataset = VideoDatasetWithResizeAndRectangleCrop(
+            video_reshape_mode=args.video_reshape_mode,
+            data_root=args.data_root,
+            dataset_file=args.dataset_file,
+            caption_column=args.caption_column,
+            video_column=args.video_column,
+            max_num_frames=args.max_num_frames,
+            id_token=args.id_token,
+            height_buckets=args.height_buckets,
+            width_buckets=args.width_buckets,
+            frame_buckets=args.frame_buckets,
+            load_tensors=args.load_tensors,
+            random_flip=args.random_flip,
+        )
 
     def collate_fn(data):
         prompts = [x["prompt"] for x in data[0]]

training/dataset.py

@@ -2,11 +2,14 @@
 from pathlib import Path
 from typing import Any, Dict, List, Optional, Tuple
 
+import numpy as np
 import pandas as pd
 import torch
+import torchvision.transforms as TT
 from accelerate.logging import get_logger
 from torch.utils.data import Dataset, Sampler
 from torchvision import transforms
+from torchvision.transforms import InterpolationMode
 from torchvision.transforms.functional import resize
 
 
@@ -281,6 +284,71 @@ def _find_nearest_resolution(self, height, 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(
+                self.frame_buckets, key=lambda x: abs(x - min(video_num_frames, self.max_num_frames))
+            )
+
+            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]
+
+
 class BucketSampler(Sampler):
     def __init__(self, data_source: VideoDataset, batch_size: int = 8, shuffle: bool = True) -> None:
         self.data_source = data_source

training/prepare_dataset.py

@@ -7,12 +7,16 @@
 import traceback
 from typing import Any, Dict, List, Optional, Tuple, Union
 
+import numpy as np
 import pandas as pd
 import torch
 import torch.distributed as dist
+import torchvision.transforms as TT
 from diffusers import AutoencoderKLCogVideoX
 from diffusers.utils import export_to_video, get_logger
 from torchvision import transforms
+from torchvision.transforms import InterpolationMode
+from torchvision.transforms.functional import resize
 from tqdm import tqdm
 from transformers import T5EncoderModel, T5Tokenizer
 
@@ -153,13 +157,51 @@ def load_dataset_from_csv(
     return prompts, video_paths
 
 
+def resize_for_rectangle_crop(arr, height, width, reshape_mode):
+    image_size = height, width
+    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 load_and_preprocess_video(
-    path: pathlib.Path, height: int, width: int, max_num_frames: int, video_transforms, num_threads: int = 0
+    path: pathlib.Path,
+    height: int,
+    width: int,
+    max_num_frames: int,
+    video_transforms,
+    num_threads: int = 0,
+    video_reshape_mode: str = "center",
 ) -> Optional[torch.Tensor]:
     frames = None
 
     try:
-        video_reader = decord.VideoReader(uri=path.as_posix(), height=height, width=width, num_threads=num_threads)
+        video_reader = decord.VideoReader(uri=path.as_posix(), num_threads=num_threads)
         video_num_frames = len(video_reader)
 
         if video_num_frames < max_num_frames:
@@ -172,6 +214,7 @@ def load_and_preprocess_video(
         frames: torch.Tensor = video_reader.get_batch(indices)
         frames = frames[:max_num_frames].float()
         frames = frames.permute(0, 3, 1, 2).contiguous()
+        frames = resize_for_rectangle_crop(frames, height, width, video_reshape_mode)
         frames = torch.stack([video_transforms(frame) for frame in frames], dim=0)
     except Exception as e:
         logger.error(f"Error: {e}. Skipping video located at `{path.as_posix()}`")