app.py

"""
ABot-World — Interactive Action-Conditioned World Rollout
gradio.Server + WebSocket live-backend edition.

Given an uploaded starting image (i2v conditioning), a scene prompt, and live
WASD / IJKL controls, the model autoregressively rolls out an action-conditioned
navigable world and streams decoded frames to the browser over a WebSocket.

This mirrors the live backend/infrastructure of
https://huggingface.co/spaces/Overworld/waypoint-1-5 (gradio.Server for
ZeroGPU-friendly start/stop + a raw WebSocket for real-time binary JPEG frame
streaming and control input), with a cleaner custom UI and image-upload seeding.

Multi-user safe: every endpoint is keyed by a per-client `session_id` so
concurrent players never share seed images, frame queues, or status messages.

ZeroGPU quota: the incoming request's ZeroGPU proxy token (the `x-ip-token` /
`x-api-token` header injected by the HF iframe) is captured per-session and
propagated into the worker thread's gradio request context, so the GPU work is
billed against the *requesting user's* quota — not the Space owner's.

Upstream: https://github.com/amap-cvlab/ABot-World
Model:    https://huggingface.co/acvlab/ABot-World-0-5B-LF  (built on Wan2.2-TI2V-5B)
"""
import os
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces  # must precede torch / CUDA-touching imports

import sys
import io
import time
import queue
import asyncio
import struct
import tempfile
import threading
import contextvars
import uuid
from collections import deque
from dataclasses import dataclass, field
from multiprocessing import Queue as MPQueue
from pathlib import Path
from typing import Dict, Optional, Set

import numpy as np
import torch
from PIL import Image
from omegaconf import OmegaConf
from huggingface_hub import snapshot_download

from fastapi import UploadFile, File, WebSocket, WebSocketDisconnect
from fastapi.responses import HTMLResponse, JSONResponse, FileResponse
from gradio import Server
from gradio.context import LocalContext

# ── Repo paths ───────────────────────────────────────────────────────────────
APP_DIR = Path(__file__).resolve().parent
if str(APP_DIR) not in sys.path:
    sys.path.insert(0, str(APP_DIR))

MODEL_ID = "acvlab/ABot-World-0-5B-LF"
CKPT_DIR = APP_DIR / "checkpoints" / "ABot-World-0-5B-LF"

# Preset starting-world images bundled with the Space (sourced from the ABot-World
# repo). Shown in the UI as clickable thumbnails that seed the i2v rollout directly.
EXAMPLES_DIR = APP_DIR / "examples"
EXAMPLE_SEEDS = [
    {"name": "desert_valley.png", "label": "Desert valley"},
    {"name": "forest_stream.png", "label": "Forest stream"},
    {"name": "mountain_meadow.png", "label": "Mountain meadow"},
    {"name": "example.png", "label": "Sample scene"},
]

# ── Stream / rollout configuration ───────────────────────────────────────────
# 704x1280 is the native training resolution used by the upstream web client.
STREAM_HEIGHT = 704
STREAM_WIDTH = 1280
JPEG_QUALITY = 82
MAX_BLOCKS_PER_SESSION = 512     # hard cap so a session can't run forever
SESSION_IDLE_TIMEOUT = 600       # seconds; janitor reaps abandoned sessions
GPU_DURATION = 90                # seconds per @spaces.GPU allocation

# ── Real-time pacing configuration ───────────────────────────────────────────
# The GPU decodes a whole block (NUM_FPB frames) at once, so all of a block's
# frames become available in a burst. If we forward them to the browser the
# instant they finish, the client sees N frames clustered together followed by a
# gap while the next block generates — the fps counter averages out fine, but
# the *felt* cadence is bursty. To deliver a steady real-time stream we pace the
# frames of each block evenly across the time we expect one block to take
# (mirroring the official ABot-World web_client's block-frame spreader), and we
# smooth the per-block generation time with an EMA so a single slow/fast block
# doesn't cause a visible speed-up/slow-down. See gpu_worker_thread().
PACING_EMA_ALPHA = 0.25          # smoothing factor for per-block generation time
MIN_PACING_SLEEP = 0.004         # don't bother sleeping for sub-4ms slices
DEFAULT_BLOCK_SECONDS = 0.5      # initial per-block estimate before first measure

