worker.py

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)