Spaces:
Running
Running
| """ | |
| DINO Hierarchies — upload an image, get its condensed tree (hover) + reveal animation. | |
| Runs DINOv3 ViT-L/16, builds the same condensed tree we precomputed for the static site, | |
| and injects it into the existing JS widgets (no server-side rendering of the video). | |
| """ | |
| import os, io, json, html, base64, sys | |
| print("[app] starting import", flush=True) | |
| import numpy as np, torch | |
| try: | |
| import gradio as gr | |
| print("[app] gradio", gr.__version__, flush=True) | |
| except ImportError: | |
| gr = None # allows importing the pipeline without gradio (local testing) | |
| 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, ViTMAEModel, CLIPVisionModel | |
| RES, MS = 224, 4 | |
| 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 = { | |
| "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), | |
| "DreamSim (human-aligned)": dict(repo="dreamsim:dino_vitb16", kind="dreamsim", patch=16, norm=None), | |
| } | |
| DEFAULT_MODEL = next(iter(MODELS)) | |
| 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, "pre": None} | |
| def get_model(spec): | |
| repo = spec["repo"] | |
| if _cur["repo"] != repo: | |
| _cur["model"] = _cur["pre"] = None | |
| import gc; gc.collect() | |
| 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) | |
| 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): | |
| spec = MODELS.get(model_label, MODELS[DEFAULT_MODEL]) | |
| model = get_model(spec) | |
| mdt = next(model.parameters()).dtype | |
| img = image.convert("RGB").resize((RES, RES), Image.BICUBIC) | |
| 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)) | |
| fn = feat / (np.linalg.norm(feat, axis=1, keepdims=True) + 1e-8) | |
| root = to_tree(linkage(squareform(1.0 - fn @ fn.T, checks=False), method="average")) | |
| nodes = [] | |
| def add(cn, depth, parent): | |
| idx = len(nodes) | |
| nodes.append({"leaves": cn.pre_order(lambda v: v.id), "children": [], "depth": depth, | |
| "split_tau": None, "parent": parent}) | |
| cur, fo = cn, [] | |
| while not cur.is_leaf(): | |
| l, r = cur.left, cur.right | |
| if l.count >= MS and r.count >= MS: break | |
| small, big = (r, l) if l.count >= r.count else (l, r); fo.append(small); cur = big | |
| if not cur.is_leaf(): | |
| nodes[idx]["split_tau"] = round(1.0 - cur.dist, 3) | |
| for ch in (cur.left, cur.right): | |
| nodes[idx]["children"].append(add(ch, depth + 1, idx)) | |
| for f in fo: | |
| fi = len(nodes); nodes.append({"leaves": f.pre_order(lambda v: v.id), "children": [], | |
| "depth": depth + 1, "split_tau": None, "parent": idx}) | |
| nodes[idx]["children"].append(fi) | |
| return idx | |
| add(root, 0, -1) | |
| xpos = {}; cnt = [0] | |
| def setx(i): | |
| ch = nodes[i]["children"] | |
| if not ch: xpos[i] = cnt[0]; cnt[0] += 1 | |
| else: | |
| for c in ch: setx(c) | |
| xpos[i] = float(np.mean([xpos[c] for c in ch])) | |
| setx(0) | |
| nleaves = cnt[0]; maxdepth = max(n["depth"] for n in nodes) | |
| tree = {"g": g, "nleaves": nleaves, "maxdepth": maxdepth, | |
| "nodes": [{"id": i, "x": round(xpos[i], 3), "depth": n["depth"], "parent": n["parent"], | |
| "leaf": len(n["children"]) == 0, "n": len(n["leaves"]), | |
| "att": round(float(att_pct[n["leaves"]].sum()), 1), # % of total attention | |
| "attm": round(float(att_pct[n["leaves"]].mean()), 4), # mean attention % per patch (reveal order) | |
| "attmv": round(float(att[n["leaves"]].mean()), 5), # raw mean attention value per patch | |
| "tau": n["split_tau"], "patches": n["leaves"]} for i, n in enumerate(nodes)]} | |
| buf = io.BytesIO(); img.save(buf, "PNG") | |
| return tree, base64.b64encode(buf.getvalue()).decode() | |
| def iframe(widget_html, tree, img_b64, height): | |
| inject = (f'<script>window.TREE={json.dumps(tree)};' | |
| f'window.IMG_SRC="data:image/png;base64,{img_b64}";</script>') | |
| doc = widget_html.replace("<body>", "<body>" + inject, 1) | |
| return (f'<iframe srcdoc="{html.escape(doc, quote=True)}" ' | |
| f'style="width:100%;height:{height}px;border:none;background:#0f1117;border-radius:8px"></iframe>') | |
| def process(image, model_label): | |
| if image is None: | |
| return "<p style='color:#999'>Upload an image first.</p>", "" | |
| tree, img_b64 = build_tree(image, model_label) | |
| return iframe(TREE_HTML, tree, img_b64, 660), iframe(REVEAL_HTML, tree, img_b64, 700) | |
| if gr is not None: | |
| with gr.Blocks(title="Visual Tree", theme=gr.themes.Base()) as demo: | |
| gr.Markdown("# 🌳 Visual Tree\nUpload an image → get its **condensed tree** (hover a node) " | |
| "and the **reveal animation** (best-first by CLS attention).") | |
| 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="Encoder") | |
| 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), " | |
| "**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; DreamSim is ~1 GB).</small>") | |
| with gr.Tab("Condensed tree (hover)"): | |
| out_tree = gr.HTML() | |
| with gr.Tab("Reveal animation"): | |
| out_rev = gr.HTML() | |
| btn.click(process, [inp, model_sel], [out_tree, out_rev]) | |
| inp.upload(process, [inp, model_sel], [out_tree, out_rev]) | |
| print("[app] launching gradio", flush=True) | |
| demo.queue().launch() # module-level: HF Spaces imports this file, so no __main__ guard | |