Qwen3-TTS on a Shared GPU

# from the autohub repo; isolated sidecars on spare ports, never the production one
PY=.venv-qwen3-tts/bin/python
$PY scripts/qwen3-tts-server.py --port 8786 --voice-ref audio/autojack-ref.wav \
    --model mlx-community/Qwen3-TTS-12Hz-1.7B-Base-bf16 --temperature 0.7 &
$PY scripts/qwen3-tts-server.py --port 8784 --voice-ref audio/autojack-ref.wav \
    --model mlx-community/Qwen3-TTS-12Hz-0.6B-Base-bf16 --temperature 0.7 &
~/.omlx/venv/bin/omlx serve --model-dir ~/.omlx/models --port 8081 &

sysctl vm.swapusage          # must read ~0 before trusting anything below
$PY bench.py                 # TTS phases + Gemma phase, writes bench.json
$PY bench.py gemma           # re-run only the Gemma phase

"""Qwen3-TTS 1.7B-bf16 vs 0.6B-bf16 on isolated sidecars (never prod :8783).

Measures, per model: client TTFA + RTF over the SSE route, idle and while Gemma 4
26B-A4B streams on oMLX; Gemma decode tok/s alone and under each TTS loop;
phys_footprint of each sidecar.
"""
import base64, json, statistics, subprocess, threading, time
from pathlib import Path
import httpx

OUT = Path(__file__).parent
MODELS = {"1.7B-bf16": "http://127.0.0.1:8786", "0.6B-bf16": "http://127.0.0.1:8784"}
OMLX = "http://127.0.0.1:8081"
GEMMA = "mlx-community--gemma-4-26b-a4b-it-4bit"
REPEATS = 3
SENTENCES = [
    "Hey Jack, you've got three PRs waiting on review, and the watch relay one just went green.",
    "Your ten thirty with Danielle moved to eleven fifteen, so you have about forty minutes free right now.",
    "Honestly? I'd skip the refactor tonight. The voice path is stable, and you're running on four hours of sleep.",
    "The hub's been up for nine days straight. Memory's fine, but Railway billed us twenty-three dollars more than last month.",
    "Okay, done! I set a reminder for Thursday at two p.m. to renew the Cloudflare certificate.",
    "Wait, that's weird. The Gemma turn took eight point four seconds, which is way slower than usual. Want me to dig in?",
]

def sse_once(c, url, text, seed):
    t0 = time.perf_counter()
    ttfa = None
    audio_s = 0.0
    with c.stream("POST", f"{url}/v1/audio/speech",
                  json={"input": text, "stream_format": "sse", "seed": seed}) as r:
        r.raise_for_status()
        for line in r.iter_lines():
            if not line.startswith("data: "):
                continue
            ev = json.loads(line[6:])
            if ev["type"] == "speech.audio.delta":
                if ttfa is None:
                    ttfa = (time.perf_counter() - t0) * 1000
                audio_s += len(base64.b64decode(ev["audio"])) / 2 / 24000
            elif ev["type"] == "speech.audio.error":
                raise RuntimeError(ev)
    wall = time.perf_counter() - t0
    return {"ttfa_ms": ttfa, "rtf": wall / audio_s}

def gemma_stream(stop, samples):
    prompt = "Write a long, detailed story about a lighthouse keeper and a cat. Keep going."
    with httpx.Client(timeout=300) as c:
        while not stop.is_set():
            usage = None
            with c.stream("POST", f"{OMLX}/v1/chat/completions", json={
                "model": GEMMA, "stream": True, "max_tokens": 300, "temperature": 0.7,
                "stream_options": {"include_usage": True},
                "messages": [{"role": "user", "content": prompt}],
            }) as r:
                for line in r.iter_lines():
                    if not line.startswith("data: ") or line.strip() == "data: [DONE]":
                        continue
                    try:
                        ch = json.loads(line[6:]).get("choices") or []
                    except json.JSONDecodeError:
                        continue
                    usage = json.loads(line[6:]).get("usage") or usage
            # oMLX bursts several tokens per SSE chunk; trust its own decode rate.
            if usage and usage.get("completion_tokens", 0) >= 100:
                samples.append(usage["generation_tokens_per_second"])

