caption.py

"""AI captioning + trigger suggestion for the Krea 2 LoRA trainer Space.

Runs on the Space itself (cpu-basic) by calling the HF Inference API for a multimodal LLM
(`google/gemma-4-31B-it`, served with vision by the **novita** provider — the default `auto`
route lands on an endpoint that returns empty text, so the provider is pinned).

The captioning token (`CAPTION_HF_TOKEN` secret) only ever calls the Inference API. It is
independent of the user's OAuth token (push/dataset) and the gated `KREA_TOKEN` (Krea weights).

Caption recipe follows the Krea 2 authors' training guidance:
  * STYLE LoRA  — describe only the *content* (subjects, poses, layout, setting), never the
    medium/technique/palette, then append the style trigger phrase (e.g. ", heavy impasto style").
  * OBJECT/CHARACTER LoRA — describe the scene with the subject referred to by its class noun,
    then append a unique trigger token (e.g. " b3@rcup").
"""

from __future__ import annotations

import base64
import io
import os

from huggingface_hub import InferenceClient
from PIL import Image

CAPTION_MODEL = "google/gemma-4-31B-it"
CAPTION_PROVIDER = "novita"
_MAX_SIDE = 768  # downscale before upload to keep the request small / fast

# Few-shot exemplars taken from the authors' reference captions (content only — the trigger is
# appended programmatically, so the examples here deliberately omit the trailing trigger).
_STYLE_EXAMPLES = [
    "A person is running forward in profile. The figure leans into the motion with their head "
    "tilted slightly down and long hair trailing horizontally behind. The arms are bent at the "
    "elbows, with one arm swung forward and the other pulled back toward the hip. One leg is "
    "extended backward, capturing a mid-stride movement. The figure is positioned centrally in a "
    "void of plain white.",
    "A fishing boat is stationed in a narrow canal between rows of multi-story buildings. The boat "
    "features a central cabin with windows and two vertical masts extending upwards. A red buoy "
    "hangs from the side of the hull. The water in the canal occupies the lower portion of the "
    "scene, while the sky is visible above the rooflines of the structures.",
    "A black sports car is positioned in the center of a wet road through a dense forest. The car "
    "faces forward, with its round headlights visible. The road surface is covered in puddles that "
    "reflect the front end. Tall coniferous trees line both sides of the road and a dense fog fills "
    "the space between the trees behind the vehicle.",
]
_OBJECT_EXAMPLES = [
    "A cup, sitting on a grainy wooden table with a grey door in the background. An iron stand has "
    "grey and black plastic containers in separate piles.",
    "A cup being held by a woman in her hand in the outdoors. The background is a textured patch of "
    "lawn grass.",
]


def _token() -> str:
    tok = os.environ.get("CAPTION_HF_TOKEN") or os.environ.get("HF_TOKEN") or ""
    if not tok:
        raise RuntimeError("AI captioning is unavailable: the CAPTION_HF_TOKEN secret is not set.")
    return tok


def _client() -> InferenceClient:
    return InferenceClient(model=CAPTION_MODEL, provider=CAPTION_PROVIDER, token=_token())


def _data_url(path: str) -> str:
    img = Image.open(path).convert("RGB")
    img.thumbnail((_MAX_SIDE, _MAX_SIDE))
    buf = io.BytesIO()
    img.save(buf, "JPEG", quality=90)
    return "data:image/jpeg;base64," + base64.b64encode(buf.getvalue()).decode()


def _ask(instruction: str, image_paths: list[str], max_tokens: int = 320,
         temperature: float = 0.4) -> str:
    content: list[dict] = [{"type": "text", "text": instruction}]
    for p in image_paths:
        content.append({"type": "image_url", "image_url": {"url": _data_url(p)}})
    r = _client().chat_completion(
        messages=[{"role": "user", "content": content}],
        max_tokens=max_tokens, temperature=temperature,
    )
    return (r.choices[0].message.content or "").strip()


def _clean(text: str) -> str:
    """Strip wrapping quotes / a leading 'Caption:' the model sometimes adds."""
    t = text.strip().strip('"').strip("'").strip()
    for prefix in ("Caption:", "caption:", "Trigger:", "trigger:"):
        if t.startswith(prefix):
            t = t[len(prefix):].strip()
    return t.rstrip()


def caption_one(image_path: str, concept_type: str, trigger: str) -> str:
    """Caption a single image for the given concept type, appending the trigger."""
    trigger = (trigger or "").strip()
    if concept_type == "custom":
        instruction = (
            "Write a concise, natural training caption that describes this image as you see it: "
            "the subjects, what they are doing, the setting, and the overall look. Write 1-3 plain "
            "declarative sentences. Return only the caption, with no preamble, labels or quotes."
        )
        cap = _clean(_ask(instruction, [image_path]))
        if trigger:
            cap = f"{trigger}, {cap}" if cap else trigger
        return cap
    if concept_type == "style":
        instruction = (
            "You are writing a training caption for a STYLE LoRA. Describe ONLY the literal "
            "content of the image: the subjects, their poses and actions, the key objects, their "
            "spatial arrangement, and the setting or background. Write 2-4 plain declarative "
            "sentences. Do NOT mention the artistic style, medium, technique, brushwork, lighting "
            "mood or palette, and do NOT use words like painting, illustration, drawing, render, "
            "sketch or photo. Match the tone of these examples:\n\n"
            + "\n\n".join(_STYLE_EXAMPLES)
            + "\n\nReturn only the caption sentence(s), with no preamble, labels or quotes."
        )
        cap = _clean(_ask(instruction, [image_path]))
        if trigger:
            cap = f"{cap.rstrip('.')}, {trigger}" if cap else trigger
        return cap
    # object / character
    instruction = (
        "You are writing a training caption for a LoRA of one specific subject. Describe the "
        "scene: where the subject is, what it is doing or how it is positioned, and the background "
        "or setting. Refer to the subject by its generic class noun (e.g. 'a cup', 'a dog'), never "
        "by a name. Write 1-3 plain declarative sentences. Match the tone of these examples:\n\n"
        + "\n\n".join(_OBJECT_EXAMPLES)
        + "\n\nReturn only the caption sentence(s), with no preamble, labels or quotes."
    )
    cap = _clean(_ask(instruction, [image_path]))
    if trigger:
        cap = f"{cap} {trigger}" if cap else trigger
    return cap


def suggest_trigger(image_paths: list[str], concept_type: str) -> str:
    """Suggest a trigger from 2-3 sample images: a style phrase, or a unique object token."""
    sample = list(image_paths)[:3]
    if not sample:
        raise gr_error("Upload images first.")
    if concept_type == "custom":
        instruction = (
            "Propose a SHORT unique trigger token for the concept shown in these images: a rare "
            "made-up token, optionally followed by a class noun. Examples: 'TOK', 'b3@rcup', "
            "'zxy style'. Return only the trigger, with no quotes or explanation."
        )
        return _clean(_ask(instruction, sample, max_tokens=16, temperature=0.7))
    if concept_type == "style":
        instruction = (
            "These images share one artistic style. Propose a SHORT distinctive trigger phrase "
            "naming that style: 2 to 5 words, ending with the word 'style'. Examples: 'heavy "
            "impasto style', 'monochrome ink wash style', 'flat pastel vector style'. Return only "
            "the phrase in lowercase, with no quotes or explanation."
        )
        return _clean(_ask(instruction, sample, max_tokens=24, temperature=0.6)).lower()
    instruction = (
        "These images show one specific subject. Propose a SHORT unique trigger for it: a rare "
        "made-up token, optionally followed by its class noun. Examples: 'b3@rcup', 'sks dog', "
        "'zxy sneaker'. Return only the trigger, with no quotes or explanation."
    )
    return _clean(_ask(instruction, sample, max_tokens=16, temperature=0.7))


def gr_error(msg: str):  # tiny indirection so this module stays importable without gradio
    try:
        import gradio as gr  # noqa: PLC0415
        return gr.Error(msg)
    except Exception:  # noqa: BLE001
        return ValueError(msg)