# Actions map to the 8-key one-hot the model was trained on (W A S D I J K L).
# The browser sends the currently-held key set; we translate to this dict.
KEY_ORDER = ["W", "A", "S", "D", "I", "J", "K", "L"]

DEFAULT_PROMPT = (
    "A realistic outdoor world scene with a navigable path, natural lighting, "
    "detailed ground texture, and stable forward motion."
)

# ── Download weights (once, at startup) ──────────────────────────────────────
print(f"[startup] downloading {MODEL_ID} weights ...", flush=True)
snapshot_download(
    repo_id=MODEL_ID,
    repo_type="model",
    local_dir=str(CKPT_DIR),
)
print("[startup] weights downloaded.", flush=True)

os.environ["TAEW2_2_CHECKPOINT"] = str(CKPT_DIR / "taew2_2.pth")

# ── Build the causal streaming pipeline (module scope, eager to CUDA) ─────────
from pipeline import CausalInferencePipeline
from utils.misc import set_seed
from utils.wan_wrapper import create_vae_from_config

_CONFIG_PATH = APP_DIR / "configs" / "long_forcing_dmd.yaml"
_DEFAULT_CFG_PATH = APP_DIR / "configs" / "default_config.yaml"


def _build_config():
    cfg = OmegaConf.merge(
        OmegaConf.load(str(_DEFAULT_CFG_PATH)),
        OmegaConf.load(str(_CONFIG_PATH)),
    )
    cfg.taew2_2_checkpoint = str(CKPT_DIR / "taew2_2.pth")
    cfg.lightvae_encoder_checkpoint = str(CKPT_DIR / "Wan2.2_VAE.pth")
    cfg.model_kwargs.model_name = str(CKPT_DIR)
    cfg.text_encoder_kwargs.tokenizer_path = str(CKPT_DIR / "google" / "umt5-xxl") + "/"
    cfg.text_encoder_kwargs.encoder_pth_path = str(CKPT_DIR / "models_t5_umt5-xxl-enc-bf16.pth")
    cfg.vae_kwargs.pretrained_path = str(CKPT_DIR / "Wan2.2_VAE.pth")
    cfg.vae_type = "taew2_2"
    cfg.use_fp8_gemm = False
    return cfg


print("[startup] building pipeline ...", flush=True)
set_seed(42)
torch.set_grad_enabled(False)

CONFIG = _build_config()
_vae = create_vae_from_config(CONFIG)
pipeline = CausalInferencePipeline(CONFIG, device=torch.device("cuda"), vae=_vae)

try:
    from wan.modules.helios_kernels import (
        replace_all_norms_with_flash_norms,
        replace_rope_with_flash_rope,
    )
    replace_all_norms_with_flash_norms(pipeline.generator.model)
    replace_rope_with_flash_rope()
    print("[startup] helios Triton kernels enabled.", flush=True)
except Exception as e:  # pragma: no cover
    print(f"[startup] helios kernels disabled ({e!r}); using eager norms/rope.", flush=True)

pipeline = pipeline.to(dtype=torch.bfloat16)
pipeline.text_encoder.to(device="cuda")
pipeline.generator.to(device="cuda")
pipeline.vae.to(device="cuda")
if pipeline.encoder is not None:
    pipeline.encoder.to(device="cuda")
pipeline.torch_dtype = torch.bfloat16
print("[startup] pipeline ready.", flush=True)

NUM_FPB = int(getattr(CONFIG, "num_frame_per_block", 3))
_vae_for_shape = pipeline.encoder if pipeline.encoder is not None else pipeline.vae
LATENT_CHANNELS = _vae_for_shape.z_dim
UPSAMPLE = getattr(_vae_for_shape, "upsampling_factor", 16)
LATENT_H = STREAM_HEIGHT // UPSAMPLE
LATENT_W = STREAM_WIDTH // UPSAMPLE

# Only one rollout may touch the shared pipeline at a time.
_infer_lock = threading.Lock()


def _keys_from_buttons(buttons) -> Dict[str, bool]:
    """Translate a set of held key names (e.g. {'W','A'}) into the model dict."""
    held = {k.upper() for k in (buttons or [])}
    return {k: True for k in KEY_ORDER if k in held}


def _decode_block_to_frames(lat_block):
    frames = []

    class _Cap:
        def append_data(self, f):
            frames.append(np.asarray(f))

    pipeline.decode_block_and_write(lat_block, _Cap())
    return frames


