Upload viz/preference/pref_score2.py with huggingface_hub

viz/preference/pref_score2.py

@@ -0,0 +1,82 @@
+import os, glob, json, math
+import torch
+from PIL import Image
+
+DPG = "/mnt/localssd/ELLA_dpg/dpg_bench/prompts"
+DEVICE = "cuda"
+
+DOMAINS = {
+  "pixel": ("/mnt/localssd/eval_work/pref_pixel/baseline",
+            "/mnt/localssd/eval_work/pref_pixel/gan_dinov2_text"),
+  "latent": ("/mnt/localssd/eval_work_sana/sana_blip3o_sft/imgs_50000",
+             "/mnt/localssd/eval_work_sana/sana_blip3o_gan/imgs_50000"),
+}
+
+def common_idxs(nogan, gan):
+    def idxs(d):
+        return {os.path.basename(f)[:-len("_0.png")] for f in glob.glob(os.path.join(d,"*_0.png"))}
+    c = idxs(nogan) & idxs(gan)
+    c = {i for i in c if os.path.exists(os.path.join(DPG, i + ".txt"))}
+    return sorted(c, key=lambda x: (0, int(x)) if x.isdigit() else (1, x))
+
+def prompt_for(i):
+    return open(os.path.join(DPG, i + ".txt")).read().strip()
+
+import ImageReward
+ir = ImageReward.load("ImageReward-v1.0", device=DEVICE)
+def ir_score(p, path):
+    return float(ir.score(p, path))
+
+from transformers import AutoProcessor, AutoModel
+pk_proc = AutoProcessor.from_pretrained('laion/CLIP-ViT-H-14-laion2B-s32B-b79K')
+pk_model = AutoModel.from_pretrained('yuvalkirstain/PickScore_v1').eval().to(DEVICE)
+@torch.no_grad()
+def pk_score(p, path):
+    img = Image.open(path).convert("RGB")
+    ii = pk_proc(images=img, return_tensors="pt").to(DEVICE)
+    ti = pk_proc(text=p, padding=True, truncation=True, max_length=77, return_tensors="pt").to(DEVICE)
+    ie = pk_model.get_image_features(**ii); ie = ie / ie.norm(dim=-1, keepdim=True)
+    te = pk_model.get_text_features(**ti); te = te / te.norm(dim=-1, keepdim=True)
+    return (pk_model.logit_scale.exp() * (te @ ie.T)).item()
+
+def binom_two_sided(k, n):
+    if n == 0:
+        return 1.0
+    kk = min(k, n - k); ln2n = n * math.log(2)
+    tail = sum(math.exp(math.lgamma(n+1) - math.lgamma(i+1) - math.lgamma(n-i+1) - ln2n) for i in range(kk+1))
+    return min(1.0, 2 * tail)
+
+MODELS = [("ImageReward-v1.0", ir_score), ("PickScore_v1", pk_score)]
+
+results = {"domains": {}}
+for dom, (ng, g) in DOMAINS.items():
+    idxs = common_idxs(ng, g)
+    results["domains"][dom] = {"n_pairs": len(idxs), "models": {}}
+    print(dom, "n_pairs", len(idxs), flush=True)
+    for mname, fn in MODELS:
+        rows = []
+        for i in idxs:
+            p = prompt_for(i)
+            s_ng = fn(p, os.path.join(ng, i + "_0.png"))
+            s_g = fn(p, os.path.join(g, i + "_0.png"))
+            rows.append((i, s_ng, s_g))
+        n = len(rows)
+        wins = sum(1 for _, a, b in rows if b > a)
+        ties = sum(1 for _, a, b in rows if b == a)
+        losses = n - wins - ties
+        mn = sum(a for _, a, _ in rows) / n
+        mg = sum(b for _, _, b in rows) / n
+        md = sum(b - a for _, a, b in rows) / n
+        eff = wins + losses
+        summ = {"n": n, "wins": wins, "ties": ties, "losses": losses,
+                "win_rate": 100 * wins / n,
+                "win_rate_excl_ties": 100 * wins / eff if eff else None,
+                "mean_nogan": mn, "mean_gan": mg, "mean_delta": md,
+                "binom_p_two_sided": binom_two_sided(wins, eff)}
+        results["domains"][dom]["models"][mname] = {
+            "summary": summ,
+            "rows": [{"idx": i, "nogan": a, "gan": b, "delta": b - a} for i, a, b in rows]}
+        print(dom, mname, json.dumps(summ), flush=True)
+
+json.dump(results, open("/mnt/localssd/pref_results2.json", "w"), indent=2)
+print("WROTE /mnt/localssd/pref_results2.json", flush=True)