[Core] Graduate MiniCPM-o 4.5 and PersonaPlex full-duplex serving out of experimental (#6196)

Signed-off-by: chickeyton <ngton2014@gmail.com>
Signed-off-by: Gao Han <hgaoaf@connect.ust.hk>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>
This commit is contained in:
authored and GitHub committed 2026-09-05 08:20:21 +00:00
1 parent d4096e01fe
commit 5d7fdf9350
179 files changed
+6462 -3832

No files matched your search

Whitespace-only changes.
@@ -0,0 +1,102 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations
import pytest
from vllm_omni.entrypoints.duplex.commit_policy import (
CommitAction,
CommitSnapshot,
decide_commit_action,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@pytest.mark.parametrize(
("case", "snapshot", "expected"),
[
(
"first idle commit",
CommitSnapshot(
auto_responds=True,
speech_since_commit=True,
active_response_id=None,
overlap_speech_ms=0,
native_response_in_progress=False,
playback_active=False,
),
CommitAction.START_AUTO_RESPONSE,
),
(
"native auto-response remains model-owned during playback",
CommitSnapshot(
auto_responds=True,
speech_since_commit=True,
active_response_id=None,
overlap_speech_ms=800,
native_response_in_progress=True,
playback_active=True,
),
CommitAction.START_AUTO_RESPONSE,
),
(
"native auto-response commit without speech",
CommitSnapshot(
auto_responds=True,
speech_since_commit=False,
active_response_id=None,
overlap_speech_ms=0,
native_response_in_progress=True,
playback_active=False,
),
CommitAction.COMMIT_ONLY,
),
(
"non-auto response overlap",
CommitSnapshot(
auto_responds=False,
speech_since_commit=True,
active_response_id=None,
overlap_speech_ms=800,
native_response_in_progress=True,
playback_active=False,
),
CommitAction.DEFER_ACTIVE_RESPONSE,
),
(
"non-auto playback overlap",
CommitSnapshot(
auto_responds=False,
speech_since_commit=True,
active_response_id=None,
overlap_speech_ms=800,
native_response_in_progress=False,
playback_active=True,
),
CommitAction.DEFER_ACTIVE_RESPONSE,
),
(
"non-auto commit without overlap",
CommitSnapshot(
auto_responds=False,
speech_since_commit=True,
active_response_id=None,
overlap_speech_ms=0,
native_response_in_progress=False,
playback_active=False,
),
CommitAction.COMMIT_ONLY,
),
],
ids=lambda value: value if isinstance(value, str) else None,
)
def test_decide_commit_action_covers_realtime_lifecycle_sequences(
case: str,
snapshot: CommitSnapshot,
expected: CommitAction,
) -> None:
del case
assert decide_commit_action(snapshot) is expected
@@ -0,0 +1,179 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import base64
import numpy as np
import pytest
from vllm_omni.model_executor.models.minicpmo_4_5.duplex.input import MiniCPMO45PcmAppendBuffer
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def pcm_payload(samples: int, *, speech: bool = True) -> dict[str, object]:
audio = np.ones(samples, dtype=np.float32).tobytes()
return {
"type": "audio",
"audio": base64.b64encode(audio).decode("ascii"),
"format": "pcm_f32le",
"sample_rate_hz": 16_000,
"is_speech": speech,
}
def test_commit_does_not_add_silence_after_incremental_audio_was_drained():
buffer = MiniCPMO45PcmAppendBuffer()
emitted = buffer.append(pcm_payload(16_000), chunk_period_ms=1_000)
committed = buffer.commit(chunk_period_ms=1_000)
assert emitted is not None
assert not buffer.has_pending()
assert committed is None
def test_append_emits_one_model_unit_when_multiple_units_are_buffered():
buffer = MiniCPMO45PcmAppendBuffer()
emitted = buffer.append(pcm_payload(32_000), chunk_period_ms=1_000)
assert emitted is not None
assert len(base64.b64decode(emitted["audio"])) == 16_000 * 4
assert buffer.pending_byte_count == 16_000 * 4
def test_speech_marker_does_not_leak_across_irregular_chunk_boundaries():
buffer = MiniCPMO45PcmAppendBuffer()
assert buffer.append(pcm_payload(15_000), chunk_period_ms=1_000) is None
speech_unit = buffer.append(pcm_payload(2_000, speech=False), chunk_period_ms=1_000)
silence_unit = buffer.append(pcm_payload(15_000, speech=False), chunk_period_ms=1_000)
assert speech_unit is not None
assert speech_unit["is_speech"] is True
assert silence_unit is not None
assert silence_unit["is_speech"] is False
def test_force_listen_marker_does_not_leak_across_irregular_chunk_boundaries():
buffer = MiniCPMO45PcmAppendBuffer()
forced = pcm_payload(15_000, speech=False)
forced["force_listen"] = True
assert buffer.append(forced, chunk_period_ms=1_000) is None
forced_unit = buffer.append(pcm_payload(2_000, speech=False), chunk_period_ms=1_000)
unforced_unit = buffer.append(pcm_payload(15_000, speech=False), chunk_period_ms=1_000)
assert forced_unit is not None
assert forced_unit["force_listen"] is True
assert unforced_unit is not None
assert unforced_unit["force_listen"] is False
def test_pcm_append_rollback_restores_per_span_speech_markers():
buffer = MiniCPMO45PcmAppendBuffer()
assert buffer.append(pcm_payload(15_000), chunk_period_ms=1_000) is None
reservation = buffer.prepare_append(
pcm_payload(2_000, speech=False),
operation_id="mixed-unit",
chunk_period_ms=1_000,
)
assert reservation is not None
reservation.rollback()
first = buffer.append(pcm_payload(15_000, speech=False), chunk_period_ms=1_000)
second = buffer.flush(chunk_period_ms=1_000)
assert first is not None
assert first["is_speech"] is True
assert second is not None
assert second["is_speech"] is False
def test_commit_without_speech_does_not_synthesize_terminal_audio():
buffer = MiniCPMO45PcmAppendBuffer()
buffer.append(pcm_payload(8_000, speech=False), chunk_period_ms=1_000)
committed = buffer.commit(chunk_period_ms=1_000)
assert committed is None
def test_commit_resets_cumulative_turn_accounting():
buffer = MiniCPMO45PcmAppendBuffer()
buffer.append(pcm_payload(16_000), chunk_period_ms=1_000)
buffer.commit(chunk_period_ms=1_000)
empty = buffer.commit(chunk_period_ms=1_000)
assert empty is None
def test_pcm_append_reservation_rollback_restores_emitted_audio():
buffer = MiniCPMO45PcmAppendBuffer()
original = pcm_payload(16_000)
reservation = buffer.prepare_append(
original,
operation_id="append-1",
chunk_period_ms=1_000,
)
assert reservation is not None
assert not buffer.has_pending()
reservation.rollback()
assert buffer.has_pending()
retried = buffer.flush(chunk_period_ms=1_000)
assert retried is not None
assert base64.b64decode(retried["audio"]) == base64.b64decode(original["audio"])
def test_pcm_commit_keeps_prior_append_reservation_active():
buffer = MiniCPMO45PcmAppendBuffer()
append_reservation = buffer.prepare_append(
pcm_payload(16_000),
operation_id="append-before-commit",
chunk_period_ms=1_000,
)
commit_reservation = buffer.prepare_commit(
operation_id="commit-after-append",
chunk_period_ms=1_000,
)
assert append_reservation is not None
assert append_reservation.active
assert commit_reservation.payload is None
append_reservation.commit()
commit_reservation.commit()
def test_pcm_commit_reservation_rollback_restores_residual_audio():
buffer = MiniCPMO45PcmAppendBuffer()
original = pcm_payload(8_000)
assert (
buffer.prepare_append(
original,
operation_id="buffer-half-chunk",
chunk_period_ms=1_000,
)
is None
)
reservation = buffer.prepare_commit(
operation_id="final-half-chunk",
chunk_period_ms=1_000,
)
assert reservation.payload is not None
assert reservation.payload["final"] is True
reservation.rollback()
retried = buffer.prepare_commit(
operation_id="retry-final-half-chunk",
chunk_period_ms=1_000,
)
assert retried.payload is not None
assert base64.b64decode(retried.payload["audio"]) == (base64.b64decode(original["audio"]) + b"\x00" * (8_000 * 4))
@@ -0,0 +1,22 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import pytest
from vllm_omni.model_executor.models.minicpmo_4_5.duplex.policy import (
MiniCPMO45DuplexPolicy,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def test_seed_tts_initial_text_is_part_of_native_duplex_context():
prefix, suffix = MiniCPMO45DuplexPolicy.session_context_texts(
"Speak exactly.",
True,
"The quick brown fox.",
)
assert prefix == "<|im_start|>system\nSpeak exactly.\n<|audio_start|>"
assert suffix == (
"<|audio_end|><|im_end|>\n<|im_start|>user\nThe quick brown fox.<|im_end|>\n<|im_start|>assistant\n"
)
@@ -61,7 +61,7 @@ class TestRegistryDeclaration:
pipeline = OMNI_PIPELINES[_PIPELINE_KEY]
assert pipeline.duplex_control_enabled is True
assert pipeline.duplex_serving_adapter == (
"vllm_omni.experimental.fullduplex.minicpmo45.serving_adapter.MiniCPMO45ServingRuntimeAdapter"
"vllm_omni.model_executor.models.minicpmo_4_5.duplex.serving_adapter.MiniCPMO45ServingRuntimeAdapter"
)
assert not hasattr(pipeline, "max_native_duplex_sessions")
@@ -0,0 +1,2 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
@@ -0,0 +1,2 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
@@ -0,0 +1,97 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations
from types import SimpleNamespace
import numpy as np
import pytest
from vllm_omni.engine.duplex.messages import DuplexFence
from vllm_omni.engine.duplex.runtime import duplex_resource_request_id
from vllm_omni.model_executor.models.nemotron_voicechat.duplex.data_plane import (
NemotronVoiceChatDataPlaneContext,
NemotronVoiceChatDataPlaneSession,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
_RUNTIME = {
"nvc_text_bos_id": 0,
"nvc_text_eos_id": 1,
"nvc_text_pad_id": 12,
"nvc_function_sotc_id": 20,
"nvc_function_eotc_id": 21,
"nvc_tokenizer_ref": "test-tokenizer",
}
def _projector(encode_audio=lambda *_: None) -> NemotronVoiceChatDataPlaneSession:
projector = NemotronVoiceChatDataPlaneSession(encode_audio)
projector.configure_runtime(_RUNTIME, tokenizer=SimpleNamespace(decode=lambda *_args, **_kwargs: "text"))
return projector
def _stage0_output(request_id: str = "req-0", *, text_token: int = 12, function_token: int | None = None) -> object:
metadata: dict[str, object] = {"nvc_text_token_ids": [text_token]}
if function_token is not None:
metadata["nvc_function_token"] = [function_token]
completion = SimpleNamespace(multimodal_output=metadata)
return SimpleNamespace(
stage_id=0,
request_output=SimpleNamespace(request_id=request_id, outputs=[completion]),
)
def test_function_channel_projects_completed_call_without_ending_speech() -> None:
projector = _projector()
projector._decode = lambda token_ids: (
'[{"name":"weather","arguments":{"city":"Shanghai"}}]' if token_ids == [99] else ""
)
context = NemotronVoiceChatDataPlaneContext(epoch=0)
events = []
for function_token in (20, 99, 21):
events.extend(projector.project_output(_stage0_output(function_token=function_token), context=context))
listen = [event for event in events if event.get("is_listen") is True]
function = [event for event in events if event.get("function_call") is True]
assert len(listen) == 3
assert len(function) == 1
assert function[0]["name"] == "weather"
assert function[0]["arguments"] == '{"city":"Shanghai"}'
projector._decode = lambda _token_ids: "not-json"
bad_events = []
for function_token in (20, 98, 21):
bad_events.extend(projector.project_output(_stage0_output(function_token=function_token)))
assert any(event.get("error_code") == "nemotron_function_call_parse_error" for event in bad_events)
def test_new_epoch_request_does_not_inherit_partial_function_or_frame_state() -> None:
projector = _projector(lambda *_: "audio")
old_request = duplex_resource_request_id(DuplexFence("sid", epoch=1), "stage0")
new_request = duplex_resource_request_id(DuplexFence("sid", epoch=2), "stage0")
projector.begin_request(old_request)
list(projector.project_output(_stage0_output(old_request, text_token=42, function_token=20)))
projector.begin_request(new_request)
# EOTC in the new epoch must not close the old epoch's partial function.
new_events = list(projector.project_output(_stage0_output(new_request, text_token=1, function_token=21)))
assert not [event for event in new_events if event.get("function_call") is True]
assert not [event for event in new_events if event.get("end_of_turn") is True]
audio = SimpleNamespace(
stage_id=2,
request_output=SimpleNamespace(
request_id=new_request,
outputs=[
SimpleNamespace(multimodal_output={"model_outputs": [np.ones(1764, dtype=np.float32)], "sr": [22050]})
],
),
)
audio_events = list(projector.project_output(audio))
assert sum(event.get("end_of_turn") is True for event in audio_events) == 1
@@ -0,0 +1,76 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations
import base64
import numpy as np
import pytest
from vllm_omni.model_executor.models.nemotron_voicechat.duplex.input import (
NemotronVoiceChatPcmAppendBuffer,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def _payload(samples: np.ndarray, *, sample_rate_hz: int = 16000, fmt: str = "pcm_f32le") -> dict[str, object]:
return {
"type": "audio",
"audio": base64.b64encode(np.asarray(samples, dtype=np.float32).tobytes()).decode("ascii"),
"format": fmt,
"sample_rate_hz": sample_rate_hz,
}
def _append(buffer, samples, operation_id):
return buffer.prepare_append(
_payload(np.asarray(samples, dtype=np.float32)),
operation_id=operation_id,
chunk_period_ms=80,
allow_emit=True,
)
def test_irregular_browser_packets_emit_exact_ordered_80ms_frame() -> None:
buffer = NemotronVoiceChatPcmAppendBuffer()
assert _append(buffer, np.arange(1000), "packet-1") is None
reservation = _append(buffer, np.arange(1000, 1280), "packet-2")
assert reservation is not None
assert reservation.operation_id == "packet-2"
decoded = np.frombuffer(base64.b64decode(reservation.payload["audio"]), dtype=np.float32)
np.testing.assert_array_equal(decoded, np.arange(1280, dtype=np.float32))
assert buffer.pending_byte_count == 0
def test_append_rejects_invalid_frame_contract() -> None:
cases = (
(1281, 16000, "pcm_f32le", 0.0, "at most 1280"),
(1, 8000, "pcm_f32le", 0.0, "sample_rate_hz"),
(1, 16000, "audio/pcm", 0.0, "format"),
(1, 16000, "pcm_f32le", np.nan, "finite"),
)
for samples, sample_rate_hz, fmt, value, message in cases:
with pytest.raises(ValueError, match=message):
NemotronVoiceChatPcmAppendBuffer().prepare_append(
_payload(np.full(samples, value, dtype=np.float32), sample_rate_hz=sample_rate_hz, fmt=fmt),
operation_id="invalid",
chunk_period_ms=80,
allow_emit=True,
)
def test_terminal_commit_rejects_more_than_one_frame_tail() -> None:
buffer = NemotronVoiceChatPcmAppendBuffer()
for index, samples in enumerate((1000, 281)):
buffer.prepare_append(
_payload(np.zeros(samples, dtype=np.float32)),
operation_id=f"tail-{index}",
chunk_period_ms=80,
allow_emit=False,
)
with pytest.raises(ValueError, match="more than one 80 ms frame"):
buffer.prepare_commit(operation_id="commit", chunk_period_ms=80)
@@ -0,0 +1,115 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
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.contracts import DuplexInputMode
from vllm_omni.engine.duplex.messages import DuplexFence
from vllm_omni.model_executor.models.nemotron_voicechat.duplex.runtime import (
NemotronVoiceChatDuplexRuntimeExtension,
)
from vllm_omni.model_executor.models.nemotron_voicechat.duplex.serving_adapter import (
NemotronVoiceChatServingRuntimeAdapter,
_render_tool_response,
)
from vllm_omni.model_executor.models.nemotron_voicechat.nemotron_voicechat_thinker import (
NemotronVoiceChatThinkerForConditionalGeneration,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def _frame() -> dict[str, object]:
raw = np.zeros(1280, dtype=np.float32).tobytes()
return {
"type": "audio",
"audio": base64.b64encode(raw).decode("ascii"),
"format": "pcm_f32le",
"sample_rate_hz": 16000,
}
def _plan(extension, *, input_seq: int):
return extension.plan_append(
request_id="req",
fence=DuplexFence("sid", incarnation=2, epoch=3),
session_config={},
runtime_config={
"nvc_prompt_token_ids": [0, 42, 1],
"nvc_text_pad_id": 12,
"nvc_max_model_len": 8192,
},
seq=input_seq,
turn_seq=input_seq,
mode=DuplexInputMode.APPEND_AUDIO_CHUNK,
payload=_frame(),
final=False,
sampling_params=SamplingParams(),
)
def test_first_append_prefills_prompt_then_each_append_consumes_one_frame() -> None:
extension = NemotronVoiceChatDuplexRuntimeExtension()
params = extension.configure_sampling_params(
runtime_config={}, defaults=(SamplingParams(temperature=0.7, max_tokens=4),)
)[0]
assert params.temperature == 0.0 and params.max_tokens == 1 and params.top_p == 1.0 and params.top_k == 0
first = _plan(extension, input_seq=1)
later = _plan(extension, input_seq=2)
assert first.prompt["prompt_token_ids"] == [0, 42, 1, 12]
assert later.prompt["prompt_token_ids"] == [12]
assert first.prompt["model_intermediate_buffer"]["duplex"]["source_input_seq"] == 1
assert later.prompt["model_intermediate_buffer"]["duplex"]["source_input_seq"] == 2
def test_append_rejects_stage0_context_overflow() -> None:
with pytest.raises(ValueError, match="max_model_len"):
_plan(NemotronVoiceChatDuplexRuntimeExtension(), input_seq=8190)
def test_function_output_becomes_versioned_nvidia_channel_tokens() -> None:
encoded: list[str] = []
adapter = NemotronVoiceChatServingRuntimeAdapter(lambda *_: None)
adapter._tokenizer = SimpleNamespace(
encode=lambda text, **_kwargs: encoded.append(text) or [31, 32, 33],
)
first = adapter.runtime_config_for_function_output(
{"type": "function_call_output", "call_id": "call-1", "output": '{"result":20}'},
{},
)
second = adapter.runtime_config_for_function_output(
{"type": "function_call_output", "call_id": "call-2", "output": "plain text"},
first,
)
assert _render_tool_response('{"result":20}') == '<TOOL_RESPONSE>[{"result":20}]</TOOL_RESPONSE>'
assert encoded == [
'<TOOL_RESPONSE>[{"result":20}]</TOOL_RESPONSE>',
'<TOOL_RESPONSE>["plain text"]</TOOL_RESPONSE>',
]
assert first["nvc_function_response_generation"] == 1
assert second["nvc_function_response_generation"] == 2
assert second["nvc_function_response_token_ids"] == [31, 32, 33]
assert second["nvc_function_response_call_id"] == "call-2"
assert second["nvc_function_response_batches"] == [
{"generation": 1, "call_id": "call-1", "token_ids": [31, 32, 33]},
{"generation": 2, "call_id": "call-2", "token_ids": [31, 32, 33]},
]
session = {"function_response_generation": 0, "forced_function_tokens": []}
NemotronVoiceChatThinkerForConditionalGeneration._sync_forced_function_response(session, second)
assert session["function_response_generation"] == 2
assert session["forced_function_token"] == 31
assert session["forced_function_tokens"] == [31, 32, 33, 31, 32, 33]
@@ -0,0 +1,2 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
@@ -0,0 +1,2 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
@@ -0,0 +1,303 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import base64
from types import SimpleNamespace
import numpy as np
import pytest
import torch
from vllm_omni.model_executor.models.personaplex.duplex.policy import (
SILENCE_TOKENS,
SINE_TOKENS,
)
from vllm_omni.model_executor.models.personaplex.duplex.stage0 import (
PersonaPlexStage0DuplexRuntime,
)
from vllm_omni.model_executor.models.personaplex.personaplex_talker import (
PersonaPlexTalkerForConditionalGeneration,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
class _FakeCodec:
def __init__(self) -> None:
self.encode_calls = 0
self.reset_calls = 0
def streaming_init(self, batch_size: int) -> None:
assert batch_size == 1
def encode_frame(self, pcm: torch.Tensor) -> torch.Tensor:
assert pcm.shape == (1, 1920)
self.encode_calls += 1
return torch.full((1, 8), self.encode_calls, dtype=torch.long)
def reset_streaming(self) -> None:
self.reset_calls += 1
class _FakeTalker:
def __init__(self) -> None:
self.device = torch.device("cpu")
self.dtype = torch.float32
self.frame_calls: list[dict[str, torch.Tensor | None]] = []
def _build_prefill_embed(
self,
tokens,
offset,
span,
device,
silence=None,
user_sine=None,
):
del offset, silence, user_sine
values = tokens[:span].to(device=device, dtype=torch.float32)
return values[:, None].expand(-1, 4).contiguous()
def _build_frame_embed(
self,
text_token,
last_agent,
prev_agent,
device,
user_d0=None,
user_d1=None,
):
self.frame_calls.append(
{
"text_token": text_token.clone(),
"last_agent": None if last_agent is None else last_agent.clone(),
"prev_agent": None if prev_agent is None else prev_agent.clone(),
"user_d0": None if user_d0 is None else user_d0.clone(),
"user_d1": None if user_d1 is None else user_d1.clone(),
}
)
value = float(user_d0[0].item()) if user_d0 is not None else -1.0
return torch.full((1, 4), value, device=device)
def _duplex_info(*, seq: int, session_id: str = "session", incarnation: int = 1):
pcm = np.zeros(1920, dtype="<f4")
return {
"data_plane": True,
"session_id": session_id,
"incarnation": incarnation,
"epoch": 0,
"seq": seq,
"payload": {
"format": "pcm_f32le",
"sample_rate_hz": 24000,
"audio": base64.b64encode(pcm.tobytes()).decode("ascii"),
},
"runtime_config": {
"personaplex_model_path": "/unused",
"personaplex_voice_prompt": "NATF2.pt",
"personaplex_persona": "Be concise.",
},
}
def _runtime(*codecs: _FakeCodec) -> PersonaPlexStage0DuplexRuntime:
voice_embeddings = torch.arange(8, dtype=torch.float32).reshape(2, 1, 1, 4)
codec_args = {"codec": codecs[0]}
if len(codecs) > 1:
available = iter(codecs)
codec_args = {
"codec_factory": lambda: next(available),
"max_sessions": len(codecs),
}
return PersonaPlexStage0DuplexRuntime(
_FakeTalker(),
model_path="/unused",
device="cpu",
**codec_args,
tokenizer=lambda _text: [7, 8, 9],
voice_loader=lambda _voice: {
"embeddings": voice_embeddings,
"cache": torch.zeros((1, 17, 4), dtype=torch.long),
},
)
def test_first_append_prepends_voice_and_persona_once() -> None:
codec = _FakeCodec()
runtime = _runtime(codec)
first = runtime.prepare_append(_duplex_info(seq=1), prompt_len=18)
second = runtime.prepare_append(_duplex_info(seq=2), prompt_len=18)
assert first.prefill_applied is True
assert first.prompt_offset == 0
assert first.user_codes.shape == (1, 8)
assert first.inputs_embeds.shape == (18, 4)
assert second.prefill_applied is False
assert second.prompt_offset == 17
assert second.user_codes.shape == (2, 8)
assert second.inputs_embeds.shape == (1, 4)
assert codec.encode_calls == 2
@pytest.mark.parametrize(
("append_count", "expected_tokens", "expected_provided"),
[
(
1,
[*SILENCE_TOKENS, 1, *SINE_TOKENS[1:]],
[False, *([True] * 15)],
),
(
2,
[*SILENCE_TOKENS, 2, *([1] * 7)],
[*([False] * 8), *([True] * 8)],
),
],
)
def test_prepare_append_preserves_native_depformer_teacher_forcing(
append_count: int,
expected_tokens: list[int],
expected_provided: list[bool],
) -> None:
runtime = _runtime(_FakeCodec())
result = None
for seq in range(1, append_count + 1):
result = runtime.prepare_append(_duplex_info(seq=seq), prompt_len=18)
assert result is not None
assert result.info_update["pplex_depformer_audio_tokens"].tolist() == expected_tokens
assert result.info_update["pplex_depformer_audio_provided"].tolist() == expected_provided
def test_frame_embed_uses_previous_effective_agent_frame() -> None:
fake_talker = SimpleNamespace(
config=SimpleNamespace(
num_audio_codebooks=16,
audio_vocab_size=2048,
),
input_embeddings=lambda stack: stack,
)
last_agent = torch.arange(10, 18)
embeds = PersonaPlexTalkerForConditionalGeneration._build_frame_embed(
fake_talker,
torch.tensor([3]),
last_agent,
torch.arange(8),
torch.device("cpu"),
)
assert embeds.shape == (1, 17)
assert torch.equal(embeds[0, 1:9], last_agent)
def test_repeated_append_identity_does_not_advance_codec() -> None:
codec = _FakeCodec()
runtime = _runtime(codec)
info = _duplex_info(seq=1)
first = runtime.prepare_append(info, prompt_len=18, request_id="req")
retry = runtime.prepare_append(info, prompt_len=18, request_id="req")
assert codec.encode_calls == 1
assert torch.equal(first.inputs_embeds, retry.inputs_embeds)
assert torch.equal(first.user_codes, retry.user_codes)
def test_next_append_uses_prior_sample_and_causally_delayed_user_frame() -> None:
runtime = _runtime(_FakeCodec())
runtime.prepare_append(_duplex_info(seq=1), prompt_len=18, request_id="req")
first_agent = torch.arange(8, dtype=torch.long)
runtime.record_sample(request_id="req", text_token=torch.tensor(101), agent_codes=first_agent)
runtime.prepare_append(_duplex_info(seq=2), prompt_len=19, request_id="req")
second_agent = torch.arange(10, 18, dtype=torch.long)
runtime.record_sample(request_id="req", text_token=torch.tensor(102), agent_codes=second_agent)
runtime.prepare_append(_duplex_info(seq=3), prompt_len=20, request_id="req")
silence = torch.tensor(SILENCE_TOKENS, dtype=torch.long)
sine = torch.tensor(SINE_TOKENS, dtype=torch.long)
first_call = runtime.stage_model.frame_calls[0]
assert all(torch.equal(first_call[key], silence) for key in ("last_agent", "prev_agent"))
assert all(torch.equal(first_call[key], sine) for key in ("user_d0", "user_d1"))
second_call = runtime.stage_model.frame_calls[1]
assert second_call["text_token"].tolist() == [101]
expected_first_effective = torch.cat([first_agent[:1], silence[1:]])
assert torch.equal(second_call["last_agent"], expected_first_effective)
assert torch.equal(second_call["prev_agent"], expected_first_effective)
assert torch.equal(second_call["user_d0"], torch.full((8,), 1, dtype=torch.long))
assert torch.equal(second_call["user_d1"], sine)
third_call = runtime.stage_model.frame_calls[2]
assert third_call["text_token"].tolist() == [102]
assert torch.equal(third_call["last_agent"], second_agent)
assert torch.equal(third_call["prev_agent"], second_agent)
assert torch.equal(third_call["user_d0"], torch.full((8,), 2, dtype=torch.long))
assert torch.equal(third_call["user_d1"], torch.full((8,), 1, dtype=torch.long))
def test_decoded_pcm_is_writable_for_torch_zero_copy() -> None:
pcm = PersonaPlexStage0DuplexRuntime._decode_pcm(_duplex_info(seq=1)["payload"])
assert pcm.flags.writeable
def _prepare_two_sessions(
runtime: PersonaPlexStage0DuplexRuntime,
) -> tuple[object, object, object, object]:
first_1 = runtime.prepare_append(_duplex_info(seq=1), prompt_len=18)
second_1 = runtime.prepare_append(
_duplex_info(seq=1, session_id="other", incarnation=2),
prompt_len=18,
)
first_2 = runtime.prepare_append(_duplex_info(seq=2), prompt_len=19)
second_2 = runtime.prepare_append(
_duplex_info(seq=2, session_id="other", incarnation=2),
prompt_len=19,
)
return first_1, second_1, first_2, second_2
def test_live_sessions_keep_independent_streaming_encoders() -> None:
first_codec = _FakeCodec()
second_codec = _FakeCodec()
runtime = _runtime(first_codec, second_codec)
first_1, second_1, first_2, second_2 = _prepare_two_sessions(runtime)
assert first_1.user_codes[:, 0].tolist() == [1]
assert second_1.user_codes[:, 0].tolist() == [1]
assert first_2.user_codes[:, 0].tolist() == [1, 2]
assert second_2.user_codes[:, 0].tolist() == [1, 2]
def test_stage0_session_capacity_fails_before_codec_state_is_shared() -> None:
runtime = _runtime(_FakeCodec(), _FakeCodec())
_prepare_two_sessions(runtime)
with pytest.raises(RuntimeError, match="capacity 2"):
runtime.prepare_append(
_duplex_info(seq=1, session_id="third", incarnation=3),
prompt_len=18,
)
def test_close_session_resets_and_reuses_released_codec() -> None:
first_codec = _FakeCodec()
second_codec = _FakeCodec()
runtime = _runtime(first_codec, second_codec)
_prepare_two_sessions(runtime)
runtime.close_session("session", 1)
assert first_codec.reset_calls == 1
assert second_codec.reset_calls == 0
replacement = runtime.prepare_append(
_duplex_info(seq=1, session_id="replacement", incarnation=4),
prompt_len=18,
)
assert replacement.user_codes[:, 0].tolist() == [3]
@@ -0,0 +1,420 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import base64
from pathlib import Path
from types import SimpleNamespace
import numpy as np
import pytest
from vllm.sampling_params import SamplingParams
from tests.e2e.online_serving import personaplex_realtime_duplex as e2e_driver
from vllm_omni.config.stage_config import (
load_deploy_config,
merge_pipeline_deploy,
)
from vllm_omni.engine.duplex.messages import DuplexFence
from vllm_omni.engine.duplex.runtime import (
DuplexInputMode,
)
from vllm_omni.entrypoints.duplex.runtime_adapter import (
ServingRuntimeConfigError,
)
from vllm_omni.model_executor.models.personaplex.duplex.config import DEFAULT_PERSONA
from vllm_omni.model_executor.models.personaplex.duplex.data_plane import (
PersonaPlexDataPlaneContext,
PersonaPlexDataPlaneSession,
)
from vllm_omni.model_executor.models.personaplex.duplex.input import (
PersonaPlexPcmAppendBuffer,
)
from vllm_omni.model_executor.models.personaplex.duplex.runtime_extension import (
PersonaPlexDuplexRuntimeExtension,
)
from vllm_omni.model_executor.models.personaplex.duplex.serving_adapter import (
PersonaPlexServingRuntimeAdapter,
)
from vllm_omni.model_executor.models.personaplex.pipeline import (
PERSONAPLEX_PIPELINE,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def _pcm_payload(samples: np.ndarray, *, sample_rate_hz: int = 24000) -> dict[str, object]:
samples = np.ascontiguousarray(samples, dtype="<f4")
return {
"type": "audio",
"format": "pcm_f32le",
"sample_rate_hz": sample_rate_hz,
"audio": base64.b64encode(samples.tobytes()).decode("ascii"),
}
def _audio_client(
frame_count: int = 10,
*,
voiced_frames: int = 0,
) -> SimpleNamespace:
frames = np.full((frame_count, e2e_driver.FRAME_SAMPLES), 7, dtype="<i2")
frames[:voiced_frames] = 1000
raw = frames.tobytes()
event = {
"type": "response.audio.delta",
"response_id": "response-1",
"sample_rate_hz": e2e_driver.SAMPLE_RATE_HZ,
"metadata": {
"vllm_omni": {
"runtime_impl": "scheduler_data_plane",
"uses_model_runner_scheduler": True,
"runner_kv_backed": True,
}
},
}
events = SimpleNamespace(
events=[event],
response_audio={"response-1": [raw]},
audio_bytes=lambda: raw,
)
return SimpleNamespace(events=events)
def test_realtime_audio_frame_stats_separate_voiced_and_silent_frames() -> None:
silent = np.zeros(e2e_driver.FRAME_SAMPLES, dtype="<i2")
voiced = np.full(e2e_driver.FRAME_SAMPLES, 1000, dtype="<i2")
stats = e2e_driver._audio_frame_stats(
silent.tobytes() + voiced.tobytes(),
input_frames=3,
voiced_frame_rms_threshold=1e-4,
)
assert stats == {
"output_frames": 2,
"voiced_frames": 1,
"silent_frames": 1,
"frame_deficit": 1,
"frame_coverage_ratio": pytest.approx(2 / 3),
}
def test_realtime_audio_frame_stats_reject_partial_codec_frame() -> None:
partial_frame = np.zeros(e2e_driver.FRAME_SAMPLES - 1, dtype="<i2")
with pytest.raises(AssertionError, match="whole 80 ms codec frames"):
e2e_driver._audio_frame_stats(
partial_frame.tobytes(),
input_frames=1,
voiced_frame_rms_threshold=1e-4,
)
@pytest.mark.parametrize("frame_count", [10, 40])
def test_e2e_driver_rejects_inaudible_sessions(frame_count: int) -> None:
args = SimpleNamespace(
max_frame_deficit=4,
voiced_frame_rms_threshold=1e-3,
min_voiced_frames=5,
)
with pytest.raises(AssertionError, match="audible speech"):
e2e_driver._session_result(
_audio_client(frame_count),
input_frames=frame_count,
args=args,
minimum_chunks=1,
)
def test_e2e_driver_uses_absolute_audible_floor() -> None:
args = SimpleNamespace(
max_frame_deficit=4,
voiced_frame_rms_threshold=1e-3,
min_voiced_frames=5,
)
e2e_driver._session_result(
_audio_client(200, voiced_frames=5),
input_frames=200,
args=args,
minimum_chunks=1,
)
def test_e2e_driver_defaults_replacement_to_full_input() -> None:
args = e2e_driver.parse_args(
[
"--model",
"/model",
"--input-wav",
"/input.wav",
]
)
assert args.replacement_frames == 0
assert args.min_voiced_frames == 5
def test_e2e_driver_records_and_validates_input_sha256(tmp_path: Path) -> None:
input_wav = tmp_path / "input.wav"
input_wav.write_bytes(b"fixture-wav")
identity = e2e_driver._input_identity(
input_wav,
expected_sha256="bade4b3c163edde390ff391207d34a887257d8a9cc3b621cc8c618b6e6761304",
)
assert identity == {
"path": str(input_wav.resolve()),
"sha256": "bade4b3c163edde390ff391207d34a887257d8a9cc3b621cc8c618b6e6761304",
}
def test_e2e_driver_rejects_unexpected_input_sha256(tmp_path: Path) -> None:
input_wav = tmp_path / "input.wav"
input_wav.write_bytes(b"fixture-wav")
with pytest.raises(ValueError, match="input WAV SHA-256 mismatch"):
e2e_driver._input_identity(input_wav, expected_sha256="0" * 64)
@pytest.mark.parametrize(
("attribute", "expected"),
[
(
"duplex_runtime_extension",
"vllm_omni.model_executor.models.personaplex.duplex.runtime_extension.PersonaPlexDuplexRuntimeExtension",
),
(
"duplex_serving_adapter",
"vllm_omni.model_executor.models.personaplex.duplex.serving_adapter.PersonaPlexServingRuntimeAdapter",
),
],
)
def test_personaplex_pipeline_registers_unified_duplex_components(
attribute: str,
expected: str,
) -> None:
assert PERSONAPLEX_PIPELINE.duplex_control_enabled is True
assert getattr(PERSONAPLEX_PIPELINE, attribute) == expected
@pytest.mark.parametrize(
("engine_arg", "expected"),
[
("skip_tokenizer_init", True),
("enable_prefix_caching", False),
],
)
def test_personaplex_stage_engine_args(
engine_arg: str,
expected: bool,
) -> None:
deploy_path = Path(__file__).parents[5] / "vllm_omni" / "deploy" / "personaplex.yaml"
stages = merge_pipeline_deploy(
PERSONAPLEX_PIPELINE,
load_deploy_config(deploy_path),
)
assert [stage.yaml_engine_args.get(engine_arg) for stage in stages] == [
expected,
expected,
]
def test_personaplex_duplex_capacity_is_propagated_to_all_model_stages() -> None:
deploy_path = Path(__file__).parents[5] / "vllm_omni" / "deploy" / "personaplex.yaml"
deploy = load_deploy_config(deploy_path)
stages = merge_pipeline_deploy(PERSONAPLEX_PIPELINE, deploy)
assert deploy.duplex_session.max_sessions == 2
assert [stage.yaml_engine_args.get("duplex_max_sessions") for stage in stages] == [2, 2]
assert "personaplex_codec_max_sessions" not in deploy.connectors["connector_of_shared_memory"]["extra"]
def test_personaplex_capabilities_are_honest() -> None:
single = PersonaPlexServingRuntimeAdapter(lambda *_: None).capabilities(max_sessions=1)
multi = PersonaPlexServingRuntimeAdapter(lambda *_: None).capabilities(max_sessions=2)
assert multi.implementation_level == "model_native_duplex"
assert multi.input_modes == ["append_audio_chunk"]
assert multi.chunk_period_ms == 80
assert single.supports_multi_session is False
assert multi.supports_multi_session is True
assert multi.supports_multi_session_same_replica is True
assert multi.supports_barge_in is False
def test_personaplex_runtime_config_update_rejects_changed_persona() -> None:
adapter = PersonaPlexServingRuntimeAdapter(lambda *_: None)
short_persona = "You are a helpful assistant."
longer_persona = (
"You are now a swashbuckling pirate captain who speaks entirely in "
"nautical slang, references buried treasure constantly, and never "
"breaks character no matter what the user asks."
)
current = {"personaplex_persona": short_persona, "personaplex_voice_prompt": "NATF2.pt"}
with pytest.raises(ServingRuntimeConfigError, match="persona"):
adapter.runtime_config_for_update(
SimpleNamespace(instructions=longer_persona, extra_body={}),
current,
)
unchanged = adapter.runtime_config_for_update(
SimpleNamespace(instructions=short_persona, extra_body={}),
current,
)
assert unchanged["personaplex_persona"] == short_persona
def test_personaplex_runtime_config_update_rejects_changed_voice() -> None:
adapter = PersonaPlexServingRuntimeAdapter(lambda *_: None)
current = {"personaplex_persona": DEFAULT_PERSONA, "personaplex_voice_prompt": "NATF2.pt"}
with pytest.raises(ServingRuntimeConfigError, match="voice"):
adapter.runtime_config_for_update(
SimpleNamespace(instructions=None, voice="OtherVoice.pt", extra_body={}),
current,
)
unchanged = adapter.runtime_config_for_update(
SimpleNamespace(instructions=None, voice="NATF2.pt", extra_body={}),
current,
)
assert unchanged["personaplex_voice_prompt"] == "NATF2.pt"
def test_personaplex_rejects_native_duplex_opt_out_flag() -> None:
# PersonaPlex is always model-native; the per-session opt-in knob is
# server-owned there, under both the canonical name and its deprecated
# model-prefixed alias.
adapter = PersonaPlexServingRuntimeAdapter(lambda *_: None)
with pytest.raises(ServingRuntimeConfigError, match="native_duplex"):
adapter.validate_client_extra_body({"native_duplex": False})
with pytest.raises(ServingRuntimeConfigError, match="minicpmo45_native_duplex"):
adapter.validate_client_extra_body({"minicpmo45_native_duplex": False})
def test_pcm_buffer_emits_one_80ms_frame_transactionally() -> None:
buffer = PersonaPlexPcmAppendBuffer()
reservation = buffer.prepare_append(
_pcm_payload(np.arange(1920, dtype=np.float32)),
operation_id="op-1",
chunk_period_ms=80,
allow_emit=True,
)
assert reservation is not None
assert reservation.byte_count == 1920 * 4
assert buffer.pending_byte_count == 0
reservation.rollback()
assert buffer.pending_byte_count == 1920 * 4
def test_pcm_buffer_rejects_changed_sample_rate() -> None:
buffer = PersonaPlexPcmAppendBuffer()
assert (
buffer.prepare_append(
_pcm_payload(np.zeros(100, np.float32)),
operation_id="op-1",
chunk_period_ms=80,
allow_emit=False,
)
is None
)
with pytest.raises(ValueError, match="sample_rate"):
buffer.prepare_append(
_pcm_payload(np.zeros(100, np.float32), sample_rate_hz=16000),
operation_id="op-2",
chunk_period_ms=80,
allow_emit=False,
)
def test_pcm_buffer_rejects_non_finite_samples() -> None:
buffer = PersonaPlexPcmAppendBuffer()
with pytest.raises(ValueError, match="finite"):
buffer.prepare_append(
_pcm_payload(np.array([0.0, np.nan], np.float32)),
operation_id="op-1",
chunk_period_ms=80,
allow_emit=False,
)
def test_runtime_extension_builds_one_frame_scheduler_append() -> None:
extension = PersonaPlexDuplexRuntimeExtension()
defaults = (
SamplingParams(temperature=0.8, top_k=10, max_tokens=10),
SamplingParams(max_tokens=1024),
)
configured = extension.configure_sampling_params(runtime_config={}, defaults=defaults)
assert configured[0].temperature == 0.0
assert configured[0].top_k == 1
assert configured[0].max_tokens == 1
assert configured[1] is defaults[1]
plan = extension.plan_append(
request_id="req",
fence=DuplexFence("session", incarnation=2, epoch=3),
session_config={"instructions": "Be concise."},
runtime_config={"personaplex_prefill_slots": 4},
seq=1,
turn_seq=1,
mode=DuplexInputMode.APPEND_AUDIO_CHUNK,
payload=_pcm_payload(np.zeros(1920, np.float32)),
final=False,
sampling_params=configured[0],
)
assert len(plan.prompt["prompt_token_ids"]) == 5
duplex = plan.prompt["model_intermediate_buffer"]["duplex"]
assert duplex["session_id"] == "session"
assert duplex["incarnation"] == 2
assert duplex["epoch"] == 3
assert duplex["seq"] == 1
assert duplex["data_plane"] is True
def test_projector_emits_only_cumulative_audio_and_text_suffixes() -> None:
encoded_sizes: list[int] = []
def encode_audio(audio, _sample_rate, _response_format, _speed):
size = int(np.asarray(audio, dtype=np.float32).size)
encoded_sizes.append(size)
return f"audio-{size}"
projector = PersonaPlexDataPlaneSession(encode_audio)
projector.begin_request("req")
context = PersonaPlexDataPlaneContext(response_format="wav")
first = SimpleNamespace(
request_id="req",
outputs=[SimpleNamespace(text="he", multimodal_output={})],
multimodal_output={"model_outputs": np.arange(4, dtype=np.float32), "sr": 24000},
finished=False,
)
second = SimpleNamespace(
request_id="req",
outputs=[SimpleNamespace(text="hello", multimodal_output={})],
multimodal_output={"model_outputs": np.arange(6, dtype=np.float32), "sr": 24000},
finished=False,
)
projected = list(
projector.project(
{"data_plane_outputs": [first, second]},
context=context,
)
)
assert encoded_sizes == [4, 2]
assert [item["text"] for item in projected] == ["he", "llo"]
assert [item["audio_data"] for item in projected] == ["audio-4", "audio-2"]
assert all(item["sample_rate_hz"] == 24000 for item in projected)
assert all(item["end_of_turn"] is False for item in projected)
@@ -0,0 +1,316 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from types import SimpleNamespace
import pytest
import torch
from torch import nn
from vllm_omni.model_executor.models.personaplex import (
personaplex_code2wav,
personaplex_mimi,
)
from vllm_omni.model_executor.models.personaplex.personaplex_code2wav import (
PersonaPlexCode2Wav,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
class _FakeStreamingMimi(nn.Module):
def __init__(self, stream_value: int = 0) -> None:
super().__init__()
self.stream_value = stream_value
self.decode_frame_calls = 0
self.reset_calls = 0
def decode(self, codes: torch.Tensor, return_dict: bool = True):
assert return_dict
frames = int(codes.shape[-1])
return SimpleNamespace(audio_values=torch.arange(frames * 4, dtype=torch.float32).reshape(1, 1, -1))
def streaming_init(self, batch_size: int) -> None:
assert batch_size == 1
def decode_frame(self, codes: torch.Tensor) -> torch.Tensor:
assert codes.shape == (1, 2)
self.decode_frame_calls += 1
value = self.stream_value + self.decode_frame_calls
return torch.full((1, 4), float(value), dtype=torch.float32)
def reset_streaming(self) -> None:
self.reset_calls += 1
class _FakeChunkedMimi(_FakeStreamingMimi):
def __init__(self) -> None:
super().__init__()
self.decode_frames_calls = 0
self.decode_frames_shapes: list[tuple[int, ...]] = []
def decode_frames(self, codes: torch.Tensor) -> torch.Tensor:
self.decode_frames_calls += 1
self.decode_frames_shapes.append(tuple(codes.shape))
return torch.arange(codes.shape[-1] * 4, dtype=torch.float32).reshape(1, -1)
class _FakeMimiModel(nn.Module):
def __init__(self, _config) -> None:
super().__init__()
self.encoder_transformer = nn.Linear(1, 1, bias=False)
self.decoder_transformer = nn.Linear(1, 1, bias=False)
self.quantizer = nn.Linear(1, 1, bias=False)
self.encoder = SimpleNamespace(layers=[])
self.downsample = SimpleNamespace(conv=nn.Conv1d(1, 1, 1))
self.upsample = SimpleNamespace(conv=nn.ConvTranspose1d(1, 1, 1))
self.decoder = SimpleNamespace(layers=[])
def load_state_dict(self, _state_dict, strict: bool = True):
assert not strict
return SimpleNamespace(
missing_keys=[
"encoder_transformer.weight",
"decoder_transformer.weight",
],
unexpected_keys=[],
)
class _FakeMimiTransformer(nn.Module):
def __init__(self) -> None:
super().__init__()
self.weight = nn.Parameter(torch.zeros(1))
def load_weights(self, _state_dict, _prefix: str) -> int:
return 80
def test_moshi_checkpoint_mapping_drops_replaced_transformer_weights(
monkeypatch: pytest.MonkeyPatch,
) -> None:
mapper = getattr(personaplex_mimi, "_map_moshi_codec_weights", None)
assert callable(mapper), "PersonaPlex Mimi must load its bundled codec weights without Hugging Face"
source = {
"encoder.model.3.conv.conv.weight": torch.ones(1),
"downsample.conv.conv.conv.weight": torch.ones(1),
"quantizer.rvq_first.vq.layers.0._codebook.embedding_sum": torch.ones(1),
"quantizer.rvq_rest.vq.layers.3._codebook._initialized": torch.ones(1),
"encoder_transformer.transformer.layers.0.norm1.weight": torch.ones(1),
}
mapped = mapper(source)
assert set(mapped) == {
"encoder.layers.3.conv.weight",
"downsample.conv.weight",
"quantizer.semantic_residual_vector_quantizer.layers.0.codebook.embed_sum",
"quantizer.acoustic_residual_vector_quantizer.layers.3.codebook.initialized",
}
assert mapped["encoder.layers.3.conv.weight"] is source["encoder.model.3.conv.conv.weight"]
monkeypatch.setattr("transformers.MimiConfig", lambda: object())
monkeypatch.setattr("transformers.MimiModel", _FakeMimiModel)
monkeypatch.setattr("safetensors.torch.load_file", lambda *_args, **_kwargs: {})
monkeypatch.setattr(personaplex_mimi, "_MimiStreamingTransformer", _FakeMimiTransformer)
codec = personaplex_mimi.PersonaPlexMimiCodec(
checkpoint="/unused",
device="cpu",
)
assert not hasattr(codec.model, "encoder_transformer")
assert not hasattr(codec.model, "decoder_transformer")
assert not any(
name.startswith(("model.encoder_transformer.", "model.decoder_transformer."))
for name, _ in codec.named_parameters()
)
def test_mimi_full_stream_reset_reuses_all_state_storage_and_clears_offsets() -> None:
codec = personaplex_mimi.PersonaPlexMimiCodec.__new__(personaplex_mimi.PersonaPlexMimiCodec)
nn.Module.__init__(codec)
codec.device = torch.device("cpu")
codec.dtype = torch.float32
codec._enc_stages = [("conv", personaplex_mimi._StreamConv1d(nn.Conv1d(1, 1, 3)))]
codec._dec_stages = []
codec._downsample = personaplex_mimi._StreamConv1d(nn.Conv1d(1, 1, 3))
codec._upsample = personaplex_mimi._StreamConvTr1d(nn.ConvTranspose1d(1, 1, 3))
codec.encoder_transformer = personaplex_mimi._MimiStreamingTransformer(
num_layers=1,
dim=4,
num_heads=1,
context=3,
)
codec.decoder_transformer = personaplex_mimi._MimiStreamingTransformer(
num_layers=1,
dim=4,
num_heads=1,
context=3,
)
codec.streaming_init(batch_size=2)
conv_states = list(codec._conv_states())
conv_buffers = [
state.prev if isinstance(state, personaplex_mimi._StreamConv1d) else state.partial for state in conv_states
]
transformers = [codec.encoder_transformer, codec.decoder_transformer]
kv_states = [kv for transformer in transformers for kv in transformer._kv]
state_tensors = [
*conv_buffers,
*(kv.cache for kv in kv_states),
*(kv.end_offset for kv in kv_states),
*(kv.start_offset for kv in kv_states),
*(transformer._offset for transformer in transformers),
]
for tensor in state_tensors:
tensor.fill_(1)
for state in conv_states:
state._fresh.zero_()
storage = [tensor.data_ptr() for tensor in state_tensors]
codec.reset_streaming()
assert storage == [tensor.data_ptr() for tensor in state_tensors]
cleared = [
*conv_buffers,
*(kv.end_offset for kv in kv_states),
*(kv.start_offset for kv in kv_states),
*(transformer._offset for transformer in transformers),
]
assert all(not tensor.any() for tensor in cleared)
assert all(state._fresh.all() for state in conv_states)
def _model(*, max_sessions: int = 1) -> tuple[PersonaPlexCode2Wav, _FakeStreamingMimi]:
mimi_config = SimpleNamespace(num_codebooks=2, sample_rate=24000, samples_per_frame=4, mimi_name=None)
config = SimpleNamespace(mimi_config=mimi_config, mimi_name=None)
vllm_config = SimpleNamespace(
model_config=SimpleNamespace(
model="/unused",
hf_config=config,
duplex_max_sessions=max_sessions,
),
device_config=SimpleNamespace(device="cpu"),
)
model = PersonaPlexCode2Wav(vllm_config=vllm_config)
mimi = _FakeStreamingMimi()
model.mimi = mimi
model._mimi_device = torch.device("cpu")
return model, mimi
def _codes(frames: int, *, start: int = 0) -> torch.Tensor:
return torch.stack([torch.arange(start + offset, start + offset + frames) for offset in (0, 100)]).reshape(-1)
def _audio(output) -> torch.Tensor:
return output.multimodal_outputs["model_outputs"][0]
@pytest.mark.parametrize("second_codes", [_codes(3), _codes(1, start=100)])
def test_resumable_codes_emit_only_new_pcm(second_codes: torch.Tensor) -> None:
model, mimi = _model()
first = model(input_ids=_codes(2), request_ids=["req"])
second = model(input_ids=second_codes, request_ids=["req"])
assert _audio(first).numel() == 8
assert _audio(second).numel() == 4
assert mimi.decode_frame_calls == 3
@pytest.mark.parametrize(
("frames", "expected_shapes"),
[
(5, [(1, 2, 5)]),
(12, [(1, 2, 5), (1, 2, 5), (1, 2, 2)]),
],
)
def test_streaming_decode_uses_bounded_multi_frame_calls(
frames: int,
expected_shapes: list[tuple[int, ...]],
) -> None:
model, _ = _model()
mimi = _FakeChunkedMimi()
model.mimi = mimi
output = model(input_ids=_codes(frames), request_ids=["req"])
assert _audio(output).numel() == frames * 4
assert mimi.decode_frames_shapes == expected_shapes
assert mimi.decode_frame_calls == 0
def test_profile_inputs_skip_decode_while_malformed_online_input_warns(
monkeypatch: pytest.MonkeyPatch,
) -> None:
model, mimi = _model()
runtime_info = [{"meta": {"personaplex_dummy_profile": True}}]
warnings: list[tuple] = []
monkeypatch.setattr(personaplex_code2wav.logger, "warning", lambda *args: warnings.append(args))
output = model(
input_ids=torch.arange(3),
runtime_additional_information=runtime_info,
)
assert _audio(output).numel() == 0
assert mimi.decode_frame_calls == 0
assert warnings == []
output = model(input_ids=torch.arange(3))
assert _audio(output).numel() == 0
assert mimi.decode_frame_calls == 0
assert len(warnings) == 1
assert "not divisible by" in warnings[0][0]
assert model.get_dummy_runtime_additional_information(2) == [
{"meta": {"personaplex_dummy_profile": True}},
{"meta": {"personaplex_dummy_profile": True}},
]
def test_request_id_falls_back_to_runtime_information() -> None:
model, _ = _model()
info = [{"request_id": "runtime-req"}]
model(input_ids=_codes(2), runtime_additional_information=info)
second = model(input_ids=_codes(3), runtime_additional_information=info)
assert _audio(second).numel() == 4
def test_decoder_slot_lifecycle_isolated_capacity_and_reuse() -> None:
model, _ = _model(max_sessions=2)
first = _FakeStreamingMimi(stream_value=10)
second = _FakeStreamingMimi(stream_value=20)
model._set_mimi_codecs([first, second])
assert model._max_codec_sessions == 2
initial = model(
input_ids=torch.cat([_codes(1), _codes(1)]),
request_ids=["first", "second"],
seq_token_counts=[2, 2],
)
continued = model(
input_ids=torch.cat([_codes(1, start=10), _codes(1, start=20)]),
request_ids=["first", "second"],
seq_token_counts=[2, 2],
)
initial_audio = initial.multimodal_outputs["model_outputs"]
continued_audio = continued.multimodal_outputs["model_outputs"]
assert initial_audio[0].tolist() == [11.0] * 4
assert initial_audio[1].tolist() == [21.0] * 4
assert continued_audio[0].tolist() == [12.0] * 4
assert continued_audio[1].tolist() == [22.0] * 4
assert model._request_codec_slots == {"first": 0, "second": 1}
with pytest.raises(RuntimeError, match="decoder capacity 2 is exhausted"):
model._codec_for_request("overflow")
model.on_requests_finished(["first"])
assert first.reset_calls == 1
assert second.reset_calls == 0
replacement = model(input_ids=_codes(1), request_ids=["replacement"])
assert _audio(replacement).tolist() == [13.0] * 4