# ── Command types (browser -> worker) ────────────────────────────────────────
@dataclass
class ControlCommand:
    buttons: Set[str]
    prompt: str


@dataclass
class StopCommand:
    pass


# ── Per-session state ────────────────────────────────────────────────────────
# NOTE on queues: the @spaces.GPU rollout runs in a forked subprocess, so any
# object it reads must cross the fork boundary. `command_queue` is therefore a
# multiprocessing Queue (browser controls / stop reach the GPU loop through it).
# `frame_queue` / `status_queue` are plain queue.Queue used only in the parent
# process (frames arrive back via the ZeroGPU generator IPC and are forwarded
# to the WebSocket by the worker thread).
@dataclass
class GameSession:
    session_id: str
    command_queue: "MPQueue"
    frame_queue: "queue.Queue"
    status_queue: "queue.Queue"
    stop_event: threading.Event
    seed_path: str
    prompt: str
    seed: int
    worker_thread: Optional[threading.Thread] = None
    frame_times: deque = field(default_factory=lambda: deque(maxlen=30))
    last_active: float = field(default_factory=time.time)

    def touch(self):
        self.last_active = time.time()

    def stop(self):
        self.stop_event.set()
        try:
            self.command_queue.put_nowait(StopCommand())
        except Exception:
            pass
        if self.worker_thread and self.worker_thread.is_alive():
            self.worker_thread.join(timeout=4.0)


_sessions: Dict[str, GameSession] = {}
_sessions_lock = threading.Lock()

# Contextvar carrying the active session's status queue (inherited by the worker
# thread via contextvars.copy_context()).
_current_status_queue: "contextvars.ContextVar[Optional[queue.Queue]]" = contextvars.ContextVar(
    "abot_status_queue", default=None
)


def broadcast_status(msg: str):
    q = _current_status_queue.get()
    if q is None:
        return
    try:
        q.put_nowait(msg)
    except queue.Full:
        pass


def _get_session(session_id: str) -> Optional[GameSession]:
    with _sessions_lock:
        return _sessions.get(session_id)


def _drop_session(session_id: str) -> Optional[GameSession]:
    with _sessions_lock:
        return _sessions.pop(session_id, None)


def _reap_idle_sessions():
    while True:
        time.sleep(60)
        now = time.time()
        to_drop = []
        with _sessions_lock:
            for sid, sess in list(_sessions.items()):
                worker_dead = sess.worker_thread is None or not sess.worker_thread.is_alive()
                idle = (now - sess.last_active) > SESSION_IDLE_TIMEOUT
                if worker_dead and idle:
                    to_drop.append(sid)
            for sid in to_drop:
                _sessions.pop(sid, None)
        if to_drop:
            print(f"Janitor reaped {len(to_drop)} idle session(s)", flush=True)


threading.Thread(target=_reap_idle_sessions, daemon=True).start()


