Add MAE and CLIP encoders alongside DINOv3 (per-model prefix/grid/normalization; MAE ids_restore un-shuffle)

app.py

@@ -15,50 +15,72 @@ from PIL import Image
 print("[app] transformers import...", flush=True)
 from scipy.cluster.hierarchy import linkage, to_tree
 from scipy.spatial.distance import squareform
-from transformers import AutoModel
+from transformers import AutoModel, ViTMAEModel, CLIPVisionModel
 
 RES, MS = 224, 4
-GRID = RES // 16          # 14x14 = 196 patches
 DEV = "cuda" if torch.cuda.is_available() else "cpu"
 TOKEN = os.environ.get("HF_TOKEN")
 
+IMNET = (torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1),
+         torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1))
+CLIPN = (torch.tensor([0.48145466, 0.4578275, 0.40821073]).view(3, 1, 1),
+         torch.tensor([0.26862954, 0.26130258, 0.27577711]).view(3, 1, 1))
+
+# kind: how to strip prefix tokens / un-shuffle;  patch: 224/patch = grid side
 MODELS = {
-    "ViT-L · 0.3B (fast)":       "facebook/dinov3-vitl16-pretrain-lvd1689m",
-    "ViT-H+ · 0.84B (slower)":   "facebook/dinov3-vith16plus-pretrain-lvd1689m",
-    "ViT-7B · 6.7B (needs GPU)": "facebook/dinov3-vit7b16-pretrain-lvd1689m",
+    "DINOv3 ViT-L · 0.3B (fast)":       dict(repo="facebook/dinov3-vitl16-pretrain-lvd1689m",     kind="dinov3", patch=16, norm=IMNET),
+    "DINOv3 ViT-H+ · 0.84B (slower)":   dict(repo="facebook/dinov3-vith16plus-pretrain-lvd1689m", kind="dinov3", patch=16, norm=IMNET),
+    "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),
 }
 DEFAULT_MODEL = next(iter(MODELS))
 
-MEAN = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1)
-STD = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1)
 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}
-def get_model(repo):
+def get_model(spec):
+    repo = spec["repo"]
     if _cur["repo"] != repo:
         _cur["model"] = 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)
-        _cur["model"] = AutoModel.from_pretrained(repo, dtype=dt, token=TOKEN,
-                                                  attn_implementation="eager").eval().to(DEV)
+        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."""
+    with torch.no_grad():
+        o = model(x, output_attentions=True)
+    h, a = o.last_hidden_state, o.attentions[-1][0]            # a: (heads, S, S)
+    if spec["kind"] == "mae":
+        ids = o.ids_restore[0]                                 # MAE shuffles tokens -> un-shuffle
+        feat, att = h[0, 1:][ids], a[:, 0, 1:][:, ids].mean(0)
+    else:
+        g = RES // spec["patch"]
+        prefix = h.shape[1] - g * g                            # strip CLS (+ registers for DINOv3)
+        feat, att = h[0, prefix:], a[:, 0, prefix:].mean(0)
+    return feat.float().cpu().numpy(), att.float().cpu().numpy()
+
+
 def build_tree(image, model_label):
-    model = get_model(MODELS.get(model_label, MODELS[DEFAULT_MODEL]))
+    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)
-    with torch.no_grad():
-        o = model(x, output_attentions=True)
-    prefix = o.last_hidden_state.shape[1] - GRID * GRID            # strip CLS + register tokens
-    feat = o.last_hidden_state[0, prefix:].float().cpu().numpy()
-    att = o.attentions[-1][0][:, 0, prefix:].mean(0).float().cpu().numpy()
+    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))
     fn = feat / (np.linalg.norm(feat, axis=1, keepdims=True) + 1e-8)
@@ -127,9 +149,12 @@ if gr is not None:
         with gr.Row():
             inp = gr.Image(type="pil", label="Upload an image", height=300)
             with gr.Column(scale=0):
-                model_sel = gr.Dropdown(choices=list(MODELS.keys()), value=DEFAULT_MODEL, label="DINOv3 model")
+                model_sel = gr.Dropdown(choices=list(MODELS.keys()), value=DEFAULT_MODEL, label="Encoder")
                 btn = gr.Button("Build tree + reveal", variant="primary")
-                gr.Markdown("<small>ViT-H+/7B are large — slow on the free CPU; **ViT-7B realistically needs a GPU Space**. "
+                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>"
+                            "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>")
         with gr.Tab("Condensed tree (hover)"):
             out_tree = gr.HTML()