diff --git a/tests/diffusion/diffusion_backend/test_diffusers_backend.py b/tests/diffusion/diffusion_backend/test_diffusers_backend.py index ec351d746..3b094c37f 100644 --- a/tests/diffusion/diffusion_backend/test_diffusers_backend.py +++ b/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) diff --git a/tests/diffusion/models/qwen_image/test_qwen_image_postprocess.py b/tests/diffusion/models/qwen_image/test_qwen_image_postprocess.py index 62d17f705..23f0bd467 100644 --- a/tests/diffusion/models/qwen_image/test_qwen_image_postprocess.py +++ b/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 diff --git a/tests/diffusion/test_diffusion_engine.py b/tests/diffusion/test_diffusion_engine.py index a0c6acb02..04d528af4 100644 --- a/tests/diffusion/test_diffusion_engine.py +++ b/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", diff --git a/tests/diffusion/test_diffusion_output_formatter.py b/tests/diffusion/test_diffusion_output_formatter.py index 28714ac97..4578d8368 100644 --- a/tests/diffusion/test_diffusion_output_formatter.py +++ b/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: diff --git a/tests/diffusion/test_diffusion_plugin_hooks.py b/tests/diffusion/test_diffusion_plugin_hooks.py index 1885d3338..65ad9ccc7 100644 --- a/tests/diffusion/test_diffusion_plugin_hooks.py +++ b/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.""" diff --git a/tests/entrypoints/test_async_omni_diffusion_config.py b/tests/entrypoints/test_async_omni_diffusion_config.py index 641729601..3864d1d67 100644 --- a/tests/entrypoints/test_async_omni_diffusion_config.py +++ b/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( diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 590afcd8c..f29c59c34 100644 --- a/vllm_omni/diffusion/data.py +++ b/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. " diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index 800b57e34..0929e9419 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/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)) diff --git a/vllm_omni/diffusion/models/diffusers_adapter/pipeline_utils.py b/vllm_omni/diffusion/models/diffusers_adapter/pipeline_utils.py index dc5ea2038..2d00881b8 100644 --- a/vllm_omni/diffusion/models/diffusers_adapter/pipeline_utils.py +++ b/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 " diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py index 70e0ef0f1..0aa00abd9 100644 --- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py +++ b/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 diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index 8a707711a..ed6a816a9 100644 --- a/vllm_omni/diffusion/registry.py +++ b/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]