Yueqian LinandSy0307
e788ef6e76
Signed-off-by: Yueqian Lin <linyueqian@outlook.com> Signed-off-by: Sy0307 <24841857+Sy0307@users.noreply.github.com> Co-authored-by: Sy0307 <24841857+Sy0307@users.noreply.github.com>
493 lines
19 KiB
Python
493 lines
19 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Strict PersonaPlex validation for the unified OpenAI Realtime duplex path."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import asyncio
|
|
import base64
|
|
import hashlib
|
|
import json
|
|
import math
|
|
import uuid
|
|
import wave
|
|
from collections.abc import Sequence
|
|
from pathlib import Path
|
|
from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
|
|
|
|
import numpy as np
|
|
|
|
from vllm_omni.experimental.fullduplex.client import (
|
|
RealtimeDuplexClient,
|
|
wait_for,
|
|
write_pcm16_wav,
|
|
)
|
|
|
|
SAMPLE_RATE_HZ = 24_000
|
|
FRAME_SAMPLES = 1_920
|
|
FRAME_PERIOD_S = FRAME_SAMPLES / SAMPLE_RATE_HZ
|
|
|
|
|
|
def _input_identity(
|
|
path: Path,
|
|
*,
|
|
expected_sha256: str | None,
|
|
) -> dict[str, str]:
|
|
resolved = path.resolve()
|
|
digest = hashlib.sha256()
|
|
with resolved.open("rb") as input_file:
|
|
for chunk in iter(lambda: input_file.read(1024 * 1024), b""):
|
|
digest.update(chunk)
|
|
actual = digest.hexdigest()
|
|
if expected_sha256 is not None and actual != expected_sha256.lower():
|
|
raise ValueError(
|
|
f"input WAV SHA-256 mismatch: expected={expected_sha256.lower()}, actual={actual}, path={resolved}"
|
|
)
|
|
return {"path": str(resolved), "sha256": actual}
|
|
|
|
|
|
def _realtime_url(base_url: str, model: str, session_id: str) -> str:
|
|
parts = urlsplit(base_url)
|
|
query = dict(parse_qsl(parts.query, keep_blank_values=True))
|
|
query.update(duplex="1", model=model, autostart="0", session_id=session_id)
|
|
return urlunsplit((parts.scheme, parts.netloc, parts.path, urlencode(query), parts.fragment))
|
|
|
|
|
|
def _read_wav_as_float32(path: Path) -> np.ndarray:
|
|
with wave.open(str(path), "rb") as wav_file:
|
|
channels = wav_file.getnchannels()
|
|
sample_width = wav_file.getsampwidth()
|
|
source_rate = wav_file.getframerate()
|
|
frames = wav_file.readframes(wav_file.getnframes())
|
|
dtypes = {1: np.uint8, 2: np.dtype("<i2"), 4: np.dtype("<i4")}
|
|
if sample_width not in dtypes:
|
|
raise ValueError(f"unsupported WAV sample width: {sample_width}")
|
|
pcm = np.frombuffer(frames, dtype=dtypes[sample_width]).astype(np.float32)
|
|
if sample_width == 1:
|
|
pcm = (pcm - 128.0) / 128.0
|
|
else:
|
|
pcm /= float(1 << (sample_width * 8 - 1))
|
|
if channels > 1:
|
|
pcm = pcm.reshape(-1, channels).mean(axis=1)
|
|
if source_rate != SAMPLE_RATE_HZ and pcm.size:
|
|
output_samples = max(1, round(pcm.size * SAMPLE_RATE_HZ / source_rate))
|
|
pcm = np.interp(
|
|
np.linspace(0, pcm.size - 1, output_samples),
|
|
np.arange(pcm.size),
|
|
pcm,
|
|
).astype(np.float32)
|
|
pcm = np.ascontiguousarray(pcm, dtype="<f4")
|
|
if not pcm.size or not np.isfinite(pcm).all():
|
|
raise ValueError("input WAV must contain finite audio")
|
|
return pcm
|
|
|
|
|
|
def _events(client: RealtimeDuplexClient, event_type: str) -> list[dict[str, object]]:
|
|
return [event for event in client.events.events if event.get("type") == event_type]
|
|
|
|
|
|
async def _open_session(
|
|
args: argparse.Namespace,
|
|
*,
|
|
session_id: str,
|
|
persona: str,
|
|
expect_error: bool = False,
|
|
) -> tuple[RealtimeDuplexClient, dict[str, object]]:
|
|
client = RealtimeDuplexClient(_realtime_url(args.url, args.model, session_id))
|
|
await client.__aenter__()
|
|
await client.send(
|
|
{
|
|
"type": "session.update",
|
|
"session": {
|
|
"session_id": session_id,
|
|
"model": args.model,
|
|
"modalities": ["audio", "text"],
|
|
"input_audio_format": "pcm_f32le",
|
|
"output_audio_format": "pcm16",
|
|
"voice": args.voice,
|
|
"instructions": persona,
|
|
"turn_detection": None,
|
|
"extra_body": {"auto_response": True},
|
|
},
|
|
}
|
|
)
|
|
event_type = "error" if expect_error else "session.created"
|
|
await wait_for(
|
|
lambda: client.events.count(event_type) > 0,
|
|
timeout_s=args.timeout_s,
|
|
label=event_type,
|
|
)
|
|
return client, _events(client, event_type)[-1]
|
|
|
|
|
|
async def _close_session(client: RealtimeDuplexClient, *, timeout_s: float) -> None:
|
|
await client.send({"type": "session.close"})
|
|
await wait_for(
|
|
lambda: client.events.count("session.closed") > 0,
|
|
timeout_s=timeout_s,
|
|
label="session.closed",
|
|
)
|
|
await client.__aexit__(None, None, None)
|
|
|
|
|
|
async def _stream_frames(
|
|
client: RealtimeDuplexClient,
|
|
pcm: np.ndarray,
|
|
*,
|
|
max_frames: int | None = None,
|
|
) -> int:
|
|
frame_count = math.ceil(pcm.size / FRAME_SAMPLES)
|
|
if max_frames is not None:
|
|
if max_frames < 0:
|
|
raise ValueError(f"max_frames must be non-negative, got {max_frames}")
|
|
if max_frames > 0:
|
|
frame_count = min(frame_count, max_frames)
|
|
for seq in range(frame_count):
|
|
frame = pcm[seq * FRAME_SAMPLES : (seq + 1) * FRAME_SAMPLES]
|
|
frame = np.pad(frame, (0, FRAME_SAMPLES - frame.size))
|
|
frame = np.ascontiguousarray(frame, dtype="<f4")
|
|
await client.send(
|
|
{
|
|
"type": "input_audio_buffer.append",
|
|
"audio": base64.b64encode(frame).decode("ascii"),
|
|
"format": "pcm_f32le",
|
|
"sample_rate_hz": SAMPLE_RATE_HZ,
|
|
"duration_ms": 80,
|
|
"audio_end_ms": (seq + 1) * 80,
|
|
"is_speech": bool(np.sqrt(np.mean(np.square(frame))) >= 1e-4),
|
|
}
|
|
)
|
|
await asyncio.sleep(FRAME_PERIOD_S)
|
|
return frame_count
|
|
|
|
|
|
def _audio_frame_stats(
|
|
raw: bytes,
|
|
*,
|
|
input_frames: int,
|
|
voiced_frame_rms_threshold: float,
|
|
) -> dict[str, int | float]:
|
|
pcm = np.frombuffer(raw, dtype="<i2")
|
|
if pcm.size % FRAME_SAMPLES:
|
|
raise AssertionError(
|
|
"PersonaPlex output must contain whole 80 ms codec frames: "
|
|
f"samples={pcm.size}, samples_per_frame={FRAME_SAMPLES}"
|
|
)
|
|
output_frames = pcm.size // FRAME_SAMPLES
|
|
frames = pcm.astype(np.float32).reshape(output_frames, FRAME_SAMPLES) / 32768.0
|
|
voiced_frames = int(np.count_nonzero(np.sqrt(np.mean(np.square(frames), axis=1)) >= voiced_frame_rms_threshold))
|
|
return {
|
|
"output_frames": int(output_frames),
|
|
"voiced_frames": voiced_frames,
|
|
"silent_frames": int(output_frames - voiced_frames),
|
|
"frame_deficit": int(input_frames - output_frames),
|
|
"frame_coverage_ratio": float(output_frames / input_frames),
|
|
}
|
|
|
|
|
|
def _response_ids(client: RealtimeDuplexClient) -> set[str]:
|
|
return {
|
|
str(event["response_id"])
|
|
for event in _events(client, "response.audio.delta")
|
|
if isinstance(event.get("response_id"), str)
|
|
}
|
|
|
|
|
|
def _session_result(
|
|
client: RealtimeDuplexClient,
|
|
*,
|
|
input_frames: int,
|
|
args: argparse.Namespace,
|
|
minimum_chunks: int,
|
|
) -> tuple[bytes, set[str], dict[str, object]]:
|
|
chunks = sum(len(items) for items in client.events.response_audio.values())
|
|
raw = client.events.audio_bytes()
|
|
if chunks < minimum_chunks or not raw or len(raw) % 2:
|
|
raise AssertionError(f"invalid PCM16 output: chunks={chunks}, bytes={len(raw)}")
|
|
pcm = np.frombuffer(raw, dtype="<i2").astype(np.float32) / 32768.0
|
|
rms = float(np.sqrt(np.mean(np.square(pcm))))
|
|
# This low floor only catches corrupt, non-finite, or effectively zero PCM.
|
|
# Audible speech is checked per frame below with the higher 1e-3 default.
|
|
if not np.isfinite(pcm).all() or rms <= 1e-5:
|
|
raise AssertionError(f"output audio is non-finite or silent: samples={pcm.size}, rms={rms}")
|
|
rates = {
|
|
int(event["sample_rate_hz"])
|
|
for event in _events(client, "response.audio.delta")
|
|
if isinstance(event.get("sample_rate_hz"), int)
|
|
}
|
|
if rates != {SAMPLE_RATE_HZ}:
|
|
raise AssertionError(f"unexpected output sample rates: {sorted(rates)}")
|
|
stats = _audio_frame_stats(
|
|
raw,
|
|
input_frames=input_frames,
|
|
voiced_frame_rms_threshold=args.voiced_frame_rms_threshold,
|
|
)
|
|
if not 0 <= int(stats["frame_deficit"]) <= args.max_frame_deficit:
|
|
raise AssertionError(f"invalid synchronous frame accounting: input={input_frames}, stats={stats}")
|
|
min_voiced_frames = int(args.min_voiced_frames)
|
|
if min_voiced_frames < 1:
|
|
raise ValueError(f"min_voiced_frames must be positive, got {min_voiced_frames}")
|
|
if int(stats["voiced_frames"]) < min_voiced_frames:
|
|
raise AssertionError(
|
|
f"output contained insufficient audible speech: minimum_voiced_frames={min_voiced_frames}, stats={stats}"
|
|
)
|
|
runtime_metadata = [
|
|
metadata["vllm_omni"]
|
|
for event in _events(client, "response.audio.delta")
|
|
if isinstance((metadata := event.get("metadata")), dict) and isinstance(metadata.get("vllm_omni"), dict)
|
|
]
|
|
if not any(
|
|
item.get("runtime_impl") == "scheduler_data_plane"
|
|
and item.get("uses_model_runner_scheduler") is True
|
|
and item.get("runner_kv_backed") is True
|
|
for item in runtime_metadata
|
|
):
|
|
raise AssertionError(f"audio events did not prove the scheduler data plane: {runtime_metadata}")
|
|
response_ids = _response_ids(client)
|
|
if len(response_ids) != 1:
|
|
raise AssertionError(f"continuous stream used multiple responses: {sorted(response_ids)}")
|
|
return (
|
|
raw,
|
|
response_ids,
|
|
{
|
|
**stats,
|
|
"audio_chunks": chunks,
|
|
"audio_samples": len(raw) // 2,
|
|
"audio_rms": rms,
|
|
"response_ids": sorted(response_ids),
|
|
},
|
|
)
|
|
|
|
|
|
def _capabilities(created: dict[str, object]) -> dict[str, object]:
|
|
session = created.get("session")
|
|
capabilities = session.get("capabilities") if isinstance(session, dict) else None
|
|
if not isinstance(capabilities, dict):
|
|
raise AssertionError(f"session.created omitted capabilities: {created}")
|
|
return capabilities
|
|
|
|
|
|
def _save(output_dir: Path, name: str, client: RealtimeDuplexClient, raw: bytes) -> None:
|
|
(output_dir / f"{name}-events.jsonl").write_text(
|
|
"".join(json.dumps(event, ensure_ascii=False) + "\n" for event in client.events.events),
|
|
encoding="utf-8",
|
|
)
|
|
write_pcm16_wav(output_dir / f"{name}-output.wav", raw, sample_rate_hz=SAMPLE_RATE_HZ)
|
|
|
|
|
|
async def run(args: argparse.Namespace) -> dict[str, object]:
|
|
input_path = Path(args.input_wav)
|
|
input_identity = _input_identity(
|
|
input_path,
|
|
expected_sha256=args.expected_input_sha256,
|
|
)
|
|
pcm = _read_wav_as_float32(input_path)
|
|
if args.tail_s > 0:
|
|
pcm = np.concatenate([pcm, np.zeros(round(args.tail_s * SAMPLE_RATE_HZ), dtype=np.float32)])
|
|
output_dir = Path(args.output_dir)
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
ids = {name: f"personaplex-{name}-{uuid.uuid4().hex}" for name in ("primary", "secondary", "replacement")}
|
|
primary, created = await _open_session(args, session_id=ids["primary"], persona=args.persona)
|
|
secondary, secondary_created = await _open_session(
|
|
args,
|
|
session_id=ids["secondary"],
|
|
persona=args.secondary_persona,
|
|
)
|
|
capabilities = _capabilities(created)
|
|
if _capabilities(secondary_created) != capabilities:
|
|
raise AssertionError("concurrent sessions returned different capabilities")
|
|
expected_capabilities = {
|
|
"implementation_level": "model_native_duplex",
|
|
"chunk_period_ms": 80,
|
|
"supports_multi_session": True,
|
|
"supports_multi_session_same_replica": True,
|
|
"supports_barge_in": False,
|
|
}
|
|
if any(capabilities.get(key) != value for key, value in expected_capabilities.items()):
|
|
raise AssertionError(f"unexpected PersonaPlex capabilities: {capabilities}")
|
|
|
|
overflow, error = await _open_session(
|
|
args,
|
|
session_id=f"personaplex-overflow-{uuid.uuid4().hex}",
|
|
persona=args.persona,
|
|
expect_error=True,
|
|
)
|
|
error_body = error.get("error")
|
|
overflow_code = error_body.get("code") if isinstance(error_body, dict) else error.get("code")
|
|
await overflow.__aexit__(None, None, None)
|
|
if overflow_code != "resource_exhausted":
|
|
raise AssertionError(f"overflow session was not rejected: {error}")
|
|
|
|
primary_frames, secondary_frames = await asyncio.gather(
|
|
_stream_frames(primary, pcm),
|
|
_stream_frames(secondary, pcm),
|
|
)
|
|
await asyncio.gather(
|
|
wait_for(
|
|
lambda: len(primary.events.audio_bytes()) // (2 * FRAME_SAMPLES) >= primary_frames - args.max_frame_deficit,
|
|
timeout_s=args.timeout_s,
|
|
label="primary frame coverage",
|
|
),
|
|
wait_for(
|
|
lambda: (
|
|
len(secondary.events.audio_bytes()) // (2 * FRAME_SAMPLES) >= secondary_frames - args.max_frame_deficit
|
|
),
|
|
timeout_s=args.timeout_s,
|
|
label="secondary frame coverage",
|
|
),
|
|
)
|
|
await asyncio.sleep(args.drain_s)
|
|
primary_audio, primary_response_ids, primary_stats = _session_result(
|
|
primary,
|
|
input_frames=primary_frames,
|
|
args=args,
|
|
minimum_chunks=args.minimum_audio_chunks,
|
|
)
|
|
secondary_audio, secondary_response_ids, _ = _session_result(
|
|
secondary,
|
|
input_frames=secondary_frames,
|
|
args=args,
|
|
minimum_chunks=args.minimum_audio_chunks,
|
|
)
|
|
if primary_response_ids & secondary_response_ids:
|
|
raise AssertionError("concurrent sessions shared a response id")
|
|
_save(output_dir, "primary", primary, primary_audio)
|
|
await _close_session(primary, timeout_s=args.timeout_s)
|
|
|
|
replacement, replacement_created = await _open_session(
|
|
args,
|
|
session_id=ids["replacement"],
|
|
persona=args.replacement_persona,
|
|
)
|
|
if _capabilities(replacement_created) != capabilities:
|
|
raise AssertionError("replacement session returned different capabilities")
|
|
continuation_frames, replacement_frames = await asyncio.gather(
|
|
_stream_frames(secondary, pcm, max_frames=args.replacement_frames),
|
|
_stream_frames(replacement, pcm, max_frames=args.replacement_frames),
|
|
)
|
|
secondary_total_frames = secondary_frames + continuation_frames
|
|
await asyncio.gather(
|
|
wait_for(
|
|
lambda: (
|
|
len(secondary.events.audio_bytes()) // (2 * FRAME_SAMPLES)
|
|
>= secondary_total_frames - args.max_frame_deficit
|
|
),
|
|
timeout_s=args.timeout_s,
|
|
label="survivor frame coverage after slot recycle",
|
|
),
|
|
wait_for(
|
|
lambda: (
|
|
len(replacement.events.audio_bytes()) // (2 * FRAME_SAMPLES)
|
|
>= replacement_frames - args.max_frame_deficit
|
|
),
|
|
timeout_s=args.timeout_s,
|
|
label="replacement frame coverage",
|
|
),
|
|
)
|
|
await asyncio.sleep(args.drain_s)
|
|
secondary_audio, secondary_after_ids, secondary_stats = _session_result(
|
|
secondary,
|
|
input_frames=secondary_total_frames,
|
|
args=args,
|
|
minimum_chunks=args.minimum_audio_chunks + 1,
|
|
)
|
|
replacement_audio, replacement_response_ids, replacement_stats = _session_result(
|
|
replacement,
|
|
input_frames=replacement_frames,
|
|
args=args,
|
|
minimum_chunks=1,
|
|
)
|
|
if secondary_after_ids != secondary_response_ids:
|
|
raise AssertionError("survivor response changed while another session slot was recycled")
|
|
if replacement_response_ids & (primary_response_ids | secondary_response_ids):
|
|
raise AssertionError("replacement session reused another session's response id")
|
|
_save(output_dir, "secondary", secondary, secondary_audio)
|
|
_save(output_dir, "replacement", replacement, replacement_audio)
|
|
await asyncio.gather(
|
|
_close_session(secondary, timeout_s=args.timeout_s),
|
|
_close_session(replacement, timeout_s=args.timeout_s),
|
|
)
|
|
|
|
errors = primary.events.errors() + secondary.events.errors() + replacement.events.errors()
|
|
result = {
|
|
"ok": not errors,
|
|
"model": args.model,
|
|
"input": {**input_identity, "tail_s": args.tail_s},
|
|
"capabilities": capabilities,
|
|
"overflow_error_code": overflow_code,
|
|
"primary": {"session_id": ids["primary"], "input_frames": primary_frames, **primary_stats},
|
|
"secondary": {
|
|
"session_id": ids["secondary"],
|
|
"input_frames": secondary_total_frames,
|
|
"continuation_frames_after_recycle": continuation_frames,
|
|
**secondary_stats,
|
|
},
|
|
"replacement": {
|
|
"session_id": ids["replacement"],
|
|
"input_frames": replacement_frames,
|
|
**replacement_stats,
|
|
},
|
|
"errors": errors,
|
|
"output_dir": str(output_dir),
|
|
}
|
|
(output_dir / "result.json").write_text(json.dumps(result, ensure_ascii=False, indent=2), encoding="utf-8")
|
|
if errors:
|
|
raise AssertionError(f"Realtime sessions emitted errors: {errors}")
|
|
return result
|
|
|
|
|
|
def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--url", default="ws://127.0.0.1:8099/v1/realtime?duplex=1")
|
|
parser.add_argument("--model", required=True)
|
|
parser.add_argument("--input-wav", required=True)
|
|
parser.add_argument(
|
|
"--expected-input-sha256",
|
|
help="Fail before opening sessions when the input WAV does not match this SHA-256.",
|
|
)
|
|
parser.add_argument("--output-dir", default="/tmp/personaplex-realtime-duplex")
|
|
parser.add_argument("--voice", default="NATF2.pt")
|
|
parser.add_argument("--persona", default="You are a concise and helpful assistant.")
|
|
parser.add_argument("--secondary-persona", default="You are a calm and precise assistant.")
|
|
parser.add_argument("--replacement-persona", default="You are a cheerful travel assistant.")
|
|
parser.add_argument("--tail-s", type=float, default=2.0)
|
|
parser.add_argument("--drain-s", type=float, default=2.0)
|
|
parser.add_argument("--timeout-s", type=float, default=180.0)
|
|
parser.add_argument(
|
|
"--max-frame-deficit",
|
|
type=int,
|
|
default=4,
|
|
help=(
|
|
"Maximum unflushed tail frames. PersonaPlex Code2Wav emits five-frame "
|
|
"chunks, so a session cutoff can leave at most four pending frames."
|
|
),
|
|
)
|
|
parser.add_argument("--voiced-frame-rms-threshold", type=float, default=1e-3)
|
|
parser.add_argument(
|
|
"--min-voiced-frames",
|
|
type=int,
|
|
default=5,
|
|
help=(
|
|
"Require at least one normal five-frame Code2Wav chunk (400 ms) above "
|
|
"the per-frame RMS threshold; this detects inaudible output, not speech quality."
|
|
),
|
|
)
|
|
parser.add_argument("--minimum-audio-chunks", type=int, default=2)
|
|
parser.add_argument(
|
|
"--replacement-frames",
|
|
type=int,
|
|
default=0,
|
|
help="Frames sent after slot reuse; 0 replays the full input so the replacement must produce audible output.",
|
|
)
|
|
return parser.parse_args(argv)
|
|
|
|
|
|
def main() -> None:
|
|
print(json.dumps(asyncio.run(run(parse_args())), ensure_ascii=False, indent=2))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|