feat: preprocess validation dataset only when exist (#734)

fastvideo/pipelines/pipeline_registry.py

@@ -25,6 +25,11 @@ _PIPELINE_NAME_TO_ARCHITECTURE_NAME: dict[str, str] = {
     "HunyuanVideoPipeline": "hunyuan",
 }
 
+_PREPROCESS_WORKLOAD_TYPE_TO_PIPELINE_NAME: dict[WorkloadType, str] = {
+    WorkloadType.I2V: "PreprocessPipelineI2V",
+    WorkloadType.T2V: "PreprocessPipelineT2V",
+}
+
 
 class PipelineType(str, Enum):
     """
@@ -68,12 +73,8 @@ class _PipelineRegistry:
     def _load_preprocess_pipeline_cls(
             self, workload_type: WorkloadType,
             arch: str) -> type[ComposedPipelineBase] | None:
-        if workload_type == WorkloadType.I2V:
-            pipeline_name = "I2VPreprocessPipeline"
-        elif workload_type == WorkloadType.T2V:
-            pipeline_name = "T2VPreprocessPipeline"
-        else:
-            raise ValueError(f"Invalid workload type: {workload_type.value}")
+        pipeline_name = _PREPROCESS_WORKLOAD_TYPE_TO_PIPELINE_NAME[
+            workload_type]
 
         return self.pipelines[
             PipelineType.PREPROCESS.value][arch][pipeline_name]

fastvideo/pipelines/preprocess/wan/wan_preprocess_pipelines.py

@@ -5,7 +5,7 @@ from fastvideo.pipelines.stages import (EncodingStage, ImageEncodingStage,
                                         TextEncodingStage)
 
 
-class I2VPreprocessPipeline(ComposedPipelineBase):
+class PreprocessPipelineI2V(ComposedPipelineBase):
     _required_config_modules = [
         "image_encoder", "image_processor", "text_encoder", "tokenizer", "vae"
     ]
@@ -25,7 +25,7 @@ class I2VPreprocessPipeline(ComposedPipelineBase):
                        stage=EncodingStage(vae=self.get_module("vae"), ))
 
 
-class T2VPreprocessPipeline(ComposedPipelineBase):
+class PreprocessPipelineT2V(ComposedPipelineBase):
     _required_config_modules = ["text_encoder", "tokenizer", "vae"]
 
     def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
@@ -50,4 +50,4 @@ class T2VPreprocessPipeline(ComposedPipelineBase):
                        stage=EncodingStage(vae=self.get_module("vae"), ))
 
 
-EntryClass = [I2VPreprocessPipeline, T2VPreprocessPipeline]
+EntryClass = [PreprocessPipelineI2V, PreprocessPipelineT2V]

fastvideo/workflow/preprocess/components.py

@@ -149,7 +149,7 @@ class VideoForwardBatchBuilder:
         return forward_batch
 
 
-def basic_record_creator(
+def t2v_record_creator(
         video_name: str,
         vae_latent: np.ndarray,
         text_embedding: np.ndarray,
@@ -230,7 +230,7 @@ class ParquetDatasetSaver:
         self.flush_frequency = flush_frequency
         self.samples_per_file = samples_per_file
         self.schema_fields = schema_fields
-        self.create_record = record_creator or basic_record_creator
+        self.create_record = record_creator or t2v_record_creator
         self.file_writer_fn: Callable[
             [tuple], int] = file_writer_fn or self._default_file_writer_fn
         self.all_tables: list[pa.Table] = []

fastvideo/workflow/preprocess/preprocess_workflow.py

@@ -6,11 +6,14 @@ from torch.utils.data import DataLoader
 
 from fastvideo.configs.configs import PreprocessConfig
 from fastvideo.fastvideo_args import FastVideoArgs, WorkloadType
+from fastvideo.logger import init_logger
 from fastvideo.pipelines.pipeline_registry import PipelineType
 from fastvideo.workflow.preprocess.components import (
     PreprocessingDataValidator, VideoForwardBatchBuilder)
 from fastvideo.workflow.workflow_base import WorkflowBase
 
+logger = init_logger(__name__)
+
 
 class PreprocessWorkflow(WorkflowBase):
 
@@ -49,16 +52,23 @@ class PreprocessWorkflow(WorkflowBase):
         )
         self.add_component("training_dataloader", training_dataloader)
 
-        # validation dataset
-        validation_dataset = load_dataset(preprocess_config.dataset_path,
-                                          split="validation")
-        validation_dataset = validation_dataset.filter(raw_data_validator)
-        validation_dataloader = DataLoader(
-            validation_dataset,
-            batch_size=preprocess_config.preprocess_video_batch_size,
-            num_workers=preprocess_config.dataloader_num_workers,
-            collate_fn=lambda x: x,
-        )
+        # try to load validation dataset if it exists
+        try:
+            validation_dataset = load_dataset(preprocess_config.dataset_path,
+                                              split="validation")
+            validation_dataset = validation_dataset.filter(raw_data_validator)
+            validation_dataloader = DataLoader(
+                validation_dataset,
+                batch_size=preprocess_config.preprocess_video_batch_size,
+                num_workers=preprocess_config.dataloader_num_workers,
+                collate_fn=lambda x: x,
+            )
+        except ValueError:
+            logger.warning(
+                "Validation dataset not found, skipping validation dataset preprocessing."
+            )
+            validation_dataloader = None
+
         self.add_component("validation_dataloader", validation_dataloader)
 
         # forward batch builder

fastvideo/workflow/preprocess/preprocess_workflow_t2v.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Optional
 
 from tqdm import tqdm
 
@@ -17,7 +17,7 @@ if TYPE_CHECKING:
 
 class PreprocessWorkflowT2V(PreprocessWorkflow):
     training_dataloader: "DataLoader"
-    validation_dataloader: "DataLoader"
+    validation_dataloader: Optional["DataLoader"]
     preprocess_pipeline: "ComposedPipelineBase"
     processed_dataset_saver: "ParquetDatasetSaver"
     video_forward_batch_builder: "VideoForwardBatchBuilder"
@@ -54,16 +54,17 @@ class PreprocessWorkflowT2V(PreprocessWorkflow):
         self.processed_dataset_saver.clean_up()
 
         # Validation dataset preprocessing
-        for batch in tqdm(self.validation_dataloader,
-                          desc="Preprocessing validation dataset",
-                          unit="batch"):
-            forward_batch = self.video_forward_batch_builder(batch)
+        if self.validation_dataloader is not None:
+            for batch in tqdm(self.validation_dataloader,
+                              desc="Preprocessing validation dataset",
+                              unit="batch"):
+                forward_batch = self.video_forward_batch_builder(batch)
 
-            forward_batch = self.preprocess_pipeline.forward(
-                forward_batch, self.fastvideo_args)
+                forward_batch = self.preprocess_pipeline.forward(
+                    forward_batch, self.fastvideo_args)
 
-            self.processed_dataset_saver.save_and_write_parquet_batch(
-                forward_batch, self.validation_dataset_output_dir)
-        self.processed_dataset_saver.flush_tables(
-            self.validation_dataset_output_dir)
-        self.processed_dataset_saver.clean_up()
+                self.processed_dataset_saver.save_and_write_parquet_batch(
+                    forward_batch, self.validation_dataset_output_dir)
+            self.processed_dataset_saver.flush_tables(
+                self.validation_dataset_output_dir)
+            self.processed_dataset_saver.clean_up()