From 7b6c8aee998f33b19dd59ef71da3608703eba9a2 Mon Sep 17 00:00:00 2001 From: Jinzhe Pan <48981407+Eigensystem@users.noreply.github.com> Date: Mon, 28 Jul 2025 14:30:17 +0800 Subject: [PATCH] [2/3][Preprocess] refactor pipeline registry & file structure (#639) --- fastvideo/configs/configs.py | 2 +- fastvideo/fastvideo_args.py | 11 +- fastvideo/models/registry.py | 4 +- fastvideo/pipelines/__init__.py | 35 ++- fastvideo/pipelines/basic/__init__.py | 6 + .../pipelines/{ => basic}/hunyuan/__init__.py | 0 .../{ => basic}/hunyuan/hunyuan_pipeline.py | 0 .../{ => basic}/stepvideo/__init__.py | 0 .../stepvideo/stepvideo_pipeline.py | 0 .../pipelines/{ => basic}/wan/__init__.py | 0 .../{ => basic}/wan/wan_dmd_pipeline.py | 0 .../{ => basic}/wan/wan_i2v_pipeline.py | 0 .../pipelines/{ => basic}/wan/wan_pipeline.py | 0 fastvideo/pipelines/pipeline_registry.py | 250 ++++++++++++++---- fastvideo/pipelines/training/__init__.py | 6 + .../training/wan_i2v_training_pipeline.py | 3 +- fastvideo/training/wan_training_pipeline.py | 2 +- 17 files changed, 245 insertions(+), 74 deletions(-) create mode 100644 fastvideo/pipelines/basic/__init__.py rename fastvideo/pipelines/{ => basic}/hunyuan/__init__.py (100%) rename fastvideo/pipelines/{ => basic}/hunyuan/hunyuan_pipeline.py (100%) rename fastvideo/pipelines/{ => basic}/stepvideo/__init__.py (100%) rename fastvideo/pipelines/{ => basic}/stepvideo/stepvideo_pipeline.py (100%) rename fastvideo/pipelines/{ => basic}/wan/__init__.py (100%) rename fastvideo/pipelines/{ => basic}/wan/wan_dmd_pipeline.py (100%) rename fastvideo/pipelines/{ => basic}/wan/wan_i2v_pipeline.py (100%) rename fastvideo/pipelines/{ => basic}/wan/wan_pipeline.py (100%) create mode 100644 fastvideo/pipelines/training/__init__.py diff --git a/fastvideo/configs/configs.py b/fastvideo/configs/configs.py index 5bc22aa35..d76b3428f 100644 --- a/fastvideo/configs/configs.py +++ b/fastvideo/configs/configs.py @@ -138,7 +138,7 @@ class PreprocessConfig: def check_preprocess_config(self) -> None: if self.dataset_path == "": - raise ValueError("dataset_path must be set for preprocessing mode") + raise ValueError("dataset_path must be set for preprocess mode") if self.samples_per_file <= 0: raise ValueError("samples_per_file must be greater than 0") if self.flush_frequency <= 0: diff --git a/fastvideo/fastvideo_args.py b/fastvideo/fastvideo_args.py index fa66be947..b451c13b5 100644 --- a/fastvideo/fastvideo_args.py +++ b/fastvideo/fastvideo_args.py @@ -26,7 +26,7 @@ class ExecutionMode(str, Enum): Inherits from str to allow string comparison for backward compatibility. """ INFERENCE = "inference" - PREPROCESSING = "preprocessing" + PREPROCESS = "preprocess" FINETUNING = "finetuning" DISTILLATION = "distillation" @@ -471,9 +471,8 @@ class FastVideoArgs: "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: + elif self.mode in [ExecutionMode.INFERENCE, ExecutionMode.PREPROCESS + ] and not self.inference_mode: logger.warning( "Mode is '%s' but inference_mode is False. Setting inference_mode to True.", self.mode) @@ -510,10 +509,10 @@ class FastVideoArgs: self.pipeline_config.check_pipeline_config() # Add preprocessing config validation if needed - if self.mode == ExecutionMode.PREPROCESSING: + if self.mode == ExecutionMode.PREPROCESS: if self.preprocess_config is None: raise ValueError( - "preprocess_config is not set in FastVideoArgs when mode is PREPROCESSING" + "preprocess_config is not set in FastVideoArgs when mode is PREPROCESS" ) if self.preprocess_config.model_path == "": self.preprocess_config.model_path = self.model_path diff --git a/fastvideo/models/registry.py b/fastvideo/models/registry.py index 8f364cd08..5c585fa16 100644 --- a/fastvideo/models/registry.py +++ b/fastvideo/models/registry.py @@ -16,7 +16,9 @@ from typing import NoReturn, TypeVar, cast import cloudpickle from torch import nn -from fastvideo.logger import logger +from fastvideo.logger import init_logger + +logger = init_logger(__name__) # huggingface class name: (component_name, fastvideo module name, fastvideo class name) _TEXT_TO_VIDEO_DIT_MODELS = { diff --git a/fastvideo/pipelines/__init__.py b/fastvideo/pipelines/__init__.py index f6ab43463..2a20ec2ce 100644 --- a/fastvideo/pipelines/__init__.py +++ b/fastvideo/pipelines/__init__.py @@ -12,7 +12,8 @@ from fastvideo.logger import init_logger from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase from fastvideo.pipelines.lora_pipeline import LoRAPipeline from fastvideo.pipelines.pipeline_batch_info import ForwardBatch, TrainingBatch -from fastvideo.pipelines.pipeline_registry import PipelineRegistry +from fastvideo.pipelines.pipeline_registry import (PipelineType, + get_pipeline_registry) from fastvideo.utils import (maybe_download_model, verify_model_config_and_directory) @@ -24,7 +25,10 @@ class PipelineWithLoRA(LoRAPipeline, ComposedPipelineBase): pass -def build_pipeline(fastvideo_args: FastVideoArgs) -> PipelineWithLoRA: +def build_pipeline( + fastvideo_args: FastVideoArgs, + pipeline_type: PipelineType | str = PipelineType.BASIC +) -> PipelineWithLoRA: """ Only works with valid hf diffusers configs. (model_index.json) We want to build a pipeline based on the inference args mode_path: @@ -37,30 +41,37 @@ def build_pipeline(fastvideo_args: FastVideoArgs) -> PipelineWithLoRA: model_path = maybe_download_model(model_path) # fastvideo_args.downloaded_model_path = model_path logger.info("Model path: %s", model_path) - config = verify_model_config_and_directory(model_path) - pipeline_architecture = config.get("_class_name") - if pipeline_architecture is None: + config = verify_model_config_and_directory(model_path) + pipeline_name = config.get("_class_name") + if pipeline_name is None: raise ValueError( "Model config does not contain a _class_name attribute. " "Only diffusers format is supported.") - pipeline_cls, pipeline_architecture = PipelineRegistry.resolve_pipeline_cls( - pipeline_architecture) + # Get the appropriate pipeline registry based on pipeline_type + logger.info( + "Building pipeline of type: %s", pipeline_type.value if isinstance( + pipeline_type, PipelineType) else pipeline_type) + pipeline_registry = get_pipeline_registry(pipeline_type) - # instantiate the pipeline + if isinstance(pipeline_type, str): + pipeline_type = PipelineType.from_string(pipeline_type) + + pipeline_cls = pipeline_registry.resolve_pipeline_cls( + pipeline_name, pipeline_type, fastvideo_args.workload_type) + + # instantiate the pipelines pipeline = pipeline_cls(model_path, fastvideo_args) - logger.info("Pipeline instantiated") - # pipeline is now initialized and ready to use + logger.info("Pipelines instantiated") + return cast(PipelineWithLoRA, pipeline) __all__ = [ "build_pipeline", - "list_available_pipelines", "ComposedPipelineBase", - "PipelineRegistry", "ForwardBatch", "LoRAPipeline", "TrainingBatch", diff --git a/fastvideo/pipelines/basic/__init__.py b/fastvideo/pipelines/basic/__init__.py new file mode 100644 index 000000000..f0f2eec35 --- /dev/null +++ b/fastvideo/pipelines/basic/__init__.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +Basic inference pipelines for fastvideo. + +This package contains basic pipelines for video and image generation. +""" diff --git a/fastvideo/pipelines/hunyuan/__init__.py b/fastvideo/pipelines/basic/hunyuan/__init__.py similarity index 100% rename from fastvideo/pipelines/hunyuan/__init__.py rename to fastvideo/pipelines/basic/hunyuan/__init__.py diff --git a/fastvideo/pipelines/hunyuan/hunyuan_pipeline.py b/fastvideo/pipelines/basic/hunyuan/hunyuan_pipeline.py similarity index 100% rename from fastvideo/pipelines/hunyuan/hunyuan_pipeline.py rename to fastvideo/pipelines/basic/hunyuan/hunyuan_pipeline.py diff --git a/fastvideo/pipelines/stepvideo/__init__.py b/fastvideo/pipelines/basic/stepvideo/__init__.py similarity index 100% rename from fastvideo/pipelines/stepvideo/__init__.py rename to fastvideo/pipelines/basic/stepvideo/__init__.py diff --git a/fastvideo/pipelines/stepvideo/stepvideo_pipeline.py b/fastvideo/pipelines/basic/stepvideo/stepvideo_pipeline.py similarity index 100% rename from fastvideo/pipelines/stepvideo/stepvideo_pipeline.py rename to fastvideo/pipelines/basic/stepvideo/stepvideo_pipeline.py diff --git a/fastvideo/pipelines/wan/__init__.py b/fastvideo/pipelines/basic/wan/__init__.py similarity index 100% rename from fastvideo/pipelines/wan/__init__.py rename to fastvideo/pipelines/basic/wan/__init__.py diff --git a/fastvideo/pipelines/wan/wan_dmd_pipeline.py b/fastvideo/pipelines/basic/wan/wan_dmd_pipeline.py similarity index 100% rename from fastvideo/pipelines/wan/wan_dmd_pipeline.py rename to fastvideo/pipelines/basic/wan/wan_dmd_pipeline.py diff --git a/fastvideo/pipelines/wan/wan_i2v_pipeline.py b/fastvideo/pipelines/basic/wan/wan_i2v_pipeline.py similarity index 100% rename from fastvideo/pipelines/wan/wan_i2v_pipeline.py rename to fastvideo/pipelines/basic/wan/wan_i2v_pipeline.py diff --git a/fastvideo/pipelines/wan/wan_pipeline.py b/fastvideo/pipelines/basic/wan/wan_pipeline.py similarity index 100% rename from fastvideo/pipelines/wan/wan_pipeline.py rename to fastvideo/pipelines/basic/wan/wan_pipeline.py diff --git a/fastvideo/pipelines/pipeline_registry.py b/fastvideo/pipelines/pipeline_registry.py index a51e0d2eb..08b57d161 100644 --- a/fastvideo/pipelines/pipeline_registry.py +++ b/fastvideo/pipelines/pipeline_registry.py @@ -6,84 +6,230 @@ import importlib import pkgutil from collections.abc import Set from dataclasses import dataclass, field +from enum import Enum from functools import lru_cache +from fastvideo.fastvideo_args import WorkloadType from fastvideo.logger import init_logger from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase from fastvideo.pipelines.lora_pipeline import LoRAPipeline logger = init_logger(__name__) +# map pipeline name to folder name +_PIPELINE_NAME_TO_ARCHITECTURE_NAME: dict[str, str] = { + "WanPipeline": "wan", + "WanImageToVideoPipeline": "wan", + "StepVideoPipeline": "stepvideo", + "HunyuanVideoPipeline": "hunyuan", +} + + +class PipelineType(str, Enum): + """ + Enumeration for different pipeline types. + + Inherits from str to allow string comparison for backward compatibility. + """ + BASIC = "basic" + PREPROCESS = "preprocess" + TRAINING = "training" + + @classmethod + def from_string(cls, value: str) -> "PipelineType": + """Convert string to PipelineType enum.""" + try: + return cls(value.lower()) + except ValueError: + raise ValueError( + f"Invalid pipeline type: {value}. Must be one of: {', '.join([t.value for t in cls])}" + ) from None + + @classmethod + def choices(cls) -> list[str]: + """Get all available choices as strings.""" + return [pipeline_type.value for pipeline_type in cls] + @dataclass class _PipelineRegistry: - # Keyed by pipeline_arch - pipelines: dict[str, type[ComposedPipelineBase] - | None] = field(default_factory=dict) + # Keyed by pipeline_type -> architecture -> pipeline_name + # pipelines[pipeline_type][architecture][pipeline_name] = pipeline_cls + pipelines: dict[str, dict[str, dict[str, type[ComposedPipelineBase] + | None]]] = field(default_factory=dict) - def get_supported_archs(self) -> Set[str]: - return self.pipelines.keys() + def get_supported_archs(self, pipeline_name_in_config: str, + pipeline_type: PipelineType) -> Set[str]: + """Get supported architectures, optionally filtered by pipeline type and workload type.""" + arch = _PIPELINE_NAME_TO_ARCHITECTURE_NAME[pipeline_name_in_config] + return set(self.pipelines[pipeline_type.value][arch].keys()) + + def _load_preprocessing_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}") + + return self.pipelines[ + PipelineType.PREPROCESS.value][arch][pipeline_name] def _try_load_pipeline_cls( - self, pipeline_arch: str) -> type[ComposedPipelineBase] | None: - if pipeline_arch not in self.pipelines: + self, pipeline_name_in_config: str, pipeline_type: PipelineType, + workload_type: WorkloadType + ) -> type[ComposedPipelineBase] | type[LoRAPipeline] | None: + """Try to load a pipeline class for the given architecture, pipeline type, and workload type.""" + arch = _PIPELINE_NAME_TO_ARCHITECTURE_NAME[pipeline_name_in_config] + + if (pipeline_type.value not in self.pipelines + or arch not in self.pipelines[pipeline_type.value]): return None - return self.pipelines[pipeline_arch] + if pipeline_type == PipelineType.PREPROCESS: + return self._load_preprocessing_pipeline_cls(workload_type, arch) + elif pipeline_type == PipelineType.BASIC: + return self.pipelines[ + pipeline_type.value][arch][pipeline_name_in_config] + elif pipeline_type == PipelineType.TRAINING: + pass + else: + raise ValueError(f"Invalid pipeline type: {pipeline_type.value}") + + return None def resolve_pipeline_cls( self, - architecture: str, - ) -> tuple[type[ComposedPipelineBase] | type[LoRAPipeline], str]: - if not architecture: + pipeline_name_in_config: str, + pipeline_type: PipelineType, + workload_type: WorkloadType, + ) -> type[ComposedPipelineBase] | type[LoRAPipeline]: + """Resolve pipeline class based on pipeline name in the config, pipeline type, and workload type.""" + if not pipeline_name_in_config: logger.warning("No pipeline architecture is specified") - pipeline_cls = self._try_load_pipeline_cls(architecture) + pipeline_cls = self._try_load_pipeline_cls(pipeline_name_in_config, + pipeline_type, workload_type) if pipeline_cls is not None: - return (pipeline_cls, architecture) - - supported_archs = self.get_supported_archs() + return pipeline_cls + supported_archs = self.get_supported_archs(pipeline_name_in_config, + pipeline_type) raise ValueError( - f"Pipeline architectures {architecture} are not supported for now. " + f"Pipeline architecture '{pipeline_name_in_config}' is not supported for pipeline type '{pipeline_type.value}' " + f"and workload type '{workload_type.value}'. " f"Supported architectures: {supported_archs}") @lru_cache -def import_pipeline_classes(): - pipeline_arch_name_to_cls = {} - package_name = "fastvideo.pipelines" - package = importlib.import_module(package_name) - for _, name, ispkg in pkgutil.iter_modules(package.__path__, - package_name + "."): - if ispkg: - if name.split(".")[-1] == "stages": - continue - sub_package_name = name - sub_package = importlib.import_module(sub_package_name) - for _, name, ispkg in pkgutil.iter_modules(sub_package.__path__, - sub_package_name + "."): - try: - module = importlib.import_module(name) - except Exception as e: - logger.warning("Ignore import error when loading %s. %s", - name, e) - continue - if hasattr(module, "EntryClass"): - entry = module.EntryClass - if isinstance( - entry, list - ): # To support multiple pipeline classes in one module - for tmp in entry: - assert ( - tmp.__name__ not in pipeline_arch_name_to_cls - ), f"Duplicated pipeline implementation for {tmp.__name__}" - pipeline_arch_name_to_cls[tmp.__name__] = tmp - else: - assert ( - entry.__name__ not in pipeline_arch_name_to_cls - ), f"Duplicated pipeline implementation for {entry.__name__}" - pipeline_arch_name_to_cls[entry.__name__] = entry - return pipeline_arch_name_to_cls +def import_pipeline_classes( + pipeline_types: list[PipelineType] | PipelineType | None = None +) -> dict[str, dict[str, dict[str, type[ComposedPipelineBase] | None]]]: + """ + Import pipeline classes based on the pipeline type and workload type. + + Args: + pipeline_types: The pipeline types to load (basic, preprocessing, training). + If None, loads all types. + + Returns: + A three-level nested dictionary: + {pipeline_type: {architecture_name: {pipeline_name: pipeline_cls}}} + e.g., {"basic": {"wan": {"WanPipeline": WanPipeline}}} + """ + type_to_arch_to_pipeline_dict: dict[str, + dict[str, + dict[str, + type[ComposedPipelineBase] + | None]]] = {} + package_name: str = "fastvideo.pipelines" + + # Determine which pipeline types to scan + if isinstance(pipeline_types, list): + pipeline_types_to_scan = [ + pipeline_type.value for pipeline_type in pipeline_types + ] + elif isinstance(pipeline_types, PipelineType): + pipeline_types_to_scan = [pipeline_types.value] + else: + pipeline_types_to_scan = [pt.value for pt in PipelineType] + + logger.info("Loading pipelines for types: %s", pipeline_types_to_scan) + + for pipeline_type_str in pipeline_types_to_scan: + arch_to_pipeline_dict: dict[str, dict[str, type[ComposedPipelineBase] + | None]] = {} + + # Try to load from pipeline-type-specific directory first + pipeline_type_package_name = f"{package_name}.{pipeline_type_str}" + + try: + pipeline_type_package = importlib.import_module( + pipeline_type_package_name) + logger.debug("Successfully imported %s", pipeline_type_package_name) + + for _, arch, ispkg in pkgutil.iter_modules( + pipeline_type_package.__path__): + pipeline_dict: dict[str, type[ComposedPipelineBase] | None] = {} + + arch_package_name = f"{pipeline_type_package_name}.{arch}" + if ispkg: + arch_package = importlib.import_module(arch_package_name) + for _, module_name, ispkg in pkgutil.walk_packages( + arch_package.__path__, arch_package_name + "."): + if not ispkg: + pipeline_module = importlib.import_module( + module_name) + if hasattr(pipeline_module, "EntryClass"): + if isinstance(pipeline_module.EntryClass, list): + for pipeline in pipeline_module.EntryClass: + pipeline_name = pipeline.__name__ + assert ( + pipeline_name not in pipeline_dict + ), f"Duplicated pipeline implementation for {pipeline_name} in {pipeline_type_str}.{arch_package_name}" + pipeline_dict[pipeline_name] = pipeline + else: + pipeline_name = pipeline_module.EntryClass.__name__ + assert ( + pipeline_name not in pipeline_dict + ), f"Duplicated pipeline implementation for {pipeline_name} in {pipeline_type_str}.{arch_package_name}" + pipeline_dict[ + pipeline_name] = pipeline_module.EntryClass + + arch_to_pipeline_dict[arch] = pipeline_dict + + except ImportError as e: + raise ImportError( + f"Could not import {pipeline_type_package_name} when importing pipeline classes: {e}" + ) from None + + type_to_arch_to_pipeline_dict[pipeline_type_str] = arch_to_pipeline_dict + + # Log summary + total_pipelines = sum( + len(pipeline_dict) + for arch_to_pipeline_dict in type_to_arch_to_pipeline_dict.values() + for pipeline_dict in arch_to_pipeline_dict.values()) + logger.info("Loaded %d pipeline classes across %d types", total_pipelines, + len(pipeline_types_to_scan)) + + return type_to_arch_to_pipeline_dict -PipelineRegistry = _PipelineRegistry(import_pipeline_classes()) +def get_pipeline_registry( + pipeline_type: PipelineType | str | None = None) -> _PipelineRegistry: + """ + Get a pipeline registry for the specified mode, pipeline type, and workload type. + + Args: + pipeline_type: Pipeline type to load. If None and mode is provided, will be derived from mode. + + Returns: + A pipeline registry instance. + """ + if isinstance(pipeline_type, str): + pipeline_type = PipelineType.from_string(pipeline_type) + + pipeline_classes = import_pipeline_classes(pipeline_type) + return _PipelineRegistry(pipeline_classes) diff --git a/fastvideo/pipelines/training/__init__.py b/fastvideo/pipelines/training/__init__.py new file mode 100644 index 000000000..e5927d838 --- /dev/null +++ b/fastvideo/pipelines/training/__init__.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +Training pipelines for fastvideo.v1. + +This package contains pipelines for training diffusion models. +""" diff --git a/fastvideo/training/wan_i2v_training_pipeline.py b/fastvideo/training/wan_i2v_training_pipeline.py index fc5774c4c..18810b0c1 100644 --- a/fastvideo/training/wan_i2v_training_pipeline.py +++ b/fastvideo/training/wan_i2v_training_pipeline.py @@ -12,8 +12,9 @@ from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs from fastvideo.logger import init_logger from fastvideo.models.schedulers.scheduling_flow_unipc_multistep import ( FlowUniPCMultistepScheduler) +from fastvideo.pipelines.basic.wan.wan_i2v_pipeline import ( + WanImageToVideoPipeline) from fastvideo.pipelines.pipeline_batch_info import ForwardBatch, TrainingBatch -from fastvideo.pipelines.wan.wan_i2v_pipeline import WanImageToVideoPipeline from fastvideo.training.training_pipeline import TrainingPipeline from fastvideo.utils import is_vsa_available, shallow_asdict diff --git a/fastvideo/training/wan_training_pipeline.py b/fastvideo/training/wan_training_pipeline.py index b304b9e6a..c9e7ab593 100644 --- a/fastvideo/training/wan_training_pipeline.py +++ b/fastvideo/training/wan_training_pipeline.py @@ -6,7 +6,7 @@ from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs from fastvideo.logger import init_logger from fastvideo.models.schedulers.scheduling_flow_unipc_multistep import ( FlowUniPCMultistepScheduler) -from fastvideo.pipelines.wan.wan_pipeline import WanPipeline +from fastvideo.pipelines.basic.wan.wan_pipeline import WanPipeline from fastvideo.training.training_pipeline import TrainingPipeline from fastvideo.utils import is_vsa_available