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 <aryan@huggingface.co>
This commit is contained in:
glide-the
2024-10-10 04:25:12 +08:00
committed by GitHub
parent b1b72c0e38
commit 00f7519f0e
7 changed files with 215 additions and 45 deletions
+32 -15
View File
@@ -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]]