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

Author: Eigensystem

fastvideo/configs/configs.py

@@ -138,7 +138,7 @@ def from_kwargs(cls, kwargs: dict[str,
 
     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 @@ def check_fastvideo_args(self) -> None:
                 "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 @@ def check_fastvideo_args(self) -> None:
         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 @@
 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.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)
+
+    if isinstance(pipeline_type, str):
+        pipeline_type = PipelineType.from_string(pipeline_type)
 
-    # instantiate the pipeline
+    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.
+"""