update

app.py

@@ -2,18 +2,25 @@
 WorldEngine Real-Time World Model Demo - ZeroGPU Edition
 
 A Gradio demo optimized for HuggingFace ZeroGPU with:
+- Generator-based GPU session that stays alive
 - Persistent compilation cache for faster cold starts
-- Warmup pass to pre-compile kernels
-- @spaces.GPU decorator support
+- Command queue for real-time control
 """
+# Check for ZeroGPU environment - must be before other imports
+try:
+    import spaces
+    IS_ZERO_GPU = True
+    print("ZeroGPU environment detected")
+except ImportError:
+    IS_ZERO_GPU = False
+    print("Running in standard GPU mode")
 
 import base64
 import os
 import random
-import threading
-import time
-from dataclasses import dataclass, field
+from dataclasses import dataclass
 from io import BytesIO
+from multiprocessing import Queue
 from pathlib import Path
 from typing import Optional, Set, Tuple
 
@@ -26,14 +33,11 @@ from diffusers import ModularPipeline
 from diffusers.utils import load_image
 
 # --- ZeroGPU Compilation Cache Setup ---
-# Set persistent cache directory BEFORE importing torch.compile
 CACHE_DIR = Path.home() / ".cache" / "world_engine_compile"
 CACHE_DIR.mkdir(parents=True, exist_ok=True)
 
 os.environ["TORCHINDUCTOR_CACHE_DIR"] = str(CACHE_DIR)
 os.environ["TORCH_COMPILE_CACHE_DIR"] = str(CACHE_DIR)
-
-# Enable FX graph caching for faster recompilation
 os.environ["TORCHINDUCTOR_FX_GRAPH_CACHE"] = "1"
 
 print(f"Compilation cache directory: {CACHE_DIR}")
@@ -42,20 +46,10 @@ torch._dynamo.config.recompile_limit = 64
 torch.set_float32_matmul_precision("medium")
 torch._dynamo.config.capture_scalar_outputs = True
 
-# Check for ZeroGPU environment
-try:
-    import spaces
-    IS_ZERO_GPU = True
-    print("ZeroGPU environment detected")
-except ImportError:
-    IS_ZERO_GPU = False
-    print("Running in standard GPU mode")
-
 # --- Configuration ---
 MODEL_ID = os.environ.get("MODEL_PATH", "diffusers-internal-dev/world-engine-modular")
-MAX_SEED = np.iinfo(np.int32).max
+USE_COMPILE = os.environ.get("COMPILE", "1") == "1"
 
-# Seed frame URLs for reset
 SEED_FRAME_URLS = [
     "https://gist.github.com/user-attachments/assets/5d91c49a-2ae9-418f-99c0-e93ae387e1de",
     "https://gist.github.com/user-attachments/assets/4adc5a3d-6980-4d1e-b6e8-9033cdf61c66",
@@ -80,297 +74,216 @@ def image_to_base64(image: Image.Image) -> str:
     return f"data:image/png;base64,{img_str}"
 
 
-# --- Control State ---
+# --- Command Types ---
+@dataclass
+class GenerateCommand:
+    """Generate next frame with given controls."""
+    buttons: Set[int]
+    mouse: Tuple[float, float]
+    prompt: str
+
+
 @dataclass
-class ControlState:
-    """Thread-safe control state."""
-
-    buttons: Set[int] = field(default_factory=set)
-    mouse: Tuple[float, float] = (0.0, 0.0)
-    _lock: threading.Lock = field(default_factory=threading.Lock)
-
-    def update(self, buttons: Set[int], mouse: Tuple[float, float]):
-        with self._lock:
-            self.buttons = buttons.copy()
-            self.mouse = mouse
-
-    def get(self) -> Tuple[Set[int], Tuple[float, float]]:
-        with self._lock:
-            return self.buttons.copy(), self.mouse
-
-
-# --- Game State with Background Generation ---
-class GameState:
-    """Global game state with background frame generation."""
-
-    def __init__(self):
-        self.pipe: Optional[ModularPipeline] = None
-        self.state = None
-        self.control_state = ControlState()
-        self.frame_count = 0
-        self.n_frames = None
-        self.is_initialized = False
-        self.is_warming_up = False
-        self.prompt = "An explorable world"
-
-        # Background generation
-        self._generation_thread: Optional[threading.Thread] = None
-        self._running = False
-        self._state_lock = threading.Lock()
-
-        # Latest frame buffer
-        self._latest_frame: Optional[Image.Image] = None
-        self._latest_frame_b64: Optional[str] = None
-        self._frame_lock = threading.Lock()
+class ResetCommand:
+    """Reset world with new seed image."""
+    seed_image: Optional[Image.Image] = None
+    seed_url: Optional[str] = None
+    prompt: str = "An explorable world"
+
+
+@dataclass
+class StopCommand:
+    """Stop the GPU session."""
+    pass
+
+
+# --- GPU Session Generator ---
+def create_gpu_game_loop(command_queue: Queue):
+    """Create GPU game loop generator with closure over command_queue."""
 
     @spaces.GPU(duration=120)
-    def initialize(
-        self,
-        model_path: str,
-        device: str = "cuda",
-        use_compile: bool = True,
-        warmup_steps: int = 1,
-    ):
-        if self.is_initialized:
-            return
-
-        print(f"Loading ModularPipeline from: {model_path}")
-        self.pipe = ModularPipeline.from_pretrained(model_path, trust_remote_code=True)
-        self.pipe.load_components(
-            device_map=device,
+    def gpu_game_loop():
+        """
+        Generator that keeps GPU allocated and processes commands.
+        Yields (frame, frame_count) tuples.
+        """
+        print("Starting GPU session...")
+
+        # Load model
+        print(f"Loading ModularPipeline from: {MODEL_ID}")
+        pipe = ModularPipeline.from_pretrained(MODEL_ID, trust_remote_code=True)
+        pipe.load_components(
+            device_map="cuda",
             torch_dtype=torch.bfloat16,
             trust_remote_code=True,
         )
 
-        self.pipe.transformer.apply_inference_patches()
-        self.pipe.transformer.quantize("fp8")
+        pipe.transformer.apply_inference_patches()
+        pipe.transformer.quantize("fp8")
+
+        # Check if we have cached kernels
+        cache_files = list(CACHE_DIR.glob("**/*.so"))
+        needs_warmup = USE_COMPILE and not cache_files
+
+        if USE_COMPILE:
+            if cache_files:
+                print(f"Found {len(cache_files)} cached kernels")
+            else:
+                print("No cached kernels - first run will compile")
 
-        if use_compile:
             print("Compiling transformer...")
-            self.pipe.transformer = torch.compile(
-                self.pipe.transformer,
+            pipe.transformer = torch.compile(
+                pipe.transformer,
                 mode="max-autotune-no-cudagraphs",
                 fullgraph=True,
                 dynamic=False,
             )
 
             print("Compiling VAE...")
-            self.pipe.vae = torch.compile(
-                self.pipe.vae,
+            pipe.vae = torch.compile(
+                pipe.vae,
                 mode="max-autotune",
                 fullgraph=True,
                 dynamic=False,
             )
 
-        self.n_frames = self.pipe.transformer.config.n_frames
-        print(f"ModularPipeline loaded! (n_frames={self.n_frames})")
+        n_frames = pipe.transformer.config.n_frames
+        print(f"Model loaded! (n_frames={n_frames})")
 
-        # Only warmup if cache is empty (first run)
-        # Subsequent cold starts can skip warmup since kernels are already compiled
-        cache_files = list(CACHE_DIR.glob("**/*.so"))
-        if use_compile and warmup_steps > 0 and not cache_files:
-            print("No cached kernels found - running warmup to compile...")
-            self._warmup_compile(warmup_steps)
-        elif cache_files:
-            print(f"Found {len(cache_files)} cached kernels - skipping warmup")
-
-        self.is_initialized = True
+        # Warmup pass if no cached kernels (triggers compilation)
+        if needs_warmup:
+            print("Running warmup pass to compile kernels...")
+            warmup_image = load_seed_frame(SEED_FRAME_URLS[0])
+            _ = pipe(
+                prompt="warmup",
+                image=warmup_image,
+                button=set(),
+                mouse=(0.0, 0.0),
+                output_type="pil",
+            )
+            print("Warmup complete! Kernels cached for future runs.")
 
-        # Generate initial state
-        print("Generating initial state...")
+        # Initialize state
         seed_image = load_seed_frame(random.choice(SEED_FRAME_URLS))
-        self.state = self.pipe(
-            prompt=self.prompt,
+        state = pipe(
+            prompt="An explorable world",
             image=seed_image,
             button=set(),
             mouse=(0.0, 0.0),
             output_type="pil",
         )
-        self.frame_count = 1
-
-        # Store initial frame
-        images = self.state.values.get("images")
-        if images is not None:
-            with self._frame_lock:
-                self._latest_frame = images
-                self._latest_frame_b64 = image_to_base64(images)
+        frame_count = 1
 
-        print("Initial state ready!")
-        self._start_generation_thread()
+        # Get initial frame
+        frame = state.values.get("images")
+        print("Initial frame generated, entering game loop...")
 
-    def _warmup_compile(self, num_steps: int = 3):
-        """Run warmup passes to trigger torch.compile before first real request."""
-        print(f"Warming up compiled model ({num_steps} steps)...")
-        self.is_warming_up = True
+        # Yield initial frame
+        yield (frame, frame_count)
 
-        try:
-            # Create a temporary state for warmup
-            seed_image = load_seed_frame(SEED_FRAME_URLS[0])
+        # Main loop - process commands and yield frames
+        while True:
+            try:
+                # Non-blocking get with short timeout for responsiveness
+                command = command_queue.get(timeout=0.01)
+            except:
+                # No command, continue idle
+                command = None
 
-            warmup_state = self.pipe(
-                prompt="warmup",
-                image=seed_image,
-                button=set(),
-                mouse=(0.0, 0.0),
-                output_type="pil",
-            )
+            if command is None:
+                continue
 
-            # Run a few more steps to ensure all kernels are compiled
-            for i in range(num_steps - 1):
-                warmup_state = self.pipe(
-                    warmup_state,
-                    prompt="warmup",
+            if isinstance(command, StopCommand):
+                print("Stop command received, ending GPU session")
+                break
+
+            elif isinstance(command, ResetCommand):
+                print("Reset command received")
+                if command.seed_image is not None:
+                    seed_img = command.seed_image.resize((640, 360), Image.BILINEAR)
+                elif command.seed_url:
+                    seed_img = load_seed_frame(command.seed_url)
+                else:
+                    seed_img = load_seed_frame(random.choice(SEED_FRAME_URLS))
+
+                state = pipe(
+                    prompt=command.prompt,
+                    image=seed_img,
                     button=set(),
                     mouse=(0.0, 0.0),
+                    output_type="pil",
+                )
+                frame_count = 1
+                frame = state.values.get("images")
+                yield (frame, frame_count)
+
+            elif isinstance(command, GenerateCommand):
+                # Generate next frame
+                state = pipe(
+                    state,
+                    prompt=command.prompt,
+                    button=command.buttons,
+                    mouse=command.mouse,
                     image=None,
                     output_type="pil",
                 )
-                print(f"  Warmup step {i + 2}/{num_steps} complete")
-
-            print("Warmup complete! Compiled kernels cached.")
-
-        except Exception as e:
-            print(f"Warmup error (non-fatal): {e}")
-
-        self.is_warming_up = False
-
-    def _start_generation_thread(self):
-        if self._generation_thread is not None and self._generation_thread.is_alive():
-            return
-
-        self._running = True
-        self._generation_thread = threading.Thread(
-            target=self._generation_loop, daemon=True
-        )
-        self._generation_thread.start()
-        print("Background generation thread started!")
-
-    def _generation_loop(self):
-        while self._running:
-            if not self.is_initialized or self.state is None or self.is_warming_up:
-                time.sleep(0.01)
-                continue
-
-            buttons, mouse = self.control_state.get()
-
-            with self._state_lock:
-                try:
-                    self.state = self.pipe(
-                        self.state,
-                        prompt=self.prompt,
-                        button=buttons,
-                        mouse=mouse,
-                        image=None,
+                frame_count += 1
+                frame = state.values.get("images")
+
+                # Auto-reset near end of context
+                if frame_count >= n_frames - 2:
+                    print(f"Auto-reset at frame {frame_count}")
+                    seed_img = load_seed_frame(random.choice(SEED_FRAME_URLS))
+                    state = pipe(
+                        prompt=command.prompt,
+                        image=seed_img,
+                        button=set(),
+                        mouse=(0.0, 0.0),
                         output_type="pil",
                     )
-                    self.frame_count += 1
-
-                    images = self.state.values.get("images")
-                    if images is not None:
-                        with self._frame_lock:
-                            self._latest_frame = images
-                            self._latest_frame_b64 = image_to_base64(images)
-
-                    if self.frame_count >= self.n_frames - 2:
-                        print(f"Auto-reset at frame {self.frame_count}")
-                        self._do_reset()
-
-                except Exception as e:
-                    import traceback
-
-                    print(f"Generation error: {e}")
-                    traceback.print_exc()
-                    time.sleep(0.5)
-
-    def _do_reset(self, seed_url: str = None, seed_image: Image.Image = None):
-        self.frame_count = 0
-
-        # Use provided image, or load from URL, or pick random
-        if seed_image is not None:
-            # Resize uploaded image to target size
-            target_size = (360, 640)  # (H, W)
-            seed_image = seed_image.resize(
-                (target_size[1], target_size[0]), Image.BILINEAR
-            )
-        else:
-            url = seed_url or random.choice(SEED_FRAME_URLS)
-            seed_image = load_seed_frame(url)
-
-        buttons, mouse = self.control_state.get()
-        self.state = self.pipe(
-            prompt=self.prompt,
-            image=seed_image,
-            button=buttons,
-            mouse=mouse,
-            output_type="pil",
-        )
-        self.frame_count = 1
-
-        images = self.state.values.get("images")
-        if images is not None:
-            with self._frame_lock:
-                self._latest_frame = images
-                self._latest_frame_b64 = image_to_base64(images)
-
-    def reset(self, seed_url: str = None, seed_image: Image.Image = None):
-        with self._state_lock:
-            self._do_reset(seed_url=seed_url, seed_image=seed_image)
-
-    def get_latest_frame(self) -> Tuple[Optional[str], int]:
-        with self._frame_lock:
-            return self._latest_frame_b64, self.frame_count
+                    frame_count = 1
+                    frame = state.values.get("images")
 
-    def update_controls(self, buttons: Set[int], mouse: Tuple[float, float]):
-        self.control_state.update(buttons, mouse)
+                yield (frame, frame_count)
 
-    def stop(self):
-        self._running = False
-        if self._generation_thread is not None:
-            self._generation_thread.join(timeout=2.0)
+        print("GPU session ended")
 
+    # Return the generator
+    return gpu_game_loop()
 
-game_state = GameState()
 
-
-# --- Control Input HTML Component ---
+# --- Gradio App ---
 CONTROL_INPUT_HTML = """
 <div id="control-input-wrapper" style="width: 100%; background: #0a0a0f; border-radius: 12px; overflow: hidden; font-family: 'JetBrains Mono', 'Fira Code', monospace;">
     <div style="padding: 16px; background: linear-gradient(180deg, rgba(22, 27, 34, 0.95) 0%, rgba(13, 17, 23, 0.98) 100%);">
