scripts/clip_align_kink_assets.py

#!/usr/bin/env python3
"""
Offline CLIP scoring: cosine similarity between each cached image and a text prompt = **kink name +
definition** (same data the app shows).

Uses open_clip (ViT-B-32 + LAION). Install::

    pip install torch open-clip-torch pillow

Writes ``data/clip_asset_scores.json`` next to the SQLite DB. Each URL maps to
``similarity``, ``accepted`` (vs ``--min-similarity``), ``kink_id``, ``prompt``, etc.

**Default ``--min-similarity``** is tuned from ``scripts/clip_threshold_report.py`` on the
existing distribution (stricter than the original 0.20 floor). Re-run that report after big
ingestion changes.

The API catalog reads this file optionally. With ``KINK_STRICT_CLIP_IMAGES=1`` and a non-empty
JSON file, assets **without** a CLIP row are dropped at catalog build (see ``backend/image_policy.py``).
Legacy behavior (no env) still allows assets missing from the file.

After ``compute``, reload the catalog (``POST /admin/reload`` or restart the server).

Examples::

    python scripts/clip_align_kink_assets.py compute --db data/store.db
    python scripts/clip_align_kink_assets.py inspect --db data/store.db --kink-id rope_bondage
"""
from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

# Repo root (parent of scripts/)
ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))


def media_path_for_url(asset_url: str, media_root: Path) -> Path | None:
    prefix = "/media/"
    if not asset_url.startswith(prefix):
        return None
    rest = asset_url[len(prefix) :].lstrip("/")
    target = (media_root / rest).resolve()
    if media_root.resolve() not in target.parents and target != media_root.resolve():
        return None
    return target if target.is_file() else None


def prompt_for_kink(kink, defs_first: dict[str, str]) -> str:
    parts = [kink.name.strip()]
    d0 = (kink.short_definition or "").strip()
    if d0:
        parts.append(d0)
    else:
        t = defs_first.get(kink.id, "").strip()
        if t:
            parts.append(t)
    return ". ".join(p for p in parts if p)


