.

dispatcher.py

@@ -0,0 +1,441 @@
+from fastapi import FastAPI, WebSocket, WebSocketDisconnect, HTTPException
+from fastapi.responses import HTMLResponse
+from fastapi.staticfiles import StaticFiles
+from typing import List, Dict, Any, Optional
+import asyncio
+import json
+import time
+import os
+from dataclasses import dataclass, asdict
+from enum import Enum
+import uuid
+import aiohttp
+import logging
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+class SessionStatus(Enum):
+    QUEUED = "queued"
+    ACTIVE = "active"
+    COMPLETED = "completed"
+    TIMEOUT = "timeout"
+
+@dataclass
+class UserSession:
+    session_id: str
+    client_id: str
+    websocket: WebSocket
+    created_at: float
+    status: SessionStatus
+    worker_id: Optional[str] = None
+    last_activity: Optional[float] = None
+    max_session_time: Optional[float] = None
+    user_has_interacted: bool = False
+
+@dataclass
+class WorkerInfo:
+    worker_id: str
+    gpu_id: int
+    endpoint: str
+    is_available: bool
+    current_session: Optional[str] = None
+    last_ping: float = 0
+
+class SessionManager:
+    def __init__(self):
+        self.sessions: Dict[str, UserSession] = {}
+        self.workers: Dict[str, WorkerInfo] = {}
+        self.session_queue: List[str] = []
+        self.active_sessions: Dict[str, str] = {}  # session_id -> worker_id
+        
+        # Configuration
+        self.IDLE_TIMEOUT = 20.0  # When no queue
+        self.QUEUE_WARNING_TIME = 10.0
+        self.MAX_SESSION_TIME_WITH_QUEUE = 60.0  # When there's a queue
+        self.QUEUE_SESSION_WARNING_TIME = 45.0  # 15 seconds before timeout
+        self.GRACE_PERIOD = 10.0
+
+    async def register_worker(self, worker_id: str, gpu_id: int, endpoint: str):
+        """Register a new worker"""
+        self.workers[worker_id] = WorkerInfo(
+            worker_id=worker_id,
+            gpu_id=gpu_id,
+            endpoint=endpoint,
+            is_available=True,
+            last_ping=time.time()
+        )
+        logger.info(f"Registered worker {worker_id} on GPU {gpu_id} at {endpoint}")
+
+    async def get_available_worker(self) -> Optional[WorkerInfo]:
+        """Get an available worker"""
+        for worker in self.workers.values():
+            if worker.is_available and time.time() - worker.last_ping < 30:  # Worker ping timeout
+                return worker
+        return None
+
+    async def add_session_to_queue(self, session: UserSession):
+        """Add a session to the queue"""
+        self.sessions[session.session_id] = session
+        self.session_queue.append(session.session_id)
+        session.status = SessionStatus.QUEUED
+        logger.info(f"Added session {session.session_id} to queue. Queue size: {len(self.session_queue)}")
+
+    async def process_queue(self):
+        """Process the session queue"""
+        while self.session_queue:
+            session_id = self.session_queue[0]
+            session = self.sessions.get(session_id)
+            
+            if not session or session.status != SessionStatus.QUEUED:
+                self.session_queue.pop(0)
+                continue
+                
+            worker = await self.get_available_worker()
+            if not worker:
+                break  # No available workers
+                
+            # Assign session to worker
+            self.session_queue.pop(0)
+            session.status = SessionStatus.ACTIVE
+            session.worker_id = worker.worker_id
+            session.last_activity = time.time()
+            
+            # Set session time limit based on queue status
+            if len(self.session_queue) > 0:
+                session.max_session_time = self.MAX_SESSION_TIME_WITH_QUEUE
+            
+            worker.is_available = False
+            worker.current_session = session_id
+            self.active_sessions[session_id] = worker.worker_id
+            
+            logger.info(f"Assigned session {session_id} to worker {worker.worker_id}")
+            
+            # Notify user that their session is starting
+            await self.notify_session_start(session)
+            
+            # Start session monitoring
+            asyncio.create_task(self.monitor_active_session(session_id))
+
+    async def notify_session_start(self, session: UserSession):
+        """Notify user that their session is starting"""
+        try:
+            await session.websocket.send_json({
+                "type": "session_start",
+                "worker_id": session.worker_id,
+                "max_session_time": session.max_session_time
+            })
+        except Exception as e:
+            logger.error(f"Failed to notify session start for {session.session_id}: {e}")
+
+    async def monitor_active_session(self, session_id: str):
+        """Monitor an active session for timeouts"""
+        session = self.sessions.get(session_id)
+        if not session:
+            return
+            
+        try:
+            while session.status == SessionStatus.ACTIVE:
+                current_time = time.time()
+                
+                # Check if session has exceeded time limit
+                if session.max_session_time:
+                    elapsed = current_time - session.last_activity if session.last_activity else 0
+                    remaining = session.max_session_time - elapsed
+                    
+                    # Send warning at 15 seconds before timeout
+                    if remaining <= 15 and remaining > 10:
+                        await session.websocket.send_json({
+                            "type": "session_warning",
+                            "time_remaining": remaining,
+                            "queue_size": len(self.session_queue)
+                        })
+                    
+                    # Grace period handling
+                    elif remaining <= 10 and remaining > 0:
+                        # Check if queue is empty - if so, extend session
+                        if len(self.session_queue) == 0:
+                            session.max_session_time = None  # Remove time limit
+                            await session.websocket.send_json({
+                                "type": "time_limit_removed",
+                                "reason": "queue_empty"
+                            })
+                        else:
+                            await session.websocket.send_json({
+                                "type": "grace_period",
+                                "time_remaining": remaining,
+                                "queue_size": len(self.session_queue)
+                            })
+                    
+                    # Timeout
+                    elif remaining <= 0:
+                        await self.end_session(session_id, SessionStatus.TIMEOUT)
+                        return
+                
+                # Check idle timeout when no queue
+                elif not session.max_session_time and session.last_activity:
+                    idle_time = current_time - session.last_activity
+                    if idle_time >= self.IDLE_TIMEOUT:
+                        await self.end_session(session_id, SessionStatus.TIMEOUT)
+                        return
+                    elif idle_time >= self.QUEUE_WARNING_TIME:
+                        await session.websocket.send_json({
+                            "type": "idle_warning",
+                            "time_remaining": self.IDLE_TIMEOUT - idle_time
+                        })
+                
+                await asyncio.sleep(1)  # Check every second
+                
+        except Exception as e:
+            logger.error(f"Error monitoring session {session_id}: {e}")
+            await self.end_session(session_id, SessionStatus.COMPLETED)
+
+    async def end_session(self, session_id: str, status: SessionStatus):
+        """End a session and free up the worker"""
+        session = self.sessions.get(session_id)
+        if not session:
+            return
+            
+        session.status = status
+        
+        # Free up the worker
+        if session.worker_id and session.worker_id in self.workers:
+            worker = self.workers[session.worker_id]
+            worker.is_available = True
+            worker.current_session = None
+            
+            # Notify worker to clean up
+            try:
+                async with aiohttp.ClientSession() as client_session:
+                    await client_session.post(f"{worker.endpoint}/end_session", 
+                                            json={"session_id": session_id})
+            except Exception as e:
+                logger.error(f"Failed to notify worker {worker.worker_id} of session end: {e}")
+        
+        # Remove from active sessions
+        if session_id in self.active_sessions:
+            del self.active_sessions[session_id]
+            
+        logger.info(f"Ended session {session_id} with status {status}")
+        
+        # Process next in queue
+        asyncio.create_task(self.process_queue())
+
+    async def update_queue_info(self):
+        """Send queue information to waiting users"""
+        for i, session_id in enumerate(self.session_queue):
+            session = self.sessions.get(session_id)
+            if session and session.status == SessionStatus.QUEUED:
+                try:
+                    # Calculate estimated wait time
+                    active_sessions_count = len(self.active_sessions)
+                    avg_session_time = self.MAX_SESSION_TIME_WITH_QUEUE if active_sessions_count > 0 else 30.0
+                    estimated_wait = (i + 1) * avg_session_time / max(len(self.workers), 1)
+                    
+                    await session.websocket.send_json({
+                        "type": "queue_update",
+                        "position": i + 1,
+                        "total_waiting": len(self.session_queue),
+                        "estimated_wait_minutes": estimated_wait / 60,
+                        "active_sessions": active_sessions_count
+                    })
+                except Exception as e:
+                    logger.error(f"Failed to send queue update to session {session_id}: {e}")
+
+    async def handle_user_activity(self, session_id: str):
+        """Update user activity timestamp"""
+        session = self.sessions.get(session_id)
+        if session:
+            session.last_activity = time.time()
+            if not session.user_has_interacted:
+                session.user_has_interacted = True
+                logger.info(f"User started interacting in session {session_id}")
+
+    async def _forward_to_worker(self, worker: WorkerInfo, session_id: str, data: dict):
+        """Forward input to worker asynchronously"""
+        try:
+            async with aiohttp.ClientSession() as client_session:
+                async with client_session.post(
+                    f"{worker.endpoint}/process_input",
+                    json={
+                        "session_id": session_id,
+                        "data": data
+                    }
+                ) as response:
+                    if response.status != 200:
+                        logger.error(f"Worker returned status {response.status}")
+                        # Optionally handle worker errors here
+        except Exception as e:
+            logger.error(f"Error forwarding to worker {worker.worker_id}: {e}")
+
+# Global session manager
+session_manager = SessionManager()
+
+app = FastAPI()
+app.mount("/static", StaticFiles(directory="static"), name="static")
+
+@app.get("/")
+async def get():
+    return HTMLResponse(open("static/index.html").read())
+
+@app.post("/register_worker")
+async def register_worker(worker_info: dict):
+    """Endpoint for workers to register themselves"""
+    await session_manager.register_worker(
+        worker_info["worker_id"],
+        worker_info["gpu_id"], 
+        worker_info["endpoint"]
+    )
+    return {"status": "registered"}
+
+@app.post("/worker_ping")
+async def worker_ping(worker_info: dict):
+    """Endpoint for workers to ping their availability"""
+    worker_id = worker_info["worker_id"]
+    if worker_id in session_manager.workers:
+        session_manager.workers[worker_id].last_ping = time.time()
+        session_manager.workers[worker_id].is_available = worker_info.get("is_available", True)
+    return {"status": "ok"}
+
+@app.post("/worker_result")
+async def worker_result(result_data: dict):
+    """Endpoint for workers to send back processing results"""
+    session_id = result_data.get("session_id")
+    worker_id = result_data.get("worker_id")
+    result = result_data.get("result")
+    
+    if not session_id or not result:
+        raise HTTPException(status_code=400, detail="Missing session_id or result")
+    
+    # Find the session and send result to the WebSocket
+    session = session_manager.sessions.get(session_id)
+    if session and session.status == SessionStatus.ACTIVE:
+        try:
+            await session.websocket.send_json(result)
+            logger.info(f"Sent result to session {session_id}")
+        except Exception as e:
+            logger.error(f"Failed to send result to session {session_id}: {e}")
+    else:
+        logger.warning(f"Could not find active session {session_id} for result")
+    
+    return {"status": "ok"}
+
+@app.websocket("/ws")
+async def websocket_endpoint(websocket: WebSocket):
+    await websocket.accept()
+    
+    # Create session
+    session_id = str(uuid.uuid4())
+    client_id = f"{int(time.time())}_{session_id[:8]}"
+    
+    session = UserSession(
+        session_id=session_id,
+        client_id=client_id,
+        websocket=websocket,
+        created_at=time.time(),
+        status=SessionStatus.QUEUED
+    )
+    
+    logger.info(f"New WebSocket connection: {client_id}")
+    
+    try:
+        # Add to queue
+        await session_manager.add_session_to_queue(session)
+        
+        # Try to process queue immediately
+        await session_manager.process_queue()
+        
+        # Send initial queue status
+        if session.status == SessionStatus.QUEUED:
+            await session_manager.update_queue_info()
+        
+        # Main message loop
+        while True:
+            try:
+                data = await websocket.receive_json()
+                
+                # Update activity
+                await session_manager.handle_user_activity(session_id)
+                
+                # Handle different message types
+                if data.get("type") == "heartbeat":
+                    await websocket.send_json({"type": "heartbeat_response"})
+                    continue
+                
+                # If session is active, forward to worker
+                if session.status == SessionStatus.ACTIVE and session.worker_id:
+                    worker = session_manager.workers.get(session.worker_id)
+                    if worker:
+                        try:
+                            # Forward message to worker (don't wait for response for regular inputs)
+                            # The worker will send results back asynchronously via /worker_result
+                            asyncio.create_task(session_manager._forward_to_worker(worker, session_id, data))
+                        except Exception as e:
+                            logger.error(f"Error forwarding to worker: {e}")
+                
+                # Handle control messages (these need synchronous responses)
+                elif data.get("type") in ["reset", "update_sampling_steps", "update_use_rnn", "get_settings"]:
+                    if session.status == SessionStatus.ACTIVE and session.worker_id:
+                        worker = session_manager.workers.get(session.worker_id)
+                        if worker:
+                            try:
+                                async with aiohttp.ClientSession() as client_session:
+                                    async with client_session.post(
+                                        f"{worker.endpoint}/process_input",
+                                        json={
+                                            "session_id": session_id,
+                                            "data": data
+                                        }
+                                    ) as response:
+                                        if response.status == 200:
+                                            result = await response.json()
+                                            await websocket.send_json(result)
+                                        else:
+                                            logger.error(f"Worker returned status {response.status}")
+                            except Exception as e:
+                                logger.error(f"Error forwarding control message: {e}")
+                    else:
+                        # Send appropriate response for queued users
+                        await websocket.send_json({
+                            "type": "error",
+                            "message": "Session not active yet. Please wait in queue."
+                        })
+                        
+            except asyncio.TimeoutError:
+                logger.info("WebSocket connection timed out")
+                break
+            except WebSocketDisconnect:
+                logger.info(f"WebSocket disconnected: {client_id}")
+                break
+                
+    except Exception as e:
+        logger.error(f"Error in WebSocket connection {client_id}: {e}")
+        import traceback
+        traceback.print_exc()
+    
+    finally:
+        # Clean up session
+        if session_id in session_manager.sessions:
+            await session_manager.end_session(session_id, SessionStatus.COMPLETED)
+            del session_manager.sessions[session_id]
+        
+        logger.info(f"WebSocket connection closed: {client_id}")
+
+# Background task to periodically update queue info
+async def periodic_queue_update():
+    while True:
+        try:
+            await session_manager.update_queue_info()
+            await asyncio.sleep(5)  # Update every 5 seconds
+        except Exception as e:
+            logger.error(f"Error in periodic queue update: {e}")
+
+@app.on_event("startup")
+async def startup_event():
+    # Start background tasks
+    asyncio.create_task(periodic_queue_update())
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8000) 