-        <!-- Status Bar -->
         <div style="display: flex; justify-content: space-between; align-items: center; margin-bottom: 16px;">
             <div id="status-indicator" style="display: flex; align-items: center; gap: 6px; background: rgba(0,0,0,0.4); padding: 6px 12px; border-radius: 20px; border: 1px solid rgba(88, 166, 255, 0.2);">
                 <div id="status-dot" style="width: 8px; height: 8px; border-radius: 50%; background: #ff6b6b; box-shadow: 0 0 8px #ff6b6b;"></div>
                 <span id="status-text" style="font-size: 11px; color: #8b949e; text-transform: uppercase; letter-spacing: 1px;">Not capturing</span>
             </div>
         </div>
-
         <div style="display: grid; grid-template-columns: 1fr 1fr; gap: 16px;">
-            <!-- Keyboard State (touch-enabled) -->
             <div style="background: rgba(0,0,0,0.3); border-radius: 12px; padding: 12px; border: 1px solid rgba(88, 166, 255, 0.1);">
                 <div style="font-size: 10px; color: #8b949e; text-transform: uppercase; letter-spacing: 1px; margin-bottom: 10px;">Movement</div>
                 <div style="display: grid; grid-template-columns: repeat(3, 1fr); gap: 4px; max-width: 120px; margin: 0 auto;">
                     <div></div>
