7b4fc7ce2e
Signed-off-by: levius <2114377220@qq.com> Signed-off-by: Yueqian Lin <linyueqian@outlook.com> Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com> Co-authored-by: Yueqian Lin <linyueqian@outlook.com>
61 lines
2.4 KiB
Python
61 lines
2.4 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
import pytest
|
|
|
|
from tests.helpers.stage_config import get_deploy_config_path
|
|
from vllm_omni.config.endpoint_policy import OmniServingCapability
|
|
from vllm_omni.config.stage_config import StageExecutionType, load_deploy_config
|
|
from vllm_omni.diffusion.registry import _DIFFUSION_MODELS
|
|
from vllm_omni.model_executor.models.mammoth_moda2.pipeline import (
|
|
MAMMOTH_MODA2_AR_PIPELINE,
|
|
MAMMOTH_MODA2_PIPELINE,
|
|
)
|
|
from vllm_omni.model_executor.models.registry import _OMNI_MODELS
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
|
|
def test_mammothmoda2_generation_stage_uses_shared_diffusion_runtime() -> None:
|
|
ar_stage, dit_stage = MAMMOTH_MODA2_PIPELINE.stages
|
|
assert ar_stage.execution_type is StageExecutionType.LLM_AR
|
|
assert dit_stage.execution_type is StageExecutionType.DIFFUSION
|
|
assert dit_stage.model_arch == "MammothModa2DiTPipeline"
|
|
assert dit_stage.custom_process_input_func.endswith(".ar2diffusion")
|
|
assert dit_stage.omni_kv_config == {"need_recv_cache": False}
|
|
assert dit_stage.input_sources == (0,)
|
|
assert dit_stage.final_output is True
|
|
assert dit_stage.final_output_type == "image"
|
|
|
|
|
|
def test_mammothmoda2_dit_is_registered_only_with_diffusion_runtime() -> None:
|
|
assert _DIFFUSION_MODELS["MammothModa2DiTPipeline"] == (
|
|
"mammoth_moda2",
|
|
"pipeline_mammothmoda2_dit",
|
|
"MammothModa2DiTPipeline",
|
|
)
|
|
assert "MammothModa2DiTPipeline" not in _OMNI_MODELS
|
|
|
|
|
|
def test_mammothmoda2_deploy_does_not_shadow_request_sampling_fields() -> None:
|
|
deploy = load_deploy_config(get_deploy_config_path("mammoth_moda2.yaml"))
|
|
|
|
assert deploy.stages[1].default_sampling_params is None
|
|
|
|
|
|
def test_mammothmoda2_rejects_unsupported_image_edits() -> None:
|
|
restrictions = {
|
|
restriction.capability: restriction.reason for restriction in MAMMOTH_MODA2_PIPELINE.endpoint_restrictions
|
|
}
|
|
|
|
assert set(restrictions) == {OmniServingCapability.IMAGE_EDITS}
|
|
assert "img2img" in restrictions[OmniServingCapability.IMAGE_EDITS]
|
|
assert MAMMOTH_MODA2_AR_PIPELINE.endpoint_restrictions == ()
|
|
|
|
|
|
def test_mammothmoda2_ar_only_topology_is_unchanged() -> None:
|
|
assert len(MAMMOTH_MODA2_AR_PIPELINE.stages) == 1
|
|
stage = MAMMOTH_MODA2_AR_PIPELINE.stages[0]
|
|
assert stage.execution_type is StageExecutionType.LLM_AR
|
|
assert stage.final_output_type == "text"
|