first tests

.gitignore

@@ -0,0 +1,3 @@
+.venv/
+__pycache__/
+smoke_out/

.python-version

@@ -0,0 +1 @@
+3.12

pyproject.toml

@@ -0,0 +1,18 @@
+[project]
+name = "lfm-voice"
+version = "0.1.0"
+description = "Streaming web voice demo for LiquidAI/LFM3-Audio-2B"
+requires-python = ">=3.11"
+dependencies = [
+    "transformers>=5.20.2",
+    "torch>=2.6",
+    "accelerate",
+    "librosa",
+    "soundfile",
+    "numpy",
+    "fastapi",
+    "uvicorn[standard]",
+]
+
+[tool.uv]
+package = false

scripts/session_test.py

@@ -0,0 +1,92 @@
+"""Check StreamingSession: streamed ASR must match offline ASR; two-turn chat from the cache; latency after release."""
+
+import sys
+import time
+from pathlib import Path
+
+import librosa
+import numpy as np
+import soundfile as sf
+import torch
+
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+import server  # noqa: E402
+from session import StreamingSession  # noqa: E402
+
+engine = server.Engine()
+tok = engine.tokenizer
+wav, rate = sf.read(Path(__file__).resolve().parent.parent / "smoke_out/tts.wav", dtype="float32")
+CHUNK = 2048  # what the browser worklet sends
+
+
+def system_block(text):
+    return engine.processor.apply_chat_template([{"role": "system", "content": text}], tokenize=False)
+
+
+def stream_in(s, audio, rate):
+    for i in range(0, len(audio), CHUNK):
+        s.feed_audio(audio[i:i + CHUNK], rate)
+
+
+def decode_text(out, suffix_len):
+    seq = out.sequences[0, suffix_len:].tolist()
+    return tok.decode([t for t in seq if t != engine.model.config.audio_out_token_id], skip_special_tokens=True)
+
+
+# offline reference
+_, inputs = engine.build_inputs([{"role": "system", "content": "Perform ASR."}, {"role": "user", "content": [{"type": "audio"}]}],
+                                [librosa.resample(wav, orig_sr=rate, target_sr=16000)], None)
+with torch.inference_mode():
+    ref = engine.model.generate(**inputs, max_new_tokens=64, generation_mode="text", return_dict_in_generate=True)
+print("offline ASR :", tok.decode(ref.sequences[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True))
+
+s = StreamingSession(engine)
+s.begin_turn(("asr", ""), system_block("Perform ASR."), keep_history=False)
+stream_in(s, wav, rate)
+t = time.perf_counter()
+s.feed_audio(np.zeros(0, np.float32), rate, final=True)
+out = s.generate("<|im_end|>\n<|im_start|>assistant\n", max_new_tokens=64, generation_mode="text")
+print(f"streamed ASR: {decode_text(out, 5)}  ({(time.perf_counter() - t) * 1000:.0f} ms after release, {s.audio_tokens} audio tokens)")
+
+
+class FirstFrame:
+    def __init__(self): self.t = None
+    def put(self, v): self.t = self.t or time.perf_counter()
+    def end(self): pass
+
+
+sysb = system_block("Respond with interleaved text and audio.")
+conv = [{"role": "system", "content": "Respond with interleaved text and audio."}, {"role": "user", "content": [{"type": "audio"}]}]
+w16 = librosa.resample(wav, orig_sr=rate, target_sr=16000)
+for label in ("warmup", "offline"):
+    ff = None
+
+    class FF:
+        t = None
+        def put(self, v): self.t = self.t or time.perf_counter()
+        def end(self): pass
+    ff = FF()
+    t = time.perf_counter()
+    _, inputs = engine.build_inputs(conv, [w16], "<|mixed_start|>")
+    with torch.inference_mode():
+        engine.model.generate(**inputs, max_new_tokens=512, audio_streamer=ff, return_dict_in_generate=True)
+    print(f"{label} chat: first_frame={(ff.t - t) * 1000:.0f}ms after release (prep + full prefill)")
+
+s = StreamingSession(engine)
+for turn in range(3):
+    info = s.begin_turn(("chat", ""), sysb, keep_history=True)
+    stream_in(s, wav, rate)
+    t = time.perf_counter()
+    s.feed_audio(np.zeros(0, np.float32), rate, final=True)
+    t_flush = time.perf_counter()
+    ff = FirstFrame()
+    out = s.generate("<|im_end|>\n<|im_start|>assistant\n<|mixed_start|>", max_new_tokens=512, audio_streamer=ff)
+    print(f"chat turn {turn}: {info} flush={(t_flush - t) * 1000:.0f}ms first_frame={(ff.t - t) * 1000:.0f}ms "
+          f"total={(time.perf_counter() - t):.1f}s ctx={s.length} close={s.needs_close}\n  {decode_text(out, 6)!r}")
+
+# snapshot / restore round trip
+snap = s.snapshot()
+n = s.length
+s.begin_turn(("chat", ""), sysb, keep_history=True)
+s.restore(snap)
+print("restore ok:", s.length == n)

scripts/smoke_test.py

@@ -0,0 +1,52 @@
+"""Headless check of the server's generation path: TTS -> ASR -> voice chat, with streaming."""
+
+import asyncio
+import sys
+import threading
+import time
+from pathlib import Path
+
+import librosa
+import soundfile as sf
+import torch
+
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+import server  # noqa: E402
+
+OUT = Path(__file__).resolve().parent.parent / "smoke_out"
+OUT.mkdir(exist_ok=True)
+
+
+def run(engine, conv, clips, prefill, mode):
+    loop = asyncio.new_event_loop()
+    queue = asyncio.Queue()
+    stats, cancel = {}, threading.Event()
+    emit = server.Emitter(loop, queue, time.perf_counter())
+    emit.send = lambda item: None
+    _, inputs = engine.build_inputs(conv, clips, prefill)
+    text = server.TextStream(engine, emit, cancel, stats)
+    audio = server.AudioStream(engine, emit, cancel, stats) if mode != "asr" else None
+    t = time.perf_counter()
+    with torch.inference_mode():
+        engine.model.generate(**inputs, max_new_tokens=512, streamer=text, audio_streamer=audio,
+                              return_dict_in_generate=True, **({"generation_mode": "text"} if mode == "asr" else {}))
+    gen = time.perf_counter() - t
+    pcm = torch.cat([torch.from_numpy(c) for c in audio.collected]).numpy() if audio and audio.collected else None
+    dur = pcm.size / server.OUT_RATE if pcm is not None else 0
+    print(f"[{mode}] prompt={inputs['input_ids'].shape[1]} gen={gen:.2f}s audio={dur:.2f}s stats={stats}\n  text={text.text!r}")
+    return text.text, pcm
+
+
+engine = server.Engine()
+print(engine.info())
+_, pcm = run(engine, [{"role": "system", "content": "Perform TTS in English using a female voice."},
+                      {"role": "user", "content": "What is the tallest mountain in Europe?"}], [], "<|audio_start|>", "tts")
+sf.write(OUT / "tts.wav", pcm, server.OUT_RATE)
+q16 = librosa.resample(pcm, orig_sr=server.OUT_RATE, target_sr=server.IN_RATE)
+run(engine, [{"role": "system", "content": "Perform ASR."}, {"role": "user", "content": [{"type": "audio"}]}], [q16], None, "asr")
+reply, pcm = run(engine, [{"role": "system", "content": "Respond with interleaved text and audio."},
+                          {"role": "user", "content": [{"type": "audio"}]}], [q16], "<|mixed_start|>", "chat")
+sf.write(OUT / "chat.wav", pcm, server.OUT_RATE)
+# sanity-check the chat reply's speech by transcribing it
+r16 = librosa.resample(pcm, orig_sr=server.OUT_RATE, target_sr=server.IN_RATE)
+run(engine, [{"role": "system", "content": "Perform ASR."}, {"role": "user", "content": [{"type": "audio"}]}], [r16], None, "asr")

scripts/tools_test.py

@@ -0,0 +1,64 @@
+"""Exercise "Voice + tools" over the WebSocket: typed requests, then a spoken one (made with the server's TTS)."""
+
+import asyncio
+import json
+
+import numpy as np
+import websockets
+
+URL = "ws://127.0.0.1:8000/ws"
+PROMPTS = [
+    "What's the weather like in Paris right now?",
+    "What is 17 times 23, divided by 4?",
+    "Please make the screen green.",
+    "What time is it in Tokyo?",
+    "Tell me a fun fact about octopuses.",  # should answer without a tool
+]
+
+
+async def turn(ws, send):
+    await send()
+    pcm, events = [], []
+    while True:
+        m = await ws.recv()
+        if isinstance(m, bytes):
+            pcm.append(np.frombuffer(m, dtype=np.float32))
+            continue
+        m = json.loads(m)
+        if m["type"] in ("tool_call", "tool_result", "ui"):
+            events.append({k: v for k, v in m.items() if k not in ("t_ms",)})
+        if m["type"] in ("done", "error"):
+            return m, events, np.concatenate(pcm) if pcm else np.zeros(0, np.float32)
+
+
+async def main():
+    async with websockets.connect(URL, max_size=None) as ws:
+        await ws.recv()
+        await ws.send(json.dumps({"type": "settings", "mode": "tools", "keep_history": False}))
+        for p in PROMPTS:
+            done, events, pcm = await turn(ws, lambda: ws.send(json.dumps({"type": "text", "text": p})))
+            print(f"\n> {p}")
+            for e in events:
+                print("  ", e)
+            s = done.get("stats", {})
+            print(f"  reply: {done.get('text')!r}\n  rounds={s.get('rounds')} audio={s.get('audio_out_s')}s "
+                  f"first_text={s.get('first_text_ms')}ms gen={s.get('gen_ms')}ms")
+
+        # spoken request: synthesize it with TTS, then send it as mic audio
+        q = "What's the weather in Stockholm?"
+        _, _, speech = await turn(ws, lambda: ws.send(json.dumps({"type": "tts", "text": q})))
+
+        async def speak():
+            await ws.send(json.dumps({"type": "start", "rate": 24000}))
+            for i in range(0, len(speech), 2048):
+                await ws.send(speech[i:i + 2048].tobytes())
+            await ws.send(json.dumps({"type": "stop"}))
+
+        done, events, _ = await turn(ws, speak)
+        print(f"\n> (spoken) {q}")
+        for e in events:
+            print("  ", e)
+        print(f"  reply: {done.get('text')!r} rounds={done['stats'].get('rounds')}")
+
+
+asyncio.run(main())

scripts/ws_client_test.py

@@ -0,0 +1,42 @@
+"""Drive the server's WebSocket like the browser does: stream a wav in real time in 2048-sample chunks, 3 turns."""
+
+import asyncio
+import json
+import sys
+
+import numpy as np
+import soundfile as sf
+import websockets
+
+
+async def main(path, url="ws://127.0.0.1:8000/ws"):
+    audio, rate = sf.read(path, dtype="float32")
+    async with websockets.connect(url, max_size=None) as ws:
+        print("<", json.loads(await ws.recv())["type"])
+        for turn in range(3):  # later turns exercise the cached history
+            await ws.send(json.dumps({"type": "start", "rate": rate}))
+            for i in range(0, len(audio), 2048):
+                await ws.send(audio[i:i + 2048].tobytes())
+                await asyncio.sleep(2048 / rate)  # real-time pacing, like a mic
+            await ws.send(json.dumps({"type": "stop"}))
+            pcm_bytes = 0
+            while True:
+                msg = await ws.recv()
+                if isinstance(msg, bytes):
+                    pcm_bytes += len(msg)
+                    continue
+                m = json.loads(msg)
+                if m["type"] in ("token", "audio_chunk", "log"):
+                    continue
+                if m["type"] == "done":
+                    keys = ("first_text_ms", "first_frame_ms", "first_audio_ms", "flush_ms", "prefill_calls",
+                            "prefill_busy_ms", "context_tokens", "context_tokens_after", "rtf", "gen_ms", "codec_decode_ms", "audio_out_s")
+                    print("< done", m["text"][:80], {k: m["stats"].get(k) for k in keys})
+                elif m["type"] != "prompt":
+                    print("<", m["type"], {k: v for k, v in m.items() if k != "type"})
+                if m["type"] in ("done", "error"):
+                    break
+            print(f"  received {pcm_bytes / 4 / 24000:.2f}s of PCM")
+
+
+asyncio.run(main(sys.argv[1]))

server.py

@@ -0,0 +1,560 @@
+"""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")))

session.py

@@ -0,0 +1,139 @@
+"""A conversation held in the model's KV cache, with audio prefilled while the user is still speaking.
+
+The cache always ends right after a complete chat-template segment, so each step only appends:
+
+    begin_turn():  [<bos>system…<|im_end|>\\n]  <|im_start|>user\\n
+    feed_audio():  audio tokens, encoded and prefilled as the mic chunks arrive
+    generate():    <|im_end|>\\n<|im_start|>assistant\\n<|mixed_start|>  → generated text + speech frames
+
+The model's own generated tokens (including its speech frames) stay in the cache, so the next turn starts from
+there instead of re-encoding the whole conversation.
+"""
+
+import copy
+import time
+
+import numpy as np
+import soxr
+import torch
+
+MAX_CONTEXT = 16_000  # tokens; past this the conversation is reset at the next turn
+
+
+class StreamingSession:
+    def __init__(self, engine):
+        self.engine = engine
+        self.model = engine.model
+        self.tok = engine.tokenizer
+        self.fe = engine.processor.feature_extractor
+        self.im_end = self.tok.convert_tokens_to_ids("<|im_end|>")
+        self.reset()
+
+    # ---- state -------------------------------------------------------------------------------------------------
+
+    def reset(self):
+        self.cache = None
+        self.mask = None
+        self.key = None  # (mode, system prompt): a change starts a fresh context
+        self.needs_close = False  # last generation stopped without <|im_end|>
+        self.turn_snapshot = None
+        self._end_audio()
+
+    def _end_audio(self):
+        self.resampler = None
+        self.feature_cache = None
+        self.encoder_cache = None
+        self.audio_in_s = 0.0
+        self.audio_tokens = 0
+
+    @property
+    def length(self) -> int:
+        return 0 if self.cache is None else self.cache.get_seq_length()
+
+    def snapshot(self):
+        return copy.deepcopy((self.cache, self.mask, self.needs_close))
+
+    def restore(self, snap):
+        self.cache, self.mask, self.needs_close = snap
+        self._end_audio()
+
+    # ---- appending -----------------------------------------------------------------------------------------------
+
+    def _append_embeds(self, embeds: torch.Tensor):
+        n = embeds.shape[1]
+        if self.cache is None:
+            from transformers import DynamicCache
+
+            self.cache = DynamicCache(config=self.model.config.text_config)
+            self.mask = torch.ones((1, 0), dtype=torch.long, device=self.model.device)
+        self.mask = torch.cat((self.mask, self.mask.new_ones((1, n))), dim=-1)
+        position_ids = self.mask.cumsum(-1).sub(1)[:, -n:]
+        self.model.model.language_model(
+            inputs_embeds=embeds,
+            past_key_values=self.cache,
+            attention_mask=self.mask,
+            position_ids=position_ids,
+            use_cache=True,
+        )
+
+    def append_text(self, text: str):
+        ids = self.tok(text, add_special_tokens=False, return_tensors="pt").input_ids.to(self.model.device)
+        self._append_embeds(self.model.get_input_embeddings()(ids))
+
+    def close_prefix(self) -> str:
+        """Text that closes the previous assistant turn, if generation left it open."""
+        return "<|im_end|>\n" if self.needs_close else "\n" if self.cache is not None else ""
+
+    # ---- a turn --------------------------------------------------------------------------------------------------
+
+    @torch.inference_mode()
+    def begin_turn(self, key, system_block: str, keep_history: bool) -> dict:
+        """Open a user turn. `system_block` is the rendered `<bos><|im_start|>system…<|im_end|>\\n` prefix."""
+        info = {}
+        if not keep_history or key != self.key or self.length > MAX_CONTEXT:
+            if self.cache is not None:
+                info["context_reset"] = "history off" if not keep_history else "settings changed" if key != self.key else "context full"
+            self.reset()
+            self.key = key
+        self._end_audio()
+        t = time.perf_counter()
+        self.turn_snapshot = self.snapshot()
+        info["snapshot_ms"] = round((time.perf_counter() - t) * 1000, 1)
+        prefix = self.close_prefix() if self.cache is not None else system_block
+        self.append_text(prefix + "<|im_start|>user\n")
+        self.needs_close = False
+        info["context_tokens"] = self.length
+        return info
+
+    @torch.inference_mode()
+    def feed_audio(self, pcm: np.ndarray, rate: int, final: bool = False):
+        """Resample, extract mel frames and prefill whatever audio tokens they complete."""
+        if self.resampler is None:
+            self.resampler = soxr.ResampleStream(rate, self.fe.sampling_rate, 1, dtype="float32")
+        pcm16 = self.resampler.resample_chunk(np.asarray(pcm, dtype=np.float32), last=final)
+        self.audio_in_s += len(pcm16) / self.fe.sampling_rate
+        features, self.feature_cache = self.fe.stream(pcm16, self.feature_cache, is_final=final)
+        mel = features["input_features"].to(self.model.device)
+        if not mel.shape[-1] and not final:
+            return
+        out = self.model.prefill_audio(
+            mel, past_key_values=self.cache, encoder_cache=self.encoder_cache, attention_mask=self.mask, is_final=final
+        )
+        self.audio_tokens += out.attention_mask.shape[1] - self.mask.shape[1]
+        self.cache, self.mask, self.encoder_cache = out.past_key_values, out.attention_mask, out.encoder_cache
+
+    @torch.inference_mode()
+    def generate(self, suffix: str, **kwargs):
+        """Append `suffix` and generate from the cache; the cache then covers the whole reply."""
+        ids = self.tok(suffix, add_special_tokens=False, return_tensors="pt").input_ids.to(self.model.device)
+        mask = torch.cat((self.mask, self.mask.new_ones((1, ids.shape[1]))), dim=-1)
+        try:
+            out = self.model.generate(
+                input_ids=ids, attention_mask=mask, past_key_values=self.cache, return_audio_eos=True,
+                return_dict_in_generate=True, **kwargs,
+            )
+        finally:
+            # the cache was advanced in place, whether generation finished or was interrupted
+            self.mask = torch.ones((1, self.length), dtype=torch.long, device=self.model.device)
+        self.needs_close = int(out.sequences[0, -1]) != self.im_end
+        return out

static/index.html

@@ -0,0 +1,404 @@
+<!doctype html>
+<html lang="en">
+<head>
+<meta charset="utf-8">
+<meta name="viewport" content="width=device-width, initial-scale=1">
+<title>LFM Voice Demo</title>
+<style>
+  :root {
+    --bg: #f6f7f9; --panel: #fff; --fg: #1b1f24; --muted: #6b7280; --line: #e3e6ea;
+    --accent: #2f6fed; --rec: #e5484d; --ok: #2e9e5b; --warn: #c27c0e; --mono-bg: #f1f3f5;
+  }
+  @media (prefers-color-scheme: dark) {
+    :root {
+      --bg: #0f1115; --panel: #171a20; --fg: #e6e8eb; --muted: #8b93a1; --line: #2a2f38;
+      --accent: #5b8cff; --rec: #ff6369; --ok: #4cc38a; --warn: #f1a10d; --mono-bg: #11141a;
+    }
+  }
+  * { box-sizing: border-box; }
+  body { margin: 0; background: var(--bg); color: var(--fg); font: 14px/1.45 system-ui, -apple-system, sans-serif; }
+  header { display: flex; align-items: center; gap: 12px; padding: 12px 16px; border-bottom: 1px solid var(--line); background: var(--panel); flex-wrap: wrap; }
+  header h1 { font-size: 16px; margin: 0; }
+  .pill { font-size: 12px; padding: 2px 8px; border-radius: 99px; background: var(--mono-bg); color: var(--muted); border: 1px solid var(--line); }
+  .pill.ok { color: var(--ok); } .pill.bad { color: var(--rec); } .pill.busy { color: var(--warn); }
+  main { display: grid; grid-template-columns: minmax(0, 1fr) minmax(0, 1.1fr); gap: 16px; padding: 16px; max-width: 1500px; margin: 0 auto; }
+  @media (max-width: 900px) { main { grid-template-columns: 1fr; } }
+  section { background: var(--panel); border: 1px solid var(--line); border-radius: 10px; padding: 14px; min-width: 0; }
+  h2 { font-size: 13px; text-transform: uppercase; letter-spacing: .04em; color: var(--muted); margin: 0 0 10px; }
+  .row { display: flex; gap: 8px; align-items: center; flex-wrap: wrap; margin-bottom: 10px; }
+  label { font-size: 12px; color: var(--muted); display: flex; flex-direction: column; gap: 2px; }
+  input, select, textarea, button { font: inherit; color: var(--fg); background: var(--bg); border: 1px solid var(--line); border-radius: 6px; padding: 6px 8px; }
+  input[type=number] { width: 80px; }
+  button { cursor: pointer; background: var(--panel); }
+  button:hover { border-color: var(--accent); }
+  button:disabled { opacity: .5; cursor: default; }
+  #talk { width: 100%; padding: 18px; font-size: 16px; font-weight: 600; background: var(--accent); color: #fff; border: none; user-select: none; touch-action: none; }
+  #talk.rec { background: var(--rec); }
+  #meter { height: 6px; background: var(--mono-bg); border-radius: 3px; overflow: hidden; margin: 8px 0; }
+  #meter div { height: 100%; width: 0; background: var(--ok); transition: width 60ms linear; }
+  #chat { display: flex; flex-direction: column; gap: 8px; max-height: 45vh; overflow-y: auto; padding-right: 4px; }
+  .msg { padding: 8px 10px; border-radius: 8px; max-width: 90%; white-space: pre-wrap; word-wrap: break-word; }
+  .msg.user { align-self: flex-end; background: var(--accent); color: #fff; }
+  .msg.bot { align-self: flex-start; background: var(--mono-bg); border: 1px solid var(--line); }
+  .msg .meta { font-size: 11px; opacity: .75; margin-top: 4px; }
+  .msg audio { display: block; margin-top: 6px; height: 32px; max-width: 100%; }
+  .mono, pre { font: 12px/1.4 ui-monospace, SFMono-Regular, Menlo, monospace; }
+  pre { background: var(--mono-bg); border: 1px solid var(--line); border-radius: 6px; padding: 8px; margin: 0; overflow: auto; white-space: pre-wrap; word-break: break-all; }
+  #stats { display: grid; grid-template-columns: repeat(auto-fill, minmax(130px, 1fr)); gap: 6px; margin-bottom: 12px; }
+  .stat { background: var(--mono-bg); border: 1px solid var(--line); border-radius: 6px; padding: 6px 8px; }
+  .stat b { display: block; font: 600 15px ui-monospace, Menlo, monospace; }
+  .stat span { font-size: 11px; color: var(--muted); }
+  #tokens { max-height: 120px; }
+  #tokens .sp { color: var(--warn); }
+  #timeline { width: 100%; height: 60px; display: block; background: var(--mono-bg); border: 1px solid var(--line); border-radius: 6px; margin-bottom: 12px; }
+  #log { max-height: 220px; }
+  details { margin-top: 10px; } summary { cursor: pointer; color: var(--muted); font-size: 12px; margin-bottom: 6px; }
+  .hint { font-size: 12px; color: var(--muted); }
+  #toolHint { margin-bottom: 8px; }
+  .msg.tool { align-self: flex-start; background: transparent; border: 1px dashed var(--warn); font: 12px/1.4 ui-monospace, Menlo, monospace; }
+  .msg.tool .res { color: var(--muted); margin-top: 4px; }
+  #toolsLog .call { color: var(--warn); }
+</style>
+</head>
+<body>
+<header>
+  <h1>LFM3-Audio-2B voice demo</h1>
+  <span id="conn" class="pill bad">disconnected</span>
+  <span id="state" class="pill">idle</span>
+  <span id="info" class="pill mono">–</span>
+</header>
+<main>
+  <section>
+    <h2>Conversation</h2>
+    <div class="row">
+      <label>Mode
+        <select id="mode">
+          <option value="chat">Voice chat (interleaved)</option>
+          <option value="tools">Voice + tools</option>
+          <option value="asr">ASR (transcribe)</option>
+        </select>
+      </label>
+      <label>Max tokens <input id="max_new_tokens" type="number" value="1024" min="16" max="4096"></label>
+      <label>Audio temp <input id="audio_temperature" type="number" value="1.0" step="0.1" min="0" max="2"></label>
+      <label>Audio top‑k <input id="audio_top_k" type="number" value="128" min="1" max="2048"></label>
+      <label>History <select id="keep_history"><option value="1">keep</option><option value="0">off</option></select></label>
+      <label title="Tools mode: make the final answer spoken even when the model would reply in plain text">Force speech <select id="force_speech"><option value="1">on</option><option value="0">off</option></select></label>
+    </div>
+    <div class="row">
+      <label style="flex:1">System prompt override (blank = mode default)
+        <input id="system" placeholder="Respond with interleaved text and audio.">
+      </label>
+    </div>
+    <button id="talk">Hold to talk (or hold Space)</button>
+    <div id="meter"><div></div></div>
+    <div class="row">
+      <button id="cancel">Cancel</button>
+      <button id="reset">Reset history</button>
+      <button id="clear">Clear view</button>
+      <span class="hint">Audio is streamed while you hold; generation starts on release.</span>
+    </div>
+    <div class="row">
+      <input id="ttsText" style="flex:1" placeholder="Type a message (Enter sends) or text for TTS…">
+      <button id="sendText">Send</button>
+      <select id="voice"><option>female</option><option>male</option></select>
+      <button id="tts">Speak (TTS)</button>
+    </div>
+    <div id="toolHint" class="hint" hidden>Tools: <span class="mono">get_current_time</span>, <span class="mono">get_weather</span> (simulated),
+      <span class="mono">calculate</span>, <span class="mono">set_theme_color</span>. Try “What's the weather in Paris?”,
+      “What's 17 times 23?” or “Make the screen green.”</div>
+    <div id="chat"></div>
+  </section>
+
+  <section>
+    <h2>Debug — last turn</h2>
+    <div id="stats"></div>
+    <canvas id="timeline" width="800" height="60"></canvas>
+    <h2>Token stream <span class="hint">(special tokens highlighted)</span></h2>
+    <pre id="tokens"></pre>
+    <h2 style="margin-top:12px">Tool calls</h2>
+    <pre id="toolsLog">–</pre>
+    <details><summary>Prompt sent to model (last round)</summary><pre id="prompt"></pre></details>
+    <details><summary>Server info</summary><pre id="serverInfo"></pre></details>
+    <h2 style="margin-top:12px">Event log</h2>
+    <pre id="log"></pre>
+  </section>
+</main>
+
+<script>
+const $ = (id) => document.getElementById(id);
+const OUT_RATE = 24000;
+let ws, audioCtx, micStream, micNode, micSource, recording = false;
+let playHead = 0, turn = null;
+
+// ---------- logging ----------
+const t0 = performance.now();
+function log(msg, cls) {
+  const line = `[${((performance.now() - t0) / 1000).toFixed(2).padStart(7)}] ${msg}\n`;
+  const el = $("log");
+  el.textContent += line;
+  if (el.textContent.length > 40000) el.textContent = el.textContent.slice(-30000);
+  el.scrollTop = el.scrollHeight;
+}
+function setState(s, cls = "") { $("state").textContent = s; $("state").className = "pill " + cls; }
+
+// ---------- websocket ----------
+function connect() {
+  ws = new WebSocket(`${location.protocol === "https:" ? "wss" : "ws"}://${location.host}/ws`);
+  ws.binaryType = "arraybuffer";
+  ws.onopen = () => { $("conn").textContent = "connected"; $("conn").className = "pill ok"; sendSettings(); };
+  ws.onclose = () => { $("conn").textContent = "disconnected"; $("conn").className = "pill bad"; log("socket closed, retrying in 2s"); setTimeout(connect, 2000); };
+  ws.onmessage = (ev) => typeof ev.data === "string" ? onJson(JSON.parse(ev.data)) : onPcm(ev.data);
+}
+const send = (o) => ws && ws.readyState === 1 && ws.send(typeof o === "string" || o instanceof ArrayBuffer ? o : JSON.stringify(o));
+
+function sendSettings() {
+  send({ type: "settings", mode: $("mode").value, max_new_tokens: +$("max_new_tokens").value,
+         audio_temperature: +$("audio_temperature").value, audio_top_k: +$("audio_top_k").value,
+         system: $("system").value, voice: $("voice").value, keep_history: $("keep_history").value === "1",
+         force_speech: $("force_speech").value === "1" });
+}
+$("mode").addEventListener("change", () => { $("toolHint").hidden = $("mode").value !== "tools"; });
+for (const id of ["mode", "max_new_tokens", "audio_temperature", "audio_top_k", "system", "voice", "keep_history", "force_speech"])
+  $(id).addEventListener("change", sendSettings);
+
+// ---------- turn state ----------
+function newTurn(kind) {
+  turn = { kind, start: performance.now(), events: [], pcm: [], bot: addMsg("bot", ""), firstAudioAt: null, played: 0 };
+  $("tokens").innerHTML = ""; $("prompt").textContent = ""; $("toolsLog").textContent = "–";
+  renderStats({});
+  setState("generating…", "busy");
+}
+function mark(name) { if (turn) turn.events.push({ name, t: performance.now() - turn.start }); drawTimeline(); }
+
+function onJson(m) {
+  switch (m.type) {
+    case "ready":
+      $("info").textContent = `${m.device} · ${m.dtype}`;
+      $("serverInfo").textContent = JSON.stringify(m, null, 2);
+      log(`model ready: ${m.model} on ${m.device}/${m.dtype}, loaded in ${m.load_s}s`);
+      break;
+    case "log": log("server: " + m.msg); break;
+    case "user_audio":
+      log(`user audio ${m.user_audio_s}s rms=${m.user_rms} peak=${m.user_peak}`);
+      if (m.user_audio_s >= 0.3) newTurn("voice");
+      break;
+    case "prompt":
+      $("prompt").textContent = m.text;
+      log(`prompt (round ${m.round}): ${m.tokens} tokens @ ${m.t_ms} ms`);
+      if (m.round > 0) { const sep = document.createElement("span"); sep.className = "sp"; sep.textContent = `\n── round ${m.round} ──\n`; $("tokens").appendChild(sep); }
+      mark(m.round ? `round ${m.round}` : "prompt");
+      break;
+    case "tool_call": {
+      log(`tool call: ${m.name}(${JSON.stringify(m.arguments)}) @ ${m.t_ms} ms`);
+      mark("tool call");
+      const card = addMsg("tool", `🔧 ${m.name}(${JSON.stringify(m.arguments)})`);
+      if (turn) { turn.toolCards = turn.toolCards || []; turn.toolCards.push(card); $("chat").appendChild(turn.bot); }
+      const tl = $("toolsLog"); if (tl.textContent === "–") tl.textContent = "";
+      tl.textContent += `round ${m.round} @ ${m.t_ms} ms → ${m.name}(${JSON.stringify(m.arguments)})\n`;
+      break;
+    }
+    case "tool_result": {
+      log(`tool result: ${m.name} → ${JSON.stringify(m.result)} (${m.ms} ms)`);
+      const card = turn && turn.toolCards && turn.toolCards.shift();
+      if (card) { const r = document.createElement("div"); r.className = "res"; r.textContent = "↳ " + JSON.stringify(m.result); card.appendChild(r); }
+      $("toolsLog").textContent += `   ↳ ${JSON.stringify(m.result)} (${m.ms} ms)\n`;
+      break;
+    }
+    case "ui":
+      if (m.action === "accent") { document.documentElement.style.setProperty("--accent", m.value); log(`UI: accent → ${m.value}`); }
+      break;
+    case "token": {
+      if (!turn) newTurn("?");
+      if (!turn.firstTok) { turn.firstTok = true; mark("1st text"); log(`first text token @ ${m.t_ms} ms (server)`); }
+      const span = document.createElement("span");
+      span.textContent = m.raw; span.title = `id ${m.id} @ ${m.t_ms} ms`;
+      if (m.special) span.className = "sp";
+      $("tokens").appendChild(span); $("tokens").scrollTop = $("tokens").scrollHeight;
+      if (m.delta) { turn.bot.querySelector(".body").textContent += m.delta; $("chat").scrollTop = $("chat").scrollHeight; }
+      break;
+    }
+    case "audio_chunk":
+      if (!turn) newTurn("?");
+      if (turn.chunks === undefined) { turn.chunks = 0; mark("1st audio"); log(`first audio chunk @ ${m.t_ms} ms (server), codes[0]=${JSON.stringify(m.codes)}`); }
+      turn.chunks++;
+      turn.lastChunk = m;
+      break;
+    case "done": {
+      const s = m.stats;
+      mark("done");
+      renderStats(s);
+      log(`done: ${JSON.stringify(s)}`);
+      setState(s.cancelled ? "cancelled" : "idle", s.cancelled ? "bad" : "");
+      if (turn) {
+        const meta = document.createElement("div"); meta.className = "meta";
+        meta.textContent = `${s.mode} · ${s.gen_ms} ms · ${s.audio_out_s}s audio${s.rtf ? " · RTF " + s.rtf : ""}`;
+        turn.bot.appendChild(meta);
+        if (turn.pcm.length) attachAudio(turn.bot, turn.pcm, OUT_RATE);
+        if (!turn.bot.querySelector(".body").textContent && s.mode === "tts") turn.bot.querySelector(".body").textContent = "(speech)";
+      }
+      turn = null;
+      break;
+    }
+    case "error": log("ERROR " + m.msg); setState("error", "bad"); if (turn) turn.bot.querySelector(".body").textContent += `\n⚠ ${m.msg}`; turn = null; break;
+  }
+}
+
+// ---------- playback ----------
+function ensureCtx() {
+  if (!audioCtx) audioCtx = new AudioContext();
+  if (audioCtx.state === "suspended") audioCtx.resume();
+  return audioCtx;
+}
+function onPcm(buf) {
+  const pcm = new Float32Array(buf);
+  if (!turn) return;
+  turn.pcm.push(pcm);
+  const ctx = ensureCtx();
+  const b = ctx.createBuffer(1, pcm.length, OUT_RATE);
+  b.copyToChannel(pcm, 0);
+  const src = ctx.createBufferSource();
+  src.buffer = b; src.connect(ctx.destination);
+  const now = ctx.currentTime;
+  if (playHead < now + 0.02) {
+    if (turn.played > 0) { turn.underruns = (turn.underruns || 0) + 1; log(`playback underrun #${turn.underruns} (gap ${((now - playHead) * 1000).toFixed(0)} ms)`); }
+    playHead = now + 0.05;  // small jitter buffer
+  }
+  if (!turn.played) { mark("play"); log(`playback starts, client latency ${(performance.now() - turn.start).toFixed(0)} ms since release`); }
+  src.start(playHead);
+  playHead += b.duration;
+  turn.played += pcm.length;
+}
+function attachAudio(el, chunks, rate) {
+  const n = chunks.reduce((a, c) => a + c.length, 0);
+  const all = new Float32Array(n); let o = 0;
+  for (const c of chunks) { all.set(c, o); o += c.length; }
+  const a = document.createElement("audio"); a.controls = true; a.src = URL.createObjectURL(wav(all, rate));
+  el.appendChild(a);
+}
+function wav(f32, rate) {
+  const buf = new ArrayBuffer(44 + f32.length * 2), v = new DataView(buf);
+  const w = (o, s) => [...s].forEach((c, i) => v.setUint8(o + i, c.charCodeAt(0)));
+  w(0, "RIFF"); v.setUint32(4, 36 + f32.length * 2, true); w(8, "WAVEfmt "); v.setUint32(16, 16, true);
+  v.setUint16(20, 1, true); v.setUint16(22, 1, true); v.setUint32(24, rate, true); v.setUint32(28, rate * 2, true);
+  v.setUint16(32, 2, true); v.setUint16(34, 16, true); w(36, "data"); v.setUint32(40, f32.length * 2, true);
+  for (let i = 0; i < f32.length; i++) v.setInt16(44 + i * 2, Math.max(-1, Math.min(1, f32[i])) * 0x7fff, true);
+  return new Blob([buf], { type: "audio/wav" });
+}
+
+// ---------- mic capture ----------
+const WORKLET = `
+class Cap extends AudioWorkletProcessor {
+  constructor() { super(); this.buf = []; this.n = 0; }
+  process(inputs) {
+    const ch = inputs[0][0];
+    if (ch) { this.buf.push(new Float32Array(ch)); this.n += ch.length;
+      if (this.n >= 2048) { const out = new Float32Array(this.n); let o = 0;
+        for (const b of this.buf) { out.set(b, o); o += b.length; }
+        this.port.postMessage(out, [out.buffer]); this.buf = []; this.n = 0; } }
+    return true;
+  }
+}
+registerProcessor("cap", Cap);`;
+
+async function initMic() {
+  const ctx = ensureCtx();
+  if (micNode) return;
+  micStream = await navigator.mediaDevices.getUserMedia({ audio: { echoCancellation: true, noiseSuppression: true, channelCount: 1 } });
+  await ctx.audioWorklet.addModule(URL.createObjectURL(new Blob([WORKLET], { type: "application/javascript" })));
+  micSource = ctx.createMediaStreamSource(micStream);
+  micNode = new AudioWorkletNode(ctx, "cap");
+  micNode.port.onmessage = (e) => {
+    const d = e.data;
+    let s = 0; for (let i = 0; i < d.length; i++) s += d[i] * d[i];
+    $("meter").firstChild.style.width = Math.min(100, Math.sqrt(s / d.length) * 400) + "%";
+    if (recording) { send(d.buffer); recBuf.push(d); }
+  };
+  micSource.connect(micNode);
+  log(`mic ready @ ${ctx.sampleRate} Hz`);
+}
+let recBuf = [], recStart = 0;
+async function startRec() {
+  if (recording) return;
+  try { await initMic(); } catch (e) { log("mic error: " + e); return; }
+  recording = true; recBuf = []; recStart = performance.now();
+  send({ type: "start", rate: audioCtx.sampleRate });
+  $("talk").classList.add("rec"); $("talk").textContent = "Listening… release to send";
+  setState("recording", "bad");
+}
+function stopRec() {
+  if (!recording) return;
+  recording = false;
+  send({ type: "stop" });
+  $("talk").classList.remove("rec"); $("talk").textContent = "Hold to talk (or hold Space)";
+  $("meter").firstChild.style.width = "0";
+  const dur = (performance.now() - recStart) / 1000;
+  if (dur >= 0.3) {
+    const m = addMsg("user", $("mode").value === "asr" ? "🎙 (transcribe)" : "🎙");
+    attachAudio(m, recBuf, audioCtx.sampleRate);
+  }
+  setState("waiting…", "busy");
+}
+const talk = $("talk");
+talk.addEventListener("pointerdown", (e) => { e.preventDefault(); talk.setPointerCapture(e.pointerId); startRec(); });
+talk.addEventListener("pointerup", stopRec);
+talk.addEventListener("pointercancel", stopRec);
+window.addEventListener("keydown", (e) => { if (e.code === "Space" && !e.repeat && !/INPUT|TEXTAREA|SELECT/.test(document.activeElement.tagName)) { e.preventDefault(); startRec(); } });
+window.addEventListener("keyup", (e) => { if (e.code === "Space" && recording) { e.preventDefault(); stopRec(); } });
+
+// ---------- other controls ----------
+$("cancel").onclick = () => send({ type: "cancel" });
+$("reset").onclick = () => { send({ type: "reset" }); addMsg("bot", "— history cleared —"); };
+$("clear").onclick = () => { $("chat").innerHTML = ""; $("log").textContent = ""; };
+$("tts").onclick = () => {
+  const text = $("ttsText").value.trim(); if (!text) return;
+  ensureCtx(); addMsg("user", "🔊 " + text); newTurn("tts");
+  send({ type: "tts", text });
+};
+$("sendText").onclick = () => {
+  const text = $("ttsText").value.trim(); if (!text) return;
+  ensureCtx(); addMsg("user", text); newTurn("text"); $("ttsText").value = "";
+  send({ type: "text", text });
+};
+$("ttsText").addEventListener("keydown", (e) => { if (e.key === "Enter") $("sendText").click(); });
+
+function addMsg(who, text) {
+  const d = document.createElement("div"); d.className = "msg " + who;
+  const b = document.createElement("div"); b.className = "body"; b.textContent = text; d.appendChild(b);
+  $("chat").appendChild(d); $("chat").scrollTop = $("chat").scrollHeight;
+  return d;
+}
+
+// ---------- debug rendering ----------
+const STAT_LABELS = [
+  ["user_audio_s", "input audio (s)"], ["prompt_tokens", "prompt tokens"], ["prep_ms", "prep (ms)"],
+  ["first_text_ms", "1st text (ms)"], ["first_frame_ms", "1st frame (ms)"], ["first_audio_ms", "1st audio chunk (ms)"],
+  ["gen_ms", "generate (ms)"], ["rounds", "rounds (tools)"], ["speech_restarts", "speech restarts"], ["text_tokens", "text tokens"], ["audio_frames", "audio frames"],
+  ["audio_out_s", "audio out (s)"], ["rtf", "RTF (gen/audio)"], ["text_tok_per_s", "text tok/s"],
+  ["frames_per_s", "frames/s (12.5=RT)"], ["codec_decode_ms", "codec decode (ms)"], ["max_decode_backlog", "decode backlog (frames)"],
+];
+function renderStats(s) {
+  $("stats").innerHTML = STAT_LABELS.map(([k, l]) => `<div class="stat"><b>${s[k] ?? "–"}</b><span>${l}</span></div>`).join("")
+    + (turn && turn.underruns ? `<div class="stat"><b>${turn.underruns}</b><span>playback underruns</span></div>` : "");
+}
+function drawTimeline() {
+  const c = $("timeline"), g = c.getContext("2d");
+  c.width = c.clientWidth * devicePixelRatio; c.height = 60 * devicePixelRatio;
+  g.scale(devicePixelRatio, devicePixelRatio);
+  const W = c.clientWidth, css = getComputedStyle(document.documentElement);
+  g.clearRect(0, 0, W, 60);
+  if (!turn || !turn.events.length) return;
+  const max = Math.max(1000, ...turn.events.map((e) => e.t)) * 1.1;
+  g.fillStyle = css.getPropertyValue("--muted"); g.font = "11px system-ui";
+  g.fillText(`client timeline (ms since release) — 0…${max.toFixed(0)}`, 6, 13);
+  turn.events.forEach((e, i) => {
+    const x = 6 + (e.t / max) * (W - 12);
+    g.fillStyle = css.getPropertyValue("--accent"); g.fillRect(x, 20, 2, 34);
+    g.fillStyle = css.getPropertyValue("--fg");
+    g.fillText(`${e.name} ${e.t.toFixed(0)}`, Math.min(x + 4, W - 90), 32 + (i % 2) * 14);
+  });
+}
+renderStats({});
+connect();
+</script>
+</body>
+</html>

tools.py

@@ -0,0 +1,149 @@
+"""Test tools for the "Voice + tools" mode, plus the parser for the model's tool-call format."""
+
+import ast
+import hashlib
+import json
+import operator
+import re
+from datetime import datetime
+from zoneinfo import ZoneInfo, ZoneInfoNotFoundError
+
+SYSTEM_PROMPT = "Respond with function calls, and interleaved text and audio.\n\nList of tools:"
+
+TOOLS = [
+    {
+        "type": "function",
+        "function": {
+            "name": "get_current_time",
+            "description": "Get the current local date and time in a time zone.",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "timezone": {"type": "string", "description": "IANA time zone, e.g. Europe/Stockholm. Defaults to UTC."}
+                },
+            },
+        },
+    },
+    {
+        "type": "function",
+        "function": {
+            "name": "get_weather",
+            "description": "Get the current weather for a city.",
+            "parameters": {
+                "type": "object",
+                "properties": {"city": {"type": "string", "description": "City name, e.g. Paris."}},
+                "required": ["city"],
+            },
+        },
+    },
+    {
+        "type": "function",
+        "function": {
+            "name": "calculate",
+            "description": "Evaluate an arithmetic expression with + - * / ** % and parentheses.",
+            "parameters": {
+                "type": "object",
+                "properties": {"expression": {"type": "string", "description": "The expression, e.g. (17 * 23) / 4."}},
+                "required": ["expression"],
+            },
+        },
+    },
+    {
+        "type": "function",
+        "function": {
+            "name": "set_theme_color",
+            "description": "Change the accent color of the user's screen.",
+            "parameters": {
+                "type": "object",
+                "properties": {"color": {"type": "string", "description": "A CSS color name or hex code, e.g. green or #ff8800."}},
+                "required": ["color"],
+            },
+        },
+    },
+]
+
+CALL_RE = re.compile(r"<tool_call>\s*<function=([^>\n]+)>(.*?)</function>\s*</tool_call>", re.DOTALL)
+PARAM_RE = re.compile(r"<parameter=([^>\n]+)>\n?(.*?)\n?</parameter>", re.DOTALL)
+OPEN_TAG = "<tool_call>"
+
+
+def parse_tool_calls(text: str) -> list[dict]:
+    calls = []
+    for name, body in CALL_RE.findall(text):
+        arguments = {}
+        for key, value in PARAM_RE.findall(body):
+            try:
+                arguments[key] = json.loads(value)
+            except json.JSONDecodeError:
+                arguments[key] = value  # plain string
+        calls.append({"type": "function", "function": {"name": name.strip(), "arguments": arguments}})
+    return calls
+
+
+def visible_text(text: str) -> str:
+    """The text with tool-call blocks removed, cut before an unfinished one, so the chat bubble never shows markup."""
+    text = CALL_RE.sub("", text)
+    cut = text.find("<tool_call")
+    if cut >= 0:
+        return text[:cut]
+    for k in range(len(OPEN_TAG) - 1, 0, -1):  # a tag still arriving token by token
+        if text.endswith(OPEN_TAG[:k]):
+            return text[:-k]
+    return text
+
+
+_OPS = {
+    ast.Add: operator.add, ast.Sub: operator.sub, ast.Mult: operator.mul, ast.Div: operator.truediv,
+    ast.Pow: operator.pow, ast.Mod: operator.mod, ast.FloorDiv: operator.floordiv,
+    ast.USub: operator.neg, ast.UAdd: operator.pos,
+}
+
+
+def _eval(node):
+    if isinstance(node, ast.Expression):
+        return _eval(node.body)
+    if isinstance(node, ast.Constant) and isinstance(node.value, (int, float)):
+        return node.value
+    if isinstance(node, ast.BinOp) and type(node.op) in _OPS:
+        if isinstance(node.op, ast.Pow) and abs(_eval(node.right)) > 100:
+            raise ValueError("exponent too large")
+        return _OPS[type(node.op)](_eval(node.left), _eval(node.right))
+    if isinstance(node, ast.UnaryOp) and type(node.op) in _OPS:
+        return _OPS[type(node.op)](_eval(node.operand))
+    raise ValueError("unsupported expression")
+
+
+def get_current_time(timezone: str = "UTC", **_):
+    try:
+        now = datetime.now(ZoneInfo(timezone))
+    except (ZoneInfoNotFoundError, ValueError):
+        return {"error": f"unknown time zone {timezone!r}"}
+    return {"timezone": timezone, "datetime": now.strftime("%Y-%m-%d %H:%M"), "weekday": now.strftime("%A")}
+
+
+def get_weather(city: str, **_):
+    # simulated but stable per city, so repeated tests give the same answer
+    h = int(hashlib.sha1(city.strip().lower().encode()).hexdigest(), 16)
+    conditions = ["sunny", "partly cloudy", "overcast", "light rain", "windy", "snowing"]
+    return {"city": city, "temperature_c": h % 35 - 5, "condition": conditions[h % len(conditions)], "simulated": True}
+
+
+def calculate(expression: str, **_):
+    try:
+        return {"expression": expression, "result": _eval(ast.parse(str(expression), mode="eval"))}
+    except Exception as e:
+        return {"expression": expression, "error": str(e)}
+
+
+def run_tool(name: str, arguments: dict, ui) -> dict:
+    """Run one call; `ui(action, value)` pushes a visible effect to the browser."""
+    if name == "set_theme_color":
+        ui("accent", str(arguments.get("color", "")))
+        return {"ok": True, "color": arguments.get("color")}
+    funcs = {"get_current_time": get_current_time, "get_weather": get_weather, "calculate": calculate}
+    if name not in funcs:
+        return {"error": f"unknown tool {name!r}"}
+    try:
+        return funcs[name](**arguments)
+    except TypeError as e:
+        return {"error": f"bad arguments: {e}"}