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)