-                    <div id="key-w" data-key="KeyW" style="aspect-ratio: 1; background: rgba(88, 166, 255, 0.1); border: 1px solid rgba(88, 166, 255, 0.3); border-radius: 6px; display: flex; align-items: center; justify-content: center; font-size: 11px; color: #58a6ff; transition: all 0.1s; cursor: pointer; user-select: none; -webkit-user-select: none; touch-action: manipulation;">W</div>
+                    <div id="key-w" data-key="KeyW" style="aspect-ratio: 1; background: rgba(88, 166, 255, 0.1); border: 1px solid rgba(88, 166, 255, 0.3); border-radius: 6px; display: flex; align-items: center; justify-content: center; font-size: 11px; color: #58a6ff; transition: all 0.1s; cursor: pointer; user-select: none;">W</div>
                     <div></div>
-                    <div id="key-a" data-key="KeyA" style="aspect-ratio: 1; background: rgba(88, 166, 255, 0.1); border: 1px solid rgba(88, 166, 255, 0.3); border-radius: 6px; display: flex; align-items: center; justify-content: center; font-size: 11px; color: #58a6ff; transition: all 0.1s; cursor: pointer; user-select: none; -webkit-user-select: none; touch-action: manipulation;">A</div>
-                    <div id="key-s" data-key="KeyS" style="aspect-ratio: 1; background: rgba(88, 166, 255, 0.1); border: 1px solid rgba(88, 166, 255, 0.3); border-radius: 6px; display: flex; align-items: center; justify-content: center; font-size: 11px; color: #58a6ff; transition: all 0.1s; cursor: pointer; user-select: none; -webkit-user-select: none; touch-action: manipulation;">S</div>
-                    <div id="key-d" data-key="KeyD" style="aspect-ratio: 1; background: rgba(88, 166, 255, 0.1); border: 1px solid rgba(88, 166, 255, 0.3); border-radius: 6px; display: flex; align-items: center; justify-content: center; font-size: 11px; color: #58a6ff; transition: all 0.1s; cursor: pointer; user-select: none; -webkit-user-select: none; touch-action: manipulation;">D</div>
+                    <div id="key-a" data-key="KeyA" style="aspect-ratio: 1; background: rgba(88, 166, 255, 0.1); border: 1px solid rgba(88, 166, 255, 0.3); border-radius: 6px; display: flex; align-items: center; justify-content: center; font-size: 11px; color: #58a6ff; transition: all 0.1s; cursor: pointer; user-select: none;">A</div>
+                    <div id="key-s" data-key="KeyS" style="aspect-ratio: 1; background: rgba(88, 166, 255, 0.1); border: 1px solid rgba(88, 166, 255, 0.3); border-radius: 6px; display: flex; align-items: center; justify-content: center; font-size: 11px; color: #58a6ff; transition: all 0.1s; cursor: pointer; user-select: none;">S</div>
+                    <div id="key-d" data-key="KeyD" style="aspect-ratio: 1; background: rgba(88, 166, 255, 0.1); border: 1px solid rgba(88, 166, 255, 0.3); border-radius: 6px; display: flex; align-items: center; justify-content: center; font-size: 11px; color: #58a6ff; transition: all 0.1s; cursor: pointer; user-select: none;">D</div>
                 </div>
                 <div style="display: flex; gap: 4px; margin-top: 8px; justify-content: center;">
