[2/3][Preprocess] refactor pipeline registry & file structure (#639)

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:

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

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 = {

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",

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.
+"""

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)

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.
+"""

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
 

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