71 lines
3.3 KiB
Python
71 lines
3.3 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
|
|
from vllm_omni.model_executor.models.qwen3_omni.first_frame_decoder import talker_first_audio_enabled
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("env", "extra", "async_chunk", "v2", "expected"),
|
|
[
|
|
(None, {}, True, True, False),
|
|
("1", {}, True, True, False),
|
|
(None, {"talker_first_audio": True, "codec_chunk_ramp": [1, 2]}, True, True, True),
|
|
("0", {"talker_first_audio": True, "codec_chunk_ramp": [1, 2]}, True, True, False),
|
|
(None, {"talker_first_audio": True, "codec_chunk_ramp": [4, 8]}, True, True, False),
|
|
(None, {"talker_first_audio": True, "initial_codec_chunk_frames": 1}, True, True, True),
|
|
(None, {"talker_first_audio": True, "initial_codec_chunk_frames": 4}, True, True, False),
|
|
(None, {"talker_first_audio": True, "codec_chunk_ramp": [1]}, True, True, False),
|
|
(None, {"talker_first_audio": True, "codec_chunk_ramp": [1, 2]}, False, True, False),
|
|
(None, {"talker_first_audio": True, "codec_chunk_ramp": [1, 2]}, True, False, False),
|
|
],
|
|
)
|
|
def test_talker_first_audio_default_and_conditions(monkeypatch, env, extra, async_chunk, v2, expected):
|
|
if env is None:
|
|
monkeypatch.delenv("VLLM_OMNI_TALKER_FIRST_AUDIO", raising=False)
|
|
else:
|
|
monkeypatch.setenv("VLLM_OMNI_TALKER_FIRST_AUDIO", env)
|
|
from vllm_omni.platforms import current_omni_platform
|
|
|
|
monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: True)
|
|
config = SimpleNamespace(
|
|
model_config=SimpleNamespace(
|
|
async_chunk=async_chunk, use_v2_model_runner=v2, stage_connector_config={"extra": extra}
|
|
),
|
|
device_config=SimpleNamespace(device="cuda"),
|
|
parallel_config=SimpleNamespace(
|
|
tensor_parallel_size=1, pipeline_parallel_size=1, distributed_executor_backend=None
|
|
),
|
|
cache_config=SimpleNamespace(enable_prefix_caching=False),
|
|
)
|
|
assert talker_first_audio_enabled(config) is expected
|
|
for section, field, value in [
|
|
("device_config", "device", "cpu"),
|
|
("parallel_config", "tensor_parallel_size", 2),
|
|
("parallel_config", "pipeline_parallel_size", 2),
|
|
("parallel_config", "distributed_executor_backend", "mp"),
|
|
("cache_config", "enable_prefix_caching", True),
|
|
]:
|
|
with monkeypatch.context() as patch:
|
|
patch.setattr(getattr(config, section), field, value)
|
|
assert not talker_first_audio_enabled(config)
|
|
monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: False)
|
|
assert not talker_first_audio_enabled(config)
|
|
|
|
|
|
@pytest.mark.parametrize("option", [False, True])
|
|
def test_tts_requires_connector_opt_in(monkeypatch, option):
|
|
from vllm_omni.model_executor.models.qwen3_tts import first_audio
|
|
from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code_predictor_vllm import (
|
|
Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM as Predictor,
|
|
)
|
|
|
|
monkeypatch.setattr(Predictor, "_stage_connector_extra_config", lambda _: {"talker_first_audio": option})
|
|
monkeypatch.setattr(first_audio, "supports_talker_first_audio", lambda _: True)
|
|
assert first_audio.talker_first_audio_enabled(None) is option
|