[Hardware] Extend diffusion engine plugin extensibility for out-of-tree hardware backends (#3239)

Signed-off-by: Yuchen Jiang <yucjiang@amazon.com>
Co-authored-by: Yuchen Jiang <yucjiang@amazon.com>
Co-authored-by: Canlin Guo <canlinguosdu@gmail.com>

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")

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)

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
 

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