[WIP][Bugfix][MiniCPM-o] Fix Full duplex bugs for minicpm-challenge (#5496)
Signed-off-by: Sy03 <1370724210@qq.com>
This commit is contained in:
60 files changed
+3696
-634
No files matched your search
@@ -165,6 +165,31 @@ steps:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: "Omni · MiniCPM-o 4.5 Duplex Test"
|
||||
source_file_dependencies:
|
||||
- tests/assets/minicpmo_4_5/response_required_16k.wav
|
||||
- tests/assets/minicpmo_4_5/soft_interrupt_16k.wav
|
||||
- tests/e2e/online_serving/test_minicpmo_4_5_duplex.py
|
||||
- tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py
|
||||
- tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py
|
||||
- tests/e2e/online_serving/minicpmo_realtime_duplex_scenarios.py
|
||||
- tests/e2e/online_serving/run_minicpmo_realtime_duplex_multi_session.py
|
||||
- tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py
|
||||
- tests/helpers/minicpmo_4_5_duplex.py
|
||||
- vllm_omni/deploy/minicpmo_4_5_duplex.yaml
|
||||
- vllm_omni/experimental/fullduplex/
|
||||
- vllm_omni/model_executor/models/minicpmo_4_5/
|
||||
- vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py
|
||||
- vllm_omni/core/sched/
|
||||
- vllm_omni/outputs/multimodal_accumulation.py
|
||||
- vllm_omni/engine/stage_init_utils.py
|
||||
- vllm_omni/engine/async_omni_engine.py
|
||||
- vllm_omni/entrypoints/cli/serve.py
|
||||
timeout_in_minutes: 50
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/online_serving/test_minicpmo_4_5_duplex.py -m 'advanced_model and cuda' --run-level 'advanced_model'
|
||||
mirror_hardwares: h100_2
|
||||
|
||||
- label: "TTS · VoxCPM2 Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/offline_inference/test_voxcpm2_tts.py
|
||||
|
||||
@@ -15,7 +15,13 @@ steps:
|
||||
- label: ":full_moon: Omni · Function Test with H100"
|
||||
timeout_in_minutes: 120
|
||||
commands:
|
||||
- pytest -sv tests/e2e --run-level "full_model" -m "full_model and H100 and omni" --ignore=tests/e2e/accuracy --ignore=tests/e2e/online_serving/test_qwen3_omni_multi_replicas.py
|
||||
- pytest -sv tests/e2e --run-level "full_model" -m "full_model and H100 and omni" --ignore=tests/e2e/accuracy --ignore=tests/e2e/online_serving/test_qwen3_omni_multi_replicas.py --ignore=tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py
|
||||
mirror_hardwares: h100_2
|
||||
|
||||
- label: ":full_moon: Omni · MiniCPM-o 4.5 Duplex Test"
|
||||
timeout_in_minutes: 50
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py -m "full_model and cuda and H100 and omni" --run-level "full_model"
|
||||
mirror_hardwares: h100_2
|
||||
|
||||
- label: ":full_moon: Omni · Doc Test with L4"
|
||||
|
||||
@@ -210,6 +210,33 @@ steps:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: "Omni · MiniCPM-o 4.5 Duplex Test"
|
||||
source_file_dependencies:
|
||||
- tests/assets/minicpmo_4_5/response_required_16k.wav
|
||||
- tests/assets/minicpmo_4_5/soft_interrupt_16k.wav
|
||||
- tests/e2e/online_serving/test_minicpmo_4_5_duplex.py
|
||||
- tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py
|
||||
- tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py
|
||||
- tests/e2e/online_serving/minicpmo_realtime_duplex_scenarios.py
|
||||
- tests/e2e/online_serving/run_minicpmo_realtime_duplex_multi_session.py
|
||||
- tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py
|
||||
- tests/helpers/minicpmo_4_5_duplex.py
|
||||
- vllm_omni/deploy/minicpmo_4_5_duplex.yaml
|
||||
- vllm_omni/experimental/fullduplex/
|
||||
- vllm_omni/model_executor/models/minicpmo_4_5/
|
||||
- vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py
|
||||
- vllm_omni/core/sched/
|
||||
- vllm_omni/outputs/multimodal_accumulation.py
|
||||
- vllm_omni/engine/stage_init_utils.py
|
||||
- vllm_omni/engine/async_omni_engine.py
|
||||
- vllm_omni/entrypoints/cli/serve.py
|
||||
commands:
|
||||
- |
|
||||
timeout 20m bash -c "
|
||||
pytest -s -v tests/e2e/online_serving/test_minicpmo_4_5_duplex.py -m 'core_model and cuda' --run-level 'core_model'
|
||||
"
|
||||
mirror_hardwares: h100_2
|
||||
|
||||
- label: "TTS · VoxCPM2 Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_voxcpm2_tts.py
|
||||
|
||||
@@ -6,6 +6,7 @@ class FullDuplexPcmPlayback extends AudioWorkletProcessor {
|
||||
this.playedFrames = 0;
|
||||
this.underrunFrames = 0;
|
||||
this.drain = null;
|
||||
this.terminalFadeFrames = 0;
|
||||
this.started = false;
|
||||
this.activeResponseId = null;
|
||||
this.initialBufferFrames = Math.round(sampleRate * 0.2);
|
||||
@@ -31,6 +32,7 @@ class FullDuplexPcmPlayback extends AudioWorkletProcessor {
|
||||
}
|
||||
} else if (message.type === 'drain') {
|
||||
this.drain = { responseId: message.responseId || null };
|
||||
this.terminalFadeFrames = Math.min(this.fadeFrames, this.bufferedFrames());
|
||||
if (!this.started && this.bufferedFrames() > 0) {
|
||||
this.bufferWaitFrames = 0;
|
||||
this.startPlayback();
|
||||
@@ -42,6 +44,7 @@ class FullDuplexPcmPlayback extends AudioWorkletProcessor {
|
||||
this.playedFrames = 0;
|
||||
this.underrunFrames = 0;
|
||||
this.drain = null;
|
||||
this.terminalFadeFrames = 0;
|
||||
this.started = false;
|
||||
this.activeResponseId = null;
|
||||
this.bufferWaitFrames = this.initialBufferFrames;
|
||||
@@ -76,6 +79,7 @@ class FullDuplexPcmPlayback extends AudioWorkletProcessor {
|
||||
this.playedFrames = 0;
|
||||
this.underrunFrames = 0;
|
||||
this.drain = null;
|
||||
this.terminalFadeFrames = 0;
|
||||
this.started = false;
|
||||
this.activeResponseId = null;
|
||||
this.bufferWaitFrames = this.initialBufferFrames;
|
||||
@@ -116,6 +120,7 @@ class FullDuplexPcmPlayback extends AudioWorkletProcessor {
|
||||
while (target < output.length && this.queue.length > 0) {
|
||||
const pcm = this.queue[0];
|
||||
const count = Math.min(output.length - target, pcm.length - this.offset);
|
||||
const remainingBeforeChunk = this.bufferedFrames();
|
||||
for (let index = 0; index < count; index += 1) {
|
||||
let sample = pcm[this.offset + index] / 32768;
|
||||
if (this.fadeInFrames > 0) {
|
||||
@@ -123,6 +128,16 @@ class FullDuplexPcmPlayback extends AudioWorkletProcessor {
|
||||
sample *= elapsed / this.fadeFrames;
|
||||
this.fadeInFrames -= 1;
|
||||
}
|
||||
const remainingFrames = remainingBeforeChunk - index;
|
||||
if (
|
||||
this.drain
|
||||
&& this.terminalFadeFrames > 0
|
||||
&& remainingFrames <= this.terminalFadeFrames
|
||||
) {
|
||||
sample *= this.terminalFadeFrames === 1
|
||||
? 0
|
||||
: (remainingFrames - 1) / (this.terminalFadeFrames - 1);
|
||||
}
|
||||
output[target + index] = sample;
|
||||
}
|
||||
target += count;
|
||||
|
||||
Binary file not shown.
Binary file not shown.
@@ -242,9 +242,52 @@ def test_explicit_streaming_payload_replaces_placeholder_prompt() -> None:
|
||||
"tts_token_ids": [10, 20],
|
||||
"meta": {"replace_streaming_prompt": True},
|
||||
}
|
||||
update.model_intermediate_buffer = {
|
||||
"ids": {"tts": [41, 42, 99]},
|
||||
"meta": {"turn_eos_token_id": 99},
|
||||
}
|
||||
|
||||
sched._update_request_as_session(session, update)
|
||||
|
||||
assert session.prompt_token_ids == [10, 20]
|
||||
assert session.additional_information == update.additional_information
|
||||
assert session.model_intermediate_buffer == {
|
||||
"ids": {"tts": [41, 42, 99]},
|
||||
"meta": {"turn_eos_token_id": 99},
|
||||
}
|
||||
assert session.status == RequestStatus.WAITING
|
||||
|
||||
|
||||
def test_model_intermediate_streaming_payload_replaces_computed_prompt() -> None:
|
||||
sched = _make_scheduler(stage_id=1)
|
||||
sched.chunk_transfer_adapter = SimpleNamespace(
|
||||
receives_chunks=False,
|
||||
segment_finished_requests=set(),
|
||||
)
|
||||
session = _make_request()
|
||||
session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
|
||||
session.prompt_token_ids = [0] * 59
|
||||
session._all_token_ids.clear()
|
||||
session._all_token_ids.extend(session.prompt_token_ids)
|
||||
session.num_prompt_tokens = 59
|
||||
session.num_computed_tokens = 59
|
||||
update = _make_update([0] * 10)
|
||||
update.additional_information = None
|
||||
update.model_intermediate_buffer = {
|
||||
"ids": {"tts": list(range(8))},
|
||||
"hidden_states": {"tts": [[0.0]] * 8},
|
||||
"meta": {
|
||||
"replace_streaming_prompt": True,
|
||||
"next_stage_prompt_len": 10,
|
||||
},
|
||||
}
|
||||
|
||||
sched._update_request_as_session(session, update)
|
||||
|
||||
assert session.prompt_token_ids == [0] * 10
|
||||
assert list(session._all_token_ids) == [0] * 10
|
||||
assert session.num_prompt_tokens == 10
|
||||
assert session.num_computed_tokens == 0
|
||||
assert session.additional_information is None
|
||||
assert session.model_intermediate_buffer == update.model_intermediate_buffer
|
||||
assert session.status == RequestStatus.WAITING
|
||||
@@ -54,6 +54,18 @@ class _SchedulerStub(OmniGenerationScheduler):
|
||||
raise AssertionError("unexpected enqueue for skipped_waiting miss")
|
||||
|
||||
|
||||
class _AsyncChunkStopSchedulerStub(OmniGenerationScheduler):
|
||||
def __init__(self) -> None:
|
||||
self.num_waiting_for_streaming_input = 0
|
||||
self.chunk_transfer_adapter = SimpleNamespace(receives_chunks=True)
|
||||
self.enqueued: list[Request] = []
|
||||
self.enqueued_statuses: list[RequestStatus] = []
|
||||
|
||||
def _enqueue_waiting_request(self, session: Request) -> None:
|
||||
self.enqueued.append(session)
|
||||
self.enqueued_statuses.append(session.status)
|
||||
|
||||
|
||||
def _make_request(**kwargs) -> Request:
|
||||
sp = SamplingParams(max_tokens=8)
|
||||
defaults = dict(
|
||||
@@ -133,6 +145,21 @@ def test_resumable_generation_stop_marks_segment_boundary() -> None:
|
||||
assert output.is_segment_finished is True
|
||||
|
||||
|
||||
def test_async_chunk_resumable_stop_rearms_connector_polling() -> None:
|
||||
sched = _AsyncChunkStopSchedulerStub()
|
||||
session = _make_request(request_id="req-async-chunk-next-segment")
|
||||
session.status = RequestStatus.FINISHED_STOPPED
|
||||
session.resumable = True
|
||||
|
||||
finished = sched._handle_stopped_request(session)
|
||||
|
||||
assert finished is False
|
||||
assert sched.enqueued == [session]
|
||||
assert sched.enqueued_statuses == [RequestStatus.WAITING]
|
||||
assert session.status == RequestStatus.WAITING
|
||||
assert sched.num_waiting_for_streaming_input == 0
|
||||
|
||||
|
||||
class TestReplaceSessionWithStreamingUpdate:
|
||||
def test_resets_tokens_and_prompt_from_update(self) -> None:
|
||||
sched = _SchedulerStub()
|
||||
|
||||
@@ -22,10 +22,7 @@ _MINICPMO_SERVER = [
|
||||
model="openbmb/MiniCPM-o-4_5",
|
||||
stage_config_path=_MINICPMO_DEPLOY,
|
||||
use_stage_cli=True,
|
||||
server_args=[
|
||||
"--trust-remote-code",
|
||||
"--no-async-chunk",
|
||||
],
|
||||
server_args=["--trust-remote-code"],
|
||||
),
|
||||
id="minicpmo_4_5",
|
||||
)
|
||||
|
||||
@@ -71,6 +71,30 @@ def test_streaming_payload_can_replace_placeholder_prompt(mocker: MockerFixture)
|
||||
request.update_block_hashes.assert_called_once_with()
|
||||
|
||||
|
||||
def test_streaming_payload_can_append_exact_prompt_length(mocker: MockerFixture) -> None:
|
||||
request = SimpleNamespace(
|
||||
_all_token_ids=[0, 0, 7, 8],
|
||||
_output_token_ids=[7, 8],
|
||||
prompt_token_ids=[0, 0],
|
||||
num_computed_tokens=4,
|
||||
num_prompt_tokens=2,
|
||||
update_block_hashes=mocker.Mock(),
|
||||
)
|
||||
payload = {
|
||||
"ids": {"prompt": [1, 2, 3]},
|
||||
"meta": {"next_stage_prompt_len": 3},
|
||||
}
|
||||
|
||||
construct_next_stage_streaming_input_prompt(payload, request)
|
||||
|
||||
assert request.prompt_token_ids == [0, 0, 7, 8, 0, 0, 0]
|
||||
assert request._all_token_ids == [0, 0, 7, 8, 0, 0, 0]
|
||||
assert request._output_token_ids == []
|
||||
assert request.num_computed_tokens == 4
|
||||
assert request.num_prompt_tokens == 7
|
||||
request.update_block_hashes.assert_called_once_with()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def build_adapter(monkeypatch, mocker: MockerFixture):
|
||||
def _build(
|
||||
@@ -252,6 +276,37 @@ def test_save_async_uses_confirmed_tokens_for_async_scheduler_watermark(build_ad
|
||||
assert len(adapter._pending_save_reqs) == 1
|
||||
|
||||
|
||||
def test_send_single_request_terminal_chunk_still_flushes_processor(build_adapter, monkeypatch):
|
||||
"""A terminal stop is not a segment boundary (#5383), but the producer-side
|
||||
processor must still receive the flush signal on the terminal chunk.
|
||||
Passing only ``is_segment_finished`` starved processors of their final
|
||||
accumulated payload once terminal stops stopped setting it (#5413: the
|
||||
downstream stage got ``meta.finished`` with the tail data missing).
|
||||
"""
|
||||
adapter, connector = build_adapter(stage_id=0)
|
||||
request = _req("req-terminal", RequestStatus.FINISHED_STOPPED, external_req_id="ext-terminal")
|
||||
|
||||
seen_flush_flags = []
|
||||
|
||||
def recording_processor(**kwargs):
|
||||
seen_flush_flags.append(kwargs["is_finished"])
|
||||
return OmniPayloadStruct(
|
||||
codes=CodesStruct(audio=torch.tensor([1, 2, 3], dtype=torch.long)),
|
||||
)
|
||||
|
||||
adapter.custom_process_next_stage_input_func = recording_processor
|
||||
monkeypatch.setattr(adapter, "cleanup", lambda *a, **kw: None)
|
||||
|
||||
adapter._send_single_request(
|
||||
{"multimodal_output": None, "request": request, "is_finished": True, "is_segment_finished": False}
|
||||
)
|
||||
|
||||
assert seen_flush_flags == [True]
|
||||
sent_payload = connector.put.call_args.kwargs["data"]
|
||||
assert bool(sent_payload.meta.finished.item()) is True
|
||||
assert bool(sent_payload.meta.is_segment_finished.item()) is False
|
||||
|
||||
|
||||
def test_send_single_request_struct_without_meta_does_not_crash(build_adapter, monkeypatch):
|
||||
"""Producer may return a struct with ``meta=None`` (e.g. payload that
|
||||
carries only ``embed`` or ``codes``). The sender's ``meta is not None``
|
||||
@@ -314,6 +369,23 @@ def test_send_single_request_struct_preserves_segment_finished(build_adapter, mo
|
||||
assert sent_payload.meta.is_segment_finished.item() is True
|
||||
|
||||
|
||||
def test_send_single_request_respects_processor_receiver_boundary(build_adapter, monkeypatch):
|
||||
adapter, connector = build_adapter(stage_id=1)
|
||||
request = _req("req-stream", RequestStatus.WAITING, external_req_id="ext-stream")
|
||||
|
||||
adapter.custom_process_next_stage_input_func = lambda **kwargs: OmniPayloadStruct(
|
||||
meta=MetaStruct(is_segment_finished=torch.tensor(False, dtype=torch.bool))
|
||||
)
|
||||
monkeypatch.setattr(adapter, "cleanup", lambda *a, **kw: None)
|
||||
|
||||
adapter._send_single_request(
|
||||
{"multimodal_output": None, "request": request, "is_finished": False, "is_segment_finished": True}
|
||||
)
|
||||
|
||||
sent_payload = connector.put.call_args.kwargs["data"]
|
||||
assert sent_payload.meta.is_segment_finished.item() is False
|
||||
|
||||
|
||||
def test_save_async_skips_stale_resumable_chunk_until_dedup_is_reset(build_adapter):
|
||||
adapter, _ = build_adapter(stage_id=1)
|
||||
request = _req("req-stream", RequestStatus.WAITING, external_req_id="ext-stream")
|
||||
@@ -463,7 +535,7 @@ def test_fifo_promotion(build_adapter):
|
||||
adapter.process_pending_chunks(waiting_queue, running_queue)
|
||||
|
||||
assert list(adapter._active_streams) == ["req-1", "req-2"]
|
||||
assert waiting_queue == reqs[2:]
|
||||
assert waiting_queue == []
|
||||
assert reqs[0].status == RequestStatus.WAITING_FOR_CHUNK
|
||||
assert reqs[1].status == RequestStatus.WAITING_FOR_CHUNK
|
||||
assert reqs[2].status == RequestStatus.WAITING
|
||||
@@ -473,6 +545,7 @@ def test_fifo_promotion(build_adapter):
|
||||
# (commit c4d95fd9 — otherwise the terminal chunk deadlocks at c=8 K=2).
|
||||
# Simulate it here so promotion can pick up the freed slot.
|
||||
adapter._evict_finished_active_streams({"req-1"})
|
||||
adapter.restore_queues(waiting_queue, running_queue)
|
||||
adapter.process_pending_chunks(waiting_queue, running_queue)
|
||||
|
||||
assert list(adapter._active_streams) == ["req-2", "req-3"]
|
||||
@@ -481,6 +554,22 @@ def test_fifo_promotion(build_adapter):
|
||||
assert reqs[2].status == RequestStatus.WAITING_FOR_CHUNK
|
||||
|
||||
|
||||
def test_non_active_waiting_request_is_held_off_scheduler(build_adapter):
|
||||
adapter, _ = build_adapter(stage_id=2, max_num_seqs=2, active_stream_window=1)
|
||||
active = _req("req-active", RequestStatus.WAITING)
|
||||
non_active = _req("req-non-active", RequestStatus.WAITING)
|
||||
waiting_queue = DummyWaitingQueue([active, non_active])
|
||||
running_queue = []
|
||||
|
||||
adapter.process_pending_chunks(waiting_queue, running_queue)
|
||||
|
||||
assert waiting_queue == []
|
||||
assert active.status == RequestStatus.WAITING_FOR_CHUNK
|
||||
assert non_active.status == RequestStatus.WAITING
|
||||
assert list(adapter._pending_load_reqs) == [active]
|
||||
assert list(adapter.waiting_for_chunk_waiting_requests) == [active, non_active]
|
||||
|
||||
|
||||
def test_legacy_k0(build_adapter):
|
||||
adapter, _ = build_adapter(stage_id=1, max_num_seqs=1, active_stream_window=0)
|
||||
waiting_req = _req("waiting", RequestStatus.WAITING)
|
||||
@@ -509,12 +598,13 @@ def test_finished_releases_slot(build_adapter):
|
||||
|
||||
adapter.process_pending_chunks(waiting_queue, running_queue)
|
||||
assert list(adapter._active_streams) == ["req-1"]
|
||||
assert waiting_queue == [req_2]
|
||||
assert waiting_queue == []
|
||||
|
||||
adapter.finished_requests.add("req-1")
|
||||
# Eviction is deferred to postprocess_scheduler_output in the runtime path
|
||||
# (commit c4d95fd9). Simulate it so promotion can pick up the freed slot.
|
||||
adapter._evict_finished_active_streams({"req-1"})
|
||||
adapter.restore_queues(waiting_queue, running_queue)
|
||||
adapter.process_pending_chunks(waiting_queue, running_queue)
|
||||
|
||||
assert list(adapter._active_streams) == ["req-2"]
|
||||
|
||||
@@ -83,7 +83,7 @@ def test_video_to_text(omni_runner, omni_runner_handler) -> None:
|
||||
@hardware_test(res={"cuda": "H100", "npu": "A2"}, num_cards=2)
|
||||
@pytest.mark.parametrize("omni_runner", test_params, indirect=True)
|
||||
def test_text_to_audio(omni_runner, omni_runner_handler) -> None:
|
||||
"""Test processing text, generating audio output through the talker token2wav path."""
|
||||
"""Test processing text and generating audio through Talker and Code2Wav."""
|
||||
request_config = {"prompts": get_question("text"), "modalities": ["audio"]}
|
||||
omni_runner_handler.send_omni_request(request_config)
|
||||
|
||||
|
||||
@@ -1174,10 +1174,13 @@ async def _send_listen_only_overlap_pair(
|
||||
before_second_listen = state.model_listen_count
|
||||
overlap_started_while_active = not state.response_done(first_response_id)
|
||||
model_unit_ready_while_active = False
|
||||
second_unit_event_index: int | None = None
|
||||
|
||||
def record_model_unit_ready() -> None:
|
||||
nonlocal model_unit_ready_while_active
|
||||
nonlocal model_unit_ready_while_active, second_unit_event_index
|
||||
model_unit_ready_while_active = not state.response_done(first_response_id)
|
||||
if model_unit_ready_while_active:
|
||||
second_unit_event_index = len(state.events)
|
||||
|
||||
await _send_pcm16(
|
||||
ws,
|
||||
@@ -1239,14 +1242,13 @@ async def _send_listen_only_overlap_pair(
|
||||
label=transcripts[1],
|
||||
)
|
||||
|
||||
second_speech_index = _nth_event_index(state, "input_audio_buffer.speech_started", 2)
|
||||
first_done_index = _event_index_for_response(state, "response.done", first_response_id)
|
||||
overlap_ok = (
|
||||
overlap_started_while_active
|
||||
and model_unit_ready_while_active
|
||||
and second_speech_index is not None
|
||||
and second_unit_event_index is not None
|
||||
and first_done_index is not None
|
||||
and second_speech_index < first_done_index
|
||||
and second_unit_event_index < first_done_index
|
||||
)
|
||||
return [first_response_id, second_response_id], ["speak", second_outcome], overlap_ok
|
||||
|
||||
@@ -1380,6 +1382,17 @@ async def run_demo(args: argparse.Namespace) -> dict[str, object]:
|
||||
await ws.send(json.dumps({"type": "session.close"}))
|
||||
await _wait_for(state, lambda: state.count("session.closed") > 0, timeout_s=20, label="session.closed")
|
||||
finally:
|
||||
if state.count("session.created") > 0 and state.count("session.closed") == 0:
|
||||
try:
|
||||
await ws.send(json.dumps({"type": "session.close"}))
|
||||
await _wait_for(
|
||||
state,
|
||||
lambda: state.count("session.closed") > 0,
|
||||
timeout_s=min(args.timeout_s, 5),
|
||||
label="session.closed during cleanup",
|
||||
)
|
||||
except (ConnectionClosed, TimeoutError):
|
||||
pass
|
||||
stop.set()
|
||||
reader.cancel()
|
||||
try:
|
||||
|
||||
@@ -19,7 +19,11 @@ SCRIPT_DIR = Path(__file__).resolve().parent
|
||||
if str(SCRIPT_DIR) not in sys.path:
|
||||
sys.path.insert(0, str(SCRIPT_DIR))
|
||||
|
||||
from minicpmo_realtime_duplex_scenarios import _url_with_model, run_demo # noqa: E402
|
||||
from minicpmo_realtime_duplex_scenarios import ( # noqa: E402
|
||||
_ref_audio_data_url,
|
||||
_url_with_model,
|
||||
run_demo,
|
||||
)
|
||||
|
||||
|
||||
def _with_resume_mode(url: str) -> str:
|
||||
@@ -114,7 +118,12 @@ async def _open_admission_session(
|
||||
args: argparse.Namespace,
|
||||
session_id: str,
|
||||
) -> tuple[object, dict[str, object]]:
|
||||
url = _url_with_model(args.url, args.model, session_id=session_id)
|
||||
url = _url_with_model(
|
||||
args.url,
|
||||
args.model,
|
||||
autostart=False if getattr(args, "ref_audio", None) else None,
|
||||
session_id=session_id,
|
||||
)
|
||||
ws = await websockets.connect(url, max_size=64 * 1024 * 1024)
|
||||
await ws.send(
|
||||
json.dumps(
|
||||
@@ -125,6 +134,7 @@ async def _open_admission_session(
|
||||
"model": args.model,
|
||||
"modalities": ["audio", "text"],
|
||||
"extra_body": {"minicpmo45_native_duplex": True},
|
||||
**({"ref_audio": _ref_audio_data_url(args.ref_audio)} if getattr(args, "ref_audio", None) else {}),
|
||||
},
|
||||
}
|
||||
)
|
||||
@@ -151,7 +161,12 @@ async def _admission_probe(args: argparse.Namespace, *, limit: int) -> dict[str,
|
||||
accepted.append(await _open_admission_session(args, f"{prefix}-accepted-{index}"))
|
||||
|
||||
overflow_id = f"{prefix}-overflow"
|
||||
overflow_url = _url_with_model(args.url, args.model, session_id=overflow_id)
|
||||
overflow_url = _url_with_model(
|
||||
args.url,
|
||||
args.model,
|
||||
autostart=False if getattr(args, "ref_audio", None) else None,
|
||||
session_id=overflow_id,
|
||||
)
|
||||
async with websockets.connect(overflow_url, max_size=64 * 1024 * 1024) as overflow:
|
||||
await overflow.send(
|
||||
json.dumps(
|
||||
@@ -162,6 +177,11 @@ async def _admission_probe(args: argparse.Namespace, *, limit: int) -> dict[str,
|
||||
"model": args.model,
|
||||
"modalities": ["audio", "text"],
|
||||
"extra_body": {"minicpmo45_native_duplex": True},
|
||||
**(
|
||||
{"ref_audio": _ref_audio_data_url(args.ref_audio)}
|
||||
if getattr(args, "ref_audio", None)
|
||||
else {}
|
||||
),
|
||||
},
|
||||
}
|
||||
)
|
||||
@@ -207,16 +227,27 @@ async def _resume_probe(
|
||||
session_id: str,
|
||||
expect_expired: bool = False,
|
||||
) -> dict[str, object]:
|
||||
url = _url_with_model(args.url, args.model, session_id=session_id)
|
||||
url = _url_with_model(
|
||||
args.url,
|
||||
args.model,
|
||||
autostart=False if getattr(args, "ref_audio", None) else None,
|
||||
session_id=session_id,
|
||||
)
|
||||
async with websockets.connect(url, max_size=64 * 1024 * 1024) as first:
|
||||
await first.send(
|
||||
json.dumps(
|
||||
{
|
||||
"type": "session.update",
|
||||
"session": {
|
||||
"session_id": session_id,
|
||||
"model": args.model,
|
||||
"modalities": ["audio", "text"],
|
||||
"extra_body": {"minicpmo45_native_duplex": True},
|
||||
**(
|
||||
{"ref_audio": _ref_audio_data_url(args.ref_audio)}
|
||||
if getattr(args, "ref_audio", None)
|
||||
else {}
|
||||
),
|
||||
},
|
||||
}
|
||||
)
|
||||
@@ -306,7 +337,12 @@ async def _takeover_probe(
|
||||
*,
|
||||
session_id: str,
|
||||
) -> dict[str, object]:
|
||||
url = _url_with_model(args.url, args.model, session_id=session_id)
|
||||
url = _url_with_model(
|
||||
args.url,
|
||||
args.model,
|
||||
autostart=False if getattr(args, "ref_audio", None) else None,
|
||||
session_id=session_id,
|
||||
)
|
||||
first = await websockets.connect(url, max_size=64 * 1024 * 1024)
|
||||
second = None
|
||||
try:
|
||||
@@ -315,9 +351,15 @@ async def _takeover_probe(
|
||||
{
|
||||
"type": "session.update",
|
||||
"session": {
|
||||
"session_id": session_id,
|
||||
"model": args.model,
|
||||
"modalities": ["audio", "text"],
|
||||
"extra_body": {"minicpmo45_native_duplex": True},
|
||||
**(
|
||||
{"ref_audio": _ref_audio_data_url(args.ref_audio)}
|
||||
if getattr(args, "ref_audio", None)
|
||||
else {}
|
||||
),
|
||||
},
|
||||
}
|
||||
)
|
||||
@@ -451,20 +493,7 @@ async def run_multi_session(args: argparse.Namespace) -> dict[str, object]:
|
||||
args,
|
||||
session_id=f"takeover-{args.takeover_session_index}-{uuid.uuid4().hex}",
|
||||
)
|
||||
expiry_result = None
|
||||
if args.expire_session_index is not None:
|
||||
if not 0 <= args.expire_session_index < args.sessions:
|
||||
raise ValueError("--expire-session-index is outside the session range")
|
||||
expiry_result = await _resume_probe(
|
||||
args,
|
||||
session_id=f"expire-{args.expire_session_index}-{uuid.uuid4().hex}",
|
||||
expect_expired=True,
|
||||
)
|
||||
admission_result = (
|
||||
await _admission_probe(args, limit=args.verify_admission_limit)
|
||||
if args.verify_admission_limit is not None
|
||||
else None
|
||||
)
|
||||
lifecycle_result = await run_lifecycle_probes(args)
|
||||
|
||||
session_results = await asyncio.gather(
|
||||
*(run_demo(_demo_args(args, index)) for index in range(args.sessions)),
|
||||
@@ -487,16 +516,15 @@ async def run_multi_session(args: argparse.Namespace) -> dict[str, object]:
|
||||
and semantic_isolation_ok
|
||||
and (resume_result is None or resume_result.get("ok") is True)
|
||||
and (takeover_result is None or takeover_result.get("ok") is True)
|
||||
and (expiry_result is None or expiry_result.get("ok") is True)
|
||||
and (admission_result is None or admission_result.get("ok") is True)
|
||||
and lifecycle_result["ok"] is True
|
||||
),
|
||||
"session_count": args.sessions,
|
||||
"identity_isolation_ok": identity_isolation_ok,
|
||||
"semantic_isolation_ok": semantic_isolation_ok,
|
||||
"resume": resume_result,
|
||||
"takeover": takeover_result,
|
||||
"expiry": expiry_result,
|
||||
"admission": admission_result,
|
||||
"expiry": lifecycle_result["expiry"],
|
||||
"admission": lifecycle_result["admission"],
|
||||
"failures": failures,
|
||||
"sessions": completed,
|
||||
}
|
||||
@@ -507,6 +535,34 @@ async def run_multi_session(args: argparse.Namespace) -> dict[str, object]:
|
||||
return result
|
||||
|
||||
|
||||
async def run_lifecycle_probes(args: argparse.Namespace) -> dict[str, object]:
|
||||
"""Run expiry and admission probes without requiring model output."""
|
||||
if args.sessions < 1:
|
||||
raise ValueError("--sessions must be positive")
|
||||
expiry_result = None
|
||||
if args.expire_session_index is not None:
|
||||
if not 0 <= args.expire_session_index < args.sessions:
|
||||
raise ValueError("--expire-session-index is outside the session range")
|
||||
expiry_result = await _resume_probe(
|
||||
args,
|
||||
session_id=f"expire-{args.expire_session_index}-{uuid.uuid4().hex}",
|
||||
expect_expired=True,
|
||||
)
|
||||
admission_result = (
|
||||
await _admission_probe(args, limit=args.verify_admission_limit)
|
||||
if args.verify_admission_limit is not None
|
||||
else None
|
||||
)
|
||||
return {
|
||||
"ok": (
|
||||
(expiry_result is None or expiry_result.get("ok") is True)
|
||||
and (admission_result is None or admission_result.get("ok") is True)
|
||||
),
|
||||
"expiry": expiry_result,
|
||||
"admission": admission_result,
|
||||
}
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--url", default="ws://127.0.0.1:8113/v1/realtime?duplex=1")
|
||||
|
||||
@@ -2,8 +2,8 @@
|
||||
|
||||
This E2E driver runs the public ``realtime_duplex_demo.py`` against a live
|
||||
duplex backend. Arbitrary audio defaults to the model-policy lifecycle contract.
|
||||
The stronger response-required mode binds the two-response contract to a known
|
||||
input checksum and expected second response.
|
||||
The stronger response-required mode binds the multi-response contract to a
|
||||
known input checksum and expected follow-up response.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -39,6 +39,17 @@ def _validate_input_sha256(path: Path, expected: str) -> str:
|
||||
return actual
|
||||
|
||||
|
||||
def _input_duration_s(path: Path) -> float:
|
||||
with wave.open(str(path), "rb") as wav_file:
|
||||
return wav_file.getnframes() / wav_file.getframerate()
|
||||
|
||||
|
||||
def _client_process_timeout_s(input_wav: Path, protocol_timeout_s: float) -> float:
|
||||
# The child first streams the WAV in real time, then can independently
|
||||
# exhaust its post-commit and session-close protocol waits.
|
||||
return _input_duration_s(input_wav) + 2 * protocol_timeout_s + 30.0
|
||||
|
||||
|
||||
def _read_json(path: Path) -> dict[str, object]:
|
||||
if not path.is_file():
|
||||
return {}
|
||||
@@ -211,7 +222,7 @@ def summarize_artifacts(
|
||||
validation_mode: str,
|
||||
min_responses: int,
|
||||
min_audio_deltas_per_response: int,
|
||||
expect_second_response_substring: str | None,
|
||||
expect_followup_response_substring: str | None,
|
||||
) -> dict[str, object]:
|
||||
if validation_mode not in {"model-policy", "response-required"}:
|
||||
raise ValueError(f"unsupported validation mode: {validation_mode}")
|
||||
@@ -274,12 +285,20 @@ def summarize_artifacts(
|
||||
)
|
||||
final_listen_after_commit = commit_index is not None and any(index > commit_index for index in listen_indices)
|
||||
transcript = "".join(str(summary.get("transcript") or "") for summary in response_summaries)
|
||||
second_response_transcript = (
|
||||
str(response_summaries[1].get("transcript") or "") if len(response_summaries) >= 2 else ""
|
||||
followup_response_transcripts = [
|
||||
str(summary.get("transcript") or "")
|
||||
for summary in response_summaries[1:]
|
||||
if commit_index is not None
|
||||
and isinstance(summary.get("created_index"), int)
|
||||
and summary["created_index"] < commit_index
|
||||
]
|
||||
followup_response_transcript_ok = any(
|
||||
bool(_normalize_text(transcript)) for transcript in followup_response_transcripts
|
||||
)
|
||||
followup_response_transcript_expectation_ok = not expect_followup_response_substring or any(
|
||||
_normalize_text(expect_followup_response_substring) in _normalize_text(transcript)
|
||||
for transcript in followup_response_transcripts
|
||||
)
|
||||
second_response_transcript_expectation_ok = not expect_second_response_substring or _normalize_text(
|
||||
expect_second_response_substring
|
||||
) in _normalize_text(second_response_transcript)
|
||||
error_events = [
|
||||
event
|
||||
for event in events
|
||||
@@ -310,9 +329,8 @@ def summarize_artifacts(
|
||||
mode_contract_ok = validation_mode == "model-policy" or (
|
||||
second_response_before_final_commit
|
||||
and listen_before_first_response
|
||||
and listen_between_responses
|
||||
and listen_after_last_done
|
||||
and second_response_transcript_expectation_ok
|
||||
and followup_response_transcript_ok
|
||||
)
|
||||
ok = common_contract_ok and mode_contract_ok
|
||||
return {
|
||||
@@ -337,9 +355,10 @@ def summarize_artifacts(
|
||||
"listen_after_response_before_commit": listen_after_response_before_commit,
|
||||
"final_listen_after_commit": final_listen_after_commit,
|
||||
"transcript": transcript,
|
||||
"second_response_transcript": second_response_transcript,
|
||||
"expect_second_response_substring": expect_second_response_substring,
|
||||
"second_response_transcript_expectation_ok": second_response_transcript_expectation_ok,
|
||||
"followup_response_transcripts": followup_response_transcripts,
|
||||
"followup_response_transcript_ok": followup_response_transcript_ok,
|
||||
"expect_followup_response_substring": expect_followup_response_substring,
|
||||
"followup_response_transcript_expectation_ok": followup_response_transcript_expectation_ok,
|
||||
"error_count": len(error_events),
|
||||
"cancelled_count": cancelled_count,
|
||||
"compact_sequence": _compact_sequence(events),
|
||||
@@ -388,7 +407,7 @@ async def run_soft_interrupt(args: argparse.Namespace) -> dict[str, object]:
|
||||
try:
|
||||
stdout_bytes, stderr_bytes = await asyncio.wait_for(
|
||||
process.communicate(),
|
||||
timeout=args.timeout_s + 30.0,
|
||||
timeout=_client_process_timeout_s(input_wav, args.timeout_s),
|
||||
)
|
||||
except TimeoutError:
|
||||
process.kill()
|
||||
@@ -406,7 +425,7 @@ async def run_soft_interrupt(args: argparse.Namespace) -> dict[str, object]:
|
||||
validation_mode=args.validation_mode,
|
||||
min_responses=args.min_responses,
|
||||
min_audio_deltas_per_response=args.min_audio_deltas_per_response,
|
||||
expect_second_response_substring=args.expect_second_response_substring,
|
||||
expect_followup_response_substring=args.expect_followup_response_substring,
|
||||
)
|
||||
summary.update(
|
||||
{
|
||||
@@ -444,14 +463,12 @@ def parse_args() -> argparse.Namespace:
|
||||
parser.add_argument("--min-responses", type=int, default=2)
|
||||
parser.add_argument("--min-audio-deltas-per-response", type=int, default=2)
|
||||
parser.add_argument("--input-sha256")
|
||||
parser.add_argument("--expect-second-response-substring")
|
||||
parser.add_argument("--expect-followup-response-substring")
|
||||
args = parser.parse_args()
|
||||
if args.validation_mode == "response-required" and args.min_responses < 2:
|
||||
parser.error("--min-responses must be at least 2")
|
||||
if args.validation_mode == "response-required" and not args.input_sha256:
|
||||
parser.error("--input-sha256 is required in response-required mode")
|
||||
if args.validation_mode == "response-required" and not args.expect_second_response_substring:
|
||||
parser.error("--expect-second-response-substring is required in response-required mode")
|
||||
if args.min_audio_deltas_per_response < 1:
|
||||
parser.error("--min-audio-deltas-per-response must be positive")
|
||||
return args
|
||||
|
||||
@@ -1,9 +1,7 @@
|
||||
"""
|
||||
E2E Online tests for MiniCPM-o 4.5 model with multimodal input and audio / text output.
|
||||
"""E2E online tests for MiniCPM-o 4.5 multimodal input and audio/text output.
|
||||
|
||||
MiniCPM-o 4.5 has ``async_chunk: false``, ``max_num_seqs: 1`` on both stages,
|
||||
and the vocoder runs in-process inside the talker stage rather than as a separate
|
||||
Code2Wav stage.
|
||||
The batching deployment uses async chunk transfer across separate Thinker,
|
||||
Talker, and Code2Wav stages.
|
||||
"""
|
||||
|
||||
import os
|
||||
@@ -26,10 +24,7 @@ test_params = [
|
||||
model=_MODEL,
|
||||
stage_config_path=_CI_DEPLOY,
|
||||
use_stage_cli=True,
|
||||
server_args=[
|
||||
"--trust-remote-code",
|
||||
"--no-async-chunk",
|
||||
],
|
||||
server_args=["--trust-remote-code"],
|
||||
),
|
||||
id="default",
|
||||
)
|
||||
@@ -99,7 +94,7 @@ def test_text_to_text_001(omni_server, openai_client) -> None:
|
||||
def test_text_to_audio_001(omni_server, openai_client) -> None:
|
||||
"""
|
||||
Test text-only input generating text + audio output via OpenAI API.
|
||||
This exercises the talker TTS region detection and in-process token2wav vocoder.
|
||||
This exercises Talker TTS region detection and the Code2Wav stage.
|
||||
Deploy Setting: default 2GPU
|
||||
Input Modal: text
|
||||
Output Modal: text + audio
|
||||
@@ -133,13 +128,14 @@ def test_audio_to_text_audio_001(omni_server, openai_client) -> None:
|
||||
messages = dummy_messages_from_mix_data(
|
||||
system_prompt=get_system_prompt(),
|
||||
audio_data_url=audio_data_url,
|
||||
content_text=get_prompt("mix"),
|
||||
content_text=get_prompt(),
|
||||
)
|
||||
|
||||
request_config = {
|
||||
"model": omni_server.model,
|
||||
"messages": messages,
|
||||
"stream": True,
|
||||
"key_words": {"text": ["Beijing"]},
|
||||
}
|
||||
|
||||
openai_client.send_omni_request(request_config, request_num=get_max_batch_size())
|
||||
|
||||
@@ -0,0 +1,147 @@
|
||||
"""CI coverage for the MiniCPM-o 4.5 native-duplex Realtime API."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
import websockets
|
||||
|
||||
from tests.e2e.online_serving.minicpmo_realtime_duplex_scenarios import (
|
||||
_ref_audio_data_url,
|
||||
run_demo,
|
||||
)
|
||||
from tests.e2e.online_serving.run_minicpmo_realtime_duplex_multi_session import (
|
||||
run_multi_session,
|
||||
)
|
||||
from tests.helpers.mark import hardware_test
|
||||
from tests.helpers.minicpmo_4_5_duplex import (
|
||||
CORE_SERVER_PARAMS,
|
||||
SERVER_PARAMS,
|
||||
demo_args,
|
||||
multi_session_args,
|
||||
realtime_url,
|
||||
resolve_ref_audio,
|
||||
validated_input_wav,
|
||||
)
|
||||
from vllm_omni.experimental.fullduplex.client import build_realtime_url
|
||||
|
||||
pytestmark = pytest.mark.omni
|
||||
|
||||
|
||||
async def _receive_protocol_events(ws, required_types: set[str], *, timeout_s: float) -> list[dict[str, object]]:
|
||||
async def receive() -> list[dict[str, object]]:
|
||||
events: list[dict[str, object]] = []
|
||||
seen: set[str] = set()
|
||||
while not required_types.issubset(seen):
|
||||
raw = await ws.recv()
|
||||
if not isinstance(raw, str):
|
||||
continue
|
||||
event = json.loads(raw)
|
||||
if not isinstance(event, dict):
|
||||
continue
|
||||
events.append(event)
|
||||
event_type = event.get("type")
|
||||
if event_type == "error":
|
||||
raise AssertionError(f"WebSocket protocol smoke received an error: {event}")
|
||||
if isinstance(event_type, str):
|
||||
seen.add(event_type)
|
||||
return events
|
||||
|
||||
return await asyncio.wait_for(receive(), timeout=timeout_s)
|
||||
|
||||
|
||||
async def _run_protocol_smoke(*, url: str, model: str, ref_audio: Path) -> list[dict[str, object]]:
|
||||
session_id = f"duplex-ci-protocol-{uuid.uuid4().hex}"
|
||||
websocket_url = build_realtime_url(url, model, autostart=False, session_id=session_id)
|
||||
async with websockets.connect(websocket_url, max_size=64 * 1024 * 1024) as ws:
|
||||
await ws.send(
|
||||
json.dumps(
|
||||
{
|
||||
"type": "session.update",
|
||||
"session": {
|
||||
"session_id": session_id,
|
||||
"model": model,
|
||||
"modalities": ["audio", "text"],
|
||||
"ref_audio": _ref_audio_data_url(str(ref_audio)),
|
||||
"extra_body": {"minicpmo45_native_duplex": True},
|
||||
},
|
||||
}
|
||||
)
|
||||
)
|
||||
events = await _receive_protocol_events(
|
||||
ws,
|
||||
{"session.created", "session.updated"},
|
||||
timeout_s=60,
|
||||
)
|
||||
await ws.send(json.dumps({"type": "session.close"}))
|
||||
events.extend(await _receive_protocol_events(ws, {"session.closed"}, timeout_s=60))
|
||||
return events
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=2)
|
||||
@pytest.mark.parametrize("omni_server", CORE_SERVER_PARAMS, indirect=True)
|
||||
def test_duplex_websocket_protocol_smoke(omni_server, model_prefix: str) -> None:
|
||||
ref_audio = resolve_ref_audio(model_prefix)
|
||||
events = asyncio.run(
|
||||
_run_protocol_smoke(
|
||||
url=realtime_url(omni_server),
|
||||
model=omni_server.model,
|
||||
ref_audio=ref_audio,
|
||||
)
|
||||
)
|
||||
event_types = [event.get("type") for event in events]
|
||||
assert "session.created" in event_types
|
||||
assert "session.updated" in event_types
|
||||
assert event_types[-1] == "session.closed"
|
||||
|
||||
|
||||
@pytest.mark.advanced_model
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=2)
|
||||
@pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True)
|
||||
def test_duplex_single_session_response_required(omni_server, model_prefix: str, tmp_path: Path) -> None:
|
||||
result = asyncio.run(
|
||||
run_demo(
|
||||
demo_args(
|
||||
omni_server=omni_server,
|
||||
input_wav=validated_input_wav(),
|
||||
ref_audio=resolve_ref_audio(model_prefix),
|
||||
output_dir=tmp_path / "single_session",
|
||||
)
|
||||
)
|
||||
)
|
||||
assert result["ok"] is True
|
||||
assert result["audio_delta_count"] > 0
|
||||
assert result["done_count"] == 1
|
||||
assert result["error_count"] == 0
|
||||
assert result["all_audio_responses_have_transcript"] is True
|
||||
assert result["transcript_delta_done_ok"] is True
|
||||
|
||||
|
||||
@pytest.mark.advanced_model
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=2)
|
||||
@pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True)
|
||||
def test_duplex_two_sessions_resume_and_takeover(omni_server, model_prefix: str, tmp_path: Path) -> None:
|
||||
result = asyncio.run(
|
||||
run_multi_session(
|
||||
multi_session_args(
|
||||
omni_server=omni_server,
|
||||
input_wav=validated_input_wav(),
|
||||
ref_audio=resolve_ref_audio(model_prefix),
|
||||
output_dir=tmp_path / "multi_session",
|
||||
response_required=True,
|
||||
)
|
||||
)
|
||||
)
|
||||
assert result["ok"] is True
|
||||
assert result["session_count"] == 2
|
||||
assert result["resume"]["ok"] is True
|
||||
assert result["takeover"]["ok"] is True
|
||||
assert not result["failures"]
|
||||
assert all(session["audio_delta_count"] > 0 for session in result["sessions"])
|
||||
assert all(session["done_count"] == 1 for session in result["sessions"])
|
||||
assert all(session["error_count"] == 0 for session in result["sessions"])
|
||||
@@ -0,0 +1,86 @@
|
||||
"""Nightly lifecycle coverage for the MiniCPM-o 4.5 native-duplex API."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.e2e.online_serving.run_minicpmo_realtime_duplex_multi_session import (
|
||||
run_lifecycle_probes,
|
||||
)
|
||||
from tests.e2e.online_serving.run_minicpmo_realtime_duplex_soft_interrupt import (
|
||||
run_soft_interrupt,
|
||||
)
|
||||
from tests.helpers.mark import hardware_test
|
||||
from tests.helpers.minicpmo_4_5_duplex import (
|
||||
SERVER_PARAMS,
|
||||
SOFT_INTERRUPT_SHA256,
|
||||
multi_session_args,
|
||||
realtime_url,
|
||||
resolve_ref_audio,
|
||||
validated_input_wav,
|
||||
validated_soft_interrupt_wav,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.full_model, pytest.mark.omni]
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=2)
|
||||
@pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True)
|
||||
def test_duplex_admission_and_expiry_reaper(omni_server, model_prefix: str, tmp_path: Path) -> None:
|
||||
args = multi_session_args(
|
||||
omni_server=omni_server,
|
||||
input_wav=validated_input_wav(),
|
||||
ref_audio=resolve_ref_audio(model_prefix),
|
||||
output_dir=tmp_path / "admission_expiry",
|
||||
response_required=False,
|
||||
)
|
||||
args.sessions = 1
|
||||
args.disconnect_session_index = None
|
||||
args.takeover_session_index = None
|
||||
args.expire_session_index = 0
|
||||
args.verify_admission_limit = 2
|
||||
result = asyncio.run(run_lifecycle_probes(args))
|
||||
|
||||
assert result["ok"] is True, json.dumps(result, ensure_ascii=False, indent=2)
|
||||
assert result["expiry"]["ok"] is True
|
||||
assert result["expiry"]["error_code"] == "session_resume_expired"
|
||||
assert result["admission"]["ok"] is True
|
||||
assert result["admission"]["overflow_error_code"] == "resource_exhausted"
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=2)
|
||||
@pytest.mark.parametrize("omni_server", SERVER_PARAMS, indirect=True)
|
||||
def test_duplex_soft_interrupt(omni_server, model_prefix: str, tmp_path: Path) -> None:
|
||||
input_wav = validated_soft_interrupt_wav()
|
||||
result = asyncio.run(
|
||||
run_soft_interrupt(
|
||||
SimpleNamespace(
|
||||
url=realtime_url(omni_server),
|
||||
model=omni_server.model,
|
||||
input_wav=str(input_wav),
|
||||
ref_audio=str(resolve_ref_audio(model_prefix)),
|
||||
output_dir=str(tmp_path / "soft_interrupt"),
|
||||
summary_output=None,
|
||||
chunk_ms=200,
|
||||
timeout_s=180.0,
|
||||
require_audio=True,
|
||||
no_realtime_pacing=False,
|
||||
validation_mode="response-required",
|
||||
min_responses=2,
|
||||
min_audio_deltas_per_response=2,
|
||||
input_sha256=SOFT_INTERRUPT_SHA256,
|
||||
expect_followup_response_substring=None,
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
assert result["ok"] is True, json.dumps(result, ensure_ascii=False, indent=2)
|
||||
assert result["error_count"] == 0
|
||||
assert result["response_lifecycle_ok"] is True
|
||||
assert result["response_audio_contract_ok"] is True
|
||||
assert result["followup_response_transcript_ok"] is True
|
||||
@@ -1,7 +1,7 @@
|
||||
"""
|
||||
E2E Online expansion tests for MiniCPM-o 4.5 covering modality combinations,
|
||||
in-process token2wav vocoder behavior, max_num_seqs=1 serial execution,
|
||||
and edge-case validation.
|
||||
"""E2E online expansion tests for MiniCPM-o 4.5.
|
||||
|
||||
These cover modality combinations, separate Code2Wav behavior, sequential
|
||||
request isolation, longer audio output, and Chinese speech.
|
||||
"""
|
||||
|
||||
import os
|
||||
@@ -94,10 +94,8 @@ def test_text_video_to_text_001(omni_server, openai_client) -> None:
|
||||
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
|
||||
def test_sequential_requests_independent(omni_server, openai_client) -> None:
|
||||
"""
|
||||
Verify that sequential requests produce independent results and state
|
||||
does not leak between requests. MiniCPM-o 4.5 has max_num_seqs=1 so
|
||||
requests are processed one-at-a-time; the second request must not
|
||||
receive the first request's audio or token state.
|
||||
Verify that sequential requests produce independent results and that the
|
||||
second request does not receive the first request's audio or token state.
|
||||
Deploy Setting: default 2GPU
|
||||
Input Modal: text (two different prompts)
|
||||
Output Modal: text + audio (both)
|
||||
@@ -138,7 +136,7 @@ def test_sequential_requests_independent(omni_server, openai_client) -> None:
|
||||
def test_text_to_audio_long_output_001(omni_server, openai_client) -> None:
|
||||
"""
|
||||
Test text input generating a longer audio output to exercise the
|
||||
token2wav decoder across multiple frames.
|
||||
Code2Wav stage across multiple frames.
|
||||
Deploy Setting: default 2GPU
|
||||
Input Modal: text (longer prompt)
|
||||
Output Modal: text + audio
|
||||
@@ -146,14 +144,14 @@ def test_text_to_audio_long_output_001(omni_server, openai_client) -> None:
|
||||
"""
|
||||
messages = dummy_messages_from_mix_data(
|
||||
system_prompt=get_system_prompt(),
|
||||
content_text="Tell me a short story about a cat in about 50 words.",
|
||||
content_text="What is the capital of China? Answer in 40 words.",
|
||||
)
|
||||
|
||||
request_config = {
|
||||
"model": omni_server.model,
|
||||
"messages": messages,
|
||||
"stream": True,
|
||||
"key_words": {"audio": ["cat"]},
|
||||
"key_words": {"audio": ["Beijing"]},
|
||||
}
|
||||
|
||||
openai_client.send_omni_request(request_config, request_num=get_max_batch_size())
|
||||
|
||||
+37
-28
@@ -12,14 +12,11 @@ import pytest
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
DEMO_PATH = Path(__file__).resolve().parents[1] / "e2e/online_serving/minicpmo_realtime_duplex_scenarios.py"
|
||||
MULTI_DEMO_PATH = (
|
||||
Path(__file__).resolve().parents[1] / "e2e/online_serving/run_minicpmo_realtime_duplex_multi_session.py"
|
||||
)
|
||||
PAIR_DEMO_PATH = Path(__file__).resolve().parents[1] / "e2e/online_serving/run_minicpmo_realtime_duplex_demo_pair.py"
|
||||
SOFT_INTERRUPT_DEMO_PATH = (
|
||||
Path(__file__).resolve().parents[1] / "e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py"
|
||||
)
|
||||
DRIVER_DIR = Path(__file__).resolve().parent
|
||||
DEMO_PATH = DRIVER_DIR / "minicpmo_realtime_duplex_scenarios.py"
|
||||
MULTI_DEMO_PATH = DRIVER_DIR / "run_minicpmo_realtime_duplex_multi_session.py"
|
||||
PAIR_DEMO_PATH = DRIVER_DIR / "run_minicpmo_realtime_duplex_demo_pair.py"
|
||||
SOFT_INTERRUPT_DEMO_PATH = DRIVER_DIR / "run_minicpmo_realtime_duplex_soft_interrupt.py"
|
||||
|
||||
|
||||
def _load_demo_module():
|
||||
@@ -338,14 +335,7 @@ def test_realtime_duplex_soft_interrupt_accepts_explicit_ref_audio(monkeypatch):
|
||||
assert args.validation_mode == "model-policy"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"extra_args",
|
||||
[
|
||||
["--expect-second-response-substring", "一加一等于二"],
|
||||
["--input-sha256", "a" * 64],
|
||||
],
|
||||
)
|
||||
def test_realtime_duplex_soft_interrupt_response_required_needs_bound_fixture(monkeypatch, extra_args):
|
||||
def test_realtime_duplex_soft_interrupt_response_required_needs_bound_fixture(monkeypatch):
|
||||
demo = _load_soft_interrupt_demo_module()
|
||||
monkeypatch.setattr(
|
||||
demo.sys,
|
||||
@@ -360,7 +350,6 @@ def test_realtime_duplex_soft_interrupt_response_required_needs_bound_fixture(mo
|
||||
"ref.wav",
|
||||
"--validation-mode",
|
||||
"response-required",
|
||||
*extra_args,
|
||||
],
|
||||
)
|
||||
|
||||
@@ -412,7 +401,6 @@ def test_realtime_duplex_soft_interrupt_accepts_multi_delta_handoff_sequence(tmp
|
||||
"_client_received_at_s": 2.6,
|
||||
},
|
||||
{"type": "response.done", "response_id": first_response_id, "_client_received_at_s": 3.0},
|
||||
{"type": "response.listen", "_client_received_at_s": 3.2},
|
||||
{"type": "response.created", "response": {"id": second_response_id}, "_client_received_at_s": 4.0},
|
||||
{"type": "response.speak", "response_id": second_response_id, "_client_received_at_s": 4.0},
|
||||
{
|
||||
@@ -458,10 +446,11 @@ def test_realtime_duplex_soft_interrupt_accepts_multi_delta_handoff_sequence(tmp
|
||||
validation_mode="response-required",
|
||||
min_responses=2,
|
||||
min_audio_deltas_per_response=2,
|
||||
expect_second_response_substring="一加一等于二",
|
||||
expect_followup_response_substring="一加一等于二",
|
||||
)
|
||||
|
||||
assert summary["ok"] is True
|
||||
assert summary["listen_between_responses"] is False
|
||||
assert summary["second_response_before_final_commit"] is True
|
||||
assert summary["final_listen_after_commit"] is True
|
||||
assert summary["response_audio_contract_ok"] is True
|
||||
@@ -498,7 +487,7 @@ def test_realtime_duplex_soft_interrupt_model_policy_accepts_single_response(tmp
|
||||
validation_mode="model-policy",
|
||||
min_responses=2,
|
||||
min_audio_deltas_per_response=2,
|
||||
expect_second_response_substring=None,
|
||||
expect_followup_response_substring=None,
|
||||
)
|
||||
|
||||
assert summary["ok"] is True
|
||||
@@ -537,14 +526,14 @@ def test_realtime_duplex_soft_interrupt_response_required_rejects_single_respons
|
||||
validation_mode="response-required",
|
||||
min_responses=2,
|
||||
min_audio_deltas_per_response=2,
|
||||
expect_second_response_substring="一加一等于二",
|
||||
expect_followup_response_substring="一加一等于二",
|
||||
)
|
||||
|
||||
assert summary["ok"] is False
|
||||
assert summary["enough_responses"] is False
|
||||
|
||||
|
||||
def test_realtime_duplex_soft_interrupt_checks_expected_text_in_second_response(tmp_path):
|
||||
def test_realtime_duplex_soft_interrupt_reports_text_expectation_without_gating(tmp_path):
|
||||
demo = _load_soft_interrupt_demo_module()
|
||||
output = tmp_path / "second_response_text"
|
||||
output.mkdir()
|
||||
@@ -611,11 +600,32 @@ def test_realtime_duplex_soft_interrupt_checks_expected_text_in_second_response(
|
||||
validation_mode="response-required",
|
||||
min_responses=2,
|
||||
min_audio_deltas_per_response=2,
|
||||
expect_second_response_substring="一加一等于二",
|
||||
expect_followup_response_substring="一加一等于二",
|
||||
)
|
||||
|
||||
assert summary["ok"] is False
|
||||
assert summary["second_response_transcript_expectation_ok"] is False
|
||||
assert summary["ok"] is True
|
||||
assert summary["followup_response_transcript_ok"] is True
|
||||
assert summary["followup_response_transcript_expectation_ok"] is False
|
||||
|
||||
events = [
|
||||
event
|
||||
for event in events
|
||||
if event.get("response_id") != second_response_id or event["type"] != "response.audio_transcript.delta"
|
||||
]
|
||||
(output / "events.jsonl").write_text(
|
||||
"".join(demo.json.dumps(event) + "\n" for event in events),
|
||||
encoding="utf-8",
|
||||
)
|
||||
missing_transcript_summary = demo.summarize_artifacts(
|
||||
output_dir=output,
|
||||
validation_mode="response-required",
|
||||
min_responses=2,
|
||||
min_audio_deltas_per_response=2,
|
||||
expect_followup_response_substring="一加一等于二",
|
||||
)
|
||||
|
||||
assert missing_transcript_summary["ok"] is False
|
||||
assert missing_transcript_summary["followup_response_transcript_ok"] is False
|
||||
|
||||
|
||||
def test_realtime_duplex_multi_session_resume_url_disables_autostart():
|
||||
@@ -1523,7 +1533,7 @@ def test_realtime_duplex_demo_waits_at_each_model_unit_and_stops_after_speak():
|
||||
assert len(ws.messages) == 10
|
||||
|
||||
|
||||
def test_realtime_duplex_demo_listen_only_overlap_sends_next_turn_before_first_done(monkeypatch):
|
||||
def test_realtime_duplex_demo_listen_only_overlap_accepts_silence_unit_before_first_done(monkeypatch):
|
||||
demo = _load_demo_module()
|
||||
state = demo.DemoState()
|
||||
send_calls = 0
|
||||
@@ -1548,7 +1558,6 @@ def test_realtime_duplex_demo_listen_only_overlap_sends_next_turn_before_first_d
|
||||
del args
|
||||
assert not state.response_done("resp-first")
|
||||
kwargs["on_model_unit_ready"]()
|
||||
state.add({"type": "input_audio_buffer.speech_started", "turn": 2})
|
||||
_add_response_transcript(state, "resp-first", transcript="第一轮完成", audio=False)
|
||||
state.add({"type": "response.created", "response": {"id": "resp-second"}})
|
||||
_add_response_transcript(state, "resp-second", transcript="第二轮完成", audio=True)
|
||||
@@ -1568,7 +1577,7 @@ def test_realtime_duplex_demo_listen_only_overlap_sends_next_turn_before_first_d
|
||||
ws,
|
||||
state,
|
||||
b"\x00\x00",
|
||||
b"\x01\x00",
|
||||
b"\x00\x00",
|
||||
transcripts=("first", "second"),
|
||||
durations_ms=(None, None),
|
||||
chunk_ms=200,
|
||||
@@ -18,6 +18,7 @@ def test_chunk_accumulation_policy_replaces_snapshots_and_drains_delta_state():
|
||||
accumulated = MultimodalPayload.from_dict(
|
||||
{
|
||||
"audio": torch.tensor([1.0]),
|
||||
"meta.segment_end": torch.tensor([0]),
|
||||
"meta.tts_is_last_chunk": torch.tensor([0]),
|
||||
"meta.turn_end": torch.tensor([0]),
|
||||
"meta.stable_request_value": "keep",
|
||||
@@ -26,6 +27,7 @@ def test_chunk_accumulation_policy_replaces_snapshots_and_drains_delta_state():
|
||||
incoming = MultimodalPayload.from_dict(
|
||||
{
|
||||
"audio": torch.tensor([2.0]),
|
||||
"meta.segment_end": torch.tensor([1]),
|
||||
"meta.tts_is_last_chunk": torch.tensor([1]),
|
||||
"meta.turn_end": torch.tensor([1]),
|
||||
}
|
||||
@@ -41,6 +43,7 @@ def test_chunk_accumulation_policy_replaces_snapshots_and_drains_delta_state():
|
||||
drain_delta_payload(merged)
|
||||
|
||||
assert "audio" not in merged
|
||||
assert "meta.segment_end" not in merged
|
||||
assert "meta.tts_is_last_chunk" not in merged
|
||||
assert "meta.turn_end" not in merged
|
||||
assert merged.metadata["meta.stable_request_value"] == "keep"
|
||||
@@ -1813,6 +1813,56 @@ async def test_stage_pool_submit_update_refreshes_output_processor_state() -> No
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_streaming_update_unknown_request_is_dropped() -> None:
|
||||
"""A streaming update for an untracked request must not resurrect it.
|
||||
|
||||
Updates always follow the first-chunk add_request through the same
|
||||
ordered submission queue, so an unknown id can only mean the request
|
||||
already finished or was aborted (e.g. the client disconnected).
|
||||
Falling back to add_request created headless sessions that kept
|
||||
cycling through the stages (issue #4271).
|
||||
"""
|
||||
|
||||
class RecordingPool:
|
||||
def __init__(self) -> None:
|
||||
self.calls: list[tuple[str, Any]] = []
|
||||
|
||||
async def submit_update(self, request_id, req_state, request, *, prompt_text=None) -> int:
|
||||
self.calls.append((request_id, prompt_text))
|
||||
return 0
|
||||
|
||||
pool = RecordingPool()
|
||||
orchestrator = object.__new__(Orchestrator)
|
||||
orchestrator.async_chunk = False
|
||||
orchestrator.request_states = {}
|
||||
orchestrator.stage_pools = [pool]
|
||||
add_request_calls: list[str] = []
|
||||
|
||||
async def record_add_request(msg) -> None:
|
||||
add_request_calls.append(msg.request_id)
|
||||
|
||||
orchestrator._handle_add_request = record_add_request
|
||||
|
||||
await orchestrator._handle_streaming_update(
|
||||
StageSubmissionMessage(
|
||||
type="streaming_update",
|
||||
request_id="req-unknown",
|
||||
prompt=SimpleNamespace(request_id="req-unknown", prompt_token_ids=[1], resumable=True),
|
||||
original_prompt={"prompt": "segment-2"},
|
||||
output_prompt_text="segment-2",
|
||||
sampling_params_list=[_sampling_params()],
|
||||
final_stage_id=0,
|
||||
preprocess_ms=0.0,
|
||||
request_timestamp=time.time(),
|
||||
enqueue_ts=time.perf_counter(),
|
||||
)
|
||||
)
|
||||
|
||||
assert pool.calls == []
|
||||
assert add_request_calls == []
|
||||
assert "req-unknown" not in orchestrator.request_states
|
||||
|
||||
|
||||
async def test_handle_streaming_update_passes_prompt_text_to_stage_pool() -> None:
|
||||
class RecordingPool:
|
||||
def __init__(self) -> None:
|
||||
|
||||
@@ -7,15 +7,24 @@ import asyncio
|
||||
import queue
|
||||
from types import SimpleNamespace
|
||||
from typing import Any
|
||||
from unittest.mock import AsyncMock
|
||||
from unittest.mock import ANY, AsyncMock, MagicMock
|
||||
|
||||
import janus
|
||||
import pytest
|
||||
from vllm.outputs import CompletionOutput, RequestOutput
|
||||
from vllm.sampling_params import SamplingParams
|
||||
|
||||
from vllm_omni.engine.orchestrator import Orchestrator, OrchestratorRequestState
|
||||
from vllm_omni.engine.orchestrator import (
|
||||
Orchestrator,
|
||||
OrchestratorRequestState,
|
||||
_OrchestratorDuplexStagePort,
|
||||
)
|
||||
from vllm_omni.engine.stage_pool import StagePool
|
||||
from vllm_omni.experimental.fullduplex.engine.contracts import (
|
||||
DuplexStageRequestContext,
|
||||
DuplexStageSubmission,
|
||||
)
|
||||
from vllm_omni.experimental.fullduplex.engine.messages import DuplexFence
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
@@ -107,11 +116,50 @@ class FakePrewarmPool:
|
||||
|
||||
async def submit_initial(self, _request_id, _req_state, request, prompt_text=None):
|
||||
self.submitted.append(request)
|
||||
return 0
|
||||
|
||||
def get_bound_replica_id(self, _request_id):
|
||||
return 0
|
||||
|
||||
|
||||
def _duplex_stage_port_submission():
|
||||
stage_pools = []
|
||||
for stage_id in range(3):
|
||||
pool = SimpleNamespace(
|
||||
stage_client=SimpleNamespace(default_sampling_params=SamplingParams(max_tokens=1)),
|
||||
stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)),
|
||||
submit_initial=AsyncMock(return_value=stage_id + 10),
|
||||
submit_update=AsyncMock(return_value=stage_id + 20),
|
||||
)
|
||||
stage_pools.append(pool)
|
||||
request_states: dict[str, OrchestratorRequestState] = {}
|
||||
prewarm = AsyncMock()
|
||||
port = _OrchestratorDuplexStagePort(
|
||||
stage_pools=stage_pools,
|
||||
request_states=request_states,
|
||||
running_counter=None,
|
||||
cleanup_request_ids=AsyncMock(),
|
||||
async_chunk=True,
|
||||
prewarm_async_chunk_stages=prewarm,
|
||||
)
|
||||
context = DuplexStageRequestContext(
|
||||
request_id="req-duplex",
|
||||
session_id="session-duplex",
|
||||
fence=DuplexFence("session-duplex"),
|
||||
stage_id=0,
|
||||
final_stage_id=2,
|
||||
config_generation=0,
|
||||
sampling_params=tuple(SamplingParams(max_tokens=1) for _ in range(3)),
|
||||
)
|
||||
port.ensure_request(context)
|
||||
submission = DuplexStageSubmission(
|
||||
context=context,
|
||||
prompt={"prompt_token_ids": [1, 2]},
|
||||
already_submitted=False,
|
||||
)
|
||||
return port, stage_pools, request_states, prewarm, submission
|
||||
|
||||
|
||||
def _request_output(request_id: str) -> RequestOutput:
|
||||
completion = CompletionOutput(
|
||||
index=0,
|
||||
@@ -199,11 +247,13 @@ async def test_async_prewarm_skips_outgoing_only_stage() -> None:
|
||||
stage2 = FakePrewarmPool("receiver")
|
||||
orchestrator.stage_pools = [stage0, stage1, stage2]
|
||||
orchestrator._emit_tx_edge = lambda **_kwargs: None
|
||||
orchestrator._record_duplex_stage_submission = MagicMock()
|
||||
req_state = OrchestratorRequestState(
|
||||
request_id="req-prewarm",
|
||||
prompt={"prompt_token_ids": [1, 2]},
|
||||
sampling_params_list=[SamplingParams(max_tokens=1) for _ in range(3)],
|
||||
final_stage_id=2,
|
||||
duplex_identity=SimpleNamespace(),
|
||||
)
|
||||
|
||||
await orchestrator._prewarm_async_chunk_stages(
|
||||
@@ -216,6 +266,23 @@ async def test_async_prewarm_skips_outgoing_only_stage() -> None:
|
||||
assert len(stage2.submitted) == 1
|
||||
assert 1 not in req_state.stage_submit_ts
|
||||
assert 2 in req_state.stage_submit_ts
|
||||
orchestrator._record_duplex_stage_submission.assert_called_once_with(
|
||||
2,
|
||||
"req-prewarm",
|
||||
0,
|
||||
req_state,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_duplex_prewarm_runs_after_first_stage0_submission() -> None:
|
||||
port, stage_pools, request_states, prewarm, submission = _duplex_stage_port_submission()
|
||||
|
||||
result = await port.submit(submission)
|
||||
|
||||
assert result.stage_id == 0
|
||||
stage_pools[0].submit_initial.assert_awaited_once()
|
||||
prewarm.assert_awaited_once_with("req-duplex", ANY, request_states["req-duplex"])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
||||
@@ -14,7 +14,7 @@ from vllm.v1.engine import FinishReason
|
||||
from vllm.v1.engine.output_processor import OutputProcessor as VLLMOutputProcessor
|
||||
|
||||
from vllm_omni.outputs import output_processor
|
||||
from vllm_omni.outputs.output_modality import OutputModalityNames
|
||||
from vllm_omni.outputs.output_modality import OutputModality, OutputModalityNames
|
||||
from vllm_omni.outputs.output_processor import MultimodalOutputProcessor, OmniRequestState
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
@@ -452,6 +452,127 @@ def test_no_detokenizer_make_request_output():
|
||||
assert AUDIO in result.outputs[0].multimodal_output
|
||||
|
||||
|
||||
def test_no_detokenizer_process_outputs_returns_nonterminal_audio_chunk(monkeypatch):
|
||||
"""Generation-stage MM-only chunks must reach the orchestrator.
|
||||
|
||||
StagePool registers generation requests without a per-request collector,
|
||||
so the processor return value is the only path from Stage2 to
|
||||
Orchestrator._handle_processed_outputs().
|
||||
"""
|
||||
state = _make_no_detok_state(RequestOutputKind.DELTA)
|
||||
processor = object.__new__(MultimodalOutputProcessor)
|
||||
processor.output_modality = OutputModality.AUDIO
|
||||
processor.request_states = {"r": state}
|
||||
upstream_result = SimpleNamespace(request_outputs=[], reqs_to_abort=[])
|
||||
monkeypatch.setattr(
|
||||
VLLMOutputProcessor,
|
||||
"process_outputs",
|
||||
lambda *_args, **_kwargs: upstream_result,
|
||||
)
|
||||
engine_output = SimpleNamespace(
|
||||
request_id="r",
|
||||
multimodal_output={
|
||||
"model_outputs": torch.arange(8, dtype=torch.float32),
|
||||
"sr": torch.tensor(24000, dtype=torch.int32),
|
||||
},
|
||||
output_type="audio",
|
||||
pooling_output=None,
|
||||
new_token_ids=[],
|
||||
finish_reason=None,
|
||||
stop_reason=None,
|
||||
kv_transfer_params=None,
|
||||
routed_experts=None,
|
||||
num_cached_tokens=0,
|
||||
)
|
||||
|
||||
result = processor.process_outputs([engine_output])
|
||||
|
||||
assert len(result.request_outputs) == 1
|
||||
assert result.request_outputs[0].request_id == "r"
|
||||
assert AUDIO in result.request_outputs[0].outputs[0].multimodal_output
|
||||
|
||||
|
||||
def _make_mm_only_output_processor(monkeypatch):
|
||||
processor = object.__new__(MultimodalOutputProcessor)
|
||||
processor.output_modality = OutputModality.AUDIO
|
||||
processor.request_states = {"r": _make_no_detok_state(RequestOutputKind.DELTA)}
|
||||
monkeypatch.setattr(
|
||||
VLLMOutputProcessor,
|
||||
"process_outputs",
|
||||
lambda *_args, **_kwargs: SimpleNamespace(
|
||||
request_outputs=[],
|
||||
reqs_to_abort=[],
|
||||
),
|
||||
)
|
||||
|
||||
def finish_request(req_state):
|
||||
processor.request_states.pop(req_state.request_id, None)
|
||||
|
||||
processor._finish_request = finish_request
|
||||
return processor
|
||||
|
||||
|
||||
def _audio_engine_output(*, is_segment_finished: bool, is_last_chunk: bool):
|
||||
return SimpleNamespace(
|
||||
request_id="r",
|
||||
multimodal_output={
|
||||
"model_outputs": torch.arange(8, dtype=torch.float32),
|
||||
"meta.tts_is_last_chunk": torch.tensor([int(is_last_chunk)]),
|
||||
},
|
||||
output_type="audio",
|
||||
pooling_output=None,
|
||||
new_token_ids=[],
|
||||
finish_reason=FinishReason.STOP,
|
||||
stop_reason=None,
|
||||
kv_transfer_params=None,
|
||||
routed_experts=None,
|
||||
num_cached_tokens=0,
|
||||
is_segment_finished=is_segment_finished,
|
||||
)
|
||||
|
||||
|
||||
def test_mm_only_segment_finish_retains_request_state_for_next_audio_chunk(monkeypatch):
|
||||
processor = _make_mm_only_output_processor(monkeypatch)
|
||||
|
||||
first = processor.process_outputs([_audio_engine_output(is_segment_finished=True, is_last_chunk=False)])
|
||||
|
||||
assert "r" in processor.request_states
|
||||
assert len(first.request_outputs) == 1
|
||||
assert first.request_outputs[0].finished is False
|
||||
|
||||
second = processor.process_outputs([_audio_engine_output(is_segment_finished=False, is_last_chunk=True)])
|
||||
|
||||
assert len(second.request_outputs) == 1
|
||||
assert second.request_outputs[0].finished is True
|
||||
assert "r" not in processor.request_states
|
||||
|
||||
|
||||
def test_mm_only_nonfinal_audio_retains_state_without_segment_marker(monkeypatch):
|
||||
processor = _make_mm_only_output_processor(monkeypatch)
|
||||
|
||||
first = processor.process_outputs([_audio_engine_output(is_segment_finished=False, is_last_chunk=False)])
|
||||
|
||||
assert "r" in processor.request_states
|
||||
assert len(first.request_outputs) == 1
|
||||
assert first.request_outputs[0].finished is False
|
||||
|
||||
second = processor.process_outputs([_audio_engine_output(is_segment_finished=False, is_last_chunk=True)])
|
||||
|
||||
assert len(second.request_outputs) == 1
|
||||
assert second.request_outputs[0].finished is True
|
||||
assert "r" not in processor.request_states
|
||||
|
||||
|
||||
def test_mm_only_terminal_finish_removes_request_state(monkeypatch):
|
||||
processor = _make_mm_only_output_processor(monkeypatch)
|
||||
|
||||
result = processor.process_outputs([_audio_engine_output(is_segment_finished=False, is_last_chunk=True)])
|
||||
|
||||
assert len(result.request_outputs) == 1
|
||||
assert result.request_outputs[0].finished is True
|
||||
assert "r" not in processor.request_states
|
||||
|
||||
|
||||
def test_no_detokenizer_make_request_output_with_routed_experts():
|
||||
"""make_request_output accepts the routed_experts arg that the multimodal
|
||||
output channel (_process_mm_only_outputs) passes positionally, and attaches
|
||||
|
||||
@@ -3,6 +3,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import socket
|
||||
from contextlib import contextmanager
|
||||
from types import SimpleNamespace
|
||||
from typing import Any
|
||||
@@ -11,6 +12,7 @@ import pytest
|
||||
from pytest_mock import MockerFixture
|
||||
from vllm.v1.engine.utils import EngineZmqAddresses
|
||||
|
||||
from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig
|
||||
from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
|
||||
from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClientBase
|
||||
from vllm_omni.engine.stage_engine_startup import (
|
||||
@@ -189,6 +191,34 @@ class TestOmniMasterServerAllocation:
|
||||
output_port = int(alloc.output_bind_address.split(":")[-1])
|
||||
assert len({handshake_port, input_port, output_port}) == 3
|
||||
|
||||
def test_route_ports_remain_os_reserved_until_master_releases_them(self):
|
||||
server = OmniMasterServer(master_address="127.0.0.1", master_port=15005, stage_ids=[0])
|
||||
alloc = server.get_allocation(0)
|
||||
ports = [
|
||||
int(alloc.handshake_bind_address.rsplit(":", 1)[-1]),
|
||||
int(alloc.input_bind_address.rsplit(":", 1)[-1]),
|
||||
int(alloc.output_bind_address.rsplit(":", 1)[-1]),
|
||||
]
|
||||
|
||||
for port in ports:
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as probe:
|
||||
with pytest.raises(OSError):
|
||||
probe.bind(("127.0.0.1", port))
|
||||
|
||||
server.release_route_port_reservations(0, handshake=True, data=False)
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as probe:
|
||||
probe.bind(("127.0.0.1", ports[0]))
|
||||
for port in ports[1:]:
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as probe:
|
||||
with pytest.raises(OSError):
|
||||
probe.bind(("127.0.0.1", port))
|
||||
|
||||
server.stop()
|
||||
|
||||
for port in ports[1:]:
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as probe:
|
||||
probe.bind(("127.0.0.1", port))
|
||||
|
||||
def test_get_zmq_addresses_returns_bind_addresses(self):
|
||||
server = OmniMasterServer(master_address="127.0.0.1", master_port=15002, stage_ids=[0])
|
||||
alloc = server.get_allocation(0)
|
||||
@@ -397,6 +427,25 @@ class TestSingleStageModeDetection:
|
||||
assert engine.single_stage_mode is False
|
||||
assert engine._single_stage_id_filter is None
|
||||
|
||||
def test_stage_configs_path_loads_duplex_runtime_config(self, mocker: MockerFixture):
|
||||
duplex_session = DuplexSessionRuntimeConfig(max_sessions=2)
|
||||
mocker.patch(
|
||||
"vllm_omni.engine.async_omni_engine.StageConfigFactory.get_pipeline_config",
|
||||
return_value=None,
|
||||
)
|
||||
load_deploy_config = mocker.patch(
|
||||
"vllm_omni.engine.async_omni_engine.load_deploy_config",
|
||||
return_value=SimpleNamespace(duplex_session=duplex_session),
|
||||
)
|
||||
|
||||
engine = self._make_engine_no_thread(
|
||||
mocker,
|
||||
stage_configs_path="/fake/duplex.yaml",
|
||||
)
|
||||
|
||||
load_deploy_config.assert_called_once_with("/fake/duplex.yaml")
|
||||
assert engine.duplex_session_config is duplex_session
|
||||
|
||||
def test_single_stage_mode_without_stage_id_has_no_filter(self, mocker: MockerFixture):
|
||||
engine = self._make_engine_no_thread(
|
||||
mocker,
|
||||
|
||||
@@ -2520,6 +2520,52 @@ def test_duplex_auto_response_tts_scheduler_eos_fallback():
|
||||
assert results[0]["abort_data_plane_request"] is True
|
||||
|
||||
|
||||
def test_duplex_auto_response_rearms_tts_boundary_for_next_segment_in_same_turn():
|
||||
data_plane = _test_data_plane()
|
||||
request_id = "duplex-sid-next-segment-same-turn-e0-stage0"
|
||||
session = DuplexSession(
|
||||
session_id="sid-next-segment-same-turn",
|
||||
config=DuplexSessionConfig(extra_body={"auto_response": True}),
|
||||
)
|
||||
session.bind_response_turn(0)
|
||||
data_plane.begin_request(request_id)
|
||||
|
||||
first_segment = _duplex_tts_output(
|
||||
request_id=request_id,
|
||||
samples=24000,
|
||||
finished=False,
|
||||
tts_is_last_chunk=True,
|
||||
turn_id=0,
|
||||
)
|
||||
first_results = _project_data_plane(
|
||||
data_plane,
|
||||
{"data_plane_outputs": [first_segment]},
|
||||
session=session,
|
||||
)
|
||||
|
||||
assert len(first_results) == 1
|
||||
assert first_results[0]["abort_data_plane_request"] is True
|
||||
|
||||
data_plane.begin_request(request_id)
|
||||
second_segment = _duplex_tts_output(
|
||||
request_id=request_id,
|
||||
samples=48000,
|
||||
finished=False,
|
||||
text="hello again",
|
||||
tts_is_last_chunk=True,
|
||||
turn_id=0,
|
||||
)
|
||||
second_results = _project_data_plane(
|
||||
data_plane,
|
||||
{"data_plane_outputs": [second_segment]},
|
||||
session=session,
|
||||
)
|
||||
|
||||
assert len(second_results) == 1
|
||||
assert second_results[0]["audio_data"] == "wav-24000"
|
||||
assert second_results[0]["abort_data_plane_request"] is True
|
||||
|
||||
|
||||
def test_duplex_turn_end_is_not_swallowed_by_finished_segment_fallback():
|
||||
data_plane = _test_data_plane()
|
||||
request_id = "duplex-sid-turn-end-after-segment-e0-stage0"
|
||||
@@ -6202,19 +6248,27 @@ async def test_minicpmo_native_auto_response_keeps_request_bound_for_segment_con
|
||||
config_timeout_s=0.1,
|
||||
idle_timeout_s=1,
|
||||
)
|
||||
ws = TimedWebSocket()
|
||||
|
||||
def on_send(ws: TimedWebSocket, data: dict[str, Any]) -> None:
|
||||
if data.get("type") != "session.created":
|
||||
return
|
||||
session = handler._registry.get("sid-native-auto-continuation")
|
||||
assert session is not None
|
||||
session.capabilities.chunk_period_ms = 50
|
||||
ws.put(
|
||||
{
|
||||
"type": "input_audio_buffer.append",
|
||||
"audio": _pcm_f32_b64(800),
|
||||
"format": "pcm_f32le",
|
||||
"sample_rate_hz": 16000,
|
||||
}
|
||||
)
|
||||
ws.put({"type": "input_audio_buffer.commit", "final": True})
|
||||
|
||||
ws = TimedWebSocket(on_send=on_send)
|
||||
event = _native_session_create("sid-native-auto-continuation")
|
||||
event["session"]["extra_body"]["auto_response"] = True
|
||||
ws.put(event)
|
||||
ws.put(
|
||||
{
|
||||
"type": "input_audio_buffer.append",
|
||||
"audio": _pcm_f32_b64(16000),
|
||||
"format": "pcm_f32le",
|
||||
"sample_rate_hz": 16000,
|
||||
}
|
||||
)
|
||||
ws.put({"type": "input_audio_buffer.commit", "final": True})
|
||||
|
||||
await handler.handle_session(ws)
|
||||
|
||||
@@ -6225,6 +6279,112 @@ async def test_minicpmo_native_auto_response_keeps_request_bound_for_segment_con
|
||||
assert final is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_minicpmo_native_auto_response_real_input_preempts_unsent_silence():
|
||||
request_id = "duplex-sid-native-real-before-silence-e0-stage0"
|
||||
|
||||
class SequencedAppendEngine(FakeEngineClient):
|
||||
def __init__(self) -> None:
|
||||
control_result = {
|
||||
"operation": "append",
|
||||
"session_id": "sid-native-real-before-silence",
|
||||
"ok": True,
|
||||
"unsupported_count": 0,
|
||||
"error_count": 0,
|
||||
"stage_results": [
|
||||
{
|
||||
"result": {
|
||||
"data_plane_append": True,
|
||||
"request_id": request_id,
|
||||
"response_stage_id": 1,
|
||||
}
|
||||
}
|
||||
],
|
||||
}
|
||||
terminal_segment = _duplex_tts_output(
|
||||
request_id=request_id,
|
||||
samples=0,
|
||||
finished=True,
|
||||
tts_is_last_chunk=True,
|
||||
token_ids=[151645],
|
||||
)
|
||||
super().__init__(
|
||||
control_result=control_result,
|
||||
collect_outputs=[[terminal_segment], [], []],
|
||||
collect_delay_s=0.01,
|
||||
)
|
||||
self.append_sequence: list[str] = []
|
||||
self.first_real_started = asyncio.Event()
|
||||
self.second_real_started = asyncio.Event()
|
||||
self.real_count = 0
|
||||
|
||||
async def append_duplex_input_async(self, session_id: str, **kwargs):
|
||||
kwargs.pop("expected_epoch", None)
|
||||
result = await super().append_duplex_input_async(session_id, **kwargs)
|
||||
payload = kwargs.get("payload")
|
||||
is_silence = (
|
||||
isinstance(payload, dict)
|
||||
and payload.get("audio") == OmniDuplexSessionHandler._NATIVE_SILENCE_UNIT_PAYLOAD_AUDIO
|
||||
)
|
||||
if is_silence:
|
||||
self.append_sequence.append("silence")
|
||||
else:
|
||||
self.real_count += 1
|
||||
self.append_sequence.append(f"real-{self.real_count}")
|
||||
if self.real_count == 1:
|
||||
self.first_real_started.set()
|
||||
elif self.real_count == 2:
|
||||
self.second_real_started.set()
|
||||
return result
|
||||
|
||||
session_created = asyncio.Event()
|
||||
engine = SequencedAppendEngine()
|
||||
handler = OmniDuplexSessionHandler(
|
||||
chat_service=FakeChatService(engine),
|
||||
config_timeout_s=0.1,
|
||||
idle_timeout_s=1,
|
||||
)
|
||||
|
||||
def on_send(ws: TimedWebSocket, data: dict[str, Any]) -> None:
|
||||
if data.get("type") != "session.created":
|
||||
return
|
||||
session = handler._registry.get("sid-native-real-before-silence")
|
||||
assert session is not None
|
||||
session.capabilities.chunk_period_ms = 50
|
||||
ws.put(
|
||||
{
|
||||
"type": "input_audio_buffer.append",
|
||||
"audio": _pcm_f32_b64(800),
|
||||
"format": "pcm_f32le",
|
||||
"sample_rate_hz": 16000,
|
||||
}
|
||||
)
|
||||
session_created.set()
|
||||
|
||||
ws = TimedWebSocket(on_send=on_send, receive_timeout_s=2.0)
|
||||
event = _native_session_create("sid-native-real-before-silence")
|
||||
event["session"]["extra_body"]["auto_response"] = True
|
||||
ws.put(event)
|
||||
handler_task = asyncio.create_task(handler.handle_session(ws))
|
||||
|
||||
await asyncio.wait_for(session_created.wait(), timeout=1)
|
||||
await asyncio.wait_for(engine.first_real_started.wait(), timeout=1)
|
||||
await asyncio.sleep(0.02)
|
||||
ws.put(
|
||||
{
|
||||
"type": "input_audio_buffer.append",
|
||||
"audio": _pcm_f32_b64(800, value=0.07),
|
||||
"format": "pcm_f32le",
|
||||
"sample_rate_hz": 16000,
|
||||
}
|
||||
)
|
||||
await asyncio.wait_for(engine.second_real_started.wait(), timeout=1)
|
||||
ws.put({"type": "session.close"})
|
||||
await asyncio.wait_for(handler_task, timeout=2)
|
||||
|
||||
assert engine.append_sequence[:2] == ["real-1", "real-2"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_minicpmo_native_auto_response_real_input_waits_for_submitted_silence_tail():
|
||||
request_id = "duplex-sid-native-silence-tail-e0-stage0"
|
||||
@@ -6294,18 +6454,26 @@ async def test_minicpmo_native_auto_response_real_input_waits_for_submitted_sile
|
||||
config_timeout_s=0.1,
|
||||
idle_timeout_s=1,
|
||||
)
|
||||
ws = TimedWebSocket(receive_timeout_s=2.0)
|
||||
|
||||
def on_send(ws: TimedWebSocket, data: dict[str, Any]) -> None:
|
||||
if data.get("type") != "session.created":
|
||||
return
|
||||
session = handler._registry.get("sid-native-silence-tail")
|
||||
assert session is not None
|
||||
session.capabilities.chunk_period_ms = 50
|
||||
ws.put(
|
||||
{
|
||||
"type": "input_audio_buffer.append",
|
||||
"audio": _pcm_f32_b64(800, value=0.05),
|
||||
"format": "pcm_f32le",
|
||||
"sample_rate_hz": 16000,
|
||||
}
|
||||
)
|
||||
|
||||
ws = TimedWebSocket(on_send=on_send, receive_timeout_s=2.0)
|
||||
event = _native_session_create("sid-native-silence-tail")
|
||||
event["session"]["extra_body"]["auto_response"] = True
|
||||
ws.put(event)
|
||||
ws.put(
|
||||
{
|
||||
"type": "input_audio_buffer.append",
|
||||
"audio": _pcm_f32_b64(16000, value=0.05),
|
||||
"format": "pcm_f32le",
|
||||
"sample_rate_hz": 16000,
|
||||
}
|
||||
)
|
||||
handler_task = asyncio.create_task(handler.handle_session(ws))
|
||||
|
||||
await asyncio.wait_for(engine.silence_append_started.wait(), timeout=1)
|
||||
@@ -6330,6 +6498,135 @@ async def test_minicpmo_native_auto_response_real_input_waits_for_submitted_sile
|
||||
assert engine.append_sequence == ["real-1", "silence-started", "silence-returned", "real-2"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_minicpmo_native_auto_response_preserves_silence_continuations_across_append_tails():
|
||||
request_id = "duplex-sid-native-continuation-tails-e0-stage0"
|
||||
|
||||
class SequencedAppendEngine(FakeEngineClient):
|
||||
def __init__(self) -> None:
|
||||
control_result = {
|
||||
"operation": "append",
|
||||
"session_id": "sid-native-continuation-tails",
|
||||
"ok": True,
|
||||
"unsupported_count": 0,
|
||||
"error_count": 0,
|
||||
"stage_results": [
|
||||
{
|
||||
"result": {
|
||||
"data_plane_append": True,
|
||||
"request_id": request_id,
|
||||
"response_stage_id": 1,
|
||||
}
|
||||
}
|
||||
],
|
||||
}
|
||||
super().__init__(control_result=control_result)
|
||||
self.real_append_started = asyncio.Event()
|
||||
self.release_real_append = asyncio.Event()
|
||||
self.first_silence_started = asyncio.Event()
|
||||
self.release_first_silence = asyncio.Event()
|
||||
self.second_silence_started = asyncio.Event()
|
||||
self.release_collect = asyncio.Event()
|
||||
self.silence_count = 0
|
||||
|
||||
async def append_duplex_input_async(self, session_id: str, **kwargs):
|
||||
kwargs.pop("expected_epoch", None)
|
||||
payload = kwargs.get("payload")
|
||||
assert isinstance(payload, dict)
|
||||
result = await super().append_duplex_input_async(session_id, **kwargs)
|
||||
if payload.get("audio") != OmniDuplexSessionHandler._NATIVE_SILENCE_UNIT_PAYLOAD_AUDIO:
|
||||
self.real_append_started.set()
|
||||
await self.release_real_append.wait()
|
||||
return result
|
||||
|
||||
self.silence_count += 1
|
||||
if self.silence_count == 1:
|
||||
self.first_silence_started.set()
|
||||
await self.release_first_silence.wait()
|
||||
else:
|
||||
self.second_silence_started.set()
|
||||
return result
|
||||
|
||||
async def collect_duplex_data_plane_outputs_async(
|
||||
self,
|
||||
request_id: str,
|
||||
*,
|
||||
response_stage_id: int | None = None,
|
||||
timeout: float | None = None,
|
||||
) -> list[object]:
|
||||
del timeout
|
||||
self.collected.append((request_id, response_stage_id))
|
||||
await self.release_collect.wait()
|
||||
return []
|
||||
|
||||
session_created = asyncio.Event()
|
||||
|
||||
def on_send(_ws: TimedWebSocket, data: dict[str, Any]) -> None:
|
||||
if data.get("type") == "session.created":
|
||||
session_created.set()
|
||||
|
||||
engine = SequencedAppendEngine()
|
||||
handler = OmniDuplexSessionHandler(
|
||||
chat_service=FakeChatService(engine),
|
||||
config_timeout_s=0.1,
|
||||
idle_timeout_s=1,
|
||||
)
|
||||
ws = TimedWebSocket(on_send=on_send, receive_timeout_s=2.0)
|
||||
event = _native_session_create("sid-native-continuation-tails")
|
||||
event["session"]["extra_body"]["auto_response"] = True
|
||||
ws.put(event)
|
||||
ws.put(
|
||||
{
|
||||
"type": "input_audio_buffer.append",
|
||||
"audio": _pcm_f32_b64(16000, value=0.05),
|
||||
"format": "pcm_f32le",
|
||||
"sample_rate_hz": 16000,
|
||||
}
|
||||
)
|
||||
handler_task = asyncio.create_task(handler.handle_session(ws))
|
||||
|
||||
await asyncio.wait_for(session_created.wait(), timeout=1)
|
||||
await asyncio.wait_for(engine.real_append_started.wait(), timeout=1)
|
||||
session = handler._registry.get("sid-native-continuation-tails")
|
||||
assert session is not None
|
||||
session.capabilities.chunk_period_ms = 50
|
||||
session.bind_request(request_id)
|
||||
response_id = session.begin_response(turn_id=session.turn_id)
|
||||
native = handler._minicpmo_session_state(session)
|
||||
scheduler = native.silence_continuation_scheduler
|
||||
assert scheduler is not None
|
||||
|
||||
payload = handler._native_silence_unit_payload()
|
||||
payload["duplex_turn_id"] = session.turn_id
|
||||
kwargs = {
|
||||
"request_id": request_id,
|
||||
"owner_id": f"response:{response_id}",
|
||||
"response_id": response_id,
|
||||
"response_owned": True,
|
||||
"expected_epoch": session.epoch,
|
||||
"expected_incarnation": session.incarnation,
|
||||
"expected_model_turn_id": None,
|
||||
"send_json": ws.send_json,
|
||||
}
|
||||
assert await scheduler(payload, **kwargs) is True
|
||||
|
||||
next_schedule = asyncio.create_task(scheduler(payload, **kwargs))
|
||||
await asyncio.sleep(0.05)
|
||||
assert not next_schedule.done()
|
||||
|
||||
engine.release_real_append.set()
|
||||
await asyncio.wait_for(engine.first_silence_started.wait(), timeout=1)
|
||||
assert not next_schedule.done()
|
||||
|
||||
engine.release_first_silence.set()
|
||||
assert await asyncio.wait_for(next_schedule, timeout=1) is True
|
||||
await asyncio.wait_for(engine.second_silence_started.wait(), timeout=1)
|
||||
ws.put({"type": "session.close"})
|
||||
await asyncio.wait_for(handler_task, timeout=2)
|
||||
|
||||
assert engine.silence_count == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_minicpmo_native_duplex_uses_segment_text_metadata_for_transcript_cursor():
|
||||
request_id = "duplex-sid-native-repeat-text-e0-stage0"
|
||||
|
||||
@@ -222,6 +222,7 @@ def test_run_headless_llm_registers_with_auto_assigned_replica_id(mocker: Mocker
|
||||
from vllm_omni.engine.stage_engine_startup import StageRegistrationResponse
|
||||
|
||||
stage_cfg = _make_stage_cfg(0, stage_type="llm")
|
||||
stage_cfg.engine_args["async_chunk"] = True
|
||||
parallel_config = SimpleNamespace(
|
||||
data_parallel_size_local=1,
|
||||
data_parallel_rank=0,
|
||||
@@ -241,7 +242,10 @@ def test_run_headless_llm_registers_with_auto_assigned_replica_id(mocker: Mocker
|
||||
"vllm_omni.distributed.omni_connectors.utils.initialization.resolve_omni_kv_config_for_stage",
|
||||
return_value=(None, None, None),
|
||||
)
|
||||
mocker.patch("vllm_omni.engine.stage_init_utils.get_stage_connector_spec", return_value={})
|
||||
mock_connector_spec = mocker.patch(
|
||||
"vllm_omni.engine.stage_init_utils.get_stage_connector_spec",
|
||||
return_value={},
|
||||
)
|
||||
mocker.patch("vllm_omni.engine.stage_init_utils.build_engine_args_dict", return_value={})
|
||||
mocker.patch(
|
||||
"vllm_omni.engine.stage_init_utils.build_vllm_config",
|
||||
@@ -277,6 +281,7 @@ def test_run_headless_llm_registers_with_auto_assigned_replica_id(mocker: Mocker
|
||||
assert kwargs["omni_stage_config"] is stage_cfg
|
||||
assert kwargs["replica_id"] is None
|
||||
assert "socket_ownership" not in kwargs
|
||||
assert mock_connector_spec.call_args.kwargs["async_chunk"] is True
|
||||
|
||||
assert mock_manager_cls.call_count == 1
|
||||
mgr_kwargs = mock_manager_cls.call_args.kwargs
|
||||
|
||||
@@ -207,3 +207,68 @@ def test_playback_worklet_waits_before_playing_and_rebuffers_after_underrun():
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
|
||||
|
||||
def test_playback_worklet_fades_terminal_drain_to_zero():
|
||||
node = shutil.which("node")
|
||||
if node is None:
|
||||
pytest.skip("node is required for the AudioWorklet regression test")
|
||||
|
||||
script = textwrap.dedent(
|
||||
"""
|
||||
const fs = require('fs');
|
||||
const vm = require('vm');
|
||||
|
||||
global.sampleRate = 1000;
|
||||
const messages = [];
|
||||
global.AudioWorkletProcessor = class {
|
||||
constructor() {
|
||||
this.port = {
|
||||
onmessage: null,
|
||||
postMessage: (message) => messages.push(message),
|
||||
};
|
||||
}
|
||||
};
|
||||
let Processor = null;
|
||||
global.registerProcessor = (_name, processor) => { Processor = processor; };
|
||||
vm.runInThisContext(fs.readFileSync(process.argv[1], 'utf8'));
|
||||
|
||||
const processor = new Processor();
|
||||
const pcm = new Int16Array(100);
|
||||
pcm.fill(16384);
|
||||
processor.handleMessage({
|
||||
type: 'audio',
|
||||
pcm,
|
||||
responseId: 'response-terminal',
|
||||
initialBufferMs: 400,
|
||||
});
|
||||
processor.handleMessage({
|
||||
type: 'drain',
|
||||
responseId: 'response-terminal',
|
||||
});
|
||||
|
||||
const output = new Float32Array(100);
|
||||
processor.process([], [[output]]);
|
||||
const assert = (condition, message) => {
|
||||
if (!condition) throw new Error(message);
|
||||
};
|
||||
|
||||
assert(output.some((sample) => sample !== 0), 'terminal audio must still play');
|
||||
assert(output[output.length - 1] === 0, 'terminal drain must fade the final sample to zero');
|
||||
const tail = Array.from(output.slice(-5), Math.abs);
|
||||
assert(
|
||||
tail.every((sample, index) => index === 0 || sample <= tail[index - 1]),
|
||||
'terminal drain fade must decrease monotonically',
|
||||
);
|
||||
assert(
|
||||
messages.filter((message) => message.type === 'playback-drained').length === 1,
|
||||
'terminal drain must be reported exactly once after playback',
|
||||
);
|
||||
"""
|
||||
)
|
||||
subprocess.run(
|
||||
[node, "-e", script, str(STATIC_ROOT / "playback_worklet.js")],
|
||||
check=True,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
@@ -43,6 +43,54 @@ def test_append_emits_one_model_unit_when_multiple_units_are_buffered():
|
||||
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)
|
||||
|
||||
@@ -0,0 +1,179 @@
|
||||
"""Shared fixtures and argument builders for MiniCPM-o 4.5 duplex E2E tests."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import uuid
|
||||
import wave
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
from tests.helpers.runtime import OmniServerParams
|
||||
from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
|
||||
|
||||
MODEL = "openbmb/MiniCPM-o-4_5"
|
||||
DEPLOY_CONFIG = modify_stage_config(
|
||||
get_deploy_config_path("minicpmo_4_5_duplex.yaml"),
|
||||
updates={
|
||||
"base_config": get_deploy_config_path("minicpmo_4_5.yaml"),
|
||||
"stages": {
|
||||
0: {"kv_cache_memory_bytes": 6 * 1024 * 1024 * 1024},
|
||||
1: {"kv_cache_memory_bytes": 512 * 1024 * 1024},
|
||||
2: {"kv_cache_memory_bytes": 256 * 1024 * 1024},
|
||||
},
|
||||
},
|
||||
)
|
||||
CORE_DEPLOY_CONFIG = modify_stage_config(
|
||||
DEPLOY_CONFIG,
|
||||
updates={
|
||||
"stages": {
|
||||
0: {"enforce_eager": True},
|
||||
1: {"enforce_eager": True},
|
||||
}
|
||||
},
|
||||
)
|
||||
ASSET_DIR = Path(__file__).resolve().parents[1] / "assets" / "minicpmo_4_5"
|
||||
RESPONSE_REQUIRED_WAV = ASSET_DIR / "response_required_16k.wav"
|
||||
RESPONSE_REQUIRED_SHA256 = "2e5fd4eb3ee434ce107ee3a0591fa624a33f7683c7462f45fe651c443c9af941"
|
||||
SOFT_INTERRUPT_WAV = ASSET_DIR / "soft_interrupt_16k.wav"
|
||||
SOFT_INTERRUPT_SHA256 = "cadae6d0ddc510310f16f8775d6379f30a5195369ccb54ff10a8e3f9a1f4a2ea"
|
||||
REF_AUDIO_RELATIVE_PATH = Path("assets") / "HT_ref_audio.wav"
|
||||
|
||||
SERVER_PARAMS = [
|
||||
pytest.param(
|
||||
OmniServerParams(
|
||||
model=MODEL,
|
||||
stage_config_path=DEPLOY_CONFIG,
|
||||
use_stage_cli=True,
|
||||
server_args=["--trust-remote-code"],
|
||||
),
|
||||
id="three-stage-single-gpu",
|
||||
)
|
||||
]
|
||||
CORE_SERVER_PARAMS = [
|
||||
pytest.param(
|
||||
OmniServerParams(
|
||||
model=MODEL,
|
||||
stage_config_path=CORE_DEPLOY_CONFIG,
|
||||
use_stage_cli=True,
|
||||
server_args=["--trust-remote-code"],
|
||||
),
|
||||
id="three-stage-single-gpu",
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
def validated_wav(path: Path, expected_sha256: str) -> Path:
|
||||
actual_sha256 = hashlib.sha256(path.read_bytes()).hexdigest()
|
||||
if actual_sha256 != expected_sha256:
|
||||
raise AssertionError(
|
||||
f"MiniCPM-o duplex fixture SHA256 mismatch: expected {expected_sha256}, got {actual_sha256}"
|
||||
)
|
||||
with wave.open(str(path), "rb") as wav_file:
|
||||
actual_format = (
|
||||
wav_file.getframerate(),
|
||||
wav_file.getnchannels(),
|
||||
wav_file.getsampwidth(),
|
||||
wav_file.getcomptype(),
|
||||
)
|
||||
expected_format = (16_000, 1, 2, "NONE")
|
||||
if actual_format != expected_format:
|
||||
raise AssertionError(f"MiniCPM-o duplex fixture must be 16 kHz mono PCM16, got {actual_format}")
|
||||
return path
|
||||
|
||||
|
||||
def validated_input_wav() -> Path:
|
||||
return validated_wav(RESPONSE_REQUIRED_WAV, RESPONSE_REQUIRED_SHA256)
|
||||
|
||||
|
||||
def validated_soft_interrupt_wav() -> Path:
|
||||
return validated_wav(SOFT_INTERRUPT_WAV, SOFT_INTERRUPT_SHA256)
|
||||
|
||||
|
||||
def resolve_ref_audio(model_prefix: str) -> Path:
|
||||
if model_prefix:
|
||||
model_root = Path(model_prefix) / MODEL
|
||||
else:
|
||||
model_root = Path(snapshot_download(MODEL, local_files_only=True))
|
||||
ref_audio = model_root / REF_AUDIO_RELATIVE_PATH
|
||||
if not ref_audio.is_file():
|
||||
raise FileNotFoundError(f"MiniCPM-o checkpoint ref audio is missing: {ref_audio}")
|
||||
return ref_audio
|
||||
|
||||
|
||||
def realtime_url(omni_server) -> str:
|
||||
return f"ws://{omni_server.host}:{omni_server.port}/v1/realtime?duplex=1"
|
||||
|
||||
|
||||
def demo_args(
|
||||
*,
|
||||
omni_server,
|
||||
input_wav: Path,
|
||||
ref_audio: Path,
|
||||
output_dir: Path,
|
||||
) -> SimpleNamespace:
|
||||
return SimpleNamespace(
|
||||
url=realtime_url(omni_server),
|
||||
model=omni_server.model,
|
||||
session_id=f"duplex-ci-single-{uuid.uuid4().hex}",
|
||||
input_wav=str(input_wav),
|
||||
ref_audio=str(ref_audio),
|
||||
turn_input_wav=[],
|
||||
output_dir=str(output_dir),
|
||||
output_audio_format="pcm16",
|
||||
chunk_ms=200,
|
||||
realtime_input=True,
|
||||
first_turn_ms=1400,
|
||||
turn_duration_ms=[],
|
||||
first_turn_transcript="duplex CI speech",
|
||||
omit_transcript_hints=False,
|
||||
validation_mode="response-required",
|
||||
temperature=0.0,
|
||||
scenario="sequential",
|
||||
require_audio=True,
|
||||
require_distinct_inputs=False,
|
||||
expect_empty_turn=[],
|
||||
short_ack_ms=350,
|
||||
turns=1,
|
||||
timeout_s=120.0,
|
||||
model_policy_settle_ms=2000,
|
||||
)
|
||||
|
||||
|
||||
def multi_session_args(
|
||||
*,
|
||||
omni_server,
|
||||
input_wav: Path,
|
||||
ref_audio: Path,
|
||||
output_dir: Path,
|
||||
response_required: bool,
|
||||
) -> SimpleNamespace:
|
||||
return SimpleNamespace(
|
||||
url=realtime_url(omni_server),
|
||||
model=omni_server.model,
|
||||
sessions=2,
|
||||
input_wav=str(input_wav),
|
||||
ref_audio=str(ref_audio),
|
||||
session_input_wav=[],
|
||||
session_expected_token=[],
|
||||
turn_input_wav=[],
|
||||
output_dir=str(output_dir),
|
||||
realtime_input=True,
|
||||
chunk_ms=200,
|
||||
turns=1,
|
||||
first_turn_ms=1400,
|
||||
turn_duration_ms=[],
|
||||
response_required=response_required,
|
||||
temperature=0.0 if response_required else None,
|
||||
disconnect_session_index=0,
|
||||
takeover_session_index=1,
|
||||
resume_after_ms=100,
|
||||
expire_session_index=None,
|
||||
expire_after_s=40.0,
|
||||
verify_admission_limit=None,
|
||||
model_policy_settle_ms=2000,
|
||||
timeout_s=180.0,
|
||||
)
|
||||
@@ -19,10 +19,12 @@ class _FakeEncoder(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.calls: list[int] = []
|
||||
self.last_chunk_calls: list[bool] = []
|
||||
|
||||
def forward_chunk(self, xs, last_chunk=False, cnn_cache=None, att_cache=None):
|
||||
batch, length, _ = xs.shape
|
||||
self.calls.append(batch)
|
||||
self.last_chunk_calls.append(last_chunk)
|
||||
old_length = 0 if att_cache is None else att_cache.shape[3]
|
||||
output = xs[:, : max(1, length - 1)]
|
||||
cnn = xs[:, :1, :].transpose(1, 2).contiguous()
|
||||
@@ -156,6 +158,35 @@ def _model():
|
||||
return model, token2wav
|
||||
|
||||
|
||||
def test_code2wav_resolves_hf_model_id_for_assets(mocker, tmp_path):
|
||||
resolved_root = tmp_path / "snapshot"
|
||||
resolved_root.mkdir()
|
||||
config = _config()
|
||||
config.model_config.model = "openbmb/MiniCPM-o-4_5"
|
||||
config.model_config.revision = "test-revision"
|
||||
config.model_config.stage_connector_config["extra"].pop("prompt_wav")
|
||||
config.load_config = SimpleNamespace(download_dir="/model-cache")
|
||||
model = MiniCPMO45Code2Wav(vllm_config=config)
|
||||
mock_download = mocker.patch(
|
||||
"vllm_omni.model_executor.model_loader.weight_utils.download_weights_from_hf_specific",
|
||||
return_value=str(resolved_root),
|
||||
)
|
||||
|
||||
assert model._resolve_model_root() == resolved_root
|
||||
assert model.model_path == str(resolved_root)
|
||||
assert model._default_prompt_wav == str(resolved_root / "assets" / "HT_ref_audio.wav")
|
||||
mock_download.assert_called_once_with(
|
||||
"openbmb/MiniCPM-o-4_5",
|
||||
"/model-cache",
|
||||
allow_patterns=[
|
||||
"assets/HT_ref_audio.wav",
|
||||
"assets/token2wav/*",
|
||||
],
|
||||
revision="test-revision",
|
||||
require_all=True,
|
||||
)
|
||||
|
||||
|
||||
def _info(
|
||||
request_id: str,
|
||||
chunk_seq: int,
|
||||
@@ -212,6 +243,18 @@ def test_adapter_runs_true_batch_cfg_and_splits_request_caches():
|
||||
assert cache1[0, 0, 0, 0, 0].item() == 20
|
||||
|
||||
|
||||
def test_fade_in_out_limits_overlap_to_available_previous_audio():
|
||||
speech = torch.arange(6, dtype=torch.float32).reshape(1, -1)
|
||||
previous = torch.full((1, 3), 2.0)
|
||||
window = torch.hamming_window(8, periodic=False)
|
||||
|
||||
actual = BatchedToken2Wav._fade_in_out(speech, previous, window)
|
||||
|
||||
expected = speech.clone()
|
||||
expected[..., :3] = speech[..., :3] * window[:3] + previous * window[-3:]
|
||||
torch.testing.assert_close(actual, expected)
|
||||
|
||||
|
||||
def test_estimator_cache_stack_split_round_trip_preserves_cfg_rows():
|
||||
token2wav = _FakeToken2Wav()
|
||||
adapter = BatchedToken2Wav(token2wav)
|
||||
@@ -259,6 +302,166 @@ def test_model_preserves_output_slots_and_prefers_runtime_codes():
|
||||
assert token2wav.flow.encoder.calls[-1] == 2
|
||||
|
||||
|
||||
def test_code2wav_projects_duplex_metadata_to_final_audio_output():
|
||||
model, token2wav = _model()
|
||||
segment = _info("duplex", 0, [10, 11])
|
||||
segment_text_utf8 = torch.tensor(list(b"hello"), dtype=torch.uint8)
|
||||
segment["meta"].update(
|
||||
{
|
||||
"duplex_epoch": 3,
|
||||
"duplex_turn_id": 7,
|
||||
"llm_output_text_utf8": segment_text_utf8,
|
||||
"tts_is_last_chunk": True,
|
||||
"turn_end": False,
|
||||
}
|
||||
)
|
||||
|
||||
segment_output = _forward(model, [segment])
|
||||
|
||||
assert segment_output.multimodal_outputs["meta.turn_end"][0].item() is False
|
||||
# A Talker unit boundary only drains pending codec tokens. The official
|
||||
# streaming path keeps Token2wav open until the assistant turn ends.
|
||||
assert token2wav.flow.encoder.last_chunk_calls[-1] is False
|
||||
assert "duplex" in model._states
|
||||
|
||||
final = _info("duplex", 1, [12, 13], last_chunk=True)
|
||||
final["meta"].update(segment["meta"])
|
||||
final["meta"]["chunk_seq"] = 1
|
||||
final["meta"]["last_chunk"] = True
|
||||
final["meta"]["turn_end"] = True
|
||||
output = _forward(model, [final])
|
||||
|
||||
payload = output.multimodal_outputs
|
||||
assert "meta" not in payload
|
||||
assert payload["meta.duplex_epoch"][0].item() == 3
|
||||
assert payload["meta.duplex_turn_id"][0].item() == 7
|
||||
torch.testing.assert_close(
|
||||
payload["meta.llm_output_text_utf8"][0],
|
||||
segment_text_utf8,
|
||||
)
|
||||
assert payload["meta.tts_is_last_chunk"][0].item() is True
|
||||
assert payload["meta.turn_end"][0].item() is True
|
||||
assert token2wav.flow.encoder.last_chunk_calls[-1] is True
|
||||
assert "duplex" not in model._states
|
||||
|
||||
|
||||
def test_initial_empty_segment_marker_initializes_stream_without_audio():
|
||||
model, token2wav = _model()
|
||||
boundary = _info("duplex", 0, [])
|
||||
boundary["meta"].update(
|
||||
{
|
||||
"code_flat_numel": 0,
|
||||
"tts_is_last_chunk": True,
|
||||
"turn_end": False,
|
||||
}
|
||||
)
|
||||
|
||||
output = _forward(model, [boundary])
|
||||
|
||||
assert output.multimodal_outputs["model_outputs"][0].numel() == 0
|
||||
assert "duplex" in model._states
|
||||
assert token2wav.hift.calls == []
|
||||
|
||||
resumed = _info(
|
||||
"duplex",
|
||||
1,
|
||||
[4218, 4218, 4218, 10, 11, 12, 13, 14],
|
||||
)
|
||||
output = _forward(model, [resumed])
|
||||
|
||||
assert output.multimodal_outputs["model_outputs"][0].numel() > 0
|
||||
assert "duplex" in model._states
|
||||
|
||||
|
||||
def test_shared_runtime_prompt_recreates_missing_file_before_second_owner(tmp_path, monkeypatch):
|
||||
monkeypatch.setattr("tempfile.gettempdir", lambda: str(tmp_path))
|
||||
model, _ = _model()
|
||||
reference = torch.tensor([0.0, 0.25, -0.25, 0.0])
|
||||
|
||||
first = _info("voice-a", 0, [10, 11])
|
||||
first["codes"]["ref"] = reference
|
||||
first["meta"]["ref_audio_sr"] = 16000
|
||||
first["meta"].pop("prompt_cache_id")
|
||||
_forward(model, [first], request_ids=["internal-a"])
|
||||
|
||||
prompt_key = model._request_prompt_keys["internal-a"]
|
||||
prompt_path = Path(model._runtime_prompts[prompt_key].path)
|
||||
prompt_path.unlink()
|
||||
|
||||
second = _info("voice-b", 0, [12, 13])
|
||||
second["codes"]["ref"] = reference
|
||||
second["meta"]["ref_audio_sr"] = 16000
|
||||
second["meta"].pop("prompt_cache_id")
|
||||
_forward(model, [second], request_ids=["internal-b"])
|
||||
|
||||
assert prompt_path.is_file()
|
||||
assert model._runtime_prompts[prompt_key].owners == {"internal-a", "internal-b"}
|
||||
|
||||
model.on_requests_finished(["internal-a"])
|
||||
assert prompt_path.is_file()
|
||||
assert model._runtime_prompts[prompt_key].owners == {"internal-b"}
|
||||
|
||||
model.on_requests_finished(["internal-b"])
|
||||
assert not prompt_path.exists()
|
||||
assert prompt_key not in model._runtime_prompts
|
||||
|
||||
|
||||
def test_runtime_prompt_write_failure_does_not_publish_partial_file(tmp_path, monkeypatch):
|
||||
monkeypatch.setattr("tempfile.gettempdir", lambda: str(tmp_path))
|
||||
model, _ = _model()
|
||||
reference = torch.tensor([0.0, 0.25, -0.25, 0.0])
|
||||
|
||||
def fail_after_partial_write(path, *_args, **_kwargs):
|
||||
Path(path).write_bytes(b"partial")
|
||||
raise OSError("simulated write failure")
|
||||
|
||||
monkeypatch.setattr(
|
||||
"vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_code2wav.sf.write",
|
||||
fail_after_partial_write,
|
||||
)
|
||||
|
||||
with pytest.raises(OSError, match="simulated write failure"):
|
||||
model._materialize_runtime_prompt(reference, 16000)
|
||||
|
||||
assert len(model._runtime_prompts) == 1
|
||||
entry = next(iter(model._runtime_prompts.values()))
|
||||
assert not Path(entry.path).exists()
|
||||
assert list(Path(entry.path).parent.iterdir()) == []
|
||||
|
||||
|
||||
def test_runtime_prompt_files_are_isolated_between_model_instances(tmp_path, monkeypatch):
|
||||
monkeypatch.setattr("tempfile.gettempdir", lambda: str(tmp_path))
|
||||
first_model, _ = _model()
|
||||
second_model, _ = _model()
|
||||
reference = torch.tensor([0.0, 0.25, -0.25, 0.0])
|
||||
|
||||
def runtime_ref_info(request_id: str):
|
||||
info = _info(request_id, 0, [10, 11])
|
||||
info["codes"]["ref"] = reference
|
||||
info["meta"]["ref_audio_sr"] = 16000
|
||||
info["meta"].pop("prompt_cache_id")
|
||||
return info
|
||||
|
||||
_forward(first_model, [runtime_ref_info("voice-a")], request_ids=["internal-a"])
|
||||
_forward(second_model, [runtime_ref_info("voice-b")], request_ids=["internal-b"])
|
||||
|
||||
first_key = first_model._request_prompt_keys["internal-a"]
|
||||
second_key = second_model._request_prompt_keys["internal-b"]
|
||||
first_path = Path(first_model._runtime_prompts[first_key].path)
|
||||
second_path = Path(second_model._runtime_prompts[second_key].path)
|
||||
assert first_key == second_key
|
||||
assert first_path != second_path
|
||||
assert first_path.is_file()
|
||||
assert second_path.is_file()
|
||||
|
||||
first_model.on_requests_finished(["internal-a"])
|
||||
assert not first_path.exists()
|
||||
assert second_path.is_file()
|
||||
|
||||
second_model.on_requests_finished(["internal-b"])
|
||||
assert not second_path.exists()
|
||||
|
||||
|
||||
def test_mixed_final_exact_buckets_keep_order_and_release_only_final_states():
|
||||
model, _ = _model()
|
||||
_forward(
|
||||
@@ -485,27 +688,37 @@ def test_reference_voice_and_duplex_metadata_follow_request_lifecycle():
|
||||
model, _ = _model()
|
||||
first = _info("voice-a", 0, [1, 2])
|
||||
first["codes"]["ref"] = torch.linspace(-0.1, 0.1, 160)
|
||||
segment_text_utf8 = torch.tensor(list(b"hello"), dtype=torch.uint8)
|
||||
first["meta"].update(
|
||||
ref_audio_sr=16000,
|
||||
native_duplex_segment_text="hello",
|
||||
llm_output_text_utf8=segment_text_utf8,
|
||||
duplex_turn_id=7,
|
||||
duplex_epoch=3,
|
||||
)
|
||||
first["meta"].pop("prompt_cache_id")
|
||||
|
||||
output = _forward(model, [first])
|
||||
prompt_cache_id, prompt_wav = model._owned_prompt_wavs["voice-a"]
|
||||
assert prompt_cache_id.startswith("voice-a:")
|
||||
prompt_key = model._request_prompt_keys["voice-a"]
|
||||
prompt = model._runtime_prompts[prompt_key]
|
||||
prompt_cache_id, prompt_wav = prompt.cache_id, prompt.path
|
||||
assert prompt_cache_id.startswith("runtime-ref-")
|
||||
assert Path(prompt_wav).is_file()
|
||||
assert bytes(output.multimodal_outputs["meta.llm_output_text_utf8"][0].tolist()).decode() == "hello"
|
||||
torch.testing.assert_close(
|
||||
output.multimodal_outputs["meta.llm_output_text_utf8"][0],
|
||||
segment_text_utf8,
|
||||
)
|
||||
assert output.multimodal_outputs["meta.duplex_turn_id"][0].item() == 7
|
||||
assert output.multimodal_outputs["meta.duplex_epoch"][0].item() == 3
|
||||
|
||||
final = _info("voice-a", 1, [3, 4], last_chunk=True)
|
||||
final["meta"].pop("prompt_cache_id")
|
||||
final["meta"]["tts_is_last_chunk"] = True
|
||||
output = _forward(model, [final])
|
||||
|
||||
assert output.multimodal_outputs["meta.tts_is_last_chunk"][0].item() is True
|
||||
assert "voice-a" not in model._owned_prompt_wavs
|
||||
assert model._request_prompt_keys["voice-a"] == prompt_key
|
||||
model.on_requests_finished(["voice-a"])
|
||||
assert "voice-a" not in model._request_prompt_keys
|
||||
assert prompt_key not in model._runtime_prompts
|
||||
assert not Path(prompt_wav).exists()
|
||||
assert (prompt_cache_id, prompt_wav) not in model.backend._prompt_features
|
||||
@@ -14,6 +14,7 @@ from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import (
|
||||
from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts import (
|
||||
MiniCPMO45OmniTTSForConditionalGeneration,
|
||||
_max_audio_tokens,
|
||||
_restore_weight_norm_weight,
|
||||
)
|
||||
from vllm_omni.utils.mm_outputs import to_payload_element
|
||||
|
||||
@@ -56,6 +57,8 @@ def _make_talker() -> MiniCPMO45OmniTTSForConditionalGeneration:
|
||||
talker._request_generators = {}
|
||||
talker._request_audio_states = {}
|
||||
talker._deferred_cleanup_ids = set()
|
||||
talker._codec_min_tokens = 50
|
||||
talker._codec_seed = 42
|
||||
return talker
|
||||
|
||||
|
||||
@@ -81,6 +84,23 @@ def test_audio_token_limit_scales_with_condition_length(
|
||||
assert _max_audio_tokens(condition_tokens) == expected
|
||||
|
||||
|
||||
def test_weight_norm_restore_matches_checkpoint_parametrization_in_bfloat16() -> None:
|
||||
generator = torch.Generator().manual_seed(42)
|
||||
weight_v = torch.randn(8, 16, generator=generator, dtype=torch.bfloat16)
|
||||
weight_g = torch.rand(8, 1, generator=generator, dtype=torch.bfloat16)
|
||||
linear = nn.utils.parametrizations.weight_norm(
|
||||
nn.Linear(16, 8, bias=False, dtype=torch.bfloat16),
|
||||
dim=0,
|
||||
)
|
||||
with torch.no_grad():
|
||||
linear.parametrizations.weight.original0.copy_(weight_g)
|
||||
linear.parametrizations.weight.original1.copy_(weight_v)
|
||||
|
||||
restored = _restore_weight_norm_weight(weight_g, weight_v)
|
||||
|
||||
assert torch.equal(restored, linear.weight)
|
||||
|
||||
|
||||
def test_talker_emits_request_aligned_codec_deltas_after_compaction(mocker) -> None:
|
||||
talker = _make_talker()
|
||||
seen: list[tuple[str, list[float], list[int]]] = []
|
||||
@@ -114,9 +134,73 @@ def test_talker_emits_request_aligned_codec_deltas_after_compaction(mocker) -> N
|
||||
assert _routed(output, 0)["codes"]["audio"].tolist() == [[2]]
|
||||
assert _routed(output, 1)["codes"]["audio"].tolist() == [[3]]
|
||||
assert _routed(output, 0)["meta"]["finished"].item() is False
|
||||
assert set(output.multimodal_outputs["meta"]) == {"finished"}
|
||||
assert talker.compute_logits(output.text_hidden_states).argmax(dim=-1).tolist() == [0, 0]
|
||||
|
||||
|
||||
def test_talker_projects_request_aligned_duplex_metadata(mocker) -> None:
|
||||
talker = _make_talker()
|
||||
mocker.patch.object(talker, "_sample_audio_code", return_value=torch.tensor(2))
|
||||
infos = [
|
||||
{
|
||||
"request_id": "req-a",
|
||||
"native_duplex": True,
|
||||
"duplex": {"epoch": 3, "turn_id": 7},
|
||||
"ids": {"tts": [41]},
|
||||
"meta": {
|
||||
"native_duplex_segment_text": "first",
|
||||
"turn_eos_token_id": 99,
|
||||
},
|
||||
"audio_codes": {"accumulated": torch.empty(0, dtype=torch.long)},
|
||||
},
|
||||
{
|
||||
"request_id": "req-b",
|
||||
"native_duplex": True,
|
||||
"duplex": {"epoch": 4, "turn_id": 8},
|
||||
"ids": {"tts": [42, 99]},
|
||||
"meta": {
|
||||
"native_duplex_segment_text": "second",
|
||||
"turn_eos_token_id": 99,
|
||||
},
|
||||
"audio_codes": {"accumulated": torch.empty(0, dtype=torch.long)},
|
||||
},
|
||||
]
|
||||
|
||||
output = talker.make_omni_output(
|
||||
torch.ones(2, 2),
|
||||
model_intermediate_buffer=infos,
|
||||
request_token_spans=[(0, 1), (1, 2)],
|
||||
)
|
||||
|
||||
meta = output.multimodal_outputs["meta"]
|
||||
assert [value.item() for value in meta["native_duplex"]] == [True, True]
|
||||
assert [value.item() for value in meta["duplex_epoch"]] == [3, 4]
|
||||
assert [value.item() for value in meta["duplex_turn_id"]] == [7, 8]
|
||||
assert "native_duplex_segment_text" not in meta
|
||||
assert [bytes(value.tolist()).decode("utf-8") for value in meta["llm_output_text_utf8"]] == [
|
||||
"first",
|
||||
"second",
|
||||
]
|
||||
assert [value.item() for value in meta["turn_end"]] == [False, True]
|
||||
|
||||
|
||||
def test_talker_rejects_native_duplex_without_fence_identity(mocker) -> None:
|
||||
talker = _make_talker()
|
||||
mocker.patch.object(talker, "_sample_audio_code", return_value=torch.tensor(2))
|
||||
info = {
|
||||
"request_id": "req-missing-fence",
|
||||
"native_duplex": True,
|
||||
"audio_codes": {"accumulated": torch.empty(0, dtype=torch.long)},
|
||||
}
|
||||
|
||||
with pytest.raises(RuntimeError, match="requires non-negative integer epoch and turn_id"):
|
||||
talker.make_omni_output(
|
||||
torch.ones(1, 2),
|
||||
model_intermediate_buffer=[info],
|
||||
request_token_spans=[(0, 1)],
|
||||
)
|
||||
|
||||
|
||||
def test_incomplete_prefill_emits_no_code_and_does_not_advance_state(mocker) -> None:
|
||||
talker = _make_talker()
|
||||
sample = mocker.patch.object(talker, "_sample_audio_code", return_value=torch.tensor(2))
|
||||
@@ -179,7 +263,7 @@ def test_eos_is_terminal_once_and_never_enters_codec_history(mocker) -> None:
|
||||
assert talker.compute_logits(second.text_hidden_states).argmax(dim=-1).tolist() == [1]
|
||||
|
||||
|
||||
def test_max_token_terminal_includes_only_new_codec_delta(mocker) -> None:
|
||||
def test_max_token_terminal_drops_unconsumed_codec_delta(mocker) -> None:
|
||||
talker = _make_talker()
|
||||
mocker.patch.object(talker, "_sample_audio_code", return_value=torch.tensor(3))
|
||||
info = {
|
||||
@@ -194,8 +278,10 @@ def test_max_token_terminal_includes_only_new_codec_delta(mocker) -> None:
|
||||
request_token_spans=[(0, 1)],
|
||||
)
|
||||
|
||||
assert info["audio_codes"]["accumulated"].tolist() == [4, 5, 3]
|
||||
assert output.multimodal_outputs["codes"]["audio"][0].tolist() == [[3]]
|
||||
# MiniCPMTTS.generate_chunk samples once at the max-token boundary to
|
||||
# advance RNG state, but the sampled code is not fed into KV or returned.
|
||||
assert info["audio_codes"]["accumulated"].tolist() == [4, 5]
|
||||
assert output.multimodal_outputs["codes"]["audio"][0].shape == (0, 1)
|
||||
assert output.multimodal_outputs["meta"]["finished"][0].item() is True
|
||||
assert talker.compute_logits(output.text_hidden_states).argmax(dim=-1).tolist() == [1]
|
||||
|
||||
@@ -268,6 +354,94 @@ def test_empty_speech_segment_finishes_without_sampling_codes() -> None:
|
||||
assert decode_embeds.shape == (1, 4)
|
||||
|
||||
|
||||
def test_chunked_prefill_tail_aligns_condition_with_prompt_length(mocker) -> None:
|
||||
talker = _make_talker()
|
||||
talker.emb_text = nn.Embedding(1, 2)
|
||||
condition = torch.arange(18, dtype=torch.float32).reshape(9, 2)
|
||||
mocker.patch.object(talker, "_build_condition_embeddings", return_value=condition)
|
||||
|
||||
_, embeds, _ = talker.preprocess(
|
||||
torch.zeros(9, dtype=torch.long),
|
||||
None,
|
||||
_omni_is_prefill=True,
|
||||
_omni_num_computed_tokens=59,
|
||||
_omni_prompt_len=68,
|
||||
request_id="req-chunked-prefill",
|
||||
tts_token_ids=torch.tensor([1]),
|
||||
tts_hidden_states=torch.ones(1, 2),
|
||||
)
|
||||
|
||||
assert torch.equal(embeds, condition)
|
||||
state = talker._request_audio_states["req-chunked-prefill"]
|
||||
assert state["min_tokens"] == 50
|
||||
assert state["max_tokens"] == 64
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("meta", "expected_min_tokens"),
|
||||
[
|
||||
({"turn_start": True}, 0),
|
||||
({}, 26),
|
||||
({"turn_end": True}, 0),
|
||||
],
|
||||
)
|
||||
def test_native_duplex_prefill_uses_official_chunk_limits(
|
||||
mocker,
|
||||
meta,
|
||||
expected_min_tokens,
|
||||
) -> None:
|
||||
talker = _make_talker()
|
||||
talker.emb_text = nn.Embedding(1, 2)
|
||||
mocker.patch.object(
|
||||
talker,
|
||||
"_build_condition_embeddings",
|
||||
return_value=torch.ones(3, 2),
|
||||
)
|
||||
|
||||
talker.preprocess(
|
||||
torch.zeros(3, dtype=torch.long),
|
||||
None,
|
||||
_omni_is_prefill=True,
|
||||
request_id="req-duplex-chunk",
|
||||
native_duplex=True,
|
||||
meta=meta,
|
||||
tts_token_ids=torch.tensor([1]),
|
||||
tts_hidden_states=torch.ones(1, 2),
|
||||
)
|
||||
|
||||
state = talker._request_audio_states["req-duplex-chunk"]
|
||||
assert state["min_tokens"] == expected_min_tokens
|
||||
assert state["max_tokens"] == 26
|
||||
|
||||
|
||||
def test_native_duplex_condition_matches_official_text_plus_audio_bos() -> None:
|
||||
talker = _make_talker()
|
||||
talker.emb_text = nn.Embedding(16, 2)
|
||||
talker.projector_semantic = nn.Identity()
|
||||
talker._normalize = False
|
||||
talker._text_eos_id = 14
|
||||
talker._tts_bos_id = 15
|
||||
with torch.no_grad():
|
||||
talker.emb_text.weight.copy_(torch.arange(32, dtype=torch.float32).reshape(16, 2))
|
||||
|
||||
token_ids = torch.tensor([2, 3])
|
||||
hidden_states = torch.tensor([[0.5, 1.0], [1.5, 2.0]])
|
||||
|
||||
condition = talker._build_condition_embeddings(
|
||||
token_ids,
|
||||
hidden_states,
|
||||
native_duplex=True,
|
||||
)
|
||||
|
||||
expected_text = talker.emb_text(token_ids) + hidden_states
|
||||
expected = torch.cat(
|
||||
[expected_text, talker.emb_text(torch.tensor([talker._tts_bos_id]))],
|
||||
dim=0,
|
||||
)
|
||||
assert torch.equal(condition, expected)
|
||||
assert condition.shape[0] == token_ids.shape[0] + 1
|
||||
|
||||
|
||||
def test_request_cleanup_evicts_ar_rng_and_decode_state() -> None:
|
||||
talker = _make_talker()
|
||||
talker._request_generators["req-done"] = torch.Generator()
|
||||
|
||||
@@ -43,6 +43,27 @@ def _delta(*codes: int):
|
||||
}
|
||||
|
||||
|
||||
def _duplex_delta(
|
||||
*codes: int,
|
||||
epoch: int = 3,
|
||||
turn_id: int = 7,
|
||||
text: str = "segment",
|
||||
turn_end: bool = False,
|
||||
):
|
||||
text_utf8 = torch.tensor(list(text.encode("utf-8")), dtype=torch.uint8)
|
||||
return {
|
||||
"codes": {"audio": torch.tensor(codes, dtype=torch.long).reshape(-1, 1)},
|
||||
"meta": {
|
||||
"finished": torch.tensor(False),
|
||||
"native_duplex": torch.tensor(True),
|
||||
"duplex_epoch": torch.tensor(epoch),
|
||||
"duplex_turn_id": torch.tensor(turn_id),
|
||||
"llm_output_text_utf8": text_utf8,
|
||||
"turn_end": torch.tensor(turn_end),
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _codes(payload) -> list[int]:
|
||||
assert payload.codes is not None
|
||||
assert isinstance(payload.codes.audio, torch.Tensor)
|
||||
@@ -108,29 +129,57 @@ def test_short_final_flushes_silence_prefix_and_tail() -> None:
|
||||
assert final.meta.finished.item() is True
|
||||
|
||||
|
||||
def test_duplex_turn_end_waits_for_terminal_codec_flush() -> None:
|
||||
manager = _manager()
|
||||
request = _request("req-duplex")
|
||||
|
||||
body = tts2code2wav_async_chunk(
|
||||
manager,
|
||||
_duplex_delta(*range(25), turn_end=True),
|
||||
request,
|
||||
False,
|
||||
)
|
||||
final = tts2code2wav_async_chunk(
|
||||
manager,
|
||||
_duplex_delta(turn_end=True),
|
||||
request,
|
||||
True,
|
||||
)
|
||||
|
||||
assert body is not None
|
||||
assert body.meta.last_chunk is False
|
||||
assert body.meta.turn_end is False
|
||||
assert final is not None
|
||||
assert _codes(final) == [22, 23, 24]
|
||||
assert final.meta.last_chunk is True
|
||||
assert final.meta.turn_end is True
|
||||
|
||||
|
||||
def test_first_chunk_forwards_reference_voice_and_duplex_identity() -> None:
|
||||
manager = _manager()
|
||||
request = _request("req")
|
||||
request.additional_information = {
|
||||
"codes": {"ref": [0.1, -0.1]},
|
||||
"meta": {
|
||||
"ref_audio_sr": 16000,
|
||||
"native_duplex_segment_text": "hello",
|
||||
"segment_end": True,
|
||||
"turn_end": True,
|
||||
},
|
||||
"duplex": {"epoch": 3, "model_turn_id": 7},
|
||||
"meta": {"ref_audio_sr": 16000},
|
||||
}
|
||||
|
||||
payload = tts2code2wav_async_chunk(manager, _delta(*range(7)), request, True)
|
||||
payload = tts2code2wav_async_chunk(
|
||||
manager,
|
||||
_duplex_delta(*range(7), text="hello", turn_end=True),
|
||||
request,
|
||||
True,
|
||||
)
|
||||
|
||||
assert payload is not None
|
||||
assert payload.codes.ref.tolist() == pytest.approx([0.1, -0.1])
|
||||
assert payload.meta.ref_audio_sr == 16000
|
||||
assert payload.meta.native_duplex_segment_text == "hello"
|
||||
torch.testing.assert_close(
|
||||
payload.meta.llm_output_text_utf8,
|
||||
torch.tensor(list(b"hello"), dtype=torch.uint8),
|
||||
)
|
||||
assert payload.meta.duplex_epoch == 3
|
||||
assert payload.meta.duplex_turn_id == 7
|
||||
assert payload.meta.segment_end is True
|
||||
assert payload.meta.tts_is_last_chunk is True
|
||||
assert payload.meta.turn_end is True
|
||||
|
||||
|
||||
@@ -211,6 +260,186 @@ def test_empty_final_releases_wait_gate_once() -> None:
|
||||
assert duplicate is None
|
||||
|
||||
|
||||
def test_empty_duplex_boundary_uses_zero_length_transport_placeholder() -> None:
|
||||
manager = _manager()
|
||||
|
||||
boundary = tts2code2wav_async_chunk(
|
||||
manager,
|
||||
_duplex_delta(text="boundary"),
|
||||
_request("req-duplex"),
|
||||
True,
|
||||
)
|
||||
|
||||
assert boundary is not None
|
||||
assert _codes(boundary) == [0]
|
||||
assert boundary.meta.code_flat_numel == 0
|
||||
assert boundary.meta.last_chunk is False
|
||||
assert boundary.meta.is_segment_finished.item() is False
|
||||
torch.testing.assert_close(
|
||||
boundary.meta.llm_output_text_utf8,
|
||||
torch.tensor(list(b"boundary"), dtype=torch.uint8),
|
||||
)
|
||||
|
||||
|
||||
def test_duplex_segments_preserve_stream_state_without_closing_turn() -> None:
|
||||
manager = _manager()
|
||||
request = _request("req-duplex")
|
||||
request.additional_information = {
|
||||
"codes": {"ref": [0.1, 0.2, 0.3]},
|
||||
"meta": {"ref_audio_sr": 16000},
|
||||
}
|
||||
|
||||
first = tts2code2wav_async_chunk(
|
||||
manager,
|
||||
_duplex_delta(10, 11, text="first"),
|
||||
request,
|
||||
True,
|
||||
)
|
||||
second = tts2code2wav_async_chunk(
|
||||
manager,
|
||||
_duplex_delta(12, 13, text="second"),
|
||||
request,
|
||||
True,
|
||||
)
|
||||
|
||||
assert first is not None
|
||||
assert second is not None
|
||||
assert first.meta.last_chunk is False
|
||||
assert second.meta.last_chunk is False
|
||||
assert first.codes is not None
|
||||
assert torch.allclose(first.codes.ref, torch.tensor([0.1, 0.2, 0.3]))
|
||||
assert first.meta.ref_audio_sr == 16000
|
||||
assert second.codes is not None
|
||||
assert second.codes.ref is None
|
||||
assert second.meta.cache_epoch == first.meta.cache_epoch
|
||||
assert second.meta.chunk_seq == first.meta.chunk_seq + 1
|
||||
assert second.meta.duplex_epoch == 3
|
||||
assert second.meta.duplex_turn_id == 7
|
||||
torch.testing.assert_close(
|
||||
second.meta.llm_output_text_utf8,
|
||||
torch.tensor(list(b"second"), dtype=torch.uint8),
|
||||
)
|
||||
assert second.meta.tts_is_last_chunk is True
|
||||
assert second.meta.turn_end is False
|
||||
assert first.meta.is_segment_finished.item() is False
|
||||
assert second.meta.is_segment_finished.item() is False
|
||||
|
||||
|
||||
def test_duplex_short_units_wait_for_minimum_stream_body() -> None:
|
||||
manager = _manager()
|
||||
request = _request("req-duplex")
|
||||
|
||||
first = tts2code2wav_async_chunk(
|
||||
manager,
|
||||
_duplex_delta(10, 11, 12, text="first"),
|
||||
request,
|
||||
True,
|
||||
)
|
||||
second = tts2code2wav_async_chunk(
|
||||
manager,
|
||||
_duplex_delta(13, 14, text="first"),
|
||||
request,
|
||||
True,
|
||||
)
|
||||
|
||||
assert first is not None
|
||||
assert _codes(first) == [0]
|
||||
assert first.meta.code_flat_numel == 0
|
||||
assert first.meta.last_chunk is False
|
||||
assert first.meta.tts_is_last_chunk is True
|
||||
assert second is not None
|
||||
assert _codes(second) == [4218, 4218, 4218, 10, 11, 12, 13, 14]
|
||||
assert second.meta.code_flat_numel == 8
|
||||
torch.testing.assert_close(
|
||||
second.meta.llm_output_text_utf8,
|
||||
torch.tensor(list(b"firstfirst"), dtype=torch.uint8),
|
||||
)
|
||||
|
||||
|
||||
def test_duplex_empty_finish_callback_does_not_replay_previous_text() -> None:
|
||||
manager = _manager()
|
||||
request = _request("req-duplex")
|
||||
|
||||
first = tts2code2wav_async_chunk(
|
||||
manager,
|
||||
_duplex_delta(*range(25), text="first"),
|
||||
request,
|
||||
False,
|
||||
)
|
||||
boundary = tts2code2wav_async_chunk(
|
||||
manager,
|
||||
_duplex_delta(text="first"),
|
||||
request,
|
||||
True,
|
||||
)
|
||||
second = tts2code2wav_async_chunk(
|
||||
manager,
|
||||
_duplex_delta(*range(25, 50), text="second"),
|
||||
request,
|
||||
False,
|
||||
)
|
||||
|
||||
assert first is not None
|
||||
assert boundary is not None
|
||||
assert boundary.meta.code_flat_numel == 0
|
||||
assert second is not None
|
||||
torch.testing.assert_close(
|
||||
second.meta.llm_output_text_utf8,
|
||||
torch.tensor(list(b"second"), dtype=torch.uint8),
|
||||
)
|
||||
|
||||
|
||||
def test_duplex_short_tail_does_not_replay_previous_segment_text() -> None:
|
||||
manager = _manager()
|
||||
request = _request("req-duplex")
|
||||
|
||||
def chunk(codes, text: str, finished: bool):
|
||||
return tts2code2wav_async_chunk(
|
||||
manager,
|
||||
_duplex_delta(*codes, text=text),
|
||||
request,
|
||||
finished,
|
||||
)
|
||||
|
||||
first = chunk(range(25), "和上海之间", False)
|
||||
first_tail = chunk([25, 26], "和上海之间", True)
|
||||
assert first is not None
|
||||
assert first.meta.llm_output_text_utf8.tolist() == list("和上海之间".encode())
|
||||
assert first_tail is not None
|
||||
assert first_tail.meta.code_flat_numel == 0
|
||||
assert chunk([27, 28, 29], "的距离大约是", False) is None
|
||||
next_flush = chunk([], "的距离大约是", True)
|
||||
assert next_flush is not None
|
||||
assert next_flush.meta.llm_output_text_utf8.tolist() == list("的距离大约是".encode())
|
||||
|
||||
|
||||
def test_duplex_turn_end_closes_epoch_and_next_turn_restarts_sequence() -> None:
|
||||
manager = _manager()
|
||||
request = _request("req-duplex")
|
||||
|
||||
turn_end = tts2code2wav_async_chunk(
|
||||
manager,
|
||||
_duplex_delta(14, turn_id=7, turn_end=True),
|
||||
request,
|
||||
True,
|
||||
)
|
||||
next_turn = tts2code2wav_async_chunk(
|
||||
manager,
|
||||
_duplex_delta(20, turn_id=8),
|
||||
request,
|
||||
True,
|
||||
)
|
||||
|
||||
assert turn_end is not None
|
||||
assert next_turn is not None
|
||||
assert turn_end.meta.last_chunk is True
|
||||
assert turn_end.meta.turn_end is True
|
||||
assert turn_end.meta.is_segment_finished.item() is True
|
||||
assert next_turn.meta.cache_epoch == turn_end.meta.cache_epoch + 1
|
||||
assert next_turn.meta.chunk_seq == 0
|
||||
assert next_turn.meta.last_chunk is False
|
||||
|
||||
|
||||
def test_staggered_requests_keep_accumulators_isolated() -> None:
|
||||
manager = _manager()
|
||||
req_a = _request("a")
|
||||
|
||||
@@ -128,13 +128,125 @@ def test_native_duplex_speak_segment_reaches_split_talker() -> None:
|
||||
info = converted["model_intermediate_buffer"]
|
||||
assert info["native_duplex"] is True
|
||||
assert info["ids"]["tts"] == [21, 22]
|
||||
assert converted["prompt_token_ids"] == [0, 0, 0, 0]
|
||||
assert converted["prompt_token_ids"] == [0, 0, 0]
|
||||
assert info["meta"]["replace_streaming_prompt"] is True
|
||||
assert info["meta"]["next_stage_prompt_len"] == 3
|
||||
assert info["meta"]["turn_start"] is True
|
||||
assert info["meta"]["segment_end"] is True
|
||||
assert info["duplex"]["epoch"] == 3
|
||||
assert info["duplex"]["turn_id"] == 7
|
||||
|
||||
|
||||
def test_native_duplex_continuation_appends_only_new_talker_condition() -> None:
|
||||
prompt_ids = [101, 102]
|
||||
token_ids = {
|
||||
"tts_bos_token_id": 9301,
|
||||
"tts_eos_token_id": 9302,
|
||||
"listen_token_id": 9303,
|
||||
"speak_token_id": 9304,
|
||||
"chunk_eos_token_id": 9308,
|
||||
"chunk_tts_eos_token_id": 9309,
|
||||
"turn_eos_token_id": 9310,
|
||||
}
|
||||
context = SimpleNamespace(
|
||||
bridge_states={
|
||||
"duplex": {
|
||||
"epoch": 3,
|
||||
"model_turn_id": 7,
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
first_ids = [9304, 21, 22, 9308]
|
||||
first = _output(
|
||||
prompt_ids=prompt_ids,
|
||||
output_ids=first_ids,
|
||||
latent=torch.arange(24, dtype=torch.float32).reshape(6, 4),
|
||||
multimodal_output={
|
||||
"duplex_prompt_token_ids": prompt_ids,
|
||||
"meta": token_ids,
|
||||
},
|
||||
)
|
||||
second_ids = [*first_ids, 9304, 23, 24, 9308]
|
||||
second = _output(
|
||||
prompt_ids=prompt_ids,
|
||||
output_ids=second_ids,
|
||||
latent=torch.arange(40, dtype=torch.float32).reshape(10, 4),
|
||||
multimodal_output={
|
||||
"duplex_prompt_token_ids": prompt_ids,
|
||||
"meta": token_ids,
|
||||
},
|
||||
)
|
||||
third_ids = [*second_ids, 9304, 25, 26, 9308]
|
||||
third = _output(
|
||||
prompt_ids=prompt_ids,
|
||||
output_ids=third_ids,
|
||||
latent=torch.arange(56, dtype=torch.float32).reshape(14, 4),
|
||||
multimodal_output={
|
||||
"duplex_prompt_token_ids": prompt_ids,
|
||||
"meta": token_ids,
|
||||
},
|
||||
)
|
||||
|
||||
first_input = llm2tts([first], prompt=[{}], _streaming_context=context)[0]
|
||||
second_input = llm2tts([second], prompt=[{}], _streaming_context=context)[0]
|
||||
context.bridge_states["duplex"]["model_turn_id"] = 8
|
||||
third_input = llm2tts([third], prompt=[{}], _streaming_context=context)[0]
|
||||
|
||||
assert first_input["model_intermediate_buffer"]["ids"]["tts"] == [21, 22]
|
||||
assert second_input["model_intermediate_buffer"]["ids"]["tts"] == [23, 24]
|
||||
assert third_input["model_intermediate_buffer"]["ids"]["tts"] == [25, 26]
|
||||
assert first_input["model_intermediate_buffer"]["meta"]["turn_start"] is True
|
||||
assert second_input["model_intermediate_buffer"]["meta"]["turn_start"] is False
|
||||
assert third_input["model_intermediate_buffer"]["meta"]["turn_start"] is True
|
||||
assert first_input["model_intermediate_buffer"]["meta"]["replace_streaming_prompt"] is True
|
||||
assert "replace_streaming_prompt" not in second_input["model_intermediate_buffer"]["meta"]
|
||||
assert third_input["model_intermediate_buffer"]["meta"]["replace_streaming_prompt"] is True
|
||||
assert second_input["model_intermediate_buffer"]["meta"]["next_stage_prompt_len"] == 3
|
||||
assert second_input["prompt_token_ids"] == [0, 0, 0]
|
||||
|
||||
|
||||
def test_native_duplex_transcript_decodes_the_talker_condition_slice() -> None:
|
||||
prompt_ids = [101, 102]
|
||||
metadata = {
|
||||
"tts_bos_token_id": 9301,
|
||||
"tts_eos_token_id": 9302,
|
||||
"listen_token_id": 9303,
|
||||
"speak_token_id": 9304,
|
||||
"chunk_eos_token_id": 9308,
|
||||
"chunk_tts_eos_token_id": 9309,
|
||||
"turn_eos_token_id": 9310,
|
||||
}
|
||||
token_text = {21: "杭州", 22: "和", 23: "州和", 24: "上海之间", 25: "大"}
|
||||
|
||||
def output(token_ids: list[int]):
|
||||
return _output(
|
||||
prompt_ids=prompt_ids,
|
||||
output_ids=token_ids,
|
||||
latent=torch.zeros((len(prompt_ids) + len(token_ids), 1)),
|
||||
multimodal_output={
|
||||
"duplex_prompt_token_ids": prompt_ids,
|
||||
"meta": metadata,
|
||||
},
|
||||
)
|
||||
|
||||
context = SimpleNamespace(
|
||||
bridge_states={"duplex": {"epoch": 3, "model_turn_id": 7}},
|
||||
source_token_decoder=lambda ids, **_: "".join(token_text.get(int(token_id), "") for token_id in ids),
|
||||
)
|
||||
first_ids = [9304, 21, 22, 9308]
|
||||
first_info = llm2tts([output(first_ids)], prompt=[{}], _streaming_context=context)[0]["model_intermediate_buffer"]
|
||||
second_info = llm2tts(
|
||||
[output([*first_ids, 23, 24, 25, 9308])],
|
||||
prompt=[{}],
|
||||
_streaming_context=context,
|
||||
)[0]["model_intermediate_buffer"]
|
||||
|
||||
assert first_info["meta"]["native_duplex_segment_text"] == "杭州和"
|
||||
assert second_info["ids"]["tts"] == [24, 25]
|
||||
assert second_info["meta"]["native_duplex_segment_text"] == "上海之间大"
|
||||
|
||||
|
||||
def test_native_duplex_requires_tokenizer_boundary_metadata() -> None:
|
||||
latent = torch.zeros((3, 4))
|
||||
source = _output(
|
||||
|
||||
@@ -4,6 +4,7 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections import defaultdict
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
@@ -227,6 +228,47 @@ def test_talker2code2wav_full_payload_keeps_all_zero_codec_rows() -> None:
|
||||
assert "code_predictor_codes" not in payload
|
||||
|
||||
|
||||
def test_talker2code2wav_async_chunk_flushes_cached_tail_on_stop_token() -> None:
|
||||
request_id = "codec_tail"
|
||||
stop_token_id = 999
|
||||
transfer_manager = SimpleNamespace(
|
||||
code_prompt_token_ids=defaultdict(list),
|
||||
put_req_chunk=defaultdict(int, {request_id: 2}),
|
||||
connector=SimpleNamespace(
|
||||
config={
|
||||
"extra": {
|
||||
"initial_codec_chunk_frames": 4,
|
||||
"codec_chunk_frames": 25,
|
||||
"codec_left_context_frames": 25,
|
||||
}
|
||||
}
|
||||
),
|
||||
)
|
||||
transfer_manager.code_prompt_token_ids[request_id] = [
|
||||
torch.tensor([[frame, frame + 100]], dtype=torch.long) for frame in range(50)
|
||||
]
|
||||
request = SimpleNamespace(
|
||||
external_req_id=request_id,
|
||||
sampling_params=SimpleNamespace(
|
||||
stop_token_ids=[stop_token_id],
|
||||
stop_token_id=None,
|
||||
),
|
||||
)
|
||||
|
||||
payload = q3.talker2code2wav_async_chunk(
|
||||
transfer_manager,
|
||||
{"codes": {"audio": torch.tensor([[stop_token_id, 0]], dtype=torch.long)}},
|
||||
request,
|
||||
is_finished=True,
|
||||
)
|
||||
|
||||
assert payload is not None
|
||||
assert payload.meta.finished.item() is True
|
||||
# 4 initial frames and one 25-frame chunk were already emitted. The final
|
||||
# payload contains 25 frames of left context plus the remaining 21 frames.
|
||||
assert payload.codes.audio.numel() == (25 + 21) * 2
|
||||
|
||||
|
||||
def test_thinker2talker_full_payload_packs_complete_tensors() -> None:
|
||||
"""Full-payload path drops the terminal thinker row before talker prefill."""
|
||||
request = SimpleNamespace(
|
||||
|
||||
@@ -1160,8 +1160,10 @@ class TestDeployConfigLoading:
|
||||
assert [stage.to_omegaconf().session_mode for stage in stages] == ["duplex", "duplex", "duplex"]
|
||||
assert [stage.yaml_engine_args["async_scheduling"] for stage in stages] == [False, False, False]
|
||||
assert all("Async" not in (stage.scheduler_cls or "") for stage in stages)
|
||||
assert [stage.devices for stage in deploy.stages] == ["0", "0", "0"]
|
||||
assert deploy.stages[1].enforce_eager is False
|
||||
assert stages[1].yaml_extras["default_sampling_params"]["max_tokens"] == 2048
|
||||
assert stages[1].yaml_extras["default_sampling_params"]["min_tokens"] == 0
|
||||
assert stages[1].yaml_extras["default_sampling_params"]["stop_token_ids"] == [1]
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
|
||||
@@ -276,16 +276,16 @@ def test_minicpmo_model_hook_same_speech_row_second_step_does_not_rearm_pending_
|
||||
next_logits[0, 7] = 30.0
|
||||
next_logits[0, 8] = -2.0
|
||||
next_logits[0, 10] = 20.0
|
||||
original_next_logits = next_logits.clone()
|
||||
|
||||
model.prepare_duplex_sampling(next_logits, SimpleNamespace(), (row,))
|
||||
|
||||
assert state.pending_speech_context is False
|
||||
assert state.current_turn_ended is False
|
||||
assert torch.isneginf(next_logits[0, 7])
|
||||
assert next_logits[0, 8].item() == 30.0
|
||||
assert torch.equal(next_logits, original_next_logits)
|
||||
|
||||
|
||||
def test_minicpmo_model_hook_mid_turn_speech_redirects_listen_to_tts_bos():
|
||||
def test_minicpmo_model_hook_mid_turn_speech_preserves_logits_for_post_sample_redirect():
|
||||
state = SimpleNamespace(
|
||||
current_turn_ended=False,
|
||||
last_terminator_token=None,
|
||||
@@ -296,12 +296,11 @@ def test_minicpmo_model_hook_mid_turn_speech_redirects_listen_to_tts_bos():
|
||||
logits[0, 7] = 30.0
|
||||
logits[0, 8] = -2.0
|
||||
logits[0, 10] = 20.0
|
||||
original_logits = logits.clone()
|
||||
|
||||
model.prepare_duplex_sampling(logits, SimpleNamespace(), (row,))
|
||||
|
||||
assert torch.isneginf(logits[0, 7])
|
||||
assert logits[0, 8].item() == 30.0
|
||||
assert logits[0, 10].item() == 20.0
|
||||
assert torch.equal(logits, original_logits)
|
||||
|
||||
|
||||
def test_minicpmo_model_hook_ignores_serving_new_user_turn_marker():
|
||||
@@ -317,12 +316,12 @@ def test_minicpmo_model_hook_ignores_serving_new_user_turn_marker():
|
||||
logits = torch.zeros((1, 16), dtype=torch.float32)
|
||||
logits[0, 7] = 5.0
|
||||
logits[0, 10] = 20.0
|
||||
original_logits = logits.clone()
|
||||
model.prepare_duplex_sampling(logits, SimpleNamespace(), (row,))
|
||||
|
||||
assert state.current_turn_ended is False
|
||||
assert state.last_terminator_token is None
|
||||
assert torch.isneginf(logits[0, 7])
|
||||
assert logits[0, 8].item() == 5.0
|
||||
assert torch.equal(logits, original_logits)
|
||||
|
||||
|
||||
def test_generic_ar_runner_builds_typed_duplex_sampling_rows():
|
||||
@@ -613,6 +612,180 @@ def test_minicpmo_stage0_rejects_unknown_special_token_fallbacks():
|
||||
runtime._require_special_token_ids()
|
||||
|
||||
|
||||
def _minicpmo45_tokenizer_stub():
|
||||
token_ids = {
|
||||
token: index
|
||||
for index, token in enumerate(
|
||||
(
|
||||
"<unit>",
|
||||
"</unit>",
|
||||
"<|listen|>",
|
||||
"<|speak|>",
|
||||
"<|tts_bos|>",
|
||||
"<|tts_eos|>",
|
||||
"<|tts_pad|>",
|
||||
"<|chunk_eos|>",
|
||||
"<|chunk_tts_eos|>",
|
||||
"<|turn_eos|>",
|
||||
),
|
||||
start=101,
|
||||
)
|
||||
}
|
||||
return SimpleNamespace(
|
||||
unk_token_id=0,
|
||||
convert_tokens_to_ids=lambda token: token_ids.get(token, 0),
|
||||
encode=lambda token, add_special_tokens=False: (
|
||||
[token_ids[token]] if not add_special_tokens and token in token_ids else [0]
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _transformers_processor_stub(auto_processor):
|
||||
class _AutoImageProcessor:
|
||||
@staticmethod
|
||||
def register(*_args, **_kwargs):
|
||||
return None
|
||||
|
||||
return SimpleNamespace(
|
||||
AutoImageProcessor=_AutoImageProcessor,
|
||||
AutoProcessor=auto_processor,
|
||||
)
|
||||
|
||||
|
||||
def test_minicpmo_stage0_loads_processor_from_hf_id(monkeypatch):
|
||||
import sys
|
||||
|
||||
from vllm_omni.experimental.fullduplex.minicpmo45.stage0 import (
|
||||
MiniCPMO45Stage0DuplexRuntime,
|
||||
)
|
||||
|
||||
model_id = "openbmb/MiniCPM-o-4_5"
|
||||
processor = SimpleNamespace(tokenizer=_minicpmo45_tokenizer_stub())
|
||||
load_calls = []
|
||||
|
||||
class _AutoImageProcessor:
|
||||
@staticmethod
|
||||
def register(config_class, image_processor_class):
|
||||
del image_processor_class
|
||||
if isinstance(config_class, str):
|
||||
raise AttributeError("'str' object has no attribute '__module__'")
|
||||
|
||||
class _AutoProcessor:
|
||||
@classmethod
|
||||
def from_pretrained(cls, model_path, *, trust_remote_code):
|
||||
load_calls.append((model_path, trust_remote_code))
|
||||
_AutoImageProcessor.register("MiniCPMVImageProcessor", object)
|
||||
return processor
|
||||
|
||||
monkeypatch.setitem(
|
||||
sys.modules,
|
||||
"transformers",
|
||||
SimpleNamespace(
|
||||
AutoImageProcessor=_AutoImageProcessor,
|
||||
AutoProcessor=_AutoProcessor,
|
||||
),
|
||||
)
|
||||
|
||||
loaded = MiniCPMO45Stage0DuplexRuntime._load_processor_from_path(model_id)
|
||||
|
||||
assert loaded is processor
|
||||
assert load_calls == [(model_id, True)]
|
||||
with pytest.raises(AttributeError, match="__module__"):
|
||||
_AutoImageProcessor.register("MiniCPMVImageProcessor", object)
|
||||
|
||||
|
||||
def test_minicpmo_stage0_processor_load_preserves_original_exception(monkeypatch):
|
||||
import sys
|
||||
|
||||
from vllm_omni.experimental.fullduplex.minicpmo45.stage0 import (
|
||||
MiniCPMO45Stage0DuplexRuntime,
|
||||
)
|
||||
|
||||
original = ValueError("invalid processor config")
|
||||
|
||||
class _AutoProcessor:
|
||||
@classmethod
|
||||
def from_pretrained(cls, model_path, *, trust_remote_code):
|
||||
del cls, model_path, trust_remote_code
|
||||
raise original
|
||||
|
||||
monkeypatch.setitem(
|
||||
sys.modules,
|
||||
"transformers",
|
||||
_transformers_processor_stub(_AutoProcessor),
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError, match="Failed to load MiniCPM-o duplex processor") as exc_info:
|
||||
MiniCPMO45Stage0DuplexRuntime._load_processor_from_path("openbmb/MiniCPM-o-4_5")
|
||||
|
||||
assert exc_info.value.__cause__ is original
|
||||
|
||||
|
||||
def test_minicpmo_stage0_processor_requires_tokenizer(monkeypatch):
|
||||
import sys
|
||||
|
||||
from vllm_omni.experimental.fullduplex.minicpmo45.stage0 import (
|
||||
MiniCPMO45Stage0DuplexRuntime,
|
||||
)
|
||||
|
||||
class _AutoProcessor:
|
||||
@classmethod
|
||||
def from_pretrained(cls, model_path, *, trust_remote_code):
|
||||
del cls, model_path, trust_remote_code
|
||||
return object()
|
||||
|
||||
monkeypatch.setitem(
|
||||
sys.modules,
|
||||
"transformers",
|
||||
_transformers_processor_stub(_AutoProcessor),
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError, match="does not expose a tokenizer"):
|
||||
MiniCPMO45Stage0DuplexRuntime._load_processor_from_path("openbmb/MiniCPM-o-4_5")
|
||||
|
||||
|
||||
def test_minicpmo_stage0_loaded_processor_validates_special_tokens(monkeypatch):
|
||||
import sys
|
||||
|
||||
from vllm_omni.experimental.fullduplex.minicpmo45.stage0 import (
|
||||
MiniCPMO45Stage0DuplexRuntime,
|
||||
)
|
||||
|
||||
processor = SimpleNamespace(tokenizer=_minicpmo45_tokenizer_stub())
|
||||
|
||||
class _AutoProcessor:
|
||||
@classmethod
|
||||
def from_pretrained(cls, model_path, *, trust_remote_code):
|
||||
del cls, model_path, trust_remote_code
|
||||
return processor
|
||||
|
||||
monkeypatch.setitem(
|
||||
sys.modules,
|
||||
"transformers",
|
||||
_transformers_processor_stub(_AutoProcessor),
|
||||
)
|
||||
|
||||
runtime = MiniCPMO45Stage0DuplexRuntime(
|
||||
SimpleNamespace(),
|
||||
model_path="openbmb/MiniCPM-o-4_5",
|
||||
device="cpu",
|
||||
)
|
||||
|
||||
assert runtime.processor is processor
|
||||
assert set(runtime._special_token_ids()) == {
|
||||
"unit_token_id",
|
||||
"unit_end_token_id",
|
||||
"listen_token_id",
|
||||
"speak_token_id",
|
||||
"tts_bos_token_id",
|
||||
"tts_eos_token_id",
|
||||
"tts_pad_token_id",
|
||||
"chunk_eos_token_id",
|
||||
"chunk_tts_eos_token_id",
|
||||
"turn_eos_token_id",
|
||||
}
|
||||
|
||||
|
||||
def test_minicpmo_stage0_data_plane_prefill_matches_official_unit_format():
|
||||
import torch
|
||||
|
||||
@@ -1150,6 +1323,88 @@ def test_minicpmo_stage0_data_plane_final_first_chunk_does_not_add_silence_unit(
|
||||
)
|
||||
|
||||
|
||||
def test_minicpmo_stage0_final_does_not_promote_first_chunk_alignment_tail_to_unit():
|
||||
import torch
|
||||
|
||||
from vllm_omni.experimental.fullduplex.minicpmo45.stage0 import (
|
||||
MiniCPMO45Stage0DuplexRuntime,
|
||||
_MiniCPMO45Stage0SessionState,
|
||||
)
|
||||
|
||||
class _StageModel(torch.nn.Module):
|
||||
first_chunk_ms = 10
|
||||
chunk_ms = 8
|
||||
sample_rate = 1000
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.embed = torch.nn.Embedding(256, 2)
|
||||
|
||||
def get_input_embeddings(self):
|
||||
return self.embed
|
||||
|
||||
def get_audio_hidden_states(self, _data):
|
||||
return [torch.tensor([[0.5, 0.5]], dtype=torch.float32)]
|
||||
|
||||
class _MelProcessor:
|
||||
sample_rate = 1000
|
||||
is_first = True
|
||||
|
||||
def get_config(self):
|
||||
return {"effective_first_chunk_ms": 9}
|
||||
|
||||
class _Processor:
|
||||
def __init__(self):
|
||||
self._streaming_mel_processor = _MelProcessor()
|
||||
|
||||
def get_streaming_chunk_size(self):
|
||||
return 9 if self._streaming_mel_processor.is_first else 8
|
||||
|
||||
def process_audio_streaming(self, audio, **_kwargs):
|
||||
self._streaming_mel_processor.is_first = False
|
||||
return {"audio_features": audio, "audio_feature_lens": [[len(audio)]]}
|
||||
|
||||
runtime = MiniCPMO45Stage0DuplexRuntime.__new__(MiniCPMO45Stage0DuplexRuntime)
|
||||
runtime.stage_model = _StageModel()
|
||||
runtime.thinker = runtime.stage_model
|
||||
runtime.tokenizer = SimpleNamespace(
|
||||
unk_token_id=0,
|
||||
convert_tokens_to_ids=lambda token: {
|
||||
"<unit>": 1,
|
||||
"</unit>": 2,
|
||||
"<|listen|>": 3,
|
||||
"<|speak|>": 4,
|
||||
"<|tts_bos|>": 5,
|
||||
"<|tts_eos|>": 6,
|
||||
"<|tts_pad|>": 7,
|
||||
"<|chunk_eos|>": 8,
|
||||
"<|chunk_tts_eos|>": 9,
|
||||
"<|turn_eos|>": 10,
|
||||
"<|audio|>": 11,
|
||||
}.get(token, 0),
|
||||
encode=lambda text, add_special_tokens=False: [],
|
||||
)
|
||||
runtime.processor = _Processor()
|
||||
runtime.device = "cpu"
|
||||
runtime._init_token_ids()
|
||||
state = _MiniCPMO45Stage0SessionState(session_id="sid-first-chunk-alignment-tail")
|
||||
|
||||
first = runtime._stage_prefill_embeddings_only(state, np.zeros(8, dtype=np.float32), seq=1)
|
||||
second = runtime._stage_prefill_embeddings_only(state, np.zeros(8, dtype=np.float32), seq=2)
|
||||
final = runtime._stage_prefill_embeddings_only(
|
||||
state,
|
||||
np.zeros(8, dtype=np.float32),
|
||||
seq=3,
|
||||
final=True,
|
||||
)
|
||||
|
||||
assert first["input_token_ids"].count(1) == 1
|
||||
assert second["input_token_ids"].count(1) == 1
|
||||
assert final["input_token_ids"].count(1) == 1
|
||||
assert state.audio_chunk_idx == 3
|
||||
assert len(state.audio_buffer) == 1
|
||||
|
||||
|
||||
def test_minicpmo_stage0_runtime_uses_loaded_vllm_embed_tokens_when_get_input_embeddings_is_broken():
|
||||
import torch
|
||||
|
||||
@@ -1319,7 +1574,7 @@ def test_minicpmo_stage0_native_sampler_penalizes_text_from_prior_chunk():
|
||||
model.thinker = SimpleNamespace(get_tokenizer=lambda: _Tokenizer())
|
||||
session_key = ("sid-cross-chunk-repetition", 0)
|
||||
state = _MiniCPMO45Stage0SessionState(session_id=session_key[0])
|
||||
state.generated_text_tokens = [198] * 8
|
||||
state.generated_tokens = [198] * 8
|
||||
model._minicpmo45_duplex_data_plane_helper = SimpleNamespace(sessions={session_key: state})
|
||||
model._minicpmo45_duplex_row_sessions = {0: session_key}
|
||||
|
||||
@@ -1381,7 +1636,7 @@ def test_minicpmo_stage0_native_sampler_matches_official_negative_logit_penalty(
|
||||
state = _MiniCPMO45Stage0SessionState(session_id=session_key[0])
|
||||
repeated = 198
|
||||
alternative = 1234
|
||||
state.generated_text_tokens = [repeated]
|
||||
state.generated_tokens = [repeated]
|
||||
model._minicpmo45_duplex_data_plane_helper = SimpleNamespace(sessions={session_key: state})
|
||||
model._minicpmo45_duplex_row_sessions = {0: session_key}
|
||||
|
||||
@@ -1405,7 +1660,7 @@ def test_minicpmo_stage0_native_sampler_matches_official_negative_logit_penalty(
|
||||
assert sampled.sampled_token_ids.tolist() == [[repeated]]
|
||||
|
||||
|
||||
def test_minicpmo_stage0_records_only_bounded_text_repetition_history():
|
||||
def test_minicpmo_stage0_records_bounded_model_policy_history():
|
||||
from vllm_omni.experimental.fullduplex.minicpmo45.stage0 import (
|
||||
_MiniCPMO45Stage0SessionState,
|
||||
)
|
||||
@@ -1432,15 +1687,15 @@ def test_minicpmo_stage0_records_only_bounded_text_repetition_history():
|
||||
}
|
||||
|
||||
for sampled in range(1000, 1520):
|
||||
model._record_minicpmo45_duplex_terminator(0, sampled, token_ids)
|
||||
model._record_minicpmo45_duplex_generation_token(0, sampled)
|
||||
|
||||
expected_history = list(range(1008, 1520))
|
||||
assert state.generated_text_tokens == expected_history
|
||||
assert state.generated_tokens == expected_history
|
||||
|
||||
for sampled in token_ids.values():
|
||||
model._record_minicpmo45_duplex_terminator(0, sampled, token_ids)
|
||||
model._record_minicpmo45_duplex_generation_token(0, sampled)
|
||||
|
||||
assert state.generated_text_tokens == expected_history
|
||||
assert state.generated_tokens == [*expected_history[len(token_ids) :], *token_ids.values()]
|
||||
|
||||
|
||||
def test_minicpmo_stage0_native_sampler_does_not_override_model_at_punctuation():
|
||||
|
||||
@@ -731,9 +731,17 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
|
||||
if self.log_stats:
|
||||
session.record_event(EngineCoreEventType.QUEUED)
|
||||
return
|
||||
update_info = getattr(update, "additional_information", None)
|
||||
update_meta = update_info.get("meta") if isinstance(update_info, dict) else None
|
||||
if isinstance(update_meta, dict) and update_meta.get("replace_streaming_prompt") is True:
|
||||
update_infos = (
|
||||
getattr(update, "model_intermediate_buffer", None),
|
||||
getattr(update, "additional_information", None),
|
||||
)
|
||||
replace_streaming_prompt = any(
|
||||
isinstance(info, dict)
|
||||
and isinstance(info.get("meta"), dict)
|
||||
and info["meta"].get("replace_streaming_prompt") is True
|
||||
for info in update_infos
|
||||
)
|
||||
if replace_streaming_prompt:
|
||||
self._replace_streaming_session(session, update)
|
||||
return
|
||||
super()._update_request_as_session(session, update)
|
||||
|
||||
@@ -54,6 +54,21 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
|
||||
)
|
||||
self._latest_omni_connector_output: OmniConnectorOutput | None = None
|
||||
|
||||
def _handle_stopped_request(self, request: Request) -> bool:
|
||||
if (
|
||||
request.resumable
|
||||
and not request.streaming_queue
|
||||
and self.chunk_transfer_adapter is not None
|
||||
and self.chunk_transfer_adapter.receives_chunks
|
||||
):
|
||||
# Downstream async-chunk stages receive the next segment from the
|
||||
# connector, not from an API StreamingUpdate. Enqueue them as
|
||||
# schedulable before the base class can park them in skipped_waiting.
|
||||
request.status = RequestStatus.WAITING
|
||||
self._enqueue_waiting_request(request)
|
||||
return False
|
||||
return super()._handle_stopped_request(request)
|
||||
|
||||
def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput:
|
||||
"""One-shot generation fast path:
|
||||
- Feed all input tokens of the request at once
|
||||
|
||||
@@ -63,6 +63,11 @@ class OmniSchedulerMixin:
|
||||
session.num_computed_tokens = 0
|
||||
session.prompt_token_ids = new_prompt
|
||||
session.additional_information = update.additional_information or None
|
||||
session.model_intermediate_buffer = getattr(
|
||||
update,
|
||||
"model_intermediate_buffer",
|
||||
None,
|
||||
)
|
||||
session.update_block_hashes()
|
||||
session.num_prompt_tokens = len(new_prompt)
|
||||
session.arrival_time = update.arrival_time
|
||||
|
||||
@@ -3,7 +3,7 @@
|
||||
# The registered minicpmo_4_5 pipeline owns the fixed
|
||||
# Thinker -> Talker -> Code2Wav topology. This overlay changes session
|
||||
# lifetime and bounds all three stages to two live duplex sessions.
|
||||
base_config: minicpmo_4_5_batching.yaml
|
||||
base_config: minicpmo_4_5.yaml
|
||||
pipeline: minicpmo_4_5
|
||||
session_mode: duplex
|
||||
active_stream_window: 1
|
||||
@@ -31,6 +31,7 @@ stages:
|
||||
# segment and turn boundaries travel as explicit output metadata.
|
||||
default_sampling_params:
|
||||
max_tokens: 2048
|
||||
min_tokens: 0
|
||||
stop_token_ids: [1]
|
||||
|
||||
- stage_id: 2
|
||||
|
||||
@@ -231,18 +231,18 @@ def construct_next_stage_streaming_input_prompt(payload_data: dict[str, Any], re
|
||||
ids = payload_data.get("ids", {})
|
||||
meta = payload_data.get("meta", {})
|
||||
replace_prompt = isinstance(meta, dict) and meta.get("replace_streaming_prompt") is True
|
||||
condition_len = meta.get("next_stage_prompt_len") if replace_prompt else None
|
||||
if isinstance(condition_len, int) and condition_len > 0:
|
||||
next_stage_prompt_len = meta.get("next_stage_prompt_len") if isinstance(meta, dict) else None
|
||||
if replace_prompt and isinstance(next_stage_prompt_len, int) and next_stage_prompt_len > 0:
|
||||
# Some downstream stages consume complete, independently conditioned
|
||||
# segments instead of extending an existing KV prefix. The producer
|
||||
# declares that transport behavior explicitly in payload metadata.
|
||||
new_prompt = [0] * condition_len
|
||||
new_prompt = [0] * next_stage_prompt_len
|
||||
request._output_token_ids.clear()
|
||||
request._all_token_ids.clear()
|
||||
request._all_token_ids.extend(new_prompt)
|
||||
request.prompt_token_ids = new_prompt
|
||||
request.num_computed_tokens = 0
|
||||
request.num_prompt_tokens = condition_len
|
||||
request.num_prompt_tokens = next_stage_prompt_len
|
||||
request.update_block_hashes()
|
||||
return
|
||||
prompt_token_ids = ids.get("prompt", None)
|
||||
@@ -255,7 +255,10 @@ def construct_next_stage_streaming_input_prompt(payload_data: dict[str, Any], re
|
||||
assert request.prompt_token_ids is not None
|
||||
# Extend prompt with kept output tokens.
|
||||
request.prompt_token_ids.extend(kept_output_tokens)
|
||||
next_prompt_len = max(1, compute_talker_prompt_ids_length(prompt_token_ids))
|
||||
if isinstance(next_stage_prompt_len, int) and next_stage_prompt_len > 0:
|
||||
next_prompt_len = next_stage_prompt_len
|
||||
else:
|
||||
next_prompt_len = max(1, compute_talker_prompt_ids_length(prompt_token_ids))
|
||||
new_prompt = [0] * next_prompt_len
|
||||
request._all_token_ids.extend(new_prompt or ())
|
||||
request.prompt_token_ids.extend(new_prompt or ())
|
||||
|
||||
@@ -283,7 +283,11 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
|
||||
|
||||
# Empty chunk with more data expected: keep polling.
|
||||
has_new_ids = bool(new_ids.numel()) if use_tensor_codes else bool(new_ids)
|
||||
if not has_new_ids and not payload_finished:
|
||||
if not has_new_ids and payload_segment_finished:
|
||||
# Preserve an explicit scheduler boundary even when it
|
||||
# contains no new codec frames.
|
||||
request.prompt_token_ids = [0]
|
||||
if not has_new_ids and not payload_finished and not payload_segment_finished:
|
||||
# The base recv loop treats False as "not ready yet" and
|
||||
# requeues the request. Do not mark an empty non-terminal
|
||||
# chunk as ready, otherwise Stage1 can consume before the
|
||||
@@ -317,7 +321,13 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
|
||||
multimodal_output=multimodal_output,
|
||||
request=request,
|
||||
# Existing processors use is_finished as a flush signal.
|
||||
is_finished=is_segment_finished,
|
||||
# Terminal stops no longer count as segment boundaries
|
||||
# (is_segment_finished is False when the request finishes,
|
||||
# see #5383), but the processor must still flush its
|
||||
# accumulated tail on the terminal chunk — otherwise the
|
||||
# downstream stage receives the finished marker without
|
||||
# the final payload (#5413).
|
||||
is_finished=is_segment_finished or is_finished,
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
@@ -332,7 +342,8 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
|
||||
if payload_data.meta is None:
|
||||
payload_data.meta = MetaStruct()
|
||||
payload_data.meta.finished = torch.tensor(is_finished, dtype=torch.bool)
|
||||
payload_data.meta.is_segment_finished = torch.tensor(is_segment_finished, dtype=torch.bool)
|
||||
if payload_data.meta.is_segment_finished is None:
|
||||
payload_data.meta.is_segment_finished = torch.tensor(is_segment_finished, dtype=torch.bool)
|
||||
|
||||
success, size, metadata = self.connector.put(
|
||||
from_stage=str(stage_id),
|
||||
@@ -725,6 +736,14 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
|
||||
queue_snapshot = list(queue)
|
||||
for request in queue_snapshot:
|
||||
if not self._ensure_active_stream(request):
|
||||
if target_status == RequestStatus.WAITING:
|
||||
# A non-active placeholder must not remain visible to the
|
||||
# scheduler: it has no connector payload yet, so running
|
||||
# it would execute the downstream model with empty
|
||||
# additional_information. Park it until restore_queues()
|
||||
# and retry admission on the next scheduler tick.
|
||||
queue.remove(request)
|
||||
waiting_for_chunk_list.append(request)
|
||||
continue
|
||||
if request.status != RequestStatus.WAITING_FOR_CHUNK:
|
||||
if request.request_id in self.requests_with_ready_chunks:
|
||||
|
||||
@@ -209,10 +209,10 @@ class AsyncOmniEngine:
|
||||
self._omni_master_port,
|
||||
)
|
||||
|
||||
# Stage resolution pops deploy_config, so get the pipeline endpoint
|
||||
# restriction beforehand. TODO (Alex) make this cleaner and refactor
|
||||
# stage config resolution to remove kwargs hacks.
|
||||
deploy_config_path = kwargs.get("deploy_config")
|
||||
# Stage resolution pops deploy_config, so get pipeline-wide settings
|
||||
# beforehand. The stage CLI exposes the same deploy YAML through
|
||||
# stage_configs_path.
|
||||
deploy_config_path = kwargs.get("deploy_config") or kwargs.get("stage_configs_path")
|
||||
pipeline_config = StageConfigFactory.get_pipeline_config(
|
||||
model=model,
|
||||
trust_remote_code=trust_remote_code,
|
||||
|
||||
@@ -13,7 +13,7 @@ from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import time as _time
|
||||
from collections.abc import Callable
|
||||
from collections.abc import Awaitable, Callable
|
||||
from dataclasses import dataclass, field
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
@@ -218,11 +218,18 @@ class _OrchestratorDuplexStagePort:
|
||||
request_states: dict[str, OrchestratorRequestState],
|
||||
running_counter: OmniRequestCounter | None,
|
||||
cleanup_request_ids: Callable[..., Any],
|
||||
async_chunk: bool,
|
||||
prewarm_async_chunk_stages: Callable[
|
||||
[str, Any, OrchestratorRequestState],
|
||||
Awaitable[None],
|
||||
],
|
||||
) -> None:
|
||||
self._stage_pools = stage_pools
|
||||
self._request_states = request_states
|
||||
self._running_counter = running_counter
|
||||
self._cleanup_request_ids = cleanup_request_ids
|
||||
self._async_chunk = async_chunk
|
||||
self._prewarm_async_chunk_stages = prewarm_async_chunk_stages
|
||||
|
||||
@property
|
||||
def stage_count(self) -> int:
|
||||
@@ -299,6 +306,12 @@ class _OrchestratorDuplexStagePort:
|
||||
replica_id = await pool.submit_update(context.request_id, request_state, request)
|
||||
else:
|
||||
replica_id = await pool.submit_initial(context.request_id, request_state, request, prompt_text=None)
|
||||
if self._async_chunk and context.stage_id == 0:
|
||||
await self._prewarm_async_chunk_stages(
|
||||
context.request_id,
|
||||
request,
|
||||
request_state,
|
||||
)
|
||||
request_state.duplex_stage_fences[context.stage_id] = context.fence
|
||||
request_state.stage_submit_ts[context.stage_id] = _time.time()
|
||||
if not request_state.running_counter_registered and self._running_counter is not None:
|
||||
@@ -388,6 +401,8 @@ class Orchestrator:
|
||||
request_states=self.request_states,
|
||||
running_counter=self._running_counter,
|
||||
cleanup_request_ids=self._cleanup_request_ids,
|
||||
async_chunk=self.async_chunk,
|
||||
prewarm_async_chunk_stages=self._prewarm_async_chunk_stages,
|
||||
),
|
||||
result_sink=self.rpc_async_queue,
|
||||
lifecycle_sink=self.output_async_queue,
|
||||
@@ -681,24 +696,16 @@ class Orchestrator:
|
||||
final_stage_id = msg.final_stage_id
|
||||
req_state = self.request_states.get(request_id)
|
||||
if req_state is None:
|
||||
# Streaming updates always follow the first-chunk add_request
|
||||
# through the same ordered submission queue, so an unknown id
|
||||
# here can only mean the request already finished or was
|
||||
# aborted (e.g. the client disconnected). Re-adding it would
|
||||
# resurrect a headless session that keeps cycling through the
|
||||
# stages with nobody consuming its outputs (issue #4271).
|
||||
logger.warning(
|
||||
"[Orchestrator] streaming_update for unknown req=%s, falling back to add_request",
|
||||
"[Orchestrator] streaming_update for unknown req=%s; dropping (request finished or aborted)",
|
||||
request_id,
|
||||
)
|
||||
fallback_msg = StageSubmissionMessage(
|
||||
type="add_request",
|
||||
request_id=msg.request_id,
|
||||
prompt=msg.prompt,
|
||||
original_prompt=msg.original_prompt,
|
||||
output_prompt_text=msg.output_prompt_text,
|
||||
sampling_params_list=msg.sampling_params_list,
|
||||
final_stage_id=msg.final_stage_id,
|
||||
final_output_stage_ids=msg.final_output_stage_ids,
|
||||
preprocess_ms=msg.preprocess_ms,
|
||||
request_timestamp=msg.request_timestamp,
|
||||
enqueue_ts=msg.enqueue_ts,
|
||||
)
|
||||
await self._handle_add_request(fallback_msg)
|
||||
return
|
||||
|
||||
if msg.sampling_params_list:
|
||||
@@ -2025,7 +2032,7 @@ class Orchestrator:
|
||||
_t_submit_start = _time.perf_counter()
|
||||
|
||||
if next_pool.stage_type == "diffusion":
|
||||
await next_pool.submit_initial(
|
||||
replica_id = await next_pool.submit_initial(
|
||||
request_id,
|
||||
req_state,
|
||||
req_state.prompt,
|
||||
@@ -2064,12 +2071,18 @@ class Orchestrator:
|
||||
resumable=downstream_resumable,
|
||||
)
|
||||
request.external_req_id = request.request_id
|
||||
await next_pool.submit_initial(
|
||||
replica_id = await next_pool.submit_initial(
|
||||
request_id,
|
||||
req_state,
|
||||
request,
|
||||
prompt_text=None,
|
||||
)
|
||||
self._record_duplex_stage_submission(
|
||||
next_stage_id,
|
||||
request_id,
|
||||
replica_id,
|
||||
req_state,
|
||||
)
|
||||
|
||||
# async_chunk pre-submit fires per stage edge (N-1 -> N). Source
|
||||
# replica is stage 0's bound replica (single-replica thinker in
|
||||
|
||||
@@ -183,6 +183,10 @@ class OmniMasterServer:
|
||||
# seeded with the registration port (and coordinator ROUTER port) that
|
||||
# are already bound on the same host.
|
||||
self._allocated_ports: set[int] = set()
|
||||
# Keep route ports bound until their ZMQ owners are ready to bind them.
|
||||
# Recording only the integer leaves a race with unrelated startup
|
||||
# consumers (notably NCCL's own ephemeral-port selection).
|
||||
self._port_reservations: dict[int, socket.socket] = {}
|
||||
if master_port:
|
||||
self._allocated_ports.add(int(master_port))
|
||||
# Coordinator ROUTER address echoed back in every registration reply
|
||||
@@ -283,9 +287,11 @@ class OmniMasterServer:
|
||||
``get_open_ports_list`` dedups only within its own call, so we redraw any
|
||||
port that collides with one already recorded in ``self._allocated_ports``
|
||||
(registration/coordinator ports plus every previously-allocated route
|
||||
port) and register the winners before returning. Callers already hold the
|
||||
relevant allocation context (``__init__`` is single-threaded;
|
||||
registration holds ``self._alloc_lock``).
|
||||
port). Each winner remains OS-bound until its ZMQ owner is ready, so
|
||||
another startup subsystem cannot claim it in the allocation-to-bind
|
||||
window. Callers already hold the relevant allocation context
|
||||
(``__init__`` is single-threaded; registration holds
|
||||
``self._alloc_lock``).
|
||||
"""
|
||||
picked: list[int] = []
|
||||
# Bounded retry budget so a pathological host (near port exhaustion)
|
||||
@@ -294,7 +300,17 @@ class OmniMasterServer:
|
||||
for port in get_open_ports_list(count=count - len(picked)):
|
||||
if port in self._allocated_ports:
|
||||
continue
|
||||
reservation = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
||||
try:
|
||||
# Reserve on every local interface. The advertised address
|
||||
# can be a routable interface IP, but wildcard binding also
|
||||
# protects the route from local NCCL/socket allocations.
|
||||
reservation.bind(("", port))
|
||||
except OSError:
|
||||
reservation.close()
|
||||
continue
|
||||
self._allocated_ports.add(port)
|
||||
self._port_reservations[port] = reservation
|
||||
picked.append(port)
|
||||
if len(picked) == count:
|
||||
return picked
|
||||
@@ -303,6 +319,33 @@ class OmniMasterServer:
|
||||
f"(host={self._address}); {len(self._allocated_ports)} already in use."
|
||||
)
|
||||
|
||||
def release_route_port_reservations(
|
||||
self,
|
||||
stage_id: int,
|
||||
replica_id: int = 0,
|
||||
*,
|
||||
handshake: bool = True,
|
||||
data: bool = True,
|
||||
) -> None:
|
||||
"""Release reserved route ports immediately before their ZMQ binds."""
|
||||
alloc = self.get_allocation(stage_id, replica_id)
|
||||
ports: list[int] = []
|
||||
if handshake:
|
||||
handshake_port = _port_from_zmq_address(alloc.handshake_bind_address)
|
||||
if handshake_port is not None:
|
||||
ports.append(handshake_port)
|
||||
if data:
|
||||
for address in (alloc.input_bind_address, alloc.output_bind_address):
|
||||
port = _port_from_zmq_address(address)
|
||||
if port is not None:
|
||||
ports.append(port)
|
||||
|
||||
with self._alloc_lock:
|
||||
for port in ports:
|
||||
reservation = self._port_reservations.pop(port, None)
|
||||
if reservation is not None:
|
||||
reservation.close()
|
||||
|
||||
def _next_free_replica_id(self, stage_id: int) -> int:
|
||||
"""Return the next replica id to assign for an auto-assign registration.
|
||||
|
||||
@@ -437,6 +480,11 @@ class OmniMasterServer:
|
||||
self._stop_event.set()
|
||||
if self._thread is not None:
|
||||
self._thread.join(timeout=10)
|
||||
with self._alloc_lock:
|
||||
reservations = list(self._port_reservations.values())
|
||||
self._port_reservations.clear()
|
||||
for reservation in reservations:
|
||||
reservation.close()
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Internal server logic
|
||||
@@ -771,20 +819,34 @@ def connect_remote_engine_cores(
|
||||
|
||||
handshake_bind_address = omni_master_server.get_allocation(stage_id, replica_id=replica_id).handshake_bind_address
|
||||
|
||||
with zmq_socket_ctx(handshake_bind_address, zmq.ROUTER, bind=True) as handshake_socket:
|
||||
yield StageReplicaResources(
|
||||
coordinator=coordinator,
|
||||
addresses=addresses,
|
||||
)
|
||||
wait_for_engine_startup(
|
||||
handshake_socket,
|
||||
addresses,
|
||||
engines_to_handshake,
|
||||
vllm_config.parallel_config,
|
||||
False, # coordinated_dp
|
||||
vllm_config.cache_config,
|
||||
None, # proc_manager (remote — no local procs)
|
||||
None, # coord_process
|
||||
omni_master_server.release_route_port_reservations(
|
||||
stage_id,
|
||||
replica_id,
|
||||
handshake=True,
|
||||
data=False,
|
||||
)
|
||||
try:
|
||||
with zmq_socket_ctx(handshake_bind_address, zmq.ROUTER, bind=True) as handshake_socket:
|
||||
yield StageReplicaResources(
|
||||
coordinator=coordinator,
|
||||
addresses=addresses,
|
||||
)
|
||||
wait_for_engine_startup(
|
||||
handshake_socket,
|
||||
addresses,
|
||||
engines_to_handshake,
|
||||
vllm_config.parallel_config,
|
||||
False, # coordinated_dp
|
||||
vllm_config.cache_config,
|
||||
None, # proc_manager (remote — no local procs)
|
||||
None, # coord_process
|
||||
)
|
||||
finally:
|
||||
omni_master_server.release_route_port_reservations(
|
||||
stage_id,
|
||||
replica_id,
|
||||
handshake=False,
|
||||
data=True,
|
||||
)
|
||||
|
||||
|
||||
@@ -802,17 +864,31 @@ def connect_remote_diffusion_proc(
|
||||
).handshake_bind_address
|
||||
|
||||
logger.info("Waiting for remote diffusion proc for stage %d replica %d", stage_id, replica_id)
|
||||
with zmq_socket_ctx(handshake_bind_address, zmq.ROUTER, bind=True) as handshake_socket:
|
||||
yield StageReplicaResources(addresses=addresses)
|
||||
wait_for_engine_startup(
|
||||
handshake_socket,
|
||||
addresses,
|
||||
[CoreEngine(index=0, local=False)],
|
||||
_single_diffusion_parallel_config(local_client=False),
|
||||
False,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
omni_master_server.release_route_port_reservations(
|
||||
stage_id,
|
||||
replica_id,
|
||||
handshake=True,
|
||||
data=False,
|
||||
)
|
||||
try:
|
||||
with zmq_socket_ctx(handshake_bind_address, zmq.ROUTER, bind=True) as handshake_socket:
|
||||
yield StageReplicaResources(addresses=addresses)
|
||||
wait_for_engine_startup(
|
||||
handshake_socket,
|
||||
addresses,
|
||||
[CoreEngine(index=0, local=False)],
|
||||
_single_diffusion_parallel_config(local_client=False),
|
||||
False,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
finally:
|
||||
omni_master_server.release_route_port_reservations(
|
||||
stage_id,
|
||||
replica_id,
|
||||
handshake=False,
|
||||
data=True,
|
||||
)
|
||||
|
||||
|
||||
@@ -918,49 +994,63 @@ def _launch_omni_core_engines(
|
||||
# Bind the pre-allocated handshake socket for this stage.
|
||||
handshake_bind_address = omni_master_server.get_allocation(stage_id, replica_id=replica_id).handshake_bind_address
|
||||
|
||||
with zmq_socket_ctx(handshake_bind_address, zmq.ROUTER, bind=True) as handshake_socket:
|
||||
if omni_coordinator_address is not None:
|
||||
# Use the omni subclass so each spawned subprocess instantiates
|
||||
# an OmniCoordClientForStage and heartbeats to the coordinator.
|
||||
from vllm_omni.engine.stage_engine_core_proc_manager import StageEngineCoreProcManager
|
||||
omni_master_server.release_route_port_reservations(
|
||||
stage_id,
|
||||
replica_id,
|
||||
handshake=True,
|
||||
data=False,
|
||||
)
|
||||
try:
|
||||
with zmq_socket_ctx(handshake_bind_address, zmq.ROUTER, bind=True) as handshake_socket:
|
||||
if omni_coordinator_address is not None:
|
||||
# Use the omni subclass so each spawned subprocess instantiates
|
||||
# an OmniCoordClientForStage and heartbeats to the coordinator.
|
||||
from vllm_omni.engine.stage_engine_core_proc_manager import StageEngineCoreProcManager
|
||||
|
||||
with scoped_spawn_device_env(stage_visible_devices, spawn_device_lock):
|
||||
local_engine_manager: CoreEngineProcManager = StageEngineCoreProcManager(
|
||||
local_engine_count=local_engine_count,
|
||||
start_index=start_index,
|
||||
local_start_index=local_start_index,
|
||||
vllm_config=vllm_config,
|
||||
local_client=True,
|
||||
handshake_address=handshake_address,
|
||||
executor_class=executor_class,
|
||||
log_stats=log_stats,
|
||||
omni_stage_id=stage_id,
|
||||
omni_coordinator_address=omni_coordinator_address,
|
||||
omni_replica_base_id=replica_id,
|
||||
)
|
||||
else:
|
||||
with scoped_spawn_device_env(stage_visible_devices, spawn_device_lock):
|
||||
local_engine_manager = CoreEngineProcManager(
|
||||
local_engine_count=local_engine_count,
|
||||
start_index=start_index,
|
||||
local_start_index=local_start_index,
|
||||
vllm_config=vllm_config,
|
||||
local_client=True,
|
||||
handshake_address=handshake_address,
|
||||
executor_class=executor_class,
|
||||
log_stats=log_stats,
|
||||
)
|
||||
with scoped_spawn_device_env(stage_visible_devices, spawn_device_lock):
|
||||
local_engine_manager: CoreEngineProcManager = StageEngineCoreProcManager(
|
||||
local_engine_count=local_engine_count,
|
||||
start_index=start_index,
|
||||
local_start_index=local_start_index,
|
||||
vllm_config=vllm_config,
|
||||
local_client=True,
|
||||
handshake_address=handshake_address,
|
||||
executor_class=executor_class,
|
||||
log_stats=log_stats,
|
||||
omni_stage_id=stage_id,
|
||||
omni_coordinator_address=omni_coordinator_address,
|
||||
omni_replica_base_id=replica_id,
|
||||
)
|
||||
else:
|
||||
with scoped_spawn_device_env(stage_visible_devices, spawn_device_lock):
|
||||
local_engine_manager = CoreEngineProcManager(
|
||||
local_engine_count=local_engine_count,
|
||||
start_index=start_index,
|
||||
local_start_index=local_start_index,
|
||||
vllm_config=vllm_config,
|
||||
local_client=True,
|
||||
handshake_address=handshake_address,
|
||||
executor_class=executor_class,
|
||||
log_stats=log_stats,
|
||||
)
|
||||
|
||||
yield local_engine_manager, coordinator, addresses
|
||||
wait_for_engine_startup(
|
||||
handshake_socket,
|
||||
addresses,
|
||||
engines_to_handshake,
|
||||
parallel_config,
|
||||
parallel_config.data_parallel_size > 1 and vllm_config.model_config.is_moe,
|
||||
vllm_config.cache_config,
|
||||
local_engine_manager,
|
||||
coordinator.proc if coordinator else None,
|
||||
yield local_engine_manager, coordinator, addresses
|
||||
wait_for_engine_startup(
|
||||
handshake_socket,
|
||||
addresses,
|
||||
engines_to_handshake,
|
||||
parallel_config,
|
||||
parallel_config.data_parallel_size > 1 and vllm_config.model_config.is_moe,
|
||||
vllm_config.cache_config,
|
||||
local_engine_manager,
|
||||
coordinator.proc if coordinator else None,
|
||||
)
|
||||
finally:
|
||||
omni_master_server.release_route_port_reservations(
|
||||
stage_id,
|
||||
replica_id,
|
||||
handshake=False,
|
||||
data=True,
|
||||
)
|
||||
|
||||
|
||||
@@ -1448,6 +1538,12 @@ def launch_diffusion_stage_replica(
|
||||
omni_stage_config=stage_config,
|
||||
replica_id=replica_id,
|
||||
)
|
||||
omni_master_server.release_route_port_reservations(
|
||||
metadata.stage_id,
|
||||
replica_id,
|
||||
handshake=True,
|
||||
data=False,
|
||||
)
|
||||
proc_manager = stage_diffusion_proc.StageDiffusionProcManager(
|
||||
model=model,
|
||||
od_config=od_config,
|
||||
@@ -1461,6 +1557,12 @@ def launch_diffusion_stage_replica(
|
||||
omni_stage_id=metadata.stage_id,
|
||||
omni_replica_id=replica_id,
|
||||
)
|
||||
omni_master_server.release_route_port_reservations(
|
||||
metadata.stage_id,
|
||||
replica_id,
|
||||
handshake=False,
|
||||
data=True,
|
||||
)
|
||||
client = StageDiffusionClient.from_addresses(
|
||||
metadata,
|
||||
request_address=proc_manager.addresses.inputs[0],
|
||||
@@ -1474,6 +1576,10 @@ def launch_diffusion_stage_replica(
|
||||
lock_fds=lock_fds,
|
||||
)
|
||||
except Exception:
|
||||
omni_master_server.release_route_port_reservations(
|
||||
metadata.stage_id,
|
||||
replica_id,
|
||||
)
|
||||
if proc_manager is not None:
|
||||
proc_manager.shutdown()
|
||||
if lock_fds:
|
||||
|
||||
@@ -628,6 +628,13 @@ class AsyncOmni(EngineClient, OmniBase):
|
||||
|
||||
logger.debug(f"[AsyncOmni] Request {request_id} completed")
|
||||
|
||||
# The input pump can outlive a normally-completed generation
|
||||
# (e.g. a realtime client that keeps streaming after the final
|
||||
# output); a live pump keeps issuing streaming updates for a
|
||||
# request the orchestrator no longer tracks.
|
||||
if input_stream_task is not None and not input_stream_task.done():
|
||||
input_stream_task.cancel()
|
||||
|
||||
self._log_summary_and_cleanup(request_id)
|
||||
|
||||
except (asyncio.CancelledError, GeneratorExit):
|
||||
@@ -1059,7 +1066,10 @@ class AsyncOmni(EngineClient, OmniBase):
|
||||
"""Submit request IDs to be aborted to the engine."""
|
||||
await self.engine.abort_async(request_ids)
|
||||
for rid in request_ids:
|
||||
self.request_states.pop(rid, None)
|
||||
state = self.request_states.pop(rid, None)
|
||||
input_stream_task = getattr(state, "input_stream_task", None)
|
||||
if input_stream_task is not None and not input_stream_task.done():
|
||||
input_stream_task.cancel()
|
||||
if self.log_stats:
|
||||
logger.info("[AsyncOmni] Aborted request(s) %s", ",".join(request_ids))
|
||||
|
||||
|
||||
@@ -892,7 +892,7 @@ def run_headless(args: TrackingNamespace) -> None:
|
||||
stage_connector_spec = get_stage_connector_spec(
|
||||
omni_transfer_config=omni_transfer_config,
|
||||
stage_id=stage_id,
|
||||
async_chunk=False,
|
||||
async_chunk=bool(stage_cfg.engine_args.get("async_chunk", False)),
|
||||
)
|
||||
|
||||
# ``runtime_cfg`` is mostly inherited from the parent's
|
||||
|
||||
@@ -138,6 +138,7 @@ class RealtimeConnection(VllmRealtimeConnection):
|
||||
is_streaming=True,
|
||||
)
|
||||
|
||||
result_gen = None
|
||||
try:
|
||||
result_gen = self.engine.generate(
|
||||
prompt=streaming_input_gen,
|
||||
@@ -196,6 +197,16 @@ class RealtimeConnection(VllmRealtimeConnection):
|
||||
logger.exception("Error in generation: %s", e)
|
||||
await self.send_error(str(e), "processing_error")
|
||||
finally:
|
||||
# Close the generator explicitly so AsyncOmni.generate's cleanup
|
||||
# (input-pump cancellation and engine-side abort) runs now rather
|
||||
# than whenever the event loop garbage-collects the async
|
||||
# generator; the delay window is where a disconnected session
|
||||
# keeps cycling through the stages (issue #4271).
|
||||
if result_gen is not None:
|
||||
try:
|
||||
await result_gen.aclose()
|
||||
except Exception:
|
||||
logger.exception("Failed to close realtime result generator")
|
||||
# Always send terminal event so clients don't hang forever.
|
||||
if self._is_connected and not audio_done_sent:
|
||||
try:
|
||||
@@ -206,4 +217,10 @@ class RealtimeConnection(VllmRealtimeConnection):
|
||||
self.audio_queue.get_nowait()
|
||||
|
||||
async def send_json(self, payload: dict):
|
||||
await self.websocket.send_text(json.dumps(payload))
|
||||
try:
|
||||
await self.websocket.send_text(json.dumps(payload))
|
||||
except Exception:
|
||||
# A failed send means the client is gone; flag it so the
|
||||
# generation loop stops instead of retrying into a dead socket.
|
||||
self._is_connected = False
|
||||
raise
|
||||
@@ -65,12 +65,11 @@ class MiniCPMO45DataPlaneSession:
|
||||
def begin_request(self, request_id: str) -> None:
|
||||
state = self._requests.setdefault(request_id, _RequestState())
|
||||
state.terminal = False
|
||||
# Preserve turn-scoped cursors for resumable requests. The legacy path
|
||||
# only reset request-scoped terminal markers on append.
|
||||
request_turn = state.turns.get(None)
|
||||
if request_turn is not None:
|
||||
request_turn.tts_eos_done = False
|
||||
request_turn.turn_eos_done = False
|
||||
# A resumable append starts a new Talker segment, even when it
|
||||
# continues the same model turn. Re-arm only the segment boundary;
|
||||
# turn EOS remains deduplicated until the model advances the turn.
|
||||
for turn_state in state.turns.values():
|
||||
turn_state.tts_eos_done = False
|
||||
|
||||
def is_terminal(self, request_id: str | None) -> bool:
|
||||
if request_id is None:
|
||||
|
||||
@@ -2,11 +2,19 @@ from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import binascii
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _PcmSpan:
|
||||
byte_count: int
|
||||
force_listen: bool
|
||||
is_speech: bool
|
||||
|
||||
|
||||
class MiniCPMO45PcmAppendReservation:
|
||||
__slots__ = (
|
||||
"_active",
|
||||
@@ -15,6 +23,7 @@ class MiniCPMO45PcmAppendReservation:
|
||||
"_owner",
|
||||
"_raw",
|
||||
"_sample_rate_hz",
|
||||
"_spans",
|
||||
"_turn_had_speech",
|
||||
"_video_frames",
|
||||
"operation_id",
|
||||
@@ -31,6 +40,7 @@ class MiniCPMO45PcmAppendReservation:
|
||||
sample_rate_hz: int,
|
||||
force_listen: bool,
|
||||
is_speech: bool,
|
||||
spans: list[_PcmSpan] | None = None,
|
||||
turn_had_speech: bool = False,
|
||||
video_frames: list[str] | None = None,
|
||||
) -> None:
|
||||
@@ -41,6 +51,7 @@ class MiniCPMO45PcmAppendReservation:
|
||||
self._sample_rate_hz = sample_rate_hz
|
||||
self._force_listen = force_listen
|
||||
self._is_speech = is_speech
|
||||
self._spans = list(spans or [])
|
||||
self._turn_had_speech = turn_had_speech
|
||||
self._video_frames = list(video_frames or [])
|
||||
self._active = True
|
||||
@@ -97,9 +108,8 @@ class MiniCPMO45PcmAppendBuffer:
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._buffer = bytearray()
|
||||
self._spans: list[_PcmSpan] = []
|
||||
self._sample_rate_hz: int | None = None
|
||||
self._force_listen = False
|
||||
self._is_speech = False
|
||||
self._turn_had_speech = False
|
||||
self._reservation_seq = 0
|
||||
self._reservations: list[MiniCPMO45PcmAppendReservation] = []
|
||||
@@ -112,14 +122,13 @@ class MiniCPMO45PcmAppendBuffer:
|
||||
reservation._active = False
|
||||
self._reservations.clear()
|
||||
self._buffer.clear()
|
||||
self._spans.clear()
|
||||
self._frame_queue.clear()
|
||||
self._sample_rate_hz = None
|
||||
self._force_listen = False
|
||||
self._is_speech = False
|
||||
self._turn_had_speech = False
|
||||
|
||||
def clear_force_listen(self) -> None:
|
||||
self._force_listen = False
|
||||
self._spans = [_PcmSpan(span.byte_count, False, span.is_speech) for span in self._spans]
|
||||
|
||||
def has_pending(self) -> bool:
|
||||
return bool(self._buffer)
|
||||
@@ -131,6 +140,56 @@ class MiniCPMO45PcmAppendBuffer:
|
||||
def pending_byte_count(self) -> int:
|
||||
return len(self._buffer)
|
||||
|
||||
def _append_span(self, span: _PcmSpan) -> None:
|
||||
if span.byte_count <= 0:
|
||||
return
|
||||
if (
|
||||
self._spans
|
||||
and self._spans[-1].force_listen == span.force_listen
|
||||
and self._spans[-1].is_speech == span.is_speech
|
||||
):
|
||||
previous = self._spans[-1]
|
||||
self._spans[-1] = _PcmSpan(
|
||||
previous.byte_count + span.byte_count,
|
||||
span.force_listen,
|
||||
span.is_speech,
|
||||
)
|
||||
return
|
||||
self._spans.append(span)
|
||||
|
||||
def _consume_spans(self, byte_count: int) -> list[_PcmSpan]:
|
||||
consumed: list[_PcmSpan] = []
|
||||
remaining = byte_count
|
||||
while remaining:
|
||||
if not self._spans:
|
||||
raise RuntimeError("MiniCPM-o PCM metadata is shorter than buffered audio")
|
||||
span = self._spans.pop(0)
|
||||
take = min(remaining, span.byte_count)
|
||||
consumed.append(_PcmSpan(take, span.force_listen, span.is_speech))
|
||||
if take < span.byte_count:
|
||||
self._spans.insert(
|
||||
0,
|
||||
_PcmSpan(
|
||||
span.byte_count - take,
|
||||
span.force_listen,
|
||||
span.is_speech,
|
||||
),
|
||||
)
|
||||
remaining -= take
|
||||
return consumed
|
||||
|
||||
def _prepend_spans(self, spans: list[_PcmSpan]) -> None:
|
||||
current = self._spans
|
||||
self._spans = []
|
||||
for span in [*spans, *current]:
|
||||
self._append_span(span)
|
||||
|
||||
def _buffer_force_listen(self) -> bool:
|
||||
return any(span.force_listen for span in self._spans)
|
||||
|
||||
def _buffer_is_speech(self) -> bool:
|
||||
return any(span.is_speech for span in self._spans)
|
||||
|
||||
def _reserve_passthrough(
|
||||
self,
|
||||
payload: dict[str, object],
|
||||
@@ -180,12 +239,17 @@ class MiniCPMO45PcmAppendBuffer:
|
||||
raise ValueError("MiniCPM-o native duplex audio append sample_rate_hz changed within a session")
|
||||
self._sample_rate_hz = sample_rate_hz
|
||||
self._buffer.extend(raw)
|
||||
self._append_span(
|
||||
_PcmSpan(
|
||||
len(raw),
|
||||
bool(payload.get("force_listen", False)),
|
||||
bool(payload.get("is_speech", False)),
|
||||
)
|
||||
)
|
||||
frames_in = payload.get("video_frames")
|
||||
if isinstance(frames_in, list):
|
||||
self._frame_queue.extend(frame for frame in frames_in if isinstance(frame, str) and frame)
|
||||
self._turn_had_speech = self._turn_had_speech or bool(payload.get("is_speech", False))
|
||||
self._force_listen = self._force_listen or bool(payload.get("force_listen", False))
|
||||
self._is_speech = self._is_speech or bool(payload.get("is_speech", False))
|
||||
if not allow_emit:
|
||||
return None
|
||||
|
||||
@@ -208,6 +272,7 @@ class MiniCPMO45PcmAppendBuffer:
|
||||
pad_samples = 0
|
||||
emit_bytes = emit_samples * 4
|
||||
reserved_raw = bytes(self._buffer[:emit_bytes])
|
||||
reserved_spans = self._consume_spans(emit_bytes)
|
||||
emit_raw = reserved_raw + b"\x00" * (pad_samples * 4)
|
||||
del self._buffer[:emit_bytes]
|
||||
|
||||
@@ -227,11 +292,8 @@ class MiniCPMO45PcmAppendBuffer:
|
||||
if emit_samples + pad_samples >= min_samples and self._frame_queue:
|
||||
attached_frames = [self._frame_queue.pop(0)]
|
||||
out["video_frames"] = attached_frames
|
||||
out["force_listen"] = self._force_listen
|
||||
out["is_speech"] = self._is_speech
|
||||
if not self._buffer:
|
||||
self._force_listen = False
|
||||
self._is_speech = False
|
||||
out["force_listen"] = any(span.force_listen for span in reserved_spans)
|
||||
out["is_speech"] = any(span.is_speech for span in reserved_spans)
|
||||
reservation = MiniCPMO45PcmAppendReservation(
|
||||
owner=self,
|
||||
operation_id=operation_id,
|
||||
@@ -240,6 +302,7 @@ class MiniCPMO45PcmAppendBuffer:
|
||||
sample_rate_hz=sample_rate_hz,
|
||||
force_listen=bool(out.get("force_listen", False)),
|
||||
is_speech=bool(out.get("is_speech", False)),
|
||||
spans=reserved_spans,
|
||||
turn_had_speech=self._turn_had_speech,
|
||||
video_frames=attached_frames,
|
||||
)
|
||||
@@ -261,8 +324,8 @@ class MiniCPMO45PcmAppendBuffer:
|
||||
"audio": "",
|
||||
"format": "pcm_f32le",
|
||||
"sample_rate_hz": self._sample_rate_hz or 16000,
|
||||
"force_listen": self._force_listen,
|
||||
"is_speech": self._is_speech,
|
||||
"force_listen": self._buffer_force_listen(),
|
||||
"is_speech": self._buffer_is_speech(),
|
||||
}
|
||||
reservation = self.prepare_append(
|
||||
payload,
|
||||
@@ -280,8 +343,8 @@ class MiniCPMO45PcmAppendBuffer:
|
||||
payload=None,
|
||||
raw=b"",
|
||||
sample_rate_hz=self._sample_rate_hz or 0,
|
||||
force_listen=self._force_listen,
|
||||
is_speech=self._is_speech,
|
||||
force_listen=self._buffer_force_listen(),
|
||||
is_speech=self._buffer_is_speech(),
|
||||
turn_had_speech=had_speech,
|
||||
)
|
||||
self._reservations.append(reservation)
|
||||
@@ -289,8 +352,6 @@ class MiniCPMO45PcmAppendBuffer:
|
||||
# Open the next physical input generation without touching already
|
||||
# reserved appends. A rollback restores this generation's metadata.
|
||||
self._sample_rate_hz = None
|
||||
self._force_listen = False
|
||||
self._is_speech = False
|
||||
self._turn_had_speech = False
|
||||
return reservation
|
||||
|
||||
@@ -335,12 +396,11 @@ class MiniCPMO45PcmAppendBuffer:
|
||||
rolled_back = self._reservations[index:]
|
||||
restored = b"".join(item._raw for item in rolled_back)
|
||||
self._buffer[:0] = restored
|
||||
self._prepend_spans([span for item in rolled_back for span in item._spans])
|
||||
restored_frames = [frame for item in rolled_back for frame in item._video_frames]
|
||||
if restored_frames:
|
||||
self._frame_queue[:0] = restored_frames
|
||||
self._sample_rate_hz = self._sample_rate_hz or reservation._sample_rate_hz
|
||||
self._force_listen = self._force_listen or any(item._force_listen for item in rolled_back)
|
||||
self._is_speech = self._is_speech or any(item._is_speech for item in rolled_back)
|
||||
self._turn_had_speech = self._turn_had_speech or any(item._turn_had_speech for item in rolled_back)
|
||||
for item in rolled_back:
|
||||
item._active = False
|
||||
@@ -354,8 +414,8 @@ class MiniCPMO45PcmAppendBuffer:
|
||||
"audio": "",
|
||||
"format": "pcm_f32le",
|
||||
"sample_rate_hz": self._sample_rate_hz or 16000,
|
||||
"force_listen": self._force_listen,
|
||||
"is_speech": self._is_speech,
|
||||
"force_listen": self._buffer_force_listen(),
|
||||
"is_speech": self._buffer_is_speech(),
|
||||
}
|
||||
return self.append(payload, chunk_period_ms=chunk_period_ms, flush=True)
|
||||
|
||||
|
||||
@@ -2,10 +2,10 @@ from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import copy
|
||||
import sys
|
||||
import time
|
||||
from contextlib import suppress
|
||||
from dataclasses import dataclass, field
|
||||
from threading import Lock
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
@@ -14,6 +14,7 @@ from vllm_omni.experimental.fullduplex.minicpmo45.policy import MiniCPMO45Duplex
|
||||
|
||||
_MINICPMO45_SPECIAL_TOKEN_FIELDS = MiniCPMO45DuplexPolicy.SPECIAL_TOKEN_FIELDS
|
||||
_MINICPMO45_OPTIONAL_TOKEN_FIELDS = MiniCPMO45DuplexPolicy.OPTIONAL_TOKEN_FIELDS
|
||||
_MINICPMO45_PROCESSOR_LOAD_LOCK = Lock()
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -35,7 +36,7 @@ class _MiniCPMO45Stage0SessionState:
|
||||
pending_speech_context: bool = False
|
||||
pending_speech_append_identity: tuple[int | None, int] | None = None
|
||||
pending_speech_response_open: bool = False
|
||||
generated_text_tokens: list[int] = field(default_factory=list)
|
||||
generated_tokens: list[int] = field(default_factory=list)
|
||||
|
||||
|
||||
class MiniCPMO45Stage0DuplexRuntime:
|
||||
@@ -58,6 +59,8 @@ class MiniCPMO45Stage0DuplexRuntime:
|
||||
else getattr(stage_model, "tokenizer", None)
|
||||
)
|
||||
self._init_token_ids()
|
||||
if self.tokenizer is not None:
|
||||
self._require_special_token_ids()
|
||||
|
||||
def _stage_runtime_ready(self) -> bool:
|
||||
return self.processor is not None and self.tokenizer is not None and self.thinker is not None
|
||||
@@ -196,19 +199,15 @@ class MiniCPMO45Stage0DuplexRuntime:
|
||||
embed_parts.extend(state.context_embeds)
|
||||
token_ids.extend(state.context_token_ids)
|
||||
|
||||
# Consume every complete chunk in the buffer so the appended span and
|
||||
# the scheduler's slot reservation for this append agree exactly. A
|
||||
# final append may zero-pad a real residual chunk, but it must not add
|
||||
# a whole silence unit after all input was already consumed: official
|
||||
# duplex generation runs once per microphone unit and client commit is
|
||||
# not an additional model decision.
|
||||
# Consume every complete processor chunk in the buffer so the appended
|
||||
# span and the scheduler's slot reservation agree exactly. The serving
|
||||
# PCM buffer pads a final real residual before it reaches Stage0. Do not
|
||||
# pad again here: the first processor chunk is hop-aligned below 1035ms
|
||||
# and intentionally leaves a small carry that is not another model unit.
|
||||
units_built = 0
|
||||
while True:
|
||||
if len(state.audio_buffer) < chunk_size:
|
||||
if not final or len(state.audio_buffer) == 0:
|
||||
break
|
||||
pad = np.zeros(chunk_size - len(state.audio_buffer), dtype=np.float32)
|
||||
state.audio_buffer = np.concatenate([state.audio_buffer, pad])
|
||||
break
|
||||
audio_chunk = state.audio_buffer[:chunk_size]
|
||||
batch_feature = self._process_streaming_audio(audio_chunk, state.audio_chunk_idx, processor=processor)
|
||||
for name, value in (
|
||||
@@ -224,6 +223,11 @@ class MiniCPMO45Stage0DuplexRuntime:
|
||||
if units_built == 0:
|
||||
return self._stage_prefill_result(False, start_time, "streaming audio embedding returned empty")
|
||||
break
|
||||
consumed_samples = self._consumed_audio_samples(
|
||||
state.audio_chunk_idx,
|
||||
chunk_size,
|
||||
processor=processor,
|
||||
)
|
||||
if state.audio_chunk_idx > 0:
|
||||
# Official duplex closes every unit (finalize_unit feeds the
|
||||
# sampled terminator + </unit>) before the next <unit> opens.
|
||||
@@ -255,9 +259,7 @@ class MiniCPMO45Stage0DuplexRuntime:
|
||||
token_ids.extend(
|
||||
[self._audio_embedding_placeholder_token_id()] * int(self._as_2d_tensor(audio_embeds).shape[0])
|
||||
)
|
||||
state.audio_buffer = state.audio_buffer[
|
||||
self._consumed_audio_samples(state.audio_chunk_idx, chunk_size, processor=processor) :
|
||||
]
|
||||
state.audio_buffer = state.audio_buffer[consumed_samples:]
|
||||
state.audio_chunk_idx += 1
|
||||
units_built += 1
|
||||
chunk_size = self._streaming_chunk_size(processor)
|
||||
@@ -610,14 +612,33 @@ class MiniCPMO45Stage0DuplexRuntime:
|
||||
def _load_processor_from_path(model_path: str | None) -> Any | None:
|
||||
if not model_path:
|
||||
return None
|
||||
if model_path not in sys.path:
|
||||
sys.path.insert(0, model_path)
|
||||
try:
|
||||
from processing_minicpmo import MiniCPMOProcessor
|
||||
from transformers import AutoImageProcessor, AutoProcessor
|
||||
|
||||
return MiniCPMOProcessor.from_pretrained(model_path, trust_remote_code=True)
|
||||
except Exception:
|
||||
return None
|
||||
original_register = AutoImageProcessor.register
|
||||
|
||||
def register_image_processor(config_class, *args, **kwargs):
|
||||
# The checkpoint's auto_map already loads this class. Its
|
||||
# legacy string registration is incompatible with some
|
||||
# Transformers versions and is otherwise redundant.
|
||||
if config_class == "MiniCPMVImageProcessor":
|
||||
return None
|
||||
return original_register(config_class, *args, **kwargs)
|
||||
|
||||
with _MINICPMO45_PROCESSOR_LOAD_LOCK:
|
||||
AutoImageProcessor.register = staticmethod(register_image_processor)
|
||||
try:
|
||||
processor = AutoProcessor.from_pretrained(
|
||||
model_path,
|
||||
trust_remote_code=True,
|
||||
)
|
||||
finally:
|
||||
AutoImageProcessor.register = staticmethod(original_register)
|
||||
except Exception as exc:
|
||||
raise RuntimeError(f"Failed to load MiniCPM-o duplex processor from {model_path!r}") from exc
|
||||
if getattr(processor, "tokenizer", None) is None:
|
||||
raise RuntimeError(f"MiniCPM-o duplex processor loaded from {model_path!r} does not expose a tokenizer")
|
||||
return processor
|
||||
|
||||
def _init_token_ids(self) -> None:
|
||||
if self.tokenizer is None:
|
||||
|
||||
@@ -503,11 +503,45 @@ class DuplexSessionRunnerMixin:
|
||||
clear_completed_pending_silence()
|
||||
pending_silence = native.pending_silence_task
|
||||
if pending_silence is not None and not pending_silence.done():
|
||||
if pending_silence is asyncio.current_task():
|
||||
return False
|
||||
try:
|
||||
if not await pending_silence:
|
||||
return False
|
||||
except asyncio.CancelledError:
|
||||
current = asyncio.current_task()
|
||||
if current is not None and current.cancelling():
|
||||
raise
|
||||
return False
|
||||
except Exception:
|
||||
return False
|
||||
clear_completed_pending_silence()
|
||||
pending_silence = native.pending_silence_task
|
||||
if pending_silence is not None and not pending_silence.done():
|
||||
return False
|
||||
append_tail = actor.native_append_tail
|
||||
if (append_tail is None or append_tail.done()) and real_native_input_waiting():
|
||||
return False
|
||||
if actor.native_append_tail is not None and not actor.native_append_tail.done():
|
||||
return False
|
||||
if real_native_input_waiting():
|
||||
return False
|
||||
continuation_delay_s = max(
|
||||
0.0,
|
||||
float(session.capabilities.chunk_period_ms or 1000) / 1000.0,
|
||||
)
|
||||
if continuation_delay_s > 0:
|
||||
await asyncio.sleep(continuation_delay_s)
|
||||
if (
|
||||
actor.native_append_tail is not append_tail
|
||||
or ((append_tail is None or append_tail.done()) and real_native_input_waiting())
|
||||
or self._native_silence_continuation_is_stale(
|
||||
session,
|
||||
request_id=request_id,
|
||||
response_id=response_id,
|
||||
response_owned=response_owned,
|
||||
expected_epoch=expected_epoch,
|
||||
expected_incarnation=expected_incarnation,
|
||||
expected_model_turn_id=expected_model_turn_id,
|
||||
)
|
||||
):
|
||||
return False
|
||||
|
||||
def _still_valid() -> bool:
|
||||
return not real_native_input_waiting() and not self._native_silence_continuation_is_stale(
|
||||
|
||||
@@ -348,9 +348,16 @@ class BatchedToken2Wav(nn.Module):
|
||||
previous: torch.Tensor,
|
||||
window: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
overlap = int(window.shape[0] // 2)
|
||||
overlap = min(
|
||||
int(window.shape[0] // 2),
|
||||
int(speech.shape[-1]),
|
||||
int(previous.shape[-1]),
|
||||
)
|
||||
result = speech.clone()
|
||||
result[..., :overlap] = result[..., :overlap] * window[:overlap] + previous[..., -overlap:] * window[overlap:]
|
||||
if overlap > 0:
|
||||
result[..., :overlap] = (
|
||||
result[..., :overlap] * window[:overlap] + previous[..., -overlap:] * window[-overlap:]
|
||||
)
|
||||
return result
|
||||
|
||||
def decode_batch(
|
||||
@@ -360,6 +367,7 @@ class BatchedToken2Wav(nn.Module):
|
||||
states: list[BatchedToken2WavState],
|
||||
*,
|
||||
last_chunk: bool,
|
||||
flush_encoder: bool = False,
|
||||
) -> tuple[list[torch.Tensor], list[BatchedToken2WavState]]:
|
||||
batch_size = int(tokens.shape[0])
|
||||
if batch_size != len(states):
|
||||
@@ -382,7 +390,7 @@ class BatchedToken2Wav(nn.Module):
|
||||
with self._autocast(tokens.device):
|
||||
hidden, conformer_cnn, conformer_att = self._encode_chunk(
|
||||
tokens,
|
||||
last_chunk=last_chunk,
|
||||
last_chunk=last_chunk or flush_encoder,
|
||||
cnn_cache=flow_cache["conformer_cnn_cache"],
|
||||
att_cache=flow_cache["conformer_att_cache"],
|
||||
)
|
||||
|
||||
@@ -4,12 +4,12 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import tempfile
|
||||
from collections.abc import Iterable, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from hashlib import sha256
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
@@ -77,6 +77,13 @@ class _RequestState:
|
||||
token2wav: BatchedToken2WavState
|
||||
|
||||
|
||||
@dataclass
|
||||
class _RuntimePrompt:
|
||||
cache_id: str
|
||||
path: str
|
||||
owners: set[str]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _WorkItem:
|
||||
output_index: int
|
||||
@@ -89,9 +96,11 @@ class _WorkItem:
|
||||
last_chunk: bool
|
||||
tokens: torch.Tensor
|
||||
previous: _RequestState | None
|
||||
segment_text: str
|
||||
duplex_turn_id: int | None
|
||||
duplex_epoch: int | None
|
||||
runtime_prompt_key: str | None
|
||||
duplex_epoch: int
|
||||
duplex_turn_id: int
|
||||
segment_text_utf8: torch.Tensor
|
||||
tts_is_last_chunk: bool
|
||||
segment_end: bool
|
||||
turn_end: bool
|
||||
has_payload: bool = True
|
||||
@@ -120,12 +129,17 @@ class MiniCPMO45Code2Wav(nn.Module):
|
||||
self.model_path = str(vllm_config.model_config.model)
|
||||
self.backend: BatchedToken2Wav | None = None
|
||||
self._states: dict[str, _RequestState] = {}
|
||||
self._owned_prompt_wavs: dict[str, tuple[str, str]] = {}
|
||||
self._runtime_prompts: dict[str, _RuntimePrompt] = {}
|
||||
self._request_prompt_keys: dict[str, str] = {}
|
||||
self._runtime_prompt_dir = tempfile.TemporaryDirectory(
|
||||
prefix="minicpmo45-runtime-prompts-",
|
||||
)
|
||||
extra = self._extra_config()
|
||||
self._min_batch_size = int(extra.get("code2wav_min_batch_size", 1))
|
||||
if self._min_batch_size < 1:
|
||||
raise ValueError("MiniCPM-o Code2Wav code2wav_min_batch_size must be >= 1")
|
||||
self._default_prompt_id = str(extra.get("prompt_cache_id", "HT_ref_audio"))
|
||||
self._prompt_wav_explicit = "prompt_wav" in extra
|
||||
self._default_prompt_wav = str(
|
||||
extra.get(
|
||||
"prompt_wav",
|
||||
@@ -133,6 +147,34 @@ class MiniCPMO45Code2Wav(nn.Module):
|
||||
)
|
||||
)
|
||||
|
||||
def _resolve_model_root(self) -> Path:
|
||||
model_root = Path(self.model_path)
|
||||
if model_root.is_dir():
|
||||
return model_root
|
||||
|
||||
from vllm_omni.model_executor.model_loader.weight_utils import (
|
||||
download_weights_from_hf_specific,
|
||||
)
|
||||
|
||||
model_config = self.vllm_config.model_config
|
||||
load_config = getattr(self.vllm_config, "load_config", None)
|
||||
model_root = Path(
|
||||
download_weights_from_hf_specific(
|
||||
self.model_path,
|
||||
getattr(load_config, "download_dir", None),
|
||||
allow_patterns=[
|
||||
"assets/HT_ref_audio.wav",
|
||||
"assets/token2wav/*",
|
||||
],
|
||||
revision=getattr(model_config, "revision", None),
|
||||
require_all=True,
|
||||
)
|
||||
)
|
||||
self.model_path = str(model_root)
|
||||
if not self._prompt_wav_explicit:
|
||||
self._default_prompt_wav = str(model_root / "assets" / "HT_ref_audio.wav")
|
||||
return model_root
|
||||
|
||||
def _extra_config(self) -> dict[str, Any]:
|
||||
model_config = getattr(self.vllm_config, "model_config", None)
|
||||
connector = getattr(model_config, "stage_connector_config", None)
|
||||
@@ -142,63 +184,122 @@ class MiniCPMO45Code2Wav(nn.Module):
|
||||
extra = getattr(connector, "extra", None)
|
||||
return dict(extra) if isinstance(extra, Mapping) else {}
|
||||
|
||||
def _release_prompt(self, state_id: str) -> None:
|
||||
owned = self._owned_prompt_wavs.pop(state_id, None)
|
||||
if owned is None:
|
||||
return
|
||||
prompt_cache_id, prompt_wav = owned
|
||||
if self.backend is not None:
|
||||
self.backend.evict_prompt(prompt_cache_id, prompt_wav)
|
||||
try:
|
||||
os.unlink(prompt_wav)
|
||||
except FileNotFoundError:
|
||||
pass
|
||||
|
||||
def _write_reference_prompt(
|
||||
self,
|
||||
state_id: str,
|
||||
reference: object,
|
||||
sample_rate: object,
|
||||
) -> tuple[str, str]:
|
||||
waveform = torch.as_tensor(reference, dtype=torch.float32).reshape(-1).cpu()
|
||||
if waveform.numel() == 0:
|
||||
return self._default_prompt_id, self._default_prompt_wav
|
||||
sample_rate_hz = int(_scalar(sample_rate, 24000))
|
||||
if sample_rate_hz <= 0:
|
||||
raise _batch_error(
|
||||
"invalid_ref_audio_sample_rate",
|
||||
request_id=state_id,
|
||||
sample_rate=sample_rate_hz,
|
||||
)
|
||||
digest = hashlib.sha256()
|
||||
digest.update(str(sample_rate_hz).encode("ascii"))
|
||||
digest.update(waveform.numpy().tobytes())
|
||||
prompt_cache_id = f"{state_id}:{digest.hexdigest()}"
|
||||
current = self._owned_prompt_wavs.get(state_id)
|
||||
if current is not None and current[0] == prompt_cache_id and os.path.exists(current[1]):
|
||||
return current
|
||||
|
||||
self._release_prompt(state_id)
|
||||
handle = tempfile.NamedTemporaryFile(prefix="minicpmo45-ref-", suffix=".wav", delete=False)
|
||||
prompt_wav = handle.name
|
||||
handle.close()
|
||||
try:
|
||||
sf.write(prompt_wav, waveform.numpy(), sample_rate_hz, format="WAV")
|
||||
except Exception:
|
||||
try:
|
||||
os.unlink(prompt_wav)
|
||||
except FileNotFoundError:
|
||||
pass
|
||||
raise
|
||||
self._owned_prompt_wavs[state_id] = (prompt_cache_id, prompt_wav)
|
||||
return prompt_cache_id, prompt_wav
|
||||
|
||||
def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor:
|
||||
return torch.zeros((input_ids.numel(), 1), device=input_ids.device, dtype=torch.float32)
|
||||
|
||||
def compute_logits(self, hidden_states: Any, sampling_metadata: Any = None) -> None:
|
||||
return None
|
||||
|
||||
def _materialize_runtime_prompt(
|
||||
self,
|
||||
ref_audio: Any,
|
||||
sample_rate: Any,
|
||||
) -> tuple[str, _RuntimePrompt]:
|
||||
sample_rate_hz = int(_scalar(sample_rate, 0))
|
||||
waveform = torch.as_tensor(ref_audio, dtype=torch.float32).reshape(-1).cpu().contiguous()
|
||||
if sample_rate_hz <= 0:
|
||||
raise _batch_error("invalid_ref_audio_sample_rate", sample_rate=sample_rate_hz)
|
||||
if waveform.numel() == 0:
|
||||
raise _batch_error("empty_ref_audio")
|
||||
if not bool(torch.isfinite(waveform).all().item()):
|
||||
raise _batch_error("non_finite_ref_audio")
|
||||
|
||||
digest = sha256()
|
||||
digest.update(waveform.numpy().tobytes())
|
||||
digest.update(str(sample_rate_hz).encode())
|
||||
cache_key = digest.hexdigest()
|
||||
cache_id = f"runtime-ref-{cache_key[:24]}-{sample_rate_hz}"
|
||||
path = str(Path(self._runtime_prompt_dir.name) / f"minicpmo45_ref_{cache_key[:24]}_{sample_rate_hz}.wav")
|
||||
entry = self._runtime_prompts.get(cache_key)
|
||||
if entry is None:
|
||||
entry = _RuntimePrompt(cache_id=cache_id, path=path, owners=set())
|
||||
self._runtime_prompts[cache_key] = entry
|
||||
prompt_path = Path(entry.path)
|
||||
if not prompt_path.is_file():
|
||||
with tempfile.NamedTemporaryFile(
|
||||
dir=prompt_path.parent,
|
||||
prefix=f".{prompt_path.stem}-",
|
||||
suffix=".wav",
|
||||
delete=False,
|
||||
) as handle:
|
||||
temporary_path = Path(handle.name)
|
||||
try:
|
||||
sf.write(
|
||||
temporary_path,
|
||||
waveform.numpy(),
|
||||
sample_rate_hz,
|
||||
format="WAV",
|
||||
)
|
||||
os.replace(temporary_path, prompt_path)
|
||||
finally:
|
||||
temporary_path.unlink(missing_ok=True)
|
||||
return cache_key, entry
|
||||
|
||||
def _resolve_prompt(
|
||||
self,
|
||||
state_id: str,
|
||||
info: Mapping[str, Any],
|
||||
meta: Mapping[str, Any],
|
||||
previous: _RequestState | None,
|
||||
) -> tuple[str, str, str | None]:
|
||||
codes = info.get("codes")
|
||||
ref_audio = codes.get("ref") if isinstance(codes, Mapping) else None
|
||||
if ref_audio is not None:
|
||||
cache_key, entry = self._materialize_runtime_prompt(
|
||||
ref_audio,
|
||||
meta.get("ref_audio_sr"),
|
||||
)
|
||||
return entry.cache_id, entry.path, cache_key
|
||||
|
||||
if previous is not None:
|
||||
return previous.prompt_cache_id, previous.prompt_wav, self._request_prompt_keys.get(state_id)
|
||||
|
||||
cache_key = self._request_prompt_keys.get(state_id)
|
||||
entry = self._runtime_prompts.get(cache_key) if cache_key is not None else None
|
||||
if entry is not None:
|
||||
return entry.cache_id, entry.path, cache_key
|
||||
|
||||
return (
|
||||
str(_scalar(meta.get("prompt_cache_id"), self._default_prompt_id)),
|
||||
str(_scalar(meta.get("prompt_wav"), self._default_prompt_wav)),
|
||||
None,
|
||||
)
|
||||
|
||||
def _release_request_prompt(self, state_id: str) -> None:
|
||||
cache_key = self._request_prompt_keys.pop(state_id, None)
|
||||
entry = self._runtime_prompts.get(cache_key) if cache_key is not None else None
|
||||
if entry is None:
|
||||
return
|
||||
entry.owners.discard(state_id)
|
||||
if entry.owners:
|
||||
return
|
||||
if self.backend is not None:
|
||||
self.backend.evict_prompt(entry.cache_id, entry.path)
|
||||
Path(entry.path).unlink(missing_ok=True)
|
||||
self._runtime_prompts.pop(cache_key, None)
|
||||
|
||||
def _commit_runtime_prompt_owners(self, items: list[_WorkItem]) -> None:
|
||||
for item in items:
|
||||
cache_key = item.runtime_prompt_key
|
||||
if cache_key is None:
|
||||
continue
|
||||
previous_key = self._request_prompt_keys.get(item.state_id)
|
||||
if previous_key != cache_key:
|
||||
self._release_request_prompt(item.state_id)
|
||||
entry = self._runtime_prompts.get(cache_key)
|
||||
if entry is not None:
|
||||
entry.owners.add(item.state_id)
|
||||
self._request_prompt_keys[item.state_id] = cache_key
|
||||
|
||||
def _prune_unowned_runtime_prompts(self) -> None:
|
||||
for cache_key, entry in list(self._runtime_prompts.items()):
|
||||
if entry.owners:
|
||||
continue
|
||||
if self.backend is not None:
|
||||
self.backend.evict_prompt(entry.cache_id, entry.path)
|
||||
Path(entry.path).unlink(missing_ok=True)
|
||||
self._runtime_prompts.pop(cache_key, None)
|
||||
|
||||
@staticmethod
|
||||
def _split_segments(input_ids: torch.Tensor, counts: Any) -> list[torch.Tensor]:
|
||||
flat = input_ids.reshape(-1)
|
||||
@@ -247,9 +348,11 @@ class MiniCPMO45Code2Wav(nn.Module):
|
||||
last_chunk=False,
|
||||
tokens=segment.new_empty(0, dtype=torch.long),
|
||||
previous=None,
|
||||
segment_text="",
|
||||
duplex_turn_id=None,
|
||||
duplex_epoch=None,
|
||||
runtime_prompt_key=None,
|
||||
duplex_epoch=-1,
|
||||
duplex_turn_id=-1,
|
||||
segment_text_utf8=torch.empty(0, dtype=torch.uint8),
|
||||
tts_is_last_chunk=False,
|
||||
segment_end=False,
|
||||
turn_end=False,
|
||||
has_payload=False,
|
||||
@@ -266,14 +369,15 @@ class MiniCPMO45Code2Wav(nn.Module):
|
||||
chunk_seq=chunk_seq,
|
||||
)
|
||||
last_chunk = bool(_scalar(meta.get("last_chunk"), False))
|
||||
tts_is_last_chunk = bool(_scalar(meta.get("tts_is_last_chunk"), False))
|
||||
codes = info.get("codes")
|
||||
audio = codes.get("audio") if isinstance(codes, Mapping) else None
|
||||
reference = codes.get("ref") if isinstance(codes, Mapping) else None
|
||||
tokens = _codec_tensor(audio, segment)
|
||||
if last_chunk and int(_scalar(meta.get("code_flat_numel"), tokens.numel())) == 0:
|
||||
if int(_scalar(meta.get("code_flat_numel"), tokens.numel())) == 0:
|
||||
# The generation scheduler reserves one placeholder token for an
|
||||
# empty terminal chunk. The producer's explicit length is the
|
||||
# authority, so do not decode that placeholder as codec data.
|
||||
# empty terminal or segment-boundary chunk. The producer's
|
||||
# explicit length is the authority, so do not decode that
|
||||
# placeholder as codec data.
|
||||
tokens = segment.new_empty(0, dtype=torch.long)
|
||||
previous = self._states.get(state_id)
|
||||
if previous is None:
|
||||
@@ -299,7 +403,6 @@ class MiniCPMO45Code2Wav(nn.Module):
|
||||
cache_epoch=cache_epoch,
|
||||
chunk_seq=chunk_seq,
|
||||
)
|
||||
self._release_prompt(state_id)
|
||||
previous = None
|
||||
elif chunk_seq != previous.chunk_seq + 1:
|
||||
raise _batch_error(
|
||||
@@ -308,22 +411,12 @@ class MiniCPMO45Code2Wav(nn.Module):
|
||||
expected=previous.chunk_seq + 1,
|
||||
actual=chunk_seq,
|
||||
)
|
||||
|
||||
explicit_prompt_id = _scalar(meta.get("prompt_cache_id"))
|
||||
explicit_prompt_wav = _scalar(meta.get("prompt_wav"))
|
||||
if previous is not None and explicit_prompt_id is None and explicit_prompt_wav is None:
|
||||
prompt_cache_id = previous.prompt_cache_id
|
||||
prompt_wav = previous.prompt_wav
|
||||
elif reference is not None:
|
||||
prompt_cache_id, prompt_wav = self._write_reference_prompt(
|
||||
state_id,
|
||||
reference,
|
||||
meta.get("ref_audio_sr"),
|
||||
)
|
||||
else:
|
||||
prompt_cache_id = str(explicit_prompt_id or self._default_prompt_id)
|
||||
prompt_wav = str(explicit_prompt_wav or self._default_prompt_wav)
|
||||
|
||||
prompt_cache_id, prompt_wav, runtime_prompt_key = self._resolve_prompt(
|
||||
state_id,
|
||||
info,
|
||||
meta,
|
||||
previous,
|
||||
)
|
||||
if previous is not None and prompt_cache_id != previous.prompt_cache_id:
|
||||
raise _batch_error(
|
||||
"prompt_changed_midstream",
|
||||
@@ -338,8 +431,9 @@ class MiniCPMO45Code2Wav(nn.Module):
|
||||
expected=previous.prompt_wav,
|
||||
actual=prompt_wav,
|
||||
)
|
||||
duplex_turn_id = _scalar(meta.get("duplex_turn_id"))
|
||||
duplex_epoch = _scalar(meta.get("duplex_epoch"))
|
||||
segment_text_utf8 = meta.get("llm_output_text_utf8")
|
||||
if not isinstance(segment_text_utf8, torch.Tensor):
|
||||
segment_text_utf8 = torch.empty(0, dtype=torch.uint8)
|
||||
return _WorkItem(
|
||||
output_index=index,
|
||||
state_id=state_id,
|
||||
@@ -351,9 +445,11 @@ class MiniCPMO45Code2Wav(nn.Module):
|
||||
last_chunk=last_chunk,
|
||||
tokens=tokens,
|
||||
previous=previous,
|
||||
segment_text=str(_scalar(meta.get("native_duplex_segment_text"), "")),
|
||||
duplex_turn_id=int(duplex_turn_id) if duplex_turn_id is not None else None,
|
||||
duplex_epoch=int(duplex_epoch) if duplex_epoch is not None else None,
|
||||
runtime_prompt_key=runtime_prompt_key,
|
||||
duplex_epoch=int(_scalar(meta.get("duplex_epoch"), -1)),
|
||||
duplex_turn_id=int(_scalar(meta.get("duplex_turn_id"), -1)),
|
||||
segment_text_utf8=segment_text_utf8,
|
||||
tts_is_last_chunk=tts_is_last_chunk,
|
||||
segment_end=bool(_scalar(meta.get("segment_end"), False)),
|
||||
turn_end=bool(_scalar(meta.get("turn_end"), False)),
|
||||
)
|
||||
@@ -371,6 +467,7 @@ class MiniCPMO45Code2Wav(nn.Module):
|
||||
int(item.tokens.numel()),
|
||||
cache_signature,
|
||||
item.last_chunk,
|
||||
item.tts_is_last_chunk,
|
||||
item.cache_epoch,
|
||||
)
|
||||
|
||||
@@ -435,26 +532,37 @@ class MiniCPMO45Code2Wav(nn.Module):
|
||||
request_ids=len(state_ids),
|
||||
)
|
||||
items: list[_WorkItem] = []
|
||||
for index, (state_id, segment, info) in enumerate(
|
||||
zip(state_ids, segments, runtime_additional_information, strict=True)
|
||||
):
|
||||
if not isinstance(info, Mapping):
|
||||
raise _batch_error(
|
||||
"invalid_runtime_info",
|
||||
output_index=index,
|
||||
value_type=type(info).__name__,
|
||||
)
|
||||
items.append(self._parse_item(index, str(state_id), segment, info))
|
||||
try:
|
||||
for index, (state_id, segment, info) in enumerate(
|
||||
zip(state_ids, segments, runtime_additional_information, strict=True)
|
||||
):
|
||||
if not isinstance(info, Mapping):
|
||||
raise _batch_error(
|
||||
"invalid_runtime_info",
|
||||
output_index=index,
|
||||
value_type=type(info).__name__,
|
||||
)
|
||||
items.append(self._parse_item(index, str(state_id), segment, info))
|
||||
except Exception:
|
||||
self._prune_unowned_runtime_prompts()
|
||||
raise
|
||||
state_ids = [item.state_id for item in items]
|
||||
if len(state_ids) != len(set(state_ids)):
|
||||
self._prune_unowned_runtime_prompts()
|
||||
raise _batch_error("duplicate_request_in_forward", request_ids=state_ids)
|
||||
outputs = [empty for _ in segments]
|
||||
sentinels = [item for item in items if item.last_chunk and item.tokens.numel() == 0]
|
||||
segment_markers = [
|
||||
item for item in items if not item.last_chunk and item.tts_is_last_chunk and item.tokens.numel() == 0
|
||||
]
|
||||
compute_items = [item for item in items if item.tokens.numel() > 0]
|
||||
invalid_empty = [
|
||||
item.request_id for item in items if item.has_payload and not item.last_chunk and item.tokens.numel() == 0
|
||||
item.request_id
|
||||
for item in items
|
||||
if item.has_payload and not item.last_chunk and not item.tts_is_last_chunk and item.tokens.numel() == 0
|
||||
]
|
||||
if invalid_empty:
|
||||
self._prune_unowned_runtime_prompts()
|
||||
raise _batch_error("empty_nonfinal_chunk", request_ids=invalid_empty)
|
||||
|
||||
buckets: dict[tuple[Any, ...], list[_WorkItem]] = {}
|
||||
@@ -470,6 +578,7 @@ class MiniCPMO45Code2Wav(nn.Module):
|
||||
if len(bucket) < self._min_batch_size
|
||||
]
|
||||
if undersized:
|
||||
self._prune_unowned_runtime_prompts()
|
||||
raise _batch_error(
|
||||
"exact_shape_bucket_below_minimum",
|
||||
minimum=self._min_batch_size,
|
||||
@@ -477,6 +586,58 @@ class MiniCPMO45Code2Wav(nn.Module):
|
||||
)
|
||||
|
||||
pending: dict[str, _RequestState | None] = {item.state_id: None for item in sentinels}
|
||||
pending.update(
|
||||
{
|
||||
item.state_id: _RequestState(
|
||||
cache_epoch=item.cache_epoch,
|
||||
chunk_seq=item.chunk_seq,
|
||||
prompt_cache_id=item.prompt_cache_id,
|
||||
prompt_wav=item.prompt_wav,
|
||||
token2wav=item.previous.token2wav,
|
||||
)
|
||||
for item in segment_markers
|
||||
if item.previous is not None
|
||||
}
|
||||
)
|
||||
initial_marker_buckets: dict[tuple[str, str], list[_WorkItem]] = {}
|
||||
for item in segment_markers:
|
||||
if item.previous is None:
|
||||
initial_marker_buckets.setdefault(
|
||||
(item.prompt_cache_id, item.prompt_wav),
|
||||
[],
|
||||
).append(item)
|
||||
for bucket in initial_marker_buckets.values():
|
||||
try:
|
||||
features = self.backend.prepare_prompt(
|
||||
bucket[0].prompt_cache_id,
|
||||
bucket[0].prompt_wav,
|
||||
)
|
||||
states = self.backend.setup_batch(features, len(bucket))
|
||||
except Exception as exc:
|
||||
self._prune_unowned_runtime_prompts()
|
||||
if isinstance(exc, RuntimeError) and str(exc).startswith("MiniCPMO45Code2WavBatchError "):
|
||||
raise
|
||||
raise _batch_error(
|
||||
"backend_unsupported_or_failed",
|
||||
request_ids=[item.request_id for item in bucket],
|
||||
error_type=type(exc).__name__,
|
||||
error=str(exc),
|
||||
) from exc
|
||||
if len(states) != len(bucket):
|
||||
self._prune_unowned_runtime_prompts()
|
||||
raise _batch_error(
|
||||
"backend_result_size_mismatch",
|
||||
expected=len(bucket),
|
||||
states=len(states),
|
||||
)
|
||||
for item, state in zip(bucket, states, strict=True):
|
||||
pending[item.state_id] = _RequestState(
|
||||
cache_epoch=item.cache_epoch,
|
||||
chunk_seq=item.chunk_seq,
|
||||
prompt_cache_id=item.prompt_cache_id,
|
||||
prompt_wav=item.prompt_wav,
|
||||
token2wav=state,
|
||||
)
|
||||
for bucket in buckets.values():
|
||||
batch_size = len(bucket)
|
||||
try:
|
||||
@@ -496,6 +657,7 @@ class MiniCPMO45Code2Wav(nn.Module):
|
||||
last_chunk=bucket[0].last_chunk,
|
||||
)
|
||||
except Exception as exc:
|
||||
self._prune_unowned_runtime_prompts()
|
||||
if isinstance(exc, RuntimeError) and str(exc).startswith("MiniCPMO45Code2WavBatchError "):
|
||||
raise
|
||||
raise _batch_error(
|
||||
@@ -505,6 +667,7 @@ class MiniCPMO45Code2Wav(nn.Module):
|
||||
error=str(exc),
|
||||
) from exc
|
||||
if len(audios) != batch_size or len(next_states) != batch_size:
|
||||
self._prune_unowned_runtime_prompts()
|
||||
raise _batch_error(
|
||||
"backend_result_size_mismatch",
|
||||
expected=batch_size,
|
||||
@@ -525,10 +688,10 @@ class MiniCPMO45Code2Wav(nn.Module):
|
||||
)
|
||||
)
|
||||
|
||||
self._commit_runtime_prompt_owners(items)
|
||||
for request_id, state in pending.items():
|
||||
if state is None:
|
||||
self._states.pop(request_id, None)
|
||||
self._release_prompt(request_id)
|
||||
else:
|
||||
self._states[request_id] = state
|
||||
sample_rate_tensor = torch.as_tensor(sample_rate, dtype=torch.int32)
|
||||
@@ -537,24 +700,15 @@ class MiniCPMO45Code2Wav(nn.Module):
|
||||
multimodal_outputs={
|
||||
"model_outputs": outputs,
|
||||
"sr": [sample_rate_tensor.clone() for _ in outputs],
|
||||
"meta.llm_output_text_utf8": [
|
||||
torch.tensor(list(item.segment_text.encode("utf-8")), dtype=torch.uint8) for item in items
|
||||
],
|
||||
"meta.tts_is_last_chunk": [torch.tensor(item.last_chunk, dtype=torch.bool) for item in items],
|
||||
# Generation runner wire payloads are flat and tensor-only.
|
||||
# Dotted metadata keys are unflattened again by the output
|
||||
# processor before the full-duplex data plane consumes them.
|
||||
"meta.duplex_epoch": [torch.tensor(item.duplex_epoch, dtype=torch.int32) for item in items],
|
||||
"meta.duplex_turn_id": [torch.tensor(item.duplex_turn_id, dtype=torch.int32) for item in items],
|
||||
"meta.llm_output_text_utf8": [item.segment_text_utf8 for item in items],
|
||||
"meta.tts_is_last_chunk": [torch.tensor(item.tts_is_last_chunk, dtype=torch.bool) for item in items],
|
||||
"meta.segment_end": [torch.tensor(item.segment_end, dtype=torch.bool) for item in items],
|
||||
"meta.turn_end": [torch.tensor(item.turn_end, dtype=torch.bool) for item in items],
|
||||
"meta.duplex_turn_id": [
|
||||
torch.tensor(item.duplex_turn_id, dtype=torch.int64)
|
||||
if item.duplex_turn_id is not None
|
||||
else torch.empty(0, dtype=torch.int64)
|
||||
for item in items
|
||||
],
|
||||
"meta.duplex_epoch": [
|
||||
torch.tensor(item.duplex_epoch, dtype=torch.int64)
|
||||
if item.duplex_epoch is not None
|
||||
else torch.empty(0, dtype=torch.int64)
|
||||
for item in items
|
||||
],
|
||||
},
|
||||
)
|
||||
|
||||
@@ -562,7 +716,7 @@ class MiniCPMO45Code2Wav(nn.Module):
|
||||
for request_id in finished_req_ids:
|
||||
state_id = str(request_id)
|
||||
self._states.pop(state_id, None)
|
||||
self._release_prompt(state_id)
|
||||
self._release_request_prompt(state_id)
|
||||
|
||||
def make_omni_output(self, model_outputs: Any, **_: Any) -> OmniOutput:
|
||||
if isinstance(model_outputs, OmniOutput):
|
||||
@@ -600,10 +754,11 @@ class MiniCPMO45Code2Wav(nn.Module):
|
||||
from stepaudio2.token2wav import Token2wav
|
||||
|
||||
extra = self._extra_config()
|
||||
model_root = self._resolve_model_root()
|
||||
prompt_path = Path(self._default_prompt_wav)
|
||||
if not prompt_path.is_file():
|
||||
raise FileNotFoundError(f"MiniCPM-o Code2Wav prompt audio not found: {prompt_path}")
|
||||
token2wav_path = Path(self.model_path) / "assets" / "token2wav"
|
||||
token2wav_path = model_root / "assets" / "token2wav"
|
||||
if not token2wav_path.is_dir():
|
||||
raise FileNotFoundError(f"MiniCPM-o Code2Wav assets not found: {token2wav_path}")
|
||||
use_float16 = bool(extra.get("token2wav_float16", False))
|
||||
|
||||
@@ -175,7 +175,6 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp
|
||||
|
||||
token_ids = self._minicpmo45_native_duplex_token_ids()
|
||||
listen_id = int(token_ids.get("listen_token_id", -1))
|
||||
tts_bos_id = int(token_ids.get("tts_bos_token_id", -1))
|
||||
turn_eos_id = int(token_ids.get("turn_eos_token_id", -1))
|
||||
if listen_id < 0 or listen_id >= logits.shape[-1]:
|
||||
return
|
||||
@@ -199,7 +198,6 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp
|
||||
continue
|
||||
force_listen = payload.get("force_listen") is True
|
||||
is_speech = payload.get("is_speech")
|
||||
redirect_listen = False
|
||||
segment_key = (row.request_id, row.seq if row.seq is not None else -1)
|
||||
session_key = (row.session_id, row.incarnation) if row.session_id is not None else None
|
||||
if turn_eos_id >= 0 and session_key is not None:
|
||||
@@ -209,8 +207,6 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp
|
||||
)
|
||||
if is_speech is True:
|
||||
if state is not None:
|
||||
if not getattr(state, "current_turn_ended", True):
|
||||
redirect_listen = True
|
||||
with suppress(Exception):
|
||||
state.last_terminator_token = None
|
||||
else:
|
||||
@@ -218,7 +214,7 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp
|
||||
if turn_ended and not pending_speech_context:
|
||||
force_listen = True
|
||||
|
||||
if not force_listen and not redirect_listen:
|
||||
if not force_listen:
|
||||
continue
|
||||
if force_listen and segment_key in force_listen_segments:
|
||||
continue
|
||||
@@ -226,10 +222,6 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp
|
||||
logits[row_idx, :] = float("-inf")
|
||||
logits[row_idx, listen_id] = 0.0
|
||||
force_listen_segments.add(segment_key)
|
||||
elif 0 <= tts_bos_id < logits.shape[-1]:
|
||||
if logits[row_idx, listen_id] > logits[row_idx, tts_bos_id]:
|
||||
logits[row_idx, tts_bos_id] = logits[row_idx, listen_id]
|
||||
logits[row_idx, listen_id] = float("-inf")
|
||||
|
||||
# -------------------- Device utilities --------------------
|
||||
@staticmethod
|
||||
@@ -730,6 +722,10 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp
|
||||
output_token_ids = getattr(sampling_metadata, "output_token_ids", None) or []
|
||||
raw_recent_tokens = output_token_ids[row_idx] if row_idx < len(output_token_ids) else []
|
||||
recent_tokens = [int(token_id) for token_id in raw_recent_tokens if isinstance(token_id, int) and token_id >= 0]
|
||||
temperature = float(self._sampling_metadata_value(sampling_metadata, "temperature", row_idx, 0.7))
|
||||
top_k = int(self._sampling_metadata_value(sampling_metadata, "top_k", row_idx, 100))
|
||||
top_p = float(self._sampling_metadata_value(sampling_metadata, "top_p", row_idx, 0.8))
|
||||
state = self._minicpmo45_duplex_state_for_row(row_idx)
|
||||
if chunk_eos_id >= 0 and chunk_eos_id < logits.shape[-1]:
|
||||
max_speak_tokens = int(
|
||||
getattr(
|
||||
@@ -765,23 +761,19 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp
|
||||
if valid_forbidden:
|
||||
logits[:, valid_forbidden] = float("-inf")
|
||||
|
||||
special_ids = self._minicpmo45_native_special_token_ids(token_ids)
|
||||
state = self._minicpmo45_duplex_state_for_row(row_idx)
|
||||
generated_text_tokens = getattr(state, "generated_text_tokens", None)
|
||||
repetition_tokens = generated_text_tokens if generated_text_tokens else recent_tokens
|
||||
generated_tokens = getattr(state, "generated_tokens", None)
|
||||
repetition_tokens = generated_tokens if generated_tokens else recent_tokens
|
||||
repetition_penalty = 1.05
|
||||
if repetition_penalty != 1.0 and repetition_tokens:
|
||||
history_size = MiniCPMO45DuplexPolicy.REPETITION_HISTORY_SIZE
|
||||
for token_id in set(repetition_tokens[-history_size:]):
|
||||
if token_id in special_ids or token_id < 0 or token_id >= logits.shape[-1]:
|
||||
if token_id < 0 or token_id >= logits.shape[-1]:
|
||||
continue
|
||||
logits[0, token_id] /= repetition_penalty
|
||||
|
||||
temperature = float(self._sampling_metadata_value(sampling_metadata, "temperature", row_idx, 0.7))
|
||||
top_k = int(self._sampling_metadata_value(sampling_metadata, "top_k", row_idx, 100))
|
||||
top_p = float(self._sampling_metadata_value(sampling_metadata, "top_p", row_idx, 0.8))
|
||||
if getattr(sampling_metadata, "all_greedy", False) or temperature <= 0:
|
||||
sampled = int(torch.argmax(logits, dim=-1).item())
|
||||
self._record_minicpmo45_duplex_generation_token(row_idx, sampled)
|
||||
sampled = self._maybe_cut_minicpmo45_native_duplex_text_chunk(
|
||||
sampled,
|
||||
recent_tokens,
|
||||
@@ -793,6 +785,7 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp
|
||||
logits = self._top_k_top_p_filter(logits, top_k=top_k, top_p=top_p)
|
||||
probs = F.softmax(logits, dim=-1)
|
||||
sampled = int(torch.multinomial(probs, num_samples=1, generator=generator).item())
|
||||
self._record_minicpmo45_duplex_generation_token(row_idx, sampled)
|
||||
sampled = self._maybe_cut_minicpmo45_native_duplex_text_chunk(
|
||||
sampled,
|
||||
recent_tokens,
|
||||
@@ -901,6 +894,28 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp
|
||||
return int(tts_bos_id)
|
||||
return int(sampled)
|
||||
|
||||
def _record_minicpmo45_duplex_generation_token(self, row_idx: int, sampled: int) -> None:
|
||||
"""Track tokens returned by the model-policy decoder.
|
||||
|
||||
The released StreamDecoder is constructed without special_token_ids, so
|
||||
every normally decoded token participates in repetition penalty. Forced
|
||||
listen bypasses that decoder, while chunk_eos boundary decisions return
|
||||
before this method is called.
|
||||
"""
|
||||
state = self._minicpmo45_duplex_state_for_row(row_idx)
|
||||
if state is None:
|
||||
return
|
||||
payload = self._minicpmo45_duplex_payload_for_row(row_idx)
|
||||
if isinstance(payload, dict) and payload.get("force_listen") is True:
|
||||
return
|
||||
generated_tokens = getattr(state, "generated_tokens", None)
|
||||
if not isinstance(generated_tokens, list):
|
||||
generated_tokens = []
|
||||
state.generated_tokens = generated_tokens
|
||||
generated_tokens.append(int(sampled))
|
||||
history_size = MiniCPMO45DuplexPolicy.REPETITION_HISTORY_SIZE
|
||||
del generated_tokens[:-history_size]
|
||||
|
||||
def _record_minicpmo45_duplex_terminator(self, row_idx: int, sampled: int, token_ids: dict[str, int]) -> None:
|
||||
"""Remember sampled unit state for the next append.
|
||||
|
||||
@@ -928,15 +943,6 @@ class MiniCPMO45OmniForConditionalGeneration(nn.Module, SupportsMultiModal, Supp
|
||||
with suppress(Exception):
|
||||
state.pending_speech_response_open = False
|
||||
return
|
||||
special_ids = self._minicpmo45_native_special_token_ids(token_ids)
|
||||
if sampled not in special_ids:
|
||||
generated_text_tokens = getattr(state, "generated_text_tokens", None)
|
||||
if not isinstance(generated_text_tokens, list):
|
||||
generated_text_tokens = []
|
||||
state.generated_text_tokens = generated_text_tokens
|
||||
generated_text_tokens.append(int(sampled))
|
||||
history_size = MiniCPMO45DuplexPolicy.REPETITION_HISTORY_SIZE
|
||||
del generated_text_tokens[:-history_size]
|
||||
if (
|
||||
sampled == tts_bos_id
|
||||
and getattr(state, "current_turn_ended", True)
|
||||
|
||||
@@ -43,6 +43,7 @@ _CODEC_TOP_K = 25
|
||||
_CODEC_TOP_P = 0.85
|
||||
_CODEC_REPETITION_PENALTY = 1.05
|
||||
_CODEC_MIN_TOKENS = 50
|
||||
_DUPLEX_CODEC_TOKENS_PER_CHUNK = 26
|
||||
|
||||
|
||||
def _max_audio_tokens(condition_tokens: int) -> int:
|
||||
@@ -61,12 +62,7 @@ def _max_audio_tokens(condition_tokens: int) -> int:
|
||||
|
||||
def _restore_weight_norm_weight(weight_g: torch.Tensor, weight_v: torch.Tensor) -> torch.Tensor:
|
||||
"""Materialize ``weight_norm(..., dim=0)`` checkpoint parameters."""
|
||||
row_norm = torch.linalg.vector_norm(
|
||||
weight_v.float(),
|
||||
dim=1,
|
||||
keepdim=True,
|
||||
).clamp_min(torch.finfo(torch.float32).tiny)
|
||||
return weight_v * (weight_g.to(dtype=weight_v.dtype) / row_norm.to(dtype=weight_v.dtype))
|
||||
return torch._weight_norm(weight_v, weight_g, dim=0)
|
||||
|
||||
|
||||
def _apply_repetition_penalty(
|
||||
@@ -218,6 +214,8 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP):
|
||||
self,
|
||||
tts_token_ids: torch.Tensor,
|
||||
tts_hidden_states: torch.Tensor,
|
||||
*,
|
||||
native_duplex: bool = False,
|
||||
) -> torch.Tensor:
|
||||
if tts_token_ids.numel() == 0 or tts_hidden_states.numel() == 0:
|
||||
# The thinker can legally emit an empty speech segment (<|tts_bos|>
|
||||
@@ -239,7 +237,12 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP):
|
||||
hidden_embeds = self.projector_semantic(hidden)
|
||||
if self._normalize:
|
||||
hidden_embeds = F.normalize(hidden_embeds, p=2, dim=-1)
|
||||
return torch.cat([text_embeds + hidden_embeds, self._boundary_embeddings()], dim=0)
|
||||
audio_bos = self.emb_text(torch.tensor([self._tts_bos_id], device=device, dtype=torch.long))
|
||||
condition = text_embeds + hidden_embeds
|
||||
if native_duplex:
|
||||
# Match MiniCPMTTS.generate_chunk's streaming condition.
|
||||
return torch.cat([condition, audio_bos], dim=0)
|
||||
return torch.cat([condition, self._boundary_embeddings()], dim=0)
|
||||
|
||||
def preprocess(
|
||||
self,
|
||||
@@ -281,18 +284,22 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP):
|
||||
"MiniCPM-o Talker received an empty condition (request %s); this request produces no audio.",
|
||||
info_dict.get("request_id"),
|
||||
)
|
||||
full_embeds = self._build_condition_embeddings(token_ids, hidden_states)
|
||||
native_duplex = bool(info_dict.get("native_duplex", False))
|
||||
full_embeds = self._build_condition_embeddings(
|
||||
token_ids,
|
||||
hidden_states,
|
||||
native_duplex=native_duplex,
|
||||
)
|
||||
offset = int(info_dict.get("_omni_num_computed_tokens", 0))
|
||||
if offset == 0 and span_len > full_embeds.shape[0]:
|
||||
# Async-chunk prewarms Stage 1 with placeholder token IDs. Two
|
||||
# Thinker handoffs can arrive before the first one executes,
|
||||
# so the scheduler prompt may already include the newer
|
||||
# placeholder prefix while additional_information still
|
||||
# carries the preceding condition. Materialize that prefix as
|
||||
# the embeddings of its actual zero token IDs and keep the
|
||||
# complete Talker condition tail-aligned (audio BOS remains
|
||||
# the final prefill token).
|
||||
prefix_len = span_len - full_embeds.shape[0]
|
||||
request_id = str(info_dict.get("request_id", "0"))
|
||||
meta = info_dict.get("meta")
|
||||
# The handoff rebuilds only the tail-aligned Talker condition.
|
||||
# Materialize zero-token embeddings for any scheduler prompt
|
||||
# prefix so chunked prefill can slice from a non-zero offset.
|
||||
prompt_len = info_dict.get("_omni_prompt_len")
|
||||
target_len = int(prompt_len) if prompt_len is not None else offset + span_len
|
||||
prefix_len = target_len - full_embeds.shape[0]
|
||||
if prefix_len > 0:
|
||||
placeholder_ids = torch.zeros(
|
||||
prefix_len,
|
||||
dtype=torch.long,
|
||||
@@ -308,13 +315,21 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP):
|
||||
f"tts_ids={token_ids.shape[0]} tts_hidden={hidden_states.shape[0]} "
|
||||
f"prompt_len={info_dict.get('_omni_prompt_len')}"
|
||||
)
|
||||
max_tokens = _max_audio_tokens(int(token_ids.numel()))
|
||||
duplex_boundary = isinstance(meta, dict) and (
|
||||
bool(meta.get("turn_start", False)) or bool(meta.get("turn_end", False))
|
||||
)
|
||||
if native_duplex:
|
||||
max_tokens = _DUPLEX_CODEC_TOKENS_PER_CHUNK
|
||||
min_tokens = 0 if duplex_boundary else _DUPLEX_CODEC_TOKENS_PER_CHUNK
|
||||
else:
|
||||
max_tokens = _max_audio_tokens(int(token_ids.numel()))
|
||||
min_tokens = self._codec_min_tokens
|
||||
state = {
|
||||
"step": 0,
|
||||
"max_tokens": max_tokens,
|
||||
"min_tokens": min_tokens,
|
||||
"finished": empty_condition,
|
||||
}
|
||||
request_id = str(info_dict.get("request_id", "0"))
|
||||
request_states = getattr(self, "_request_audio_states", None)
|
||||
if request_states is None:
|
||||
request_states = {}
|
||||
@@ -370,7 +385,12 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP):
|
||||
penalty=self._codec_repetition_penalty,
|
||||
window_size=_REPETITION_WINDOW,
|
||||
)
|
||||
if step < self._codec_min_tokens:
|
||||
request_states = getattr(self, "_request_audio_states", {})
|
||||
state = request_states.get(request_id)
|
||||
min_tokens = (
|
||||
int(state.get("min_tokens", self._codec_min_tokens)) if isinstance(state, dict) else self._codec_min_tokens
|
||||
)
|
||||
if step < min_tokens:
|
||||
logits[..., eos_id] = float("-inf")
|
||||
logits = _apply_top_k_top_p(
|
||||
logits,
|
||||
@@ -404,12 +424,61 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP):
|
||||
raise RuntimeError(
|
||||
f"MiniCPM-o continuous Talker received {len(sample_eligible)} sampling flags for {len(infos)} requests"
|
||||
)
|
||||
emit_duplex_metadata = any(isinstance(info, dict) and info.get("native_duplex") is True for info in infos)
|
||||
|
||||
stop_rows: list[torch.Tensor] = []
|
||||
codec_deltas: list[torch.Tensor] = []
|
||||
terminal_flags: list[torch.Tensor] = []
|
||||
native_duplex_flags: list[torch.Tensor] = []
|
||||
duplex_epochs: list[torch.Tensor] = []
|
||||
duplex_turn_ids: list[torch.Tensor] = []
|
||||
segment_texts_utf8: list[torch.Tensor] = []
|
||||
turn_end_flags: list[torch.Tensor] = []
|
||||
empty_delta = hidden.new_empty((0, 1), dtype=torch.long)
|
||||
for index, info in enumerate(infos):
|
||||
info_dict = info if isinstance(info, dict) else {}
|
||||
native_duplex = info_dict.get("native_duplex") is True
|
||||
if emit_duplex_metadata:
|
||||
duplex_info = info_dict.get("duplex")
|
||||
if not isinstance(duplex_info, dict):
|
||||
duplex_info = {}
|
||||
epoch = duplex_info.get("epoch", -1)
|
||||
turn_id = duplex_info.get("turn_id", -1)
|
||||
if native_duplex and not all(
|
||||
isinstance(value, int) and not isinstance(value, bool) and value >= 0 for value in (epoch, turn_id)
|
||||
):
|
||||
raise RuntimeError(
|
||||
"MiniCPM-o native duplex Talker requires non-negative integer "
|
||||
f"epoch and turn_id, got epoch={epoch!r}, turn_id={turn_id!r}"
|
||||
)
|
||||
meta_info = info_dict.get("meta")
|
||||
if not isinstance(meta_info, dict):
|
||||
meta_info = {}
|
||||
segment_text = meta_info.get("native_duplex_segment_text", "") if native_duplex else ""
|
||||
if not isinstance(segment_text, str):
|
||||
segment_text = ""
|
||||
turn_eos_id = meta_info.get("turn_eos_token_id")
|
||||
ids_info = info_dict.get("ids")
|
||||
tts_ids = ids_info.get("tts") if native_duplex and isinstance(ids_info, dict) else None
|
||||
if isinstance(tts_ids, torch.Tensor):
|
||||
contains_turn_eos = isinstance(turn_eos_id, int) and bool(
|
||||
torch.any(tts_ids.reshape(-1) == turn_eos_id).item()
|
||||
)
|
||||
elif isinstance(tts_ids, (list, tuple)):
|
||||
contains_turn_eos = isinstance(turn_eos_id, int) and turn_eos_id in tts_ids
|
||||
else:
|
||||
contains_turn_eos = False
|
||||
native_duplex_flags.append(torch.tensor(native_duplex, dtype=torch.bool))
|
||||
duplex_epochs.append(torch.tensor(epoch if isinstance(epoch, int) else -1, dtype=torch.long))
|
||||
duplex_turn_ids.append(torch.tensor(turn_id if isinstance(turn_id, int) else -1, dtype=torch.long))
|
||||
segment_texts_utf8.append(
|
||||
torch.tensor(
|
||||
list(segment_text.encode("utf-8")),
|
||||
dtype=torch.uint8,
|
||||
)
|
||||
)
|
||||
turn_end_flags.append(torch.tensor(native_duplex and contains_turn_eos, dtype=torch.bool))
|
||||
|
||||
if not isinstance(info, dict):
|
||||
stop_rows.append(hidden.new_tensor([0.0, float("-inf")]))
|
||||
codec_deltas.append(empty_delta)
|
||||
@@ -453,12 +522,15 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP):
|
||||
codes = codes.to(device=hidden.device, dtype=torch.long).reshape(-1)
|
||||
step = int(state.get("step", 0))
|
||||
sampled = self._sample_audio_code(hidden[end - 1 : end], codes, request_id, step)
|
||||
is_eos = int(sampled.item()) == self._num_audio_tokens - 1
|
||||
sampled_id = int(sampled.item())
|
||||
is_eos = sampled_id == self._num_audio_tokens - 1
|
||||
state["step"] = int(state.get("step", 0)) + 1
|
||||
reached_limit = int(state["step"]) >= int(state.get("max_tokens", 2048))
|
||||
finished = is_eos or reached_limit
|
||||
state["finished"] = finished
|
||||
if not is_eos:
|
||||
# MiniCPMTTS.generate_chunk consumes the boundary sample but
|
||||
# returns only codes that were fed into the retained KV state.
|
||||
if not is_eos and not reached_limit:
|
||||
codes = torch.cat([codes[-(_REPETITION_WINDOW - 1) :], sampled.reshape(1)])
|
||||
delta = sampled.reshape(1, 1)
|
||||
else:
|
||||
@@ -476,9 +548,20 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP):
|
||||
self._batch_stop_logits = torch.stack(stop_rows, dim=0) if stop_rows else hidden.new_empty((0, 2))
|
||||
# Lists are deliberate: the runner routes element i to request i,
|
||||
# preserving compaction alignment while emitting only this step's code.
|
||||
meta_outputs = {"finished": terminal_flags}
|
||||
if emit_duplex_metadata:
|
||||
meta_outputs.update(
|
||||
{
|
||||
"native_duplex": native_duplex_flags,
|
||||
"duplex_epoch": duplex_epochs,
|
||||
"duplex_turn_id": duplex_turn_ids,
|
||||
"llm_output_text_utf8": segment_texts_utf8,
|
||||
"turn_end": turn_end_flags,
|
||||
}
|
||||
)
|
||||
multimodal_outputs: dict[str, Any] = {
|
||||
"codes": {"audio": codec_deltas},
|
||||
"meta": {"finished": terminal_flags},
|
||||
"meta": meta_outputs,
|
||||
}
|
||||
return OmniOutput(
|
||||
text_hidden_states=hidden,
|
||||
|
||||
@@ -21,6 +21,7 @@ logger = logging.getLogger(__name__)
|
||||
_MINICPMO45_ASYNC_STATE = "_minicpmo45_async_codec_state"
|
||||
_MINICPMO45_STREAM_RECORD = "_minicpmo45_async_stream_record"
|
||||
_MINICPMO45_SILENCE_CODE = 4218
|
||||
_MINICPMO45_MIN_STREAM_BODY_FRAMES = 5
|
||||
|
||||
|
||||
class _MiniCPMO45MetaStruct(MetaStruct):
|
||||
@@ -28,6 +29,7 @@ class _MiniCPMO45MetaStruct(MetaStruct):
|
||||
|
||||
ref_audio_sr: int | None = None
|
||||
native_duplex_segment_text: str | None = None
|
||||
llm_output_text_utf8: torch.Tensor | None = None
|
||||
duplex_turn_id: int | None = None
|
||||
duplex_epoch: int | None = None
|
||||
segment_end: bool | None = None
|
||||
@@ -212,6 +214,16 @@ def tts2code2wav_async_chunk(
|
||||
internal_id = getattr(request, "request_id", None)
|
||||
request_id = str(external_id if external_id is not None else internal_id)
|
||||
internal_id = str(internal_id if internal_id is not None else request_id)
|
||||
output_meta = multimodal_output.get("meta") if isinstance(multimodal_output, Mapping) else None
|
||||
output_meta = output_meta if isinstance(output_meta, Mapping) else {}
|
||||
native_duplex = bool(_coerce_int(output_meta.get("native_duplex")))
|
||||
duplex_epoch = _coerce_int(output_meta.get("duplex_epoch"))
|
||||
duplex_turn_id = _coerce_int(output_meta.get("duplex_turn_id"))
|
||||
segment_text_utf8 = output_meta.get("llm_output_text_utf8")
|
||||
if not isinstance(segment_text_utf8, torch.Tensor):
|
||||
segment_text_utf8 = None
|
||||
turn_end = bool(_coerce_int(output_meta.get("turn_end")))
|
||||
duplex_turn_key = (duplex_epoch, duplex_turn_id)
|
||||
|
||||
request_payload = getattr(transfer_manager, "request_payload", None)
|
||||
if request_payload is None:
|
||||
@@ -227,7 +239,9 @@ def tts2code2wav_async_chunk(
|
||||
record = {
|
||||
"internal_id": internal_id,
|
||||
"cache_epoch": 0,
|
||||
"chunk_seq": 0,
|
||||
"retired_internal_ids": set(),
|
||||
"last_terminal_turn": None,
|
||||
}
|
||||
container[_MINICPMO45_STREAM_RECORD] = record
|
||||
elif internal_id in record["retired_internal_ids"]:
|
||||
@@ -236,6 +250,8 @@ def tts2code2wav_async_chunk(
|
||||
record["retired_internal_ids"].add(record["internal_id"])
|
||||
record["internal_id"] = internal_id
|
||||
record["cache_epoch"] = int(record["cache_epoch"]) + 1
|
||||
record["chunk_seq"] = 0
|
||||
record["last_terminal_turn"] = None
|
||||
_drop_codec_state(transfer_manager, request_id)
|
||||
|
||||
if _is_aborted(request):
|
||||
@@ -243,44 +259,46 @@ def tts2code2wav_async_chunk(
|
||||
_drop_codec_state(transfer_manager, request_id)
|
||||
return None
|
||||
|
||||
if native_duplex and turn_end and record.get("last_terminal_turn") == duplex_turn_key:
|
||||
return None
|
||||
|
||||
state = container.get(_MINICPMO45_ASYNC_STATE)
|
||||
if not isinstance(state, dict):
|
||||
request_info = getattr(request, "additional_information", None)
|
||||
if not isinstance(request_info, Mapping):
|
||||
request_info = {}
|
||||
codes_info = request_info.get("codes")
|
||||
if not isinstance(codes_info, Mapping):
|
||||
codes_info = {}
|
||||
meta_info = request_info.get("meta")
|
||||
if not isinstance(meta_info, Mapping):
|
||||
meta_info = {}
|
||||
duplex_info = request_info.get("duplex")
|
||||
if not isinstance(duplex_info, Mapping):
|
||||
duplex_info = {}
|
||||
state = {
|
||||
"internal_id": internal_id,
|
||||
"pending": [],
|
||||
"pending_text_utf8": [],
|
||||
"segment_text_recorded": False,
|
||||
"left_context": [],
|
||||
"codec_end": 0,
|
||||
"ref_audio": codes_info.get("ref"),
|
||||
"ref_audio_sr": meta_info.get("ref_audio_sr"),
|
||||
"segment_text": meta_info.get("native_duplex_segment_text"),
|
||||
"segment_end": bool(meta_info.get("segment_end", False)),
|
||||
"turn_end": bool(meta_info.get("turn_end", False)),
|
||||
"duplex_turn_id": duplex_info.get("model_turn_id", duplex_info.get("turn_id")),
|
||||
"duplex_epoch": duplex_info.get("epoch"),
|
||||
}
|
||||
container[_MINICPMO45_ASYNC_STATE] = state
|
||||
|
||||
pending = state["pending"]
|
||||
pending.extend(_extract_codec_delta(multimodal_output, request_id))
|
||||
pending_text_utf8 = state.setdefault("pending_text_utf8", [])
|
||||
current_text_utf8 = (
|
||||
segment_text_utf8.detach().to(device="cpu", dtype=torch.uint8).reshape(-1).tolist()
|
||||
if isinstance(segment_text_utf8, torch.Tensor)
|
||||
else []
|
||||
)
|
||||
if native_duplex and current_text_utf8 and not state.get("segment_text_recorded", False):
|
||||
# Talker repeats the unit text on every codec step. Queue it once for
|
||||
# the first Code2Wav payload that can carry audio for this unit.
|
||||
pending_text_utf8.extend(current_text_utf8)
|
||||
state["segment_text_recorded"] = True
|
||||
request_finished = getattr(request, "is_finished", None)
|
||||
finished = bool(is_finished or (callable(request_finished) and request_finished()))
|
||||
chunk_frames, left_context_frames = _codec_config(transfer_manager)
|
||||
if not finished and len(pending) < chunk_frames:
|
||||
flush_pending = finished
|
||||
last_chunk = bool(flush_pending and (not native_duplex or turn_end))
|
||||
if not flush_pending and len(pending) < chunk_frames:
|
||||
return None
|
||||
|
||||
new_token_count = len(pending) if finished else chunk_frames
|
||||
hold_short_unit = (
|
||||
native_duplex and flush_pending and not last_chunk and 0 < len(pending) < _MINICPMO45_MIN_STREAM_BODY_FRAMES
|
||||
)
|
||||
new_token_count = 0 if hold_short_unit else (len(pending) if flush_pending else chunk_frames)
|
||||
new_codes = pending[:new_token_count]
|
||||
del pending[:new_token_count]
|
||||
codec_start = int(state["codec_end"])
|
||||
@@ -292,53 +310,79 @@ def tts2code2wav_async_chunk(
|
||||
else:
|
||||
context = list(state["left_context"])
|
||||
output_codes = [*context, *new_codes]
|
||||
history = [*state["left_context"], *new_codes]
|
||||
history = output_codes
|
||||
state["left_context"] = history[-left_context_frames:] if left_context_frames else []
|
||||
elif finished and codec_start > 0 and state["left_context"]:
|
||||
elif last_chunk and codec_start > 0 and state["left_context"]:
|
||||
context = list(state["left_context"])
|
||||
output_codes = context
|
||||
else:
|
||||
context = []
|
||||
output_codes = []
|
||||
state["codec_end"] = codec_end
|
||||
code_flat_numel = len(output_codes)
|
||||
if native_duplex and code_flat_numel > 0:
|
||||
segment_text_utf8 = torch.tensor(pending_text_utf8, dtype=torch.uint8)
|
||||
pending_text_utf8.clear()
|
||||
if native_duplex and finished:
|
||||
state["segment_text_recorded"] = False
|
||||
if flush_pending and not last_chunk and code_flat_numel == 0:
|
||||
# Keep the generic generation connector model-agnostic: a real token
|
||||
# makes this control-only TTS boundary schedulable, while the explicit
|
||||
# zero length tells Code2Wav to discard the placeholder.
|
||||
output_codes = [0]
|
||||
|
||||
last_chunk = bool(finished and not pending)
|
||||
if last_chunk:
|
||||
if last_chunk and not native_duplex:
|
||||
record["retired_internal_ids"].add(internal_id)
|
||||
_drop_codec_state(transfer_manager, request_id)
|
||||
|
||||
chunk_seq = int(getattr(transfer_manager, "put_req_chunk", {}).get(request_id, 0))
|
||||
chunk_seq = int(record["chunk_seq"])
|
||||
record["chunk_seq"] = chunk_seq + 1
|
||||
ref_audio = None
|
||||
ref_audio_sr = None
|
||||
if int(record["cache_epoch"]) == 0 and chunk_seq == 0:
|
||||
request_info = getattr(request, "additional_information", None)
|
||||
if isinstance(request_info, Mapping):
|
||||
codes_info = request_info.get("codes")
|
||||
meta_info = request_info.get("meta")
|
||||
raw_ref_audio = codes_info.get("ref") if isinstance(codes_info, Mapping) else None
|
||||
raw_ref_audio_sr = meta_info.get("ref_audio_sr") if isinstance(meta_info, Mapping) else None
|
||||
ref_audio_sr = _coerce_int(raw_ref_audio_sr)
|
||||
if raw_ref_audio is not None:
|
||||
ref_audio = torch.as_tensor(raw_ref_audio, dtype=torch.float32).reshape(-1).cpu()
|
||||
finished_tensor = torch.tensor(last_chunk, dtype=torch.bool)
|
||||
ref_audio = state.get("ref_audio") if codec_start == 0 else None
|
||||
return OmniPayloadStruct(
|
||||
payload = OmniPayloadStruct(
|
||||
codes=CodesStruct(
|
||||
audio=torch.tensor(output_codes, dtype=torch.long),
|
||||
ref=torch.as_tensor(ref_audio, dtype=torch.float32).reshape(-1) if ref_audio is not None else None,
|
||||
ref=ref_audio,
|
||||
),
|
||||
meta=_MiniCPMO45MetaStruct(
|
||||
request_id=request_id,
|
||||
chunk_seq=chunk_seq,
|
||||
cache_epoch=int(record["cache_epoch"]),
|
||||
code_flat_numel=len(output_codes),
|
||||
code_flat_numel=code_flat_numel,
|
||||
codec_chunk_frames=new_token_count,
|
||||
codec_left_context_frames=len(context),
|
||||
left_context_size=len(context),
|
||||
last_chunk=last_chunk,
|
||||
stream_finished=finished_tensor,
|
||||
finished=finished_tensor,
|
||||
is_segment_finished=finished_tensor,
|
||||
req_id=[request_id],
|
||||
ref_audio_sr=_coerce_int(state.get("ref_audio_sr")),
|
||||
native_duplex_segment_text=(
|
||||
str(state["segment_text"]) if isinstance(state.get("segment_text"), str) else None
|
||||
),
|
||||
duplex_turn_id=_coerce_int(state.get("duplex_turn_id")),
|
||||
duplex_epoch=_coerce_int(state.get("duplex_epoch")),
|
||||
segment_end=bool(state.get("segment_end", False)),
|
||||
turn_end=bool(state.get("turn_end", False)),
|
||||
tts_is_last_chunk=last_chunk,
|
||||
duplex_epoch=duplex_epoch,
|
||||
duplex_turn_id=duplex_turn_id,
|
||||
llm_output_text_utf8=segment_text_utf8,
|
||||
tts_is_last_chunk=flush_pending,
|
||||
turn_end=turn_end and last_chunk,
|
||||
ref_audio_sr=ref_audio_sr,
|
||||
),
|
||||
request_id=request_id,
|
||||
)
|
||||
if last_chunk and native_duplex:
|
||||
record["last_terminal_turn"] = duplex_turn_key
|
||||
record["cache_epoch"] = int(record["cache_epoch"]) + 1
|
||||
record["chunk_seq"] = 0
|
||||
_drop_codec_state(transfer_manager, request_id)
|
||||
return payload
|
||||
|
||||
|
||||
def tts2code2wav_full_payload(
|
||||
@@ -496,13 +540,33 @@ def _native_tts_boundary_token_ids(special_token_ids):
|
||||
}
|
||||
|
||||
|
||||
def _decode_native_duplex_token_ids(
|
||||
token_ids: list[int],
|
||||
streaming_context,
|
||||
*,
|
||||
request_id: str,
|
||||
) -> str | None:
|
||||
decode_token_ids = getattr(streaming_context, "source_token_decoder", None)
|
||||
if not isinstance(decode_token_ids, Callable):
|
||||
return None
|
||||
decode_ids = [int(token_id) for token_id in token_ids]
|
||||
try:
|
||||
try:
|
||||
return str(decode_token_ids(decode_ids, skip_special_tokens=True))
|
||||
except TypeError:
|
||||
return str(decode_token_ids(decode_ids))
|
||||
except Exception:
|
||||
logger.exception("Failed to decode MiniCPM-o duplex token delta for request_id=%s", request_id)
|
||||
return ""
|
||||
|
||||
|
||||
def _native_duplex_segment_output_ids(
|
||||
output_ids: list[int],
|
||||
output_text: str,
|
||||
streaming_context,
|
||||
*,
|
||||
request_id: str,
|
||||
) -> tuple[list[int], str]:
|
||||
) -> tuple[list[int], str, bool]:
|
||||
"""Slice the cumulative thinker output down to the current segment.
|
||||
|
||||
Tracks how many output tokens were already handed to the talker in the
|
||||
@@ -511,7 +575,7 @@ def _native_duplex_segment_output_ids(
|
||||
"""
|
||||
bridge_states = getattr(streaming_context, "bridge_states", None)
|
||||
if not isinstance(bridge_states, dict):
|
||||
return output_ids, output_text
|
||||
return output_ids, output_text, True
|
||||
state = bridge_states.setdefault("minicpmo45_tts_handoff", {})
|
||||
duplex_state = bridge_states.get("duplex")
|
||||
turn_id = duplex_state.get("model_turn_id", duplex_state.get("turn_id")) if isinstance(duplex_state, dict) else None
|
||||
@@ -521,17 +585,12 @@ def _native_duplex_segment_output_ids(
|
||||
state["request_id"] = request_id
|
||||
state["sent_output_len"] = 0
|
||||
state["sent_output_ids"] = []
|
||||
state["acc_tts_ids"] = []
|
||||
state["acc_tts_hidden"] = []
|
||||
previous_turn_id = state.get("turn_id")
|
||||
sent_len = state.get("sent_output_len", 0)
|
||||
prev_output_ids = state.get("sent_output_ids", [])
|
||||
prev_turn_id = state.get("turn_id")
|
||||
if not isinstance(sent_len, int) or sent_len < 0 or sent_len > len(output_ids):
|
||||
# Shrunken cumulative output = epoch reset after barge-in: the talker
|
||||
# condition history is stale too.
|
||||
# Shrunken cumulative output = epoch reset after barge-in.
|
||||
sent_len = 0
|
||||
state["acc_tts_ids"] = []
|
||||
state["acc_tts_hidden"] = []
|
||||
elif sent_len and isinstance(prev_output_ids, list):
|
||||
prev_prefix = prev_output_ids[:sent_len]
|
||||
current_prefix = output_ids[:sent_len]
|
||||
@@ -540,27 +599,15 @@ def _native_duplex_segment_output_ids(
|
||||
# others keep returning cumulative ids. Detect restart from the
|
||||
# token prefix instead of clearing the cursor at turn_eos.
|
||||
sent_len = 0
|
||||
state["acc_tts_ids"] = []
|
||||
state["acc_tts_hidden"] = []
|
||||
if isinstance(prev_turn_id, int) and isinstance(turn_id, int) and prev_turn_id != turn_id:
|
||||
# The thinker output list is cumulative across clean turns, so the
|
||||
# output cursor must stay put. The talker condition is per assistant
|
||||
# turn, though; carrying it across turn_id boundaries makes stage 1
|
||||
# replay the previous turn after its consumed cursor has been reset.
|
||||
state["acc_tts_ids"] = []
|
||||
state["acc_tts_hidden"] = []
|
||||
turn_start = sent_len == 0 or previous_turn_id != turn_id
|
||||
segment_ids = output_ids[sent_len:]
|
||||
decode_token_ids = getattr(streaming_context, "source_token_decoder", None)
|
||||
if isinstance(decode_token_ids, Callable):
|
||||
decode_ids = [int(token_id) for token_id in segment_ids]
|
||||
try:
|
||||
try:
|
||||
segment_text = str(decode_token_ids(decode_ids, skip_special_tokens=True))
|
||||
except TypeError:
|
||||
segment_text = str(decode_token_ids(decode_ids))
|
||||
except Exception:
|
||||
logger.exception("Failed to decode MiniCPM-o duplex token delta for request_id=%s", request_id)
|
||||
segment_text = ""
|
||||
decoded_segment_text = _decode_native_duplex_token_ids(
|
||||
segment_ids,
|
||||
streaming_context,
|
||||
request_id=request_id,
|
||||
)
|
||||
if decoded_segment_text is not None:
|
||||
segment_text = decoded_segment_text
|
||||
elif sent_len == 0:
|
||||
# Non-orchestrator unit tests and legacy callers may not install a
|
||||
# decoder. Never slice cumulative text with a stale character cursor.
|
||||
@@ -575,22 +622,7 @@ def _native_duplex_segment_output_ids(
|
||||
state["sent_output_len"] = len(output_ids)
|
||||
state["sent_output_ids"] = list(output_ids)
|
||||
state["turn_id"] = turn_id
|
||||
return segment_ids, segment_text
|
||||
|
||||
|
||||
def _reset_native_tts_handoff(streaming_context) -> None:
|
||||
bridge_states = getattr(streaming_context, "bridge_states", None)
|
||||
if not isinstance(bridge_states, dict):
|
||||
return
|
||||
state = bridge_states.get("minicpmo45_tts_handoff")
|
||||
if not isinstance(state, dict):
|
||||
return
|
||||
# Keep the output cursor: the thinker can keep reporting cumulative
|
||||
# output after turn_eos on the same resumable request. If a runtime really
|
||||
# restarts token ids, _native_duplex_segment_output_ids detects the prefix
|
||||
# mismatch and resets the slice cursor there.
|
||||
state["acc_tts_ids"] = []
|
||||
state["acc_tts_hidden"] = []
|
||||
return segment_ids, segment_text, turn_start
|
||||
|
||||
|
||||
def _native_duplex_data_plane_metadata(streaming_context) -> dict[str, object] | None:
|
||||
@@ -623,32 +655,6 @@ def _native_duplex_data_plane_metadata(streaming_context) -> dict[str, object] |
|
||||
return metadata
|
||||
|
||||
|
||||
def _accumulate_native_tts_handoff(streaming_context, new_ids, new_hidden):
|
||||
"""Hand the talker the FULL accumulated condition on every handoff.
|
||||
|
||||
The runner's streaming buffer update is not merge-safe for a resumable
|
||||
stage-1 request: in-place updates merge sub-keys, but a resume prefill
|
||||
REPLACES the buffer, silently dropping every earlier segment's tts
|
||||
tokens/hiddens (observed losing alternating reply segments — the talker
|
||||
vocalized text it never saw between islands it did). Accumulate here so
|
||||
the latest handoff always carries the complete history and downstream
|
||||
replace semantics are lossless; the talker consumes by cursor.
|
||||
"""
|
||||
bridge_states = getattr(streaming_context, "bridge_states", None)
|
||||
if not isinstance(bridge_states, dict):
|
||||
return new_ids, new_hidden
|
||||
state = bridge_states.setdefault("minicpmo45_tts_handoff", {})
|
||||
acc_ids = state.setdefault("acc_tts_ids", [])
|
||||
acc_hidden = state.setdefault("acc_tts_hidden", [])
|
||||
if new_ids:
|
||||
acc_ids.extend(int(t) for t in new_ids)
|
||||
if new_hidden:
|
||||
acc_hidden.extend(new_hidden)
|
||||
if not acc_ids:
|
||||
return None, None
|
||||
return list(acc_ids), list(acc_hidden)
|
||||
|
||||
|
||||
def _build_tts_scheduler_prompt_token_ids(
|
||||
tts_token_ids: torch.Tensor | None,
|
||||
llm_output_ids: list[int],
|
||||
@@ -723,6 +729,7 @@ def llm2tts(
|
||||
# buffer overflows.
|
||||
llm_output_ids = list(llm_output_ids)
|
||||
thinker_text = getattr(output, "text", "") or ""
|
||||
native_turn_start = False
|
||||
if _has_native_duplex_prompt_metadata(mm_output):
|
||||
# The thinker's resumable duplex request reports cumulative
|
||||
# output ids/text, but earlier segments are already folded into
|
||||
@@ -732,7 +739,7 @@ def llm2tts(
|
||||
# alignment holds, the talker prompt grows linearly with new
|
||||
# tokens instead of quadratically, and downstream transcripts
|
||||
# carry per-unit deltas instead of re-sending the whole reply.
|
||||
llm_output_ids, thinker_text = _native_duplex_segment_output_ids(
|
||||
llm_output_ids, thinker_text, native_turn_start = _native_duplex_segment_output_ids(
|
||||
llm_output_ids,
|
||||
thinker_text,
|
||||
_streaming_context,
|
||||
@@ -875,6 +882,18 @@ def llm2tts(
|
||||
.to(torch.float32)
|
||||
.contiguous()
|
||||
)
|
||||
handoff_ids = _coerce_token_id_list(tts_token_ids_slice) if tts_token_ids_slice is not None else None
|
||||
if is_native_duplex_handoff and handoff_ids:
|
||||
handoff_text = _decode_native_duplex_token_ids(
|
||||
handoff_ids,
|
||||
_streaming_context,
|
||||
request_id=str(llm_output.request_id),
|
||||
)
|
||||
if handoff_text is not None:
|
||||
# Match the released streaming_generate contract: transcript
|
||||
# text comes from total_ids_in_unit, the same slice that
|
||||
# conditions the Talker.
|
||||
thinker_text = handoff_text
|
||||
model_intermediate_buffer = build_duplex_intermediate_buffer(
|
||||
request_id=str(llm_output.request_id),
|
||||
prompt_token_ids=prompt_token_ids,
|
||||
@@ -900,6 +919,7 @@ def llm2tts(
|
||||
# The talker detects turn end from <|turn_eos|> inside the
|
||||
# handed condition (official conditions on its embedding).
|
||||
meta["turn_eos_token_id"] = int(turn_eos_id)
|
||||
meta["turn_start"] = native_turn_start
|
||||
if native_segment_end:
|
||||
meta["segment_end"] = True
|
||||
req_mm_data = multi_modal_data.get(llm_output.request_id)
|
||||
@@ -911,30 +931,27 @@ def llm2tts(
|
||||
if ref_audio is not None:
|
||||
ref_waveform, ref_sr = ref_audio
|
||||
set_ref_audio(model_intermediate_buffer, _to_transport_list(ref_waveform), ref_sr)
|
||||
handoff_ids = _coerce_token_id_list(tts_token_ids_slice) if tts_token_ids_slice is not None else None
|
||||
handoff_hidden = _to_transport_list(tts_hidden_slice) if tts_hidden_slice is not None else None
|
||||
native_turn_end_handoff = False
|
||||
if is_native_duplex_handoff:
|
||||
turn_eos_id = special_token_ids.get("turn_eos_token_id")
|
||||
native_turn_end_handoff = turn_eos_id is not None and handoff_ids is not None and turn_eos_id in handoff_ids
|
||||
handoff_ids, handoff_hidden = _accumulate_native_tts_handoff(
|
||||
_streaming_context,
|
||||
handoff_ids,
|
||||
handoff_hidden,
|
||||
)
|
||||
if not handoff_ids:
|
||||
continue
|
||||
set_tts_handoff(model_intermediate_buffer, handoff_ids, handoff_hidden)
|
||||
if native_turn_end_handoff:
|
||||
model_intermediate_buffer.setdefault("meta", {})["turn_end"] = True
|
||||
_reset_native_tts_handoff(_streaming_context)
|
||||
|
||||
if handoff_ids is not None and handoff_hidden is not None:
|
||||
condition_length = max(len(handoff_ids), len(handoff_hidden)) + 2
|
||||
condition_suffix_length = 1 if is_native_duplex_handoff else 2
|
||||
condition_length = max(len(handoff_ids), len(handoff_hidden)) + condition_suffix_length
|
||||
scheduler_prompt_token_ids = [0] * condition_length
|
||||
handoff_meta = model_intermediate_buffer.setdefault("meta", {})
|
||||
handoff_meta["replace_streaming_prompt"] = True
|
||||
handoff_meta["next_stage_prompt_len"] = condition_length
|
||||
# Native duplex resumes one Talker request within a turn, but a new
|
||||
# assistant turn must discard the previous turn's prompt and KV.
|
||||
if not is_native_duplex_handoff or native_turn_start:
|
||||
handoff_meta["replace_streaming_prompt"] = True
|
||||
else:
|
||||
scheduler_prompt_token_ids = _build_tts_scheduler_prompt_token_ids(
|
||||
tts_token_ids_slice,
|
||||
|
||||
@@ -685,19 +685,32 @@ def talker2code2wav_async_chunk(
|
||||
"""
|
||||
Multimodal output version.
|
||||
"""
|
||||
if not isinstance(multimodal_output, Mapping):
|
||||
return None
|
||||
talker_codes = multimodal_output.get("codes", {})
|
||||
if not isinstance(talker_codes, dict):
|
||||
return None
|
||||
code_predictor_codes = talker_codes.get("audio")
|
||||
if code_predictor_codes is None:
|
||||
return None
|
||||
request_id = request.external_req_id
|
||||
code_predictor_codes = None
|
||||
if isinstance(multimodal_output, Mapping):
|
||||
talker_codes = multimodal_output.get("codes", {})
|
||||
if isinstance(talker_codes, dict):
|
||||
code_predictor_codes = talker_codes.get("audio")
|
||||
|
||||
if code_predictor_codes.numel() == 0:
|
||||
return None
|
||||
sampling_params = getattr(request, "sampling_params", None)
|
||||
stop_token_ids = set(getattr(sampling_params, "stop_token_ids", None) or [])
|
||||
stop_token_id = getattr(sampling_params, "stop_token_id", None)
|
||||
if stop_token_id is not None:
|
||||
stop_token_ids.add(stop_token_id)
|
||||
|
||||
if not code_predictor_codes.any():
|
||||
append_codes = (
|
||||
isinstance(code_predictor_codes, torch.Tensor)
|
||||
and code_predictor_codes.numel() > 0
|
||||
and bool(code_predictor_codes.any())
|
||||
)
|
||||
if append_codes:
|
||||
first_codebook = int(code_predictor_codes[0, 0].item())
|
||||
if first_codebook in stop_token_ids:
|
||||
logger.debug("skip stop-token codec frame: first_codebook=%s", first_codebook)
|
||||
append_codes = False
|
||||
if append_codes:
|
||||
transfer_manager.code_prompt_token_ids[request_id].append(code_predictor_codes)
|
||||
elif not is_finished:
|
||||
return None
|
||||
|
||||
connector = getattr(transfer_manager, "connector", None)
|
||||
@@ -707,20 +720,10 @@ def talker2code2wav_async_chunk(
|
||||
left_context_size_config = int(cfg.get("codec_left_context_frames", 25))
|
||||
configured_initial_chunk_size = int(cfg.get("initial_codec_chunk_frames") or 0)
|
||||
|
||||
sampling_params = getattr(request, "sampling_params", None)
|
||||
stop_token_ids = set(getattr(sampling_params, "stop_token_ids", None) or [])
|
||||
stop_token_id = getattr(sampling_params, "stop_token_id", None)
|
||||
if stop_token_id is not None:
|
||||
stop_token_ids.add(stop_token_id)
|
||||
first_codebook = int(code_predictor_codes[0, 0].item())
|
||||
if first_codebook in stop_token_ids:
|
||||
logger.debug("skip stop-token codec frame: first_codebook=%s", first_codebook)
|
||||
return None
|
||||
|
||||
request_id = request.external_req_id
|
||||
chunk_id = transfer_manager.put_req_chunk[request_id]
|
||||
transfer_manager.code_prompt_token_ids[request_id].append(code_predictor_codes)
|
||||
length = len(transfer_manager.code_prompt_token_ids[request_id])
|
||||
if length <= 0:
|
||||
return None
|
||||
|
||||
if configured_initial_chunk_size > 0:
|
||||
if chunk_id == 0:
|
||||
@@ -732,6 +735,9 @@ def talker2code2wav_async_chunk(
|
||||
if chunk_length != 0 and not is_finished:
|
||||
return None
|
||||
|
||||
if is_finished and not append_codes and chunk_length == 0:
|
||||
return None
|
||||
|
||||
context_length = chunk_length if chunk_length != 0 else chunk_size_config
|
||||
# ensure left context does not exceed available length
|
||||
if configured_initial_chunk_size > 0 and chunk_id == 1:
|
||||
|
||||
@@ -11,6 +11,7 @@ CHUNK_METADATA_KEYS: frozenset[str] = frozenset(
|
||||
"duplex_epoch",
|
||||
"duplex_turn_id",
|
||||
"llm_output_text_utf8",
|
||||
"segment_end",
|
||||
"tts_is_last_chunk",
|
||||
"turn_end",
|
||||
}
|
||||
|
||||
@@ -523,24 +523,27 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
|
||||
upstream_outputs.append(eco)
|
||||
|
||||
# Handle multimodal-only outputs (generation stages) locally.
|
||||
self._process_mm_only_outputs(mm_only_outputs)
|
||||
mm_request_outputs = self._process_mm_only_outputs(mm_only_outputs)
|
||||
|
||||
# Delegate text/pooling outputs to upstream.
|
||||
return super().process_outputs(
|
||||
processed = super().process_outputs(
|
||||
upstream_outputs,
|
||||
engine_core_timestamp=engine_core_timestamp,
|
||||
iteration_stats=iteration_stats,
|
||||
)
|
||||
processed.request_outputs.extend(mm_request_outputs)
|
||||
return processed
|
||||
|
||||
def _process_mm_only_outputs(
|
||||
self,
|
||||
engine_core_outputs: list[EngineCoreOutput],
|
||||
) -> None:
|
||||
) -> list[OmniRequestOutput | PoolingRequestOutput]:
|
||||
"""Handle outputs from generation stages that have no detokenizer.
|
||||
|
||||
These cannot go through upstream process_outputs because it asserts
|
||||
detokenizer is not None when pooling_output is None.
|
||||
"""
|
||||
request_outputs: list[OmniRequestOutput | PoolingRequestOutput] = []
|
||||
for eco in engine_core_outputs:
|
||||
req_state = self.request_states.get(eco.request_id)
|
||||
if req_state is None or not isinstance(req_state, OmniRequestState):
|
||||
@@ -554,6 +557,11 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
|
||||
req_state.num_cached_tokens = eco.num_cached_tokens
|
||||
req_state.is_prefilling = False
|
||||
|
||||
is_non_final_audio_chunk = (
|
||||
finish_reason is not None
|
||||
and req_state.output_kind == RequestOutputKind.DELTA
|
||||
and is_non_final_delta_audio_chunk(req_state.mm_accumulated, req_state.mm_type)
|
||||
)
|
||||
if request_output := req_state.make_request_output(
|
||||
new_token_ids,
|
||||
None, # pooling_output
|
||||
@@ -564,9 +572,13 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
|
||||
):
|
||||
if req_state.queue is not None:
|
||||
req_state.queue.put(request_output)
|
||||
else:
|
||||
request_outputs.append(request_output)
|
||||
|
||||
if finish_reason is not None:
|
||||
is_segment_finished = bool(getattr(eco, "is_segment_finished", False))
|
||||
if finish_reason is not None and not is_segment_finished and not is_non_final_audio_chunk:
|
||||
self._finish_request(req_state)
|
||||
return request_outputs
|
||||
|
||||
def _update_stats_from_output(
|
||||
self,
|
||||
|
||||
Reference in new issue
Block a user