561 lines
24 KiB
Python
561 lines
24 KiB
Python
"""Streaming web voice demo for LiquidAI/LFM3-Audio-2B.
|
||
|
||
One WebSocket per browser tab. The browser streams mic PCM (float32, native rate) while the
|
||
talk button is held, and the server encodes and prefills it into the model's KV cache as it
|
||
arrives (see session.py). On release only the last audio token and the assistant prefix remain,
|
||
then generation streams text tokens (JSON) and decoded 24 kHz speech (binary float32 PCM) back
|
||
frame by frame, together with debug events and timings. The conversation stays in the cache
|
||
between turns.
|
||
"""
|
||
|
||
import asyncio
|
||
import json
|
||
import os
|
||
import queue
|
||
import re
|
||
import sys
|
||
import threading
|
||
import time
|
||
from concurrent.futures import ThreadPoolExecutor
|
||
from pathlib import Path
|
||
|
||
import numpy as np
|
||
import torch
|
||
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
|
||
from fastapi.responses import FileResponse
|
||
from transformers import AutoModelForMultimodalLM, AutoProcessor
|
||
from transformers.generation.streamers import BaseStreamer
|
||
|
||
import tools
|
||
from session import StreamingSession
|
||
|
||
MODEL_ID = os.environ.get("LFM_MODEL", "LiquidAI/LFM3-Audio-2B")
|
||
IN_RATE = 16_000
|
||
OUT_RATE = 24_000
|
||
FRAME_SAMPLES = 1920 # one codec frame = 80 ms at 24 kHz
|
||
CHUNK_FRAMES = int(os.environ.get("LFM_CHUNK_FRAMES", "1")) # frames per decoded audio chunk sent to the browser
|
||
MAX_TOOL_ROUNDS = 4
|
||
MIXED_START = "<|mixed_start|>"
|
||
|
||
MODES = {
|
||
"chat": {"system": "Respond with interleaved text and audio.", "prefill": "<|mixed_start|>"},
|
||
"asr": {"system": "Perform ASR.", "prefill": None},
|
||
"tts": {"system": "Perform TTS in English using a female voice.", "prefill": "<|audio_start|>"},
|
||
# no prefill: the model chooses between <tool_call> and <|mixed_start|>
|
||
"tools": {"system": tools.SYSTEM_PROMPT, "prefill": None},
|
||
}
|
||
|
||
|
||
def pick_device() -> tuple[str, torch.dtype]:
|
||
device = os.environ.get("LFM_DEVICE")
|
||
if not device:
|
||
device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
|
||
default_dtype = "float32" if device == "cpu" else "bfloat16"
|
||
return device, getattr(torch, os.environ.get("LFM_DTYPE", default_dtype))
|
||
|
||
|
||
class Engine:
|
||
def __init__(self):
|
||
self.device, self.dtype = pick_device()
|
||
t0 = time.perf_counter()
|
||
self.model = AutoModelForMultimodalLM.from_pretrained(
|
||
MODEL_ID, dtype=self.dtype, device_map=self.device, trust_remote_code=True
|
||
).eval()
|
||
self.processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
|
||
# the codec is ~87x realtime on one CPU thread; keeping it on CPU avoids contending with the LM
|
||
self.codec_device = os.environ.get("LFM_CODEC_DEVICE", "cpu")
|
||
self.codec = self.processor.audio_tokenizer.to(device=self.codec_device, dtype=torch.float32).eval()
|
||
self.load_s = time.perf_counter() - t0
|
||
self.tokenizer = self.processor.tokenizer
|
||
self.lock = threading.Lock() # one generation at a time across all connections
|
||
|
||
def info(self) -> dict:
|
||
return {
|
||
"model": MODEL_ID,
|
||
"device": self.device,
|
||
"dtype": str(self.dtype).removeprefix("torch."),
|
||
"codec_device": self.codec_device,
|
||
"load_s": round(self.load_s, 2),
|
||
"chunk_frames": CHUNK_FRAMES,
|
||
"interleaved_n_text": self.model.config.interleaved_n_text,
|
||
"interleaved_n_audio": self.model.config.interleaved_n_audio,
|
||
"torch": torch.__version__,
|
||
}
|
||
|
||
def build_inputs(self, conversation: list[dict], clips: list[np.ndarray], prefill: str | None, tool_defs=None):
|
||
if prefill:
|
||
conversation = conversation + [{"role": "assistant", "content": prefill}]
|
||
prompt = self.processor.apply_chat_template(
|
||
conversation,
|
||
tokenize=False,
|
||
continue_final_message=bool(prefill),
|
||
add_generation_prompt=not prefill,
|
||
tools=tool_defs,
|
||
)
|
||
enc = self.processor(
|
||
text=prompt,
|
||
audio=clips or None,
|
||
sampling_rate=IN_RATE,
|
||
add_special_tokens=False,
|
||
return_tensors="pt",
|
||
)
|
||
keys = ("input_ids", "attention_mask", "input_features", "feature_attention_mask")
|
||
inputs = {k: enc[k].to(self.model.device) for k in keys if k in enc}
|
||
if "input_features" in inputs:
|
||
inputs["input_features"] = inputs["input_features"].to(self.dtype)
|
||
return prompt, inputs
|
||
|
||
|
||
class Cancelled(Exception):
|
||
pass
|
||
|
||
|
||
class PlainTextStart(Exception):
|
||
"""The reply opened with plain text where a tool call or speech marker was required."""
|
||
|
||
|
||
class Emitter:
|
||
"""Thread-safe bridge from the generation thread to the connection's asyncio send queue."""
|
||
|
||
def __init__(self, loop: asyncio.AbstractEventLoop, queue: asyncio.Queue, t0: float):
|
||
self.loop, self.queue, self.t0 = loop, queue, t0
|
||
|
||
def ms(self) -> float:
|
||
return round((time.perf_counter() - self.t0) * 1000, 1)
|
||
|
||
def send(self, item):
|
||
self.loop.call_soon_threadsafe(self.queue.put_nowait, item)
|
||
|
||
|
||
class TextStream(BaseStreamer):
|
||
def __init__(self, engine: Engine, emit: Emitter, cancel: threading.Event, stats: dict, first_allowed=None):
|
||
self.tok, self.emit, self.cancel, self.stats = engine.tokenizer, emit, cancel, stats
|
||
self.first_allowed = first_allowed # token ids the reply may open with; None allows anything
|
||
self.special = set(self.tok.all_special_ids)
|
||
self.ids: list[int] = []
|
||
self.text = "" # everything generated, tool-call markup included
|
||
self.visible = "" # what the chat bubble shows
|
||
self.prompt_seen = False
|
||
|
||
def put(self, value):
|
||
if self.cancel.is_set():
|
||
raise Cancelled()
|
||
if not self.prompt_seen: # generate() first hands the streamer the prompt
|
||
self.prompt_seen = True
|
||
return
|
||
tid = int(value.reshape(-1)[-1])
|
||
if self.first_allowed is not None and not self.ids and tid not in self.first_allowed:
|
||
raise PlainTextStart()
|
||
self.first_allowed = None
|
||
t = self.emit.ms()
|
||
self.stats.setdefault("first_text_ms", t)
|
||
self.stats["text_tokens"] = self.stats.get("text_tokens", 0) + 1
|
||
raw = self.tok.decode([tid], skip_special_tokens=False)
|
||
delta = ""
|
||
if tid not in self.special:
|
||
self.ids.append(tid)
|
||
full = self.tok.decode(self.ids, skip_special_tokens=True)
|
||
if not full.endswith("�"): # wait for the rest of a split multi-byte char
|
||
self.text = full
|
||
visible = tools.visible_text(full)
|
||
if visible.startswith(self.visible):
|
||
delta, self.visible = visible[len(self.visible):], visible
|
||
self.emit.send({"type": "token", "id": tid, "raw": raw, "delta": delta, "special": tid in self.special, "t_ms": t})
|
||
|
||
def end(self):
|
||
pass
|
||
|
||
|
||
class AudioStream(BaseStreamer):
|
||
"""Collects speech frames and decodes them incrementally with the codec's streaming cache.
|
||
|
||
Decoding runs on its own thread so the codec never stalls the language model's sampling loop.
|
||
"""
|
||
|
||
def __init__(self, engine: Engine, emit: Emitter, cancel: threading.Event, stats: dict):
|
||
self.engine, self.emit, self.cancel, self.stats = engine, emit, cancel, stats
|
||
# the cache class lives in the codec's remote-code module
|
||
codec_module = sys.modules[type(engine.codec).__module__]
|
||
self.cache = codec_module.LiquidCodecStreamingCache(engine.codec.config)
|
||
self.frames = 0
|
||
self.decode_s = 0.0
|
||
self.max_backlog = 0
|
||
self.collected: list[np.ndarray] = []
|
||
self.inbox: queue.SimpleQueue = queue.SimpleQueue()
|
||
self.worker = threading.Thread(target=self._decode_loop, daemon=True)
|
||
self.worker.start()
|
||
|
||
def put(self, codes):
|
||
if self.cancel.is_set():
|
||
raise Cancelled()
|
||
self.stats.setdefault("first_frame_ms", self.emit.ms())
|
||
self.inbox.put(codes.detach().to("cpu"))
|
||
self.frames += 1
|
||
|
||
def end(self):
|
||
self.inbox.put(None)
|
||
self.worker.join()
|
||
self.stats["audio_frames"] = self.frames
|
||
self.stats["codec_decode_ms"] = round(self.decode_s * 1000, 1)
|
||
self.stats["max_decode_backlog"] = self.max_backlog
|
||
|
||
def _decode_loop(self):
|
||
pending: list[torch.Tensor] = []
|
||
done = False
|
||
while not done:
|
||
item = self.inbox.get()
|
||
# drain whatever else is ready so a slow decode catches up in one bigger chunk
|
||
items = [item]
|
||
while not self.inbox.empty():
|
||
items.append(self.inbox.get())
|
||
done = items[-1] is None
|
||
pending.extend(i for i in items if i is not None)
|
||
self.max_backlog = max(self.max_backlog, len(pending))
|
||
if done:
|
||
if self.cache.started or pending:
|
||
self._flush(pending, finalize=True)
|
||
elif len(pending) >= CHUNK_FRAMES:
|
||
self._flush(pending, finalize=False)
|
||
pending = []
|
||
|
||
@torch.no_grad()
|
||
def _flush(self, pending: list[torch.Tensor], finalize: bool):
|
||
n = len(pending)
|
||
codes = (
|
||
torch.stack(pending, dim=-1)[None].to(self.engine.codec_device)
|
||
if n
|
||
else torch.zeros((1, self.engine.codec.config.n_q, 0), dtype=torch.long)
|
||
)
|
||
first_codes = pending[0].tolist() if n else None
|
||
t = time.perf_counter()
|
||
pcm = self.engine.codec.decode(codes, streaming_cache=self.cache, finalize=finalize).audio_values[0]
|
||
self.decode_s += time.perf_counter() - t
|
||
pcm = pcm.cpu().numpy().astype(np.float32)
|
||
if not pcm.size:
|
||
return
|
||
self.collected.append(pcm)
|
||
t_ms = self.emit.ms()
|
||
self.stats.setdefault("first_audio_ms", t_ms)
|
||
self.emit.send(
|
||
{"type": "audio_chunk", "frames": n, "samples": int(pcm.size), "total_frames": self.frames,
|
||
"codes": first_codes, "final": finalize, "t_ms": t_ms}
|
||
)
|
||
self.emit.send(pcm.tobytes())
|
||
|
||
|
||
app = FastAPI()
|
||
engine: Engine | None = None
|
||
STATIC = Path(__file__).parent / "static"
|
||
|
||
|
||
@app.on_event("startup")
|
||
def load():
|
||
global engine
|
||
engine = Engine()
|
||
print("model ready:", json.dumps(engine.info()))
|
||
|
||
|
||
@app.get("/")
|
||
def index():
|
||
return FileResponse(STATIC / "index.html")
|
||
|
||
|
||
@app.websocket("/ws")
|
||
async def ws_endpoint(ws: WebSocket):
|
||
await ws.accept()
|
||
loop = asyncio.get_running_loop()
|
||
out: asyncio.Queue = asyncio.Queue()
|
||
# every model call for this connection runs on one thread, in arrival order: prefill chunks, then generation
|
||
executor = ThreadPoolExecutor(max_workers=1)
|
||
|
||
def run(fn, *args):
|
||
return loop.run_in_executor(executor, fn, *args)
|
||
|
||
async def sender():
|
||
while True:
|
||
item = await out.get()
|
||
if isinstance(item, (bytes, bytearray)):
|
||
await ws.send_bytes(item)
|
||
else:
|
||
await ws.send_text(json.dumps(item))
|
||
|
||
send_task = asyncio.create_task(sender())
|
||
await out.put({"type": "ready", **engine.info()})
|
||
|
||
session = StreamingSession(engine)
|
||
settings = {"mode": "chat", "max_new_tokens": 1024, "audio_temperature": 1.0, "audio_top_k": 128,
|
||
"system": "", "voice": "female", "keep_history": True, "force_speech": True}
|
||
cancel = threading.Event()
|
||
gen_task: asyncio.Task | None = None
|
||
recording = False
|
||
rec: list[np.ndarray] = []
|
||
rec_rate = IN_RATE
|
||
mic_inbox: queue.SimpleQueue = queue.SimpleQueue() # chunks waiting to be prefilled
|
||
speech = {} # prefill stats for the utterance being recorded
|
||
|
||
def log(msg, **kw):
|
||
out.put_nowait({"type": "log", "msg": msg, **kw})
|
||
|
||
def voice_mode():
|
||
return settings["mode"] if settings["mode"] in ("chat", "asr", "tools") else "chat"
|
||
|
||
def system_text(mode):
|
||
if settings["system"].strip():
|
||
return settings["system"].strip()
|
||
if mode == "tts":
|
||
return f"Perform TTS in English using a {settings['voice']} voice."
|
||
return MODES[mode]["system"]
|
||
|
||
# ---- jobs that run on the executor thread ----------------------------------------------------------------
|
||
|
||
def open_turn(mode):
|
||
"""Close the previous turn and open a user turn in the cached context."""
|
||
sys_txt = system_text(mode)
|
||
block = engine.processor.apply_chat_template(
|
||
[{"role": "system", "content": sys_txt}], tools=tools.TOOLS if mode == "tools" else None, tokenize=False
|
||
)
|
||
with engine.lock:
|
||
t = time.perf_counter()
|
||
info = session.begin_turn((mode, sys_txt), block, settings["keep_history"] and mode != "asr")
|
||
info["open_ms"] = round((time.perf_counter() - t) * 1000, 1)
|
||
if "context_reset" in info:
|
||
emit_now({"type": "log", "msg": f"context reset ({info['context_reset']})"})
|
||
speech.clear()
|
||
speech.update(info, prefill_calls=0, prefill_busy_ms=0.0)
|
||
|
||
def feed_pending(final=False):
|
||
chunks = []
|
||
while not mic_inbox.empty():
|
||
chunks.append(mic_inbox.get())
|
||
if not chunks and not final:
|
||
return
|
||
pcm = np.concatenate(chunks) if chunks else np.zeros(0, np.float32)
|
||
t = time.perf_counter()
|
||
with engine.lock:
|
||
session.feed_audio(pcm, rec_rate, final=final)
|
||
ms = (time.perf_counter() - t) * 1000
|
||
if final:
|
||
speech["flush_ms"] = round(ms, 1)
|
||
else:
|
||
speech["prefill_calls"] += 1
|
||
speech["prefill_busy_ms"] += ms
|
||
|
||
def emit_now(item):
|
||
loop.call_soon_threadsafe(out.put_nowait, item)
|
||
|
||
def run_round(gen, prompt_text, prompt_tokens, mode, emit, round_no, first_allowed=None):
|
||
stats: dict = {"prompt_tokens": prompt_tokens}
|
||
emit.send({"type": "prompt", "text": prompt_text, "tokens": prompt_tokens, "round": round_no, "t_ms": emit.ms()})
|
||
text_stream = TextStream(engine, emit, cancel, stats, first_allowed)
|
||
audio_stream = AudioStream(engine, emit, cancel, stats) if mode != "asr" else None
|
||
kwargs = dict(
|
||
max_new_tokens=int(settings["max_new_tokens"]),
|
||
audio_temperature=float(settings["audio_temperature"]),
|
||
audio_top_k=int(settings["audio_top_k"]),
|
||
streamer=text_stream,
|
||
audio_streamer=audio_stream,
|
||
)
|
||
if mode == "asr":
|
||
kwargs.update(generation_mode="text")
|
||
stats["cancelled"] = False
|
||
with engine.lock, torch.inference_mode():
|
||
gen_t0 = time.perf_counter()
|
||
try:
|
||
gen(**kwargs)
|
||
except Cancelled:
|
||
stats["cancelled"] = True
|
||
except PlainTextStart:
|
||
stats["plain_start"] = True
|
||
stats["gen_ms"] = round((time.perf_counter() - gen_t0) * 1000, 1)
|
||
audio_s = sum(c.size for c in audio_stream.collected) / OUT_RATE if audio_stream else 0.0
|
||
stats["audio_out_s"] = round(audio_s, 2)
|
||
return text_stream, stats
|
||
|
||
def session_round(suffix, mode, emit, round_no, first_allowed=None):
|
||
cached = session.length
|
||
n = len(engine.tokenizer(suffix, add_special_tokens=False).input_ids)
|
||
shown = f"[{cached} cached tokens]" + suffix
|
||
return run_round(lambda **kw: session.generate(suffix, **kw), shown, cached + n, mode, emit, round_no, first_allowed)
|
||
|
||
def run_session_turn(user_text, mode, t0, user_meta):
|
||
"""Generate the reply to the open user turn; in tools mode, run calls and continue until the model answers."""
|
||
emit = Emitter(loop, out, t0)
|
||
if user_text is None: # spoken: flush the encoder's last token
|
||
feed_pending(final=True)
|
||
open_suffix = (user_text or "") + "<|im_end|>\n<|im_start|>assistant\n"
|
||
prefill = MODES[mode]["prefill"] or ""
|
||
force_speech = mode == "tools" and settings["force_speech"]
|
||
speech_markers = {engine.tokenizer.convert_tokens_to_ids(t) for t in ("<tool_call>", "<|mixed_start|>", "<|audio_start|>")}
|
||
rounds: list[dict] = []
|
||
suffix = open_suffix + prefill
|
||
try:
|
||
for round_no in range(MAX_TOOL_ROUNDS + 1):
|
||
# the model sometimes answers in plain text only; force speech by prefilling the interleaved marker
|
||
# after a tool result, and by restarting a first round that opens with plain text
|
||
if force_speech and round_no:
|
||
suffix += MIXED_START
|
||
guard = speech_markers if force_speech and not round_no else None
|
||
snap = session.snapshot()
|
||
text_stream, r = session_round(suffix, mode, emit, round_no, guard)
|
||
if r.get("plain_start"):
|
||
emit.send({"type": "log", "msg": f"round {round_no} opened with plain text; restarting with speech"})
|
||
rounds.append(r)
|
||
session.restore(snap)
|
||
text_stream, r = session_round(suffix + MIXED_START, mode, emit, round_no)
|
||
rounds.append(r)
|
||
if r["cancelled"]:
|
||
break
|
||
calls = tools.parse_tool_calls(text_stream.text) if mode == "tools" else []
|
||
if not calls or round_no == MAX_TOOL_ROUNDS:
|
||
break
|
||
results = []
|
||
for call in calls:
|
||
fn = call["function"]
|
||
emit.send({"type": "tool_call", "name": fn["name"], "arguments": fn["arguments"], "round": round_no,
|
||
"t_ms": emit.ms()})
|
||
t = time.perf_counter()
|
||
result = tools.run_tool(fn["name"], fn["arguments"],
|
||
lambda action, value: emit.send({"type": "ui", "action": action, "value": value}))
|
||
results.append(result)
|
||
emit.send({"type": "tool_result", "name": fn["name"], "result": result,
|
||
"ms": round((time.perf_counter() - t) * 1000, 1), "t_ms": emit.ms()})
|
||
suffix = (session.close_prefix() + "<|im_start|>tool\n<tool_response>\n" + json.dumps(results)
|
||
+ "\n</tool_response><|im_end|>\n<|im_start|>assistant\n")
|
||
except BaseException:
|
||
session.restore(session.turn_snapshot)
|
||
raise
|
||
if rounds[-1]["cancelled"]:
|
||
session.restore(session.turn_snapshot) # a cancelled turn leaves no trace in the context
|
||
stats = summarize(rounds, mode, emit, user_meta)
|
||
stats.update({k: (round(v, 1) if isinstance(v, float) else v) for k, v in speech.items()},
|
||
context_tokens_after=session.length)
|
||
return text_stream.visible, stats
|
||
|
||
def run_tts(text, t0):
|
||
emit = Emitter(loop, out, t0)
|
||
conv = [{"role": "system", "content": system_text("tts")}, {"role": "user", "content": text}]
|
||
prompt, inputs = engine.build_inputs(conv, [], MODES["tts"]["prefill"])
|
||
text_stream, r = run_round(lambda **kw: engine.model.generate(**inputs, return_dict_in_generate=True, **kw),
|
||
prompt, int(inputs["input_ids"].shape[1]), "tts", emit, 0)
|
||
return text_stream.visible, summarize([r], "tts", emit, {})
|
||
|
||
# ---- event-loop side -------------------------------------------------------------------------------------
|
||
|
||
async def generate(job, *args):
|
||
cancel.clear()
|
||
try:
|
||
text, stats = await run(job, *args)
|
||
await out.put({"type": "done", "text": text, "stats": stats})
|
||
except Exception as e: # surface model errors in the UI instead of killing the socket
|
||
import traceback
|
||
|
||
traceback.print_exc()
|
||
await out.put({"type": "error", "msg": f"{type(e).__name__}: {e}"})
|
||
|
||
try:
|
||
while True:
|
||
msg = await ws.receive()
|
||
if msg["type"] == "websocket.disconnect":
|
||
break
|
||
if msg.get("bytes") is not None:
|
||
if recording:
|
||
chunk = np.frombuffer(msg["bytes"], dtype=np.float32)
|
||
rec.append(chunk)
|
||
mic_inbox.put(chunk)
|
||
run(feed_pending) # queued behind open_turn; drains whatever has arrived by then
|
||
continue
|
||
data = json.loads(msg["text"])
|
||
kind = data.get("type")
|
||
busy = gen_task is not None and not gen_task.done()
|
||
|
||
if kind == "settings":
|
||
settings.update({k: v for k, v in data.items() if k in settings})
|
||
log("settings updated", settings=settings)
|
||
elif kind == "start":
|
||
if busy:
|
||
log("busy: still generating, utterance ignored")
|
||
continue
|
||
recording, rec, rec_rate = True, [], int(data.get("rate", IN_RATE))
|
||
run(open_turn, voice_mode())
|
||
log(f"recording @ {rec_rate} Hz, prefilling while you speak")
|
||
elif kind == "stop":
|
||
if not recording:
|
||
continue
|
||
recording = False
|
||
t0 = time.perf_counter()
|
||
audio = np.concatenate(rec) if rec else np.zeros(0, np.float32)
|
||
dur = audio.size / rec_rate
|
||
rms = float(np.sqrt(np.mean(audio**2))) if audio.size else 0.0
|
||
peak = float(np.abs(audio).max()) if audio.size else 0.0
|
||
user_meta = {"user_audio_s": round(dur, 2), "user_rms": round(rms, 4), "user_peak": round(peak, 3)}
|
||
await out.put({"type": "user_audio", **user_meta})
|
||
if dur < 0.3:
|
||
log("utterance too short, ignored")
|
||
run(lambda: session.restore(session.turn_snapshot))
|
||
continue
|
||
gen_task = asyncio.create_task(generate(run_session_turn, None, voice_mode(), t0, user_meta))
|
||
elif kind == "text":
|
||
if busy or recording:
|
||
log("busy")
|
||
continue
|
||
t0 = time.perf_counter()
|
||
run(open_turn, voice_mode())
|
||
gen_task = asyncio.create_task(generate(run_session_turn, data["text"], voice_mode(), t0, {}))
|
||
elif kind == "tts":
|
||
if busy or recording:
|
||
log("busy")
|
||
continue
|
||
gen_task = asyncio.create_task(generate(run_tts, data["text"], time.perf_counter()))
|
||
elif kind == "cancel":
|
||
cancel.set()
|
||
log("cancel requested")
|
||
elif kind == "reset":
|
||
run(session.reset)
|
||
log("context cleared")
|
||
except WebSocketDisconnect:
|
||
pass
|
||
finally:
|
||
cancel.set()
|
||
send_task.cancel()
|
||
executor.shutdown(wait=False, cancel_futures=True)
|
||
|
||
|
||
def summarize(rounds: list[dict], mode: str, emit: Emitter, user_meta: dict) -> dict:
|
||
gen_s = sum(r["gen_ms"] for r in rounds) / 1000
|
||
audio_s = sum(r["audio_out_s"] for r in rounds)
|
||
total = lambda key: sum(r.get(key, 0) for r in rounds) # noqa: E731
|
||
first = lambda key: next((r[key] for r in rounds if key in r), None) # noqa: E731
|
||
return {
|
||
"mode": mode, **user_meta,
|
||
"rounds": len(rounds),
|
||
"speech_restarts": sum(1 for r in rounds if r.get("plain_start")),
|
||
"prompt_tokens": rounds[-1]["prompt_tokens"],
|
||
"first_text_ms": first("first_text_ms"),
|
||
"first_frame_ms": first("first_frame_ms"),
|
||
"first_audio_ms": first("first_audio_ms"),
|
||
"text_tokens": total("text_tokens"),
|
||
"audio_frames": total("audio_frames"),
|
||
"codec_decode_ms": round(total("codec_decode_ms"), 1),
|
||
"max_decode_backlog": max((r.get("max_decode_backlog", 0) for r in rounds), default=0),
|
||
"cancelled": rounds[-1]["cancelled"],
|
||
"gen_ms": round(gen_s * 1000, 1),
|
||
"total_ms": emit.ms(),
|
||
"audio_out_s": round(audio_s, 2),
|
||
"rtf": round(gen_s / audio_s, 3) if audio_s else None,
|
||
"text_tok_per_s": round(total("text_tokens") / gen_s, 1) if gen_s else None,
|
||
"frames_per_s": round(total("audio_frames") / gen_s, 1) if gen_s else None,
|
||
"per_round": rounds,
|
||
}
|
||
|
||
|
||
def collapse_placeholders(prompt: str) -> str:
|
||
"""Show runs of audio placeholders as <|audio|>×N so the prompt stays readable."""
|
||
|
||
return re.sub(r"((<\|audio(?:_out)?\|>))\1+", lambda m: f"{m.group(2)}×{len(m.group(0)) // len(m.group(2))}", prompt)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
import uvicorn
|
||
|
||
uvicorn.run(app, host=os.environ.get("HOST", "127.0.0.1"), port=int(os.environ.get("PORT", "8000")))
|