refactor: use websocket again

tmp/server.py

@@ -0,0 +1,166 @@
+import threading
+import os
+from time import sleep
+import requests
+
+from contextlib import asynccontextmanager
+from fastapi import (
+    FastAPI,
+    WebSocket,
+    status,
+)
+from fastapi.responses import FileResponse
+from fastapi.staticfiles import StaticFiles
+from listener_counter import ListenerCounter
+from logger import log_info, log_warn
+from websocket_connection_manager import WebSocketConnectionManager
+
+# the index of the current audio track from 0 to 9
+current_index = -1
+# the timer that periodically advances the current audio track
+t = None
+inference_url = ""
+api_key = ""
+ws_connection_manager = WebSocketConnectionManager()
+listener_counter = ListenerCounter()
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    global ws, inference_url, api_key
+
+    inference_url = os.environ.get("INFERENCE_SERVER_URL")
+    api_key = os.environ.get("API_KEY")
+
+    if not inference_url:
+        inference_url = "http://localhost:8001"
+
+    advance()
+
+    yield
+
+    if t:
+        t.cancel()
+
+
+def generate_new_audio():
+    if not inference_url:
+        return
+
+    global current_index
+
+    offset = 0
+    if current_index == 0:
+        offset = 5
+    elif current_index == 5:
+        offset = 0
+    else:
+        return
+
+    log_info("requesting new audio...")
+
+    try:
+        requests.post(
+            f"{inference_url}/generate",
+            headers={"Authorization": f"key {api_key}"},
+        )
+    except:
+        log_warn(
+            "inference server potentially unreachable. recycling cached audio for now."
+        )
+        return
+
+    is_available = False
+    while not is_available:
+        try:
+            res = requests.post(
+                f"{inference_url}/clips/0",
+                stream=True,
+                headers={"Authorization": f"key {api_key}"},
+            )
+        except:
+            log_warn(
+                "inference server potentially unreachable. recycling cached audio for now."
+            )
+            return
+
+        if res.status_code != status.HTTP_200_OK:
+            print(res.status_code)
+            print("still generating...")
+            sleep(30)
+            continue
+
+        print("inference complete! downloading new clips")
+
+        is_available = True
+        with open(f"{offset}.mp3", "wb") as f:
+            for chunk in res.iter_content(chunk_size=128):
+                f.write(chunk)
+
+    for i in range(4):
+        res = requests.post(
+            f"{inference_url}/clips/{i + 1}",
+            stream=True,
+            headers={"Authorization": f"key {api_key}"},
+        )
+
+        if res.status_code != status.HTTP_200_OK:
+            continue
+
+        with open(f"{i + 1 + offset}.mp3", "wb") as f:
+            for chunk in res.iter_content(chunk_size=128):
+                f.write(chunk)
+
+    log_info("audio generated.")
+
+
+def advance():
+    global current_index, t
+
+    if current_index == 9:
+        current_index = 0
+    else:
+        current_index = current_index + 1
+    threading.Thread(target=generate_new_audio).start()
+
+    t = threading.Timer(60, advance)
+    t.start()
+
+
+app = FastAPI(lifespan=lifespan)
+
+
+@app.get("/current.mp3")
+def get_current_audio():
+    return FileResponse(f"{current_index}.mp3")
+
+
+@app.websocket("/ws")
+async def ws_endpoint(ws: WebSocket):
+    await ws_connection_manager.connect(ws)
+
+    addr = ""
+    if ws.client:
+        addr, _ = ws.client
+    else:
+        await ws.close()
+        ws_connection_manager.disconnect(ws)
+        return
+
+    try:
+        while True:
+            msg = await ws.receive_text()
+            match msg:
+                case "listening":
+                    listener_counter.add_listener(addr)
+                    await ws_connection_manager.broadcast(f"{listener_counter.count()}")
+                case "paused":
+                    listener_counter.remove_listener(addr)
+                    await ws_connection_manager.broadcast(f"{listener_counter.count()}")
+    except:
+        listener_counter.remove_listener(addr)
+        ws_connection_manager.disconnect(ws)
+        await ws_connection_manager.broadcast(f"{listener_counter.count()}")
+
+
+app.mount("/", StaticFiles(directory="web", html=True), name="web")