173 lines
7.4 KiB
Python
173 lines
7.4 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, 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. Two backends, selected by env:
|
|
|
|
SERVER mode (fast — model loaded ONCE): set LLAMACPP_SERVER_BIN to the `llama-server` binary.
|
|
CLI mode (simple — reloads per frame): set LLAMACPP_MTMD_BIN to the `llama-mtmd-cli` binary.
|
|
(server mode wins if both are set.)
|
|
|
|
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)
|
|
LLAMACPP_SERVER_PORT loopback port for server mode (optional, default 8123)
|
|
|
|
CPU is forced (`-ngl 0 --no-mmproj-offload`): this machine's integrated GPU times out on the vision encoder
|
|
(Metal command-buffer timeout). Server mode keeps the model resident and POSTs base64 frames to the
|
|
OpenAI-compatible /v1/chat/completions endpoint on 127.0.0.1 — far faster than reloading per frame.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import base64
|
|
import json
|
|
import os
|
|
import re
|
|
import subprocess
|
|
import sys
|
|
import time
|
|
import urllib.error
|
|
import urllib.request
|
|
|
|
|
|
def _clean(text: str) -> str:
|
|
if not text:
|
|
return ""
|
|
text = text.replace("\x1b[0m", "")
|
|
# CLI mode echoes the chat-templated conversation; take the LAST assistant turn. (Server mode returns
|
|
# clean content already, in which case there is no 'assistant' marker and the whole string is kept.)
|
|
idx = text.rfind("assistant")
|
|
if idx >= 0:
|
|
text = text[idx + len("assistant"):]
|
|
text = re.sub(r"<\|[^>]*\|>", " ", text).replace("[end of text]", " ")
|
|
lines = [ln.strip() for ln in text.splitlines() if ln.strip()]
|
|
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()
|
|
|
|
|
|
# --- CLI mode (per-frame reload) --------------------------------------------------------------------------
|
|
|
|
def _caption_cli(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",
|
|
"-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)
|
|
|
|
|
|
# --- SERVER mode (model loaded once) ----------------------------------------------------------------------
|
|
|
|
def _start_server(binary: str, model: str, mmproj: str, port: int):
|
|
command = [
|
|
binary, "-m", model, "--mmproj", mmproj, "--host", "127.0.0.1", "--port", str(port),
|
|
"-ngl", "0", "--no-mmproj-offload", "-t", "4", "-c", "4096",
|
|
]
|
|
proc = subprocess.Popen(command, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
|
|
# Wait for readiness (model load on CPU can take a while).
|
|
deadline = time.time() + 300
|
|
while time.time() < deadline:
|
|
if proc.poll() is not None:
|
|
raise RuntimeError("llama-server exited during startup (code %s)" % proc.returncode)
|
|
try:
|
|
with urllib.request.urlopen("http://127.0.0.1:%d/health" % port, timeout=3) as resp:
|
|
if resp.status == 200:
|
|
return proc
|
|
except Exception:
|
|
time.sleep(2)
|
|
proc.terminate()
|
|
raise RuntimeError("llama-server did not become healthy within timeout")
|
|
|
|
|
|
def _caption_server(port: int, image: str, question: str, ntokens: int) -> str:
|
|
with open(image, "rb") as fh:
|
|
b64 = base64.b64encode(fh.read()).decode("ascii")
|
|
payload = {
|
|
"messages": [{"role": "user", "content": [
|
|
{"type": "text", "text": question},
|
|
{"type": "image_url", "image_url": {"url": "data:image/jpeg;base64," + b64}},
|
|
]}],
|
|
"temperature": 0,
|
|
"max_tokens": ntokens,
|
|
}
|
|
req = urllib.request.Request(
|
|
"http://127.0.0.1:%d/v1/chat/completions" % port,
|
|
data=json.dumps(payload).encode("utf-8"),
|
|
headers={"Content-Type": "application/json"},
|
|
)
|
|
with urllib.request.urlopen(req, timeout=900) as resp:
|
|
body = json.loads(resp.read().decode("utf-8"))
|
|
return _clean(body["choices"][0]["message"]["content"])
|
|
|
|
|
|
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()
|
|
|
|
model = os.environ.get("LLAMACPP_VLM_MODEL", "")
|
|
mmproj = os.environ.get("LLAMACPP_VLM_MMPROJ", "")
|
|
ntokens = int(os.environ.get("LLAMACPP_VLM_NTOKENS", "64"))
|
|
server_bin = os.environ.get("LLAMACPP_SERVER_BIN", "")
|
|
cli_bin = os.environ.get("LLAMACPP_MTMD_BIN", "")
|
|
port = int(os.environ.get("LLAMACPP_SERVER_PORT", "8123"))
|
|
|
|
for name, val in (("LLAMACPP_VLM_MODEL", model), ("LLAMACPP_VLM_MMPROJ", mmproj)):
|
|
if not val or not os.path.exists(val):
|
|
sys.stderr.write("vision_caption_llamacpp: missing/invalid %s\n" % name)
|
|
return 2
|
|
use_server = bool(server_bin and os.path.exists(server_bin))
|
|
if not use_server and not (cli_bin and os.path.exists(cli_bin)):
|
|
sys.stderr.write("vision_caption_llamacpp: set LLAMACPP_SERVER_BIN or LLAMACPP_MTMD_BIN to a real path\n")
|
|
return 2
|
|
|
|
with open(args.manifest, "r", encoding="utf-8") as handle:
|
|
items = json.load(handle)
|
|
|
|
server = None
|
|
results = []
|
|
try:
|
|
if use_server:
|
|
server = _start_server(server_bin, model, mmproj, port)
|
|
for item in items:
|
|
answer = ""
|
|
try:
|
|
if use_server:
|
|
answer = _caption_server(port, item["image"], item["question"], ntokens)
|
|
else:
|
|
answer = _caption_cli(cli_bin, 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})
|
|
finally:
|
|
if server is not None:
|
|
server.terminate()
|
|
try:
|
|
server.wait(timeout=10)
|
|
except Exception:
|
|
server.kill()
|
|
|
|
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())
|