rest_vlm.py

"""VLM (Visual Language Model) inference via REST API.

Usage:
    python examples/rest_vlm.py --image examples/images/coco_cat_remote.jpg --prompt "What do you see?"
    python examples/rest_vlm.py --image examples/images/coco_cat_remote.jpg --prompt "List all objects you can identify in this image." --output_mode "detect"

    # With extra generation options
    python examples/rest_vlm.py --image examples/images/coco_cat_remote.jpg \
        --prompt "List every object visible." \
        --max-tokens 256 --temperature 0.3
"""

import argparse
import base64
from pathlib import Path

import requests

BASE_URL = "http://127.0.0.1:8110/v1"
DEFAULT_MODEL = "Qwen/Qwen3-VL-2B-Instruct"


def infer_vlm(
    model: str,
    image_path: Path,
    prompt: str,
    max_tokens: int | None = None,
    temperature: float | None = None,
    output_mode: str | None = None,
) -> dict:
    """POST /v1/infer for a VLM task and return the parsed response."""
    image_b64 = base64.b64encode(image_path.read_bytes()).decode()

    params = {"prompt": prompt}
    if max_tokens is not None:
        params["max_tokens"] = max_tokens
    if temperature is not None:
        params["temperature"] = temperature
    if output_mode is not None:
        params["output_mode"] = output_mode

    payload = {
        "model": model,
        "task": "vlm",
        "image": image_b64,
        "params": params,
    }

    resp = requests.post(f"{BASE_URL}/infer", json=payload, timeout=120)
    resp.raise_for_status()
    return resp.json()


def main(
    image_path: Path,
    model: str,
    prompt: str,
    max_tokens: int | None,
    temperature: float | None,
    output_mode: str,
) -> None:
    print("\n--- VLM Inference ---")
    print(f"  Model  : {model}")
    print(f"  Prompt : {prompt!r}")
    print(f"  Output Mode : {output_mode}")
    try:
        result = infer_vlm(model, image_path, prompt, max_tokens, temperature, output_mode)
    except requests.HTTPError as exc:
        print(f"  ERROR {exc.response.status_code}: {exc.response.text}")
        return

    answer = result.get("text") or "(no text output)"
    print(f"\n  Response:\n{answer}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="REST VLM inference example")
    parser.add_argument("--image", type=Path, required=True, help="Path to an image file")
    parser.add_argument("--prompt", default="Describe this image.", help="Question or instruction")
    parser.add_argument("--model", default=DEFAULT_MODEL, help="VLM model ref")
    parser.add_argument("--max-tokens", type=int, default=None, help="Max tokens to generate")
    parser.add_argument("--temperature", type=float, default=None, help="Sampling temperature")
    parser.add_argument("--output_mode", default="text", help="Output mode of the VLM")
    parser.add_argument("--host", default="127.0.0.1")
    parser.add_argument("--port", type=int, default=8110)
    args = parser.parse_args()

    BASE_URL = f"http://{args.host}:{args.port}/v1"

    if not args.image.exists():
        print(f"ERROR: image not found: {args.image}")
        raise SystemExit(1)

    main(args.image, args.model, args.prompt, args.max_tokens, args.temperature, args.output_mode)