# ── GPU worker ───────────────────────────────────────────────────────────────
def gpu_worker_thread(session: "GameSession"):
    """Parent-thread driver: consumes frames yielded by the ZeroGPU generator,
    computes FPS, and forwards frames to the WebSocket via `frame_queue`.

    Status/stop live in the parent process; the GPU loop is steered purely
    through the (picklable, cross-fork) `command_queue`.
    """
    try:
        broadcast_status("GPU allocated — starting world…")
        gen = create_gpu_rollout_loop(
            session.command_queue, session.seed_path, session.prompt, session.seed,
        )
        first = True
        # Steady send clock: `next_send` is the monotonic time at which the next
        # frame *should* be delivered. Each frame's slot is one smoothed
        # inter-frame interval after the previous, so frames leave at a constant
        # cadence regardless of the bursty block boundaries. `block_seconds` is
        # an EMA of measured per-block generation time (frames/block ÷ that gives
        # the target inter-frame interval).
        block_seconds = DEFAULT_BLOCK_SECONDS
        next_send = None
        while not session.stop_event.is_set():
            try:
                frame, block_idx, frame_idx, frames_in_block, block_elapsed = next(gen)
            except StopIteration:
                print("Rollout generator exhausted", flush=True)
                break
            except Exception as e:
                if "aborted" in str(e).lower() or "duration" in str(e).lower():
                    print(f"GPU time expired: {e}", flush=True)
                else:
                    print(f"Worker error: {e}", flush=True)
                    broadcast_status(f"error:{e}")
                break

            if first:
                broadcast_status("Rolling out — use WASD / IJKL to steer.")
                first = False

            # Update the smoothed per-block time on the first frame of each block
            # (block_elapsed is constant across a block's frames).
            if frame_idx == 0 and block_elapsed > 0:
                block_seconds = (
                    PACING_EMA_ALPHA * block_elapsed
                    + (1.0 - PACING_EMA_ALPHA) * block_seconds
                )
            fpb = max(1, frames_in_block)
            interval = block_seconds / fpb        # target seconds between frames

            # ── Steady-cadence gate ──────────────────────────────────────────
            # Hold each frame until its scheduled slot so the parent emits at a
            # constant interval instead of dumping a whole block at once.
            now = time.time()
            if next_send is None:
                next_send = now
            sleep_for = next_send - now
            if sleep_for > MIN_PACING_SLEEP:
                # Wake early if a stop is requested so we stay responsive.
                if session.stop_event.wait(timeout=sleep_for):
                    break
                now = time.time()
            # Advance the schedule; if we've fallen far behind (e.g. a long GPU
            # stall), resync to now so we don't try to "catch up" in a burst.
            next_send = max(now, next_send + interval)

            now = time.time()
            session.frame_times.append(now)
            fps = 0.0
            if len(session.frame_times) >= 2:
                elapsed = session.frame_times[-1] - session.frame_times[0]
                fps = (len(session.frame_times) - 1) / elapsed if elapsed > 0 else 0.0
            # Keep only the freshest frame if the consumer fell behind: coalesce
            # stale frames rather than letting them queue up and flush in a burst.
            while session.frame_queue.qsize() > 1:
                try:
                    session.frame_queue.get_nowait()
                except queue.Empty:
                    break
            try:
                session.frame_queue.put_nowait((frame, block_idx, round(fps, 1)))
            except queue.Full:
                pass
    finally:
        session.stop_event.set()
        print("Worker thread finished", flush=True)


def create_gpu_rollout_loop(command_queue, seed_path, prompt_text, seed):
    """Return a ZeroGPU generator that rolls the world out block-by-block.

    Only picklable primitives + the multiprocessing `command_queue` cross the
    fork boundary. Live controls (held key set) and stop arrive via that queue.
    """
    @spaces.GPU(duration=GPU_DURATION)
    def gpu_rollout():
        device = torch.device("cuda")
        prompt = (prompt_text or DEFAULT_PROMPT).strip() or DEFAULT_PROMPT
        set_seed(int(seed))

        with _infer_lock:
            noise = torch.randn(
                [1, NUM_FPB, LATENT_CHANNELS, LATENT_H, LATENT_W],
                device=device, dtype=torch.bfloat16,
            )

            pipeline.set_prompts([prompt], device=device)
            pipeline.set_ref_latent_mask_from_exists_paths(
                ref_dir=str(APP_DIR / "__no_such_ref_dir__"), device=device,
            )
            pipeline.reset_stream(batch_size=1, dtype=torch.bfloat16,
                                  device=device, initial_latent=None)
            pipeline.set_first_frame_latent(
                seed_path, height=STREAM_HEIGHT, width=STREAM_WIDTH, device=device,
            )

            current_keys: Dict[str, bool] = {"W": True}  # default: forward
            stop_requested = False
            try:
                for b in range(MAX_BLOCKS_PER_SESSION):
                    # Drain control commands: newest held-key set wins.
                    while True:
                        try:
                            cmd = command_queue.get_nowait()
                        except Exception:
                            break
                        if isinstance(cmd, StopCommand):
                            stop_requested = True
                            break
                        if isinstance(cmd, ControlCommand):
                            current_keys = _keys_from_buttons(cmd.buttons)
                    if stop_requested:
                        break

                    # Time the full generate+decode of one block so the parent
                    # thread can pace this block's frames over that duration.
                    block_start = time.time()
                    pipeline.set_act(current_keys, height=STREAM_HEIGHT, width=STREAM_WIDTH,
                                     num_frames=NUM_FPB, device=device)
                    lat_block = pipeline.generate_next_block(noise)
                    noise = torch.randn_like(noise)
                    frames = _decode_block_to_frames(lat_block)
                    block_elapsed = time.time() - block_start
                    n = len(frames)
                    for i, f in enumerate(frames):
                        # (frame, block_idx, frame_idx_in_block, frames_in_block,
                        #  block_elapsed) — the pacing metadata lets the parent
                        # spread this block's frames evenly rather than bursting.
                        yield (f, b, i, n, block_elapsed)
            finally:
                try:
                    pipeline.reset_stream(batch_size=1, dtype=torch.bfloat16,
                                          device=device, initial_latent=None)
                except Exception:
                    pass
                try:
                    if hasattr(pipeline.vae, "model") and hasattr(pipeline.vae.model, "clear_cache"):
                        pipeline.vae.model.clear_cache()
                except Exception:
                    pass

    return gpu_rollout()


