caption.py

from PIL import Image
import torch

def trim_to_n_sentences(text: str, n: int = 6) -> str:
    s = [x.strip() for x in text.replace("\n", " ").split(".") if x.strip()]
    return ". ".join(s[:n]) + "."

def caption_image(
    processor,
    model,
    device,
    image: Image.Image,
    prompt: str,
    max_new_tokens: int = 120,
):
    messages = [{
        "role": "user",
        "content": [
            {"type": "image"},
            {"type": "text", "text": prompt},
        ],
    }]

    text = processor.apply_chat_template(
        messages,
        add_generation_prompt=True,
    )

    inputs = processor(
        text=text,
        images=image.convert("RGB"),
        return_tensors="pt",
    )

    inputs = {k: v.to(device) for k, v in inputs.items()}

    with torch.inference_mode():
        out = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=False,
            temperature=0.0,
            repetition_penalty=1.2,
        )

    raw = processor.decode(out[0], skip_special_tokens=True)
    caption = raw.split("assistant")[-1].strip()
    return trim_to_n_sentences(caption)