From 0a0cdebaba3eed5aa32fa8d0092cbcade99716cc Mon Sep 17 00:00:00 2001 From: kevin314 Date: Thu, 5 Feb 2026 04:02:49 +0000 Subject: [PATCH] Add gpu selection --- ui/world_model/client/src/App.svelte | 40 +++++++++++++++++++++++++--- ui/world_model/client/vite.config.js | 11 +++++++- ui/world_model/server/gpu_pool.py | 12 ++++++--- ui/world_model/server/main.py | 13 +++++++-- 4 files changed, 66 insertions(+), 10 deletions(-) diff --git a/ui/world_model/client/src/App.svelte b/ui/world_model/client/src/App.svelte index cce843970..e77c6ec7e 100644 --- a/ui/world_model/client/src/App.svelte +++ b/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 @@ -->
- + {#if !sessionStarted}
Click "Join Session" to start
- {:else if !canvas || frameBuffer.length === 0 && !isPlayingBuffer} + {:else if !loadingAnimation && (!canvas || (frameBuffer.length === 0 && !isPlayingBuffer && !canvasInitialized))}
Waiting for frame...
{/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%; diff --git a/ui/world_model/client/vite.config.js b/ui/world_model/client/vite.config.js index 690caedc3..e98d07a7a 100644 --- a/ui/world_model/client/vite.config.js +++ b/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', + }, + }, } }) diff --git a/ui/world_model/server/gpu_pool.py b/ui/world_model/server/gpu_pool.py index b3ecb9937..154a906c9 100644 --- a/ui/world_model/server/gpu_pool.py +++ b/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() diff --git a/ui/world_model/server/main.py b/ui/world_model/server/main.py index 6f1e9dc65..5fb97017a 100644 --- a/ui/world_model/server/main.py +++ b/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)