[Preprocess][Feat] support torchvision to load video in new preprocessing (#761)

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 \

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(

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)

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)
 

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}")