# ── App (gradio.Server) ──────────────────────────────────────────────────────
app = Server()


@app.api(name="start_game")
def start_game(session_id: str = "", seed_path: str = "",
               prompt: str = "", seed: int = 42) -> str:
    """Start a new interactive world rollout for `session_id`.

    Args:
        session_id: per-client id (UUID) isolating this player's stream.
        seed_path: filepath (uploaded via /upload) of the starting frame image
            that seeds the i2v world rollout.
        prompt: scene description.
        seed: RNG seed for reproducibility.

    Returns:
        The session_id actually used.
    """
    if not session_id:
        session_id = str(uuid.uuid4())

    prior = _drop_session(session_id)
    if prior is not None:
        prior.stop()

    if not seed_path:
        raise ValueError("A starting image is required — please upload one first.")

    command_queue = MPQueue()  # crosses the ZeroGPU fork boundary
    frame_queue: "queue.Queue" = queue.Queue(maxsize=4)
    status_queue: "queue.Queue" = queue.Queue(maxsize=32)
    stop_event = threading.Event()

    session = GameSession(
        session_id=session_id,
        command_queue=command_queue,
        frame_queue=frame_queue,
        status_queue=status_queue,
        stop_event=stop_event,
        seed_path=seed_path,
        prompt=prompt or DEFAULT_PROMPT,
        seed=int(seed),
    )
    with _sessions_lock:
        _sessions[session_id] = session

    # Capture the *incoming request* — HF has already injected this user's
    # ZeroGPU proxy token (x-ip-token / x-api-token) into its headers. We
    # re-set it into the worker thread's gradio LocalContext so that
    # @spaces.GPU bills GPU time against THIS user's quota, not the owner's.
    gradio_request = LocalContext.request.get(None)
    status_token = _current_status_queue.set(status_queue)
    try:
        broadcast_status("Requesting GPU from ZeroGPU…")

        def _thread_entry():
            # Re-establish the request context inside the worker thread so the
            # ZeroGPU scheduler reads the requesting user's token.
            if gradio_request is not None:
                try:
                    LocalContext.request.set(gradio_request)
                except Exception:
                    pass
            gpu_worker_thread(session)

        ctx = contextvars.copy_context()
        worker = threading.Thread(target=ctx.run, args=(_thread_entry,), daemon=True)
        session.worker_thread = worker
        worker.start()
    finally:
        _current_status_queue.reset(status_token)

    return session_id


@app.api(name="stop_game")
def stop_game(session_id: str = "") -> str:
    """Stop the active rollout for the given client."""
    if not session_id:
        return "no_session"
    session = _drop_session(session_id)
    if session is not None:
        session.stop()
    return "stopped"


