neiam /dms

#!/usr/bin/env python3
"""HTTP wrapper around the rknn-llm multimodal_demo binary.

Wire contract documented in README.md. Each /describe request:
  1. saves the uploaded image to a tmp file
  2. spawns ./demo with that image; demo loads vision encoder + LLM, encodes
     the image once, then enters an interactive REPL on stdin
  3. we feed it `<image><prompt>\\n` and `exit\\n`, capture stdout, and
     extract the text after the "robot: " marker

This reloads the model per request -- fine for first-cut validation, slow
for production. Replace with a long-lived ctypes wrapper around
librkllmrt.so once the integration is proven end-to-end.
"""
from __future__ import annotations

import argparse
import os
import re
import shutil
import subprocess
import tempfile
import time
from pathlib import Path

import yaml
from fastapi import FastAPI, File, Form, HTTPException, UploadFile
from fastapi.responses import JSONResponse
from PIL import Image
from PIL.Image import UnidentifiedImageError
import uvicorn

# Optional codecs — register if present so HEIC/AVIF uploads work.
try:
    import pillow_heif  # type: ignore
    pillow_heif.register_heif_opener()
    if hasattr(pillow_heif, "register_avif_opener"):
        pillow_heif.register_avif_opener()
except Exception:
    pass


def load_models(path: str) -> dict:
    with open(path) as f:
        cfg = yaml.safe_load(f) or {}
    if not cfg:
        raise SystemExit(f"no models defined in {path}")
    for name, m in cfg.items():
        for key in ("vision", "llm"):
            if not Path(m[key]).exists():
                raise SystemExit(f"model {name}: {key} missing: {m[key]}")
    return cfg


# multimodal_demo prints `robot: <text>\n` after each prompt; we want the text
# from the *last* `robot:` (after our own prompt, not the startup banner).
ROBOT_RE = re.compile(r"robot:\s*(.*?)(?=\nuser:|\Z)", re.DOTALL)


def run_demo(demo_bin: str, demo_dir: str, image_path: str, model_cfg: dict,
             prompt: str, max_new_tokens: int, max_context_len: int,
             timeout: int) -> tuple[str, str]:
    abs_demo = os.path.abspath(demo_bin)
    cmd = [
        abs_demo,
        image_path,
        model_cfg["vision"],
        model_cfg["llm"],
        str(max_new_tokens),
        str(max_context_len),
        str(model_cfg.get("rknn_core_num", 3)),
        model_cfg.get("img_start", "<|vision_start|>"),
        model_cfg.get("img_end", "<|vision_end|>"),
        model_cfg.get("img_content", "<|image_pad|>"),
    ]
    env = dict(os.environ)
    # cwd is changed before exec, so use absolute lib path or the dynamic
    # linker can't find librkllmrt.so / librknnrt.so.
    lib_dir = str(Path(abs_demo).parent / "lib")
    env["LD_LIBRARY_PATH"] = lib_dir + ":" + env.get("LD_LIBRARY_PATH", "")

    # Demo's REPL: send `<image>{prompt}\n`, then `exit\n` to break the loop.
    # The pre-canned numeric shortcuts (`0`, `1`) are ignored because our
    # input is plain text starting with `<image>`.
    stdin_payload = f"<image>{prompt}\nexit\n"
    proc = subprocess.run(
        cmd,
        input=stdin_payload,
        capture_output=True,
        text=True,
        cwd=demo_dir,
        env=env,
        timeout=timeout,
    )
    return proc.stdout, proc.stderr


def extract_robot_reply(stdout: str) -> str:
    # The interactive prompt prints "user: " before each line, then "robot: "
    # followed by streamed tokens, then a newline at finish. Take the LAST
    # robot block to ignore any startup chatter.
    matches = ROBOT_RE.findall(stdout)
    if not matches:
        return ""
    return matches[-1].strip()


def make_app(models: dict, demo_bin: str, default_prompt: str) -> FastAPI:
    app = FastAPI()
    default_model = next(iter(models))
    demo_dir = str(Path(demo_bin).parent)

    @app.get("/healthz")
    def healthz():
        return {"ok": True, "models": list(models.keys()), "default": default_model}

    @app.post("/describe")
    async def describe(
        image: UploadFile = File(...),
        prompt: str = Form(default_prompt),
        model: str = Form(default_model),
        max_new_tokens: int = Form(256),
    ):
        if model not in models:
            raise HTTPException(404, f"unknown model {model!r}; have {list(models)}")
        m = models[model]
        max_context_len = int(m.get("max_context_len", 4096))

        with tempfile.TemporaryDirectory() as td:
            # Demo binary uses OpenCV 3.4.5, which has a narrow codec set
            # (no HEIC/AVIF, flaky WebP). Decode via Pillow first and
            # re-encode as plain JPEG so the demo only ever sees a format
            # it can handle.
            img_path = Path(td) / "input.jpg"
            try:
                with Image.open(image.file) as im:
                    im.load()
                    if im.mode not in ("RGB", "L"):
                        im = im.convert("RGB")
                    im.save(img_path, format="JPEG", quality=92)
            except (UnidentifiedImageError, OSError) as e:
                return JSONResponse(
                    {"error": f"could not decode image: {e}"},
                    status_code=400,
                )

            t0 = time.monotonic()
            try:
                stdout, stderr = run_demo(
                    demo_bin, demo_dir, str(img_path), m,
                    prompt, max_new_tokens, max_context_len,
                    timeout=int(os.environ.get("RKNN_DEMO_TIMEOUT", "300")),
                )
            except subprocess.TimeoutExpired:
                return JSONResponse({"error": "demo timeout"}, status_code=504)
            elapsed = time.monotonic() - t0

        text = extract_robot_reply(stdout)
        if not text:
            return JSONResponse(
                {"error": "no reply extracted",
                 "stdout_tail": stdout[-2000:],
                 "stderr_tail": stderr[-2000:]},
                status_code=500,
            )
        return {
            "description": text,
            "model": model,
            "elapsed_seconds": round(elapsed, 2),
        }

    return app


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--models-config", default="models.yaml")
    ap.add_argument("--demo-bin", required=True,
                    help="path to the built rknn-llm multimodal demo binary")
    ap.add_argument("--bind", default="0.0.0.0:8080")
    ap.add_argument("--default-prompt", default=(
        "Describe this image in one or two concise sentences. "
        "List the salient objects, setting, and any visible text."))
    args = ap.parse_args()

    models = load_models(args.models_config)
    app = make_app(models, args.demo_bin, args.default_prompt)
    host, port = args.bind.split(":")
    uvicorn.run(app, host=host, port=int(port))


if __name__ == "__main__":
    main()