Fix per user isolation (#4)

- Isolate per-user state to fix cross-user seed/frame race (713796d5dbe6d59915d1b4eea3d6ad7a8cca6959)
- Switch to Biome seed catalog and drop user uploads (fd21ff2ccd7da0855bbe378452c1590c04d2e9cc)
- World selector: 3-up grid and click-to-start (6352b84a1cd840fbe382d72a58531c75e96ed103)


Co-authored-by: Apolinário from multimodal AI art <multimodalart@users.noreply.huggingface.co>

app.py

@@ -3,6 +3,9 @@ Waypoint v1.5 Real-Time World Model Demo — gradio.Server + WebSocket Edition
 
 Uses gradio.Server for ZeroGPU-compatible game start/stop, and a raw WebSocket
 for real-time binary JPEG frame streaming + control input. No Gradio UI polling.
+
+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.
 """
 # Check for ZeroGPU environment - must be before other imports
 try:
@@ -22,17 +25,16 @@ import struct
 import contextvars
 import threading
 import time
+import uuid
 from collections import deque
 from dataclasses import dataclass, field
 from multiprocessing import Queue
-from typing import Optional, Set, Tuple
+from typing import Dict, Optional, Set, Tuple
 
 import torch
-import numpy as np
 from PIL import Image
-import tempfile
-from fastapi import WebSocket, WebSocketDisconnect, UploadFile, File
-from fastapi.responses import HTMLResponse, JSONResponse
+from fastapi import WebSocket, WebSocketDisconnect
+from fastapi.responses import HTMLResponse
 from gradio import Server
 
 from diffusers.modular_pipelines import ModularPipeline
@@ -44,6 +46,7 @@ IMAGE_WIDTH = 1024
 IMAGE_HEIGHT = 512
 MAX_FRAMES_BEFORE_RESET = 4096
 JPEG_QUALITY = 80
+SESSION_IDLE_TIMEOUT = 600  # seconds; janitor reaps abandoned sessions
 
 torch.set_float32_matmul_precision("medium")
 
@@ -66,36 +69,52 @@ constants_path = hf_hub_download(repo_id=MODEL_ID, filename="transformer/transfo
 print("Pipeline and AOT artifacts downloaded.")
 
 _aot_initialized = False
+_aot_init_lock = threading.Lock()
 
 
 def _ensure_aot_on_gpu():
     global _aot_initialized
     if _aot_initialized:
         return
-    broadcast_status("Moving pipeline to GPU...")
-    print("Initializing AOT model on GPU...")
-    pipe.to("cuda")
-    broadcast_status("Setting up KV cache...")
-    pipe.blocks.sub_blocks["before_denoise"].sub_blocks["setup_kv_cache"]._setup_kv_cache(
-        pipe.transformer, torch.device("cuda"), torch.bfloat16
-    )
-    broadcast_status("Loading AOT-compiled model...")
-    aot_model = torch._inductor.aoti_load_package(pt2_path)
-    constants_map = torch.load(constants_path, map_location="cuda", weights_only=True)
-    aot_model.load_constants(constants_map, check_full_update=True, user_managed=True)
-    pipe.transformer.forward = aot_model
-    _aot_initialized = True
-    broadcast_status("AOT model ready!")
-    print("AOT model loaded successfully!")
-
-
-SEED_FRAME_URLS = [
-    "https://huggingface.co/spaces/Overworld/waypoint-1-small/resolve/main/starter_18.png",
-    "https://huggingface.co/spaces/Overworld/waypoint-1-small/resolve/main/starter_9.png",
-    "https://huggingface.co/spaces/Overworld/waypoint-1-small/resolve/main/starter_22.png",
-    "https://huggingface.co/spaces/Overworld/waypoint-1-small/resolve/main/starter_14.png",
-    "https://huggingface.co/spaces/Overworld/waypoint-1-small/resolve/main/starter_21.png",
+    with _aot_init_lock:
+        if _aot_initialized:
+            return
+        broadcast_status("Moving pipeline to GPU...")
+        print("Initializing AOT model on GPU...")
+        pipe.to("cuda")
+        broadcast_status("Setting up KV cache...")
+        pipe.blocks.sub_blocks["before_denoise"].sub_blocks["setup_kv_cache"]._setup_kv_cache(
+            pipe.transformer, torch.device("cuda"), torch.bfloat16
+        )
+        broadcast_status("Loading AOT-compiled model...")
+        aot_model = torch._inductor.aoti_load_package(pt2_path)
+        constants_map = torch.load(constants_path, map_location="cuda", weights_only=True)
+        aot_model.load_constants(constants_map, check_full_update=True, user_managed=True)
+        pipe.transformer.forward = aot_model
+        _aot_initialized = True
+        broadcast_status("AOT model ready!")
+        print("AOT model loaded successfully!")
+
+
+_SEED_BASE = "https://github.com/Overworldai/Biome/blob/main/seeds"
+_SEED_NAMES = [
+    "abandoned_city_runner.jpg",
+    "alien_command_center.jpg",
+    "ancient_standing_stones.jpg",
+    "default.jpg",
+    "desert_outpost_hangar.jpg",
+    "enchanted_swamp_torch.jpg",
+    "frozen_crystal_cavern.jpg",
+    "frozen_valley_sniper.jpg",
+    "highland_castle_loch.jpg",
+    "mountain_ruins_gun.jpg",
+    "shattered_cockpit_nebula.jpg",
+    "shipwreck_shore_revolver.jpg",
+    "snowy_forest_tracks.jpg",
+    "stormy_countryside_rifle.jpg",
+    "sunken_city_depths.jpg",
 ]
+SEED_FRAME_URLS = [f"{_SEED_BASE}/{name}?raw=true" for name in _SEED_NAMES]
 
 
 def load_seed_frame(url: str, target_size: Tuple[int, int] = (IMAGE_HEIGHT, IMAGE_WIDTH)) -> Image.Image:
@@ -113,7 +132,6 @@ class GenerateCommand:
 
 @dataclass
 class ResetCommand:
-    seed_image: Optional[Image.Image] = None
     seed_url: Optional[str] = None
     prompt: str = "An explorable world"
 
@@ -126,40 +144,86 @@ class StopCommand:
 # --- Session ---
 @dataclass
 class GameSession:
+    session_id: str
     command_queue: Queue
     frame_queue: queue.Queue
     stop_event: threading.Event
+    status_queue: queue.Queue
     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()
-        self.command_queue.put(StopCommand())
+        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=3.0)
 
 
-# Global session (one player at a time on ZeroGPU)
-_session: Optional[GameSession] = None
-_session_lock = threading.Lock()
+# Per-client state, keyed by session_id (UUID generated by the browser).
+_sessions: Dict[str, GameSession] = {}
+_sessions_lock = threading.Lock()
 
-# Status queue for broadcasting init progress to WebSocket clients
-_status_queue: queue.Queue = queue.Queue(maxsize=20)
+# Contextvar carrying the active session's status queue. The worker thread
+# inherits the value via contextvars.copy_context(), so broadcast_status()
+# always lands in the right session's queue without explicit threading.
+_current_status_queue: contextvars.ContextVar[Optional[queue.Queue]] = contextvars.ContextVar(
+    "waypoint_status_queue", default=None
+)
 
 
 def broadcast_status(msg: str):
-    """Push a status message for WebSocket clients to pick up."""
+    """Push a status message to the current session's WebSocket consumer."""
+    q = _current_status_queue.get()
+    if q is None:
+        return
     try:
-        _status_queue.put_nowait(msg)
+        q.put_nowait(msg)
     except queue.Full:
         pass
 
 
-def gpu_worker_thread(seed_url, prompt, command_queue, frame_queue, stop_event, frame_times, seed_image=None):
+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():
+    """Background janitor: stops sessions whose worker died or that went idle."""
+    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)")
+
+
+threading.Thread(target=_reap_idle_sessions, daemon=True).start()
+
+
+def gpu_worker_thread(seed_url, prompt, command_queue, frame_queue, stop_event, frame_times):
     try:
         gen = create_gpu_game_loop(
             command_queue,
-            initial_seed_image=seed_image,
             initial_seed_url=seed_url,
             initial_prompt=prompt,
         )
