pipeline/captions.py

"""Step 7: caption each generated image with a vision-language model."""
from __future__ import annotations

import base64
import mimetypes
from pathlib import Path
from typing import List

from huggingface_hub import InferenceClient

from . import config
from .factual_accuracy import FACTUAL_ACCURACY_GUIDELINES

_CAPTION_SYSTEM = (
    "You are writing image captions for a published blog post.\n\n" + FACTUAL_ACCURACY_GUIDELINES
)

_CAPTION_PROMPT = (
    "Write a concise, engaging one-sentence caption for this blog illustration. "
    "Describe what is shown; do not start with 'This image' or 'A picture of'. "
    "Return only the caption."
)


def _data_uri(path: Path) -> str:
    mime = mimetypes.guess_type(str(path))[0] or "image/png"
    b64 = base64.b64encode(path.read_bytes()).decode("utf-8")
    return f"data:{mime};base64,{b64}"


def caption_images(client: InferenceClient, images: List[dict]) -> List[dict]:
    """Add a 'caption' key to each image dict that has a valid 'path'."""
    for item in images:
        path = item.get("path")
        if not path or not Path(path).exists():
            item["caption"] = ""
            continue
        try:
            resp = client.chat.completions.create(
                model=config.MODEL_VISION,
                messages=[
                    {"role": "system", "content": _CAPTION_SYSTEM},
                    {
                        "role": "user",
                        "content": [
                            {"type": "text", "text": _CAPTION_PROMPT},
                            {"type": "image_url", "image_url": {"url": _data_uri(Path(path))}},
                        ],
                    },
                ],
                max_tokens=80,
                temperature=0.5,
            )
            item["caption"] = (resp.choices[0].message.content or "").strip().strip('"')
        except Exception:
            # fall back to the scene description if captioning fails
            item["caption"] = item.get("scene", "")
    return images