diff --git a/tests/diffusion/test_diffusion_plugin_hooks.py b/tests/diffusion/test_diffusion_plugin_hooks.py new file mode 100644 index 000000000..724266939 --- /dev/null +++ b/tests/diffusion/test_diffusion_plugin_hooks.py @@ -0,0 +1,110 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +""" +Unit tests for diffusion engine plugin extensibility hooks. + +This module tests: +- Platform hooks: get_diffusion_worker_cls, get_diffusion_model_runner_cls +- Registry API: register_diffusion_model +- Worker integration: model runner resolved via platform hook +""" + +from unittest.mock import patch + +import pytest + +from vllm_omni.diffusion.registry import ( + _DIFFUSION_MODELS, + _DIFFUSION_POST_PROCESS_FUNCS, + _DIFFUSION_PRE_PROCESS_FUNCS, + register_diffusion_model, +) +from vllm_omni.platforms.interface import OmniPlatform, OmniPlatformEnum + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class TestPlatformDiffusionHooks: + """Test OmniPlatform diffusion hook defaults.""" + + def test_get_diffusion_worker_cls_default(self): + """Test default diffusion worker class path.""" + result = OmniPlatform.get_diffusion_worker_cls() + assert result == "vllm_omni.diffusion.worker.diffusion_worker.DiffusionWorker" + + def test_get_diffusion_model_runner_cls_default(self): + """Test default diffusion model runner class path.""" + result = OmniPlatform.get_diffusion_model_runner_cls() + assert result == "vllm_omni.diffusion.worker.diffusion_model_runner.DiffusionModelRunner" + + def test_oot_enum_exists(self): + """Test OOT is a valid platform enum value.""" + assert OmniPlatformEnum.OOT.value == "oot" + + def test_is_out_of_tree(self): + """Test is_out_of_tree() returns True for OOT platform.""" + platform = OmniPlatform.__new__(OmniPlatform) + platform._omni_enum = OmniPlatformEnum.OOT + assert platform.is_out_of_tree() is True + assert platform.is_cuda() is False + + +class TestRegisterDiffusionModel: + """Test register_diffusion_model public API.""" + + @pytest.fixture(autouse=True) + def cleanup_registry(self): + """Restore global registry dicts after each test.""" + original_models = _DIFFUSION_MODELS.copy() + original_pre = _DIFFUSION_PRE_PROCESS_FUNCS.copy() + original_post = _DIFFUSION_POST_PROCESS_FUNCS.copy() + yield + _DIFFUSION_MODELS.clear() + _DIFFUSION_MODELS.update(original_models) + _DIFFUSION_PRE_PROCESS_FUNCS.clear() + _DIFFUSION_PRE_PROCESS_FUNCS.update(original_pre) + _DIFFUSION_POST_PROCESS_FUNCS.clear() + _DIFFUSION_POST_PROCESS_FUNCS.update(original_post) + + def test_register_new_model(self): + """Test registering a new diffusion model with pre/post process functions.""" + register_diffusion_model( + model_arch="TestPipeline", + module_name="test_plugin.diffusion.pipeline", + class_name="TestPipeline", + pre_process_func_name="test_pre_process", + post_process_func_name="test_post_process", + ) + assert "TestPipeline" in _DIFFUSION_MODELS + assert _DIFFUSION_MODELS["TestPipeline"] == ( + "test_plugin.diffusion.pipeline", + "", + "TestPipeline", + ) + assert _DIFFUSION_PRE_PROCESS_FUNCS["TestPipeline"] == "test_pre_process" + assert _DIFFUSION_POST_PROCESS_FUNCS["TestPipeline"] == "test_post_process" + + +class TestWorkerUsesHook: + """Test that DiffusionWorker resolves model runner via platform hook.""" + + @patch("vllm_omni.diffusion.worker.diffusion_worker.resolve_obj_by_qualname") + @patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform") + def test_model_runner_resolved_via_platform(self, mock_platform, mock_resolve): + """Test model runner class is resolved from platform hook return value.""" + from unittest.mock import Mock + + from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker + + mock_runner_instance = Mock() + mock_runner_cls = Mock(return_value=mock_runner_instance) + mock_platform.get_diffusion_model_runner_cls.return_value = "custom.path" + mock_resolve.return_value = mock_runner_cls + + with patch.object(DiffusionWorker, "init_device"): + worker = DiffusionWorker(local_rank=0, rank=0, od_config=Mock(), skip_load_model=True) + + assert worker.model_runner is mock_runner_instance + mock_platform.get_diffusion_model_runner_cls.assert_called_once() + mock_resolve.assert_called_once_with("custom.path") diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index 59497c778..14225e5ce 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -487,10 +487,65 @@ _DIFFUSION_PRE_PROCESS_FUNCS = { } +def register_diffusion_model( + model_arch: str, + module_name: str, + class_name: str, + pre_process_func_name: str | None = None, + post_process_func_name: str | None = None, +) -> None: + """Register a diffusion model pipeline from an out-of-tree plugin. + + This can be used to add new model architectures or to replace an + existing built-in pipeline with a platform-optimised implementation + (same ``model_arch`` key). + + Args: + model_arch: Architecture name (e.g. ``"WanPipeline"``). + module_name: Fully qualified module path + (e.g. ``"my_plugin.diffusion.pipeline_wan"``). + class_name: Class name within *module_name*. + pre_process_func_name: Optional name of the pre-process function + located in *module_name*. Pass ``None`` to keep the existing + entry when replacing a built-in model. + post_process_func_name: Optional name of the post-process function + located in *module_name*. Pass ``None`` to keep the existing + entry when replacing a built-in model. + """ + # Register model class in DiffusionModelRegistry + DiffusionModelRegistry.register_model( + model_arch, + f"{module_name}:{class_name}", + ) + + # Store in _DIFFUSION_MODELS so _load_process_func can resolve the + # module. Convention: when mod_relname is empty the mod_folder field + # stores a *full* module path instead of a relative folder. + _DIFFUSION_MODELS[model_arch] = (module_name, "", class_name) + + # Optionally register pre/post process funcs. + if pre_process_func_name is not None: + _DIFFUSION_PRE_PROCESS_FUNCS[model_arch] = pre_process_func_name + if post_process_func_name is not None: + _DIFFUSION_POST_PROCESS_FUNCS[model_arch] = post_process_func_name + + logger.info( + "Registered diffusion model %s -> %s.%s", + model_arch, + module_name, + class_name, + ) + + def _load_process_func(od_config: OmniDiffusionConfig, func_name: str): """Load and return a process function from the appropriate module.""" mod_folder, mod_relname, _ = _DIFFUSION_MODELS[od_config.model_class_name] - module_name = f"vllm_omni.diffusion.models.{mod_folder}.{mod_relname}" + if mod_relname == "": + # Full module path (registered via register_diffusion_model) + module_name = mod_folder + else: + # Built-in model (relative path convention) + module_name = f"vllm_omni.diffusion.models.{mod_folder}.{mod_relname}" module = importlib.import_module(module_name) func = getattr(module, func_name) return func(od_config) diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index b25addb46..5140e89b7 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/vllm_omni/diffusion/worker/diffusion_worker.py @@ -84,8 +84,10 @@ class DiffusionWorker: self.lora_manager: DiffusionLoRAManager | None = None self.stage_id = getattr(od_config, "stage_id", 0) self.init_device() - # Create model runner - self.model_runner = DiffusionModelRunner( + # Create model runner using the platform-specified class + model_runner_cls_path = current_omni_platform.get_diffusion_model_runner_cls() + model_runner_cls = resolve_obj_by_qualname(model_runner_cls_path) + self.model_runner = model_runner_cls( vllm_config=self.vllm_config, od_config=self.od_config, device=self.device, @@ -567,11 +569,14 @@ class WorkerProc: custom_pipeline_args: dict[str, Any] | None = None, ) -> DiffusionWorker: """Create a worker instance. Override in subclasses for different worker types.""" + worker_cls_path = current_omni_platform.get_diffusion_worker_cls() + base_worker_class = resolve_obj_by_qualname(worker_cls_path) wrapper = WorkerWrapperBase( gpu_id=gpu_id, od_config=od_config, worker_extension_cls=worker_extension_cls, custom_pipeline_args=custom_pipeline_args, + base_worker_class=base_worker_class, ) return wrapper diff --git a/vllm_omni/platforms/interface.py b/vllm_omni/platforms/interface.py index 16a33f9be..0e0762b59 100644 --- a/vllm_omni/platforms/interface.py +++ b/vllm_omni/platforms/interface.py @@ -21,6 +21,7 @@ class OmniPlatformEnum(Enum): NPU = "npu" XPU = "xpu" MUSA = "musa" + OOT = "oot" UNSPECIFIED = "unspecified" @@ -50,6 +51,9 @@ class OmniPlatform(Platform): def is_musa(self) -> bool: return self._omni_enum == OmniPlatformEnum.MUSA + def is_out_of_tree(self) -> bool: + return self._omni_enum == OmniPlatformEnum.OOT + @classmethod def get_omni_ar_worker_cls(cls) -> str: raise NotImplementedError @@ -110,6 +114,25 @@ class OmniPlatform(Platform): """Check if a Flash Attention package is available and usable on this platform.""" return False + @classmethod + def get_diffusion_worker_cls(cls) -> str: + """Get the diffusion worker class path for this platform. + + Returns a fully qualified class path string that will be resolved + and instantiated by WorkerWrapperBase. The class must be compatible + with the DiffusionWorker interface. + """ + return "vllm_omni.diffusion.worker.diffusion_worker.DiffusionWorker" + + @classmethod + def get_diffusion_model_runner_cls(cls) -> str: + """Get the diffusion model runner class path for this platform. + + Returns a fully qualified class path string. The class must be + compatible with the DiffusionModelRunner interface. + """ + return "vllm_omni.diffusion.worker.diffusion_model_runner.DiffusionModelRunner" + @classmethod def get_torch_device(cls, local_rank: int | None = None) -> torch.device: raise NotImplementedError