main.py

from fastapi import FastAPI, Form, HTTPException
from fastapi.responses import JSONResponse
from pydantic import BaseModel
from PIL import Image
from io import BytesIO
import base64
import torch
import re
import logging
import asyncio
from contextlib import asynccontextmanager

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Initialize global variables
model = None
processor = None
tokenizer = None
model_name = "microsoft/GUI-Actor-2B-Qwen2-VL"
model_loaded = False

async def load_model():
    """Load model with proper error handling"""
    global model, processor, tokenizer, model_loaded
    
    try:
        logger.info("Starting model loading...")
        
        # Import required modules
        from transformers import AutoProcessor, AutoModelForCausalLM
        
        logger.info("Loading processor...")
        # Use AutoProcessor for better compatibility
        processor = AutoProcessor.from_pretrained(
            model_name, 
            trust_remote_code=True
        )
        logger.info("Processor loaded successfully")
        
        tokenizer = processor.tokenizer
        
        logger.info("Loading model...")
        # Use AutoModelForCausalLM for better compatibility
        model = AutoModelForCausalLM.from_pretrained(
            model_name,
            torch_dtype=torch.float32,
            device_map=None,  # CPU only
            trust_remote_code=True,
            low_cpu_mem_usage=True  # For better memory management
        ).eval()
        
        logger.info("Model loaded successfully!")
        model_loaded = True
        return True
        
    except Exception as e:
        logger.error(f"Error loading model: {e}")
        model_loaded = False
        return False

@asynccontextmanager
async def lifespan(app: FastAPI):
    # Startup
    logger.info("Starting up GUI-Actor API...")
    await load_model()
    yield
    # Shutdown
    logger.info("Shutting down GUI-Actor API...")

# Initialize FastAPI app with lifespan
app = FastAPI(
    title="GUI-Actor API", 
    version="1.0.0",
    lifespan=lifespan
)

class Base64Request(BaseModel):
    image_base64: str
    instruction: str

def extract_coordinates(text):
    """
    Extract coordinates from model output text
    """
    # Pattern untuk mencari koordinat dalam berbagai format
    patterns = [
        r'click\s*\(\s*(\d+(?:\.\d+)?)\s*,\s*(\d+(?:\.\d+)?)\s*\)',  # click(x, y)
        r'\[\s*(\d+(?:\.\d+)?)\s*,\s*(\d+(?:\.\d+)?)\s*\]',          # [x, y]
        r'(\d+(?:\.\d+)?)\s*,\s*(\d+(?:\.\d+)?)',                    # x, y
        r'point:\s*\(\s*(\d+(?:\.\d+)?)\s*,\s*(\d+(?:\.\d+)?)\s*\)', # point: (x, y)
    ]
    
    for pattern in patterns:
        matches = re.findall(pattern, text.lower())
        if matches:
            try:
                x, y = float(matches[0][0]), float(matches[0][1])
                # Normalize jika koordinat > 1 (asumsi pixel coordinates)
                if x > 1 or y > 1:
                    # Asumsi resolusi 1920x1080 untuk normalisasi
                    x = x / 1920 if x > 1 else x
                    y = y / 1080 if y > 1 else y
                return [(x, y)]
            except (ValueError, IndexError):
                continue
    
    # Default ke center jika tidak ditemukan
    return [(0.5, 0.5)]

def cpu_inference(conversation, model, tokenizer, processor):
    """
    Inference function untuk CPU
    """
    try:
        # Apply chat template
        text = processor.apply_chat_template(
            conversation, 
            tokenize=False, 
            add_generation_prompt=True
        )
        
        # Get image from conversation
        image = conversation[1]["content"][0]["image"]
        
        # Process inputs
        inputs = processor(
            text=[text], 
            images=[image], 
            return_tensors="pt"
        )
        
        # Generate response
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=256,
                do_sample=True,
                temperature=0.3,
                top_p=0.8,
                pad_token_id=tokenizer.eos_token_id
            )
        
        # Decode response
        generated_ids = outputs[0][inputs["input_ids"].shape[1]:]
        response = tokenizer.decode(generated_ids, skip_special_tokens=True)
        
        # Extract coordinates
        coordinates = extract_coordinates(response)
        
        return {
            "topk_points": coordinates,
            "response": response,
            "success": True
        }
        
    except Exception as e:
        logger.error(f"Inference error: {e}")
        return {
            "topk_points": [(0.5, 0.5)],
            "response": f"Error during inference: {str(e)}",
            "success": False
        }

@app.get("/")
async def root():
    return {
        "message": "GUI-Actor API is running", 
        "status": "healthy",
        "model_loaded": model_loaded
    }

@app.post("/click/base64")
async def predict_click_base64(data: Base64Request):
    if not model_loaded:
        raise HTTPException(
            status_code=503,
            detail="Model not loaded properly"
        )
    
    try:
        # Decode base64 to image
        try:
            # Handle data URL format
            if "," in data.image_base64:
                image_data = base64.b64decode(data.image_base64.split(",")[-1])
            else:
                image_data = base64.b64decode(data.image_base64)
        except Exception as e:
            raise HTTPException(status_code=400, detail=f"Invalid base64 image: {e}")
        
        try:
            pil_image = Image.open(BytesIO(image_data)).convert("RGB")
        except Exception as e:
            raise HTTPException(status_code=400, detail=f"Invalid image format: {e}")

        conversation = [
            {
                "role": "system",
                "content": [
                    {
                        "type": "text",
                        "text": "You are a GUI agent. You are given a task and a screenshot of the screen. You need to perform a series of pyautogui actions to complete the task. Please provide the click coordinates.",
                    }
                ]
            },
            {
                "role": "user",
                "content": [
                    {
                        "type": "image",
                        "image": pil_image,
                    },
                    {
                        "type": "text",
                        "text": data.instruction,
                    },
                ],
            },
        ]

        # Run inference
        pred = cpu_inference(conversation, model, tokenizer, processor)
        px, py = pred["topk_points"][0]
        
        return JSONResponse(content={
            "x": round(px, 4), 
            "y": round(py, 4),
            "response": pred["response"],
            "success": pred["success"]
        })
        
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Prediction error: {e}")
        raise HTTPException(
            status_code=500,
            detail=f"Internal server error: {str(e)}"
        )

@app.get("/health")
async def health_check():
    return {
        "status": "healthy" if model_loaded else "unhealthy",
        "model": model_name,
        "device": "cpu",
        "torch_dtype": "float32",
        "model_loaded": model_loaded
    }

@app.post("/click/form")
async def predict_click_form(
    image_base64: str = Form(...),
    instruction: str = Form(...)
):
    data = Base64Request(image_base64=image_base64, instruction=instruction)
    return await predict_click_base64(data)