@@ -194,7 +258,7 @@ def gpu_worker_thread(seed_url, prompt, command_queue, frame_queue, stop_event,
         print("Worker thread finished")
 
 
-def create_gpu_game_loop(command_queue, initial_seed_image=None, initial_seed_url=None, initial_prompt="An explorable world"):
+def create_gpu_game_loop(command_queue, initial_seed_url=None, initial_prompt="An explorable world"):
     @spaces.GPU(duration=120)
     def gpu_game_loop():
         broadcast_status("GPU allocated! Loading model...")
@@ -203,9 +267,7 @@ def create_gpu_game_loop(command_queue, initial_seed_image=None, initial_seed_ur
         print(f"Model loaded! (n_frames={n_frames})")
 
         broadcast_status("Loading seed image...")
-        if initial_seed_image is not None:
-            seed_image = initial_seed_image.resize((IMAGE_WIDTH, IMAGE_HEIGHT), Image.BILINEAR)
-        elif initial_seed_url is not None:
+        if initial_seed_url:
             seed_image = load_seed_frame(initial_seed_url)
         else:
             seed_image = load_seed_frame(random.choice(SEED_FRAME_URLS))
@@ -243,9 +305,7 @@ def create_gpu_game_loop(command_queue, initial_seed_image=None, initial_seed_ur
                 break
 
             if reset_command is not None:
-                if reset_command.seed_image is not None:
-                    seed_img = reset_command.seed_image.resize((IMAGE_WIDTH, IMAGE_HEIGHT), Image.BILINEAR)
-                elif reset_command.seed_url:
+                if reset_command.seed_url:
                     seed_img = load_seed_frame(reset_command.seed_url)
                 else:
                     seed_img = load_seed_frame(random.choice(SEED_FRAME_URLS))
@@ -281,91 +341,74 @@ app = Server()
 
 
 @app.api(name="start_game")
-def start_game(seed_url: str = "", prompt: str = "An explorable world") -> str:
-    """Start a new game session. Returns 'ok' immediately; GPU init happens in background."""
-    global _session
-    with _session_lock:
-        if _session is not None:
-            _session.stop()
-            _session = None
-
-    # Drain any stale status messages
-    while not _status_queue.empty():
-        try:
-            _status_queue.get_nowait()
-        except queue.Empty:
-            break
-
-    broadcast_status("Requesting GPU from ZeroGPU...")
+def start_game(session_id: str = "", seed_url: str = "", prompt: str = "An explorable world") -> str:
+    """Start a new game session for `session_id`. Returns the session_id used."""
+    if not session_id:
+        session_id = str(uuid.uuid4())
+
+    # Tear down any prior session for this client (e.g. a reload mid-game).
+    prior = _drop_session(session_id)
+    if prior is not None:
+        prior.stop()
+
+    # If no seed selected, pick a random one from the curated set so the
+    # player always lands somewhere instead of falling through to the model's
+    # default.
+    if not seed_url:
+        seed_url = random.choice(SEED_FRAME_URLS)
 
     command_queue = Queue()
     frame_queue = queue.Queue(maxsize=2)
     stop_event = threading.Event()
+    status_queue = queue.Queue(maxsize=20)
     frame_times = deque(maxlen=30)
 
-    # Use uploaded seed image if available, otherwise use URL
-    seed_image = _uploaded_seed_image
-
-    # GPU init + first frame generation happens in the worker thread,
-    # not here — avoids blocking the Gradio SSE queue (which caused 404 timeouts).
-    ctx = contextvars.copy_context()
-    worker = threading.Thread(
-        target=ctx.run,
-        args=(gpu_worker_thread,
-              seed_url if seed_url and not seed_image else None,
-              prompt or "An explorable world",
-              command_queue, frame_queue, stop_event, frame_times),
-        kwargs={"seed_image": seed_image},
-        daemon=True,
-    )
-    worker.start()
-
     session = GameSession(
+        session_id=session_id,
         command_queue=command_queue,
         frame_queue=frame_queue,
         stop_event=stop_event,
-        worker_thread=worker,
+        status_queue=status_queue,
         frame_times=frame_times,
     )
 
-    with _session_lock:
-        _session = session
+    with _sessions_lock:
+        _sessions[session_id] = session
 
-    return "ok"
+    # Bind the status queue into the contextvar so the worker thread (and the
+    # @spaces.GPU function it spawns) inherits it via copy_context().
+    token = _current_status_queue.set(status_queue)
+    try:
+        broadcast_status("Requesting GPU from ZeroGPU...")
+
+        ctx = contextvars.copy_context()
+        worker = threading.Thread(
+            target=ctx.run,
+            args=(gpu_worker_thread,
+                  seed_url,
+                  prompt or "An explorable world",
+                  command_queue, frame_queue, stop_event, frame_times),
+            daemon=True,
+        )
+        session.worker_thread = worker
+        worker.start()
+    finally:
+        _current_status_queue.reset(token)
+
+    return session_id
 
 
 @app.api(name="stop_game")
-def stop_game() -> str:
-    """Stop the active game session."""
-    global _session
-    with _session_lock:
-        if _session is not None:
-            _session.stop()
-            _session = None
+def stop_game(session_id: str = "") -> str:
+    """Stop the active game session 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"
 
 
-# Store for uploaded seed image (one at a time, single-player)
-_uploaded_seed_image: Optional[Image.Image] = None
-
-
-@app.post("/upload_seed_image")
-async def upload_seed_image(file: UploadFile = File(...)):
-    """Upload a custom seed image. Returns confirmation."""
-    global _uploaded_seed_image
-    data = await file.read()
-    _uploaded_seed_image = Image.open(io.BytesIO(data)).convert("RGB")
-    return JSONResponse({"status": "ok"})
-
-
-@app.post("/clear_seed_image")
-async def clear_seed_image():
-    """Clear the uploaded seed image."""
-    global _uploaded_seed_image
-    _uploaded_seed_image = None
-    return JSONResponse({"status": "ok"})
-
-
 @app.api(name="get_worlds")
 def get_worlds() -> list:
     """Return the list of seed world URLs."""
@@ -373,34 +416,37 @@ def get_worlds() -> list:
 
 
 @app.websocket("/ws")
-async def game_ws(websocket: WebSocket):
+async def game_ws(websocket: WebSocket, session_id: str = ""):
     """
-    Real-time game WebSocket.
-    - Server sends: binary messages = 8-byte header (uint32 frame_count, uint32 fps*10) + JPEG bytes
-    - Server sends: JSON {"type": "session_ended"} when GPU session expires
-    - Client sends: JSON {"type": "control", "buttons": [...], "mouse_x": f, "mouse_y": f, "prompt": "..."}
-    - Client sends: JSON {"type": "reset", "seed_url": "...", "prompt": "..."}
+    Real-time game WebSocket. Requires `?session_id=...` query param matching
+    the value passed to /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:
-            # Relay any status messages (init progress / errors)
-            try:
-                status_msg = _status_queue.get_nowait()
-                if status_msg.startswith("error:"):
-                    await websocket.send_json({"type": "error", "message": status_msg[6:]})
+            session = _get_session(session_id)
+
+            # Drain any status messages for this session (init progress / errors).
+            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
-                await websocket.send_json({"type": "status", "message": status_msg})
-            except queue.Empty:
-                pass
-            except (WebSocketDisconnect, RuntimeError):
-                break
 
-            with _session_lock:
-                session = _session
             if session is None:
                 await asyncio.sleep(0.05)
                 continue
@@ -423,6 +469,7 @@ async def game_ws(websocket: WebSocket):
                 jpeg_bytes = buf.getvalue()
                 header = struct.pack(">II", count, int(fps * 10))
                 await websocket.send_bytes(header + jpeg_bytes)
+                session.touch()
             except queue.Empty:
                 pass
             except (WebSocketDisconnect, RuntimeError):
@@ -432,10 +479,10 @@ async def game_ws(websocket: WebSocket):
         while True:
             try:
                 data = await websocket.receive_json()
-                with _session_lock:
-                    session = _session
+                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", []))
@@ -443,13 +490,10 @@ async def game_ws(websocket: WebSocket):
                     prompt = data.get("prompt", "An explorable world")
                     session.command_queue.put(GenerateCommand(buttons=buttons, mouse=mouse, prompt=prompt))
                 elif msg_type == "reset":
-                    seed_url = data.get("seed_url")
+                    seed_url = data.get("seed_url") or random.choice(SEED_FRAME_URLS)
                     prompt = data.get("prompt", "An explorable world")
-                    use_custom = data.get("use_custom_image", False)
-                    seed_image = _uploaded_seed_image if use_custom else None
                     session.command_queue.put(ResetCommand(
-                        seed_image=seed_image,
-                        seed_url=seed_url if not seed_image else None,
+                        seed_url=seed_url,
                         prompt=prompt,
                     ))
             except WebSocketDisconnect:

index.html

@@ -389,12 +389,13 @@
 
   /* --- World Selector --- */
   .world-grid {
-    display: flex;
+    display: grid;
+    grid-template-columns: repeat(3, 1fr);
     gap: 6px;
-    flex-wrap: wrap;
   }
   .world-thumb {
-    width: 56px; height: 28px;
+    width: 100%;
+    aspect-ratio: 16 / 9;
     border-radius: 4px;
     border: 2px solid transparent;
     cursor: pointer;
@@ -404,48 +405,6 @@
   .world-thumb:hover { border-color: var(--accent); transform: scale(1.05); }
   .world-thumb.selected { border-color: var(--accent); box-shadow: 0 0 8px rgba(88,166,255,0.4); }
 
-  /* --- Custom Image Upload --- */
-  .upload-area {
-    border: 2px dashed rgba(88,166,255,0.25);
-    border-radius: 8px;
-    padding: 12px;
-    text-align: center;
-    cursor: pointer;
-    transition: border-color 0.15s, background 0.15s;
-    position: relative;
-    min-height: 60px;
-    display: flex;
-    flex-direction: column;
-    align-items: center;
-    justify-content: center;
-    gap: 4px;
-  }
-  .upload-area:hover { border-color: var(--accent); background: rgba(88,166,255,0.05); }
-  .upload-area.has-image { border-style: solid; padding: 4px; }
-  .upload-area img {
-    max-width: 100%;
-    max-height: 80px;
-    object-fit: contain;
-    border-radius: 4px;
-  }
-  .upload-area .placeholder { font-size: 11px; color: var(--text-dim); }
-  .upload-area .icon { font-size: 20px; opacity: 0.5; }
-  .upload-clear {
-    position: absolute;
-    top: 4px; right: 4px;
-    background: rgba(255,107,107,0.8);
-    border: none;
-    color: #fff;
-    width: 18px; height: 18px;
-    border-radius: 50%;
-    font-size: 11px;
-    cursor: pointer;
-    display: flex;
-    align-items: center;
-    justify-content: center;
-    line-height: 1;
-  }
-
   /* --- Prompt --- */
   #prompt-input {
     width: 100%;
@@ -516,7 +475,7 @@
     .key-wide { min-height: 34px; padding: 6px 14px; font-size: 11px; }
     .joystick-ring { width: 90px; height: 90px; min-width: 90px; }
 
-    .world-thumb { width: 44px; height: 22px; }
+    .world-thumb { aspect-ratio: 16 / 9; }
   }
 </style>
 </head>
@@ -636,16 +595,6 @@
     <div class="world-grid" id="world-selector"></div>
   </div>
 
-  <!-- Custom Image Upload -->
-  <div class="sidebar-section">
-    <div class="section-label">Custom Start Image</div>
-    <div class="upload-area" id="upload-area">
-      <div class="icon">+</div>
-      <div class="placeholder">Click or drop image</div>
-    </div>
-    <input type="file" id="file-input" accept="image/*" style="display:none">
-  </div>
-
   <!-- Prompt -->
   <div class="sidebar-section">
     <div class="section-label">World Prompt</div>
@@ -657,7 +606,13 @@
 import { Client } from "https://cdn.jsdelivr.net/npm/@gradio/client/dist/index.min.js";
 
 // --- State ---
+// One session_id per browser tab — keeps every server-side queue, status
+// stream, and uploaded seed isolated from other concurrent users.
+const SESSION_ID = (crypto.randomUUID && crypto.randomUUID()) ||
+  (`s-${Date.now().toString(36)}-${Math.random().toString(36).slice(2, 10)}`);
+
 const state = {
+  sessionId: SESSION_ID,
   playing: false,
   capturing: false,
   pressedKeys: new Set(),
@@ -665,7 +620,6 @@ const state = {
   mouseVelocity: { x: 0, y: 0 },
   lastMouseMove: 0,
   selectedSeedUrl: "",
-  useCustomImage: false,
   prompt: "An explorable world",
   ws: null,
   client: null,
@@ -705,8 +659,6 @@ const ctrlStatusText = document.getElementById("ctrl-status-text");
 const activeTags = document.getElementById("active-tags");
 const mouseXVal = document.getElementById("mouse-x-val");
 const mouseYVal = document.getElementById("mouse-y-val");
-const uploadArea = document.getElementById("upload-area");
-const fileInput = document.getElementById("file-input");
 const endedOverlay = document.getElementById("ended-overlay");
 const restartBtn = document.getElementById("restart-btn");
 
@@ -720,7 +672,7 @@ async function initClient() {
 // --- WebSocket ---
 function connectWS() {
   const proto = location.protocol === "https:" ? "wss:" : "ws:";
-  const ws = new WebSocket(`${proto}//${location.host}/ws`);
+  const ws = new WebSocket(`${proto}//${location.host}/ws?session_id=${encodeURIComponent(state.sessionId)}`);
   ws.binaryType = "arraybuffer";
 
   ws.onmessage = (e) => {
@@ -1027,81 +979,45 @@ joystick.addEventListener("touchend", (e) => { e.preventDefault(); resetJoystick
 joystick.addEventListener("touchcancel", (e) => { e.preventDefault(); resetJoystick(); }, { passive: false });
 
 // --- World Selector ---
-const SEED_URLS = [
-  "https://huggingface.co/spaces/Overworld/waypoint-1-small/resolve/main/starter_18.png",
-  "https://huggingface.co/spaces/Overworld/waypoint-1-small/resolve/main/starter_9.png",
-  "https://huggingface.co/spaces/Overworld/waypoint-1-small/resolve/main/starter_22.png",
-  "https://huggingface.co/spaces/Overworld/waypoint-1-small/resolve/main/starter_14.png",
-  "https://huggingface.co/spaces/Overworld/waypoint-1-small/resolve/main/starter_21.png",
+const SEED_BASE = "https://github.com/Overworldai/Biome/blob/main/seeds";
+const SEED_NAMES = [
+  "abandoned_city_runner.jpg",
+  "alien_command_center.jpg",
+  "ancient_standing_stones.jpg",
+  "default.jpg",
+  "desert_outpost_hangar.jpg",
+  "enchanted_swamp_torch.jpg",
+  "frozen_crystal_cavern.jpg",
+  "frozen_valley_sniper.jpg",
+  "highland_castle_loch.jpg",
+  "mountain_ruins_gun.jpg",
+  "shattered_cockpit_nebula.jpg",
+  "shipwreck_shore_revolver.jpg",
+  "snowy_forest_tracks.jpg",
+  "stormy_countryside_rifle.jpg",
+  "sunken_city_depths.jpg",
 ];
+const SEED_URLS = SEED_NAMES.map((name) => `${SEED_BASE}/${name}?raw=true`);
 
 SEED_URLS.forEach((url, i) => {
   const img = document.createElement("img");
   img.src = url;
   img.className = "world-thumb";
-  img.title = `World ${i + 1}`;
+  img.title = SEED_NAMES[i].replace(/\.jpg$/, "").replace(/_/g, " ");
   img.addEventListener("click", () => {
     document.querySelectorAll(".world-thumb").forEach(t => t.classList.remove("selected"));
     img.classList.add("selected");
     state.selectedSeedUrl = url;
-    state.useCustomImage = false;
-    // If playing, reset to this world
+    // If playing, reset to this world; otherwise kick off a new game.
     if (state.playing && state.ws?.readyState === WebSocket.OPEN) {
       state.ws.send(JSON.stringify({ type: "reset", seed_url: url, prompt: state.prompt }));
+    } else if (!state.playing && !startBtn.disabled) {
+      startGame();
     }
   });
   worldSelector.appendChild(img);
 });
 
-// --- Custom Image Upload ---
-uploadArea.addEventListener("click", () => fileInput.click());
-uploadArea.addEventListener("dragover", (e) => { e.preventDefault(); uploadArea.style.borderColor = "var(--accent)"; });
-uploadArea.addEventListener("dragleave", () => { uploadArea.style.borderColor = ""; });
-uploadArea.addEventListener("drop", (e) => {
-  e.preventDefault();
-  uploadArea.style.borderColor = "";
-  const file = e.dataTransfer?.files?.[0];
-  if (file && file.type.startsWith("image/")) handleImageUpload(file);
-});
-fileInput.addEventListener("change", () => {
-  if (fileInput.files?.[0]) handleImageUpload(fileInput.files[0]);
-});
-
-async function handleImageUpload(file) {
-  // Show preview
-  const reader = new FileReader();
-  reader.onload = (e) => {
-    uploadArea.innerHTML = `<img src="${e.target.result}" alt="Custom seed"><button class="upload-clear" id="clear-upload">&times;</button>`;
-    uploadArea.classList.add("has-image");
-    document.getElementById("clear-upload").addEventListener("click", (ev) => {
-      ev.stopPropagation();
-      clearUpload();
-    });
-  };
-  reader.readAsDataURL(file);
-
-  // Upload to server
-  const formData = new FormData();
-  formData.append("file", file);
-  try {
-    await fetch("/upload_seed_image", { method: "POST", body: formData });
-    state.useCustomImage = true;
-    // Deselect world thumbs
-    document.querySelectorAll(".world-thumb").forEach(t => t.classList.remove("selected"));
-    state.selectedSeedUrl = "";
-  } catch (err) {
-    console.error("Upload failed:", err);
-  }
-}
-
-async function clearUpload() {
-  uploadArea.innerHTML = '<div class="icon">+</div><div class="placeholder">Click or drop image</div>';
-  uploadArea.classList.remove("has-image");
-  state.useCustomImage = false;
-  fileInput.value = "";
-  try { await fetch("/clear_seed_image", { method: "POST" }); } catch {}
-}
-
 // --- Prompt ---
 promptInput.addEventListener("input", () => {
   state.prompt = promptInput.value || "An explorable world";
@@ -1122,6 +1038,7 @@ async function startGame() {
 
   try {
     await state.client.predict("/start_game", {
+      session_id: state.sessionId,
       seed_url: state.selectedSeedUrl || "",
       prompt: state.prompt,
     });
@@ -1138,7 +1055,7 @@ restartBtn.addEventListener("click", startGame);
 
 stopBtn.addEventListener("click", async () => {
   if (!state.client) return;
-  try { await state.client.predict("/stop_game"); } catch {}
+  try { await state.client.predict("/stop_game", { session_id: state.sessionId }); } catch {}
   setPlaying(false);
   endedOverlay.classList.add("hidden");
   setStatus("idle", "Idle");
@@ -1153,7 +1070,6 @@ resetBtn.addEventListener("click", () => {
     state.ws.send(JSON.stringify({
       type: "reset",
       seed_url: state.selectedSeedUrl || "",
-      use_custom_image: state.useCustomImage,
       prompt: state.prompt,
     }));
   }