97 lines
3.6 KiB
Python
97 lines
3.6 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
"""L1 tests for checkpoint-to-pipeline routing of Boogu-Image Turbo."""
|
|
|
|
import json
|
|
|
|
import pytest
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
|
|
|
|
|
def test_turbo_registry_entry_loads_the_turbo_class():
|
|
from vllm_omni.diffusion.models.boogu_image import BooguImageTurboPipeline
|
|
from vllm_omni.diffusion.registry import DiffusionModelRegistry
|
|
|
|
assert DiffusionModelRegistry._try_load_model_cls("BooguImageTurboPipeline") is BooguImageTurboPipeline
|
|
|
|
|
|
def test_turbo_entry_disables_request_batching():
|
|
from vllm_omni.diffusion.models.boogu_image import BooguImagePipeline, BooguImageTurboPipeline
|
|
|
|
assert BooguImagePipeline.supports_request_batch
|
|
assert not BooguImageTurboPipeline.supports_request_batch
|
|
|
|
|
|
def test_turbo_process_funcs_resolve_to_the_shared_boogu_helpers():
|
|
from vllm_omni.diffusion.models import boogu_image
|
|
from vllm_omni.diffusion.registry import _DIFFUSION_POST_PROCESS_FUNCS, _DIFFUSION_PRE_PROCESS_FUNCS
|
|
|
|
module = boogu_image.pipeline_boogu_image
|
|
for funcs in (_DIFFUSION_PRE_PROCESS_FUNCS, _DIFFUSION_POST_PROCESS_FUNCS):
|
|
name = funcs["BooguImageTurboPipeline"]
|
|
assert name == funcs["BooguImagePipeline"]
|
|
assert callable(getattr(module, name))
|
|
|
|
|
|
def test_turbo_inherits_the_base_multimodal_metadata():
|
|
from vllm_omni.diffusion.model_metadata import get_diffusion_model_metadata
|
|
|
|
metadata = get_diffusion_model_metadata("BooguImageTurboPipeline")
|
|
|
|
assert metadata == get_diffusion_model_metadata("BooguImagePipeline")
|
|
assert metadata.supports_multimodal_inputs
|
|
assert metadata.max_multimodal_image_inputs == 1
|
|
|
|
|
|
def test_turbo_model_index_resolves_and_enriches(tmp_path):
|
|
from vllm_omni.diffusion.data import OmniDiffusionConfig, resolve_model_class_name
|
|
|
|
(tmp_path / "model_index.json").write_text(
|
|
json.dumps({"_class_name": "BooguImageTurboPipeline"}),
|
|
encoding="utf-8",
|
|
)
|
|
|
|
assert resolve_model_class_name(str(tmp_path)) == "BooguImageTurboPipeline"
|
|
|
|
config = OmniDiffusionConfig(model=str(tmp_path))
|
|
config.enrich_config()
|
|
|
|
assert config.model_class_name == "BooguImageTurboPipeline"
|
|
assert config.supports_multimodal_inputs
|
|
assert config.max_multimodal_image_inputs == 1
|
|
|
|
|
|
def test_edit_turbo_override_selects_the_turbo_class(tmp_path):
|
|
"""Edit-Turbo publishes ``BooguImagePipeline``; ``--model-class-name`` must win."""
|
|
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
|
from vllm_omni.diffusion.registry import DiffusionModelRegistry
|
|
|
|
(tmp_path / "model_index.json").write_text(
|
|
json.dumps({"_class_name": "BooguImagePipeline"}),
|
|
encoding="utf-8",
|
|
)
|
|
|
|
config = OmniDiffusionConfig(model=str(tmp_path), model_class_name="BooguImageTurboPipeline")
|
|
config.enrich_config()
|
|
|
|
assert config.model_class_name == "BooguImageTurboPipeline"
|
|
assert DiffusionModelRegistry._try_load_model_cls(config.model_class_name)._is_turbo
|
|
|
|
|
|
def test_edit_turbo_without_override_keeps_the_base_class(tmp_path):
|
|
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
|
from vllm_omni.diffusion.registry import DiffusionModelRegistry
|
|
|
|
(tmp_path / "model_index.json").write_text(
|
|
json.dumps({"_class_name": "BooguImagePipeline"}),
|
|
encoding="utf-8",
|
|
)
|
|
|
|
config = OmniDiffusionConfig(model=str(tmp_path))
|
|
config.enrich_config()
|
|
|
|
assert config.model_class_name == "BooguImagePipeline"
|
|
assert not DiffusionModelRegistry._try_load_model_cls(config.model_class_name)._is_turbo
|