-                    <div id="key-shift" data-key="ShiftLeft" style="padding: 4px 8px; background: rgba(88, 166, 255, 0.1); border: 1px solid rgba(88, 166, 255, 0.3); border-radius: 4px; font-size: 9px; color: #58a6ff; transition: all 0.1s; cursor: pointer; user-select: none; -webkit-user-select: none; touch-action: manipulation;">SHIFT</div>
-                    <div id="key-space" data-key="Space" style="padding: 4px 16px; background: rgba(88, 166, 255, 0.1); border: 1px solid rgba(88, 166, 255, 0.3); border-radius: 4px; font-size: 9px; color: #58a6ff; transition: all 0.1s; cursor: pointer; user-select: none; -webkit-user-select: none; touch-action: manipulation;">SPACE</div>
-                    <div id="key-e" data-key="KeyE" style="padding: 4px 8px; background: rgba(88, 166, 255, 0.1); border: 1px solid rgba(88, 166, 255, 0.3); border-radius: 4px; font-size: 9px; color: #58a6ff; transition: all 0.1s; cursor: pointer; user-select: none; -webkit-user-select: none; touch-action: manipulation;">E</div>
+                    <div id="key-shift" data-key="ShiftLeft" style="padding: 4px 8px; background: rgba(88, 166, 255, 0.1); border: 1px solid rgba(88, 166, 255, 0.3); border-radius: 4px; font-size: 9px; color: #58a6ff; cursor: pointer; user-select: none;">SHIFT</div>
+                    <div id="key-space" data-key="Space" style="padding: 4px 16px; background: rgba(88, 166, 255, 0.1); border: 1px solid rgba(88, 166, 255, 0.3); border-radius: 4px; font-size: 9px; color: #58a6ff; cursor: pointer; user-select: none;">SPACE</div>
+                    <div id="key-e" data-key="KeyE" style="padding: 4px 8px; background: rgba(88, 166, 255, 0.1); border: 1px solid rgba(88, 166, 255, 0.3); border-radius: 4px; font-size: 9px; color: #58a6ff; cursor: pointer; user-select: none;">E</div>
                 </div>
             </div>
-
-            <!-- Mouse/Look Joystick (touch-enabled) -->
             <div style="background: rgba(0,0,0,0.3); border-radius: 12px; padding: 12px; border: 1px solid rgba(88, 166, 255, 0.1);">
                 <div style="font-size: 10px; color: #8b949e; text-transform: uppercase; letter-spacing: 1px; margin-bottom: 10px;">Look (drag)</div>
                 <div style="display: flex; align-items: center; justify-content: center; gap: 16px;">
                     <div id="mouse-joystick" style="width: 100px; height: 100px; background: rgba(88, 166, 255, 0.05); border: 2px solid rgba(88, 166, 255, 0.3); border-radius: 50%; position: relative; cursor: pointer; touch-action: none;">
-                        <div id="mouse-dot" style="width: 24px; height: 24px; background: #58a6ff; border-radius: 50%; position: absolute; top: 50%; left: 50%; transform: translate(-50%, -50%); box-shadow: 0 0 12px rgba(88, 166, 255, 0.6); transition: all 0.05s ease-out; pointer-events: none;"></div>
-                        <div style="position: absolute; top: 50%; left: 0; right: 0; height: 1px; background: rgba(88, 166, 255, 0.2); pointer-events: none;"></div>
-                        <div style="position: absolute; left: 50%; top: 0; bottom: 0; width: 1px; background: rgba(88, 166, 255, 0.2); pointer-events: none;"></div>
+                        <div id="mouse-dot" style="width: 24px; height: 24px; background: #58a6ff; border-radius: 50%; position: absolute; top: 50%; left: 50%; transform: translate(-50%, -50%); box-shadow: 0 0 12px rgba(88, 166, 255, 0.6); pointer-events: none;"></div>
                     </div>
                     <div style="text-align: left;">
                         <div style="font-size: 10px; color: #8b949e; margin-bottom: 4px;">Velocity</div>
@@ -380,7 +293,6 @@ CONTROL_INPUT_HTML = """
                 </div>
             </div>
         </div>
-
         <div style="margin-top: 12px; background: rgba(0,0,0,0.3); border-radius: 8px; padding: 8px 12px; border: 1px solid rgba(88, 166, 255, 0.1);">
             <div style="display: flex; align-items: center; gap: 8px;">
                 <span style="font-size: 10px; color: #8b949e; text-transform: uppercase; letter-spacing: 1px;">Active:</span>
@@ -404,26 +316,14 @@ CONTROL_INPUT_JS = """
     const mouseJoystick = element.querySelector('#mouse-joystick');
 
     let isCapturing = false;
