app.py

"""Iconclass-9B ZeroGPU demo — gradio.Server pattern with a custom Tufte frontend.

Upload an artwork -> the 9B VLM (davanstrien/qwen35-9b-iconclass-sft-multitask-2ep)
predicts Iconclass codes -> each code is shown with its decoded meaning (via the
`iconclass` package) and a link to https://iconclass.org/{code}.

The multitask model supports three prompt modes (templates must stay byte-identical
to training — see model-training/iconclass-qwen35/train_sft_brill.py):
  1. standard      — "classify this image"
  2. N-conditioned — "...containing exactly N codes"   (cataloguing-depth knob)
  3. completion    — "image already has codes [...]; add (exactly N | any) more"
                     (catalogue-densification: existing codes condition the model)

Inference recipe (empirically de-risked, do not change):
  enable_thinking=False, max_new_tokens=768, do_sample=False, repetition_penalty=1.1
The repetition_penalty is ESSENTIAL: without it greedy decoding degenerates into a
runaway sibling-code enumeration that blows the token budget and yields invalid JSON.
"""
from __future__ import annotations

import json
import os
import re
import threading
import time
import urllib.request
import uuid
from functools import lru_cache
from pathlib import Path

import spaces
import torch
from fastapi.responses import HTMLResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from gradio import Server
from gradio.data_classes import FileData
from iconclass import init as ic_init
from iconclass import split_on_colon
from PIL import Image
from pydantic import BaseModel
from transformers import AutoModelForImageTextToText, AutoProcessor

MODEL = "davanstrien/qwen35-9b-iconclass-sft-multitask-2ep"
SEED_REPO = "davanstrien/iconclass-annotation-seed"

# --- Prompt templates: KEEP IN SYNC with train_sft_brill.py (training format) ---
INSTRUCTION = (
    "Classify this image using Iconclass codes. "
    "Return a JSON object with key 'iconclass-codes' containing a list of codes."
)
INSTRUCTION_N = (
    "Classify this image using Iconclass codes. "
    "Return a JSON object with key 'iconclass-codes' containing exactly {n} codes."
)
INSTRUCTION_COMPLETE_N = (
    "This image already has the following Iconclass codes: {existing}. "
    "Add exactly {n} more Iconclass codes for aspects not yet covered. "
    "Return a JSON object with key 'iconclass-codes' containing only the new codes."
)
INSTRUCTION_COMPLETE_OPEN = (
    "This image already has the following Iconclass codes: {existing}. "
    "Add any further applicable Iconclass codes for aspects not yet covered. "
    "Return a JSON object with key 'iconclass-codes' containing only the new codes."
)

HERE = Path(__file__).resolve().parent

# ---------------------------------------------------------------------------
# Iconclass decoding
# ---------------------------------------------------------------------------
_ic = ic_init()


def _base(code: str) -> str:
    """Strip modifiers/keys to the base notation used for lookup (e.g. 61B(+52) -> 61B)."""
    return code.split("(")[0].strip() if isinstance(code, str) else ""


@lru_cache(maxsize=100_000)
def _decode_one(code: str) -> str:
    """Decode a single (non-composite) Iconclass notation to English."""
    base = _base(code)
    try:
        node = _ic[base]
        txt = node() if callable(node) else str(node)
        return txt or base or code
    except Exception:
        return base or code


@lru_cache(maxsize=100_000)
def decode_meaning(code: str) -> str:
    """Decoded English meaning for an Iconclass code.

    Composite codes joined with a colon (e.g. 31D15:61BB) are decoded part-by-part
    and joined with a middle dot. Falls back to the raw code on any failure.
    """
    try:
        parts = split_on_colon(code)
    except Exception:
        parts = [code]
    return " · ".join(_decode_one(p) for p in parts) if parts else _decode_one(code)


def code_url(code: str) -> str:
    return f"https://iconclass.org/{code}"


@lru_cache(maxsize=100_000)
def _decode_exact(code: str) -> str:
    """Decode a notation EXACTLY as given (keyed/named forms resolve in the lib);
    falls back to the base-stripped decode, then to the raw code."""
    try:
        node = _ic[code]
        txt = node() if callable(node) else str(node)
        if txt:
            # the lib doubles the name on '(NAME)' forms: "... (AZOR) (AZOR)"
            return re.sub(r"(\([^()]+\))\s*\1$", r"\1", txt)
    except Exception:
        pass
    return _decode_one(code)


@lru_cache(maxsize=100_000)
def hierarchy_path(code: str) -> tuple:
    """Decoded root->leaf ladder(s) for a code, for non-expert display.

    Returns a tuple of ladders (one per ':'-composite part); each ladder is a
    tuple of {"code", "text"} rungs. Uses the iconclass lib's Notation.path()
    (handles letter steps, '(+key)' rungs and '(NAME)' forms natively); drops
    the '(...)' placeholder rung when a named leaf follows it.
    """
    try:
        parts = split_on_colon(code) or [code]
    except Exception:
        parts = [code]
    ladders = []
    for part in parts:
        try:
            rungs = [str(r) for r in _ic[part].path()]
        except Exception:
            rungs = []
        if not rungs:
            rungs = [part]
        if rungs[-1] != part:
            rungs.append(part)  # e.g. named forms whose path ends at '(...)'
        rungs = [
            r for i, r in enumerate(rungs)
            if not (r.endswith("(...)") and i + 1 < len(rungs))
        ]
        ladder, prev_txt = [], None
        for r in rungs:
            txt = _decode_exact(r)
            if txt == prev_txt:
                continue  # skip rungs that add no new meaning
            prev_txt = txt
            ladder.append({"code": r, "text": txt})
        ladders.append(tuple(ladder))
    return tuple(ladders)


def _parse_existing(existing_codes: str) -> list[str]:
    """Parse the user's 'existing codes' input (comma/semicolon/newline separated)."""
    if not existing_codes:
        return []
    raw = existing_codes.replace(";", ",").replace("\n", ",").split(",")
    seen, out = set(), []
    for c in raw:
        c = c.strip()
        if c and c not in seen:
            seen.add(c)
            out.append(c)
    return out


# ---------------------------------------------------------------------------
# Model (loaded once at startup on CPU; moved to cuda inside the GPU function)
# ---------------------------------------------------------------------------
print("loading processor + model (CPU)…", flush=True)
processor = AutoProcessor.from_pretrained(MODEL)
model = AutoModelForImageTextToText.from_pretrained(MODEL, torch_dtype=torch.bfloat16)
model.eval()
print("loaded.", flush=True)


def _build_instruction(num_codes: int, existing: list[str]) -> str:
    """Pick the training-format instruction for the requested mode.

    num_codes <= 0 means 'auto' (model decides how many codes).
    """
    if existing:
        if num_codes and num_codes > 0:
            return INSTRUCTION_COMPLETE_N.format(existing=json.dumps(existing), n=num_codes)
        return INSTRUCTION_COMPLETE_OPEN.format(existing=json.dumps(existing))
    if num_codes and num_codes > 0:
        return INSTRUCTION_N.format(n=num_codes)
    return INSTRUCTION


def _run_model(image: Image.Image, instruction: str) -> str:
    """Run the validated inference recipe and return the raw decoded text."""
    inner = getattr(processor, "tokenizer", processor)
    msgs = [
        {"role": "user", "content": [{"type": "image"}, {"type": "text", "text": instruction}]}
    ]
    text = inner.apply_chat_template(
        msgs, add_generation_prompt=True, tokenize=False, enable_thinking=False
    )
    inputs = processor(text=text, images=[image.convert("RGB")], return_tensors="pt").to("cuda")
    out = model.generate(
        **inputs,
        max_new_tokens=768,
        do_sample=False,
        repetition_penalty=1.1,
    )
    return processor.decode(out[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)


def _parse_codes(raw: str) -> list[str]:
    try:
        obj = json.loads(raw)
    except Exception:
        return []
    codes = obj.get("iconclass-codes", []) if isinstance(obj, dict) else []
    seen, out = set(), []
    for c in codes:
        if isinstance(c, str) and c and c not in seen:
            seen.add(c)
            out.append(c)
    return out


# ---------------------------------------------------------------------------
# Annotation mode: seed (precomputed predictions, no GPU) + verdict storage
# ---------------------------------------------------------------------------
_seed_cache: dict = {}
_seed_lock = threading.Lock()


def get_seed() -> dict:
    """Server-side proxy for seed.json (same-origin for the frontend, cached)."""
    with _seed_lock:
        if not _seed_cache:
            url = f"https://huggingface.co/datasets/{SEED_REPO}/resolve/main/seed.json"
            try:
                with urllib.request.urlopen(url, timeout=30) as r:
                    _seed_cache.update(json.loads(r.read()))
            except Exception as exc:
                return {"error": f"seed unavailable: {exc}", "items": []}
    return _seed_cache


# Verdicts: JSONL appended locally, synced to a PRIVATE HF Bucket every couple
# of minutes (object-store semantics: each flush re-uploads changed files; one
# uniquely-named file per app instance so concurrent replicas never collide).
ANNOTATIONS_BUCKET = "davanstrien/iconclass-annotations"
_ann_dir = Path("annotations")
_ann_dir.mkdir(exist_ok=True)
_ann_file = _ann_dir / f"events_{uuid.uuid4().hex[:12]}.jsonl"
_ann_lock = threading.Lock()
_bucket_ok = False


def _start_bucket_flusher():
    global _bucket_ok
    if not os.environ.get("HF_TOKEN"):
        print("HF_TOKEN not set; annotations stored locally only (lost on restart)", flush=True)
        return
    try:
        from huggingface_hub import HfApi

        api = HfApi()
        api.create_bucket(ANNOTATIONS_BUCKET, private=True, exist_ok=True)
        _bucket_ok = True
        print(f"annotation bucket flusher -> hf://buckets/{ANNOTATIONS_BUCKET}/raw", flush=True)
    except Exception as exc:
        print(f"bucket unavailable ({exc}); annotations stored locally only", flush=True)
        return

    def _flush_loop():
        last_size = -1
        while True:
            time.sleep(120)
            try:
                size = _ann_file.stat().st_size if _ann_file.exists() else 0
                if size and size != last_size:
                    with _ann_lock:
                        api.sync_bucket(
                            source=str(_ann_dir),
                            dest=f"hf://buckets/{ANNOTATIONS_BUCKET}/raw",
                            quiet=True,
                        )
                    last_size = size
            except Exception as exc:  # transient errors: retry next tick
                print(f"bucket flush failed (will retry): {exc}", flush=True)

    threading.Thread(target=_flush_loop, daemon=True, name="bucket-flusher").start()


_start_bucket_flusher()


class Annotation(BaseModel):
    seed_id: int
    code: str
    verdict: str  # simple ui: "yes" | "no" | "skip"; ladder ui: "depth" | "none" | "skip"
    depth_index: int | None = None  # 0-based rung index when verdict == "depth"
    rung_code: str | None = None
    ladder_index: int | None = None  # which ':'-part, for composite codes
    max_depth: int | None = None
    gt: bool | None = None
    judge: str | None = None
    session: str = ""
    ui: str | None = None  # "simple-v1" | "ladder-v1" — distinguishes the two voting UIs


app = Server()


@app.get("/seed")
async def seed_endpoint():
    return JSONResponse(get_seed())


@app.post("/annotate")
async def annotate_endpoint(ann: Annotation):
    event = ann.model_dump()
    event["model_version"] = MODEL
    event["ts"] = time.time()
    with _ann_lock:
        with _ann_file.open("a") as f:
            f.write(json.dumps(event) + "\n")
    return JSONResponse({"ok": True, "persisted": _bucket_ok})


@app.api(name="classify")
@spaces.GPU(duration=120)
def classify(image_path: FileData, num_codes: int = 0, existing_codes: str = "") -> dict:
    """Predict Iconclass codes for an uploaded artwork.

    Modes (selected by the optional params):
      - default: model decides which / how many codes
      - num_codes > 0: ask for exactly that many codes (cataloguing-depth knob)
      - existing_codes non-empty: completion mode — the given codes condition the
        model, which proposes codes for aspects *not yet covered* (additions only)

    Returns {"codes": [...], "existing": [...], "mode": str, "raw": str, "error": str|None}
    where each code entry is {"code", "meaning", "url"}.
    """
    path = image_path["path"] if isinstance(image_path, dict) else image_path
    image = Image.open(path).convert("RGB")

    existing = _parse_existing(existing_codes)
    n = int(num_codes) if num_codes else 0
    instruction = _build_instruction(n, existing)
    mode = "completion" if existing else ("exact-n" if n > 0 else "standard")

    model.to("cuda")
    raw = _run_model(image, instruction)

    codes = _parse_codes(raw)
    if existing:  # never echo a given code back as an "addition"
        given = set(existing)
        codes = [c for c in codes if c not in given]
    decoded = [
        {
            "code": c,
            "meaning": decode_meaning(c),
            "url": code_url(c),
            "path": [list(ladder) for ladder in hierarchy_path(c)],
        }
        for c in codes
    ]
    existing_decoded = [
        {"code": c, "meaning": decode_meaning(c), "url": code_url(c)} for c in existing
    ]
    return {
        "codes": decoded,
        "existing": existing_decoded,
        "mode": mode,
        "raw": raw,
        "error": None if decoded else "No valid Iconclass codes parsed from the model output.",
    }


# ---------------------------------------------------------------------------
# Frontend
# ---------------------------------------------------------------------------
@app.get("/", response_class=HTMLResponse)
async def homepage():
    return (HERE / "index.html").read_text(encoding="utf-8")


# Serve example artworks (and any static assets) under /static
_static_dir = HERE / "examples"
if _static_dir.is_dir():
    app.mount("/static", StaticFiles(directory=str(_static_dir)), name="static")


if __name__ == "__main__":
    app.launch(show_error=True)