start_workers.py

@@ -0,0 +1,152 @@
+#!/usr/bin/env python3
+"""
+Script to start multiple GPU workers for the neural OS demo.
+Usage: python start_workers.py --num-gpus 4
+"""
+
+import argparse
+import subprocess
+import time
+import sys
+import signal
+import os
+from typing import List
+
+class WorkerManager:
+    def __init__(self, num_gpus: int, dispatcher_url: str = "http://localhost:8000"):
+        self.num_gpus = num_gpus
+        self.dispatcher_url = dispatcher_url
+        self.processes: List[subprocess.Popen] = []
+        
+    def start_workers(self):
+        """Start all worker processes"""
+        print(f"Starting {self.num_gpus} GPU workers...")
+        
+        for gpu_id in range(self.num_gpus):
+            try:
+                port = 8001 + gpu_id
+                print(f"Starting worker for GPU {gpu_id} on port {port}...")
+                
+                # Start worker process
+                cmd = [
+                    sys.executable, "worker.py",
+                    "--gpu-id", str(gpu_id),
+                    "--dispatcher-url", self.dispatcher_url
+                ]
+                
+                process = subprocess.Popen(
+                    cmd,
+                    stdout=subprocess.PIPE,
+                    stderr=subprocess.STDOUT,
+                    universal_newlines=True,
+                    bufsize=1
+                )
+                
+                self.processes.append(process)
+                print(f"✓ Started worker {gpu_id} (PID: {process.pid})")
+                
+                # Small delay between starts
+                time.sleep(1)
+                
+            except Exception as e:
+                print(f"✗ Failed to start worker for GPU {gpu_id}: {e}")
+                self.cleanup()
+                return False
+        
+        print(f"\n✓ All {self.num_gpus} workers started successfully!")
+        print("Workers are running on ports:", [8001 + i for i in range(self.num_gpus)])
+        return True
+    
+    def monitor_workers(self):
+        """Monitor worker processes and print their output"""
+        print("\nMonitoring workers (Ctrl+C to stop)...")
+        print("-" * 50)
+        
+        try:
+            while True:
+                # Check if any process has died
+                for i, process in enumerate(self.processes):
+                    if process.poll() is not None:
+                        print(f"⚠️  Worker {i} (PID: {process.pid}) has died!")
+                        # Optionally restart it
+                        
+                # Print output from processes
+                for i, process in enumerate(self.processes):
+                    if process.stdout and process.stdout.readable():
+                        try:
+                            line = process.stdout.readline()
+                            if line:
+                                print(f"[GPU {i}] {line.strip()}")
+                        except:
+                            pass
+                
+                time.sleep(0.1)
+                
+        except KeyboardInterrupt:
+            print("\n\nReceived interrupt signal, shutting down workers...")
+            self.cleanup()
+    
+    def cleanup(self):
+        """Clean up all worker processes"""
+        print("Stopping all workers...")
+        
+        for i, process in enumerate(self.processes):
+            if process.poll() is None:  # Process is still running
+                print(f"Stopping worker {i} (PID: {process.pid})...")
+                try:
+                    process.terminate()
+                    # Wait for graceful shutdown
+                    process.wait(timeout=5)
+                    print(f"✓ Worker {i} stopped gracefully")
+                except subprocess.TimeoutExpired:
+                    print(f"⚠️  Force killing worker {i}...")
+                    process.kill()
+                    process.wait()
+                except Exception as e:
+                    print(f"Error stopping worker {i}: {e}")
+        
+        print("✓ All workers stopped")
+
+def main():
+    parser = argparse.ArgumentParser(description="Start multiple GPU workers")
+    parser.add_argument("--num-gpus", type=int, required=True, 
+                       help="Number of GPU workers to start")
+    parser.add_argument("--dispatcher-url", type=str, default="http://localhost:8000",
+                       help="URL of the dispatcher service")
+    parser.add_argument("--no-monitor", action="store_true",
+                       help="Start workers but don't monitor them")
+    
+    args = parser.parse_args()
+    
+    if args.num_gpus < 1:
+        print("Error: Number of GPUs must be at least 1")
+        sys.exit(1)
+    
+    # Check if worker.py exists
+    if not os.path.exists("worker.py"):
+        print("Error: worker.py not found in current directory")
+        sys.exit(1)
+    
+    manager = WorkerManager(args.num_gpus, args.dispatcher_url)
+    
+    # Set up signal handlers for clean shutdown
+    def signal_handler(sig, frame):
+        print(f"\nReceived signal {sig}, shutting down...")
+        manager.cleanup()
+        sys.exit(0)
+    
+    signal.signal(signal.SIGINT, signal_handler)
+    signal.signal(signal.SIGTERM, signal_handler)
+    
+    # Start workers
+    if not manager.start_workers():
+        sys.exit(1)
+    
+    if not args.no_monitor:
+        manager.monitor_workers()
+    else:
+        print("Workers started. Use 'ps aux | grep worker.py' to check status.")
+        print("To stop workers, use: pkill -f 'python.*worker.py'")
+
+if __name__ == "__main__":
+    main() 