-    let isTouchActive = false;
     let pressedKeys = new Set();
-    let touchPressedKeys = new Set();
     let mouseVelocity = { x: 0, y: 0 };
     let lastMouseMove = Date.now();
 
-    // Windows Virtual Key codes
     const BUTTON_MAP = {
         'KeyW': 87, 'KeyA': 65, 'KeyS': 83, 'KeyD': 68,
         'KeyQ': 81, 'KeyE': 69, 'KeyR': 82, 'KeyF': 70,
-        'KeyT': 84, 'KeyG': 71, 'KeyZ': 90, 'KeyX': 88,
-        'KeyC': 67, 'KeyV': 86, 'KeyB': 66, 'KeyN': 78,
-        'KeyM': 77, 'KeyH': 72, 'KeyJ': 74, 'KeyK': 75,
-        'KeyL': 76, 'KeyI': 73, 'KeyO': 79, 'KeyP': 80,
-        'KeyU': 85, 'KeyY': 89,
-        'Digit1': 49, 'Digit2': 50, 'Digit3': 51, 'Digit4': 52,
-        'Digit5': 53, 'Digit6': 54, 'Digit7': 55, 'Digit8': 56,
-        'Digit9': 57, 'Digit0': 48,
-        'Space': 32,
-        'ShiftLeft': 16, 'ShiftRight': 16,
+        'Space': 32, 'ShiftLeft': 16, 'ShiftRight': 16,
     };
 
     const KEY_DISPLAY_MAP = {
@@ -431,8 +331,6 @@ CONTROL_INPUT_JS = """
         'ShiftLeft': 'key-shift', 'Space': 'key-space', 'KeyE': 'key-e',
     };
 
-    const isTouchDevice = 'ontouchstart' in window || navigator.maxTouchPoints > 0;
-
     function updateKeyDisplay(code, pressed) {
         const elementId = KEY_DISPLAY_MAP[code];
         if (elementId) {
@@ -440,7 +338,6 @@ CONTROL_INPUT_JS = """
             if (keyEl) {
                 keyEl.style.background = pressed ? 'rgba(88, 166, 255, 0.4)' : 'rgba(88, 166, 255, 0.1)';
                 keyEl.style.borderColor = pressed ? '#58a6ff' : 'rgba(88, 166, 255, 0.3)';
-                keyEl.style.boxShadow = pressed ? '0 0 8px rgba(88, 166, 255, 0.4)' : 'none';
             }
         }
     }
@@ -448,9 +345,8 @@ CONTROL_INPUT_JS = """
     function updateMouseDisplay() {
         const displayX = Math.max(-1, Math.min(1, mouseVelocity.x / 20));
         const displayY = Math.max(-1, Math.min(1, mouseVelocity.y / 20));
-        const maxOffset = 40;
-        mouseDot.style.left = (50 + displayX * maxOffset) + '%';
-        mouseDot.style.top = (50 + displayY * maxOffset) + '%';
+        mouseDot.style.left = (50 + displayX * 40) + '%';
+        mouseDot.style.top = (50 + displayY * 40) + '%';
         mouseXValue.textContent = mouseVelocity.x.toFixed(1);
         mouseYValue.textContent = mouseVelocity.y.toFixed(1);
     }
@@ -460,18 +356,12 @@ CONTROL_INPUT_JS = """
             activeButtonsDisplay.innerHTML = '<span style="font-size: 11px; color: #484f58;">None</span>';
         } else {
             activeButtonsDisplay.innerHTML = Array.from(pressedKeys)
-                .map(code => code.replace('Key', '').replace('Digit', '').replace('Left', ''))
+                .map(code => code.replace('Key', '').replace('Left', ''))
                 .map(name => `<span style="font-size: 10px; background: rgba(88, 166, 255, 0.2); color: #58a6ff; padding: 2px 6px; border-radius: 4px;">${name}</span>`)
                 .join('');
         }
     }
 
-    function updateStatus(capturing) {
-        statusDot.style.background = capturing ? '#3fb950' : '#ff6b6b';
-        statusDot.style.boxShadow = capturing ? '0 0 8px #3fb950' : '0 0 8px #ff6b6b';
-        statusText.textContent = capturing ? 'Capturing - ESC to release' : 'Not capturing';
-    }
-
     function triggerUpdate() {
         const buttonIds = Array.from(pressedKeys)
             .filter(code => BUTTON_MAP[code] !== undefined)
@@ -482,7 +372,8 @@ CONTROL_INPUT_JS = """
 
     document.addEventListener('pointerlockchange', () => {
         isCapturing = document.pointerLockElement !== null;
-        updateStatus(isCapturing);
+        statusDot.style.background = isCapturing ? '#3fb950' : '#ff6b6b';
+        statusText.textContent = isCapturing ? 'Capturing - ESC to release' : 'Not capturing';
         if (!isCapturing) {
             pressedKeys.clear();
             mouseVelocity = { x: 0, y: 0 };
@@ -515,19 +406,17 @@ CONTROL_INPUT_JS = """
         }
     });
 
