Kaiky Maia
revert GPU duration to default (free ZeroGPU caps <120s); keep 512-token cap
0f66bac verified | """ | |
| space/server_app.py — One for All, custom frontend on gr.Server. | |
| gr.Server = FastAPI with Gradio's API engine: @app.api functions become | |
| queued/streaming endpoints (SSE), and we serve our own HTML/JS at "/" — | |
| the browser talks to the endpoints via @gradio/client (required for ZeroGPU). | |
| Run locally: | |
| cd space && VIZ_DATA_PATH=/path/to/viz_data.json python server_app.py | |
| The legacy Blocks UI (app.py) keeps working — both share _boot.py. | |
| """ | |
| from __future__ import annotations | |
| from pathlib import Path | |
| import spaces | |
| from fastapi.responses import FileResponse, HTMLResponse | |
| from gradio import Server | |
| import _boot | |
| RT = _boot.load_runtime() | |
| FRONTEND = Path(__file__).parent / "frontend" | |
| app = Server() | |
| # ── API: static viz payload (soul space, CKA, curves, meta) ──────────────── | |
| def viz() -> dict: | |
| return _boot.viz_payload(RT.viz, RT.coords3d, RT.backend, RT.model_ready) | |
| # ── API: streaming probe — text + live gates, final soul-space point ─────── | |
| def probe(text: str) -> dict: # annotation = type of each streamed chunk | |
| if not text.strip(): | |
| yield {"error": "empty prompt", "done": True} | |
| return | |
| _boot.to_device(RT) # lazy-loads the student here (real CUDA inside @spaces.GPU) | |
| if not RT.model_ready: | |
| yield {"error": RT.load_error or "model not ready", "done": True} | |
| return | |
| partial, gates = "", [] | |
| for partial, gates in _boot.stream(RT, text): | |
| yield {"text": partial, "gates": gates, "done": False} | |
| point = None | |
| if RT.reducer is not None: | |
| point, gates = _boot.final_probe(RT, text) | |
| yield {"text": partial, "gates": gates, "point": point, "done": True} | |
| # ── API: streaming arena — base (LoRA off) vs deku, interleaved ──────────── | |
| def arena(text: str) -> dict: # annotation = type of each streamed chunk | |
| if not text.strip(): | |
| yield {"error": "empty prompt", "done": True} | |
| return | |
| _boot.to_device(RT) # lazy-loads the student here (real CUDA inside @spaces.GPU) | |
| if not RT.model_ready: | |
| yield {"error": RT.load_error or "model not ready", "done": True} | |
| return | |
| base = deku = "" | |
| gates: list[float] = [] | |
| for base, deku, gates in _boot.stream_pair(RT, text): | |
| yield {"base": base, "deku": deku, "gates": gates, "done": False} | |
| yield {"base": base, "deku": deku, "gates": gates, "done": True} | |
| # ── Custom frontend (overrides Gradio's default UI at "/") ──────────────── | |
| async def home(): | |
| return (FRONTEND / "index.html").read_text() | |
| async def static(name: str): | |
| target = (FRONTEND / name).resolve() | |
| if target.parent != FRONTEND.resolve() or not target.exists(): | |
| return HTMLResponse("not found", status_code=404) | |
| return FileResponse(target) | |
| if __name__ == "__main__": | |
| app.launch() | |