@app.websocket("/ws")
async def game_ws(websocket: WebSocket, session_id: str = ""):
    """Real-time rollout WebSocket. Requires `?session_id=...` matching /start_game."""
    await websocket.accept()
    if not session_id:
        await websocket.send_json({"type": "error", "message": "missing session_id"})
        await websocket.close(code=1008)
        return

    loop = asyncio.get_event_loop()

    async def send_frames():
        session_ended_sent = False
        while True:
            session = _get_session(session_id)

            if session is not None:
                try:
                    status_msg = session.status_queue.get_nowait()
                    if status_msg.startswith("error:"):
                        await websocket.send_json({"type": "error", "message": status_msg[6:]})
                        break
                    await websocket.send_json({"type": "status", "message": status_msg})
                except queue.Empty:
                    pass
                except (WebSocketDisconnect, RuntimeError):
                    break

            if session is None:
                await asyncio.sleep(0.05)
                continue
            if session.stop_event.is_set() and session.frame_queue.empty():
                if not session_ended_sent:
                    try:
                        await websocket.send_json({"type": "session_ended"})
                    except (WebSocketDisconnect, RuntimeError):
                        break
                    session_ended_sent = True
                await asyncio.sleep(0.4)
                continue
            try:
                result = await loop.run_in_executor(
                    None, lambda s=session: s.frame_queue.get(timeout=0.1)
                )
                frame, count, fps = result
                img = Image.fromarray(frame)
                buf = io.BytesIO()
                img.save(buf, format="JPEG", quality=JPEG_QUALITY)
                jpeg_bytes = buf.getvalue()
                header = struct.pack(">II", int(count), int(fps * 10))
                await websocket.send_bytes(header + jpeg_bytes)
                session.touch()
            except queue.Empty:
                pass
            except (WebSocketDisconnect, RuntimeError):
                break

    async def receive_controls():
        while True:
            try:
                data = await websocket.receive_json()
                session = _get_session(session_id)
                if session is None:
                    continue
                session.touch()
                msg_type = data.get("type", "control")
                if msg_type == "control":
                    buttons = set(data.get("buttons", []))
                    prompt = data.get("prompt", session.prompt)
                    try:
                        session.command_queue.put_nowait(
                            ControlCommand(buttons=buttons, prompt=prompt)
                        )
                    except queue.Full:
                        pass
                elif msg_type == "stop":
                    session.stop()
            except WebSocketDisconnect:
                break
            except Exception:
                break

    try:
        await asyncio.gather(send_frames(), receive_controls())
    except WebSocketDisconnect:
        pass


@app.post("/upload_seed")
async def upload_seed(file: UploadFile = File(...)):
    """Accept a user-uploaded starting image and stash it server-side.

    Returns the temp filepath, which the browser then passes to /start_game as
    `seed_path` to seed the image-to-video (i2v) world rollout. Only an image is
    needed — there is no video upload.
    """
    try:
        raw = await file.read()
        img = Image.open(io.BytesIO(raw)).convert("RGB")
    except Exception:
        return JSONResponse({"error": "Could not read image file."}, status_code=400)

    tmp = tempfile.NamedTemporaryFile(prefix="abot_seed_", suffix=".png", delete=False)
    img.save(tmp.name, format="PNG")
    return {"seed_path": tmp.name}


def _safe_example_path(name: str) -> Optional[Path]:
    """Resolve `name` to a bundled example image, guarding against traversal."""
    if not any(name == e["name"] for e in EXAMPLE_SEEDS):
        return None
    path = (EXAMPLES_DIR / name).resolve()
    if EXAMPLES_DIR.resolve() not in path.parents or not path.is_file():
        return None
    return path


@app.get("/example_seeds")
async def example_seeds():
    """List the preset starting-world images available as clickable thumbnails."""
    return {"examples": [e for e in EXAMPLE_SEEDS if (EXAMPLES_DIR / e["name"]).is_file()]}


@app.get("/example_thumb")
async def example_thumb(name: str = ""):
    """Serve a preset starting-world image (for thumbnail display in the UI)."""
    path = _safe_example_path(name)
    if path is None:
        return JSONResponse({"error": "unknown example"}, status_code=404)
    return FileResponse(str(path), media_type="image/png")


@app.get("/example_seed")
async def example_seed(name: str = ""):
    """Seed the i2v rollout from a bundled preset image (no upload required).

    Copies the chosen example into a server-side temp file and returns its path,
    mirroring /upload_seed so the browser can pass it to /start_game as seed_path.
    """
    path = _safe_example_path(name)
    if path is None:
        return JSONResponse({"error": "unknown example"}, status_code=404)
    try:
        img = Image.open(path).convert("RGB")
    except Exception:
        return JSONResponse({"error": "could not read example image"}, status_code=500)
    tmp = tempfile.NamedTemporaryFile(prefix="abot_seed_", suffix=".png", delete=False)
    img.save(tmp.name, format="PNG")
    return {"seed_path": tmp.name}


@app.get("/", response_class=HTMLResponse)
async def homepage():
    html_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "index.html")
    with open(html_path, "r", encoding="utf-8") as f:
        return f.read()


# Avoid ZeroGPU "no GPU function" error at boot.
spaces.GPU(lambda: None)

app.launch(server_name="0.0.0.0", server_port=7860)