def cmd_compute(args: argparse.Namespace) -> int:
    try:
        import torch
        import open_clip
    except ImportError as e:
        print("Missing deps. Install: pip install torch open-clip-torch pillow", file=sys.stderr)
        raise SystemExit(1) from e

    from PIL import Image
    from sqlmodel import Session, create_engine, select

    from models import Asset, Definition, Kink

    db_path = Path(args.db).resolve()
    media_root = Path(args.media_root).resolve()
    engine = create_engine(f"sqlite:///{db_path}")

    model_name = args.model
    pretrained = args.pretrained
    device = args.device or ("cuda" if torch.cuda.is_available() else "cpu")
    batch_size = max(1, int(args.batch_size))
    min_s = float(args.min_similarity)

    model, _, preprocess = open_clip.create_model_and_transforms(model_name, pretrained=pretrained)
    model = model.to(device).eval()
    tokenizer = open_clip.get_tokenizer(model_name)

    defs_first: dict[str, str] = {}
    with Session(engine) as session:
        for row in session.exec(select(Definition)).all():
            if row.kink_id not in defs_first:
                defs_first[row.kink_id] = (row.text or "").strip()

        kinks = {k.id: k for k in session.exec(select(Kink)).all()}
        assets = session.exec(select(Asset)).all()

    work: list[tuple[str, str, Path, str, str]] = []
    for row in assets:
        kink = kinks.get(row.kink_id)
        if not kink:
            continue
        path = media_path_for_url(row.asset_url, media_root)
        if not path:
            continue
        text = prompt_for_kink(kink, defs_first)
        if not text:
            continue
        try:
            rel_path = str(path.relative_to(media_root))
        except ValueError:
            rel_path = str(path)
        work.append((row.asset_url, row.kink_id, path, text, rel_path))

    total = len(work)
    print(f"CLIP compute: {total} assets with local files + prompts, device={device}, batch={batch_size}", flush=True)

    out: dict[str, dict] = {}
    processed = 0

    for start in range(0, total, batch_size):
        chunk = work[start : start + batch_size]
        tensors = []
        texts: list[str] = []
        metas: list[tuple[str, str, str, str]] = []
        for asset_url, kid, path, text, rel_path in chunk:
            try:
                tensors.append(preprocess(Image.open(path).convert("RGB")))
            except OSError:
                continue
            texts.append(text)
            metas.append((asset_url, kid, text, rel_path))

        if not tensors:
            continue

        image_tensor = torch.stack(tensors).to(device)
        tokens = tokenizer(texts).to(device)

        with torch.inference_mode():
            image_features = model.encode_image(image_tensor)
            text_features = model.encode_text(tokens)
            image_features = image_features / image_features.norm(dim=-1, keepdim=True)
            text_features = text_features / text_features.norm(dim=-1, keepdim=True)
            sims = (image_features * text_features).sum(dim=-1).cpu()

        for i, sim_v in enumerate(sims.tolist()):
            asset_url, kid, text, rel_path = metas[i]
            accepted = float(sim_v) >= min_s
            out[asset_url] = {
                "kink_id": kid,
                "similarity": float(sim_v),
                "accepted": accepted,
                "min_similarity_used": min_s,
                "prompt": text[:500],
                "local_path": rel_path,
            }

        processed += len(tensors)
        if start == 0 or (start // batch_size) % 50 == 0 or start + batch_size >= total:
            print(f"  … {processed}/{total} scored", flush=True)

    out_path = db_path.parent / "clip_asset_scores.json"
    out_path.write_text(json.dumps(out, separators=(",", ":")) + "\n", encoding="utf-8")
    n_ok = sum(1 for v in out.values() if v.get("accepted"))
    print(f"Wrote {len(out)} scored assets to {out_path} ({n_ok} accepted, min_similarity={min_s})")
    return 0


def cmd_inspect(args: argparse.Namespace) -> int:
    db_path = Path(args.db).resolve()
    path = db_path.parent / "clip_asset_scores.json"
    if not path.is_file():
        print(f"No scores file at {path}. Run: {sys.argv[0]} compute --db {args.db}", file=sys.stderr)
        return 1
    data = json.loads(path.read_text(encoding="utf-8"))
    rows = []
    for url, meta in data.items():
        if args.kink_id and meta.get("kink_id") != args.kink_id:
            continue
        rows.append((float(meta.get("similarity", -999)), url, meta))
    rows.sort(key=lambda r: -r[0])
    if args.limit:
        rows = rows[: args.limit]
    for sim, url, meta in rows:
        acc = meta.get("accepted", "")
        print(f"{sim:8.4f}  {str(acc):>5}  {meta.get('kink_id', '')}  {url}")
        if args.verbose and meta.get("prompt"):
            print(f"         prompt: {meta['prompt'][:200]}...")
    return 0


def main() -> int:
    ap = argparse.ArgumentParser(description="CLIP alignment scores for kink images")
    sub = ap.add_subparsers(dest="cmd", required=True)

    c = sub.add_parser("compute", help="Score all local assets and write clip_asset_scores.json")
    c.add_argument("--db", default=str(ROOT / "data" / "store.db"))
    c.add_argument("--media-root", default=str(ROOT / "data" / "cached_assets"))
    c.add_argument("--model", default="ViT-B-32")
    c.add_argument("--pretrained", default="laion2b_s34b_b79k")
    c.add_argument("--device", default="")
    c.add_argument(
        "--batch-size",
        type=int,
        default=0,
        help="Images per forward pass (default: 64 cuda, 24 cpu).",
    )
    c.add_argument(
        "--min-similarity",
        type=float,
        default=0.24,
        help="Offline acceptance threshold (image vs kink name+definition embedding similarity). Default 0.24; use clip_threshold_report.py to compare cuts.",
    )

    i = sub.add_parser("inspect", help="Print sorted scores (optionally one kink)")
    i.add_argument("--db", default=str(ROOT / "data" / "store.db"))
    i.add_argument("--kink-id", default="")
    i.add_argument("--limit", type=int, default=40)
    i.add_argument("-v", "--verbose", action="store_true")

    args = ap.parse_args()
    if args.cmd == "compute":
        try:
            import torch

            if args.batch_size <= 0:
                args.batch_size = 64 if torch.cuda.is_available() else 24
        except ImportError:
            args.batch_size = 24
        return cmd_compute(args)
    return cmd_inspect(args)


if __name__ == "__main__":
    raise SystemExit(main())