Signed-off-by: chickeyton <ngton2014@gmail.com> Signed-off-by: Yueqian Lin <linyueqian@outlook.com> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com> Co-authored-by: Yueqian Lin <linyueqian@outlook.com>
This commit is contained in:
49 files changed
+3056
-1715
No files matched your search
@@ -0,0 +1,194 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""Transactional framing of client PCM into fixed model units.
|
||||
|
||||
A frame-locked model (PersonaPlex: 1920 samples every 80 ms at 24 kHz) takes
|
||||
exactly one frame per append. Clients send arbitrary chunks, so the session
|
||||
runner buffers them here and takes one whole frame out at a time as a
|
||||
*reservation*: committed once the stage accepted the append, rolled back
|
||||
(the bytes go back to the front of the buffer) when it did not.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pybase64 as base64
|
||||
|
||||
from vllm_omni.engine.duplex.plugin import PcmAppendBuffer, PcmAppendReservation
|
||||
from vllm_omni.model_executor.common.audio.pcm import (
|
||||
PCM_F32LE_BYTES_PER_SAMPLE,
|
||||
decode_pcm_f32le_base64,
|
||||
)
|
||||
|
||||
|
||||
class FixedFramePcmAppendReservation(PcmAppendReservation):
|
||||
__slots__ = ("_active", "_owner", "_raw", "operation_id", "payload")
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
owner: FixedFramePcmAppendBuffer,
|
||||
operation_id: str,
|
||||
payload: dict[str, object] | None,
|
||||
raw: bytes,
|
||||
) -> None:
|
||||
self._owner = owner
|
||||
self.operation_id = operation_id
|
||||
self.payload = payload
|
||||
self._raw = raw
|
||||
self._active = True
|
||||
|
||||
@property
|
||||
def active(self) -> bool:
|
||||
return self._active
|
||||
|
||||
@property
|
||||
def byte_count(self) -> int:
|
||||
return len(self._raw)
|
||||
|
||||
def commit(self) -> None:
|
||||
self._owner._commit_reservation(self)
|
||||
|
||||
def rollback(self) -> None:
|
||||
self._owner._rollback_reservation(self)
|
||||
|
||||
|
||||
class FixedFramePcmAppendBuffer(PcmAppendBuffer):
|
||||
"""Frame ``pcm_f32le`` at one sample rate into ``frame_samples`` units, transactionally."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
sample_rate_hz: int,
|
||||
frame_samples: int,
|
||||
chunk_period_ms: int,
|
||||
model: str = "duplex",
|
||||
) -> None:
|
||||
if sample_rate_hz <= 0 or frame_samples <= 0 or chunk_period_ms <= 0:
|
||||
raise ValueError("FixedFramePcmAppendBuffer needs positive sample_rate_hz, frame_samples, chunk_period_ms")
|
||||
self.sample_rate_hz = int(sample_rate_hz)
|
||||
self.frame_samples = int(frame_samples)
|
||||
self.chunk_period_ms = int(chunk_period_ms)
|
||||
self.model = model
|
||||
self._buffer = bytearray()
|
||||
self._reservations: list[FixedFramePcmAppendReservation] = []
|
||||
|
||||
@property
|
||||
def frame_bytes(self) -> int:
|
||||
return self.frame_samples * PCM_F32LE_BYTES_PER_SAMPLE
|
||||
|
||||
@property
|
||||
def pending_byte_count(self) -> int:
|
||||
return len(self._buffer)
|
||||
|
||||
def has_pending(self) -> bool:
|
||||
return bool(self._buffer)
|
||||
|
||||
def has_reserved(self) -> bool:
|
||||
return any(reservation.active for reservation in self._reservations)
|
||||
|
||||
def clear_force_listen(self) -> None:
|
||||
return
|
||||
|
||||
def clear(self) -> None:
|
||||
for reservation in self._reservations:
|
||||
reservation._active = False
|
||||
self._reservations.clear()
|
||||
self._buffer.clear()
|
||||
|
||||
def prepare_append(
|
||||
self,
|
||||
payload: dict[str, object],
|
||||
*,
|
||||
operation_id: str,
|
||||
chunk_period_ms: int,
|
||||
allow_emit: bool,
|
||||
) -> FixedFramePcmAppendReservation | None:
|
||||
self._require_chunk_period(chunk_period_ms)
|
||||
if any(reservation.active and reservation.operation_id == operation_id for reservation in self._reservations):
|
||||
raise ValueError(f"{self.model} duplicate active operation_id: {operation_id}")
|
||||
self._buffer.extend(self._decode_payload(payload))
|
||||
if not allow_emit or len(self._buffer) < self.frame_bytes:
|
||||
return None
|
||||
return self._reserve_frame(payload, operation_id=operation_id, flush=False)
|
||||
|
||||
def prepare_commit(self, *, operation_id: str, chunk_period_ms: int) -> FixedFramePcmAppendReservation:
|
||||
self._require_chunk_period(chunk_period_ms)
|
||||
if not self._buffer:
|
||||
reservation = FixedFramePcmAppendReservation(owner=self, operation_id=operation_id, payload=None, raw=b"")
|
||||
self._reservations.append(reservation)
|
||||
return reservation
|
||||
payload: dict[str, object] = {
|
||||
"type": "audio",
|
||||
"format": "pcm_f32le",
|
||||
"sample_rate_hz": self.sample_rate_hz,
|
||||
"audio": "",
|
||||
"final": True,
|
||||
}
|
||||
return self._reserve_frame(payload, operation_id=operation_id, flush=True)
|
||||
|
||||
def flush(self, *, chunk_period_ms: int) -> dict[str, object] | None:
|
||||
reservation = self.prepare_commit(operation_id=f"{self.model}-flush", chunk_period_ms=chunk_period_ms)
|
||||
payload = reservation.payload
|
||||
reservation.commit()
|
||||
return payload
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# Internals #
|
||||
# ------------------------------------------------------------------ #
|
||||
|
||||
def _require_chunk_period(self, chunk_period_ms: int) -> None:
|
||||
if chunk_period_ms != self.chunk_period_ms:
|
||||
raise ValueError(f"{self.model} chunk_period_ms must be {self.chunk_period_ms}")
|
||||
|
||||
def _decode_payload(self, payload: dict[str, object]) -> bytes:
|
||||
if payload.get("format") != "pcm_f32le":
|
||||
raise ValueError(f"{self.model} input format must be pcm_f32le")
|
||||
if payload.get("sample_rate_hz") != self.sample_rate_hz:
|
||||
raise ValueError(f"{self.model} sample_rate_hz must be {self.sample_rate_hz}")
|
||||
return decode_pcm_f32le_base64(payload.get("audio"), model=self.model)
|
||||
|
||||
def _reserve_frame(
|
||||
self,
|
||||
source_payload: dict[str, object],
|
||||
*,
|
||||
operation_id: str,
|
||||
flush: bool,
|
||||
) -> FixedFramePcmAppendReservation:
|
||||
frame_bytes = self.frame_bytes
|
||||
consumed_bytes = min(len(self._buffer), frame_bytes) if flush else frame_bytes
|
||||
raw = bytes(self._buffer[:consumed_bytes])
|
||||
del self._buffer[:consumed_bytes]
|
||||
encoded_raw = raw + b"\x00" * (frame_bytes - consumed_bytes)
|
||||
payload = dict(source_payload)
|
||||
payload["type"] = "audio"
|
||||
payload["format"] = "pcm_f32le"
|
||||
payload["sample_rate_hz"] = self.sample_rate_hz
|
||||
payload["audio"] = base64.b64encode(encoded_raw).decode("ascii")
|
||||
reservation = FixedFramePcmAppendReservation(owner=self, operation_id=operation_id, payload=payload, raw=raw)
|
||||
self._reservations.append(reservation)
|
||||
return reservation
|
||||
|
||||
def _commit_reservation(self, reservation: FixedFramePcmAppendReservation) -> None:
|
||||
if not reservation.active:
|
||||
return
|
||||
reservation._active = False
|
||||
self._reservations.remove(reservation)
|
||||
|
||||
def _rollback_reservation(self, reservation: FixedFramePcmAppendReservation) -> None:
|
||||
if not reservation.active:
|
||||
return
|
||||
try:
|
||||
index = self._reservations.index(reservation)
|
||||
except ValueError:
|
||||
reservation._active = False
|
||||
return
|
||||
# Later reservations were taken after this one: they go back too, in order.
|
||||
rolled_back = self._reservations[index:]
|
||||
restored = b"".join(item._raw for item in rolled_back if item.active)
|
||||
self._buffer[:0] = restored
|
||||
for item in rolled_back:
|
||||
item._active = False
|
||||
del self._reservations[index:]
|
||||
|
||||
|
||||
__all__ = ["FixedFramePcmAppendBuffer", "FixedFramePcmAppendReservation"]
|
||||
Reference in new issue
Block a user