Upload caption.py with huggingface_hub

caption.py

@@ -0,0 +1,154 @@
+"""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 == "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 == "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)