diff --git a/tests/e2e/offline_inference/custom_pipeline/test_worker_extension.py b/tests/e2e/offline_inference/custom_pipeline/test_worker_extension.py index 653b35d7e..16dddde40 100644 --- a/tests/e2e/offline_inference/custom_pipeline/test_worker_extension.py +++ b/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" ) diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py index d59c7df0a..c9f81b64d 100644 --- a/vllm_omni/diffusion/model_loader/diffusers_loader.py +++ b/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) diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index b18d8fe97..636ac7f60 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/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. diff --git a/vllm_omni/entrypoints/utils.py b/vllm_omni/entrypoints/utils.py index 86a0d8182..031527895 100644 --- a/vllm_omni/entrypoints/utils.py +++ b/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: