merge: SSE refactor + ZeroGPU img2img watercolour develop

app.py

@@ -61,8 +61,8 @@ import threading
 import time
 import uuid
 
-from fastapi import HTTPException, Request, WebSocket, WebSocketDisconnect
-from fastapi.responses import FileResponse, HTMLResponse, Response
+from fastapi import HTTPException, Request
+from fastapi.responses import FileResponse, HTMLResponse, Response, StreamingResponse
 
 from gradio import Server
 
@@ -2333,7 +2333,7 @@ async def put_room_state(room_id: str, request: Request):
     return {"version": new_version, "conflict": False}
 
 
-# --------------------------- WebSocket sync -------------------------------- #
+# --------------------------- Realtime sync (SSE) --------------------------- #
 
 _PALETTE = [
     "#7c5cff", "#e9698f", "#2bb673", "#f2994a", "#3aa0ff",
@@ -2378,80 +2378,140 @@ NON_RELAY_EVENTS = {
 }
 
 
+# NOTE: This transport was a raw WebSocket (/ws/{room_id}). Hugging Face
+# Spaces' router silently drops custom (non-Gradio) WebSocket UPGRADE requests
+# at the proxy — the handshake never reaches this process, so every connection
+# loops in "reconnecting". Plain HTTP routes on this same app work fine, so we
+# mirror what Gradio itself does on Spaces: a one-way Server-Sent Events stream
+# for server→client messages, plus ordinary POST for client→server messages.
+#
+# Each client is identified by its session_id and holds one open SSE GET. The
+# Manager keys connections by an asyncio.Queue per client instead of a socket;
+# "sending" to a client just puts a message on its queue, which the SSE
+# generator drains to the browser. Presence, join-approval, the relay
+# allow-list, and the DB writes are all unchanged from the WebSocket version.
+#
+# SINGLE-REPLICA ASSUMPTION: a POST and the matching SSE stream must land on the
+# same process. Your Space runs one replica today, so this holds. If you ever
+# scale to multiple replicas, move the queues behind Redis pub/sub so a POST on
+# replica A reaches a stream on replica B.
+
+
+# The four fields that are safe to broadcast to other participants. `info`
+# never contains anything else now, but we filter through this everywhere it
+# crosses the wire so a stray non-JSON value (e.g. an asyncio.Event) can never
+# crash the SSE stream's json.dumps.
+_PUBLIC_FIELDS = ("session_id", "name", "participant_id", "color")
+
+
+def _public(info):
+    """Project a participant's info down to only the JSON-safe public fields."""
+    return {k: info[k] for k in _PUBLIC_FIELDS}
+
+
+class _Client:
+    """One connected participant: their info + the queue feeding their SSE."""
+    def __init__(self, info):
+        self.info = info
+        self.queue: asyncio.Queue = asyncio.Queue()
+        self.alive = True
+
+    async def put(self, message):
+        if self.alive:
+            await self.queue.put(message)
+
+
+class _PendingJoin:
+    """A joiner awaiting host approval. Holds the decision Event SEPARATELY from
+    `info` so `info` stays pure-JSON and can be broadcast safely."""
+    def __init__(self, info):
+        self.info = info
+        self.event = asyncio.Event()
+        self.admitted = False
+
+
 class Manager:
     def __init__(self):
-        self.rooms = {}     # room_id -> {ws: info}  (admitted participants)
-        self.pending = {}   # room_id -> {ws: info}  (awaiting host approval)
+        self.rooms = {}     # room_id -> {session_id: _Client}  (admitted)
+        self.pending = {}   # room_id -> {session_id: info}     (awaiting host)
 
-    async def connect(self, room_id, ws, info):
-        await ws.accept()
-        self.rooms.setdefault(room_id, {})[ws] = info
+    # ---- admitted participants -------------------------------------------- #
+    def add_client(self, room_id, client):
+        self.rooms.setdefault(room_id, {})[client.info["session_id"]] = client
 
-    def admit(self, room_id, ws, info):
-        """Promote a pending socket to a full participant."""
-        self.pending.get(room_id, {}).pop(ws, None)
-        self.rooms.setdefault(room_id, {})[ws] = info
+    def get_client(self, room_id, session_id):
+        return self.rooms.get(room_id, {}).get(session_id)
+
+    def disconnect(self, room_id, session_id):
+        peers = self.rooms.get(room_id)
+        if peers and session_id in peers:
+            peers[session_id].alive = False
+            del peers[session_id]
+            if not peers:
+                self.rooms.pop(room_id, None)
+        self.remove_pending(room_id, session_id)
 
-    def add_pending(self, room_id, ws, info):
-        self.pending.setdefault(room_id, {})[ws] = info
+    # ---- pending (awaiting host approval) --------------------------------- #
+    def admit(self, room_id, info):
+        """Promote a pending session to a full participant; returns its _Client."""
+        self.pending.get(room_id, {}).pop(info["session_id"], None)
+        client = _Client(info)
+        self.add_client(room_id, client)
+        return client
 
-    def remove_pending(self, room_id, ws):
+    def add_pending(self, room_id, pending):
+        self.pending.setdefault(room_id, {})[pending.info["session_id"]] = pending
+
+    def remove_pending(self, room_id, session_id):
         p = self.pending.get(room_id)
-        if p and ws in p:
-            del p[ws]
+        if p and session_id in p:
+            del p[session_id]
             if not p:
                 self.pending.pop(room_id, None)
 
     def find_pending(self, room_id, session_id):
-        """Return (ws, info) for a pending session_id, or (None, None)."""
-        for ws, info in list(self.pending.get(room_id, {}).items()):
-            if info["session_id"] == session_id:
-                return ws, info
-        return None, None
-
-    def host_ws(self, room_id, owner_session_id):
-        """The socket that should approve joins: the owner if connected, else
+        """Return the _PendingJoin for a session, or None."""
+        return self.pending.get(room_id, {}).get(session_id)
+
+    # ---- host resolution (who approves joins) ----------------------------- #
+    def host_client(self, room_id, owner_session_id):
+        """The client that should approve joins: the owner if connected, else
         the earliest-connected participant (so approval never deadlocks)."""
         peers = self.rooms.get(room_id, {})
         if not peers:
             return None
-        if owner_session_id:
-            for ws, info in peers.items():
-                if info["session_id"] == owner_session_id:
-                    return ws
-        # fall back to the first admitted participant
-        return next(iter(peers), None)
-
-    def disconnect(self, room_id, ws):
-        peers = self.rooms.get(room_id)
-        if peers and ws in peers:
-            del peers[ws]
-            if not peers:
-                self.rooms.pop(room_id, None)
-        self.remove_pending(room_id, ws)
+        if owner_session_id and owner_session_id in peers:
+            return peers[owner_session_id]
+        return next(iter(peers.values()), None)
 
+    # ---- presence --------------------------------------------------------- #
     def presence(self, room_id):
         peers = self.rooms.get(room_id, {})
         seen, out = set(), []
-        for info in peers.values():
+        for client in peers.values():
+            info = client.info
             if info["session_id"] in seen:
                 continue
             seen.add(info["session_id"])
-            out.append({k: info[k] for k in ("session_id", "name", "participant_id", "color")})
+            out.append(_public(info))
         return out
 
+    # ---- delivery --------------------------------------------------------- #
     async def broadcast(self, room_id, message, exclude=None):
-        for ws in list(self.rooms.get(room_id, {})):
-            if ws is exclude:
+        for sid, client in list(self.rooms.get(room_id, {}).items()):
+            if sid == exclude:
                 continue
             try:
-                await ws.send_json(message)
+                await client.put(message)
             except Exception:
                 pass
 
-    async def send_to(self, ws, message):
+    async def send_to(self, room_id, session_id, message):
+        client = self.get_client(room_id, session_id)
+        if client is None:
+            return False
         try:
-            await ws.send_json(message)
+            await client.put(message)
             return True
         except Exception:
             return False
@@ -2459,124 +2519,175 @@ class Manager:
 
 manager = Manager()
 
+# Pending joiners don't have a _Client/queue yet (they're not admitted), but
+# they still need to receive join_pending / join_declined over their own SSE
+# stream. We give every SSE connection a queue up front, keyed by session_id,
+# so the stream exists from the moment the browser connects — even while the
+# session is still waiting for host approval.
+_sse_queues = {}  # (room_id, session_id) -> asyncio.Queue
+
+
+def _sse_queue(room_id, session_id):
+    key = (room_id, session_id)
+    q = _sse_queues.get(key)
+    if q is None:
+        q = asyncio.Queue()
+        _sse_queues[key] = q
+    return q
+
+
+async def _relay_incoming(room_id, session_id, data):
+    """Apply one client→server event (formerly a WS receive_json). Shared by the
+    POST /send endpoint. `data` already has a 'type'."""
+    client = manager.get_client(room_id, session_id)
+    info = client.info if client else None
+
+    t = data.get("type")
+
+    if t == "rename" and info is not None:
+        info["name"] = sanitize_name(data.get("name"))
+        with DB_LOCK:
+            _conn.execute(
+                "UPDATE participants SET display_name=?, last_seen=? WHERE participant_id=?",
+                (info["name"], time.time(), info["participant_id"]),
+            )
+            _conn.commit()
+        await manager.broadcast(room_id, {"type": "presence_update", "participants": manager.presence(room_id)})
+        return
+
+    # ---- host approves / declines a pending joiner ----
+    if t in ("admit_join", "decline_join"):
+        target_sid = data.get("session_id")
+        pend = manager.find_pending(room_id, target_sid)
+        if pend is not None:
+            pend.admitted = (t == "admit_join")
+            pend.event.set()
+        return
+
+    # explicitly known (RELAY_EVENTS) or simply not a server-owned type
+    # (NON_RELAY_EVENTS) — this way the frontend can introduce new interaction
+    # events without the server silently dropping them.
+    if t and t not in NON_RELAY_EVENTS and info is not None:
+        data["from"] = _public(info)
+        await manager.broadcast(room_id, data, exclude=session_id)
+        # cursors + transient interaction pings are too chatty to keep the room
+        # alive on; everything else counts as real activity.
+        if t not in ("cursor_position", "interaction"):
+            touch_room(room_id)
 
-@app.websocket("/ws/{room_id}")
-async def ws_room(ws: WebSocket, room_id: str):
+
+@app.get("/api/rooms/{room_id}/stream")
+async def room_stream(room_id: str, request: Request):
+    """Server-Sent Events stream: server→client messages for one participant.
+
+    Replaces the inbound half of the old WebSocket. The browser opens this with
+    EventSource(...). Query params: session_id, name."""
     if room_row(room_id) is None:
-        await ws.accept()
-        await ws.send_json({"type": "error", "message": "Room not found or expired."})
-        await ws.close(code=4404)
-        return
+        # mimic the old 4404: emit a single error event then end the stream.
+        async def _gone():
+            yield 'data: {"type":"error","message":"Room not found or expired."}\n\n'
+        return StreamingResponse(_gone(), media_type="text/event-stream")
 
-    session_id = (ws.query_params.get("session_id") or secrets.token_hex(8))[:64]
+    session_id = (request.query_params.get("session_id") or secrets.token_hex(8))[:64]
     info = {
         "session_id": session_id,
-        "name": sanitize_name(ws.query_params.get("name")),
+        "name": sanitize_name(request.query_params.get("name")),
         "participant_id": secrets.token_hex(6),
         "color": color_for(session_id),
     }
 
-    # ---- join approval ----
-    # The owner is auto-admitted. Anyone else needs the host to admit them.
-    # If no host is currently connected, we auto-admit (so a room is never
-    # un-joinable / approval never deadlocks). Returning sessions that the host
-    # already admitted earlier in this room are also let straight back in.
+    q = _sse_queue(room_id, session_id)
+
+    # ---- join approval (same policy as the WS version) ----
     row = room_row(room_id)
     owner_session = row["owner_session_id"] if row else None
     is_owner = bool(owner_session) and session_id == owner_session
-
-    await ws.accept()
-    host = manager.host_ws(room_id, owner_session)
+    host = manager.host_client(room_id, owner_session)
     needs_approval = (not is_owner) and (host is not None)
 
     if needs_approval:
-        decision = asyncio.Event()
-        info["_decision_event"] = decision
-        info["_admitted"] = False
-        manager.add_pending(room_id, ws, info)
-        # tell the joiner they're waiting
-        await manager.send_to(ws, {"type": "join_pending", "you": {k: info[k] for k in ("session_id", "name", "participant_id", "color")}})
-        # ask the host to admit / decline
-        await manager.send_to(host, {
-            "type": "join_request",
-            "participant": {k: info[k] for k in ("session_id", "name", "participant_id", "color")},
-        })
-        # wait for the host's decision (with a timeout → auto-decline)
+        pend = _PendingJoin(info)
+        manager.add_pending(room_id, pend)
+        await q.put({"type": "join_pending", "you": _public(info)})
+        await host.put({"type": "join_request", "participant": _public(info)})
         try:
-            await asyncio.wait_for(decision.wait(), timeout=120)
-        except asyncio.TimeoutError:
-            pass
+            await asyncio.wait_for(pend.event.wait(), timeout=120)
         except Exception:
             pass
-        if not info.get("_admitted"):
-            manager.remove_pending(room_id, ws)
-            await manager.send_to(ws, {"type": "join_declined"})
-            try:
-                await ws.close(code=4403)
-            except Exception:
-                pass
-            return
-        manager.admit(room_id, ws, info)
+        if not pend.admitted:
+            manager.remove_pending(room_id, session_id)
+            await q.put({"type": "join_declined"})
+            await q.put({"__close__": True})
+            # fall through to the generator, which will flush join_declined then end
+            client = None
+        else:
+            client = manager.admit(room_id, info)
+            client.queue = q  # reuse the queue the SSE stream is already draining
     else:
-        # owner, or no host present → straight in
-        manager.rooms.setdefault(room_id, {})[ws] = info
+        client = _Client(info)
+        client.queue = q
+        manager.add_client(room_id, client)
+
+    if client is not None:
+        now = time.time()
+        with DB_LOCK:
+            _conn.execute(
+                "INSERT OR REPLACE INTO participants (participant_id, room_id, "
+                "session_id, display_name, joined_at, last_seen) VALUES (?,?,?,?,?,?)",
+                (info["participant_id"], room_id, session_id, info["name"], now, now),
+            )
+            _conn.commit()
+        touch_room(room_id)
+        await q.put({"type": "init", "you": _public(info), "participants": manager.presence(room_id)})
+        await manager.broadcast(room_id, {"type": "user_joined", "participant": _public(info)}, exclude=session_id)
+        await manager.broadcast(room_id, {"type": "presence_update", "participants": manager.presence(room_id)})
 
-    now = time.time()
-    with DB_LOCK:
-        _conn.execute(
-            "INSERT OR REPLACE INTO participants (participant_id, room_id, "
-            "session_id, display_name, joined_at, last_seen) VALUES (?,?,?,?,?,?)",
-            (info["participant_id"], room_id, session_id, info["name"], now, now),
-        )
-        _conn.commit()
-    touch_room(room_id)
+    async def event_gen():
+        try:
+            while True:
+                # heartbeat keeps the connection alive through idle proxies
+                try:
+                    msg = await asyncio.wait_for(q.get(), timeout=20)
+                except asyncio.TimeoutError:
+                    yield ": keep-alive\n\n"
+                    continue
+                if isinstance(msg, dict) and msg.get("__close__"):
+                    break
+                yield f"data: {json.dumps(msg, separators=(',', ':'))}\n\n"
+        except asyncio.CancelledError:
+            pass
+        finally:
+            _sse_queues.pop((room_id, session_id), None)
+            if manager.get_client(room_id, session_id) is not None:
+                manager.disconnect(room_id, session_id)
+                await manager.broadcast(room_id, {"type": "user_left", "participant": _public(info)})
+                await manager.broadcast(room_id, {"type": "presence_update", "participants": manager.presence(room_id)})
+
+    return StreamingResponse(event_gen(), media_type="text/event-stream", headers={
+        "Cache-Control": "no-cache",
+        "X-Accel-Buffering": "no",   # disable proxy buffering so events flush live
+        "Connection": "keep-alive",
+    })
 
-    await ws.send_json({"type": "init", "you": info, "participants": manager.presence(room_id)})
-    await manager.broadcast(room_id, {"type": "user_joined", "participant": info}, exclude=ws)
-    await manager.broadcast(room_id, {"type": "presence_update", "participants": manager.presence(room_id)})
 
+@app.post("/api/rooms/{room_id}/send")
+async def room_send(room_id: str, request: Request):
+    """Client→server messages: one event the client used to send over the
+    socket (object_updated, chat_message, cursor_position, admit_join, ...)."""
+    if room_row(room_id) is None:
+        raise HTTPException(status_code=404, detail="Room not found or expired.")
     try:
-        while True:
-            data = await ws.receive_json()
-            t = data.get("type")
-            if t == "rename":
-                info["name"] = sanitize_name(data.get("name"))
-                with DB_LOCK:
-                    _conn.execute(
-                        "UPDATE participants SET display_name=?, last_seen=? WHERE participant_id=?",
-                        (info["name"], time.time(), info["participant_id"]),
-                    )
-                    _conn.commit()
-                await manager.broadcast(room_id, {"type": "presence_update", "participants": manager.presence(room_id)})
-                continue
-            # ---- host approves / declines a pending joiner ----
-            if t in ("admit_join", "decline_join"):
-                target_sid = data.get("session_id")
-                pend_ws, pend_info = manager.find_pending(room_id, target_sid)
-                if pend_info is not None:
-                    pend_info["_admitted"] = (t == "admit_join")
-                    ev = pend_info.get("_decision_event")
-                    if ev is not None:
-                        ev.set()
-                continue
-            # explicitly known (RELAY_EVENTS) or simply not a server-owned type
-            # (NON_RELAY_EVENTS) — this way the frontend can introduce new
-            # interaction events without the server silently dropping them.
-            if t and t not in NON_RELAY_EVENTS:
-                data["from"] = {k: info[k] for k in ("session_id", "name", "participant_id", "color")}
-                await manager.broadcast(room_id, data, exclude=ws)
-                # cursors + transient interaction pings are too chatty to keep
-                # the room alive on; everything else counts as real activity.
-                if t not in ("cursor_position", "interaction"):
-                    touch_room(room_id)
-    except WebSocketDisconnect:
-        pass
+        body = await request.json()
     except Exception:
-        pass
-    finally:
-        manager.disconnect(room_id, ws)
-        await manager.broadcast(room_id, {"type": "user_left", "participant": info})
-        await manager.broadcast(room_id, {"type": "presence_update", "participants": manager.presence(room_id)})
+        raise HTTPException(status_code=400, detail="Invalid JSON.")
+    session_id = (body or {}).get("session_id")
+    if not session_id:
+        raise HTTPException(status_code=400, detail="Missing session_id.")
+    data = (body or {}).get("event") or {}
+    if not isinstance(data, dict) or not data.get("type"):
+        raise HTTPException(status_code=400, detail="Missing event.type.")
+    await _relay_incoming(room_id, session_id, data)
+    return {"ok": True}
 
 
 # ------------------------------ Assets ------------------------------------- #

index.html

@@ -361,12 +361,16 @@
   </script>
 
   <!-- ----------------------------------------------------------------- -->
-  <!-- Collaboration transport (plain JS): a thin WebSocket pub/sub over   -->
-  <!-- the gradio.Server /ws/{room_id} endpoint. No-ops cleanly in solo.   -->
+  <!-- Collaboration transport (plain JS): SSE (inbound) + POST (outbound)   -->
+  <!-- over /api/rooms/{id}/stream and /send. No-ops cleanly in solo.        -->
   <!-- ----------------------------------------------------------------- -->
   <script>
     window.Collab = (function () {
-      let ws = null, me = null, participants = [], reconnectT = null;
+      // Transport: Server-Sent Events (inbound) + POST (outbound). HF Spaces'
+      // proxy drops custom WebSocket upgrades, so we use the same SSE+POST
+      // channel Gradio itself uses on Spaces. EventSource auto-reconnects on
+      // its own, so there's no manual reconnect loop here.
+      let es = null, me = null, participants = [];
       const listeners = {};
 
       function roomId() {
@@ -385,27 +389,38 @@
       }
       function emit(type, d) { (listeners[type] || []).forEach((fn) => { try { fn(d); } catch (e) {} }); }
       function send(type, payload) {
-        try { if (ws && ws.readyState === 1) ws.send(JSON.stringify(Object.assign({ type }, payload || {}))); } catch (e) {}
-      }
-      function scheduleReconnect() {
-        if (reconnectT) return;
-        reconnectT = setTimeout(() => { reconnectT = null; if (roomId()) connect(); }, 1500);
+        const id = roomId();
+        if (!id) return;
+        const event = Object.assign({ type }, payload || {});
+        try {
+          fetch(`/api/rooms/${encodeURIComponent(id)}/send`, {
+            method: "POST",
+            headers: { "Content-Type": "application/json" },
+            // keepalive lets the final send (e.g. on tab close) still go out
+            keepalive: type === "rename",
+            body: JSON.stringify({ session_id: sid(), event }),
+          }).catch(() => {});
+        } catch (e) {}
       }
       function connect() {
         const id = roomId();
         if (!id) { emit("status", { connected: false, solo: true }); return; }
-        const proto = location.protocol === "https:" ? "wss" : "ws";
-        const url = `${proto}://${location.host}/ws/${encodeURIComponent(id)}?session_id=${encodeURIComponent(sid())}&name=${encodeURIComponent(name() || "Guest")}`;
-        try { ws = new WebSocket(url); } catch (e) { scheduleReconnect(); return; }
-        ws.onopen = () => emit("status", { connected: true, solo: false });
-        ws.onmessage = (e) => {
+        try { if (es) es.close(); } catch (_) {}
+        const url = `/api/rooms/${encodeURIComponent(id)}/stream?session_id=${encodeURIComponent(sid())}&name=${encodeURIComponent(name() || "Guest")}`;
+        try { es = new EventSource(url); } catch (e) { return; }
+        es.onopen = () => emit("status", { connected: true, solo: false });
+        es.onmessage = (e) => {
           let m; try { m = JSON.parse(e.data); } catch (_) { return; }
           if (m.type === "init") { me = m.you; participants = m.participants || []; emit("me", me); emit("presence", participants); }
           else if (m.type === "presence_update") { participants = m.participants || []; emit("presence", participants); }
           emit(m.type, m); emit("any", m);
         };
-        ws.onclose = () => { emit("status", { connected: false, solo: false }); scheduleReconnect(); };
-        ws.onerror = () => { try { ws.close(); } catch (_) {} };
+        // EventSource fires onerror on transient drops and reconnects itself;
+        // we just surface the status. A closed (readyState 2) stream means the
+        // server ended it (e.g. join_declined) — don't fight it.
+        es.onerror = () => {
+          emit("status", { connected: false, solo: false });
+        };
       }
       function rename(n) { localStorage.setItem("sozai_name", n); if (me) me.name = n; send("rename", { name: n }); }
 

packages.txt

@@ -1 +1,5 @@
+ffmpeg
+libsndfile1
+sox
+libsox-fmt-all
 libheif-dev

requirements.txt

@@ -1,5 +1,5 @@
 gradio
-uvicorn[standard]
+websockets>=12.0
 transformers
 # Flux2KleinPipeline (img2img watercolour develop) is not in a released diffusers
 # yet — it must come from git main. See backend/hf/3_infer_img2img.py.