"""
DOOM-FlyWire Hugging Face Space Entry Point
============================================
Edge-to-edge DOOM-FlyWire Connectome Autonomous Agent & 3D Visualizer.
Game routes (/doom, /ws/game, assets) are injected directly at the head of
Gradio's FastAPI router before demo.launch() is called.
"""
import os
import sys
# Disable Gradio 6 Node.js SSR server so FastAPI handles all requests directly
os.environ["GRADIO_SSR_MODE"] = "false"
os.environ["GRADIO_NODE_PATH"] = ""
try:
import gradio.node_server
gradio.node_server.start_node_server = lambda *a, **k: (None, None, None)
except Exception:
pass
try:
import gradio.routes
gradio.routes.start_node_server = lambda *a, **k: (None, None, None)
except Exception:
pass
# Support both ZeroGPU and standard CPU environments
try:
import spaces
_gpu_decorator = spaces.GPU
except Exception:
spaces = None
def _gpu_decorator(fn): return fn
import asyncio
import json
import threading
import time
import gradio as gr
from gradio.routes import App as GradioApp
from starlette.responses import FileResponse, JSONResponse
from starlette.routing import Route, WebSocketRoute
from starlette.websockets import WebSocket, WebSocketDisconnect
# --- ZeroGPU Stub ---
# HF detects @spaces.GPU at scan/import time. No UI elements required.
@_gpu_decorator
def _gpu_ready():
import torch
return torch.cuda.is_available()
WEB_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "web")
os.makedirs(WEB_DIR, exist_ok=True)
# --- Fullscreen Edge-to-Edge Visualizer ---
with gr.Blocks(title="DOOM-FlyWire Connectome", fill_width=True, fill_height=True) as demo:
gr.HTML("""
""")
# --- Game Agent (Lazy Initialization) ---
_agent = None
_agent_lock = threading.Lock()
_agent_error = None
def _get_agent():
global _agent, _agent_error
with _agent_lock:
if _agent is not None:
return _agent
try:
from src.doom_agent import DoomConnectomeAgent
print("[Agent] Initializing DoomConnectomeAgent on cpu...")
_agent = DoomConnectomeAgent(
scenario_name="defend_the_center.cfg",
window_visible=False,
device="cpu",
)
print("[Agent] Ready.")
return _agent
except Exception as exc:
_agent_error = str(exc)
print(f"[Agent] Init failed: {exc}")
return None
# --- Singleton Broadcaster State ---
_active_websockets = set()
_game_loop_task = None
_loop_lock = asyncio.Lock()
async def _singleton_game_loop():
"""
Dedicated single-threaded game propagation loop:
- Guarantees agent.step() is called sequentially (thread-safe for ViZDoom & PyTorch).
- Idles gracefully when no clients are connected to save cloud CPU/RAM.
- Broadcasts the exact same live simulation frame to all concurrent visitors.
"""
agent = _get_agent()
if agent is None:
return
fps_base = 30.0
frame_interval = 1.0 / fps_base
last_log_time = 0.0
last_frame_time = time.perf_counter()
while True:
try:
if not _active_websockets:
# No active visitors: sleep to conserve resources
await asyncio.sleep(0.5)
last_frame_time = time.perf_counter()
continue
t0 = time.perf_counter()
data = await asyncio.to_thread(agent.step)
if not data:
await asyncio.sleep(0.01)
continue
# Real continuous delivery FPS matching native DOOM 35 FPS
now = time.perf_counter()
dt_stream = now - last_frame_time
last_frame_time = now
fps = round(1.0 / max(dt_stream, 1e-4), 1)
data["instant_fps"] = min(fps, 35.0)
# Live telemetry logger (flushes every 2s so Hugging Face web log console displays it in real time)
now = time.perf_counter()
if now - last_log_time >= 2.0:
step = data.get("step", 0)
act = data.get("action_taken", ["IDLE"])
print(f"[DOOM-FlyWire] Step {step:,} | FPS: {fps:.1f} | Action: {act} | Active Viewers: {len(_active_websockets)}", flush=True)
last_log_time = now
# Broadcast frame to all active connections
dead = set()
for ws in list(_active_websockets):
try:
await ws.send_json(data)
except Exception:
dead.add(ws)
if dead:
_active_websockets.difference_update(dead)
elapsed = time.perf_counter() - t0
slack = frame_interval - elapsed
if slack > 0:
await asyncio.sleep(slack)
else:
await asyncio.sleep(0.001)
except Exception as exc:
print(f"[GameLoop Error] {exc}")
await asyncio.sleep(0.05)
# --- FastAPI Route Injection ---
_orig_create_app = GradioApp.__dict__["create_app"] # staticmethod descriptor
@staticmethod
def _patched_create_app(*args, **kwargs):
app = _orig_create_app.__func__(*args, **kwargs)
print(f"[FastAPI] Injected game routes at head of router. WEB_DIR={WEB_DIR}")
def _serve(rel, media=""):
f = os.path.join(WEB_DIR, rel) if rel else None
if f and os.path.exists(f):
headers = {
"Cache-Control": "no-cache, no-store, must-revalidate, max-age=0",
"Pragma": "no-cache",
"Expires": "0"
}
return FileResponse(f, media_type=media, headers=headers) if media else FileResponse(f, headers=headers)
return JSONResponse({"error": f"{rel} not found"}, status_code=404)
async def doom_page(request): return _serve("index.html")
async def health_page(request): return JSONResponse({"status": "ok"})
async def ba_page(request): return _serve("brain_anatomy.json", "application/json")
async def ba3_page(request): return _serve("brain_anatomy_3d.json", "application/json")
async def bpos_page(request): return _serve("brain_139k_pos.bin", "application/octet-stream")
async def bcol_page(request): return _serve("brain_139k_col.bin", "application/octet-stream")
async def syn_page(request): return _serve("synapses_top5k.bin", "application/octet-stream")
async def three_page(request): return _serve(os.path.join("js", "three.min.js"), "application/javascript")
async def orbit_page(request): return _serve(os.path.join("js", "OrbitControls.js"), "application/javascript")
# --- WebSocket Game Stream (Multiplexed Broadcaster) ---
async def ws_game(websocket: WebSocket):
global _game_loop_task
await websocket.accept()
print("[WS] Client connected.")
agent = _get_agent()
if agent is None:
await websocket.send_json({"type": "error", "message": _agent_error or "init failed"})
await websocket.close()
return
_active_websockets.add(websocket)
# Ensure background broadcast game loop is running
async with _loop_lock:
if _game_loop_task is None or _game_loop_task.done():
_game_loop_task = asyncio.create_task(_singleton_game_loop())
try:
# Client listener loop for user commands (reset, parameter updates)
while True:
raw = await websocket.receive_text()
try:
msg = json.loads(raw)
kind = msg.get("type", "")
if kind == "update_params":
agent.update_sandbox(msg.get("params", {}))
elif kind in ("reset", "start"):
agent.reset_episode()
except Exception:
pass
except WebSocketDisconnect:
pass
except Exception:
pass
finally:
_active_websockets.discard(websocket)
print("[WS] Client cleanly disconnected.")
# Inject routes at the head of app.router.routes so they take precedence over everything
custom_routes = [
Route("/doom", endpoint=doom_page, methods=["GET", "HEAD"]),
Route("/doom/", endpoint=doom_page, methods=["GET", "HEAD"]),
Route("/health", endpoint=health_page, methods=["GET"]),
Route("/brain_anatomy.json", endpoint=ba_page, methods=["GET"]),
Route("/brain_anatomy_3d.json", endpoint=ba3_page, methods=["GET"]),
Route("/brain_139k_pos.bin", endpoint=bpos_page, methods=["GET"]),
Route("/brain_139k_col.bin", endpoint=bcol_page, methods=["GET"]),
Route("/synapses_top5k.bin", endpoint=syn_page, methods=["GET"]),
Route("/js/three.min.js", endpoint=three_page, methods=["GET"]),
Route("/js/OrbitControls.js", endpoint=orbit_page, methods=["GET"]),
WebSocketRoute("/ws/game", endpoint=ws_game),
]
for r in reversed(custom_routes):
app.router.routes.insert(0, r)
return app
GradioApp.create_app = _patched_create_app
# --- Application Launch ---
if __name__ == "__main__":
port = int(os.environ.get("PORT", 7860))
print(f"[HF Space] Launching DOOM-FlyWire on 0.0.0.0:{port}...")
demo.launch(server_name="0.0.0.0", server_port=port, quiet=False)