From 4aa307be55ea1893e1d6377fed6b2c991eb6a537 Mon Sep 17 00:00:00 2001 From: Jinzhe Pan <48981407+Eigensystem@users.noreply.github.com> Date: Tue, 2 Sep 2025 14:37:01 +0800 Subject: [PATCH] [Preprocess][Feat] support torchvision to load video in new preprocessing (#761) --- .../crush_smol/preprocess_wan_data_t2v_new.sh | 1 + fastvideo/configs/configs.py | 34 +++++++++++++++++++ fastvideo/pipelines/pipeline_batch_info.py | 2 +- .../pipelines/preprocess/preprocess_stages.py | 13 ++++++- fastvideo/workflow/preprocess/components.py | 6 ++-- 5 files changed, 52 insertions(+), 4 deletions(-) diff --git a/examples/training/finetune/wan_t2v_1.3B/crush_smol/preprocess_wan_data_t2v_new.sh b/examples/training/finetune/wan_t2v_1.3B/crush_smol/preprocess_wan_data_t2v_new.sh index 262f39ec9..8b850f78a 100644 --- a/examples/training/finetune/wan_t2v_1.3B/crush_smol/preprocess_wan_data_t2v_new.sh +++ b/examples/training/finetune/wan_t2v_1.3B/crush_smol/preprocess_wan_data_t2v_new.sh @@ -10,6 +10,7 @@ torchrun --nproc_per_node=$GPU_NUM \ --model_path $MODEL_PATH \ --mode preprocess \ --workload_type t2v \ + --preprocess.video_loader_type torchvision \ --preprocess.dataset_type merged \ --preprocess.dataset_path $DATASET_PATH \ --preprocess.dataset_output_dir $OUTPUT_DIR \ diff --git a/fastvideo/configs/configs.py b/fastvideo/configs/configs.py index 07a7630b2..139300a14 100644 --- a/fastvideo/configs/configs.py +++ b/fastvideo/configs/configs.py @@ -32,6 +32,29 @@ class DatasetType(str, Enum): return [dataset_type.value for dataset_type in cls] +class VideoLoaderType(str, Enum): + """ + Enumeration for different video loaders. + """ + TORCHCODEC = "torchcodec" + TORCHVISION = "torchvision" + + @classmethod + def from_string(cls, value: str) -> "VideoLoaderType": + """Convert string to VideoLoader enum.""" + try: + return cls(value.lower()) + except ValueError: + raise ValueError( + f"Invalid video loader: {value}. Must be one of: {', '.join([m.value for m in cls])}" + ) from None + + @classmethod + def choices(cls) -> list[str]: + """Get all available choices as strings for argparse.""" + return [video_loader.value for video_loader in cls] + + @dataclasses.dataclass class PreprocessConfig: """Configuration for preprocessing operations.""" @@ -51,6 +74,7 @@ class PreprocessConfig: flush_frequency: int = 256 # Video processing parameters + video_loader_type: VideoLoaderType = VideoLoaderType.TORCHCODEC max_height: int = 480 max_width: int = 848 num_frames: int = 163 @@ -120,6 +144,12 @@ class PreprocessConfig: help="How often to save to parquet files") # Video processing parameters + preprocess_args.add_argument( + f"--{prefix_with_dot}video-loader-type", + type=str, + choices=VideoLoaderType.choices(), + default=PreprocessConfig.video_loader_type.value, + help="Type of the video loader") preprocess_args.add_argument(f"--{prefix_with_dot}max-height", type=int, default=PreprocessConfig.max_height, @@ -174,6 +204,10 @@ class PreprocessConfig: if 'dataset_type' in kwargs and isinstance(kwargs['dataset_type'], str): kwargs['dataset_type'] = DatasetType.from_string( kwargs['dataset_type']) + if 'video_loader_type' in kwargs and isinstance( + kwargs['video_loader_type'], str): + kwargs['video_loader_type'] = VideoLoaderType.from_string( + kwargs['video_loader_type']) preprocess_config = cls() if not update_config_from_args( diff --git a/fastvideo/pipelines/pipeline_batch_info.py b/fastvideo/pipelines/pipeline_batch_info.py index 7fc9fa56c..453387033 100644 --- a/fastvideo/pipelines/pipeline_batch_info.py +++ b/fastvideo/pipelines/pipeline_batch_info.py @@ -241,5 +241,5 @@ class TrainingBatch: @dataclass class PreprocessBatch(ForwardBatch): - video_loader: list["VideoDecoder"] = field(default_factory=list) + video_loader: list["VideoDecoder"] | list[str] = field(default_factory=list) video_file_name: list[str] = field(default_factory=list) diff --git a/fastvideo/pipelines/preprocess/preprocess_stages.py b/fastvideo/pipelines/preprocess/preprocess_stages.py index b6ed30a77..7d82eaea8 100644 --- a/fastvideo/pipelines/preprocess/preprocess_stages.py +++ b/fastvideo/pipelines/preprocess/preprocess_stages.py @@ -4,9 +4,11 @@ from typing import cast import numpy as np import torch +import torchvision from einops import rearrange from torchvision import transforms +from fastvideo.configs.configs import VideoLoaderType from fastvideo.dataset.transform import (CenterCropResizeVideo, TemporalRandomCrop) from fastvideo.fastvideo_args import FastVideoArgs, WorkloadType @@ -61,7 +63,16 @@ class VideoTransformStage(PipelineStage): else: frame_indices = frame_indices[:self.num_frames] - video = batch.video_loader[i].get_frames_at(frame_indices).data + if fastvideo_args.preprocess_config.video_loader_type == VideoLoaderType.TORCHCODEC: + video = batch.video_loader[i].get_frames_at(frame_indices).data + elif fastvideo_args.preprocess_config.video_loader_type == VideoLoaderType.TORCHVISION: + video, _, _ = torchvision.io.read_video(batch.video_loader[i], + output_format="TCHW") + video = video[frame_indices] + else: + raise ValueError( + f"Invalid video loader type: {fastvideo_args.preprocess_config.video_loader_type}" + ) video = self.video_transform(video) video_pixel_batch.append(video) diff --git a/fastvideo/workflow/preprocess/components.py b/fastvideo/workflow/preprocess/components.py index aa9bdbbcd..1dc546613 100644 --- a/fastvideo/workflow/preprocess/components.py +++ b/fastvideo/workflow/preprocess/components.py @@ -13,7 +13,8 @@ import pyarrow.parquet as pq import torch from datasets import Dataset, Video, load_dataset -from fastvideo.configs.configs import DatasetType, PreprocessConfig +from fastvideo.configs.configs import (DatasetType, PreprocessConfig, + VideoLoaderType) from fastvideo.distributed.parallel_state import get_world_rank, get_world_size from fastvideo.logger import init_logger from fastvideo.pipelines.pipeline_batch_info import PreprocessBatch @@ -431,7 +432,8 @@ def build_dataset(preprocess_config: PreprocessConfig, split: str, return item dataset = dataset.map(add_video_column) - dataset = dataset.cast_column("video", Video()) + if preprocess_config.video_loader_type == VideoLoaderType.TORCHCODEC: + dataset = dataset.cast_column("video", Video()) else: raise ValueError( f"Invalid dataset type: {preprocess_config.dataset_type}")