From 3a38f7dfdc245ff952d6cb4446dc2b67f4f2de4c Mon Sep 17 00:00:00 2001 From: Jinzhe Pan <48981407+Eigensystem@users.noreply.github.com> Date: Sat, 26 Jul 2025 05:27:12 +0800 Subject: [PATCH] [1/3][Preprocess] refactor preprocessing configs (#638) --- fastvideo/configs/configs.py | 145 +++++++++++++++++++++++++++++ fastvideo/configs/utils.py | 18 +++- fastvideo/fastvideo_args.py | 171 +++++++++++++++++++++++++++++++++-- 3 files changed, 323 insertions(+), 11 deletions(-) create mode 100644 fastvideo/configs/configs.py diff --git a/fastvideo/configs/configs.py b/fastvideo/configs/configs.py new file mode 100644 index 000000000..5bc22aa35 --- /dev/null +++ b/fastvideo/configs/configs.py @@ -0,0 +1,145 @@ +import dataclasses +from typing import Any, Optional + +from fastvideo.configs.utils import update_config_from_args +from fastvideo.utils import FlexibleArgumentParser, StoreBoolean + + +@dataclasses.dataclass +class PreprocessConfig: + """Configuration for preprocessing operations.""" + + # Model and dataset configuration + model_path: str = "" + dataset_path: str = "" + dataset_output_dir: str = "./output" + + # Dataloader configuration + dataloader_num_workers: int = 1 + preprocess_video_batch_size: int = 2 + + # Saver configuration + samples_per_file: int = 64 + flush_frequency: int = 256 + + # Video processing parameters + max_height: int = 480 + max_width: int = 848 + num_frames: int = 163 + video_length_tolerance_range: float = 2.0 + train_fps: int = 30 + speed_factor: float = 1.0 + drop_short_ratio: float = 1.0 + do_temporal_sample: bool = False + + # Model configuration + training_cfg_rate: float = 0.0 + + @staticmethod + def add_cli_args(parser: FlexibleArgumentParser, + prefix: str = "preprocess") -> FlexibleArgumentParser: + """Add preprocessing configuration arguments to the parser.""" + prefix_with_dot = f"{prefix}." if (prefix.strip() != "") else "" + + preprocess_args = parser.add_argument_group("Preprocessing Arguments") + # Model & Dataset + preprocess_args.add_argument(f"--{prefix_with_dot}model-path", + type=str, + default=PreprocessConfig.model_path, + help="Path to the model for preprocessing") + preprocess_args.add_argument( + f"--{prefix_with_dot}dataset-path", + type=str, + default=PreprocessConfig.dataset_path, + help="Path to the dataset directory for preprocessing") + preprocess_args.add_argument( + f"--{prefix_with_dot}dataset-output-dir", + type=str, + default=PreprocessConfig.dataset_output_dir, + help="The output directory where the dataset will be written.") + + # Dataloader + preprocess_args.add_argument( + f"--{prefix_with_dot}dataloader-num-workers", + type=int, + default=PreprocessConfig.dataloader_num_workers, + help= + "Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process." + ) + preprocess_args.add_argument( + f"--{prefix_with_dot}preprocess-video-batch-size", + type=int, + default=PreprocessConfig.preprocess_video_batch_size, + help="Batch size (per device) for the training dataloader.") + + # Saver + preprocess_args.add_argument(f"--{prefix_with_dot}samples-per-file", + type=int, + default=PreprocessConfig.samples_per_file, + help="Number of samples per output file") + preprocess_args.add_argument(f"--{prefix_with_dot}flush-frequency", + type=int, + default=PreprocessConfig.flush_frequency, + help="How often to save to parquet files") + + # Video processing parameters + preprocess_args.add_argument(f"--{prefix_with_dot}max-height", + type=int, + default=PreprocessConfig.max_height, + help="Maximum height for video processing") + preprocess_args.add_argument(f"--{prefix_with_dot}max-width", + type=int, + default=PreprocessConfig.max_width, + help="Maximum width for video processing") + preprocess_args.add_argument(f"--{prefix_with_dot}num-frames", + type=int, + default=PreprocessConfig.num_frames, + help="Number of frames to process") + preprocess_args.add_argument( + f"--{prefix_with_dot}video-length-tolerance-range", + type=float, + default=PreprocessConfig.video_length_tolerance_range, + help="Video length tolerance range") + preprocess_args.add_argument(f"--{prefix_with_dot}train-fps", + type=int, + default=PreprocessConfig.train_fps, + help="Training FPS") + preprocess_args.add_argument(f"--{prefix_with_dot}speed-factor", + type=float, + default=PreprocessConfig.speed_factor, + help="Speed factor for video processing") + preprocess_args.add_argument(f"--{prefix_with_dot}drop-short-ratio", + type=float, + default=PreprocessConfig.drop_short_ratio, + help="Ratio for dropping short videos") + preprocess_args.add_argument( + f"--{prefix_with_dot}do-temporal-sample", + action=StoreBoolean, + default=PreprocessConfig.do_temporal_sample, + help="Whether to do temporal sampling") + + # Model Training configuration + preprocess_args.add_argument(f"--{prefix_with_dot}training-cfg-rate", + type=float, + default=PreprocessConfig.training_cfg_rate, + help="Training CFG rate") + + return parser + + @classmethod + def from_kwargs(cls, kwargs: dict[str, + Any]) -> Optional["PreprocessConfig"]: + """Create PreprocessConfig from keyword arguments.""" + preprocess_config = cls() + if not update_config_from_args( + preprocess_config, kwargs, prefix="preprocess", pop_args=True): + return None + return preprocess_config + + def check_preprocess_config(self) -> None: + if self.dataset_path == "": + raise ValueError("dataset_path must be set for preprocessing mode") + if self.samples_per_file <= 0: + raise ValueError("samples_per_file must be greater than 0") + if self.flush_frequency <= 0: + raise ValueError("flush_frequency must be greater than 0") diff --git a/fastvideo/configs/utils.py b/fastvideo/configs/utils.py index c6d8047db..a2d47b306 100644 --- a/fastvideo/configs/utils.py +++ b/fastvideo/configs/utils.py @@ -1,10 +1,11 @@ +import argparse from typing import Any def update_config_from_args(config: Any, args_dict: dict[str, Any], prefix: str = "", - pop_args: bool = False) -> None: + pop_args: bool = False) -> bool: """ Update configuration object from arguments dictionary. @@ -43,3 +44,18 @@ def update_config_from_args(config: Any, for key in args_to_remove: if key not in args_not_to_remove: args_dict.pop(key) + + return len(args_to_remove) > 0 + + +def clean_cli_args(args: argparse.Namespace) -> dict[str, Any]: + """ + Clean the arguments by removing the ones that not explicitly provided by the user. + """ + provided_args = {} + for k, v in vars(args).items(): + if (v is not None and hasattr(args, '_provided') + and k in args._provided): + provided_args[k] = v + + return provided_args diff --git a/fastvideo/fastvideo_args.py b/fastvideo/fastvideo_args.py index 8e11fb366..d1aeaa3bf 100644 --- a/fastvideo/fastvideo_args.py +++ b/fastvideo/fastvideo_args.py @@ -6,9 +6,12 @@ import argparse import dataclasses from contextlib import contextmanager from dataclasses import field +from enum import Enum from typing import Any +from fastvideo.configs.configs import PreprocessConfig from fastvideo.configs.pipelines.base import PipelineConfig, STA_Mode +from fastvideo.configs.utils import clean_cli_args from fastvideo.logger import init_logger from fastvideo.platforms import current_platform from fastvideo.utils import FlexibleArgumentParser, StoreBoolean @@ -16,17 +19,58 @@ from fastvideo.utils import FlexibleArgumentParser, StoreBoolean logger = init_logger(__name__) -def clean_cli_args(args: argparse.Namespace) -> dict[str, Any]: +class ExecutionMode(str, Enum): """ - Clean the arguments by removing the ones that not explicitly provided by the user. + Enumeration for different pipeline modes. + + Inherits from str to allow string comparison for backward compatibility. """ - provided_args = {} - for k, v in vars(args).items(): - if (v is not None and hasattr(args, '_provided') - and k in args._provided): - provided_args[k] = v + INFERENCE = "inference" + PREPROCESSING = "preprocessing" + FINETUNING = "finetuning" + DISTILLATION = "distillation" - return provided_args + @classmethod + def from_string(cls, value: str) -> "ExecutionMode": + """Convert string to ExecutionMode enum.""" + try: + return cls(value.lower()) + except ValueError: + raise ValueError( + f"Invalid mode: {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 [mode.value for mode in cls] + + +class WorkloadType(str, Enum): + """ + Enumeration for different workload types. + + Inherits from str to allow string comparison for backward compatibility. + """ + I2V = "i2v" # Image to Video + T2V = "t2v" # Text to Video + T2I = "t2i" # Text to Image + I2I = "i2i" # Image to Image + + @classmethod + def from_string(cls, value: str) -> "WorkloadType": + """Convert string to WorkloadType enum.""" + try: + return cls(value.lower()) + except ValueError: + raise ValueError( + f"Invalid workload type: {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 [workload.value for workload in cls] # args for fastvideo framework @@ -35,6 +79,12 @@ class FastVideoArgs: # Model and path configuration (for convenience) model_path: str + # Running mode + mode: ExecutionMode = ExecutionMode.INFERENCE + + # Workload type + workload_type: WorkloadType = WorkloadType.T2V + # Cache strategy cache_strategy: str = "none" @@ -56,6 +106,7 @@ class FastVideoArgs: dist_timeout: int | None = None # timeout for torch.distributed pipeline_config: PipelineConfig = field(default_factory=PipelineConfig) + preprocess_config: PreprocessConfig | None = None # LoRA parameters # (Wenxuan) prefer to keep it here instead of in pipeline config to not make it complicated. @@ -120,6 +171,24 @@ class FastVideoArgs: help="Directory containing StepVideo model", ) + # Running mode + parser.add_argument( + "--mode", + type=str, + choices=ExecutionMode.choices(), + default=FastVideoArgs.mode.value, + help="The mode to run FastVideo", + ) + + # Workload type + parser.add_argument( + "--workload-type", + type=str, + choices=WorkloadType.choices(), + default=FastVideoArgs.workload_type.value, + help="The workload type", + ) + # distributed_executor_backend parser.add_argument( "--distributed-executor-backend", @@ -276,6 +345,9 @@ class FastVideoArgs: # Add pipeline configuration arguments PipelineConfig.add_cli_args(parser) + # Add preprocessing configuration arguments + PreprocessConfig.add_cli_args(parser) + return parser @classmethod @@ -285,11 +357,26 @@ class FastVideoArgs: attrs = [attr.name for attr in dataclasses.fields(cls)] # Create a dictionary of attribute values, with defaults for missing attributes - kwargs = {} + kwargs: dict[str, Any] = {} for attr in attrs: if attr == 'pipeline_config': pipeline_config = PipelineConfig.from_kwargs(provided_args) - kwargs[attr] = pipeline_config + kwargs['pipeline_config'] = pipeline_config + elif attr == 'preprocess_config': + preprocess_config = PreprocessConfig.from_kwargs(provided_args) + kwargs['preprocess_config'] = preprocess_config + elif attr == 'mode': + # Convert string to ExecutionMode enum + mode_value = getattr(args, attr, FastVideoArgs.mode.value) + kwargs['mode'] = ExecutionMode.from_string( + mode_value) if isinstance(mode_value, str) else mode_value + elif attr == 'workload_type': + # Convert string to WorkloadType enum + workload_type_value = getattr(args, 'workload_type', + FastVideoArgs.workload_type.value) + kwargs['workload_type'] = WorkloadType.from_string( + workload_type_value) if isinstance( + workload_type_value, str) else workload_type_value # Use getattr with default value from the dataclass for potentially missing attributes else: # Get the field to check if it has a default_factory @@ -308,7 +395,18 @@ class FastVideoArgs: @classmethod def from_kwargs(cls, **kwargs: Any) -> "FastVideoArgs": + # Convert mode string to enum if necessary + if 'mode' in kwargs and isinstance(kwargs['mode'], str): + kwargs['mode'] = ExecutionMode.from_string(kwargs['mode']) + + # Convert workload_type string to enum if necessary + if 'workload_type' in kwargs and isinstance(kwargs['workload_type'], + str): + kwargs['workload_type'] = WorkloadType.from_string( + kwargs['workload_type']) + kwargs['pipeline_config'] = PipelineConfig.from_kwargs(kwargs) + kwargs['preprocess_config'] = PreprocessConfig.from_kwargs(kwargs) return cls(**kwargs) def check_fastvideo_args(self) -> None: @@ -316,6 +414,34 @@ class FastVideoArgs: if current_platform.is_mps(): self.use_fsdp_inference = False + # Validate mode and inference_mode consistency + assert isinstance( + self.mode, ExecutionMode + ), f"Mode must be an ExecutionMode enum, got {type(self.mode)}" + assert self.mode in ExecutionMode.choices( + ), f"Invalid execution mode: {self.mode}" + + # Validate workload type + assert isinstance( + self.workload_type, WorkloadType + ), f"Workload type must be a WorkloadType enum, got {type(self.workload_type)}" + assert self.workload_type in WorkloadType.choices( + ), f"Invalid workload type: {self.workload_type}" + + if self.mode in [ExecutionMode.DISTILLATION, ExecutionMode.FINETUNING + ] and self.inference_mode: + logger.warning( + "Mode is 'training' but inference_mode is True. Setting inference_mode to False." + ) + self.inference_mode = False + elif self.mode in [ + ExecutionMode.INFERENCE, ExecutionMode.PREPROCESSING + ] and not self.inference_mode: + logger.warning( + "Mode is '%s' but inference_mode is False. Setting inference_mode to True.", + self.mode) + self.inference_mode = True + if not self.inference_mode: assert self.hsdp_replicate_dim != -1, "hsdp_replicate_dim must be set for training" assert self.hsdp_shard_dim != -1, "hsdp_shard_dim must be set for training" @@ -346,6 +472,18 @@ class FastVideoArgs: self.pipeline_config.check_pipeline_config() + # Add preprocessing config validation if needed + if self.mode == ExecutionMode.PREPROCESSING: + if self.preprocess_config is None: + raise ValueError( + "preprocess_config is not set in FastVideoArgs when mode is PREPROCESSING" + ) + if self.preprocess_config.model_path == "": + self.preprocess_config.model_path = self.model_path + if not self.pipeline_config.vae_config.load_encoder: + self.pipeline_config.vae_config.load_encoder = True + self.preprocess_config.check_preprocess_config() + _current_fastvideo_args = None @@ -503,6 +641,19 @@ class TrainingArgs(FastVideoArgs): if attr == 'pipeline_config': pipeline_config = PipelineConfig.from_kwargs(provided_args) kwargs[attr] = pipeline_config + elif attr == 'mode': + # Convert string to ExecutionMode enum + mode_value = getattr(args, attr, ExecutionMode.FINETUNING.value) + kwargs[attr] = ExecutionMode.from_string( + mode_value) if isinstance(mode_value, str) else mode_value + elif attr == 'workload_type': + # Convert string to WorkloadType enum + workload_type_value = getattr(args, 'workload_type', + WorkloadType.T2V.value) + kwargs[attr] = WorkloadType.from_string( + workload_type_value) if isinstance( + workload_type_value, str) else workload_type_value + # Use getattr with default value from the dataclass for potentially missing attributes else: # Get the field to check its default value field = dataclasses.fields(cls)[next(