update

app.py

@@ -1,7 +1,860 @@
+"""
+WorldEngine Real-Time World Model Demo - ZeroGPU Edition
+
+A Gradio demo optimized for HuggingFace ZeroGPU with:
+- Persistent compilation cache for faster cold starts
+- Warmup pass to pre-compile kernels
+- @spaces.GPU decorator support
+"""
+
+import base64
+import os
+import random
+import threading
+import time
+from dataclasses import dataclass, field
+from io import BytesIO
+from pathlib import Path
+from typing import Optional, Set, Tuple
+
 import gradio as gr
+import numpy as np
+import torch
+from PIL import Image
+
+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}")
+
+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
+
+# 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",
+    "https://gist.github.com/user-attachments/assets/ae398747-de4c-4d43-bac4-54fe61ab0ca8",
+    "https://gist.github.com/user-attachments/assets/9d7336fa-5cec-4c7d-bb65-eaebac0a6336",
+    "https://gist.github.com/user-attachments/assets/55dae2d3-00e3-4d03-bb6c-2c7e0ac70f5f",
+]
+
+
+def load_seed_frame(url: str, target_size: Tuple[int, int] = (360, 640)) -> Image.Image:
+    """Load and resize seed frame to target size."""
+    img = load_image(url)
+    img = img.resize((target_size[1], target_size[0]), Image.BILINEAR)
+    return img
+
+
+def image_to_base64(image: Image.Image) -> str:
+    """Convert PIL image to base64 data URL."""
+    buffered = BytesIO()
+    image.save(buffered, format="PNG")
+    img_str = base64.b64encode(buffered.getvalue()).decode()
+    return f"data:image/png;base64,{img_str}"
+
+
+# --- Control State ---
+@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()
+
+    @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,
+            torch_dtype=torch.bfloat16,
+            trust_remote_code=True,
+        )
+
+        self.pipe.transformer.apply_inference_patches()
+        self.pipe.transformer.quantize("fp8")
+
+        if use_compile:
+            print("Compiling transformer...")
+            self.pipe.transformer = torch.compile(
+                self.pipe.transformer,
+                mode="max-autotune-no-cudagraphs",
+                fullgraph=True,
+                dynamic=False,
+            )
+
+            print("Compiling VAE...")
+            self.pipe.vae = torch.compile(
+                self.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})")
+
+        # 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
+
+        # Generate initial state
+        print("Generating initial state...")
+        seed_image = load_seed_frame(random.choice(SEED_FRAME_URLS))
+        self.state = self.pipe(
+            prompt=self.prompt,
+            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)
+
+        print("Initial state ready!")
+        self._start_generation_thread()
+
+    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
+
+        try:
+            # Create a temporary state for warmup
+            seed_image = load_seed_frame(SEED_FRAME_URLS[0])
+
+            warmup_state = self.pipe(
+                prompt="warmup",
+                image=seed_image,
+                button=set(),
+                mouse=(0.0, 0.0),
+                output_type="pil",
+            )
+
+            # 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",
+                    button=set(),
+                    mouse=(0.0, 0.0),
+                    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,
+                        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
+
+    def update_controls(self, buttons: Set[int], mouse: Tuple[float, float]):
+        self.control_state.update(buttons, mouse)
+
+    def stop(self):
+        self._running = False
+        if self._generation_thread is not None:
+            self._generation_thread.join(timeout=2.0)
+
+
+game_state = GameState()
+
+
+# --- Control Input HTML Component ---
+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></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>
+                <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>
+            </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>
+                    <div style="text-align: left;">
+                        <div style="font-size: 10px; color: #8b949e; margin-bottom: 4px;">Velocity</div>
+                        <div style="font-size: 12px; color: #58a6ff;">X: <span id="mouse-x-value">0.0</span></div>
+                        <div style="font-size: 12px; color: #58a6ff;">Y: <span id="mouse-y-value">0.0</span></div>
+                    </div>
+                </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>
+                <div id="active-buttons" style="display: flex; gap: 4px; flex-wrap: wrap; min-height: 20px;">
+                    <span style="font-size: 11px; color: #484f58;">None</span>
+                </div>
+            </div>
+        </div>
+    </div>
+</div>
+"""
+
+CONTROL_INPUT_JS = """
+(() => {
+    const statusDot = element.querySelector('#status-dot');
+    const statusText = element.querySelector('#status-text');
+    const mouseDot = element.querySelector('#mouse-dot');
+    const mouseXValue = element.querySelector('#mouse-x-value');
+    const mouseYValue = element.querySelector('#mouse-y-value');
+    const activeButtonsDisplay = element.querySelector('#active-buttons');
+    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,
+    };
+
+    const KEY_DISPLAY_MAP = {
+        'KeyW': 'key-w', 'KeyA': 'key-a', 'KeyS': 'key-s', 'KeyD': 'key-d',
+        '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) {
+            const keyEl = element.querySelector('#' + elementId);
+            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';
+            }
+        }
+    }
+
+    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) + '%';
+        mouseXValue.textContent = mouseVelocity.x.toFixed(1);
+        mouseYValue.textContent = mouseVelocity.y.toFixed(1);
+    }
+
+    function updateActiveButtonsDisplay() {
+        if (pressedKeys.size === 0) {
+            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(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)
+            .map(code => BUTTON_MAP[code]);
+        props.value = { buttons: buttonIds, mouse_x: mouseVelocity.x, mouse_y: mouseVelocity.y };
+        trigger('change', props.value);
+    }
+
+    document.addEventListener('pointerlockchange', () => {
+        isCapturing = document.pointerLockElement !== null;
+        updateStatus(isCapturing);
+        if (!isCapturing) {
+            pressedKeys.clear();
+            mouseVelocity = { x: 0, y: 0 };
+            Object.keys(KEY_DISPLAY_MAP).forEach(code => updateKeyDisplay(code, false));
+            updateMouseDisplay();
+            updateActiveButtonsDisplay();
+            triggerUpdate();
+        }
+    });
+
+    document.addEventListener('keydown', (e) => {
+        if (!isCapturing) return;
+        if (e.code === 'Escape') { document.exitPointerLock(); return; }
+        if (BUTTON_MAP[e.code] !== undefined && !pressedKeys.has(e.code)) {
+            pressedKeys.add(e.code);
+            updateKeyDisplay(e.code, true);
+            updateActiveButtonsDisplay();
+            triggerUpdate();
+        }
+        e.preventDefault();
+    });
+
+    document.addEventListener('keyup', (e) => {
+        if (!isCapturing) return;
+        if (pressedKeys.has(e.code)) {
+            pressedKeys.delete(e.code);
+            updateKeyDisplay(e.code, false);
+            updateActiveButtonsDisplay();
+            triggerUpdate();
+        }
+    });
+
+    const MOUSE_SENSITIVITY = 1.5;
+
+    document.addEventListener('mousemove', (e) => {
+        if (!isCapturing) return;
+        mouseVelocity.x = e.movementX * MOUSE_SENSITIVITY;
+        mouseVelocity.y = e.movementY * MOUSE_SENSITIVITY;
+        updateMouseDisplay();
+        triggerUpdate();
+        lastMouseMove = Date.now();
+    });
+
+    setInterval(() => {
+        if ((isCapturing || isTouchActive) && Date.now() - lastMouseMove > 50) {
+            mouseVelocity.x *= 0.8;
+            mouseVelocity.y *= 0.8;
+            if (Math.abs(mouseVelocity.x) < 0.01) mouseVelocity.x = 0;
+            if (Math.abs(mouseVelocity.y) < 0.01) mouseVelocity.y = 0;
+            updateMouseDisplay();
+            triggerUpdate();
+        }
+    }, 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
+    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;
+        updateMouseDisplay();
+        triggerUpdate();
+        lastMouseMove = Date.now();
+    };
+
+    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();
+        }
+    });
+
+    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; }
+#video-output img { width: 100%; height: 100%; object-fit: contain; }
+"""
+
+
+def create_app():
+    with gr.Blocks(css=css, theme=gr.themes.Soft(), title="WorldEngine") as demo:
+        gr.Markdown("""
+        # 🌍 WorldEngine — Real-Time World Model
+
+        Interactive frame-by-frame world generation.
+
+        **Controls:** WASD to move • Mouse to look • Space to jump • Shift to sprint • E to interact • ESC to release
+        """)
+
+        with gr.Row():
+            with gr.Column(scale=3):
+                video_output = gr.Image(
+                    label="Game View",
+                    elem_id="video-output",
+                    streaming=True,
+                    width=640,
+                    height=360,
+                    show_label=False,
+                )
+                frame_display = gr.Number(label="Frame", value=0, interactive=False)
+
+            with gr.Column(scale=1):
+                control_input = gr.HTML(
+                    value={"buttons": [], "mouse_x": 0.0, "mouse_y": 0.0},
+                    html_template=CONTROL_INPUT_HTML,
+                    js_on_load=CONTROL_INPUT_JS,
+                )
+
+                prompt_input = gr.Textbox(
+                    label="World Prompt",
+                    value="An explorable world",
+                    lines=2,
+                )
+
+                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))],
+                        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":
+                idx = int(seed_choice.split()[-1]) - 1
+                game_state.reset(seed_url=SEED_FRAME_URLS[idx])
+            else:
+                game_state.reset()
+
+        def get_frame():
+            with game_state._frame_lock:
+                return game_state._latest_frame, game_state.frame_count
+
+        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])
+
+        timer = gr.Timer(value=1 / 30)
+        timer.tick(fn=get_frame, outputs=[video_output, frame_display])
+
+        demo.load(fn=None, js=CAPTURE_JS)
+
+    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}")
+
+    # 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,
+    )
+
+    demo = create_app()
+    demo.launch(server_name="0.0.0.0", server_port=7860, share=False)
 
-def greet(name):
-    return "Hello " + name + "!!"
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -0,0 +1,19 @@
+# Core dependencies
+torch>=2.4.0
+torchvision>=0.19.0
+
+# Gradio for web interface
+gradio>=5.0.0
+
+# Diffusers with modular pipeline support
+diffusers>=0.36.0
+
+# WorldEngine modular dependencies
+transformers>=4.40.0
+einops>=0.8.0
+tensordict>=0.5.0
+regex>=2024.0.0
+ftfy>=6.0.0
+
+# HuggingFace Hub
+huggingface_hub>=0.20.0