[diffusion][bugfix] Separate Diffusers hooks from model metadata (#6749)

Signed-off-by: david6666666 <530634352@qq.com>
Signed-off-by: Gao Han <hgaoaf@connect.ust.hk>
Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>

tests/diffusion/diffusion_backend/test_diffusers_backend.py

@@ -9,6 +9,7 @@ import torch
 from diffusers import DiffusionPipeline  # pyright: ignore[reportPrivateImportUsage]
 from PIL import Image
 
+import vllm_omni.diffusion.data as diffusion_data
 from tests.helpers.mark import hardware_test
 from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient, dummy_messages_from_mix_data
 from vllm_omni.diffusion.data import (
@@ -629,6 +630,46 @@ class TestPipelineArgumentsHandling:
         with pytest.raises(ValueError):
             pipeline.forward(DiffusionRequestBatch(requests=[problematic_request]))
 
+        interpolation_request = _make_request(
+            sampling_params=OmniDiffusionSamplingParams(
+                enable_frame_interpolation=True,
+            ),
+        )
+        with pytest.raises(ValueError, match="Frame interpolation is not supported with the Diffusers backend"):
+            pipeline.forward(DiffusionRequestBatch(requests=[interpolation_request]))
+
+    def test_diffusers_config_preserves_checkpoint_class_for_metadata(self, mocker):
+        class MockWanImageToVideoPipeline:
+            pass
+
+        od_config = _make_od_config(model_class_name=None)
+        mocker.patch(
+            "vllm_omni.diffusion.utils.hf_utils.get_diffusion_model_index",
+            return_value={"_class_name": "WanImageToVideoPipeline"},
+        )
+        mocker.patch.object(
+            diffusion_data,
+            "diffusers",
+            SimpleNamespace(WanImageToVideoPipeline=MockWanImageToVideoPipeline),
+        )
+        mocker.patch.object(OmniDiffusionConfig, "update_multimodal_support")
+
+        od_config.enrich_config()
+
+        assert od_config.model_class_name == "WanImageToVideoPipeline"
+        assert od_config.diffusers_pipeline_cls is MockWanImageToVideoPipeline
+
+    def test_diffusers_config_falls_back_to_adapter_without_pipeline_index(self, mocker):
+        od_config = _make_od_config(model_class_name=None)
+        mocker.patch(
+            "vllm_omni.diffusion.utils.hf_utils.get_diffusion_model_index",
+            return_value=None,
+        )
+
+        od_config.enrich_config()
+
+        assert od_config.model_class_name == "DiffusersAdapterPipeline"
+
 
 @pytest.mark.advanced_model
 @hardware_test(res={"cuda": "L4"}, num_cards=1)

tests/diffusion/models/qwen_image/test_qwen_image_postprocess.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 import json
 from types import SimpleNamespace
 

tests/diffusion/test_diffusion_engine.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 import asyncio
 import queue
@@ -197,6 +197,59 @@ class TestRequestBatchCapability:
 
         assert diffusion_engine_module.supports_request_batch(od_config) is True
 
+    def test_engine_rejects_multi_seq_diffusers_using_adapter_capability(
+        self,
+        monkeypatch: pytest.MonkeyPatch,
+        mocker: MockerFixture,
+    ) -> None:
+        od_config = SimpleNamespace(
+            model_class_name="QwenImagePipeline",
+            custom_pipeline_args=None,
+            diffusion_load_format="diffusers",
+            streaming_output=False,
+            max_num_seqs=2,
+        )
+        registry_load = mocker.Mock(
+            side_effect=lambda model_class_name: (
+                _SingleRequestPipeline if model_class_name == "DiffusersAdapterPipeline" else _BatchCapablePipeline
+            )
+        )
+        monkeypatch.setattr(
+            diffusion_engine_module.DiffusionModelRegistry,
+            "_try_load_model_cls",
+            registry_load,
+        )
+
+        engine = DiffusionEngine.__new__(DiffusionEngine)
+        with pytest.raises(ValueError, match="max_num_seqs=1"):
+            engine._resolve_execution_mode(od_config)
+
+        registry_load.assert_called_once_with("DiffusersAdapterPipeline")
+
+    def test_engine_prefers_batch_capable_custom_pipeline_over_diffusers_adapter(
+        self,
+        monkeypatch: pytest.MonkeyPatch,
+        mocker: MockerFixture,
+    ) -> None:
+        od_config = SimpleNamespace(
+            model_class_name="QwenImagePipeline",
+            custom_pipeline_args={"pipeline_class": _BatchCapablePipeline},
+            diffusion_load_format="diffusers",
+            streaming_output=False,
+            max_num_seqs=2,
+        )
+        registry_load = mocker.Mock(return_value=_SingleRequestPipeline)
+        monkeypatch.setattr(
+            diffusion_engine_module.DiffusionModelRegistry,
+            "_try_load_model_cls",
+            registry_load,
+        )
+
+        engine = DiffusionEngine.__new__(DiffusionEngine)
+
+        assert engine._resolve_execution_mode(od_config) is DiffusionExecutionMode.REQUEST_BATCH
+        registry_load.assert_not_called()
+
     def test_supports_request_batch_uses_custom_pipeline_class(self, monkeypatch: pytest.MonkeyPatch) -> None:
         od_config = SimpleNamespace(
             model_class_name="SinglePipeline",

tests/diffusion/test_diffusion_output_formatter.py

@@ -328,6 +328,25 @@ def test_formatter_preserves_audio_model_video_audio_and_actions(
     }
 
 
+def test_formatter_uses_native_model_identity_for_diffusers_video_frames(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(output_formatter, "supports_audio_output", lambda _: False)
+    frames = [["frame-0", "frame-1"]]
+    postprocess_output = normalize_diffusion_postprocess_output(frames)
+
+    [result] = format_diffusion_outputs(
+        request=_request("make a video"),
+        od_config=_config("WanImageToVideoPipeline"),
+        diffusion_output=DiffusionOutput(output=frames),
+        output_data=frames,
+        postprocess_output=postprocess_output,
+    )
+
+    assert result.images == frames
+    assert result.final_output_type == "video"
+
+
 def test_formatter_preserves_audio_only_postprocess_dict(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:

tests/diffusion/test_diffusion_plugin_hooks.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 """
 Unit tests for diffusion engine plugin extensibility hooks.
@@ -20,6 +20,8 @@ from vllm_omni.diffusion.registry import (
     _DIFFUSION_MODELS,
     _DIFFUSION_POST_PROCESS_FUNCS,
     _DIFFUSION_PRE_PROCESS_FUNCS,
+    get_diffusion_post_process_func,
+    get_diffusion_pre_process_func,
     register_diffusion_model,
 )
 from vllm_omni.platforms.interface import OmniPlatform, OmniPlatformEnum
@@ -112,6 +114,49 @@ class TestRegisterDiffusionModel:
         assert _DIFFUSION_POST_PROCESS_FUNCS["LegacyActionPipeline"] == "test_post_process"
 
 
+@pytest.mark.parametrize(
+    ("model_class_name", "get_process_func"),
+    [
+        ("QwenImagePipeline", get_diffusion_post_process_func),
+        ("WanImageToVideoPipeline", get_diffusion_post_process_func),
+        ("WanImageToVideoPipeline", get_diffusion_pre_process_func),
+    ],
+)
+def test_diffusers_backend_skips_native_process_hooks(model_class_name, get_process_func):
+    od_config = SimpleNamespace(
+        model_class_name=model_class_name,
+        custom_pipeline_args=None,
+        diffusion_load_format="diffusers",
+    )
+
+    with patch("vllm_omni.diffusion.registry._load_process_func") as load_process_func:
+        assert get_process_func(od_config) is None
+
+    load_process_func.assert_not_called()
+
+
+@pytest.mark.parametrize(
+    ("model_class_name", "get_process_func"),
+    [
+        ("QwenImagePipeline", get_diffusion_post_process_func),
+        ("WanImageToVideoPipeline", get_diffusion_post_process_func),
+        ("WanImageToVideoPipeline", get_diffusion_pre_process_func),
+    ],
+)
+def test_custom_pipeline_preserves_native_process_hooks(model_class_name, get_process_func):
+    process_func = Mock()
+    od_config = SimpleNamespace(
+        model_class_name=model_class_name,
+        custom_pipeline_args={"pipeline_class": "custom.Pipeline"},
+        diffusion_load_format="diffusers",
+    )
+
+    with patch("vllm_omni.diffusion.registry._load_process_func", return_value=process_func) as load_process_func:
+        assert get_process_func(od_config) is process_func
+
+    load_process_func.assert_called_once()
+
+
 class TestWorkerUsesHook:
     """Test that DiffusionWorker resolves model runner via platform hook."""
 

tests/entrypoints/test_async_omni_diffusion_config.py

@@ -627,6 +627,36 @@ def test_default_stage_resolves_video_output_from_checkpoint(mocker):
     assert captured["final_output_type"] == "video"
 
 
+def test_default_diffusers_stage_preserves_video_model_identity(mocker):
+    captured = {}
+
+    def resolve_with_default(*args, default_stage_cfg_factory, **kwargs):
+        del args, kwargs
+        captured.update(default_stage_cfg_factory()[0])
+        stage = SimpleNamespace(stage_type="diffusion", engine_args=SimpleNamespace())
+        return ("", [stage], None)
+
+    mocker.patch(
+        "vllm_omni.diffusion.utils.hf_utils.get_diffusion_model_index",
+        return_value={"_class_name": "WanImageToVideoPipeline"},
+    )
+    mocker.patch(
+        "vllm_omni.engine.async_omni_engine.load_and_resolve_stage_configs",
+        side_effect=resolve_with_default,
+    )
+    engine = AsyncOmniEngine.__new__(AsyncOmniEngine)
+    engine._strip_single_engine_args = lambda kwargs: kwargs
+
+    engine._resolve_stage_configs(
+        "/models/Wan2.2-I2V",
+        {"diffusion_load_format": "diffusers"},
+        trust_remote_code=False,
+    )
+
+    assert captured["engine_args"]["model_class_name"] == "WanImageToVideoPipeline"
+    assert captured["final_output_type"] == "video"
+
+
 def test_resolve_stage_configs_injects_additional_config_into_diffusion_stage(mocker):
     """Ensure YAML/deploy stage resolution forwards top-level additional_config."""
     fake_diffusion_stage = SimpleNamespace(

vllm_omni/diffusion/data.py

@@ -685,6 +685,19 @@ def resolve_model_class_name(
     return None
 
 
+def uses_diffusers_adapter(od_config: object) -> bool:
+    """Return whether execution uses the Diffusers adapter backend.
+
+    A custom pipeline takes precedence over ``diffusion_load_format`` in the
+    worker, so adapter-specific behavior must only apply when no custom
+    pipeline override is configured.
+    """
+    return (
+        getattr(od_config, "custom_pipeline_args", None) is None
+        and getattr(od_config, "diffusion_load_format", "default") == "diffusers"
+    )
+
+
 @dataclass
 class OmniDiffusionConfig:
     # Model and path configuration (for convenience)
@@ -1317,10 +1330,6 @@ class OmniDiffusionConfig:
 
         from vllm_omni.diffusion.utils.hf_utils import get_diffusion_model_index
 
-        # Default model_class_name for diffusers adapter
-        if self.model_class_name is None and self.diffusion_load_format == "diffusers":
-            self.model_class_name = "DiffusersAdapterPipeline"
-
         assert self.model is not None
         try:
             config_dict = get_diffusion_model_index(
@@ -1330,6 +1339,8 @@ class OmniDiffusionConfig:
             if config_dict is not None:
                 if self.model_class_name is None:
                     self.model_class_name = config_dict.get("_class_name", None)
+                    if self.model_class_name is None and self.diffusion_load_format == "diffusers":
+                        self.model_class_name = "DiffusersAdapterPipeline"
                 self.update_multimodal_support()
 
                 # Skip transformer config loading for diffusers adapter
@@ -1370,6 +1381,8 @@ class OmniDiffusionConfig:
             # Skip transformer config loading for diffusers adapter
             # (non-DiT models don't have a separate transformer folder/config)
             if self.diffusion_load_format == "diffusers":
+                if self.model_class_name is None:
+                    self.model_class_name = "DiffusersAdapterPipeline"
                 self.set_tf_model_config(TransformerConfig())
                 logger.warning(
                     "Could not find a valid pipeline index per Diffusers format. "

vllm_omni/diffusion/diffusion_engine.py

@@ -29,6 +29,7 @@ from vllm_omni.diffusion.data import (
     DiffusionOutput,
     DiffusionRequestAbortedError,
     OmniDiffusionConfig,
+    uses_diffusers_adapter,
 )
 from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode, is_scheduler_paged_kv_mode
 from vllm_omni.diffusion.diffusion_kv.initialization import initialize_diffusion_kv_control_plane
@@ -157,8 +158,14 @@ def _resolve_custom_pipeline_cls(custom_pipeline_args: dict[str, Any] | None) ->
 
 def supports_request_batch(od_config: OmniDiffusionConfig) -> bool:
     model_cls = _resolve_custom_pipeline_cls(getattr(od_config, "custom_pipeline_args", None))
-    if model_cls is None:
-        model_cls = DiffusionModelRegistry._try_load_model_cls(getattr(od_config, "model_class_name", None))
+    if model_cls is not None:
+        return bool(getattr(model_cls, "supports_request_batch", False))
+
+    if uses_diffusers_adapter(od_config):
+        model_cls = DiffusionModelRegistry._try_load_model_cls("DiffusersAdapterPipeline")
+        return bool(getattr(model_cls, "supports_request_batch", False))
+
+    model_cls = DiffusionModelRegistry._try_load_model_cls(getattr(od_config, "model_class_name", None))
     if model_cls is None:
         return False
     return bool(getattr(model_cls, "supports_request_batch", False))

vllm_omni/diffusion/models/diffusers_adapter/pipeline_utils.py

@@ -29,7 +29,11 @@ class BasePipelineUtils:
         pass
 
     def validate_runtime_sampling_params(self, sampling: OmniDiffusionSamplingParams) -> None:
-        pass
+        if sampling.enable_frame_interpolation:
+            raise ValueError(
+                "Frame interpolation is not supported with the Diffusers backend because "
+                "Diffusers returns already postprocessed frames."
+            )
 
     def update_call_kwargs(
         self,
@@ -99,6 +103,7 @@ class WanPipelineUtils(BasePipelineUtils):
             )
 
     def validate_runtime_sampling_params(self, sampling: OmniDiffusionSamplingParams) -> None:
+        super().validate_runtime_sampling_params(sampling)
         if sampling.boundary_ratio is not None:
             raise ValueError(
                 "Boundary ratio is not supported at runtime with the diffusers backend for Wan models. Please set "

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 import copy
 import inspect

vllm_omni/diffusion/registry.py

@@ -9,7 +9,7 @@ from vllm.model_executor.model_loader.utils import configure_quant_config
 from vllm.model_executor.models.registry import _LazyRegisteredModel, _ModelRegistry
 
 from vllm_omni.diffusion.config import set_current_diffusion_config
-from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.data import OmniDiffusionConfig, uses_diffusers_adapter
 from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import DistributedVaeMixin
 from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelConfig, get_sp_plan_from_model
 from vllm_omni.diffusion.forward_context import get_forward_context
@@ -725,6 +725,10 @@ def _load_process_func(od_config: OmniDiffusionConfig, func_name: str):
 
 
 def get_diffusion_post_process_func(od_config: OmniDiffusionConfig):
+    # Keep the checkpoint's native class name for modality/capability metadata,
+    # but do not run its tensor postprocessor on Diffusers' decoded outputs.
+    if uses_diffusers_adapter(od_config):
+        return None
     if od_config.model_class_name not in _DIFFUSION_POST_PROCESS_FUNCS:
         return None
     func_name = _DIFFUSION_POST_PROCESS_FUNCS[od_config.model_class_name]
@@ -739,6 +743,9 @@ def get_diffusion_ir_op_priority_func(od_config: OmniDiffusionConfig):
 
 
 def get_diffusion_pre_process_func(od_config: OmniDiffusionConfig):
+    # The adapter translates requests to Diffusers call arguments itself.
+    if uses_diffusers_adapter(od_config):
+        return None
     if od_config.model_class_name not in _DIFFUSION_PRE_PROCESS_FUNCS:
         return None  # Return None if no pre-processing function is registered (for backward compatibility)
     func_name = _DIFFUSION_PRE_PROCESS_FUNCS[od_config.model_class_name]