first tests

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")))