From 7e7a0dba9d78dca36dbf6b15227fc5d84fbf880e Mon Sep 17 00:00:00 2001 From: Jinzhe Pan <48981407+Eigensystem@users.noreply.github.com> Date: Tue, 12 Aug 2025 17:16:31 +0800 Subject: [PATCH] feat: preprocess validation dataset only when exist (#734) --- fastvideo/pipelines/pipeline_registry.py | 13 ++++---- .../wan/wan_preprocess_pipelines.py | 6 ++-- fastvideo/workflow/preprocess/components.py | 4 +-- .../preprocess/preprocess_workflow.py | 30 ++++++++++++------- .../preprocess/preprocess_workflow_t2v.py | 27 +++++++++-------- 5 files changed, 46 insertions(+), 34 deletions(-) diff --git a/fastvideo/pipelines/pipeline_registry.py b/fastvideo/pipelines/pipeline_registry.py index fcaa7edcb..d80fd40b5 100644 --- a/fastvideo/pipelines/pipeline_registry.py +++ b/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] diff --git a/fastvideo/pipelines/preprocess/wan/wan_preprocess_pipelines.py b/fastvideo/pipelines/preprocess/wan/wan_preprocess_pipelines.py index c0349c45f..b54b3deb5 100644 --- a/fastvideo/pipelines/preprocess/wan/wan_preprocess_pipelines.py +++ b/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] diff --git a/fastvideo/workflow/preprocess/components.py b/fastvideo/workflow/preprocess/components.py index fcd34201b..23f4d2e0c 100644 --- a/fastvideo/workflow/preprocess/components.py +++ b/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] = [] diff --git a/fastvideo/workflow/preprocess/preprocess_workflow.py b/fastvideo/workflow/preprocess/preprocess_workflow.py index e93321cb6..51f2fadbb 100644 --- a/fastvideo/workflow/preprocess/preprocess_workflow.py +++ b/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 diff --git a/fastvideo/workflow/preprocess/preprocess_workflow_t2v.py b/fastvideo/workflow/preprocess/preprocess_workflow_t2v.py index 54dee8fe2..4ea08e088 100644 --- a/fastvideo/workflow/preprocess/preprocess_workflow_t2v.py +++ b/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()