[Model] Integrate AURA into Full-Duplex Runtime (#7633)

This commit is contained in:
NumberWan authored and GitHub committed 2026-09-21 09:26:54 -07:00
1 parent d397a2f547
commit 62b15142bd
77 files changed
+7695 -457

No files matched your search

@@ -0,0 +1,165 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""CPU tests for AuraDataPlaneSession (unwrap runner envelope; do not drop audio)."""
from __future__ import annotations
from types import SimpleNamespace
import numpy as np
from vllm_omni.model_executor.models.aura_omni.duplex.data_plane import AuraDataPlaneSession
def _encode_audio(audio: object, sample_rate: int, fmt: str, speed: float | None) -> str | None:
del sample_rate, fmt, speed
if audio is None:
return None
return "encoded-audio"
def test_project_unwraps_data_plane_outputs_and_emits_audio_then_done() -> None:
plane = AuraDataPlaneSession(_encode_audio)
request_id = "duplex-s.abc.e.0.r.stage0-turn1"
plane.begin_request(request_id)
audio = np.zeros(16, dtype=np.float32)
output = SimpleNamespace(
request_id=request_id,
finished=True,
stage_id=3,
outputs=[
SimpleNamespace(
text="",
cumulative_text="",
finished=True,
multimodal_output={"audio": audio, "sr": 24000},
)
],
multimodal_output={"audio": audio, "sr": 24000},
)
events = list(plane.project({"data_plane_outputs": [output]}))
assert events, "project must not silently drop the runner envelope"
assert any(event.get("audio") == "encoded-audio" for event in events)
assert events[-1].get("end_of_turn") is True
assert events[0].get("data_plane_request_id") == request_id
def test_project_sends_only_the_new_tail_of_a_growing_waveform() -> None:
seen: list[int] = []
def encode(audio: object, sample_rate: int, fmt: str, speed: float | None) -> str:
del sample_rate, fmt, speed
count = int(np.asarray(audio).reshape(-1).size)
seen.append(count)
return f"pcm-{count}"
plane = AuraDataPlaneSession(encode)
request_id = "duplex-s.abc.e.0.r.stage3-turn1"
plane.begin_request(request_id)
def chunk(n: int, *, finished: bool) -> SimpleNamespace:
audio = np.arange(n, dtype=np.float32)
return SimpleNamespace(
request_id=request_id,
finished=finished,
stage_id=3,
outputs=[
SimpleNamespace(
text="",
cumulative_text="",
finished=False,
multimodal_output={"audio": audio, "sr": 24000},
)
],
)
first = list(plane.project_output(chunk(8, finished=False)))
second = list(plane.project_output(chunk(12, finished=True)))
assert seen == [8, 4]
assert first[0]["audio_duration_ms"] == round(8 * 1000 / 24000)
assert second[0]["audio_duration_ms"] == round(4 * 1000 / 24000)
assert second[-1].get("end_of_turn") is True
def test_project_holds_silent_token_prefix() -> None:
plane = AuraDataPlaneSession(_encode_audio)
request_id = "duplex-s.abc.e.0.r.stage1-turn1"
plane.begin_request(request_id)
partial = SimpleNamespace(
request_id=request_id,
finished=False,
stage_id=1,
outputs=[SimpleNamespace(text="<|sil", cumulative_text="<|sil", finished=False)],
)
assert list(plane.project({"data_plane_outputs": [partial]})) == []
def test_project_emits_chinese_brackets_immediately() -> None:
"""[沉默] is ordinary text — do not hold as a <|silent|> prefix."""
plane = AuraDataPlaneSession(_encode_audio)
request_id = "duplex-s.abc.e.0.r.stage1_t2"
plane.begin_request(request_id)
partial = SimpleNamespace(
request_id=request_id,
finished=False,
stage_id=1,
outputs=[SimpleNamespace(text="[", cumulative_text="[", finished=False)],
)
events = list(plane.project({"data_plane_outputs": [partial]}))
assert len(events) == 1
assert events[0].get("text") == "["
def test_project_flushes_held_silent_prefix_when_text_becomes_speech() -> None:
plane = AuraDataPlaneSession(_encode_audio)
request_id = "duplex-s.abc.e.0.r.stage1-turn3"
plane.begin_request(request_id)
partial = SimpleNamespace(
request_id=request_id,
finished=False,
stage_id=1,
outputs=[SimpleNamespace(text="<|", cumulative_text="<|", finished=False)],
)
assert list(plane.project({"data_plane_outputs": [partial]})) == []
spoken = SimpleNamespace(
request_id=request_id,
finished=False,
stage_id=1,
outputs=[SimpleNamespace(text="<|hello", cumulative_text="<|hello", finished=False)],
)
events = list(plane.project({"data_plane_outputs": [spoken]}))
assert len(events) == 1
assert events[0].get("text") == "<|hello"
def test_project_marks_finished_silent_text_as_listen() -> None:
plane = AuraDataPlaneSession(_encode_audio)
request_id = "duplex-s.abc.e.0.r.stage1_t4"
plane.begin_request(request_id)
final = SimpleNamespace(
request_id=request_id,
finished=True,
stage_id=1,
outputs=[SimpleNamespace(text="<|silent|>", cumulative_text="<|silent|>", finished=True)],
)
events = list(plane.project({"data_plane_outputs": [final]}))
assert any(e.get("silent") and e.get("is_listen") for e in events)
assert any(e.get("abort_data_plane_request") is True for e in events)
def test_project_finished_spoken_turn_carries_model_context() -> None:
plane = AuraDataPlaneSession(_encode_audio)
request_id = "duplex-s.abc.e.0.r.stage1-turn5"
plane.begin_request(request_id)
final = SimpleNamespace(
request_id=request_id,
finished=True,
stage_id=1,
outputs=[SimpleNamespace(text="你好。", cumulative_text="你好。", finished=True)],
)
events = list(plane.project({"data_plane_outputs": [final]}))
context = [event for event in events if event.get("model_context_text")]
assert context and context[-1]["model_context_text"] == "你好。"
assert context[-1].get("is_listen") is False
assert context[-1].get("text") == ""
@@ -0,0 +1,115 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Aura duplex data-plane must surface Code2Wav MultimodalPayload audio."""
from __future__ import annotations
from types import SimpleNamespace
import torch
from vllm_omni.model_executor.models.aura_omni.duplex.data_plane import (
AuraDataPlaneSession,
_audio_value,
_multimodal,
)
from vllm_omni.outputs.mm_outputs import MultimodalPayload
def _encode(audio: object, sample_rate: int, fmt: str, speed: float | None) -> str | None:
del sample_rate, fmt, speed
if audio is None:
return None
return "encoded-audio"
def test_audio_value_accepts_from_raw_audio_key() -> None:
raw = {"model_outputs": [torch.zeros(8)], "sr": [torch.tensor(24000)]}
mm = MultimodalPayload.from_raw(raw, "audio")
assert mm is not None
assert "audio" in mm
audio = _audio_value(dict(mm))
assert isinstance(audio, torch.Tensor)
assert int(audio.numel()) == 8
def test_audio_value_accepts_wrongly_tagged_text_key() -> None:
# If stage output_modality is wrong, from_raw remaps model_outputs → "text".
raw = {"model_outputs": torch.zeros(4), "sr": torch.tensor(24000)}
mm = MultimodalPayload.from_raw(raw, "text")
assert mm is not None
audio = _audio_value(dict(mm))
assert isinstance(audio, torch.Tensor)
def test_multimodal_skips_empty_outer_dict_and_reads_completion() -> None:
mm = MultimodalPayload.from_dict({"audio": torch.zeros(3), "sr": torch.tensor(24000)})
completion = SimpleNamespace(multimodal_output=mm)
# Empty Mapping on the outer object must not hide completion payload.
outer = SimpleNamespace(multimodal_output={})
got = _multimodal(outer, completion)
assert "audio" in got
def test_project_emits_audio_from_multimodal_completion() -> None:
plane = AuraDataPlaneSession(_encode)
request_id = "duplex-s.test.e.0.r.stage0-turn0"
plane.begin_request(request_id)
mm = MultimodalPayload.from_raw(
{"model_outputs": [torch.zeros(16)], "sr": [torch.tensor(24000)]},
"audio",
)
completion = SimpleNamespace(
text="",
cumulative_text="",
finished=False,
multimodal_output=mm,
)
output = SimpleNamespace(
request_id=request_id,
finished=False,
stage_id=3,
outputs=[completion],
multimodal_output={},
)
events = list(plane.project({"data_plane_outputs": [output]}))
assert events, "project must emit TTS audio for MultimodalPayload"
assert any(event.get("audio") == "encoded-audio" for event in events)
def test_streaming_chunk_completion_finished_is_not_turn_eos() -> None:
"""Code2Wav chunks set CompletionOutput.finished; that must not EOT.
Runner._send_one_model_output_event drops is_terminal request ids before
emit. Marking terminal (or end_of_turn) on the first 8-frame chunk made
WS response.* never leave the engine.
"""
plane = AuraDataPlaneSession(_encode)
request_id = "duplex-s.test.e.0.r.stage0-turn0"
plane.begin_request(request_id)
mm = MultimodalPayload.from_raw(
{"model_outputs": [torch.zeros(16)], "sr": [torch.tensor(24000)]},
"audio",
)
completion = SimpleNamespace(
text="",
cumulative_text="",
finished=True,
multimodal_output=mm,
)
output = SimpleNamespace(
request_id=request_id,
finished=False,
stage_id=3,
outputs=[completion],
multimodal_output={},
)
events: list[dict[str, object]] = []
for event in plane.project({"data_plane_outputs": [output]}):
assert plane.is_terminal(request_id) is False
events.append(event)
assert events
assert events[0].get("audio") == "encoded-audio"
assert events[0].get("end_of_turn") is False
assert plane.is_terminal(request_id) is False
@@ -0,0 +1,598 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""CPU contract tests for AURA duplex plugin seams."""
from __future__ import annotations
import base64
from types import SimpleNamespace
import numpy as np
import pytest
from vllm.sampling_params import SamplingParams
from vllm_omni.engine.duplex.config import DuplexSessionConfig
from vllm_omni.engine.duplex.contracts import (
DuplexFence,
DuplexStageRequestContext,
DuplexStageSubmission,
duplex_ephemeral_stage_request_id,
duplex_resource_request_id,
)
from vllm_omni.engine.duplex.plugin import load_duplex_plugin
from vllm_omni.model_executor.models.aura_omni.duplex.data_plane import (
AuraDataPlaneContext,
AuraDataPlaneSession,
)
from vllm_omni.model_executor.models.aura_omni.duplex.history import (
drop_session_history,
get_or_create_session_history,
)
from vllm_omni.model_executor.models.aura_omni.duplex.input import AuraPcmAppendBuffer
from vllm_omni.model_executor.models.aura_omni.duplex.plugin import (
AURA_SILENT_TOKEN_ID,
AuraDuplexPlugin,
)
from vllm_omni.model_executor.stage_input_processors.aura_omni import SILENT_TEXT
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def _encode_audio(audio: object, sample_rate: int, fmt: str, speed: float | None) -> str | None:
del audio, sample_rate, fmt, speed
return "ZmFrZQ=="
def test_load_aura_duplex_plugin_and_sampling_arity() -> None:
plugin = load_duplex_plugin(
"vllm_omni.model_executor.models.aura_omni.duplex.plugin.AuraDuplexPlugin",
_encode_audio,
)
assert isinstance(plugin, AuraDuplexPlugin)
assert plugin.plugin_id == "aura"
defaults = (
SamplingParams(max_tokens=16),
SamplingParams(max_tokens=16),
SamplingParams(max_tokens=4096),
SamplingParams(max_tokens=16),
)
configured = plugin.configure_sampling_params(runtime_config={}, defaults=defaults)
assert len(configured) == 4
assert configured[2].stop_token_ids == [2150]
assert configured[2].max_tokens == 4096
stage1_stops = set(configured[1].stop_token_ids or [])
assert {151669, 151645}.issubset(stage1_stops)
assert 248070 not in stage1_stops
caps = plugin.capabilities(max_sessions=1)
assert caps.supports_turn_commit_only is True
assert caps.supports_core_resumable_request is False
assert caps.supports_concurrent_turn_requests is True
assert caps.required_input_modalities == frozenset({"video"})
assert caps.optional_input_modalities == frozenset({"audio"})
assert caps.allows_video_without_audio() is True
assert caps.validate_append_modalities(has_audio=True, has_video=False) == (
"This duplex model requires video_frames on input_audio_buffer.append"
)
assert caps.validate_append_modalities(has_audio=True, has_video=True) is None
assert caps.validate_append_modalities(has_audio=False, has_video=True) is None
def test_commit_only_buffer_emits_on_commit() -> None:
buf = AuraPcmAppendBuffer()
samples = np.zeros(1600, dtype="<f4")
payload = {
"type": "audio",
"format": "pcm_f32le",
"sample_rate_hz": 16000,
"audio": base64.b64encode(samples.tobytes()).decode("ascii"),
"is_speech": True,
"video_frames": ["aGVsbG8="],
}
reservation = buf.prepare_append(
payload,
operation_id="op1",
chunk_period_ms=1000,
allow_emit=True,
)
assert reservation is None
assert buf.has_pending()
commit = buf.prepare_commit(operation_id="c1", chunk_period_ms=1000)
assert commit.payload is not None
assert commit.payload.get("final") is True
assert commit.payload.get("aura_turn_commit") is True
assert commit.payload.get("video_frames") == ["aGVsbG8="]
commit.commit()
def test_commit_reservation_includes_video_bytes() -> None:
buf = AuraPcmAppendBuffer()
pcm = np.zeros(1600, dtype="<f4").tobytes()
frames = ["aGVsbG8=", "d29ybGQ="]
payload = {
"type": "audio",
"format": "pcm_f32le",
"sample_rate_hz": 16000,
"audio": base64.b64encode(pcm).decode("ascii"),
"is_speech": True,
"video_frames": frames,
}
assert (
buf.prepare_append(
payload,
operation_id="op1",
chunk_period_ms=1000,
allow_emit=True,
)
is None
)
pending = buf.pending_byte_count
assert pending == len(pcm) + sum(len(frame) for frame in frames)
commit = buf.prepare_commit(operation_id="c1", chunk_period_ms=1000)
assert buf.pending_byte_count == 0
assert commit.byte_count == pending
commit.commit()
assert buf.pending_byte_count == 0
assert not buf.has_reserved()
def test_plan_append_commit_builds_stage0_prompt() -> None:
plugin = AuraDuplexPlugin(_encode_audio)
samples = np.zeros(800, dtype="<f4")
payload = {
"type": "audio",
"format": "pcm_f32le",
"sample_rate_hz": 16000,
"audio": base64.b64encode(samples.tobytes()).decode("ascii"),
"final": True,
"aura_turn_commit": True,
"is_speech": True,
}
fence = DuplexFence("sess-1", epoch=2, turn_id=5)
plan = plugin.plan_append(
request_id="req",
fence=fence,
session_config={},
runtime_config={"aura_system_prompt": "sys"},
seq=1,
turn_seq=1,
payload=payload,
final=True,
sampling_params=SamplingParams(max_tokens=8),
)
info = plan.prompt["additional_information"]
assert info["aura_duplex"] is True
assert info["session_id"] == "sess-1"
assert "audio" in plan.prompt["multi_modal_data"]
assert "prompt_token_ids" not in plan.prompt
def test_video_frames_to_mm_packs_one_video() -> None:
"""Two JPEG frames are one video clip, not two image items."""
from io import BytesIO
from PIL import Image
from vllm_omni.model_executor.models.aura_omni.duplex.plugin import _video_frames_to_mm
frames: list[str] = []
for color in ((255, 0, 0), (0, 255, 0)):
buf = BytesIO()
Image.new("RGB", (8, 8), color).save(buf, format="JPEG")
frames.append(base64.b64encode(buf.getvalue()).decode("ascii"))
mm = _video_frames_to_mm(frames)
assert list(mm.keys()) == ["video"]
assert len(mm["video"]) == 1
video, metadata = mm["video"][0]
assert np.asarray(video).shape == (2, 8, 8, 3)
assert metadata["total_num_frames"] == 2
assert metadata["do_sample_frames"] is False
def test_plan_append_vision_empty_audio_never_leaves_empty_prompt() -> None:
"""Frames + empty/near-silent audio must still get a Stage0 ASR pad."""
from io import BytesIO
from PIL import Image
plugin = AuraDuplexPlugin(_encode_audio)
img = Image.new("RGB", (32, 32), (8, 16, 24))
buf = BytesIO()
img.save(buf, format="JPEG")
frame = base64.b64encode(buf.getvalue()).decode("ascii")
fence = DuplexFence("sess-v", epoch=0, turn_id=1)
near = np.linspace(1e-5, -1e-5, 80, dtype="<f4")
cases = [
("empty_speech_true", True, ""),
("empty_speech_false", False, ""),
("near_silent", False, base64.b64encode(near.tobytes()).decode("ascii")),
]
for name, is_speech, audio in cases:
payload = {
"type": "audio",
"format": "pcm_f32le",
"sample_rate_hz": 16000,
"audio": audio,
"final": True,
"aura_turn_commit": True,
"is_speech": is_speech,
"video_frames": [frame],
}
plan = plugin.plan_append(
request_id="req",
fence=fence,
session_config={},
runtime_config={},
seq=1,
turn_seq=1,
payload=payload,
final=True,
sampling_params=SamplingParams(max_tokens=8),
)
assert plan.prompt.get("prompt"), f"{name}: Stage0 prompt must be non-empty"
assert "audio" in plan.prompt["multi_modal_data"], f"{name}: padded audio required"
assert "deferred_multi_modal_data" in plan.prompt["additional_information"]
def test_project_intermediate_output_targets_stage1_only() -> None:
plugin = AuraDuplexPlugin(_encode_audio)
assert plugin.project_intermediate_output(stage_id=1, output=object(), context=object()) is True
assert plugin.project_intermediate_output(stage_id=0, output=object(), context=object()) is False
assert plugin.project_intermediate_output(stage_id=3, output=object(), context=object()) is False
def test_release_concurrent_turn_requests_on_stage1_final() -> None:
plugin = AuraDuplexPlugin(_encode_audio)
assert plugin.release_concurrent_turn_requests(stage_id=1, segment_finished=True, output=object(), context=object())
assert not plugin.release_concurrent_turn_requests(
stage_id=1, segment_finished=False, output=object(), context=object()
)
finished = type("Out", (), {"finished": True})()
assert plugin.release_concurrent_turn_requests(
stage_id=1, segment_finished=False, output=finished, context=object()
)
assert not plugin.release_concurrent_turn_requests(
stage_id=2, segment_finished=True, output=object(), context=object()
)
def test_configure_sampling_keeps_silent_stop_visible() -> None:
plugin = AuraDuplexPlugin(_encode_audio)
defaults = (
SamplingParams(max_tokens=16),
SamplingParams(max_tokens=16),
SamplingParams(max_tokens=4096),
SamplingParams(max_tokens=16),
)
configured = plugin.configure_sampling_params(runtime_config={}, defaults=defaults)
stage1 = configured[1]
assert isinstance(stage1, SamplingParams)
assert stage1.include_stop_str_in_output is True
assert stage1.skip_special_tokens is False
assert AURA_SILENT_TOKEN_ID in (stage1.stop_token_ids or [])
def test_configure_sampling_keeps_explicit_talker_limits() -> None:
plugin = AuraDuplexPlugin(_encode_audio)
defaults = (
SamplingParams(max_tokens=16),
SamplingParams(max_tokens=16),
SamplingParams(max_tokens=500, stop_token_ids=[7]),
SamplingParams(max_tokens=16),
)
configured = plugin.configure_sampling_params(runtime_config={}, defaults=defaults)
stage2 = configured[2]
assert isinstance(stage2, SamplingParams)
assert stage2.max_tokens == 500
assert stage2.stop_token_ids == [7]
unset = plugin.configure_sampling_params(
runtime_config={},
defaults=(
SamplingParams(max_tokens=16),
SamplingParams(max_tokens=16),
SamplingParams(max_tokens=None),
SamplingParams(max_tokens=16),
),
)
assert unset[2].max_tokens == 240
assert unset[2].stop_token_ids == [2150]
def test_decide_output_silent_short_circuits() -> None:
plugin = AuraDuplexPlugin(_encode_audio)
class _Completion:
text = SILENT_TEXT
token_ids = [AURA_SILENT_TOKEN_ID]
finished = True
class _Output:
outputs = [_Completion()]
decision = plugin.decide_output(
stage_id=1,
final_stage_id=3,
segment_finished=True,
segment_token_ids=(AURA_SILENT_TOKEN_ID,),
segment_output_metadata={},
output=_Output(),
)
assert decision is not None
assert decision.metadata.get("model_listen") is True
def test_decide_output_chinese_silence_is_not_special_token() -> None:
plugin = AuraDuplexPlugin(_encode_audio)
class _Completion:
text = "[沉默]"
token_ids = [58, 107107, 60]
finished = True
class _Output:
outputs = [_Completion()]
assert (
plugin.decide_output(
stage_id=1,
final_stage_id=3,
segment_finished=True,
segment_token_ids=(58, 107107, 60),
segment_output_metadata={},
output=_Output(),
)
is None
)
def test_ephemeral_request_id_includes_turn() -> None:
fence = DuplexFence("s", epoch=1, turn_id=7)
resumable = duplex_resource_request_id(fence, "stage0")
ephemeral = duplex_ephemeral_stage_request_id(fence, stage_id=0)
assert ephemeral.endswith("stage0-turn7")
assert resumable != ephemeral
fence2 = DuplexFence("s", epoch=1, turn_id=8)
assert duplex_ephemeral_stage_request_id(fence, stage_id=0) != duplex_ephemeral_stage_request_id(fence2, stage_id=0)
def test_stage_submission_defaults_resumable_true() -> None:
ctx = DuplexStageRequestContext(
request_id="r",
session_id="s",
fence=DuplexFence("s"),
stage_id=0,
final_stage_id=3,
config_generation=0,
sampling_params=(SamplingParams(max_tokens=1),),
)
submission = DuplexStageSubmission(
context=ctx,
prompt={"prompt_token_ids": [0]},
already_submitted=False,
)
assert submission.resumable is True
def test_data_plane_close_session_drops_history() -> None:
from vllm_omni.model_executor.models.aura_omni.duplex.data_plane import AuraDataPlaneSession
from vllm_omni.model_executor.models.aura_omni.duplex.history import get_or_create_session_history
drop_session_history("close-hist")
get_or_create_session_history("close-hist").begin_user_turn("hi")
AuraDataPlaneSession(encode_audio=_encode_audio).close_session("close-hist")
from vllm_omni.model_executor.models.aura_omni.duplex.history import _STORE
assert "close-hist" not in _STORE
def test_session_history_strips_oldest_silent_videos() -> None:
drop_session_history("hist-ut")
history = get_or_create_session_history("hist-ut")
history.max_video_rounds = 2
history.video_rounds_to_remove = 1
for index in range(3):
history.begin_user_turn("", video=(f"clip-{index}",))
history.commit_turn(SILENT_TEXT)
kept = [message["video"] for message in history.messages if message.get("video") is not None]
assert kept == [("clip-1",), ("clip-2",)]
assert not any(message.get("video") == ("clip-0",) for message in history.messages)
drop_session_history("hist-ut")
def test_pipeline_declares_aura_duplex_plugin() -> None:
from vllm_omni.model_executor.models.aura_omni.pipeline import AURA_OMNI_PIPELINE
assert AURA_OMNI_PIPELINE.duplex_plugin == (
"vllm_omni.model_executor.models.aura_omni.duplex.plugin.AuraDuplexPlugin"
)
def test_project_output_uses_requested_format_and_chunk_duration() -> None:
seen: dict[str, object] = {}
def encode(audio: object, sample_rate: int, fmt: str, speed: float | None) -> str:
seen["fmt"] = fmt
seen["rate"] = sample_rate
seen["speed"] = speed
del audio
return "AAAA"
plane = AuraDataPlaneSession(encode)
class _Output:
request_id = "aura-turn"
finished = False
stage_id = 3
multimodal_output = {"audio": np.zeros(4800, dtype=np.float32), "sr": 24000}
events = list(plane.project_output(_Output(), context=AuraDataPlaneContext(response_format="pcm16", speed=1.0)))
assert seen["fmt"] == "pcm16"
assert events[0]["audio_format"] == "pcm16"
assert events[0]["audio_duration_ms"] == 200
assert events[0]["sample_rate_hz"] == 24000
def encode_wav(audio: object, sample_rate: int, fmt: str, speed: float | None) -> str:
del audio, sample_rate, speed
assert fmt == "wav"
return "V0FW"
plane_wav = AuraDataPlaneSession(encode_wav)
wav_events = list(plane_wav.project_output(_Output(), context=AuraDataPlaneContext(response_format="wav")))
assert wav_events[0]["audio_format"] == "wav"
def test_instructions_update_replaces_effective_prompt() -> None:
plugin = AuraDuplexPlugin(_encode_audio)
current = {"aura_system_prompt": "creation default", "instructions": "creation default"}
updated = plugin.runtime_config_for_update(
DuplexSessionConfig(model="aurateam/AURA", instructions="be brief"),
current,
)
assert updated["aura_system_prompt"] == "be brief"
assert updated["instructions"] == "be brief"
kept = plugin.runtime_config_for_update(
DuplexSessionConfig(
model="aurateam/AURA",
extra_body={"aura_system_prompt": "explicit"},
instructions="ignored when explicit",
),
updated,
)
assert kept["aura_system_prompt"] == "explicit"
def test_plan_partial_stage_output_hands_a_sentence_to_talker() -> None:
def aura2tts() -> None:
return None
class _Pool:
stage_client = SimpleNamespace(custom_process_input_func=aura2tts)
orchestrator = SimpleNamespace(
stage_pools={2: _Pool()},
_stage_receives_async_chunks=lambda stage_id: False,
)
req_state = SimpleNamespace(
session_owned=True,
final_stage_id=3,
request_id="req-1",
prompt={"additional_information": {}},
streaming=SimpleNamespace(bridge_states={}),
)
sentence = "这是一段用来直接测试句子级语音交接而且长度已经超过三十个字的内容。"
output = SimpleNamespace(
finished=False,
request_id="req-1",
outputs=[SimpleNamespace(cumulative_text=sentence)],
)
plan = AuraDuplexPlugin(_encode_audio).plan_partial_stage_output(orchestrator, 1, 0, output, req_state)
assert plan is not None
assert plan.close_only is False
assert plan.is_final_update is False
assert plan.output.text == sentence.strip()
assert req_state.prompt["additional_information"]["aura_tts_partial"] is True
assert AuraDuplexPlugin(_encode_audio).plan_partial_stage_output(orchestrator, 0, 0, output, req_state) is None
def _sentence_plan_fixture() -> tuple[SimpleNamespace, SimpleNamespace]:
def aura2tts() -> None:
return None
class _Pool:
stage_client = SimpleNamespace(custom_process_input_func=aura2tts)
orchestrator = SimpleNamespace(
stage_pools={2: _Pool()},
_stage_receives_async_chunks=lambda stage_id: False,
)
req_state = SimpleNamespace(
session_owned=True,
final_stage_id=3,
request_id="req-1",
prompt={"additional_information": {}},
streaming=SimpleNamespace(bridge_states={}),
)
return orchestrator, req_state
def _stage1_output(text: str, *, finished: bool) -> SimpleNamespace:
return SimpleNamespace(
finished=finished,
request_id="req-1",
outputs=[SimpleNamespace(cumulative_text=text)],
)
def test_finished_remainder_stays_resumable_until_close_sentinel() -> None:
"""A second sentence at Stage1 finish must not be the Talker end sentinel.
The first sentence is already a resumable Talker request. Marking the
remainder ``is_final_update`` makes ``StreamingUpdate.from_request``
return None, which aborts that in-flight sentence before Code2Wav has a
full codec group.
"""
orchestrator, req_state = _sentence_plan_fixture()
plugin = AuraDuplexPlugin(_encode_audio)
first = "我看到你身后是一个很明亮的办公室环境,天花板上能看到白色的管道和通风口,"
full = first + "远处还有办公桌和绿植。"
opened = plugin.plan_partial_stage_output(orchestrator, 1, 0, _stage1_output(first, finished=False), req_state)
assert opened is not None
assert opened.is_final_update is False
assert opened.queue_close_after is False
assert req_state.prompt["additional_information"]["aura_tts_close_only"] is False
tail = plugin.plan_partial_stage_output(orchestrator, 1, 0, _stage1_output(full, finished=True), req_state)
assert tail is not None
assert tail.close_only is False
assert tail.queue_close_after is True
assert tail.is_final_update is False
assert tail.output.text == "远处还有办公桌和绿植。"
assert req_state.prompt["additional_information"]["aura_tts_close_only"] is False
followup = plugin.partial_stage_followup(tail, req_state)
assert followup is not None
assert followup.is_final_update is True
assert followup.close_only is True
assert followup.output.text == ""
assert req_state.prompt["additional_information"]["aura_tts_close_only"] is True
def test_single_finished_sentence_is_one_non_resumable_submit() -> None:
orchestrator, req_state = _sentence_plan_fixture()
plugin = AuraDuplexPlugin(_encode_audio)
sentence = "你好,我看到一个戴眼镜、穿黑T恤的男生正对着镜头说话呢。"
plan = plugin.plan_partial_stage_output(orchestrator, 1, 0, _stage1_output(sentence, finished=True), req_state)
assert plan is not None
assert plan.queue_close_after is False
assert plan.is_final_update is True
assert plan.close_only is False
assert plugin.partial_stage_followup(plan, req_state) is None
def test_aura_draining_stages_are_declared_by_the_plugin() -> None:
plugin = AuraDuplexPlugin(_encode_audio)
assert plugin.draining_stage_ids(stage_count=4) == frozenset({2, 3})
assert plugin.draining_stage_ids(stage_count=2) == frozenset()
def test_spoken_stage0_text_is_the_user_transcript() -> None:
plugin = AuraDuplexPlugin(_encode_audio)
output = SimpleNamespace(
finished=True,
outputs=[
SimpleNamespace(text="language Chinese<asr_text>你好", cumulative_text="language Chinese<asr_text>你好")
],
)
prompt = {"additional_information": {"is_speech": True}}
assert plugin.user_transcript(stage_id=0, output=output, prompt=prompt, finished=True) == "你好"
def test_vision_follow_stage0_text_is_not_shown_as_user_speech() -> None:
plugin = AuraDuplexPlugin(_encode_audio)
output = SimpleNamespace(finished=True, text="嗯")
prompt = {"additional_information": {"is_speech": False}}
assert plugin.user_transcript(stage_id=0, output=output, prompt=prompt, finished=True) is None
assert plugin.user_transcript(stage_id=1, output=output, prompt=prompt, finished=True) is None
assert plugin.user_transcript(stage_id=0, output=output, prompt=prompt, finished=False) is None
@@ -0,0 +1,62 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Growing Stage1 token ids must decode to the same text as one full decode."""
from __future__ import annotations
from vllm_omni.model_executor.models.aura_omni.duplex.sentence_tts import decode_growing_token_ids
def _assert_matches_full(decode, ids: list[int]) -> None:
cache: dict[str, object] = {}
for end in range(1, len(ids) + 1):
got = decode_growing_token_ids(decode, ids[:end], cache)
assert got == decode(ids[:end])
def test_growing_decode_matches_full_decode_for_concat_pieces() -> None:
vocab = {1: "你", 2: "好", 3: "。", 4: "世", 5: "界"}
def decode(token_ids: list[int], skip_special_tokens: bool = False) -> str:
del skip_special_tokens
return "".join(vocab[token_id] for token_id in token_ids)
_assert_matches_full(decode, [1, 2, 3, 4, 5])
def test_growing_decode_matches_full_decode_when_the_previous_piece_changes() -> None:
"""Token 9 rewrites the previous character. The tail window must reject a bad splice."""
def decode(token_ids: list[int], skip_special_tokens: bool = False) -> str:
del skip_special_tokens
chars: list[str] = []
for token_id in token_ids:
if token_id == 9:
if chars:
chars.pop()
chars.append("X")
else:
chars.append(chr(96 + token_id))
return "".join(chars)
_assert_matches_full(decode, [1, 2, 9, 3])
def test_growing_decode_matches_wordpiece_spacing() -> None:
vocab = {1: "hello", 2: "##world", 3: "。"}
def decode(token_ids: list[int], skip_special_tokens: bool = False) -> str:
del skip_special_tokens
parts: list[str] = []
for index, token_id in enumerate(token_ids):
piece = vocab[token_id]
if piece.startswith("##"):
parts.append(piece[2:])
else:
if index:
parts.append(" ")
parts.append(piece)
return "".join(parts)
_assert_matches_full(decode, [1, 2, 3, 1])
@@ -389,6 +389,38 @@ def test_decode_batch_preprocess_matches_decode_state_updates(mrv2):
assert updates[1]["hidden_states"]["trailing_text"].numel() == 0
def test_decode_batch_zero_fills_missing_last_hidden(monkeypatch):
tts_pad = torch.full((1, 4), -1.0, dtype=torch.bfloat16)
model = _make_minimal_talker(tts_pad_embed=tts_pad)
def fake_embed_input_ids(input_ids):
return input_ids.to(torch.float32).reshape(-1, 1, 1).expand(-1, 1, 4)
model.embed_input_ids = fake_embed_input_ids
warnings: list[tuple[object, ...]] = []
monkeypatch.setattr(
"vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker.logger.warning_once",
lambda *args, **kwargs: warnings.append(args),
)
trailing = torch.arange(8, dtype=torch.float32).reshape(2, 4)
_, _, past_hidden, _, _ = model.preprocess_decode_batch(
input_ids=torch.tensor([101], dtype=torch.long),
req_infos=[
{
"text": ["hello"],
"task_type": ["CustomVoice"],
"hidden_states": {"trailing_text": trailing},
"meta": {"talker_text_offset": 0},
}
],
)
assert torch.equal(past_hidden.cpu(), torch.zeros((1, 4), dtype=torch.bfloat16))
assert len(warnings) == 1
assert "zero-filling" in str(warnings[0][0])
def _stub_text_embedding(device_param: torch.nn.Parameter):
"""Build a lambda that emulates ``nn.Embedding`` for the ``_device()`` helper.