scripts/shootout.py

"""VLM Shootout: compare Qwen2.5-VL-3B, SmolVLM, and Gemma 3 4B.

Sends the same image + prompt to each model and measures:
  - Response quality (valid JSON, garment count, attribute completeness)
  - Inference speed (tokens/second)
  - VRAM usage (peak)

Usage:
    python scripts/shootout.py --image resources/sample.jpg
    python scripts/shootout.py --image resources/sample.jpg --model qwen2.5-vl-3b
"""

import argparse
import json
import time
import subprocess
import base64
from pathlib import Path

MODELS_DIR = Path(__file__).parent.parent / "models"

PROMPT = """Analyze this image of clothing items. For EACH visible garment or accessory, return a JSON array.

Each item must have these fields:
- "type": garment type (e.g. "sweater", "shirt", "jeans", "boots", "hat", "bag")
- "color": primary color
- "material": fabric/material if identifiable (e.g. "knit", "denim", "leather"), otherwise "unknown"
- "pattern": pattern type (e.g. "solid", "checkered", "striped"), otherwise "solid"
- "season": most suitable season ("spring", "summer", "autumn", "winter", "all")
- "formality": style level ("casual", "smart-casual", "formal")

Return ONLY a valid JSON array. No markdown fences, no explanation."""

MODEL_CONFIGS = {
    "qwen2.5-vl-3b": {
        "model_file": "Qwen2.5-VL-3B-Instruct.Q4_K_M.gguf",
        "mmproj_file": "Qwen2.5-VL-3B-Instruct.mmproj-fp16.gguf",
        "chat_handler": "qwen25vl",
    },
    "smolvlm-2b": {
        "model_file": "SmolVLM-Instruct-Q4_K_M.gguf",
        "mmproj_file": "mmproj-SmolVLM-Instruct-f16.gguf",
        "chat_handler": "mtmd",
    },
    "gemma-3-4b": {
        "model_file": "gemma-3-4b-it-Q4_K_M.gguf",
        "mmproj_file": "mmproj-model-f16.gguf",
        "chat_handler": "mtmd",
    },
}


def get_vram_usage_mb() -> float:
    """Get current VRAM usage in MB via nvidia-smi."""
    try:
        result = subprocess.run(
            ["nvidia-smi", "--query-gpu=memory.used", "--format=csv,noheader,nounits"],
            capture_output=True, text=True, timeout=5,
        )
        return float(result.stdout.strip())
    except Exception:
        return 0.0


def image_to_data_uri(image_path: str, max_pixels: int = 512) -> str:
    """Resize image to fit within max_pixels on longest side, then convert to base64 data URI."""
    from PIL import Image
    import io

    img = Image.open(image_path)
    original_size = img.size
    img.thumbnail((max_pixels, max_pixels), Image.LANCZOS)
    print(f"  Image resized: {original_size} -> {img.size}")

    buffer = io.BytesIO()
    img.save(buffer, format="JPEG", quality=85)
    b64 = base64.b64encode(buffer.getvalue()).decode("utf-8")
    return f"data:image/jpeg;base64,{b64}"