ttt.py

@@ -0,0 +1,8 @@
+import pyautogui
+import time
+
+time.sleep(3)  # Gives you 3 seconds to switch to another window (e.g., a text editor)
+
+pyautogui.press(' ')       # Sends a space
+pyautogui.press('space')   # Sends another space
+

worker.py

@@ -0,0 +1,635 @@
+from fastapi import FastAPI, HTTPException
+from typing import List, Tuple, Dict, Any, Optional
+import numpy as np
+from PIL import Image, ImageDraw
+import base64
+import io
+import json
+import asyncio
+import time
+import torch
+import os
+import logging
+from utils import initialize_model, sample_frame
+from ldm.models.diffusion.ddpm import LatentDiffusion, DDIMSampler
+import concurrent.futures
+import aiohttp
+import argparse
+import uuid
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# GPU settings
+torch.backends.cuda.matmul.allow_tf32 = True
+torch.backends.cudnn.allow_tf32 = True
+
+class GPUWorker:
+    def __init__(self, gpu_id: int, dispatcher_url: str = "http://localhost:8000"):
+        self.gpu_id = gpu_id
+        self.dispatcher_url = dispatcher_url
+        self.worker_id = f"worker_{gpu_id}_{uuid.uuid4().hex[:8]}"
+        self.device = torch.device(f'cuda:{gpu_id}' if torch.cuda.is_available() else 'cpu')
+        self.current_session: Optional[str] = None
+        self.session_data: Dict[str, Any] = {}
+        
+        # Model configuration from main.py
+        self.DEBUG_MODE = False
+        self.DEBUG_MODE_2 = False
+        self.NUM_MAX_FRAMES = 1
+        self.TIMESTEPS = 1000
+        self.SCREEN_WIDTH = 512
+        self.SCREEN_HEIGHT = 384
+        self.NUM_SAMPLING_STEPS = 32
+        self.USE_RNN = False
+        
+        self.MODEL_NAME = "yuntian-deng/computer-model-s-newnewd-freezernn-origunet-nospatial-online-x0-joint-onlineonly-222222k7-06k"
+        
+        # Initialize model
+        self._initialize_model()
+        
+        # Thread executor for heavy computation
+        self.thread_executor = concurrent.futures.ThreadPoolExecutor(max_workers=1)
+        
+        # Load keyboard mappings
+        self._load_keyboard_mappings()
+        
+        logger.info(f"GPU Worker {self.worker_id} initialized on GPU {gpu_id}")
+
+    def _initialize_model(self):
+        """Initialize the model on the specified GPU"""
+        logger.info(f"Initializing model on GPU {self.gpu_id}")
+        
+        # Load latent stats
+        with open('latent_stats.json', 'r') as f:
+            latent_stats = json.load(f)
+        self.DATA_NORMALIZATION = {
+            'mean': torch.tensor(latent_stats['mean']).to(self.device), 
+            'std': torch.tensor(latent_stats['std']).to(self.device)
+        }
+        self.LATENT_DIMS = (16, self.SCREEN_HEIGHT // 8, self.SCREEN_WIDTH // 8)
+        
+        # Initialize model based on model name
+        if 'origunet' in self.MODEL_NAME:
+            if 'x0' in self.MODEL_NAME:
+                if 'ddpm32' in self.MODEL_NAME:
+                    self.TIMESTEPS = 32
+                    self.model = initialize_model("config_final_model_origunet_nospatial_x0_ddpm32.yaml", self.MODEL_NAME)
+                else:
+                    self.model = initialize_model("config_final_model_origunet_nospatial_x0.yaml", self.MODEL_NAME)
+            else:
+                if 'ddpm32' in self.MODEL_NAME:
+                    self.TIMESTEPS = 32
+                    self.model = initialize_model("config_final_model_origunet_nospatial_ddpm32.yaml", self.MODEL_NAME)
+                else:
+                    self.model = initialize_model("config_final_model_origunet_nospatial.yaml", self.MODEL_NAME)
+        else:
+            self.model = initialize_model("config_final_model.yaml", self.MODEL_NAME)
+        
+        self.model = self.model.to(self.device)
+        
+        # Create padding image
+        self.padding_image = torch.zeros(*self.LATENT_DIMS).unsqueeze(0).to(self.device)
+        self.padding_image = (self.padding_image - self.DATA_NORMALIZATION['mean'].view(1, -1, 1, 1)) / self.DATA_NORMALIZATION['std'].view(1, -1, 1, 1)
+        
+        logger.info(f"Model initialized successfully on GPU {self.gpu_id}")
+
+    def _load_keyboard_mappings(self):
+        """Load keyboard mappings from main.py"""
+        self.KEYS = ['\t', '\n', '\r', ' ', '!', '"', '#', '$', '%', '&', "'", '(',
+                    ')', '*', '+', ',', '-', '.', '/', '0', '1', '2', '3', '4', '5', '6', '7',
+                    '8', '9', ':', ';', '<', '=', '>', '?', '@', '[', '\\', ']', '^', '_', '`',
+                    'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o',
+                    'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z', '{', '|', '}', '~',
+                    'accept', 'add', 'alt', 'altleft', 'altright', 'apps', 'backspace',
+                    'browserback', 'browserfavorites', 'browserforward', 'browserhome',
+                    'browserrefresh', 'browsersearch', 'browserstop', 'capslock', 'clear',
+                    'convert', 'ctrl', 'ctrlleft', 'ctrlright', 'decimal', 'del', 'delete',
+                    'divide', 'down', 'end', 'enter', 'esc', 'escape', 'execute', 'f1', 'f10',
+                    'f11', 'f12', 'f13', 'f14', 'f15', 'f16', 'f17', 'f18', 'f19', 'f2', 'f20',
+                    'f21', 'f22', 'f23', 'f24', 'f3', 'f4', 'f5', 'f6', 'f7', 'f8', 'f9',
+                    'final', 'fn', 'hanguel', 'hangul', 'hanja', 'help', 'home', 'insert', 'junja',
+                    'kana', 'kanji', 'launchapp1', 'launchapp2', 'launchmail',
+                    'launchmediaselect', 'left', 'modechange', 'multiply', 'nexttrack',
+                    'nonconvert', 'num0', 'num1', 'num2', 'num3', 'num4', 'num5', 'num6',
+                    'num7', 'num8', 'num9', 'numlock', 'pagedown', 'pageup', 'pause', 'pgdn',
+                    'pgup', 'playpause', 'prevtrack', 'print', 'printscreen', 'prntscrn',
+                    'prtsc', 'prtscr', 'return', 'right', 'scrolllock', 'select', 'separator',
+                    'shift', 'shiftleft', 'shiftright', 'sleep', 'space', 'stop', 'subtract', 'tab',
+                    'up', 'volumedown', 'volumemute', 'volumeup', 'win', 'winleft', 'winright', 'yen',
+                    'command', 'option', 'optionleft', 'optionright']
+
+        self.KEYMAPPING = {
+            'arrowup': 'up',
+            'arrowdown': 'down',
+            'arrowleft': 'left',
+            'arrowright': 'right',
+            'meta': 'command',
+            'contextmenu': 'apps',
+            'control': 'ctrl',
+        }
+
+        self.INVALID_KEYS = ['f13', 'f14', 'f15', 'f16', 'f17', 'f18', 'f19', 'f20',
+                            'f21', 'f22', 'f23', 'f24', 'select', 'separator', 'execute']
+        self.VALID_KEYS = [key for key in self.KEYS if key not in self.INVALID_KEYS]
+        self.itos = self.VALID_KEYS
+        self.stoi = {key: i for i, key in enumerate(self.itos)}
+
+    async def register_with_dispatcher(self):
+        """Register this worker with the dispatcher"""
+        try:
+            async with aiohttp.ClientSession() as session:
+                await session.post(f"{self.dispatcher_url}/register_worker", json={
+                    "worker_id": self.worker_id,
+                    "gpu_id": self.gpu_id,
+                    "endpoint": f"http://localhost:{8001 + self.gpu_id}"
+                })
+            logger.info(f"Successfully registered worker {self.worker_id} with dispatcher")
+        except Exception as e:
+            logger.error(f"Failed to register with dispatcher: {e}")
+
+    async def ping_dispatcher(self):
+        """Periodically ping the dispatcher to maintain connection"""
+        while True:
+            try:
+                async with aiohttp.ClientSession() as session:
+                    await session.post(f"{self.dispatcher_url}/worker_ping", json={
+                        "worker_id": self.worker_id,
+                        "is_available": self.current_session is None
+                    })
+                await asyncio.sleep(10)  # Ping every 10 seconds
+            except Exception as e:
+                logger.error(f"Failed to ping dispatcher: {e}")
+                await asyncio.sleep(5)  # Retry after 5 seconds on error
+
+    def prepare_model_inputs(
+        self, 
+        previous_frame: torch.Tensor,
+        hidden_states: Any,
+        x: int,
+        y: int,
+        right_click: bool,
+        left_click: bool,
+        keys_down: List[str],
+        time_step: int
+    ) -> Dict[str, torch.Tensor]:
+        """Prepare inputs for the model (from main.py)"""
+        # Clamp coordinates to valid ranges
+        x = min(max(0, x), self.SCREEN_WIDTH - 1) if x is not None else 0
+        y = min(max(0, y), self.SCREEN_HEIGHT - 1) if y is not None else 0
+        
+        if self.DEBUG_MODE:
+            logger.info('DEBUG MODE, SETTING TIME STEP TO 0')
+            time_step = 0
+        if self.DEBUG_MODE_2:
+            if time_step > self.NUM_MAX_FRAMES-1:
+                logger.info('DEBUG MODE_2, SETTING TIME STEP TO 0')
+                time_step = 0
+        
+        inputs = {
+            'image_features': previous_frame.to(self.device),
+            'is_padding': torch.BoolTensor([time_step == 0]).to(self.device),
+            'x': torch.LongTensor([x]).unsqueeze(0).to(self.device),
+            'y': torch.LongTensor([y]).unsqueeze(0).to(self.device),
+            'is_leftclick': torch.BoolTensor([left_click]).unsqueeze(0).to(self.device),
+            'is_rightclick': torch.BoolTensor([right_click]).unsqueeze(0).to(self.device),
+            'key_events': torch.zeros(len(self.itos), dtype=torch.long).to(self.device)
+        }
+        
+        for key in keys_down:
+            key = key.lower()
+            if key in self.KEYMAPPING:
+                key = self.KEYMAPPING[key]
+            if key in self.stoi:
+                inputs['key_events'][self.stoi[key]] = 1
+            else:
+                logger.warning(f'Key {key} not found in stoi')
+        
+        if hidden_states is not None:
+            inputs['hidden_states'] = hidden_states
+            
+        if self.DEBUG_MODE:
+            logger.info('DEBUG MODE, REMOVING INPUTS')
+            if 'hidden_states' in inputs:
+                del inputs['hidden_states']
+                
+        if self.DEBUG_MODE_2:
+            if time_step > self.NUM_MAX_FRAMES-1:
+                logger.info('DEBUG MODE_2, REMOVING HIDDEN STATES')
+                if 'hidden_states' in inputs:
+                    del inputs['hidden_states']
+                    
+        logger.info(f'Time step: {time_step}')
+        return inputs
+
+    @torch.no_grad()
+    async def process_frame(
+        self,
+        inputs: Dict[str, torch.Tensor],
+        use_rnn: bool = False,
+        num_sampling_steps: int = 32
+    ) -> Tuple[torch.Tensor, np.ndarray, Any, Dict[str, float]]:
+        """Process a single frame through the model"""
+        # Run the heavy computation in a separate thread
+        loop = asyncio.get_running_loop()
+        return await loop.run_in_executor(
+            self.thread_executor,
+            lambda: self._process_frame_sync(inputs, use_rnn, num_sampling_steps)
+        )
+
+    def _process_frame_sync(self, inputs, use_rnn, num_sampling_steps):
+        """Synchronous version of process_frame that runs in a thread"""
+        timing = {}
+        
+        # Temporal encoding
+        start = time.perf_counter()
+        output_from_rnn, hidden_states = self.model.temporal_encoder.forward_step(inputs)
+        timing['temporal_encoder'] = time.perf_counter() - start
+        
+        # UNet sampling
+        start = time.perf_counter()
+        logger.info(f"model.clip_denoised: {self.model.clip_denoised}")
+        self.model.clip_denoised = False
+        logger.info(f"USE_RNN: {use_rnn}, NUM_SAMPLING_STEPS: {num_sampling_steps}")
+        
+        if use_rnn:
+            sample_latent = output_from_rnn[:, :16]
+        else:
+            if num_sampling_steps >= self.TIMESTEPS:
+                sample_latent = self.model.p_sample_loop(
+                    cond={'c_concat': output_from_rnn}, 
+                    shape=[1, *self.LATENT_DIMS], 
+                    return_intermediates=False, 
+                    verbose=True
+                )
+            else:
+                if num_sampling_steps == 1:
+                    x = torch.randn([1, *self.LATENT_DIMS], device=self.device)
+                    t = torch.full((1,), self.TIMESTEPS-1, device=self.device, dtype=torch.long)
+                    sample_latent = self.model.apply_model(x, t, {'c_concat': output_from_rnn})
+                else:
+                    sampler = DDIMSampler(self.model)
+                    sample_latent, _ = sampler.sample(
+                        S=num_sampling_steps,
+                        conditioning={'c_concat': output_from_rnn},
+                        batch_size=1,
+                        shape=self.LATENT_DIMS,
+                        verbose=False
+                    )
+        timing['unet'] = time.perf_counter() - start
+        
+        # Decoding
+        start = time.perf_counter()
+        sample = sample_latent * self.DATA_NORMALIZATION['std'].view(1, -1, 1, 1) + self.DATA_NORMALIZATION['mean'].view(1, -1, 1, 1)
+        sample = self.model.decode_first_stage(sample)
+        sample = sample.squeeze(0).clamp(-1, 1)
+        timing['decode'] = time.perf_counter() - start
+        
+        # Convert to image
+        sample_img = ((sample[:3].transpose(0,1).transpose(1,2).cpu().float().numpy() + 1) * 127.5).astype(np.uint8)
+        
+        timing['total'] = sum(timing.values())
+        
+        return sample_latent, sample_img, hidden_states, timing
+
+    def initialize_session(self, session_id: str):
+        """Initialize a new session"""
+        self.current_session = session_id
+        self.session_data[session_id] = {
+            'previous_frame': self.padding_image,
+            'hidden_states': None,
+            'keys_down': set(),
+            'frame_num': -1,
+            'client_settings': {
+                'use_rnn': self.USE_RNN,
+                'sampling_steps': self.NUM_SAMPLING_STEPS
+            },
+            'input_queue': asyncio.Queue(),
+            'is_processing': False
+        }
+        logger.info(f"Initialized session {session_id}")
+        
+        # Start processing task for this session
+        asyncio.create_task(self._process_session_queue(session_id))
+
+    def end_session(self, session_id: str):
+        """End a session and clean up"""
+        if session_id in self.session_data:
+            # Clear any remaining items in the queue
+            session = self.session_data[session_id]
+            while not session['input_queue'].empty():
+                try:
+                    session['input_queue'].get_nowait()
+                    session['input_queue'].task_done()
+                except asyncio.QueueEmpty:
+                    break
+            del self.session_data[session_id]
+        if self.current_session == session_id:
+            self.current_session = None
+        logger.info(f"Ended session {session_id}")
+
+    async def _process_session_queue(self, session_id: str):
+        """Process the input queue for a specific session with interesting input filtering"""
+        while session_id in self.session_data:
+            try:
+                session = self.session_data[session_id]
+                input_queue = session['input_queue']
+                
+                # Wait for input to be available
+                if input_queue.empty():
+                    await asyncio.sleep(0.01)  # Small delay to prevent busy waiting
+                    continue
+                
+                # If already processing, skip
+                if session['is_processing']:
+                    await asyncio.sleep(0.01)
+                    continue
+                
+                # Set processing flag
+                session['is_processing'] = True
+                
+                try:
+                    # Process queue with interesting input filtering
+                    await self._process_next_input(session_id)
+                finally:
+                    session['is_processing'] = False
+                    
+            except Exception as e:
+                logger.error(f"Error in session queue processing for {session_id}: {e}")
+                import traceback
+                traceback.print_exc()
+                await asyncio.sleep(1)  # Prevent tight error loop
+        
+        logger.info(f"Session queue processor ended for {session_id}")
+
+    async def _process_next_input(self, session_id: str):
+        """Process next input with interesting input filtering (from main.py logic)"""
+        session = self.session_data[session_id]
+        input_queue = session['input_queue']
+        
+        if input_queue.empty():
+            return
+        
+        queue_size = input_queue.qsize()
+        logger.info(f"Processing next input for session {session_id}. Queue size: {queue_size}")
+        
+        try:
+            # Initialize variables to track progress
+            skipped = 0
+            latest_input = None
+            
+            # Process the queue one item at a time
+            while not input_queue.empty():
+                current_input = await input_queue.get()
+                input_queue.task_done()
+                
+                # Always update the latest input
+                latest_input = current_input
+                
+                # Check if this is an interesting event
+                is_interesting = (current_input.get("is_left_click") or 
+                                  current_input.get("is_right_click") or 
+                                  (current_input.get("keys_down") and len(current_input.get("keys_down")) > 0) or 
+                                  (current_input.get("keys_up") and len(current_input.get("keys_up")) > 0))
+                
+                # Process immediately if interesting
+                if is_interesting:
+                    logger.info(f"Found interesting input for session {session_id} (skipped {skipped} events)")
+                    await self._process_single_input(session_id, current_input)
+                    return
+                
+                # Otherwise, continue to the next item
+                skipped += 1
+                
+                # If this is the last item and no interesting inputs were found
+                if input_queue.empty():
+                    logger.info(f"No interesting inputs for session {session_id}, processing latest movement (skipped {skipped-1} events)")
+                    await self._process_single_input(session_id, latest_input)
+                    return
+                    
+        except Exception as e:
+            logger.error(f"Error in _process_next_input for session {session_id}: {e}")
+            import traceback
+            traceback.print_exc()
+
+    async def process_input(self, session_id: str, data: dict) -> dict:
+        """Process input for a session - adds to queue or handles control messages"""
+        if session_id not in self.session_data:
+            self.initialize_session(session_id)
+        
+        session = self.session_data[session_id]
+        
+        # Handle control messages immediately (don't queue these)
+        if data.get("type") == "reset":
+            logger.info(f"Received reset command for session {session_id}")
+            # Clear the queue
+            while not session['input_queue'].empty():
+                try:
+                    session['input_queue'].get_nowait()
+                    session['input_queue'].task_done()
+                except asyncio.QueueEmpty:
+                    break
+            session['previous_frame'] = self.padding_image
+            session['hidden_states'] = None
+            session['keys_down'] = set()
+            session['frame_num'] = -1
+            return {"type": "reset_confirmed"}
+        
+        elif data.get("type") == "update_sampling_steps":
+            steps = data.get("steps", 32)
+            if steps < 1:
+                return {"type": "error", "message": "Invalid sampling steps value"}
+            session['client_settings']['sampling_steps'] = steps
+            logger.info(f"Updated sampling steps to {steps} for session {session_id}")
+            return {"type": "steps_updated", "steps": steps}
+        
+        elif data.get("type") == "update_use_rnn":
+            use_rnn = data.get("use_rnn", False)
+            session['client_settings']['use_rnn'] = use_rnn
+            logger.info(f"Updated USE_RNN to {use_rnn} for session {session_id}")
+            return {"type": "rnn_updated", "use_rnn": use_rnn}
+        
+        elif data.get("type") == "get_settings":
+            return {
+                "type": "settings",
+                "sampling_steps": session['client_settings']['sampling_steps'],
+                "use_rnn": session['client_settings']['use_rnn']
+            }
+        
+        elif data.get("type") == "heartbeat":
+            return {"type": "heartbeat_response"}
+        
+        # For regular input data, add to queue and return immediately
+        # The actual processing will happen asynchronously in the queue processor
+        await session['input_queue'].put(data)
+        queue_size = session['input_queue'].qsize()
+        logger.info(f"Added input to queue for session {session_id}. Queue size: {queue_size}")
+        
+        # Return a placeholder response - the real response will be sent via WebSocket
+        return {"type": "queued", "queue_size": queue_size}
+
+    async def _process_single_input(self, session_id: str, data: dict):
+        """Process a single input for a session (the actual processing logic)"""
+        session = self.session_data[session_id]
+        
+        # Process regular input
+        try:
+            session['frame_num'] += 1
+            
+            # Extract input data
+            x = max(0, min(data.get("x", 0), self.SCREEN_WIDTH - 1))
+            y = max(0, min(data.get("y", 0), self.SCREEN_HEIGHT - 1))
+            is_left_click = data.get("is_left_click", False)
+            is_right_click = data.get("is_right_click", False)
+            keys_down_list = data.get("keys_down", [])
+            keys_up_list = data.get("keys_up", [])
+            
+            # Update keys_down set
+            for key in keys_down_list:
+                key = key.lower()
+                if key in self.KEYMAPPING:
+                    key = self.KEYMAPPING[key]
+                session['keys_down'].add(key)
+            
+            for key in keys_up_list:
+                key = key.lower()
+                if key in self.KEYMAPPING:
+                    key = self.KEYMAPPING[key]
+                session['keys_down'].discard(key)
+            
+            # Handle debug modes
+            if self.DEBUG_MODE:
+                logger.info("DEBUG MODE, REMOVING HIDDEN STATES")
+                session['previous_frame'] = self.padding_image
+            
+            if self.DEBUG_MODE_2:
+                if session['frame_num'] > self.NUM_MAX_FRAMES-1:
+                    logger.info("DEBUG MODE_2, REMOVING HIDDEN STATES")
+                    session['previous_frame'] = self.padding_image
+                    session['frame_num'] = 0
+            
+            # Prepare model inputs
+            inputs = self.prepare_model_inputs(
+                session['previous_frame'],
+                session['hidden_states'],
+                x, y, is_right_click, is_left_click,
+                list(session['keys_down']),
+                session['frame_num']
+            )
+            
+            # Process frame
+            logger.info(f"Processing frame {session['frame_num']} for session {session_id}")
+            sample_latent, sample_img, hidden_states, timing_info = await self.process_frame(
+                inputs,
+                use_rnn=session['client_settings']['use_rnn'],
+                num_sampling_steps=session['client_settings']['sampling_steps']
+            )
+            
+            # Update session state
+            session['previous_frame'] = sample_latent
+            session['hidden_states'] = hidden_states
+            
+            # Convert image to base64
+            img = Image.fromarray(sample_img)
+            buffered = io.BytesIO()
+            img.save(buffered, format="PNG")
+            img_str = base64.b64encode(buffered.getvalue()).decode()
+            
+            # Log timing
+            logger.info(f"Frame {session['frame_num']} processed in {timing_info['total']:.4f}s (FPS: {1.0/timing_info['total']:.2f})")
+            
+            # Send result back to dispatcher
+            await self._send_result_to_dispatcher(session_id, {"image": img_str})
+            
+        except Exception as e:
+            logger.error(f"Error processing input for session {session_id}: {e}")
+            import traceback
+            traceback.print_exc()
+            await self._send_result_to_dispatcher(session_id, {"type": "error", "message": str(e)})
+
+    async def _send_result_to_dispatcher(self, session_id: str, result: dict):
+        """Send processing result back to dispatcher"""
+        try:
+            async with aiohttp.ClientSession() as client_session:
+                await client_session.post(f"{self.dispatcher_url}/worker_result", json={
+                    "session_id": session_id,
+                    "worker_id": self.worker_id,
+                    "result": result
+                })
+        except Exception as e:
+            logger.error(f"Failed to send result to dispatcher: {e}")
+
+# FastAPI app for the worker
+app = FastAPI()
+
+# Global worker instance
+worker: Optional[GPUWorker] = None
+
+@app.post("/process_input")
+async def process_input_endpoint(request: dict):
+    """Process input from dispatcher"""
+    if not worker:
+        raise HTTPException(status_code=500, detail="Worker not initialized")
+    
+    session_id = request.get("session_id")
+    data = request.get("data")
+    
+    if not session_id or not data:
+        raise HTTPException(status_code=400, detail="Missing session_id or data")
+    
+    result = await worker.process_input(session_id, data)
+    return result
+
+@app.post("/end_session")
+async def end_session_endpoint(request: dict):
+    """End a session"""
+    if not worker:
+        raise HTTPException(status_code=500, detail="Worker not initialized")
+    
+    session_id = request.get("session_id")
+    if not session_id:
+        raise HTTPException(status_code=400, detail="Missing session_id")
+    
+    worker.end_session(session_id)
+    return {"status": "session_ended"}
+
+@app.get("/health")
+async def health_check():
+    """Health check endpoint"""
+    return {
+        "status": "healthy",
+        "worker_id": worker.worker_id if worker else None,
+        "gpu_id": worker.gpu_id if worker else None,
+        "current_session": worker.current_session if worker else None
+    }
+
+async def startup_worker(gpu_id: int, dispatcher_url: str):
+    """Initialize the worker"""
+    global worker
+    worker = GPUWorker(gpu_id, dispatcher_url)
+    
+    # Register with dispatcher
+    await worker.register_with_dispatcher()
+    
+    # Start ping task
+    asyncio.create_task(worker.ping_dispatcher())
+
+if __name__ == "__main__":
+    import uvicorn
+    
+    # Parse command line arguments
+    parser = argparse.ArgumentParser(description="GPU Worker for Neural OS")
+    parser.add_argument("--gpu-id", type=int, required=True, help="GPU ID to use")
+    parser.add_argument("--dispatcher-url", type=str, default="http://localhost:8000", help="Dispatcher URL")
+    args = parser.parse_args()
+    
+    # Calculate port based on GPU ID
+    port = 8001 + args.gpu_id
+    
+    @app.on_event("startup")
+    async def startup_event():
+        await startup_worker(args.gpu_id, args.dispatcher_url)
+    
+    logger.info(f"Starting worker on GPU {args.gpu_id}, port {port}")
+    uvicorn.run(app, host="0.0.0.0", port=port)