def pct(vals, q):
    vals = sorted(vals)
    return vals[min(len(vals) - 1, int(round(q * (len(vals) - 1))))]

def summarize(rows):
    t = [r["ttfa_ms"] for r in rows]
    f = [r["rtf"] for r in rows]
    return {"n": len(rows), "ttfa_p50": round(statistics.median(t)), "ttfa_p90": round(pct(t, 0.9)),
            "ttfa_max": round(max(t)), "rtf_p50": round(statistics.median(f), 3),
            "rtf_p90": round(pct(f, 0.9), 3), "rtf_max": round(max(f), 3)}

def footprint(url):
    pid = str(httpx.get(f"{url}/health").json()["pid"])
    out = subprocess.run(["footprint", "-p", pid], capture_output=True, text=True).stdout
    for line in out.splitlines():
        if "phys_footprint:" in line:
            return line.split("phys_footprint:")[1].strip()
    return "?"

def main():
    import sys
    if sys.argv[1:] == ["gemma"]:
        results = json.loads((OUT / "bench.json").read_text())
        return gemma_phase(results)
    results = {}
    with httpx.Client(timeout=300) as c:
        for url in MODELS.values():
            for _ in range(2):
                sse_once(c, url, SENTENCES[0], 1)
        # Load Gemma before any timed phase so model load never lands in a sample.
        with c.stream("POST", f"{OMLX}/v1/chat/completions", json={
            "model": GEMMA, "stream": True, "max_tokens": 16,
            "messages": [{"role": "user", "content": "hi"}]}) as r:
            for _ in r.iter_lines():
                pass

        for condition in ["idle", "gemma_busy"]:
            stop = threading.Event()
            gemma = []
            th = None
            if condition == "gemma_busy":
                th = threading.Thread(target=gemma_stream, args=(stop, gemma), daemon=True)
                th.start()
                time.sleep(3)
            rows = {k: [] for k in MODELS}
            names = list(MODELS)
            for rep in range(REPEATS):
                for si, text in enumerate(SENTENCES):
                    order = names if (rep + si) % 2 == 0 else names[::-1]
                    for name in order:
                        rows[name].append(sse_once(c, MODELS[name], text, 100 + si + rep))
            stop.set()
            if th:
                th.join(timeout=120)
            results[condition] = {k: summarize(v) for k, v in rows.items()}
            print(condition, json.dumps(results[condition]), flush=True)

    gemma_phase(results)

def gemma_phase(results):
    gem = {}
    for label, url in [("alone", None)] + list(MODELS.items()):
        stop, tts_stop = threading.Event(), threading.Event()
        samples = []

        def tts_loop(u=url):
            with httpx.Client(timeout=300) as c2:
                i = 0
                while not tts_stop.is_set():
                    sse_once(c2, u, SENTENCES[i % len(SENTENCES)], 7)
                    i += 1

        tts_th = threading.Thread(target=tts_loop, daemon=True) if url else None
        if tts_th:
            tts_th.start()
            time.sleep(1)
        g = threading.Thread(target=gemma_stream, args=(stop, samples), daemon=True)
        g.start()
        while len(samples) < 6:
            time.sleep(0.5)
        stop.set()
        tts_stop.set()
        g.join(timeout=120)
        if tts_th:
            tts_th.join(timeout=120)
        gem[label] = round(statistics.median(samples), 1)
        print("gemma tok/s", label, gem[label], flush=True)
    results["gemma_tok_s"] = gem
    results["footprint"] = {n: footprint(u) for n, u in MODELS.items()}
    (OUT / "bench.json").write_text(json.dumps(results, indent=2))
    print(json.dumps(results, indent=2))

if __name__ == "__main__":
    main()