def load_and_test(model_name: str, config: dict, image_path: str) -> dict:
    """Load a model, run inference, return results."""
    from llama_cpp import Llama
    from llama_cpp.llama_chat_format import Qwen25VLChatHandler, MTMDChatHandler

    model_dir = MODELS_DIR / model_name
    model_path = str(model_dir / config["model_file"])
    mmproj_path = str(model_dir / config["mmproj_file"])

    if not Path(model_path).exists():
        return {"error": f"Model file not found: {model_path}"}
    if not Path(mmproj_path).exists():
        return {"error": f"Mmproj file not found: {mmproj_path}"}

    print(f"\n--- Loading {model_name} ---")
    vram_before = get_vram_usage_mb()

    handler_cls = Qwen25VLChatHandler if config["chat_handler"] == "qwen25vl" else MTMDChatHandler
    chat_handler = handler_cls(clip_model_path=mmproj_path)

    llm = Llama(
        model_path=model_path,
        chat_handler=chat_handler,
        n_gpu_layers=-1,
        n_ctx=4096,
        verbose=False,
    )

    vram_after_load = get_vram_usage_mb()
    print(f"  VRAM: {vram_before:.0f} -> {vram_after_load:.0f} MB (+{vram_after_load - vram_before:.0f} MB)")

    data_uri = image_to_data_uri(image_path)

    messages = [
        {
            "role": "user",
            "content": [
                {"type": "text", "text": PROMPT},
                {"type": "image_url", "image_url": {"url": data_uri}},
            ],
        }
    ]

    print(f"  Running inference...")
    start = time.perf_counter()

    response = llm.create_chat_completion(
        messages=messages,
        max_tokens=2048,
        temperature=0.1,
    )

    elapsed = time.perf_counter() - start
    vram_peak = get_vram_usage_mb()

    raw_text = response["choices"][0]["message"]["content"]
    usage = response.get("usage", {})
    completion_tokens = usage.get("completion_tokens", 0)
    tokens_per_sec = completion_tokens / elapsed if elapsed > 0 else 0

    garments = parse_json_response(raw_text)

    del llm
    del chat_handler
    import gc
    gc.collect()

    return {
        "model": model_name,
        "raw_response": raw_text,
        "garments": garments,
        "garment_count": len(garments) if isinstance(garments, list) else 0,
        "valid_json": isinstance(garments, list),
        "elapsed_sec": round(elapsed, 2),
        "completion_tokens": completion_tokens,
        "tokens_per_sec": round(tokens_per_sec, 1),
        "vram_model_mb": round(vram_after_load - vram_before),
        "vram_peak_mb": round(vram_peak),
    }


def parse_json_response(text: str) -> list | str:
    """Try to extract a JSON array from the model response."""
    cleaned = text.strip()

    if cleaned.startswith("```"):
        lines = cleaned.split("\n")
        lines = lines[1:]  # remove opening fence
        if lines and lines[-1].strip() == "```":
            lines = lines[:-1]
        cleaned = "\n".join(lines).strip()

    try:
        parsed = json.loads(cleaned)
        if isinstance(parsed, list):
            return parsed
        if isinstance(parsed, dict):
            return [parsed]
        return cleaned
    except json.JSONDecodeError:
        start = cleaned.find("[")
        end = cleaned.rfind("]")
        if start != -1 and end != -1 and end > start:
            try:
                return json.loads(cleaned[start:end + 1])
            except json.JSONDecodeError:
                pass
        return cleaned


def print_results(results: list[dict]):
    """Print a comparison table of all results."""
    print(f"\n{'='*80}")
    print("SHOOTOUT RESULTS")
    print(f"{'='*80}")

    for r in results:
        if "error" in r:
            print(f"\n{r['model']}: ERROR - {r['error']}")
            continue

        print(f"\n--- {r['model']} ---")
        print(f"  Valid JSON:    {'YES' if r['valid_json'] else 'NO'}")
        print(f"  Garments:      {r['garment_count']}")
        print(f"  Time:          {r['elapsed_sec']}s")
        print(f"  Tokens/sec:    {r['tokens_per_sec']}")
        print(f"  VRAM (model):  {r['vram_model_mb']} MB")
        print(f"  VRAM (peak):   {r['vram_peak_mb']} MB")

        if r["valid_json"] and r["garments"]:
            print(f"  First garment: {json.dumps(r['garments'][0], indent=4)}")

        if not r["valid_json"]:
            print(f"  Raw response (first 500 chars):")
            print(f"    {r['raw_response'][:500]}")

    print(f"\n{'='*80}")

    results_path = Path(__file__).parent.parent / "data" / "shootout_results.json"
    results_path.parent.mkdir(parents=True, exist_ok=True)
    with open(results_path, "w") as f:
        json.dump(results, f, indent=2, ensure_ascii=False)
    print(f"Results saved to: {results_path}")


def main():
    parser = argparse.ArgumentParser(description="VLM Shootout")
    parser.add_argument("--image", required=True, help="Path to test image")
    parser.add_argument(
        "--model",
        choices=list(MODEL_CONFIGS.keys()) + ["all"],
        default="all",
        help="Which model to test (default: all)",
    )
    args = parser.parse_args()

    if not Path(args.image).exists():
        print(f"Image not found: {args.image}")
        return

    targets = MODEL_CONFIGS if args.model == "all" else {args.model: MODEL_CONFIGS[args.model]}

    results = []
    for name, config in targets.items():
        result = load_and_test(name, config, args.image)
        results.append(result)

    print_results(results)


if __name__ == "__main__":
    main()