-    const MOUSE_SENSITIVITY = 1.5;
-
     document.addEventListener('mousemove', (e) => {
         if (!isCapturing) return;
-        mouseVelocity.x = e.movementX * MOUSE_SENSITIVITY;
-        mouseVelocity.y = e.movementY * MOUSE_SENSITIVITY;
+        mouseVelocity.x = e.movementX * 1.5;
+        mouseVelocity.y = e.movementY * 1.5;
         updateMouseDisplay();
         triggerUpdate();
         lastMouseMove = Date.now();
     });
 
     setInterval(() => {
-        if ((isCapturing || isTouchActive) && Date.now() - lastMouseMove > 50) {
+        if (isCapturing && Date.now() - lastMouseMove > 50) {
             mouseVelocity.x *= 0.8;
             mouseVelocity.y *= 0.8;
             if (Math.abs(mouseVelocity.x) < 0.01) mouseVelocity.x = 0;
@@ -537,203 +426,37 @@ CONTROL_INPUT_JS = """
         }
     }, 100);
 
-    // Touch controls
-    const touchableKeys = element.querySelectorAll('[data-key]');
-
-    touchableKeys.forEach(keyEl => {
-        const keyCode = keyEl.dataset.key;
-
-        const handleTouchStart = (e) => {
-            e.preventDefault();
-            isTouchActive = true;
-            if (!pressedKeys.has(keyCode)) {
-                pressedKeys.add(keyCode);
-                touchPressedKeys.add(keyCode);
-                updateKeyDisplay(keyCode, true);
-                updateActiveButtonsDisplay();
-                triggerUpdate();
-            }
-        };
-
-        const handleTouchEnd = (e) => {
-            e.preventDefault();
-            if (touchPressedKeys.has(keyCode)) {
-                pressedKeys.delete(keyCode);
-                touchPressedKeys.delete(keyCode);
-                updateKeyDisplay(keyCode, false);
-                updateActiveButtonsDisplay();
-                triggerUpdate();
-            }
-        };
-
-        keyEl.addEventListener('touchstart', handleTouchStart, { passive: false });
-        keyEl.addEventListener('touchend', handleTouchEnd, { passive: false });
-        keyEl.addEventListener('touchcancel', handleTouchEnd, { passive: false });
-
-        keyEl.addEventListener('mousedown', (e) => {
-            e.preventDefault();
-            isTouchActive = true;
-            if (!pressedKeys.has(keyCode)) {
-                pressedKeys.add(keyCode);
-                touchPressedKeys.add(keyCode);
-                updateKeyDisplay(keyCode, true);
-                updateActiveButtonsDisplay();
-                triggerUpdate();
-            }
-        });
-
-        keyEl.addEventListener('mouseup', (e) => {
-            if (touchPressedKeys.has(keyCode)) {
-                pressedKeys.delete(keyCode);
-                touchPressedKeys.delete(keyCode);
-                updateKeyDisplay(keyCode, false);
-                updateActiveButtonsDisplay();
-                triggerUpdate();
-            }
-        });
-
-        keyEl.addEventListener('mouseleave', (e) => {
-            if (touchPressedKeys.has(keyCode)) {
-                pressedKeys.delete(keyCode);
-                touchPressedKeys.delete(keyCode);
-                updateKeyDisplay(keyCode, false);
-                updateActiveButtonsDisplay();
-                triggerUpdate();
-            }
-        });
-    });
-
-    // Touch joystick
+    // Touch/click controls for joystick
     let joystickActive = false;
-    const JOYSTICK_SENSITIVITY = 0.5;
-    const JOYSTICK_RADIUS = 50;
-
     const getJoystickCenter = () => {
         const rect = mouseJoystick.getBoundingClientRect();
         return { x: rect.left + rect.width / 2, y: rect.top + rect.height / 2 };
     };
-
     const handleJoystickMove = (clientX, clientY) => {
         const center = getJoystickCenter();
-        let dx = clientX - center.x;
-        let dy = clientY - center.y;
-        const distance = Math.sqrt(dx * dx + dy * dy);
-        if (distance > JOYSTICK_RADIUS) {
-            dx = (dx / distance) * JOYSTICK_RADIUS;
-            dy = (dy / distance) * JOYSTICK_RADIUS;
-        }
-        mouseVelocity.x = (dx / JOYSTICK_RADIUS) * 20 * JOYSTICK_SENSITIVITY;
-        mouseVelocity.y = (dy / JOYSTICK_RADIUS) * 20 * JOYSTICK_SENSITIVITY;
+        let dx = clientX - center.x, dy = clientY - center.y;
+        const dist = Math.sqrt(dx*dx + dy*dy);
+        if (dist > 50) { dx = dx/dist*50; dy = dy/dist*50; }
+        mouseVelocity.x = (dx/50) * 10;
+        mouseVelocity.y = (dy/50) * 10;
         updateMouseDisplay();
         triggerUpdate();
-        lastMouseMove = Date.now();
     };
+    mouseJoystick.addEventListener('mousedown', (e) => { joystickActive = true; handleJoystickMove(e.clientX, e.clientY); });
+    document.addEventListener('mousemove', (e) => { if (joystickActive) handleJoystickMove(e.clientX, e.clientY); });
+    document.addEventListener('mouseup', () => { if (joystickActive) { joystickActive = false; mouseVelocity = {x:0,y:0}; updateMouseDisplay(); triggerUpdate(); }});
 
-    mouseJoystick.addEventListener('touchstart', (e) => {
-        e.preventDefault();
-        joystickActive = true;
-        isTouchActive = true;
-        handleJoystickMove(e.touches[0].clientX, e.touches[0].clientY);
-    }, { passive: false });
-
-    mouseJoystick.addEventListener('touchmove', (e) => {
-        e.preventDefault();
-        if (joystickActive) handleJoystickMove(e.touches[0].clientX, e.touches[0].clientY);
-    }, { passive: false });
-
-    mouseJoystick.addEventListener('touchend', (e) => {
-        e.preventDefault();
-        joystickActive = false;
-        mouseVelocity = { x: 0, y: 0 };
-        updateMouseDisplay();
-        triggerUpdate();
-    }, { passive: false });
-
-    mouseJoystick.addEventListener('touchcancel', (e) => {
-        joystickActive = false;
-        mouseVelocity = { x: 0, y: 0 };
-        updateMouseDisplay();
-        triggerUpdate();
-    }, { passive: false });
-
-    mouseJoystick.addEventListener('mousedown', (e) => {
-        e.preventDefault();
-        joystickActive = true;
-        isTouchActive = true;
-        handleJoystickMove(e.clientX, e.clientY);
-    });
-
-    document.addEventListener('mousemove', (e) => {
-        if (joystickActive) handleJoystickMove(e.clientX, e.clientY);
-    });
-
-    document.addEventListener('mouseup', () => {
-        if (joystickActive) {
-            joystickActive = false;
-            mouseVelocity = { x: 0, y: 0 };
-            updateMouseDisplay();
-            triggerUpdate();
-        }
+    // Touch controls for keys
+    element.querySelectorAll('[data-key]').forEach(keyEl => {
+        const keyCode = keyEl.dataset.key;
+        keyEl.addEventListener('touchstart', (e) => { e.preventDefault(); pressedKeys.add(keyCode); updateKeyDisplay(keyCode, true); updateActiveButtonsDisplay(); triggerUpdate(); }, {passive:false});
+        keyEl.addEventListener('touchend', (e) => { e.preventDefault(); pressedKeys.delete(keyCode); updateKeyDisplay(keyCode, false); updateActiveButtonsDisplay(); triggerUpdate(); }, {passive:false});
     });
 
-    if (isTouchDevice) {
-        statusText.textContent = 'Touch controls ready';
-        statusDot.style.background = '#3fb950';
-        statusDot.style.boxShadow = '0 0 8px #3fb950';
-    } else {
-        updateStatus(false);
-    }
-
     updateMouseDisplay();
 })();
 """
 
-CAPTURE_JS = """
-() => {
-    const isTouchDevice = 'ontouchstart' in window || navigator.maxTouchPoints > 0;
-    if (isTouchDevice) return;
-
-    const insertButton = () => {
-        const controlPanel = document.querySelector('#control-input-wrapper');
-        if (!controlPanel) { setTimeout(insertButton, 100); return; }
-
-        const btn = document.createElement('button');
-        btn.id = 'capture-btn';
-        btn.textContent = '🎮 Click to Capture Controls';
-        btn.style.cssText = `
-            width: 100%; padding: 12px 24px; margin-bottom: 12px;
-            font-size: 14px; font-weight: bold;
-            background: linear-gradient(135deg, #58a6ff 0%, #a371f7 100%);
-            color: white; border: none; border-radius: 8px; cursor: pointer;
-            box-shadow: 0 4px 12px rgba(0,0,0,0.3); transition: all 0.2s;
-        `;
-        btn.onmouseenter = () => btn.style.transform = 'scale(1.02)';
-        btn.onmouseleave = () => btn.style.transform = 'scale(1)';
-        btn.onclick = async () => {
-            try { await document.body.requestPointerLock(); }
-            catch (e) { console.log('Pointer lock failed:', e); }
-        };
-        controlPanel.parentNode.insertBefore(btn, controlPanel);
-    };
-
-    insertButton();
-
-    document.addEventListener('pointerlockchange', () => {
-        const btn = document.querySelector('#capture-btn');
-        if (!btn) return;
-        if (document.pointerLockElement) {
-            btn.textContent = '🔒 Capturing (ESC to release)';
-            btn.style.background = 'linear-gradient(135deg, #3fb950 0%, #2ea043 100%)';
-        } else {
-            btn.textContent = '🎮 Click to Capture Controls';
-            btn.style.background = 'linear-gradient(135deg, #58a6ff 0%, #a371f7 100%)';
-        }
-    });
-}
-"""
-
-
-# --- Gradio App ---
 css = """
 #col-container { max-width: 1200px; margin: 0 auto; }
 #video-output { aspect-ratio: 16/9; max-width: 640px; }
@@ -743,12 +466,23 @@ css = """
 
 def create_app():
     with gr.Blocks(css=css, theme=gr.themes.Soft(), title="WorldEngine") as demo:
+        # State: (generator, command_queue) or empty tuple
+        session_state = gr.State(())
+
+        # Current controls (updated by JS)
+        current_controls = gr.State({"buttons": [], "mouse_x": 0.0, "mouse_y": 0.0})
+        current_prompt = gr.State("An explorable world")
+
+        # Latest frame for display
+        latest_frame = gr.State(None)
+        latest_frame_count = gr.State(0)
+
         gr.Markdown("""
         # 🌍 WorldEngine — Real-Time World Model
 
-        Interactive frame-by-frame world generation.
+        Interactive frame-by-frame world generation on ZeroGPU.
 
-        **Controls:** WASD to move • Mouse to look • Space to jump • Shift to sprint • E to interact • ESC to release
+        **Controls:** Click "Start GPU" → Click game view to capture → WASD to move • Mouse to look • ESC to release
         """)
 
         with gr.Row():
@@ -761,7 +495,10 @@ def create_app():
                     height=360,
                     show_label=False,
                 )
-                frame_display = gr.Number(label="Frame", value=0, interactive=False)
+                with gr.Row():
+                    frame_display = gr.Number(label="Frame", value=0, interactive=False)
+                    start_btn = gr.Button("🎮 Start Game", variant="primary")
+                    stop_btn = gr.Button("⏹ End Game", interactive=False)
 
             with gr.Column(scale=1):
                 control_input = gr.HTML(
@@ -777,83 +514,166 @@ def create_app():
                 )
 
                 with gr.Accordion("Seed Selection", open=False):
-                    seed_image_upload = gr.Image(
-                        label="Upload Seed Image (optional)",
-                        type="pil",
-                        sources=["upload"],
-                        height=150,
-                    )
-                    gr.Markdown("*Or select a preset seed:*", elem_classes=["text-sm"])
                     seed_dropdown = gr.Dropdown(
-                        choices=["Random"]
-                        + [f"Seed {i + 1}" for i in range(len(SEED_FRAME_URLS))],
+                        choices=["Random"] + [f"Seed {i+1}" for i in range(len(SEED_FRAME_URLS))],
                         value="Random",
                         label="Preset Seeds",
                     )
-                    reset_btn = gr.Button("Reset World", variant="primary")
-
-        # Event handlers
-        def on_controller_change(controller_value):
-            if not controller_value or not game_state.is_initialized:
-                return
-            buttons = set(controller_value.get("buttons", []))
-            mouse_x = controller_value.get("mouse_x", 0.0)
-            mouse_y = controller_value.get("mouse_y", 0.0)
-            game_state.update_controls(buttons, (mouse_x, mouse_y))
-
-        def on_prompt_change(prompt):
-            game_state.prompt = prompt
-
-        def on_reset(uploaded_image, seed_choice):
-            if uploaded_image is not None:
-                game_state.reset(seed_image=uploaded_image)
-            elif seed_choice and seed_choice != "Random":
+                    reset_btn = gr.Button("Reset World", variant="secondary")
+
+        # --- Event Handlers ---
+
+        def on_start():
+            """Start GPU session - returns generator and queue."""
+            command_queue = Queue()
+            gen = create_gpu_game_loop(command_queue)
+
+            # Get initial frame
+            frame, frame_count = next(gen)
+
+            return (
+                (gen, command_queue),  # session_state
+                frame,                  # latest_frame
+                frame_count,            # latest_frame_count
+                frame,                  # video_output
+                frame_count,            # frame_display
+                gr.update(interactive=False),  # start_btn
+                gr.update(interactive=True),   # stop_btn
+            )
+
+        def on_stop(state):
+            """Stop GPU session."""
+            if len(state) == 0:
+                return ((), None, 0, None, 0,
+                        gr.update(interactive=True), gr.update(interactive=False))
+
+            gen, command_queue = state
+            command_queue.put(StopCommand())
+
+            # Drain generator
+            try:
+                while True:
+                    next(gen)
+            except StopIteration:
+                pass
+
+            return (
+                (),
+                None,
+                0,
+                None,
+                0,
+                gr.update(interactive=True),
+                gr.update(interactive=False),
+            )
+
+        def on_generate_tick(state, controls, prompt, current_frame, current_count):
+            """Called by timer - send generate command and get next frame."""
+            if len(state) == 0:
+                return current_frame, current_count, current_frame, current_count
+
+            gen, command_queue = state
+
+            # Send generate command
+            buttons = set(controls.get("buttons", []))
+            mouse = (controls.get("mouse_x", 0.0), controls.get("mouse_y", 0.0))
+            command_queue.put(GenerateCommand(buttons=buttons, mouse=mouse, prompt=prompt))
+
+            # Get next frame
+            try:
+                frame, frame_count = next(gen)
+                return frame, frame_count, frame, frame_count
+            except StopIteration:
+                return current_frame, current_count, current_frame, current_count
+
+        def on_reset(state, seed_choice, prompt):
+            """Reset world with new seed."""
+            if len(state) == 0:
+                return None, 0, None, 0
+
+            gen, command_queue = state
+
+            seed_url = None
+            if seed_choice and seed_choice != "Random":
                 idx = int(seed_choice.split()[-1]) - 1
-                game_state.reset(seed_url=SEED_FRAME_URLS[idx])
-            else:
-                game_state.reset()
+                seed_url = SEED_FRAME_URLS[idx]
+
+            command_queue.put(ResetCommand(seed_url=seed_url, prompt=prompt))
+
+            try:
+                frame, frame_count = next(gen)
+                return frame, frame_count, frame, frame_count
+            except StopIteration:
+                return None, 0, None, 0
+
+        def on_controller_change(value):
+            """Update current controls state."""
+            return value or {"buttons": [], "mouse_x": 0.0, "mouse_y": 0.0}
+
+        def on_prompt_change(value):
+            """Update current prompt state."""
+            return value
+
+        # Wire up events
+        start_btn.click(
+            fn=on_start,
+            outputs=[session_state, latest_frame, latest_frame_count,
+                    video_output, frame_display, start_btn, stop_btn],
+        )
+
+        stop_btn.click(
+            fn=on_stop,
+            inputs=[session_state],
+            outputs=[session_state, latest_frame, latest_frame_count,
+                    video_output, frame_display, start_btn, stop_btn],
+        )
 
-        def get_frame():
-            with game_state._frame_lock:
-                return game_state._latest_frame, game_state.frame_count
+        reset_btn.click(
+            fn=on_reset,
+            inputs=[session_state, seed_dropdown, current_prompt],
+            outputs=[latest_frame, latest_frame_count, video_output, frame_display],
+        )
 
-        control_input.change(fn=on_controller_change, inputs=[control_input])
-        prompt_input.change(fn=on_prompt_change, inputs=[prompt_input])
-        reset_btn.click(fn=on_reset, inputs=[seed_image_upload, seed_dropdown])
+        control_input.change(fn=on_controller_change, inputs=[control_input], outputs=[current_controls])
+        prompt_input.change(fn=on_prompt_change, inputs=[prompt_input], outputs=[current_prompt])
 
-        timer = gr.Timer(value=1 / 30)
-        timer.tick(fn=get_frame, outputs=[video_output, frame_display])
+        # Timer for continuous generation
+        timer = gr.Timer(value=1/15)  # 15 FPS target
+        timer.tick(
+            fn=on_generate_tick,
+            inputs=[session_state, current_controls, current_prompt, latest_frame, latest_frame_count],
+            outputs=[latest_frame, latest_frame_count, video_output, frame_display],
+        )
 
-        demo.load(fn=None, js=CAPTURE_JS)
+        # Pointer lock JS
+        demo.load(fn=None, js="""
+        () => {
+            const insertButton = () => {
+                const output = document.querySelector('#video-output');
+                if (!output) { setTimeout(insertButton, 100); return; }
+                output.style.cursor = 'pointer';
+                output.onclick = () => document.body.requestPointerLock();
+            };
+            insertButton();
+        }
+        """)
 
     return demo
 
 
-def main():
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    use_compile = os.environ.get("COMPILE", "1") == "1"
-    warmup_steps = int(os.environ.get("WARMUP_STEPS", "1"))
-
-    print(f"Model path: {MODEL_ID}")
-    print(f"Device: {device}")
-    print(f"Compile: {use_compile}")
-    print(f"Warmup steps: {warmup_steps}")
-    print(f"Cache dir: {CACHE_DIR}")
+# Avoid ZeroGPU "no GPU function" error
+if IS_ZERO_GPU:
+    spaces.GPU(lambda: None)
 
-    # Check if cache exists from previous run
-    cache_files = list(CACHE_DIR.glob("**/*.json")) + list(CACHE_DIR.glob("**/*.py"))
-    if cache_files:
-        print(f"Found {len(cache_files)} cached compilation files")
 
-    game_state.initialize(
-        model_path=MODEL_ID,
-        device=device,
-        use_compile=use_compile,
-        warmup_steps=warmup_steps,
-    )
+def main():
+    print(f"Model: {MODEL_ID}")
+    print(f"Compile: {USE_COMPILE}")
+    print(f"Cache dir: {CACHE_DIR}")
+    print(f"ZeroGPU: {IS_ZERO_GPU}")
 
     demo = create_app()
-    demo.launch(server_name="0.0.0.0", server_port=7860, share=False)
+    demo.launch(server_name="0.0.0.0", server_port=7860)
 
 
 if __name__ == "__main__":