[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>

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)