Moongrep Standalone

moongrep

#!/usr/bin/env -S uv run --script
# /// script
# requires-python = ">=3.13"
# dependencies = [
#    "einops>=0.8.1",
#    "pillow>=11.2.1",
#    "pyvips>=3.0.0",
#    "torch>=2.7.1",
#    "transformers>=4.52.4"
# ]
# ///

import argparse
import os

from transformers import AutoModelForCausalLM
from PIL import Image

def main(args):
    model = AutoModelForCausalLM.from_pretrained(
        "vikhyatk/moondream2",
        revision="2025-01-09",
        trust_remote_code=True
    )

    for file in args['files']:
        path = os.path.join(str(file))

        if path.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif')):
            image = Image.open(path)
            query = model.query(image, args['prompt'][0])
            answer = query['answer']
            print(f"file: {path} answer: {answer}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Moongrep: grep with AI"
    )
    parser.add_argument(
        'prompt', type=str, nargs=1,
        help="Prompt for the model."
    )
    parser.add_argument(
        'files', type=str, nargs='+'
    )
    args = vars(parser.parse_args())
    main(args)