added eval server code

README.md

@@ -41,6 +41,19 @@ Then open the URL (default: http://localhost:7860).
 3. Click **Run Inference**
 4. View the overlay image and predicted trace points
 
+### Eval server
+
+Run a FastAPI server for batch evaluation (e.g. from scripts or the Gradio app):
+
+```bash
+python eval_server.py --model-id mihirgrao/trace-model --port 8001
+```
+
+Endpoints:
+- `POST /predict` – single image + instruction
+- `POST /predict_batch` – batch of `{image_path?|image_base64?, instruction}` samples
+- `GET /health`, `GET /model_info`
+
 ### CLI script
 
 ```bash

app.py

@@ -8,12 +8,14 @@ overlays the trace on the image, and displays the predicted coordinates.
 Model: https://huggingface.co/mihirgrao/trace-model
 """
 
+import base64
 import os
 import tempfile
 import logging
-from typing import Optional, Tuple
+from typing import List, Optional, Tuple
 
 import gradio as gr
+import requests
 import numpy as np
 import torch
 from PIL import Image
@@ -38,6 +40,150 @@ TRACE_FORMAT = (
     "Use normalized coordinates between 0 and 1."
 )
 
+# Global server state (eval server mode)
+_server_state = {"server_url": None, "base_url": "http://localhost"}
+
+
+def discover_available_models(
+    base_url: str = "http://localhost",
+    port_range: Tuple[int, int] = (8000, 8010),
+) -> List[Tuple[str, str]]:
+    """Discover trace eval servers by pinging /health on ports. Returns [(server_url, model_name), ...]."""
+    available = []
+    start_port, end_port = port_range
+    for port in range(start_port, end_port + 1):
+        server_url = f"{base_url.rstrip('/')}:{port}"
+        try:
+            r = requests.get(f"{server_url}/health", timeout=2.0)
+            if r.status_code == 200:
+                try:
+                    info = requests.get(f"{server_url}/model_info", timeout=2.0).json()
+                    name = info.get("model_id", f"Trace @ port {port}")
+                except Exception:
+                    name = f"Trace @ port {port}"
+                available.append((server_url, name))
+        except requests.exceptions.RequestException:
+            continue
+    return available
+
+
+def get_model_info_for_url(server_url: str) -> Optional[str]:
+    """Get formatted model info for a trace eval server."""
+    if not server_url:
+        return None
+    try:
+        r = requests.get(f"{server_url.rstrip('/')}/model_info", timeout=5.0)
+        if r.status_code == 200:
+            return format_trace_model_info(r.json())
+    except Exception as e:
+        logger.warning(f"Could not fetch model info: {e}")
+    return None
+
+
+def format_trace_model_info(info: dict) -> str:
+    """Format trace model info as markdown."""
+    lines = ["## Model Information\n"]
+    lines.append(f"**Model ID:** `{info.get('model_id', 'Unknown')}`\n")
+    if "model_class" in info:
+        lines.append(f"**Model Class:** `{info.get('model_class')}`\n")
+    if "total_parameters" in info:
+        lines.append(f"**Parameters:** {info.get('total_parameters', 0):,}\n")
+    if "error" in info:
+        lines.append(f"**Error:** {info['error']}\n")
+    return "".join(lines)
+
+
+def check_server_health(server_url: str) -> Tuple[str, Optional[dict], Optional[str]]:
+    """Check trace eval server health. Returns (status_msg, health_data, model_info_text)."""
+    if not server_url:
+        return "Please provide a server URL.", None, None
+    try:
+        r = requests.get(f"{server_url.rstrip('/')}/health", timeout=5.0)
+        r.raise_for_status()
+        data = r.json()
+        info = get_model_info_for_url(server_url)
+        _server_state["server_url"] = server_url
+        return f"Server connected: {data.get('status', 'ok')}", data, info
+    except requests.exceptions.RequestException as e:
+        return f"Error connecting to server: {str(e)}", None, None
+
+
+PREPROCESS_SIZE = (128, 128)
+
+
+def run_inference_via_server(
+    image_path: str,
+    instruction: str,
+    server_url: str,
+) -> Tuple[str, Optional[str], str]:
+    """Run inference via trace eval server. Returns (prediction, overlay_path, trace_points_text)."""
+    with open(image_path, "rb") as f:
+        image_b64 = base64.b64encode(f.read()).decode("utf-8")
+    r = requests.post(
+        f"{server_url.rstrip('/')}/predict",
+        json={"image_base64": image_b64, "instruction": instruction},
+        timeout=120.0,
+    )
+    r.raise_for_status()
+    data = r.json()
+    if "error" in data:
+        return data["error"], None, ""
+    prediction = data.get("prediction", "")
+    trajectory = data.get("trajectory", [])
+    trace_points_text = format_trace_points(trajectory)
+
+    overlay_path = None
+    if trajectory and len(trajectory) >= 2:
+        _, preprocessed_path = preprocess_image_for_trace(image_path)
+        try:
+            with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as f:
+                overlay_path = f.name
+            img_arr = visualize_trajectory_on_image(
+                trajectory=trajectory,
+                image_path=preprocessed_path,
+                output_path=overlay_path,
+                normalized=True,
+            )
+            if img_arr is None:
+                visualize_trajectory_on_image(
+                    trajectory=trajectory,
+                    image_path=preprocessed_path,
+                    output_path=overlay_path,
+                    normalized=False,
+                )
+        finally:
+            if os.path.exists(preprocessed_path):
+                try:
+                    os.unlink(preprocessed_path)
+                except Exception:
+                    pass
+    return prediction, overlay_path, trace_points_text
+
+
+def center_crop_resize(
+    image: "Image.Image",
+    size: Tuple[int, int] = PREPROCESS_SIZE,
+) -> "Image.Image":
+    """Center crop to square then resize to size (default 128x128)."""
+    w, h = image.size
+    min_dim = min(w, h)
+    left = (w - min_dim) // 2
+    top = (h - min_dim) // 2
+    cropped = image.crop((left, top, left + min_dim, top + min_dim))
+    return cropped.resize(size, Image.Resampling.LANCZOS)
+
+
+def preprocess_image_for_trace(image_path: str) -> Tuple["Image.Image", Optional[str]]:
+    """
+    Load image, center crop and resize to 128x128.
+    Returns (preprocessed PIL Image, path to temp file for downstream use).
+    """
+    img = Image.open(image_path).convert("RGB")
+    img = center_crop_resize(img, PREPROCESS_SIZE)
+    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
+    img.save(tmp.name)
+    return img, tmp.name
+
 
 def build_prompt(instruction: str = "") -> str:
     """Build the full prompt from task instruction + trace format."""
@@ -111,12 +257,11 @@ def run_inference(image_path: str, prompt: str, model_id: str) -> Tuple[str, Opt
     if image_path is None or not os.path.exists(image_path):
         return "Please provide a valid image.", None, ""
 
+    preprocessed_path = None
     try:
-        # Ensure file:// format for qwen_vl_utils
-        if not image_path.startswith("file://") and not image_path.startswith("http"):
-            image_uri = f"file://{os.path.abspath(image_path)}"
-        else:
-            image_uri = image_path
+        # Preprocess: center crop and resize to 128x128
+        _, preprocessed_path = preprocess_image_for_trace(image_path)
+        image_uri = f"file://{os.path.abspath(preprocessed_path)}"
 
         messages = [
             {
@@ -201,18 +346,17 @@ def run_inference(image_path: str, prompt: str, model_id: str) -> Tuple[str, Opt
         if trajectories and len(trajectories) >= 2:
             with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as f:
                 overlay_path = f.name
-            # Try normalized first (common for VLMs)
+            # Overlay on preprocessed (128x128) image
             img_arr = visualize_trajectory_on_image(
                 trajectory=trajectories,
-                image_path=image_path,
+                image_path=preprocessed_path,
                 output_path=overlay_path,
                 normalized=True,
             )
             if img_arr is None:
-                # Fallback: pixel coordinates
                 visualize_trajectory_on_image(
                     trajectory=trajectories,
-                    image_path=image_path,
+                    image_path=preprocessed_path,
                     output_path=overlay_path,
                     normalized=False,
                 )
@@ -222,6 +366,12 @@ def run_inference(image_path: str, prompt: str, model_id: str) -> Tuple[str, Opt
     except Exception as e:
         logger.exception("Inference failed")
         return f"Error: {str(e)}", None, ""
+    finally:
+        if preprocessed_path and os.path.exists(preprocessed_path):
+            try:
+                os.unlink(preprocessed_path)
+            except Exception:
+                pass
 
 
 def format_trace_points(trajectories) -> str:
@@ -251,12 +401,109 @@ with demo:
         """
         # Trace Model Visualizer
 
-        Upload an image to predict the trajectory/trace using [mihirgrao/trace-model](https://huggingface.co/mihirgrao/trace-model).
+        Upload an image and provide a natural language task instruction to predict the trajectory/trace using [mihirgrao/trace-model](https://huggingface.co/mihirgrao/trace-model).
 
-        The model predicts coordinate points; they are overlaid on the image (green → red gradient) and listed below.
+        The model predicts coordinate points from your instruction; they are overlaid on the image (green → red gradient) and listed below.
         """
     )
 
+    server_url_state = gr.State(value=None)
+    model_url_mapping_state = gr.State(value={})
+
+    def discover_and_select_models(base_url: str):
+        if not base_url:
+            return (
+                gr.update(choices=[], value=None),
+                gr.update(value="Please provide a base URL", visible=True),
+                gr.update(value="", visible=True),
+                None,
+                {},
+            )
+        _server_state["base_url"] = base_url
+        models = discover_available_models(base_url, port_range=(8000, 8010))
+        if not models:
+            return (
+                gr.update(choices=[], value=None),
+                gr.update(
+                    value="❌ No trace eval servers found on ports 8000-8010.",
+                    visible=True,
+                ),
+                gr.update(value="", visible=True),
+                None,
+                {},
+            )
+        choices = []
+        url_map = {}
+        for url, name in models:
+            choices.append(name)
+            url_map[name] = url
+        selected = choices[0] if choices else None
+        selected_url = url_map.get(selected) if selected else None
+        model_info_text = get_model_info_for_url(selected_url) if selected_url else ""
+        status = f"✅ Found {len(models)} server(s). Auto-selected first."
+        _server_state["server_url"] = selected_url
+        return (
+            gr.update(choices=choices, value=selected),
+            gr.update(value=status, visible=True),
+            gr.update(value=model_info_text, visible=True),
+            selected_url,
+            url_map,
+        )
+
+    def on_model_selected(model_choice: str, url_mapping: dict):
+        if not model_choice:
+            return gr.update(value="No model selected", visible=True), gr.update(value="", visible=True), None
+        server_url = url_mapping.get(model_choice) if url_mapping else None
+        if not server_url:
+            return (
+                gr.update(value="Could not find server URL. Please rediscover.", visible=True),
+                gr.update(value="", visible=True),
+                None,
+            )
+        model_info_text = get_model_info_for_url(server_url) or ""
+        status, _, _ = check_server_health(server_url)
+        _server_state["server_url"] = server_url
+        return gr.update(value=status, visible=True), gr.update(value=model_info_text, visible=True), server_url
+
+    with gr.Sidebar():
+        gr.Markdown("### 🔧 Model Configuration")
+
+        base_url_input = gr.Textbox(
+            label="Base Server URL",
+            placeholder="http://localhost",
+            value="http://localhost",
+            interactive=True,
+        )
+        discover_btn = gr.Button("🔍 Discover Eval Servers", variant="primary", size="lg")
+        model_dropdown = gr.Dropdown(
+            label="Select Eval Server",
+            choices=[],
+            value=None,
+            interactive=True,
+            info="Discover trace eval servers on ports 8000-8010",
+        )
+        server_status = gr.Markdown("Click 'Discover Eval Servers' or use local model below")
+        gr.Markdown("---")
+        gr.Markdown("### 📋 Model Information")
+        model_info_display = gr.Markdown("")
+
+        discover_btn.click(
+            fn=discover_and_select_models,
+            inputs=[base_url_input],
+            outputs=[
+                model_dropdown,
+                server_status,
+                model_info_display,
+                server_url_state,
+                model_url_mapping_state,
+            ],
+        )
+        model_dropdown.change(
+            fn=on_model_selected,
+            inputs=[model_dropdown, model_url_mapping_state],
+            outputs=[server_status, model_info_display, server_url_state],
+        )
+
     with gr.Row():
         with gr.Column(scale=1):
             image_input = gr.Image(
@@ -265,21 +512,26 @@ with demo:
                 height=400,
             )
             instruction_input = gr.Textbox(
-                label="Task / Language instruction",
-                placeholder="e.g. Pick up the red block and place it on the table",
+                label="Natural language instruction",
+                placeholder="e.g. Pick up the red block and place it on the table. Stack the cube on top of the block.",
                 value="",
-                lines=2,
-                info="Describe the task. The model will predict the trace for this instruction.",
+                lines=4,
+                info="Enter a task description in natural language. The model predicts the trace for this instruction.",
             )
+            gr.Markdown("### Local model (if no eval server selected)")
             model_id_input = gr.Textbox(
                 label="Model ID",
                 value=DEFAULT_MODEL_ID,
-                info="Hugging Face model ID",
+                info="Hugging Face model ID (used when no eval server is selected)",
             )
             load_model_btn = gr.Button("Load Model", variant="secondary")
             run_btn = gr.Button("Run Inference", variant="primary")
 
         with gr.Column(scale=1):
+            prompt_display = gr.Markdown(
+                f"**Prompt sent to model:**\n\n```\n{build_prompt('')}\n```",
+                label="Model prompt",
+            )
             overlay_output = gr.Image(
                 label="Image with Trace Overlay",
                 height=400,
@@ -292,19 +544,35 @@ with demo:
                 label="Extracted Trace Points",
             )
 
-    status_md = gr.Markdown("Click 'Load Model' to load the trace model, then 'Run Inference' on an image.")
+    status_md = gr.Markdown(
+        "Select an eval server from the sidebar, or load a local model and run inference."
+    )
 
     def on_load_model(model_id: str):
         _, msg = load_model(model_id)
         return f"**Status:** {msg}"
 
-    def on_run_inference(image_path, instruction, model_id):
+    def on_run_inference(image_path, instruction, model_id, server_url):
         if image_path is None:
-            return "Please upload an image first.", None, "", "**Status:** Please upload an image."
+            return (
+                "",
+                "Please upload an image first.",
+                None,
+                "",
+                "**Status:** Please upload an image.",
+            )
         prompt = build_prompt(instruction)
-        pred, overlay_path, trace_text = run_inference(image_path, prompt, model_id)
+        prompt_md = f"**Prompt sent to model:**\n\n```\n{prompt}\n```"
+
+        if server_url:
+            pred, overlay_path, trace_text = run_inference_via_server(
+                image_path, instruction, server_url
+            )
+        else:
+            pred, overlay_path, trace_text = run_inference(image_path, prompt, model_id)
+
         status = "**Status:** Inference complete." if overlay_path else f"**Status:** {pred}"
-        return pred, overlay_path, trace_text, status
+        return prompt_md, pred, overlay_path, trace_text, status
 
     load_model_btn.click(
         fn=on_load_model,
@@ -312,10 +580,26 @@ with demo:
         outputs=[status_md],
     )
 
+    def update_prompt_display(instruction: str):
+        prompt = build_prompt(instruction)
+        return f"**Prompt sent to model:**\n\n```\n{prompt}\n```"
+
+    instruction_input.change(
+        fn=update_prompt_display,
+        inputs=[instruction_input],
+        outputs=[prompt_display],
+    )
+
     run_btn.click(
         fn=on_run_inference,
-        inputs=[image_input, instruction_input, model_id_input],
+        inputs=[
+            image_input,
+            instruction_input,
+            model_id_input,
+            server_url_state,
+        ],
         outputs=[
+            prompt_display,
             prediction_output,
             overlay_output,
             trace_points_output,

eval_server.py

@@ -0,0 +1,280 @@
+#!/usr/bin/env python3
+"""
+FastAPI server for Trace Model inference.
+
+Usage:
+    python eval_server.py --model-id mihirgrao/trace-model --port 8001
+
+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 logging
+import os
+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 app import DEFAULT_MODEL_ID, build_prompt, load_model, run_inference
+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 = "",
+    ) -> 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:
+                image_bytes = base64.b64decode(image_base64)
+                with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as f:
+                    f.write(image_bytes)
+                    image_path = f.name
+                    temp_file_path = image_path
+            except Exception as e:
+                return {"error": f"Invalid base64 image: {e}"}
+
+        try:
+            prompt = build_prompt(instruction)
+            prediction, overlay_path, _ = 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)
+        return {
+            "prediction": prediction,
+            "trajectory": trajectory,
+        }
+
+    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", ""),
+            )
+            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:
+            from app import _model_state
+
+            model = _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
+        """
+        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", ""),
+        )
+
+    @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()

predict_trace.py

@@ -37,7 +37,7 @@ def main():
         "--instruction",
         type=str,
         default="",
-        help="Task / language instruction (e.g. 'Pick up the red block')",
+        help="Natural language task instruction (e.g. 'Pick up the red block and place it on the table')",
     )
     parser.add_argument(
         "-p",

requirements.txt

@@ -1,7 +1,10 @@
+fastapi>=0.100.0
 gradio>=4.0.0
+uvicorn>=0.22.0
 torch>=2.0.0
 transformers>=4.45.0
 accelerate>=0.25.0
 Pillow>=9.0.0
 numpy>=1.20.0
+requests>=2.28.0
 qwen-vl-utils>=0.0.8