From 00f7519f0e1aea4d7718a3b3feeeba6cbd779b59 Mon Sep 17 00:00:00 2001 From: glide-the Date: Thu, 10 Oct 2024 04:25:12 +0800 Subject: [PATCH] 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 --- requirements.txt | 4 ++ training/args.py | 6 ++ training/cogvideox_image_to_video_lora.py | 47 +++++++++++----- training/cogvideox_text_to_video_lora.py | 44 ++++++++++----- training/cogvideox_text_to_video_sft.py | 44 ++++++++++----- training/dataset.py | 68 +++++++++++++++++++++++ training/prepare_dataset.py | 47 +++++++++++++++- 7 files changed, 215 insertions(+), 45 deletions(-) diff --git a/requirements.txt b/requirements.txt index e5ddcf3..b15d79b 100644 --- a/requirements.txt +++ b/requirements.txt @@ -11,3 +11,7 @@ pandas torch torchvision torchao +sentencepiece +imageio-ffmpeg +imageio +numpy==2.1.1 \ No newline at end of file diff --git a/training/args.py b/training/args.py index 0bef33f..c20a471 100644 --- a/training/args.py +++ b/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", diff --git a/training/cogvideox_image_to_video_lora.py b/training/cogvideox_image_to_video_lora.py index 9c0b8da..8055687 100644 --- a/training/cogvideox_image_to_video_lora.py +++ b/training/cogvideox_image_to_video_lora.py @@ -55,7 +55,7 @@ from transformers import AutoTokenizer, T5EncoderModel 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 main(args): ) # 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]] diff --git a/training/cogvideox_text_to_video_lora.py b/training/cogvideox_text_to_video_lora.py index c0e739e..199f525 100644 --- a/training/cogvideox_text_to_video_lora.py +++ b/training/cogvideox_text_to_video_lora.py @@ -54,7 +54,7 @@ from transformers import AutoTokenizer, T5EncoderModel 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 main(args): ) # 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]] diff --git a/training/cogvideox_text_to_video_sft.py b/training/cogvideox_text_to_video_sft.py index eddbf95..0dcddce 100644 --- a/training/cogvideox_text_to_video_sft.py +++ b/training/cogvideox_text_to_video_sft.py @@ -53,7 +53,7 @@ from transformers import AutoTokenizer, T5EncoderModel 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 main(args): ) # 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]] diff --git a/training/dataset.py b/training/dataset.py index 4545073..2a6ac6c 100644 --- a/training/dataset.py +++ b/training/dataset.py @@ -2,11 +2,14 @@ import random 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 @@ class VideoDatasetWithResizing(VideoDataset): 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 diff --git a/training/prepare_dataset.py b/training/prepare_dataset.py index 041fc69..9f1ed6d 100644 --- a/training/prepare_dataset.py +++ b/training/prepare_dataset.py @@ -7,12 +7,16 @@ import pathlib 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()}`")