app.py

import gradio as gr
import base64
from PIL import Image
import io
import json
import requests

OPENROUTER_URL = "https://openrouter.ai/api/v1/chat/completions"
OPENROUTER_KEY = "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"  # OpenRouter key

def call_vlm_api(img: Image):
    # encode image to bytes
    buf = io.BytesIO()
    img.save(buf, format="JPEG")
    img_bytes = buf.getvalue()
    img_b64 = base64.b64encode(img_bytes).decode("utf-8")

    headers = {"Authorization": f"Bearer {OPENROUTER_KEY}"}
    payload = {
        "model": "qwen/qwen2.5-vl-32b-instruct:free",
        "messages": [
            {
                "role": "user",
                "content": [
                    {"type": "input_text", "text": "Describe the image in detail."},
                    {"type": "input_image", "image_base64": img_b64}
                ]
            }
        ],
        "temperature": 0.2
    }

    resp = requests.post(OPENROUTER_URL, headers=headers, json=payload, timeout=60)
    if resp.status_code == 200:
        try:
            return resp.json()["choices"][0]["message"]["content"][0]["text"]
        except Exception:
            return "Error parsing VLM response"
    else:
        return f"OpenRouter API error: {resp.status_code}"

def process(payload: dict):
    try:
        img_bytes = base64.b64decode(payload["image_b64"])
        img = Image.open(io.BytesIO(img_bytes)).convert("RGB")

        vlm_text = call_vlm_api(img)

        reply = {
            "received": True,
            "robot_id": payload.get("robot_id", "unknown"),
            "size": img.size,
            "vllm_analysis": vlm_text
        }
        return reply
    except Exception as e:
        return {"error": str(e)}

demo = gr.Interface(
    fn=process,
    inputs=gr.JSON(label="Input Payload (Dict format)"),
    outputs=gr.JSON(label="Reply to Jetson"),
    api_name="predict"
)

if __name__ == "__main__":
    demo.launch(mcp_server=True)