Add gpu selection

ui/world_model/client/src/App.svelte

@@ -37,6 +37,7 @@
   let lastFrameTime = 0;
   let canvas;
   let ctx;
+  let loadingAnimation = false;
 
   let benchmarkStats = {
     decodeTime: 0,
@@ -80,6 +81,7 @@
           canvas.width = bitmap.width;
           canvas.height = bitmap.height;
           canvasInitialized = true;
+          loadingAnimation = false;
         }
 
         const renderStart = performance.now();
@@ -125,6 +127,7 @@
     // Clear canvas
     frameBuffer = [];
     canvasInitialized = false;
+    loadingAnimation = false;
   }
 
   function connectWebSocket() {
@@ -141,7 +144,8 @@
     connected = false;
 
     try {
-      ws = new WebSocket('ws://localhost:8000/ws');
+      const wsProtocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
+      ws = new WebSocket(`${wsProtocol}//${window.location.host}/ws`);
 
       ws.onopen = () => {
         connected = true;
@@ -252,6 +256,7 @@
           nextFrameToAdd = 0;
           isPlayingBuffer = false;
           canvasInitialized = false;
+          loadingAnimation = false;
         } else if (data.type === 'reset_complete') {
           console.log('Reset complete');
           resetting = false;
@@ -267,6 +272,23 @@
           timeLeft = data.session_timeout;
           console.log(`GPU ${data.gpu_id} assigned, session timeout: ${sessionTimeout}s`);
 
+          // Load initial image onto canvas with animation
+          if (data.image_url && canvas) {
+            const img = new Image();
+            img.crossOrigin = 'anonymous';
+            img.onload = () => {
+              if (!canvasInitialized) {
+                canvas.width = img.width;
+                canvas.height = img.height;
+                if (ctx) {
+                  ctx.drawImage(img, 0, 0);
+                }
+                loadingAnimation = true;
+              }
+            };
+            img.src = data.image_url;
+          }
+
           // Start countdown timer
           if (countdownInterval) clearInterval(countdownInterval);
           countdownInterval = setInterval(() => {
@@ -338,6 +360,8 @@
       countdownInterval = null;
     }
     frameBuffer = [];
+    canvasInitialized = false;
+    loadingAnimation = false;
 
     // Reconnect
     setTimeout(connectWebSocket, 100);
@@ -479,10 +503,10 @@
     </div> -->
 
     <div class="canvas">
-    <canvas bind:this={canvas} class:blurred={resetting} style="display: block; max-width: 100%; height: auto;"></canvas>
+    <canvas bind:this={canvas} class:blurred={resetting} class:loading-animation={loadingAnimation} style="display: block; max-width: 100%; height: auto;"></canvas>
     {#if !sessionStarted}
       <div class="placeholder">Click "Join Session" to start</div>
-    {:else if !canvas || frameBuffer.length === 0 && !isPlayingBuffer}
+    {:else if !loadingAnimation && (!canvas || (frameBuffer.length === 0 && !isPlayingBuffer && !canvasInitialized))}
       <div class="placeholder">Waiting for frame...</div>
     {/if}
 
@@ -861,6 +885,16 @@
     transition: filter 0.3s ease;
   }
 
+  @keyframes warp-in {
+    0%   { transform: scale(1.05); filter: blur(6px) brightness(0.5); }
+    50%  { transform: scale(1.02); filter: blur(3px) brightness(0.8); }
+    100% { transform: scale(1.0);  filter: blur(0px) brightness(1.0); }
+  }
+
+  canvas.loading-animation {
+    animation: warp-in 1.5s ease-out forwards;
+  }
+
   .placeholder {
     width: 100%;
     height: 100%;

ui/world_model/client/vite.config.js

@@ -4,6 +4,15 @@ import { svelte } from '@sveltejs/vite-plugin-svelte'
 export default defineConfig({
   plugins: [svelte()],
   server: {
-    port: 5173
+    port: 5173,
+    proxy: {
+      '/ws': {
+        target: 'http://localhost:8000',
+        ws: true,
+      },
+      '/status': {
+        target: 'http://localhost:8000',
+      },
+    },
   }
 })

ui/world_model/server/gpu_pool.py

@@ -45,6 +45,7 @@ class Response:
 
 def gpu_worker_process(
     gpu_id: int,
+    cuda_device: str,
     command_queue: Queue,
     response_queue: Queue,
 ):
@@ -54,7 +55,7 @@ def gpu_worker_process(
     This function runs in a subprocess with CUDA_VISIBLE_DEVICES set to a single GPU.
     """
     # Set CUDA_VISIBLE_DEVICES BEFORE importing torch or any CUDA code
-    os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
+    os.environ["CUDA_VISIBLE_DEVICES"] = cuda_device
     os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "FLASH_ATTN"
 
     # Now import the generator (this will initialize CUDA with the single visible GPU)
@@ -235,8 +236,9 @@ def gpu_worker_process(
 class GPUSlot:
     """Manages a single GPU worker subprocess."""
 
-    def __init__(self, gpu_id: int):
+    def __init__(self, gpu_id: int, cuda_device: str):
         self.gpu_id = gpu_id
+        self.cuda_device = cuda_device
         self.client_id: Optional[str] = None
         self.process: Optional[Process] = None
         self.command_queue: Optional[Queue] = None
@@ -256,7 +258,7 @@ class GPUSlot:
 
         self.process = ctx.Process(
             target=gpu_worker_process,
-            args=(self.gpu_id, self.command_queue, self.response_queue),
+            args=(self.gpu_id, self.cuda_device, self.command_queue, self.response_queue),
             daemon=False,  # Must be False so executor can spawn child workers
         )
         loop = asyncio.get_event_loop()
@@ -319,7 +321,9 @@ class GPUPool:
 
     def __init__(self, gpu_ids: list[int]):
         self.gpu_ids = gpu_ids
-        self.slots: dict[int, GPUSlot] = {gpu_id: GPUSlot(gpu_id) for gpu_id in gpu_ids}
+        self.slots: dict[int, GPUSlot] = {
+            gpu_id: GPUSlot(gpu_id, str(i)) for i, gpu_id in enumerate(gpu_ids)
+        }
         self.waiting_list: list[tuple[str, asyncio.Event, "WebSocket"]] = []
         self.client_gpu_map: dict[str, int] = {}
         self._pool_lock = asyncio.Lock()

ui/world_model/server/main.py

@@ -1,15 +1,17 @@
 from fastapi import FastAPI, WebSocket, WebSocketDisconnect
 from fastapi.middleware.cors import CORSMiddleware
+from fastapi.staticfiles import StaticFiles
 from contextlib import asynccontextmanager
 import asyncio
 import base64
 import cv2
+import os
 import time
 import uuid
 
 from config import (
     KEYBOARD_MAP, CAMERA_MAP, MAX_BLOCKS, JPEG_QUALITY, BATCH_SIZE,
-    SESSION_TIMEOUT_SECONDS
+    SESSION_TIMEOUT_SECONDS, MODEL_CONFIG
 )
 from gpu_pool import GPUPool, GPUSlot, get_available_gpus
 
@@ -122,7 +124,8 @@ async def websocket_endpoint(websocket: WebSocket):
         await websocket.send_json({
             "type": "gpu_assigned",
             "gpu_id": gpu_id,
-            "session_timeout": SESSION_TIMEOUT_SECONDS
+            "session_timeout": SESSION_TIMEOUT_SECONDS,
+            "image_url": MODEL_CONFIG.get("image_url")
         })
 
         # Generate and send initial frame (GPU is already reset and ready)
@@ -202,6 +205,12 @@ async def websocket_endpoint(websocket: WebSocket):
             await gpu_pool.release(client_id)
 
 
+# Serve built frontend (must be after API/WebSocket routes)
+static_dir = os.path.join(os.path.dirname(__file__), "..", "client", "dist")
+if os.path.isdir(static_dir):
+    app.mount("/", StaticFiles(directory=static_dir, html=True), name="static")
+
+
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=8000)