Live front-end: gradio.Server + WebSocket paced frame stream (shell from acvlab/abot-world-interactive), engine = pipe.stream(action_source=...)

README.md

@@ -4,7 +4,7 @@ emoji: 🕹️
 colorFrom: yellow
 colorTo: green
 sdk: gradio
-sdk_version: 6.5.1
+sdk_version: 6.25.0
 python_version: '3.12'
 app_file: app.py
 pinned: false
@@ -12,8 +12,12 @@ short_description: Drive a live world model with Modular Diffusers
 ---
 
 Interactive [ABot-World](https://huggingface.co/YiYiXu/ABot-World-0-5B-LF-Diffusers) (a real-time,
-action-conditioned world model, Wan2.2-TI2V-5B) running natively on Modular Diffusers: the whole app is a
-`pipe.stream(action_source=...)` consumer — the pipeline polls the held control once per generated block.
+action-conditioned world model, Wan2.2-TI2V-5B) running natively on Modular Diffusers. The engine is one
+`pipe.stream(action_source=...)` loop over the `ABotWorldStreamingBlocks` preset
+([huggingface/diffusers#14159](https://github.com/huggingface/diffusers/pull/14159)): the pipeline polls the
+held keys once per generated block and yields every block's decoded frames.
 
-Runs on ZeroGPU: one `@spaces.GPU` generator owns the whole session (the rollout keeps a K/V cache alive
-across blocks) and the held control reaches it through a queue. ~20 GB VRAM in bf16.
+The serving shell — `gradio.Server` start/stop endpoints, a WebSocket that streams binary JPEG frames with
+steady pacing, per-session control queues that cross the ZeroGPU fork, and the front-end — is reused from
+[acvlab/abot-world-interactive](https://huggingface.co/spaces/acvlab/abot-world-interactive); only the
+~40-line GPU rollout body is ours.

app.py

@@ -1,131 +1,616 @@
-import spaces  # must come before torch: on ZeroGPU it patches CUDA so the model can be staged at import
+"""
+ABot-World on Modular Diffusers — interactive action-conditioned world rollout.
+gradio.Server + WebSocket live-backend edition.
 
-import multiprocessing
+The engine is one `pipe.stream(action_source=...)` 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 gradio as gr
 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
-from diffusers.utils import load_image
 
+# ── Repo paths ───────────────────────────────────────────────────────────────
+APP_DIR = Path(__file__).resolve().parent
 
-REPO = "YiYiXu/ABot-World-0-5B-LF-Diffusers"
+MODEL_ID = "YiYiXu/ABot-World-0-5B-LF-Diffusers"
 
-pipe = ABotWorldStreamingBlocks().init_pipeline(REPO)
-pipe.load_components(dtype=torch.bfloat16)
-pipe.to("cuda")  # on ZeroGPU this only packs the weights; they land on the GPU inside the @spaces.GPU call
+# 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 (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."
 )
 
-CONTROLS = {
-    "▲ forward (W)": [1, 0, 0, 0, 0, 0, 0, 0],
-    "◀ left (A)": [0, 1, 0, 0, 0, 0, 0, 0],
-    "▼ back (S)": [0, 0, 1, 0, 0, 0, 0, 0],
-    "▶ right (D)": [0, 0, 0, 1, 0, 0, 0, 0],
-    "cam ↑ (I)": [0, 0, 0, 0, 1, 0, 0, 0],
-    "cam ◀ (J)": [0, 0, 0, 0, 0, 1, 0, 0],
-    "cam ↓ (K)": [0, 0, 0, 0, 0, 0, 1, 0],
-    "cam ▶ (L)": [0, 0, 0, 0, 0, 0, 0, 1],
-    "· idle": [0] * 8,
-}
+# ── 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()
 
-# One live world at a time. On ZeroGPU the rollout runs in a forked worker process, so the held control
-# crosses over a multiprocessing queue created before the fork: buttons put an action, Stop puts None,
-# and the rollout drains the queue once per block.
-commands = multiprocessing.Queue()
-run_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]
 
-def hold(control):
-    commands.put(CONTROLS[control])
-    return f"holding: {control}"
 
+# ── Command types (browser -> worker) ────────────────────────────────────────
+@dataclass
+class ControlCommand:
+    buttons: Set[str]
+    prompt: str
 
-def stop():
-    commands.put(None)
-    return "stopping after this block…"
 
+@dataclass
+class StopCommand:
+    pass
 
-@spaces.GPU(duration=lambda image, prompt, max_blocks: 30 + 6 * int(max_blocks))
-def rollout(image, prompt, max_blocks):
-    held = CONTROLS["▲ forward (W)"]
 
-    def action_source(block_index):
-        nonlocal held
-        while True:
+# ── 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:
-                command = commands.get_nowait()
-            except queue.Empty:
+                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 command is None:
+
+            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
-            held = command
-        return held if block_index < max_blocks else None
-
-    for event in pipe.stream(
-        prompt=(prompt or DEFAULT_PROMPT).strip() or DEFAULT_PROMPT,
-        image=image,
-        action_source=action_source,
-        generator=torch.Generator("cpu").manual_seed(42),
-    ):
-        if event.path != "denoise.rollout":
-            continue
-        frames = event.state.get("frames")
-        yield (frames[-1] * 255).clip(0, 255).astype(np.uint8)
-
-
-def play(image, prompt, max_blocks):
-    if not run_lock.acquire(blocking=False):
-        raise gr.Error("A session is already running — stop it first.")
+            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,
+                action_source=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:
-        while True:  # drop controls left over from the previous session
+        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:
-                commands.get_nowait()
+                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
-        if image is None:
-            image = load_image("examples/desert_valley.png")
-        yield from rollout(image, prompt, max_blocks)
-    finally:
-        run_lock.release()
 
+    try:
+        await asyncio.gather(send_frames(), receive_controls())
+    except WebSocketDisconnect:
+        pass
 
-with gr.Blocks(title="ABot-World Interactive") as demo:
-    gr.Markdown(
-        """# ABot-World — interactive world, Modular Diffusers
-Upload a scene, press Start, then hold a control: each ~1 s block is generated live by
-[ABot-World](https://huggingface.co/YiYiXu/ABot-World-0-5B-LF-Diffusers) conditioned on the held action.
-The whole app is one `pipe.stream(action_source=...)` loop — the pipeline polls your current input once
-per block ([huggingface/diffusers#14159](https://github.com/huggingface/diffusers/pull/14159))."""
-    )
-    with gr.Row():
-        with gr.Column(scale=1):
-            image = gr.Image(label="starting scene", type="pil", height=220)
-            prompt = gr.Textbox(label="scene prompt", value=DEFAULT_PROMPT, lines=2)
-            max_blocks = gr.Slider(label="max blocks (~seconds)", minimum=4, maximum=40, step=1, value=20)
-            with gr.Row():
-                start_button = gr.Button("Start", variant="primary")
-                stop_button = gr.Button("Stop")
-            gr.Examples(
-                examples=[["examples/desert_valley.png"], ["examples/forest_stream.png"],
-                          ["examples/mountain_meadow.png"], ["examples/example.png"]],
-                inputs=[image],
-            )
-        with gr.Column(scale=2):
-            view = gr.Image(label="world", show_label=False, interactive=False)
-            status = gr.Textbox(value="holding: ▲ forward (W)", show_label=False, container=False)
-            with gr.Row():
-                buttons = [gr.Button(name, size="sm") for name in CONTROLS]
-
-    for button in buttons:
-        button.click(hold, inputs=gr.State(button.value), outputs=status, api_name="hold")
-    start_button.click(play, inputs=[image, prompt, max_blocks], outputs=view, api_name="play")
-    stop_button.click(stop, outputs=status, api_name="stop")
-
-demo.launch(show_error=True)
+
+@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)

index.html

@@ -0,0 +1,634 @@
+<!DOCTYPE html>
+<html lang="en">
+<head>
+<meta charset="UTF-8" />
+<meta name="viewport" content="width=device-width, initial-scale=1.0, maximum-scale=1.0, user-scalable=no" />
+<title>ABot-World · Interactive World Rollout · Modular Diffusers</title>
+<style>
+  :root{
+    --bg:#0b0e14; --panel:#121722; --panel-2:#171d2b; --edge:#232c3d;
+    --text:#e6ecf5; --muted:#8a97ad; --accent:#5b8cff; --accent-2:#7c5cff;
+    --good:#3ecf8e; --warn:#ffb454; --bad:#ff5c6c; --radius:16px;
+  }
+  *{box-sizing:border-box}
+  html,body{margin:0;height:100%}
+  body{
+    font-family:'Inter',system-ui,-apple-system,Segoe UI,Roboto,Helvetica,Arial,sans-serif;
+    background:radial-gradient(1200px 700px at 80% -10%, #1a2138 0%, var(--bg) 60%) fixed;
+    color:var(--text); -webkit-font-smoothing:antialiased; min-height:100%;
+  }
+  .wrap{max-width:1240px; margin:0 auto; padding:28px 20px 48px}
+  header{display:flex; align-items:center; justify-content:space-between; gap:16px; flex-wrap:wrap; margin-bottom:22px}
+  .brand{display:flex; align-items:center; gap:12px}
+  .logo{width:42px;height:42px;border-radius:12px;
+    background:linear-gradient(135deg,var(--accent),var(--accent-2));
+    display:grid;place-items:center;font-size:22px;box-shadow:0 6px 24px rgba(91,140,255,.35)}
+  .brand h1{font-size:19px;margin:0;letter-spacing:.2px}
+  .brand p{margin:2px 0 0;font-size:12.5px;color:var(--muted)}
+  .links{display:flex;gap:8px;flex-wrap:wrap}
+  .links a{font-size:12.5px;color:var(--muted);text-decoration:none;border:1px solid var(--edge);
+    padding:6px 11px;border-radius:999px;transition:.15s}
+  .links a:hover{color:var(--text);border-color:var(--accent);background:rgba(91,140,255,.08)}
+
+  .layout{display:grid;grid-template-columns:minmax(0,1fr) 360px;gap:20px}
+  @media (max-width:960px){.layout{grid-template-columns:1fr}}
+
+  .card{background:linear-gradient(180deg,var(--panel),var(--panel-2));
+    border:1px solid var(--edge);border-radius:var(--radius);padding:16px}
+
+  /* Stage */
+  .stage{position:relative;aspect-ratio:1280/704;background:#05070c;border-radius:12px;
+    overflow:hidden;border:1px solid var(--edge);display:grid;place-items:center}
+  .stage canvas{width:100%;height:100%;object-fit:cover;display:none;image-rendering:auto}
+  .stage.playing canvas{display:block}
+  .overlay{position:absolute;inset:0;display:grid;place-items:center;text-align:center;
+    padding:24px;background:rgba(5,7,12,.55);backdrop-filter:blur(2px)}
+  .overlay.hidden{display:none}
+  .spinner{width:34px;height:34px;border-radius:50%;border:3px solid rgba(255,255,255,.15);
+    border-top-color:var(--accent);animation:spin 1s linear infinite;margin:0 auto 14px}
+  @keyframes spin{to{transform:rotate(360deg)}}
+  .overlay h3{margin:0 0 6px;font-size:16px}
+  .overlay p{margin:0;color:var(--muted);font-size:13px;max-width:420px}
+
+  .hud{position:absolute;top:10px;left:10px;right:10px;display:flex;justify-content:space-between;
+    pointer-events:none;font-size:12px}
+  .chip{background:rgba(9,12,20,.7);border:1px solid var(--edge);border-radius:999px;
+    padding:5px 11px;display:flex;align-items:center;gap:7px;backdrop-filter:blur(4px)}
+  .dot{width:8px;height:8px;border-radius:50%;background:var(--muted)}
+  .dot.playing{background:var(--good);box-shadow:0 0 8px var(--good)}
+  .dot.loading{background:var(--warn);box-shadow:0 0 8px var(--warn);animation:pulse 1s infinite}
+  .dot.error,.dot.ended{background:var(--bad)}
+  @keyframes pulse{50%{opacity:.35}}
+
+  .capture-hint{position:absolute;bottom:10px;left:50%;transform:translateX(-50%);
+    background:rgba(9,12,20,.72);border:1px solid var(--edge);border-radius:999px;
+    padding:6px 14px;font-size:12px;color:var(--muted);pointer-events:none;transition:.2s;opacity:0}
+  .stage.playing .capture-hint{opacity:1}
+
+  /* Uploader */
+  .drop{border:1.5px dashed var(--edge);border-radius:12px;padding:16px;text-align:center;
+    cursor:pointer;transition:.15s;background:rgba(255,255,255,.015)}
+  .drop:hover,.drop.drag{border-color:var(--accent);background:rgba(91,140,255,.07)}
+  .drop .icon{font-size:26px}
+  .drop .t{font-size:13.5px;margin-top:6px}
+  .drop .s{font-size:11.5px;color:var(--muted);margin-top:2px}
+  .thumb{display:none;margin-top:10px;border-radius:10px;overflow:hidden;border:1px solid var(--edge)}
+  .thumb img{width:100%;display:block;max-height:170px;object-fit:cover}
+  .thumb.show{display:block}
+
+  /* Preset starting-world thumbnails */
+  .presets{display:grid;grid-template-columns:repeat(4,1fr);gap:8px;margin-top:10px}
+  .preset{position:relative;border:1px solid var(--edge);border-radius:9px;overflow:hidden;
+    cursor:pointer;background:#0d1220;aspect-ratio:16/9;transition:.15s}
+  .preset:hover{border-color:var(--accent);box-shadow:0 4px 14px rgba(91,140,255,.25)}
+  .preset.selected{border-color:var(--accent);box-shadow:0 0 0 2px rgba(91,140,255,.45)}
+  .preset img{width:100%;height:100%;object-fit:cover;display:block}
+  .preset .pl{position:absolute;left:0;right:0;bottom:0;font-size:10px;color:#fff;
+    padding:3px 5px;background:linear-gradient(0deg,rgba(5,7,12,.85),transparent);
+    white-space:nowrap;overflow:hidden;text-overflow:ellipsis}
+  .or-sep{font-size:11px;color:var(--muted);text-align:center;margin:12px 0 8px;
+    text-transform:uppercase;letter-spacing:.7px}
+
+  .field{margin-top:14px}
+  .label{font-size:11.5px;text-transform:uppercase;letter-spacing:.7px;color:var(--muted);margin-bottom:6px}
+  textarea,input[type=number]{width:100%;background:#0d1220;color:var(--text);border:1px solid var(--edge);
+    border-radius:10px;padding:10px 12px;font-size:13.5px;font-family:inherit;resize:vertical}
+  textarea:focus,input:focus{outline:none;border-color:var(--accent)}
+  .row{display:flex;gap:10px}
+  .row>*{flex:1}
+
+  .btns{display:flex;gap:10px;margin-top:16px}
+  button{border:none;border-radius:10px;padding:11px 14px;font-size:13.5px;font-weight:600;
+    cursor:pointer;font-family:inherit;transition:.15s;flex:1}
+  .primary{background:linear-gradient(135deg,var(--accent),var(--accent-2));color:#fff;
+    box-shadow:0 6px 20px rgba(91,140,255,.32)}
+  .primary:hover{filter:brightness(1.07)}
+  .ghost{background:transparent;color:var(--muted);border:1px solid var(--edge)}
+  .ghost:hover{color:var(--text);border-color:var(--accent)}
+  button:disabled{opacity:.45;cursor:not-allowed;filter:none}
+
+  /* Controls section (below the stage) */
+  .controls-section{margin-top:20px}
+  .controls-grid{display:grid;grid-template-columns:auto 1fr;gap:28px;align-items:center;margin-top:12px}
+  @media (max-width:640px){.controls-grid{grid-template-columns:1fr;gap:20px}}
+
+  /* Visual on-screen WASD keyboard */
+  .keyboard{display:grid;grid-template-columns:repeat(3,56px);grid-template-rows:repeat(2,56px);
+    gap:8px;justify-content:start}
+  .keycap{display:grid;place-items:center;border-radius:12px;background:#0d1220;
+    border:1px solid var(--edge);font-size:18px;color:var(--muted);font-weight:700;
+    transition:.08s;user-select:none}
+  .keycap.spacer{background:transparent;border:none}
+  .keycap.active{background:linear-gradient(135deg,var(--accent),var(--accent-2));color:#fff;
+    border-color:transparent;box-shadow:0 4px 16px rgba(91,140,255,.45);transform:translateY(1px)}
+  /* WASD cross placement inside the 3x2 grid */
+  .keycap[data-k="W"]{grid-column:2;grid-row:1}
+  .keycap[data-k="A"]{grid-column:1;grid-row:2}
+  .keycap[data-k="S"]{grid-column:2;grid-row:2}
+  .keycap[data-k="D"]{grid-column:3;grid-row:2}
+  /* IJKL look cross placement inside its own 3x2 grid */
+  .keycap[data-k="I"]{grid-column:2;grid-row:1}
+  .keycap[data-k="J"]{grid-column:1;grid-row:2}
+  .keycap[data-k="K"]{grid-column:2;grid-row:2}
+  .keycap[data-k="L"]{grid-column:3;grid-row:2}
+
+  .look-info{font-size:12.5px;color:var(--muted);line-height:1.6}
+  .look-info b{color:var(--text)}
+  .look-badge{display:inline-flex;align-items:center;gap:6px;font-size:11.5px;color:var(--good);
+    border:1px solid var(--edge);border-radius:999px;padding:4px 10px;margin-top:8px}
+  .look-badge .d{width:7px;height:7px;border-radius:50%;background:var(--muted)}
+  .look-badge.active .d{background:var(--good);box-shadow:0 0 8px var(--good)}
+  .hint-line{font-size:11.5px;color:var(--muted);margin-top:14px;line-height:1.55}
+  .active-tags{font-size:12px;color:var(--good);min-height:16px;margin-top:8px}
+  .footnote{font-size:11px;color:var(--muted);margin-top:12px;line-height:1.5}
+</style>
+</head>
+<body>
+<div class="wrap">
+  <header>
+    <div class="brand">
+      <div class="logo">🌍</div>
+      <div>
+        <h1>ABot-World — Interactive World Rollout on Modular Diffusers</h1>
+        <p>Upload a starting frame, then steer a live action-conditioned world with WASD + IJKL.</p>
+      </div>
+    </div>
+    <div class="links">
+      <a href="https://huggingface.co/YiYiXu/ABot-World-0-5B-LF-Diffusers" target="_blank">Model</a>
+      <a href="https://github.com/huggingface/diffusers/pull/14159" target="_blank">Modular streaming API</a>
+      <a href="https://amap-cvlab.github.io/ABot-World/" target="_blank">Project</a>
+      <a href="https://github.com/amap-cvlab/ABot-World" target="_blank">Code</a>
+    </div>
+  </header>
+
+  <div class="layout">
+    <!-- Stage -->
+    <div class="card">
+      <div class="stage" id="stage" tabindex="0">
+        <canvas id="view" width="1280" height="704"></canvas>
+        <div class="hud">
+          <div class="chip"><span class="dot" id="status-dot"></span><span id="status-label">Idle</span></div>
+          <div class="chip"><span id="fps-value">0</span> fps · <span id="frame-value">0</span> blk</div>
+        </div>
+        <div class="capture-hint" id="capture-hint">Click to capture · WASD to move · IJKL to look · Esc to release</div>
+        <div class="overlay" id="overlay">
+          <div>
+            <div class="spinner" id="spinner" style="display:none"></div>
+            <h3 id="overlay-title">Upload an image to begin</h3>
+            <p id="overlay-sub">Pick a starting frame on the right, add a scene prompt, and press <b>Start world</b>.</p>
+          </div>
+        </div>
+      </div>
+
+      <!-- Controls (directly below the video output) -->
+      <div class="controls-section">
+        <div class="controls-grid">
+          <div>
+            <div class="label" style="margin-bottom:8px">Movement</div>
+            <div class="keyboard" id="keys">
+              <span class="keycap" data-k="W">W</span>
+              <span class="keycap" data-k="A">A</span>
+              <span class="keycap" data-k="S">S</span>
+              <span class="keycap" data-k="D">D</span>
+            </div>
+          </div>
+          <div>
+            <div class="label" style="margin-bottom:8px">Look / pan</div>
+            <div class="keyboard" id="look-keys">
+              <span class="keycap" data-k="I">I</span>
+              <span class="keycap" data-k="J">J</span>
+              <span class="keycap" data-k="K">K</span>
+              <span class="keycap" data-k="L">L</span>
+            </div>
+            <div class="look-badge" id="look-badge"><span class="d"></span><span id="look-badge-text">Look idle</span></div>
+          </div>
+        </div>
+        <div class="hint-line">
+          <b>W A S D</b> — move &amp; turn (hold to steer the next block) ·
+          <b>I J K L</b> — look / pan. Click the video to capture input.
+        </div>
+        <div class="active-tags" id="active-tags"></div>
+      </div>
+    </div>
+
+    <!-- Controls -->
+    <div class="card">
+      <div class="label">Starting image (i2v seed)</div>
+      <div class="presets" id="presets"></div>
+      <div class="or-sep">or upload your own</div>
+      <div class="drop" id="drop">
+        <div class="icon">🖼️</div>
+        <div class="t">Drop an image or click to upload</div>
+        <div class="s">PNG / JPG · becomes the world's first frame</div>
+      </div>
+      <input type="file" id="file-input" accept="image/*" style="display:none" />
+      <div class="thumb" id="thumb"><img id="thumb-img" alt="seed preview" /></div>
+
+      <div class="field">
+        <div class="label">Scene prompt</div>
+        <textarea id="prompt" rows="3">A realistic outdoor world scene with a navigable path, natural lighting, detailed ground texture, and stable forward motion.</textarea>
+      </div>
+
+      <div class="field row">
+        <div>
+          <div class="label">Seed</div>
+          <input type="number" id="seed" value="42" />
+        </div>
+      </div>
+
+      <div class="btns">
+        <button class="primary" id="start-btn" disabled>Start world</button>
+        <button class="ghost" id="stop-btn" disabled>Stop</button>
+      </div>
+    </div>
+  </div>
+</div>
+
+<script type="module">
+import { Client } from "https://cdn.jsdelivr.net/npm/@gradio/client/dist/index.min.js";
+
+// One session id per tab keeps every server-side queue isolated per user.
+const SESSION_ID = (crypto.randomUUID && crypto.randomUUID()) ||
+  (`s-${Date.now().toString(36)}-${Math.random().toString(36).slice(2,10)}`);
+
+// Movement keys are driven by the keyboard (WASD); look/pan is driven by the
+// keyboard too (I J K L) — both are handled as held-key controls.
+const MOVE_KEYS = new Set(["W","A","S","D"]);
+const LOOK_KEYS = new Set(["I","J","K","L"]);
+const KEY_FROM_CODE = {
+  KeyW:"W", KeyA:"A", KeyS:"S", KeyD:"D",
+  KeyI:"I", KeyJ:"J", KeyK:"K", KeyL:"L",
+};
+
+const state = {
+  sessionId: SESSION_ID,
+  client: null,
+  ws: null,
+  playing: false,
+  capturing: false,
+  pressed: new Set(),
+  seedPath: "",
+  prompt: "",
+};
+
+// --- DOM ---
+const $ = (id) => document.getElementById(id);
+const stage = $("stage"), view = $("view"), overlay = $("overlay");
+const viewCtx = view.getContext("2d", { alpha: false, desynchronized: true });
+const spinner = $("spinner"), overlayTitle = $("overlay-title"), overlaySub = $("overlay-sub");
+const statusDot = $("status-dot"), statusLabel = $("status-label");
+const fpsValue = $("fps-value"), frameValue = $("frame-value");
+const drop = $("drop"), fileInput = $("file-input"), thumb = $("thumb"), thumbImg = $("thumb-img");
+const presetsEl = $("presets");
+const promptEl = $("prompt"), seedEl = $("seed");
+const startBtn = $("start-btn"), stopBtn = $("stop-btn");
+const activeTags = $("active-tags");
+const keycaps = Array.from(document.querySelectorAll(".keycap[data-k]"));
+const lookBadge = $("look-badge"), lookBadgeText = $("look-badge-text");
+
+// --- UI helpers ---
+function setStatus(kind, label){
+  statusDot.className = "dot " + kind;
+  statusLabel.textContent = label;
+}
+function showOverlay(title, sub, loading){
+  overlay.classList.remove("hidden");
+  overlayTitle.textContent = title;
+  overlaySub.textContent = sub || "";
+  spinner.style.display = loading ? "block" : "none";
+}
+function hideOverlay(){ overlay.classList.add("hidden"); }
+function setPlaying(on){
+  state.playing = on;
+  stage.classList.toggle("playing", on);
+  stopBtn.disabled = !on;
+}
+function refreshStartEnabled(){
+  startBtn.disabled = !state.seedPath || state.playing;
+}
+
+// --- Upload ---
+async function uploadFile(file){
+  if(!file || !file.type.startsWith("image/")){ return; }
+  clearPresetSelection();
+  // Local preview
+  const reader = new FileReader();
+  reader.onload = (e)=>{ thumbImg.src = e.target.result; thumb.classList.add("show"); };
+  reader.readAsDataURL(file);
+
+  drop.querySelector(".t").textContent = "Uploading…";
+  const fd = new FormData();
+  fd.append("file", file, file.name || "seed.png");
+  try{
+    const res = await fetch(`${window.location.origin}/upload_seed`, { method:"POST", body:fd });
+    const data = await res.json();
+    if(data.seed_path){
+      state.seedPath = data.seed_path;
+      drop.querySelector(".t").textContent = "Image ready ✓ — click to replace";
+      showOverlay("Ready to explore", "Press Start world to begin the rollout.", false);
+    }else{
+      drop.querySelector(".t").textContent = "Upload failed — try another image";
+    }
+  }catch(err){
+    console.error(err);
+    drop.querySelector(".t").textContent = "Upload failed — try again";
+  }
+  refreshStartEnabled();
+}
+
+drop.addEventListener("click", ()=>fileInput.click());
+fileInput.addEventListener("change", (e)=>uploadFile(e.target.files[0]));
+["dragenter","dragover"].forEach(ev=>drop.addEventListener(ev,(e)=>{e.preventDefault();drop.classList.add("drag");}));
+["dragleave","drop"].forEach(ev=>drop.addEventListener(ev,(e)=>{e.preventDefault();drop.classList.remove("drag");}));
+drop.addEventListener("drop",(e)=>{ if(e.dataTransfer.files.length) uploadFile(e.dataTransfer.files[0]); });
+
+// --- Preset starting-world thumbnails (seed rollout directly, no upload) ---
+function clearPresetSelection(){
+  presetsEl.querySelectorAll(".preset").forEach(p=>p.classList.remove("selected"));
+}
+async function pickPreset(name, el){
+  clearPresetSelection();
+  el.classList.add("selected");
+  // Show the chosen preset in the seed preview too.
+  thumbImg.src = `${window.location.origin}/example_thumb?name=${encodeURIComponent(name)}`;
+  thumb.classList.add("show");
+  drop.querySelector(".t").textContent = "Loading preset…";
+  try{
+    const res = await fetch(`${window.location.origin}/example_seed?name=${encodeURIComponent(name)}`);
+    const data = await res.json();
+    if(data.seed_path){
+      state.seedPath = data.seed_path;
+      drop.querySelector(".t").textContent = "Preset ready ✓ — or upload to replace";
+      showOverlay("Ready to explore", "Press Start world to begin the rollout.", false);
+    }else{
+      drop.querySelector(".t").textContent = "Could not load preset — try another";
+      el.classList.remove("selected");
+    }
+  }catch(err){
+    console.error(err);
+    drop.querySelector(".t").textContent = "Could not load preset — try again";
+    el.classList.remove("selected");
+  }
+  refreshStartEnabled();
+}
+async function loadPresets(){
+  try{
+    const res = await fetch(`${window.location.origin}/example_seeds`);
+    const data = await res.json();
+    (data.examples || []).forEach(ex=>{
+      const div = document.createElement("div");
+      div.className = "preset";
+      div.title = ex.label || ex.name;
+      div.innerHTML =
+        `<img loading="lazy" alt="${ex.label || ex.name}" ` +
+        `src="${window.location.origin}/example_thumb?name=${encodeURIComponent(ex.name)}" />` +
+        `<span class="pl">${ex.label || ex.name}</span>`;
+      div.addEventListener("click", ()=>pickPreset(ex.name, div));
+      presetsEl.appendChild(div);
+    });
+  }catch(err){ console.error("presets load failed", err); }
+}
+
+// --- Gradio client (auto-forwards the HF iframe x-api-token for quota) ---
+async function initClient(){
+  if(!state.client) state.client = await Client.connect(window.location.origin);
+}
+
+// --- Frame decode + jitter buffer + steady render clock ---
+// Rendering is DECOUPLED from network arrival. WebSocket messages only decode
+// frames into a small jitter buffer; a steady render clock (requestAnimationFrame)
+// then pulls the freshest decoded frame at a smoothed target interval and paints
+// it. This absorbs network jitter and the server's block-boundary bursts so the
+// visible cadence stays even instead of "burst of N frames, then a gap".
+//
+// The buffer is intentionally shallow (JITTER_MAX): if frames arrive faster than
+// we render, older ones are dropped (coalesced) rather than queued, so we never
+// accumulate buffering delay that would later flush as a burst of stale frames.
+const JITTER_MAX = 4;              // max decoded frames held before we drop-oldest
+const JITTER_MIN = 2;              // frames to buffer before the render clock starts
+                                   // (a 2-frame lead absorbs network gaps so the
+                                   //  render clock rarely starves between bursts)
+let jitterBuf = [];                // [{bitmap, blk, fps}] decoded, awaiting display
+let renderPrimed = false;          // becomes true once JITTER_MIN frames buffered
+let renderInterval = 1000 / 16;    // ms between paints; adapted to server fps (EMA)
+let lastPaint = 0;                 // performance.now() of last painted frame
+let lastArrival = 0;               // performance.now() of last WS frame (for interval est.)
+let arrivalEma = 0;                // EMA of inter-arrival gap (ms)
+
+async function ingestFrame(bytes, blk, fps){
+  // Estimate the true incoming cadence from arrival timing and adapt the render
+  // interval toward it (bounded), so the steady clock matches the real fps.
+  const now = performance.now();
+  if(lastArrival){
+    const gap = now - lastArrival;
+    arrivalEma = arrivalEma ? (0.2 * gap + 0.8 * arrivalEma) : gap;
+    const tgt = Math.min(200, Math.max(20, arrivalEma));   // clamp 5–50 fps
+    renderInterval = 0.2 * tgt + 0.8 * renderInterval;
+  }
+  lastArrival = now;
+
+  let bitmap = null;
+  try{
+    bitmap = await createImageBitmap(new Blob([bytes], {type:"image/jpeg"}));
+  }catch(err){
+    // Fallback path for browsers without createImageBitmap JPEG support:
+    // paint straight to canvas (bypasses buffer but keeps the stream alive).
+    await paintViaImage(bytes, blk, fps);
+    return;
+  }
+  jitterBuf.push({ bitmap, blk, fps });
+  // Drop-oldest to keep the buffer shallow (avoid stale-frame burst on flush).
+  while(jitterBuf.length > JITTER_MAX){
+    const stale = jitterBuf.shift();
+    if(stale.bitmap && stale.bitmap.close) stale.bitmap.close();
+  }
+  if(!renderPrimed && jitterBuf.length >= JITTER_MIN){
+    renderPrimed = true;
+    lastPaint = performance.now() - renderInterval;   // paint first frame promptly
+  }
+}
+
+function paintBitmap(entry){
+  viewCtx.drawImage(entry.bitmap, 0, 0, view.width, view.height);
+  if(entry.bitmap && entry.bitmap.close) entry.bitmap.close();
+  if(!state.playing){ hideOverlay(); setPlaying(true); setStatus("playing","Playing"); }
+  frameValue.textContent = entry.blk;
+  fpsValue.textContent = entry.fps.toFixed(1);
+}
+
+// Steady render clock: runs continuously via rAF while playing; pulls at most one
+// buffered frame per render-interval so display timing is even, not arrival-tied.
+function renderClock(){
+  requestAnimationFrame(renderClock);
+  if(!renderPrimed || jitterBuf.length === 0) return;
+  const now = performance.now();
+  if(now - lastPaint < renderInterval) return;   // not yet this frame's slot
+  const entry = jitterBuf.shift();               // freshest frames kept; oldest painted in order
+  lastPaint = now;
+  paintBitmap(entry);
+}
+requestAnimationFrame(renderClock);
+
+function resetJitterBuffer(){
+  for(const e of jitterBuf){ if(e.bitmap && e.bitmap.close) e.bitmap.close(); }
+  jitterBuf = [];
+  renderPrimed = false;
+  lastArrival = 0;
+  arrivalEma = 0;
+}
+
+function paintViaImage(bytes, blk, fps){
+  return new Promise((resolve)=>{
+    const url = URL.createObjectURL(new Blob([bytes], {type:"image/jpeg"}));
+    const im = new Image();
+    im.onload = ()=>{
+      viewCtx.drawImage(im, 0, 0, view.width, view.height);
+      URL.revokeObjectURL(url);
+      if(!state.playing){ hideOverlay(); setPlaying(true); setStatus("playing","Playing"); }
+      if(typeof blk !== "undefined") frameValue.textContent = blk;
+      if(typeof fps !== "undefined") fpsValue.textContent = fps.toFixed(1);
+      resolve(true);
+    };
+    im.onerror = ()=>{ URL.revokeObjectURL(url); resolve(false); };
+    im.src = url;
+  });
+}
+
+// --- WebSocket frame stream ---
+function connectWS(){
+  const proto = location.protocol === "https:" ? "wss:" : "ws:";
+  const ws = new WebSocket(`${proto}//${location.host}/ws?session_id=${encodeURIComponent(state.sessionId)}`);
+  ws.binaryType = "arraybuffer";
+  ws.onmessage = (e)=>{
+    if(e.data instanceof ArrayBuffer){
+      const dv = new DataView(e.data);
+      const blk = dv.getUint32(0);
+      const fps = dv.getUint32(4)/10;
+      // Decode into the jitter buffer; the steady render clock paints on its own
+      // cadence, so display timing is decoupled from network arrival jitter.
+      ingestFrame(e.data.slice(8), blk, fps);
+    }else{
+      try{
+        const msg = JSON.parse(e.data);
+        if(msg.type === "status"){ overlayTitle.textContent = msg.message; }
+        else if(msg.type === "session_ended"){
+          setStatus("ended","Session ended"); setPlaying(false);
+          showOverlay("Session ended","GPU time limit reached — press Start world to continue.",false);
+        }else if(msg.type === "error"){
+          setStatus("error","Error"); setPlaying(false);
+          showOverlay("Something went wrong", msg.message || "Unknown error", false);
+        }
+      }catch{}
+    }
+  };
+  ws.onclose = ()=>{ if(state.playing) setTimeout(connectWS, 1000); };
+  state.ws = ws;
+}
+
+// --- Control loop: send held keys ~10x/s ---
+let ctrlTimer = null;
+function startCtrlLoop(){
+  if(ctrlTimer) return;
+  ctrlTimer = setInterval(()=>{
+    if(state.ws?.readyState === WebSocket.OPEN && state.playing){
+      state.ws.send(JSON.stringify({
+        type:"control",
+        buttons: Array.from(state.pressed),
+        prompt: state.prompt,
+      }));
+    }
+  }, 100);
+}
+function stopCtrlLoop(){ if(ctrlTimer){ clearInterval(ctrlTimer); ctrlTimer = null; } }
+
+function renderKeys(){
+  // Highlight the on-screen WASD + IJKL keys that are currently held.
+  keycaps.forEach(c=>c.classList.toggle("active", state.pressed.has(c.dataset.k)));
+  // Look/pan status badge reflects the held look keys (I J K L).
+  const look = Array.from(state.pressed).filter(k=>LOOK_KEYS.has(k));
+  const looking = look.length > 0;
+  lookBadge.classList.toggle("active", looking);
+  const LOOK_LABEL = { I:"up", K:"down", J:"left", L:"right" };
+  lookBadgeText.textContent = looking
+    ? "Looking " + look.map(k=>LOOK_LABEL[k]).join(" + ")
+    : (state.capturing ? "Look ready" : "Look idle");
+  const held = Array.from(state.pressed);
+  activeTags.textContent = held.length ? "Active: " + held.join(" + ") : "";
+}
+
+// --- Keyboard capture: WASD (move) + IJKL (look/pan) ---
+// Clicking the stage captures input; held keys are sent ~10x/s to the server.
+function requestCapture(){
+  state.capturing = true;
+  stage.focus();
+  renderKeys();
+}
+stage.addEventListener("click", requestCapture);
+
+document.addEventListener("keydown",(e)=>{
+  if(!state.capturing) return;
+  if(e.code === "Escape"){
+    state.capturing=false; state.pressed.clear();
+    renderKeys(); return;
+  }
+  const k = KEY_FROM_CODE[e.code];
+  if(k && (MOVE_KEYS.has(k) || LOOK_KEYS.has(k))){ e.preventDefault(); state.pressed.add(k); renderKeys(); }
+});
+document.addEventListener("keyup",(e)=>{
+  const k = KEY_FROM_CODE[e.code];
+  if(k){ state.pressed.delete(k); renderKeys(); }
+});
+
+promptEl.addEventListener("input", ()=>{ state.prompt = promptEl.value; });
+
+// --- Start / Stop ---
+async function startWorld(){
+  if(!state.seedPath) return;
+  state.prompt = promptEl.value;
+  await initClient();
+  setStatus("loading","Starting…");
+  showOverlay("Initializing GPU & world…","Requesting a ZeroGPU slot — first run can take a moment.", true);
+  startBtn.disabled = true;
+
+  connectWS();
+  startCtrlLoop();
+  try{
+    await state.client.predict("/start_game", {
+      session_id: state.sessionId,
+      seed_path: state.seedPath,
+      prompt: state.prompt || "",
+      seed: parseInt(seedEl.value || "42", 10),
+    });
+    // Auto-capture keyboard so the user can steer immediately.
+    state.capturing = true; stage.focus();
+  }catch(err){
+    console.error("start failed", err);
+    setStatus("error","Failed to start");
+    showOverlay("Failed to start", String(err?.message || err), false);
+    startBtn.disabled = false;
+  }
+}
+
+async function stopWorld(){
+  try{ if(state.client) await state.client.predict("/stop_game", { session_id: state.sessionId }); }catch{}
+  if(state.ws){ try{ state.ws.send(JSON.stringify({type:"stop"})); }catch{} state.ws.close(); state.ws=null; }
+  resetJitterBuffer();
+  stopCtrlLoop();
+  setPlaying(false);
+  state.pressed.clear(); renderKeys();
+  setStatus("idle","Idle");
+  showOverlay("Stopped","Press Start world to roll out again.", false);
+  refreshStartEnabled();
+}
+
+startBtn.addEventListener("click", startWorld);
+stopBtn.addEventListener("click", stopWorld);
+
+// --- Init ---
+setStatus("idle","Idle");
+state.prompt = promptEl.value;
+loadPresets();
+initClient().catch(console.error);
+</script>
+</body>
+</html>

requirements.txt

@@ -3,3 +3,4 @@ git+https://github.com/huggingface/diffusers.git@ffee5d2d9ccf86db5c21190e223ea49
 torch
 transformers
 ftfy
+python-multipart