Upload app.py with huggingface_hub

app.py

@@ -11,6 +11,8 @@ from gradio import Server
 from fastapi import Request
 from fastapi.responses import HTMLResponse
 from huggingface_hub import hf_hub_download
+import asyncio, threading, queue as _q, contextvars, struct
+from fastapi import WebSocket, WebSocketDisconnect
 
 from magenta_rt import paths
 
@@ -145,9 +147,8 @@ async def banks(session_id: str = ""):
     return {"bankStatus": [os.path.exists(os.path.join(SESSION_DIR, f"{sid}_bank{i}.pt")) for i in range(3)]}
 
 
-@app.api(name="stream")
 @spaces.GPU(duration=90)
-def stream(session_id: str) -> str:
+def gpu_stream(session_id):
     """Continuous gen; switches model live when the dropdown changes."""
     from magenta_rt.torch.system import make_sampler, discretize_cfg, _float_to_int16, FRAME_SAMPLES
     if style_model.device != "cuda":
@@ -284,8 +285,8 @@ def stream(session_id: str) -> str:
         emitted_samples += audio.shape[1]
         frame_ms = (time.time() - gen_t) * 1000.0 / CHUNK     # real per-frame inference time
         if audio.shape[1] > 0:
-            b64 = base64.b64encode(_float_to_int16(audio[0].float().cpu().numpy()).astype("<i2").tobytes()).decode("ascii")
-            yield f"{frame_ms:.1f}|{b64}"
+            ab = _float_to_int16(audio[0].float().cpu().numpy()).astype("<i2").tobytes()
+            yield (ab, frame_ms)
         buf = (0.15, 0.3, 0.5)[max(0, min(2, int(c.get("buffer", 0))))]
         ahead = (emitted_samples / SR) - (time.time() - t0)
         if ahead < -0.3:                              # fell behind (cold-start/hiccup) -> resync, no catch-up burst
@@ -304,6 +305,70 @@ def stream(session_id: str) -> str:
             time.sleep(min(sleep_for, 0.5))
 
 
+_sessions = {}
+_slock = threading.Lock()
+
+
+def _worker(session_id, audio_q, stop_ev):
+    try:
+        while not stop_ev.is_set():
+            try:
+                for ab, fm in gpu_stream(session_id):           # ~55s GPU grant; re-grants on loop
+                    if stop_ev.is_set():
+                        break
+                    if audio_q.full():
+                        try: audio_q.get_nowait()
+                        except _q.Empty: pass
+                    try: audio_q.put_nowait((ab, fm))
+                    except _q.Full: pass
+            except Exception as e:
+                if "abort" in str(e).lower() or "duration" in str(e).lower():
+                    continue                                    # grant expired -> re-grant
+                print("[worker]", repr(e), flush=True); break
+    finally:
+        stop_ev.set()
+
+
+@app.api(name="start")
+def start(session_id: str = ""):
+    if not session_id:
+        return {"ok": False}
+    with _slock:
+        prev = _sessions.pop(session_id, None)
+    if prev:
+        prev["stop"].set()
+    audio_q = _q.Queue(maxsize=48); stop_ev = threading.Event()
+    ctx = contextvars.copy_context()
+    t = threading.Thread(target=ctx.run, args=(_worker, session_id, audio_q, stop_ev), daemon=True)
+    with _slock:
+        _sessions[session_id] = {"audio": audio_q, "stop": stop_ev}
+    t.start()
+    return {"ok": True}
+
+
+@app.websocket("/ws")
+async def audio_ws(websocket: WebSocket, session_id: str = ""):
+    await websocket.accept()
+    sess = _sessions.get(session_id)
+    if not sess:
+        await websocket.close(code=1008); return
+    audio_q, stop_ev = sess["audio"], sess["stop"]
+    loop = asyncio.get_event_loop()
+    try:
+        while not stop_ev.is_set():
+            try:
+                ab, fm = await loop.run_in_executor(None, lambda: audio_q.get(timeout=1.0))
+            except Exception:
+                continue
+            await websocket.send_bytes(struct.pack("<f", fm) + ab)
+    except WebSocketDisconnect:
+        pass
+    finally:
+        stop_ev.set()
+        with _slock:
+            _sessions.pop(session_id, None)
+
+
 @app.get("/")
 async def index():
     with open(os.path.join(os.path.dirname(os.path.abspath(__file__)), "index.html")) as f: