forked from jsl/video_editing_poc
perf(vision): add resident llama-server backend (model loads once)
The CLI backend reloads the ~3GB model per frame (~2min/frame). Add a SERVER mode: when LLAMACPP_SERVER_BIN is set, start llama-server once, POST base64 frames to its OpenAI /v1/chat/completions endpoint on loopback, stop it at the end. Same answers, ~3x faster on a full clip (measured: 3 frames 181s incl. one-time load vs ~6-8min). CLI mode (LLAMACPP_MTMD_BIN) remains the simple fallback. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
parent
1ea6640da0
commit
8c0ddc1796
|
|
@ -3,91 +3,164 @@
|
|||
|
||||
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.
|
||||
(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. 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)
|
||||
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)
|
||||
|
||||
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.
|
||||
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:
|
||||
# The CLI echoes the chat-templated conversation; the real answer is the LAST assistant turn.
|
||||
if not text:
|
||||
return ""
|
||||
text = text.replace("\x1b[0m", "")
|
||||
marker = "assistant"
|
||||
idx = text.rfind(marker)
|
||||
# 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(marker):]
|
||||
text = re.sub(r"<\|[^>]*\|>", " ", text) # strip chat control tokens like <|im_end|>
|
||||
text = text.replace("[end of text]", " ")
|
||||
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()]
|
||||
# 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:
|
||||
# --- 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",
|
||||
# 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:]))
|
||||
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()
|
||||
|
||||
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))
|
||||
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:
|
||||
answer = _caption(binary, model, mmproj, item["image"], item["question"], ntokens)
|
||||
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:
|
||||
|
|
|
|||
Loading…
Reference in New Issue