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