Add DreamSim (human-aligned) encoder: LoRA-tuned DINO ViT-B/16 patch tokens + CLS attention

app.py

@@ -33,6 +33,7 @@ MODELS = {
     "DINOv3 ViT-7B · 6.7B (needs GPU)": dict(repo="facebook/dinov3-vit7b16-pretrain-lvd1689m",    kind="dinov3", patch=16, norm=IMNET),
     "MAE ViT-L (pixel-reconstruction)": dict(repo="facebook/vit-mae-large",                       kind="mae",    patch=16, norm=IMNET),
     "CLIP ViT-L/14 (language-aligned)": dict(repo="openai/clip-vit-large-patch14",                kind="clip",   patch=14, norm=CLIPN),
+    "DreamSim (human-aligned)":         dict(repo="dreamsim:dino_vitb16",                        kind="dreamsim", patch=16, norm=None),
 }
 DEFAULT_MODEL = next(iter(MODELS))
 
@@ -40,25 +41,38 @@ TREE_HTML = open("tree_widget.html").read()
 REVEAL_HTML = open("reveal_widget.html").read()
 
 # single-slot cache: keep only the most-recently-used model (avoids OOM when switching)
-_cur = {"repo": None, "model": None}
+_cur = {"repo": None, "model": None, "pre": None}
 def get_model(spec):
     repo = spec["repo"]
     if _cur["repo"] != repo:
-        _cur["model"] = None
+        _cur["model"] = _cur["pre"] = None
         import gc; gc.collect()
-        dt = torch.bfloat16 if "vit7b16" in repo else torch.float32   # 7B in bf16 to fit
-        print(f"[app] loading {repo} ({dt})", flush=True)
-        cls = {"dinov3": AutoModel, "mae": ViTMAEModel, "clip": CLIPVisionModel}[spec["kind"]]
-        mdl = cls.from_pretrained(repo, dtype=dt, token=TOKEN, attn_implementation="eager")
-        if spec["kind"] == "mae":
-            mdl.config.mask_ratio = 0.0        # keep ALL patches (MAE masks 75% by default)
-        _cur["model"] = mdl.eval().to(DEV)
+        print(f"[app] loading {repo}", flush=True)
+        if spec["kind"] == "dreamsim":
+            from dreamsim import dreamsim as _ds          # lazy: heavy dep, only when selected
+            cache = os.environ.get("DREAMSIM_CACHE", "/tmp/dreamsim_cache")
+            os.makedirs(cache, exist_ok=True)
+            mdl, pre = _ds(pretrained=True, dreamsim_type="dino_vitb16", cache_dir=cache, device=DEV)
+            _cur["model"], _cur["pre"] = mdl.eval(), pre
+        else:
+            dt = torch.bfloat16 if "vit7b16" in repo else torch.float32   # 7B in bf16 to fit
+            cls = {"dinov3": AutoModel, "mae": ViTMAEModel, "clip": CLIPVisionModel}[spec["kind"]]
+            mdl = cls.from_pretrained(repo, dtype=dt, token=TOKEN, attn_implementation="eager")
+            if spec["kind"] == "mae":
+                mdl.config.mask_ratio = 0.0    # keep ALL patches (MAE masks 75% by default)
+            _cur["model"] = mdl.eval().to(DEV)
         _cur["repo"] = repo
     return _cur["model"]
 
 
 def encode(model, spec, x):
     """-> (feat[P,C], att[P]) with patches in spatial row-major order."""
+    if spec["kind"] == "dreamsim":                        # LoRA-tuned DINO ViT-B/16 inside DreamSim
+        vit = model.extractor_list[0].model
+        with torch.no_grad():
+            tok = vit.get_intermediate_layers(x, 1)[0]    # (1, 1+P, C)
+            a = vit.get_last_selfattention(x)             # (1, heads, S, S)
+        return tok[0, 1:].float().cpu().numpy(), a[0, :, 0, 1:].mean(0).float().cpu().numpy()
     with torch.no_grad():
         o = model(x, output_attentions=True)
     h, a = o.last_hidden_state, o.attentions[-1][0]            # a: (heads, S, S)
@@ -76,10 +90,13 @@ def build_tree(image, model_label):
     spec = MODELS.get(model_label, MODELS[DEFAULT_MODEL])
     model = get_model(spec)
     mdt = next(model.parameters()).dtype
-    mean, std = spec["norm"]
     img = image.convert("RGB").resize((RES, RES), Image.BICUBIC)
-    arr = np.asarray(img).astype(np.float32) / 255
-    x = ((torch.from_numpy(arr).permute(2, 0, 1) - mean) / std).unsqueeze(0).to(DEV, mdt)
+    if spec["kind"] == "dreamsim":
+        x = _cur["pre"](img).to(DEV)                      # DreamSim ships its own transform
+    else:
+        mean, std = spec["norm"]
+        arr = np.asarray(img).astype(np.float32) / 255
+        x = ((torch.from_numpy(arr).permute(2, 0, 1) - mean) / std).unsqueeze(0).to(DEV, mdt)
     feat, att = encode(model, spec, x)
     att_pct = 100.0 * att / att.sum()
     P = feat.shape[0]; g = int(round(P ** 0.5))
@@ -153,9 +170,10 @@ if gr is not None:
                 btn = gr.Button("Build tree + reveal", variant="primary")
                 gr.Markdown("<small>Compare encoder families: **DINOv3** (self-supervised, semantic), "
                             "**MAE** (pixel-reconstruction — more texture/appearance driven), "
-                            "**CLIP** (language-aligned — 16×16 grid, patch-14).<br>"
+                            "**CLIP** (language-aligned — 16×16 grid, patch-14), "
+                            "**DreamSim** (tuned on *human* similarity judgments).<br>"
                             "ViT-H+/7B are large — slow on the free CPU; **ViT-7B realistically needs a GPU Space**. "
-                            "First use of each model downloads its weights (one-time).</small>")
+                            "First use of each model downloads its weights (one-time; DreamSim is ~1 GB).</small>")
         with gr.Tab("Condensed tree (hover)"):
             out_tree = gr.HTML()
         with gr.Tab("Reveal animation"):

requirements.txt

@@ -5,3 +5,4 @@ scipy
 pillow
 numpy
 gradio
+dreamsim