diffusers-internal-dev/abot-world-interactive
0
1"""2ABot-World on Modular Diffusers — interactive action-conditioned world rollout.3gradio.Server + WebSocket live-backend edition.4 5The engine is one `pipe.stream(actions=<callable>)` loop over the Modular Diffusers6`ABotWorldStreamingBlocks` preset (huggingface/diffusers#14159): the pipeline polls the7held keys once per generated block and yields every block's decoded frames.8The serving shell (gradio.Server, WebSocket frame stream, pacing, per-session queues,9front-end) is reused from https://huggingface.co/spaces/acvlab/abot-world-interactive.10 11Given an uploaded starting image (i2v conditioning), a scene prompt, and live12WASD / IJKL controls, the model autoregressively rolls out an action-conditioned13navigable world and streams decoded frames to the browser over a WebSocket.14 15This mirrors the live backend/infrastructure of16https://huggingface.co/spaces/Overworld/waypoint-1-5 (gradio.Server for17ZeroGPU-friendly start/stop + a raw WebSocket for real-time binary JPEG frame18streaming and control input), with a cleaner custom UI and image-upload seeding.19 20Multi-user safe: every endpoint is keyed by a per-client `session_id` so21concurrent players never share seed images, frame queues, or status messages.22 23ZeroGPU quota: the incoming request's ZeroGPU proxy token (the `x-ip-token` /24`x-api-token` header injected by the HF iframe) is captured per-session and25propagated into the worker thread's gradio request context, so the GPU work is26billed against the *requesting user's* quota — not the Space owner's.27 28Upstream: https://github.com/amap-cvlab/ABot-World29Model: https://huggingface.co/YiYiXu/ABot-World-0-5B-LF-Diffusers (built on Wan2.2-TI2V-5B)30"""31import os32os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")33 34import spaces # must precede torch / CUDA-touching imports35 36import io37import time38import queue39import asyncio40import struct41import tempfile42import threading43import contextvars44import uuid45from collections import deque46from dataclasses import dataclass, field47from multiprocessing import Queue as MPQueue48from pathlib import Path49from typing import Dict, Optional, Set50 51import numpy as np52import torch53from PIL import Image54 55from fastapi import UploadFile, File, WebSocket, WebSocketDisconnect56from fastapi.responses import HTMLResponse, JSONResponse, FileResponse57from gradio import Server58from gradio.context import LocalContext59 60from diffusers.modular_pipelines import ABotWorldStreamingBlocks61 62# ── Repo paths ───────────────────────────────────────────────────────────────63APP_DIR = Path(__file__).resolve().parent64 65MODEL_ID = "YiYiXu/ABot-World-0-5B-LF-Diffusers"66 67# Preset starting-world images bundled with the Space (sourced from the ABot-World68# repo). Shown in the UI as clickable thumbnails that seed the i2v rollout directly.69EXAMPLES_DIR = APP_DIR / "examples"70EXAMPLE_SEEDS = [71 {"name": "desert_valley.png", "label": "Desert valley"},72 {"name": "forest_stream.png", "label": "Forest stream"},73 {"name": "mountain_meadow.png", "label": "Mountain meadow"},74 {"name": "example.png", "label": "Sample scene"},75]76 77# ── Stream / rollout configuration ───────────────────────────────────────────78# 704x1280 is the native training resolution used by the upstream web client.79STREAM_HEIGHT = 70480STREAM_WIDTH = 128081JPEG_QUALITY = 8282MAX_BLOCKS_PER_SESSION = 512 # hard cap so a session can't run forever83SESSION_IDLE_TIMEOUT = 600 # seconds; janitor reaps abandoned sessions84GPU_DURATION = 150 # seconds per @spaces.GPU allocation (one session)85 86# ── Real-time pacing configuration ───────────────────────────────────────────87# The GPU decodes a whole block (12 frames) at once, so all of a block's88# frames become available in a burst. If we forward them to the browser the89# instant they finish, the client sees N frames clustered together followed by a90# gap while the next block generates — the fps counter averages out fine, but91# the *felt* cadence is bursty. To deliver a steady real-time stream we pace the92# frames of each block evenly across the time we expect one block to take93# (mirroring the official ABot-World web_client's block-frame spreader), and we94# smooth the per-block generation time with an EMA so a single slow/fast block95# doesn't cause a visible speed-up/slow-down. See gpu_worker_thread().96PACING_EMA_ALPHA = 0.25 # smoothing factor for per-block generation time97MIN_PACING_SLEEP = 0.004 # don't bother sleeping for sub-4ms slices98DEFAULT_BLOCK_SECONDS = 0.5 # initial per-block estimate before first measure99 100# Actions map to the 8-key one-hot the model was trained on (W A S D I J K L).101# The browser sends the currently-held key set; we translate to this dict.102KEY_ORDER = ["W", "A", "S", "D", "I", "J", "K", "L"]103 104DEFAULT_PROMPT = (105 "A realistic outdoor world scene with a navigable path, natural lighting, "106 "detailed ground texture, and stable forward motion."107)108 109# ── Build the Modular Diffusers streaming pipeline (module scope) ────────────110# On ZeroGPU `pipe.to("cuda")` at import only packs the weights; they land on the111# GPU inside the @spaces.GPU rollout.112print(f"[startup] loading {MODEL_ID} ...", flush=True)113torch.set_grad_enabled(False)114pipe = ABotWorldStreamingBlocks().init_pipeline(MODEL_ID)115pipe.load_components(dtype=torch.bfloat16)116pipe.to("cuda")117print("[startup] pipeline ready.", flush=True)118 119# Only one rollout may touch the shared pipeline at a time.120_infer_lock = threading.Lock()121 122 123def _action_from_buttons(buttons):124 """Translate a set of held key names (e.g. {'W','A'}) into the model's 8-key multi-hot action."""125 held = {k.upper() for k in (buttons or [])}126 return [int(k in held) for k in KEY_ORDER]127 128 129# ── Command types (browser -> worker) ────────────────────────────────────────130@dataclass131class ControlCommand:132 buttons: Set[str]133 prompt: str134 135 136@dataclass137class StopCommand:138 pass139 140 141# ── Per-session state ────────────────────────────────────────────────────────142# NOTE on queues: the @spaces.GPU rollout runs in a forked subprocess, so any143# object it reads must cross the fork boundary. `command_queue` is therefore a144# multiprocessing Queue (browser controls / stop reach the GPU loop through it).145# `frame_queue` / `status_queue` are plain queue.Queue used only in the parent146# process (frames arrive back via the ZeroGPU generator IPC and are forwarded147# to the WebSocket by the worker thread).148@dataclass149class GameSession:150 session_id: str151 command_queue: "MPQueue"152 frame_queue: "queue.Queue"153 status_queue: "queue.Queue"154 stop_event: threading.Event155 seed_path: str156 prompt: str157 seed: int158 worker_thread: Optional[threading.Thread] = None159 frame_times: deque = field(default_factory=lambda: deque(maxlen=30))160 last_active: float = field(default_factory=time.time)161 162 def touch(self):163 self.last_active = time.time()164 165 def stop(self):166 self.stop_event.set()167 try:168 self.command_queue.put_nowait(StopCommand())169 except Exception:170 pass171 if self.worker_thread and self.worker_thread.is_alive():172 self.worker_thread.join(timeout=4.0)173 174 175_sessions: Dict[str, GameSession] = {}176_sessions_lock = threading.Lock()177 178# Contextvar carrying the active session's status queue (inherited by the worker179# thread via contextvars.copy_context()).180_current_status_queue: "contextvars.ContextVar[Optional[queue.Queue]]" = contextvars.ContextVar(181 "abot_status_queue", default=None182)183 184 185def broadcast_status(msg: str):186 q = _current_status_queue.get()187 if q is None:188 return189 try:190 q.put_nowait(msg)191 except queue.Full:192 pass193 194 195def _get_session(session_id: str) -> Optional[GameSession]:196 with _sessions_lock:197 return _sessions.get(session_id)198 199 200def _drop_session(session_id: str) -> Optional[GameSession]:201 with _sessions_lock:202 return _sessions.pop(session_id, None)203 204 205def _reap_idle_sessions():206 while True:207 time.sleep(60)208 now = time.time()209 to_drop = []210 with _sessions_lock:211 for sid, sess in list(_sessions.items()):212 worker_dead = sess.worker_thread is None or not sess.worker_thread.is_alive()213 idle = (now - sess.last_active) > SESSION_IDLE_TIMEOUT214 if worker_dead and idle:215 to_drop.append(sid)216 for sid in to_drop:217 _sessions.pop(sid, None)218 if to_drop:219 print(f"Janitor reaped {len(to_drop)} idle session(s)", flush=True)220 221 222threading.Thread(target=_reap_idle_sessions, daemon=True).start()223 224 225# ── GPU worker ───────────────────────────────────────────────────────────────226def gpu_worker_thread(session: "GameSession"):227 """Parent-thread driver: consumes frames yielded by the ZeroGPU generator,228 computes FPS, and forwards frames to the WebSocket via `frame_queue`.229 230 Status/stop live in the parent process; the GPU loop is steered purely231 through the (picklable, cross-fork) `command_queue`.232 """233 try:234 broadcast_status("GPU allocated — starting world…")235 gen = create_gpu_rollout_loop(236 session.command_queue, session.seed_path, session.prompt, session.seed,237 )238 first = True239 # Steady send clock: `next_send` is the monotonic time at which the next240 # frame *should* be delivered. Each frame's slot is one smoothed241 # inter-frame interval after the previous, so frames leave at a constant242 # cadence regardless of the bursty block boundaries. `block_seconds` is243 # an EMA of measured per-block generation time (frames/block ÷ that gives244 # the target inter-frame interval).245 block_seconds = DEFAULT_BLOCK_SECONDS246 next_send = None247 while not session.stop_event.is_set():248 try:249 frame, block_idx, frame_idx, frames_in_block, block_elapsed = next(gen)250 except StopIteration:251 print("Rollout generator exhausted", flush=True)252 break253 except Exception as e:254 if "aborted" in str(e).lower() or "duration" in str(e).lower():255 print(f"GPU time expired: {e}", flush=True)256 else:257 print(f"Worker error: {e}", flush=True)258 broadcast_status(f"error:{e}")259 break260 261 if first:262 broadcast_status("Rolling out — use WASD / IJKL to steer.")263 first = False264 265 # Update the smoothed per-block time on the first frame of each block266 # (block_elapsed is constant across a block's frames).267 if frame_idx == 0 and block_elapsed > 0:268 block_seconds = (269 PACING_EMA_ALPHA * block_elapsed270 + (1.0 - PACING_EMA_ALPHA) * block_seconds271 )272 fpb = max(1, frames_in_block)273 interval = block_seconds / fpb # target seconds between frames274 275 # ── Steady-cadence gate ──────────────────────────────────────────276 # Hold each frame until its scheduled slot so the parent emits at a277 # constant interval instead of dumping a whole block at once.278 now = time.time()279 if next_send is None:280 next_send = now281 sleep_for = next_send - now282 if sleep_for > MIN_PACING_SLEEP:283 # Wake early if a stop is requested so we stay responsive.284 if session.stop_event.wait(timeout=sleep_for):285 break286 now = time.time()287 # Advance the schedule; if we've fallen far behind (e.g. a long GPU288 # stall), resync to now so we don't try to "catch up" in a burst.289 next_send = max(now, next_send + interval)290 291 now = time.time()292 session.frame_times.append(now)293 fps = 0.0294 if len(session.frame_times) >= 2:295 elapsed = session.frame_times[-1] - session.frame_times[0]296 fps = (len(session.frame_times) - 1) / elapsed if elapsed > 0 else 0.0297 # Keep only the freshest frame if the consumer fell behind: coalesce298 # stale frames rather than letting them queue up and flush in a burst.299 while session.frame_queue.qsize() > 1:300 try:301 session.frame_queue.get_nowait()302 except queue.Empty:303 break304 try:305 session.frame_queue.put_nowait((frame, block_idx, round(fps, 1)))306 except queue.Full:307 pass308 finally:309 session.stop_event.set()310 print("Worker thread finished", flush=True)311 312 313def create_gpu_rollout_loop(command_queue, seed_path, prompt_text, seed):314 """Return a ZeroGPU generator that rolls the world out block-by-block.315 316 Only picklable primitives + the multiprocessing `command_queue` cross the317 fork boundary. Live controls (held key set) and stop arrive via that queue.318 """319 @spaces.GPU(duration=GPU_DURATION)320 def gpu_rollout():321 prompt = (prompt_text or DEFAULT_PROMPT).strip() or DEFAULT_PROMPT322 image = Image.open(seed_path).convert("RGB")323 state = {"action": _action_from_buttons({"W"}), "block_start": time.time()} # default: forward324 325 def action_source(block_index):326 """Polled by the pipeline once per block: newest held-key set wins, None stops the rollout."""327 if block_index >= MAX_BLOCKS_PER_SESSION:328 return None329 while True:330 try:331 cmd = command_queue.get_nowait()332 except Exception:333 break334 if isinstance(cmd, StopCommand):335 return None336 if isinstance(cmd, ControlCommand):337 state["action"] = _action_from_buttons(cmd.buttons)338 state["block_start"] = time.time()339 return state["action"]340 341 with _infer_lock:342 events = pipe.stream(343 prompt=prompt,344 image=image,345 height=STREAM_HEIGHT,346 width=STREAM_WIDTH,347 actions=action_source,348 generator=torch.Generator("cpu").manual_seed(int(seed)),349 )350 for event in events:351 if event.path != "denoise.rollout":352 continue # inner per-denoise-step events353 # Time the full generate+decode of one block so the parent thread354 # can pace this block's frames over that duration.355 block_elapsed = time.time() - state["block_start"]356 b = event.loop_kwargs["k"]357 frames = (event.state.get("frames") * 255).clip(0, 255).astype(np.uint8)358 n = len(frames)359 for i, f in enumerate(frames):360 # (frame, block_idx, frame_idx_in_block, frames_in_block,361 # block_elapsed) — the pacing metadata lets the parent362 # spread this block's frames evenly rather than bursting.363 yield (f, b, i, n, block_elapsed)364 365 return gpu_rollout()366 367 368# ── App (gradio.Server) ──────────────────────────────────────────────────────369app = Server()370 371 372@app.api(name="start_game")373def start_game(session_id: str = "", seed_path: str = "",374 prompt: str = "", seed: int = 42) -> str:375 """Start a new interactive world rollout for `session_id`.376 377 Args:378 session_id: per-client id (UUID) isolating this player's stream.379 seed_path: filepath (uploaded via /upload) of the starting frame image380 that seeds the i2v world rollout.381 prompt: scene description.382 seed: RNG seed for reproducibility.383 384 Returns:385 The session_id actually used.386 """387 if not session_id:388 session_id = str(uuid.uuid4())389 390 prior = _drop_session(session_id)391 if prior is not None:392 prior.stop()393 394 if not seed_path:395 raise ValueError("A starting image is required — please upload one first.")396 397 command_queue = MPQueue() # crosses the ZeroGPU fork boundary398 frame_queue: "queue.Queue" = queue.Queue(maxsize=4)399 status_queue: "queue.Queue" = queue.Queue(maxsize=32)400 stop_event = threading.Event()401 402 session = GameSession(403 session_id=session_id,404 command_queue=command_queue,405 frame_queue=frame_queue,406 status_queue=status_queue,407 stop_event=stop_event,408 seed_path=seed_path,409 prompt=prompt or DEFAULT_PROMPT,410 seed=int(seed),411 )412 with _sessions_lock:413 _sessions[session_id] = session414 415 # Capture the *incoming request* — HF has already injected this user's416 # ZeroGPU proxy token (x-ip-token / x-api-token) into its headers. We417 # re-set it into the worker thread's gradio LocalContext so that418 # @spaces.GPU bills GPU time against THIS user's quota, not the owner's.419 gradio_request = LocalContext.request.get(None)420 status_token = _current_status_queue.set(status_queue)421 try:422 broadcast_status("Requesting GPU from ZeroGPU…")423 424 def _thread_entry():425 # Re-establish the request context inside the worker thread so the426 # ZeroGPU scheduler reads the requesting user's token.427 if gradio_request is not None:428 try:429 LocalContext.request.set(gradio_request)430 except Exception:431 pass432 gpu_worker_thread(session)433 434 ctx = contextvars.copy_context()435 worker = threading.Thread(target=ctx.run, args=(_thread_entry,), daemon=True)436 session.worker_thread = worker437 worker.start()438 finally:439 _current_status_queue.reset(status_token)440 441 return session_id442 443 444@app.api(name="stop_game")445def stop_game(session_id: str = "") -> str:446 """Stop the active rollout for the given client."""447 if not session_id:448 return "no_session"449 session = _drop_session(session_id)450 if session is not None:451 session.stop()452 return "stopped"453 454 455@app.websocket("/ws")456async def game_ws(websocket: WebSocket, session_id: str = ""):457 """Real-time rollout WebSocket. Requires `?session_id=...` matching /start_game."""458 await websocket.accept()459 if not session_id:460 await websocket.send_json({"type": "error", "message": "missing session_id"})461 await websocket.close(code=1008)462 return463 464 loop = asyncio.get_event_loop()465 466 async def send_frames():467 session_ended_sent = False468 while True:469 session = _get_session(session_id)470 471 if session is not None:472 try:473 status_msg = session.status_queue.get_nowait()474 if status_msg.startswith("error:"):475 await websocket.send_json({"type": "error", "message": status_msg[6:]})476 break477 await websocket.send_json({"type": "status", "message": status_msg})478 except queue.Empty:479 pass480 except (WebSocketDisconnect, RuntimeError):481 break482 483 if session is None:484 await asyncio.sleep(0.05)485 continue486 if session.stop_event.is_set() and session.frame_queue.empty():487 if not session_ended_sent:488 try:489 await websocket.send_json({"type": "session_ended"})490 except (WebSocketDisconnect, RuntimeError):491 break492 session_ended_sent = True493 await asyncio.sleep(0.4)494 continue495 try:496 result = await loop.run_in_executor(497 None, lambda s=session: s.frame_queue.get(timeout=0.1)498 )499 frame, count, fps = result500 img = Image.fromarray(frame)501 buf = io.BytesIO()502 img.save(buf, format="JPEG", quality=JPEG_QUALITY)503 jpeg_bytes = buf.getvalue()504 header = struct.pack(">II", int(count), int(fps * 10))505 await websocket.send_bytes(header + jpeg_bytes)506 session.touch()507 except queue.Empty:508 pass509 except (WebSocketDisconnect, RuntimeError):510 break511 512 async def receive_controls():513 while True:514 try:515 data = await websocket.receive_json()516 session = _get_session(session_id)517 if session is None:518 continue519 session.touch()520 msg_type = data.get("type", "control")521 if msg_type == "control":522 buttons = set(data.get("buttons", []))523 prompt = data.get("prompt", session.prompt)524 try:525 session.command_queue.put_nowait(526 ControlCommand(buttons=buttons, prompt=prompt)527 )528 except queue.Full:529 pass530 elif msg_type == "stop":531 session.stop()532 except WebSocketDisconnect:533 break534 except Exception:535 break536 537 try:538 await asyncio.gather(send_frames(), receive_controls())539 except WebSocketDisconnect:540 pass541 542 543@app.post("/upload_seed")544async def upload_seed(file: UploadFile = File(...)):545 """Accept a user-uploaded starting image and stash it server-side.546 547 Returns the temp filepath, which the browser then passes to /start_game as548 `seed_path` to seed the image-to-video (i2v) world rollout. Only an image is549 needed — there is no video upload.550 """551 try:552 raw = await file.read()553 img = Image.open(io.BytesIO(raw)).convert("RGB")554 except Exception:555 return JSONResponse({"error": "Could not read image file."}, status_code=400)556 557 tmp = tempfile.NamedTemporaryFile(prefix="abot_seed_", suffix=".png", delete=False)558 img.save(tmp.name, format="PNG")559 return {"seed_path": tmp.name}560 561 562def _safe_example_path(name: str) -> Optional[Path]:563 """Resolve `name` to a bundled example image, guarding against traversal."""564 if not any(name == e["name"] for e in EXAMPLE_SEEDS):565 return None566 path = (EXAMPLES_DIR / name).resolve()567 if EXAMPLES_DIR.resolve() not in path.parents or not path.is_file():568 return None569 return path570 571 572@app.get("/example_seeds")573async def example_seeds():574 """List the preset starting-world images available as clickable thumbnails."""575 return {"examples": [e for e in EXAMPLE_SEEDS if (EXAMPLES_DIR / e["name"]).is_file()]}576 577 578@app.get("/example_thumb")579async def example_thumb(name: str = ""):580 """Serve a preset starting-world image (for thumbnail display in the UI)."""581 path = _safe_example_path(name)582 if path is None:583 return JSONResponse({"error": "unknown example"}, status_code=404)584 return FileResponse(str(path), media_type="image/png")585 586 587@app.get("/example_seed")588async def example_seed(name: str = ""):589 """Seed the i2v rollout from a bundled preset image (no upload required).590 591 Copies the chosen example into a server-side temp file and returns its path,592 mirroring /upload_seed so the browser can pass it to /start_game as seed_path.593 """594 path = _safe_example_path(name)595 if path is None:596 return JSONResponse({"error": "unknown example"}, status_code=404)597 try:598 img = Image.open(path).convert("RGB")599 except Exception:600 return JSONResponse({"error": "could not read example image"}, status_code=500)601 tmp = tempfile.NamedTemporaryFile(prefix="abot_seed_", suffix=".png", delete=False)602 img.save(tmp.name, format="PNG")603 return {"seed_path": tmp.name}604 605 606@app.get("/", response_class=HTMLResponse)607async def homepage():608 html_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "index.html")609 with open(html_path, "r", encoding="utf-8") as f:610 return f.read()611 612 613# Avoid ZeroGPU "no GPU function" error at boot.614spaces.GPU(lambda: None)615 616app.launch(server_name="0.0.0.0", server_port=7860, ssr_mode=False) # the SSR Node proxy does not forward the /ws upgrade617