100 lines
4.5 KiB
Python
100 lines
4.5 KiB
Python
#!/usr/bin/env python3
|
|
"""Tier-2 vision captioner backed by llama.cpp (GGUF) instead of PyTorch/transformers.
|
|
|
|
Drop-in replacement for tools/vision_caption.py: SAME manifest/output contract, so the Java
|
|
HighlightVisionDirector is unchanged — only the configured script path differs. Uses a stronger local VLM
|
|
(e.g. Qwen2.5-VL-3B-Instruct) through llama.cpp's multimodal CLI (`llama-mtmd-cli --mmproj`), which runs on
|
|
this Intel x86 CPU via AVX SIMD and sidesteps the torch==2.2.2 / transformers 4.x version trap entirely.
|
|
|
|
Manifest JSON: [{"id": "...", "image": "/abs/path.jpg", "question": "..."}]
|
|
Output JSON: [{"id": "...", "answer": "..."}]
|
|
|
|
Fully offline. Paths come from env vars (so the Java side needs no backend-specific args):
|
|
LLAMACPP_MTMD_BIN path to the llama.cpp `llama-mtmd-cli` binary (required)
|
|
LLAMACPP_VLM_MODEL path to the VLM GGUF weights, e.g. Qwen2.5-VL-3B-Instruct-Q4_K_M.gguf (required)
|
|
LLAMACPP_VLM_MMPROJ path to the matching multimodal projector GGUF (mmproj-*.gguf) (required)
|
|
LLAMACPP_VLM_NTOKENS max tokens to generate per answer (optional, default 64)
|
|
|
|
Per-frame invocation keeps the worker simple and serverless (no loopback). The model is memory-mapped, so
|
|
after the first call the OS page cache keeps repeated loads cheap. For heavy batches a persistent
|
|
`llama-server` backend would be faster — a documented future optimisation, not required here.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import re
|
|
import subprocess
|
|
import sys
|
|
|
|
|
|
def _clean(text: str) -> str:
|
|
# The CLI echoes the chat-templated conversation; the real answer is the LAST assistant turn.
|
|
text = text.replace("\x1b[0m", "")
|
|
marker = "assistant"
|
|
idx = text.rfind(marker)
|
|
if idx >= 0:
|
|
text = text[idx + len(marker):]
|
|
text = re.sub(r"<\|[^>]*\|>", " ", text) # strip chat control tokens like <|im_end|>
|
|
text = text.replace("[end of text]", " ")
|
|
lines = [ln.strip() for ln in text.splitlines() if ln.strip()]
|
|
# Drop any leaked log/timing lines; keep the natural-language answer.
|
|
lines = [ln for ln in lines
|
|
if not re.match(r"^(llama_|main:|mtmd_|clip_|ggml_|build:|load|encoding|decoding|\d[\d.:]*\s)",
|
|
ln, re.IGNORECASE)]
|
|
return " ".join(lines).strip()
|
|
|
|
|
|
def _caption(binary: str, model: str, mmproj: str, image: str, question: str, ntokens: int) -> str:
|
|
command = [
|
|
binary, "-m", model, "--mmproj", mmproj, "--image", image,
|
|
"-p", question, "--temp", "0", "-n", str(ntokens), "-t", "4",
|
|
# Force CPU: this machine's integrated GPU times out (Metal command-buffer) on the vision encoder.
|
|
"-ngl", "0", "--no-mmproj-offload",
|
|
]
|
|
proc = subprocess.run(command, capture_output=True, text=True, timeout=900)
|
|
if proc.returncode != 0:
|
|
sys.stderr.write("vision_caption_llamacpp: cli failed (%s): %s\n" % (proc.returncode, proc.stderr[-500:]))
|
|
return ""
|
|
return _clean(proc.stdout)
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(description="Local VLM frame captioner (llama.cpp / GGUF)")
|
|
parser.add_argument("--manifest", required=True)
|
|
parser.add_argument("--output", required=True)
|
|
args = parser.parse_args()
|
|
|
|
binary = os.environ.get("LLAMACPP_MTMD_BIN", "")
|
|
model = os.environ.get("LLAMACPP_VLM_MODEL", "")
|
|
mmproj = os.environ.get("LLAMACPP_VLM_MMPROJ", "")
|
|
ntokens = int(os.environ.get("LLAMACPP_VLM_NTOKENS", "64"))
|
|
missing = [name for name, val in
|
|
(("LLAMACPP_MTMD_BIN", binary), ("LLAMACPP_VLM_MODEL", model), ("LLAMACPP_VLM_MMPROJ", mmproj))
|
|
if not val or not os.path.exists(val)]
|
|
if missing:
|
|
sys.stderr.write("vision_caption_llamacpp: missing/invalid env paths: %s\n" % ", ".join(missing))
|
|
return 2
|
|
|
|
with open(args.manifest, "r", encoding="utf-8") as handle:
|
|
items = json.load(handle)
|
|
|
|
results = []
|
|
for item in items:
|
|
answer = ""
|
|
try:
|
|
answer = _caption(binary, model, mmproj, item["image"], item["question"], ntokens)
|
|
except Exception as exc: # a single bad frame must never fail the batch
|
|
sys.stderr.write("vision_caption_llamacpp: frame %s failed: %s\n" % (item.get("id"), exc))
|
|
results.append({"id": item.get("id"), "answer": answer})
|
|
|
|
os.makedirs(os.path.dirname(os.path.abspath(args.output)), exist_ok=True)
|
|
with open(args.output, "w", encoding="utf-8") as handle:
|
|
json.dump(results, handle, ensure_ascii=False, indent=2)
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|