#!/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())