#!/usr/bin/env python3 """ FastAPI server for Trace Model inference. Usage: python eval_server.py --model-id mihirgrao/trace-model --port 8000 Endpoints: POST /predict - Single image + instruction POST /predict_batch - Batch of (image, instruction) pairs GET /health - Health check GET /model_info - Model information """ import argparse import base64 import io import logging import os import re import tempfile import time from concurrent.futures import ThreadPoolExecutor from threading import Lock from typing import Any, Dict, List, Optional import uvicorn from fastapi import FastAPI, Request from fastapi.middleware.cors import CORSMiddleware from trace_inference import ( DEFAULT_MODEL_ID, build_prompt, load_model, run_inference, ) from trace_inference import _model_state as _trace_model_state from trajectory_viz import extract_trajectory_from_text logger = logging.getLogger(__name__) # --- Trace Eval Server --- class TraceEvalServer: """Inference server for the trace model.""" def __init__( self, model_id: str = DEFAULT_MODEL_ID, max_workers: int = 1, ): self.model_id = model_id self.max_workers = max_workers self._job_counter = 0 self._completed_jobs = 0 self._lock = Lock() self.executor = ThreadPoolExecutor(max_workers=max_workers) logger.info(f"Loading trace model: {model_id}") success, msg = load_model(model_id) if not success: raise RuntimeError(f"Failed to load model: {msg}") logger.info(msg) def predict_one( self, image_path: Optional[str] = None, image_base64: Optional[str] = None, instruction: str = "", is_oxe: bool = False, ) -> Dict[str, Any]: """ Run inference on a single image. Provide either image_path (file path) or image_base64 (base64-encoded image). """ if image_path is None and image_base64 is None: return {"error": "Provide image_path or image_base64"} temp_file_path = None if image_path is None: try: # Strip data URL prefix if present (e.g. "data:image/png;base64,") b64_str = image_base64.strip() if b64_str.startswith("data:"): match = re.match(r"data:image/[^;]+;base64,(.+)", b64_str, re.DOTALL) if match: b64_str = match.group(1) image_bytes = base64.b64decode(b64_str, validate=False) # Load via BytesIO to validate and get proper format, then save from PIL import Image img = Image.open(io.BytesIO(image_bytes)).convert("RGB") with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as f: img.save(f.name, format="PNG") image_path = f.name temp_file_path = image_path except Exception as e: return {"error": f"Invalid image data: {e}"} try: prompt = build_prompt(instruction, is_oxe=is_oxe) prediction, _, _ = run_inference(image_path, prompt, self.model_id) finally: if temp_file_path and os.path.exists(temp_file_path): try: os.unlink(temp_file_path) except Exception: pass if prediction.startswith("Error:") or prediction.startswith("Please "): return {"error": prediction} trajectory = extract_trajectory_from_text(prediction) result: Dict[str, Any] = { "prediction": prediction, "trajectory": trajectory, } return result def predict_batch( self, samples: List[Dict[str, Any]], ) -> Dict[str, Any]: """Process a batch of (image_path or image_base64, instruction) samples.""" results = [] for sample in samples: with self._lock: self._job_counter += 1 job_id = self._job_counter start = time.time() result = self.predict_one( image_path=sample.get("image_path"), image_base64=sample.get("image_base64"), instruction=sample.get("instruction", ""), is_oxe=sample.get("is_oxe", False), ) elapsed = time.time() - start with self._lock: self._completed_jobs += 1 logger.debug(f"[job {job_id}] completed in {elapsed:.3f}s") results.append(result) return {"results": results} def get_status(self) -> Dict[str, Any]: """Get server status.""" return { "model_id": self.model_id, "max_workers": self.max_workers, "completed_jobs": self._completed_jobs, "job_counter": self._job_counter, } def get_model_info(self) -> Dict[str, Any]: """Get model information.""" try: model = _trace_model_state.get("model") if model is None: return {"model_id": self.model_id, "status": "not_loaded"} all_params = sum(p.numel() for p in model.parameters()) return { "model_id": self.model_id, "model_class": model.__class__.__name__, "total_parameters": all_params, } except Exception as e: return {"model_id": self.model_id, "error": str(e)} def shutdown(self): """Shutdown the executor.""" self.executor.shutdown(wait=True) def create_app( model_id: str = DEFAULT_MODEL_ID, max_workers: int = 1, server: Optional[TraceEvalServer] = None, ) -> FastAPI: app = FastAPI(title="Trace Model Evaluation Server") app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) trace_server = server or TraceEvalServer(model_id=model_id, max_workers=max_workers) @app.post("/predict") async def predict(request: Request) -> Dict[str, Any]: """ Predict trace for a single image. JSON body: - image_path: (optional) path to image file - image_base64: (optional) base64-encoded image - instruction: natural language task description - is_oxe: (optional) if true, use OXE prompt format """ body = await request.json() return trace_server.predict_one( image_path=body.get("image_path"), image_base64=body.get("image_base64"), instruction=body.get("instruction", ""), is_oxe=body.get("is_oxe", False), ) @app.post("/predict_batch") async def predict_batch(request: Request) -> Dict[str, Any]: """ Predict trace for a batch of images. JSON body: - samples: list of {image_path?, image_base64?, instruction} """ body = await request.json() samples = body.get("samples", []) if not samples: return {"error": "samples list is required", "results": []} return trace_server.predict_batch(samples) @app.post("/evaluate_batch") async def evaluate_batch(request: Request) -> Dict[str, Any]: """ Alias for /predict_batch for compatibility with RFM-style clients. Accepts same format as /predict_batch. """ return await predict_batch(request) @app.get("/health") def health() -> Dict[str, Any]: """Health check.""" status = trace_server.get_status() return { "status": "healthy", "model_id": status["model_id"], } @app.get("/model_info") def model_info() -> Dict[str, Any]: """Get model information.""" return trace_server.get_model_info() @app.get("/gpu_status") def gpu_status() -> Dict[str, Any]: """Get server status (RFM-compatible endpoint name).""" return trace_server.get_status() @app.on_event("shutdown") async def shutdown_event(): trace_server.shutdown() return app def main(): parser = argparse.ArgumentParser(description="Trace Model Evaluation Server") parser.add_argument( "--model-id", type=str, default=DEFAULT_MODEL_ID, help=f"Model ID (default: {DEFAULT_MODEL_ID})", ) parser.add_argument( "--host", type=str, default="0.0.0.0", help="Server host", ) parser.add_argument( "--port", type=int, default=8001, help="Server port", ) parser.add_argument( "--max-workers", type=int, default=1, help="Max worker threads for batch processing", ) args = parser.parse_args() logging.basicConfig(level=logging.INFO) app = create_app(model_id=args.model_id, max_workers=args.max_workers) print(f"Trace eval server starting on {args.host}:{args.port}") print(f"Model: {args.model_id}") uvicorn.run(app, host=args.host, port=args.port) if __name__ == "__main__": main()