Spaces:
Sleeping
Sleeping
File size: 2,032 Bytes
31fa536 fbe8f5e 31fa536 fbe8f5e 31fa536 fbe8f5e 31fa536 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 | """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
|