[minor, fix] Allow passing class interface as custom pipeline argument (#2973)

Signed-off-by: knlnguyen1802 <knlnguyen1802@gmail.com>
Signed-off-by: Samit <285365963@qq.com>
Co-authored-by: Samit <285365963@qq.com>

tests/e2e/offline_inference/custom_pipeline/test_worker_extension.py

@@ -7,6 +7,9 @@ from contextlib import ExitStack
 
 import pytest
 
+from tests.e2e.offline_inference.custom_pipeline.qwen_image_pipeline_with_logprob import (
+    QwenImagePipelineWithLogProbForTest,
+)
 from tests.e2e.offline_inference.custom_pipeline.worker_extension import (
     vLLMOmniColocateWorkerExtensionForTest,
 )
@@ -14,9 +17,7 @@ from tests.helpers.mark import hardware_test
 from vllm_omni.diffusion.worker.diffusion_worker import CustomPipelineWorkerExtension
 from vllm_omni.entrypoints.async_omni import AsyncOmni
 
-CUSTOM_PIPELINE_CLASS = (
-    "tests.e2e.offline_inference.custom_pipeline.qwen_image_pipeline_with_logprob.QwenImagePipelineWithLogProbForTest"
-)
+CUSTOM_PIPELINE_CLASS = QwenImagePipelineWithLogProbForTest
 WORKER_EXTENSION_CLASS = (
     "tests.e2e.offline_inference.custom_pipeline.worker_extension.vLLMOmniColocateWorkerExtensionForTest"
 )

vllm_omni/diffusion/model_loader/diffusers_loader.py

@@ -55,6 +55,23 @@ TRANSFORMER_WEIGHTS_INDEX = "model.safetensors.index.json"
 INDEX_FILES = [DIFFUSION_MODEL_WEIGHTS_INDEX, TRANSFORMER_WEIGHTS_INDEX]
 
 
+def _resolve_custom_pipeline_cls(custom_pipeline_name: str | type | None) -> type:
+    """Resolve a custom pipeline reference to a class.
+
+    Accepts either a fully qualified name string (resolved via import) or an
+    already-imported class object (returned as-is).
+    """
+    if custom_pipeline_name is None:
+        raise ValueError("custom_pipeline_name is required for load_format='custom_pipeline'")
+    if isinstance(custom_pipeline_name, str):
+        return resolve_obj_by_qualname(custom_pipeline_name)
+    if isinstance(custom_pipeline_name, type):
+        return custom_pipeline_name
+    raise TypeError(
+        f"custom_pipeline_name must be a qualified name string or a class, got {type(custom_pipeline_name).__name__}"
+    )
+
+
 class DiffusersPipelineLoader:
     """Model loader that can load diffusers pipeline components from disk."""
 
@@ -294,7 +311,7 @@ class DiffusersPipelineLoader:
         od_config: OmniDiffusionConfig,
         load_device: str,
         load_format: str | None = "default",
-        custom_pipeline_name: str | None = None,
+        custom_pipeline_name: str | type[nn.Module] | None = None,
         device: torch.device | None = None,
     ) -> nn.Module:
         """Load a model with the given configurations."""
@@ -343,7 +360,7 @@ class DiffusersPipelineLoader:
                     # tracked by CuMemAllocator.
                     from vllm_omni.diffusion.config import set_current_diffusion_config
 
-                    model_cls = resolve_obj_by_qualname(custom_pipeline_name)
+                    model_cls = _resolve_custom_pipeline_cls(custom_pipeline_name)
                     with set_current_diffusion_config(od_config):
                         model = model_cls(od_config=od_config)
                     if target_device.type != "cpu":
@@ -597,7 +614,7 @@ class DiffusersPipelineLoader:
         od_config: OmniDiffusionConfig,
         target_device: torch.device,
         load_format: str = "default",
-        custom_pipeline_name: str | None = None,
+        custom_pipeline_name: str | type[nn.Module] | None = None,
     ) -> nn.Module:
         """Load model with HSDP sharding for inference.
 
@@ -625,9 +642,9 @@ class DiffusersPipelineLoader:
         if load_format == "default":
             model = initialize_model(od_config)
         elif load_format == "custom_pipeline":
+            model_cls = _resolve_custom_pipeline_cls(custom_pipeline_name)
             from vllm_omni.diffusion.config import set_current_diffusion_config
 
-            model_cls = resolve_obj_by_qualname(custom_pipeline_name)
             with set_current_diffusion_config(od_config):
                 model = model_cls(od_config=od_config)
         self.load_weights(model)

vllm_omni/diffusion/worker/diffusion_model_runner.py

@@ -102,7 +102,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
         self,
         memory_pool_context_fn: callable | None = None,
         load_format: str | None = None,
-        custom_pipeline_name: str | None = None,
+        custom_pipeline_name: str | type | None = None,
     ) -> None:
         """
         Load the diffusion model, apply compilation and offloading.

vllm_omni/entrypoints/utils.py

@@ -183,7 +183,17 @@ def _filter_dict_like_object(obj: dict | Any) -> dict:
     result = {}
     filtered_keys = []
     for k, v in obj.items():
-        if _is_callable_value(v):
+        # Preserve class objects by converting to a fully qualified name string
+        # so callers that resolve via import path (e.g. custom_pipeline_args.pipeline_class)
+        # still work after OmegaConf round-trip.
+        if isinstance(v, type):
+            module = getattr(v, "__module__", None)
+            qualname = getattr(v, "__qualname__", getattr(v, "__name__", None))
+            if module and qualname and module != "builtins":
+                result[k] = f"{module}.{qualname}"
+            else:
+                result[k] = qualname
+        elif _is_callable_value(v):
             filtered_keys.append(str(k))
         else:
             result[k] = _convert_dataclasses_to_dict(v)
@@ -240,6 +250,13 @@ def _convert_dataclasses_to_dict(obj: Any) -> Any:
     # Note: This must come AFTER Counter check since Counter is a dict subclass
     if isinstance(obj, dict):
         return _filter_dict_like_object(obj)
+    # Preserve class objects by converting to a fully qualified name string.
+    if isinstance(obj, type):
+        module = getattr(obj, "__module__", None)
+        qualname = getattr(obj, "__qualname__", getattr(obj, "__name__", None))
+        if module and qualname and module != "builtins":
+            return f"{module}.{qualname}"
+        return qualname
     # Handle callable objects (functions, methods, etc.) - skip them
     # Note: This comes after dict/list checks to avoid misclassifying dict-like objects
     if callable(obj):
@@ -617,7 +634,7 @@ def load_and_resolve_stage_configs(
         if not stage_configs:
             if default_stage_cfg_factory is not None:
                 default_stage_cfg = default_stage_cfg_factory()
-                stage_configs = create_config(default_stage_cfg)
+                stage_configs = create_config(_convert_dataclasses_to_dict(default_stage_cfg))
             else:
                 stage_configs = []
     elif stage_configs_path is None:
@@ -630,7 +647,7 @@ def load_and_resolve_stage_configs(
         if not stage_configs:
             if default_stage_cfg_factory is not None:
                 default_stage_cfg = default_stage_cfg_factory()
-                stage_configs = create_config(default_stage_cfg)
+                stage_configs = create_config(_convert_dataclasses_to_dict(default_stage_cfg))
             else:
                 stage_configs = []
     else: