"""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)