forked from jsl/video_editing_poc
71 lines
2.7 KiB
Python
71 lines
2.7 KiB
Python
#!/usr/bin/env python3
|
|
"""Tier-2 vision captioner for the highlight director.
|
|
|
|
Loads a local vision-language model (moondream2) ONCE and answers a question for each frame listed in a
|
|
manifest, writing the answers to a JSON file. Runs fully offline from the local Hugging Face cache — it
|
|
performs no network access (HF_HUB_OFFLINE is forced on).
|
|
|
|
Manifest JSON: [{"id": "...", "image": "/abs/path.jpg", "question": "..."}]
|
|
Output JSON: [{"id": "...", "answer": "..."}]
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import sys
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(description="Local VLM frame captioner (moondream2)")
|
|
parser.add_argument("--manifest", required=True)
|
|
parser.add_argument("--output", required=True)
|
|
parser.add_argument("--model-id", default="vikhyatk/moondream2")
|
|
parser.add_argument("--revision", default="2024-08-26")
|
|
args = parser.parse_args()
|
|
|
|
# Hard offline: never reach the network. The model must already be in the local cache.
|
|
os.environ["HF_HUB_OFFLINE"] = "1"
|
|
os.environ["TRANSFORMERS_OFFLINE"] = "1"
|
|
os.environ.setdefault("HF_HUB_DISABLE_XET", "1")
|
|
|
|
try:
|
|
from PIL import Image
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
except Exception as exc: # pragma: no cover - dependency guard
|
|
print(f"vision_caption: missing dependency: {exc}", file=sys.stderr)
|
|
return 3
|
|
|
|
with open(args.manifest, "r", encoding="utf-8") as handle:
|
|
items = json.load(handle)
|
|
|
|
try:
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
args.model_id, revision=args.revision, trust_remote_code=True, local_files_only=True)
|
|
tokenizer = AutoTokenizer.from_pretrained(
|
|
args.model_id, revision=args.revision, local_files_only=True)
|
|
model.eval()
|
|
except Exception as exc:
|
|
print(f"vision_caption: unable to load local model: {exc}", file=sys.stderr)
|
|
return 2
|
|
|
|
results = []
|
|
for item in items:
|
|
try:
|
|
image = Image.open(item["image"]).convert("RGB")
|
|
encoded = model.encode_image(image)
|
|
answer = model.answer_question(encoded, item["question"], tokenizer).strip()
|
|
except Exception as exc: # keep going; a single bad frame must not fail the batch
|
|
print(f"vision_caption: frame {item.get('id')} failed: {exc}", file=sys.stderr)
|
|
answer = ""
|
|
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())
|