app.py

"""
ABot-World on Modular Diffusers — interactive action-conditioned world rollout.
gradio.Server + WebSocket live-backend edition.

The engine is one `pipe.stream(actions=<callable>)` loop over the Modular Diffusers
`ABotWorldStreamingBlocks` preset (huggingface/diffusers#14159): the pipeline polls the
held keys once per generated block and yields every block's decoded frames.
The serving shell (gradio.Server, WebSocket frame stream, pacing, per-session queues,
front-end) is reused from https://huggingface.co/spaces/acvlab/abot-world-interactive.

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/YiYiXu/ABot-World-0-5B-LF-Diffusers  (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 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 fastapi import UploadFile, File, WebSocket, WebSocketDisconnect
from fastapi.responses import HTMLResponse, JSONResponse, FileResponse
from gradio import Server
from gradio.context import LocalContext

from diffusers.modular_pipelines import ABotWorldStreamingBlocks

# ── Repo paths ───────────────────────────────────────────────────────────────
APP_DIR = Path(__file__).resolve().parent

MODEL_ID = "YiYiXu/ABot-World-0-5B-LF-Diffusers"

# 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 = 150               # seconds per @spaces.GPU allocation (one session)

# ── Real-time pacing configuration ───────────────────────────────────────────
# The GPU decodes a whole block (12 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."
)

# ── Build the Modular Diffusers streaming pipeline (module scope) ────────────
# On ZeroGPU `pipe.to("cuda")` at import only packs the weights; they land on the
# GPU inside the @spaces.GPU rollout.
print(f"[startup] loading {MODEL_ID} ...", flush=True)
torch.set_grad_enabled(False)
pipe = ABotWorldStreamingBlocks().init_pipeline(MODEL_ID)
pipe.load_components(dtype=torch.bfloat16)
pipe.to("cuda")
print("[startup] pipeline ready.", flush=True)

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


def _action_from_buttons(buttons):
    """Translate a set of held key names (e.g. {'W','A'}) into the model's 8-key multi-hot action."""
    held = {k.upper() for k in (buttons or [])}
    return [int(k in held) for k in KEY_ORDER]


# ── 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():
        prompt = (prompt_text or DEFAULT_PROMPT).strip() or DEFAULT_PROMPT
        image = Image.open(seed_path).convert("RGB")
        state = {"action": _action_from_buttons({"W"}), "block_start": time.time()}  # default: forward

        def action_source(block_index):
            """Polled by the pipeline once per block: newest held-key set wins, None stops the rollout."""
            if block_index >= MAX_BLOCKS_PER_SESSION:
                return None
            while True:
                try:
                    cmd = command_queue.get_nowait()
                except Exception:
                    break
                if isinstance(cmd, StopCommand):
                    return None
                if isinstance(cmd, ControlCommand):
                    state["action"] = _action_from_buttons(cmd.buttons)
            state["block_start"] = time.time()
            return state["action"]

        with _infer_lock:
            events = pipe.stream(
                prompt=prompt,
                image=image,
                height=STREAM_HEIGHT,
                width=STREAM_WIDTH,
                actions=action_source,
                generator=torch.Generator("cpu").manual_seed(int(seed)),
            )
            for event in events:
                if event.path != "denoise.rollout":
                    continue  # inner per-denoise-step events
                # Time the full generate+decode of one block so the parent thread
                # can pace this block's frames over that duration.
                block_elapsed = time.time() - state["block_start"]
                b = event.loop_kwargs["k"]
                frames = (event.state.get("frames") * 255).clip(0, 255).astype(np.uint8)
                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)

    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, ssr_mode=False)  # the SSR Node proxy does not forward the /ws upgrade