mirror of
https://github.com/storytold/FineTrainers-Conditioning.git
synced 2026-10-25 08:29:35 +00:00
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:
@@ -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]]
|
||||
|
||||
Reference in New Issue
Block a user