"""ViT-Up: Faithful Feature Upsampling for Vision Transformers. Interactive demo that loads the ViT-Up feature upsampler, extracts dense features from an input image at a user-selected output resolution, and visualises them via a 3-component PCA projection to RGB. The DINOv3 backbone checkpoint on Hugging Face is gated, so this demo loads the equivalent pretrained weights from the non-gated timm mirror (`timm/vit_small_plus_patch16_dinov3.lvd1689m`) and maps them into the same ``DINOv3ViT`` module structure the ViT-Up code expects. The ViT-Up LoRA adapters and upsampler head are then loaded from ``Krispin/vit-up``. """ import os os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") import spaces # MUST come before torch / any CUDA-touching import import sys import math from pathlib import Path from typing import Any, Dict, List, Optional import torch import torch.nn as nn import torch.nn.functional as F import numpy as np from PIL import Image, ImageOps import gradio as gr from huggingface_hub import hf_hub_download from safetensors.torch import load_file as load_safetensors import torchvision.transforms.v2 as T # --------------------------------------------------------------------------- # Config constants — DINOv3-S+ variant # --------------------------------------------------------------------------- BACKBONE_TIMM_REPO = "timm/vit_small_plus_patch16_dinov3.lvd1689m" VITUP_WEIGHTS_REPO = "Krispin/vit-up" VITUP_WEIGHTS_FILE = "vit_up_dinov3_splus.safetensors" HIDDEN_SIZE = 384 NUM_LAYERS = 12 NUM_HEADS = 6 INTERMEDIATE_SIZE = 1536 PATCH_SIZE = 16 NUM_REGISTER_TOKENS = 4 IMAGE_SIZE = 448 LAYER_INDICES = [0, 2, 4, 6, 8, 10, 12] RESNET_MEAN = torch.tensor([0.485, 0.456, 0.406]) RESNET_STD = torch.tensor([0.229, 0.224, 0.225]) # vit_up package is included in the Space repo root from transformers import DINOv3ViTConfig from vit_up.layers.backbones.dinov3_vit import DINOv3ViT from vit_up.model.vit_up import ViTUp from vit_up.utils.state_dict_migration import migrate_vit_up_state_dict_keys from peft import LoraConfig, get_peft_model # --------------------------------------------------------------------------- # Weight mapping: timm -> DINOv3ViT # --------------------------------------------------------------------------- def _map_timm_to_dinov3(timm_sd: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: """Convert a timm ViT state-dict to the DINOv3ViT module key names.""" mapped: Dict[str, torch.Tensor] = {} for key, val in timm_sd.items(): if key == "cls_token": mapped["embeddings.cls_token"] = val elif key == "reg_token": mapped["embeddings.register_tokens"] = val elif key == "patch_embed.proj.weight": mapped["embeddings.patch_embeddings.weight"] = val elif key == "patch_embed.proj.bias": mapped["embeddings.patch_embeddings.bias"] = val elif key.startswith("blocks.") and key.endswith(".attn.qkv.weight"): idx = int(key.split(".")[1]) qkv = val # (3*hidden, hidden) but timm uses fused qkv q, k, v = qkv.chunk(3, dim=0) mapped[f"layer.{idx}.attention.q_proj.weight"] = q mapped[f"layer.{idx}.attention.k_proj.weight"] = k mapped[f"layer.{idx}.attention.v_proj.weight"] = v elif key.startswith("blocks.") and key.endswith(".attn.qkv.bias"): idx = int(key.split(".")[1]) qkv = val if val is not None and val.numel() > 0: q, k, v = qkv.chunk(3, dim=0) mapped[f"layer.{idx}.attention.q_proj.bias"] = q mapped[f"layer.{idx}.attention.k_proj.bias"] = k mapped[f"layer.{idx}.attention.v_proj.bias"] = v elif key.startswith("blocks.") and ".attn.proj." in key: idx = int(key.split(".")[1]) suffix = key.split(".attn.proj.")[-1] # weight or bias mapped[f"layer.{idx}.attention.o_proj.{suffix}"] = val elif key.startswith("blocks.") and ".norm1." in key: idx = int(key.split(".")[1]) suffix = key.split(".norm1.")[-1] mapped[f"layer.{idx}.norm1.{suffix}"] = val elif key.startswith("blocks.") and ".norm2." in key: idx = int(key.split(".")[1]) suffix = key.split(".norm2.")[-1] mapped[f"layer.{idx}.norm2.{suffix}"] = val elif key.startswith("blocks.") and ".mlp.fc1_g." in key: idx = int(key.split(".")[1]) suffix = key.split(".mlp.fc1_g.")[-1] mapped[f"layer.{idx}.mlp.gate_proj.{suffix}"] = val elif key.startswith("blocks.") and ".mlp.fc1_x." in key: idx = int(key.split(".")[1]) suffix = key.split(".mlp.fc1_x.")[-1] mapped[f"layer.{idx}.mlp.up_proj.{suffix}"] = val elif key.startswith("blocks.") and ".mlp.fc2." in key: idx = int(key.split(".")[1]) suffix = key.split(".mlp.fc2.")[-1] mapped[f"layer.{idx}.mlp.down_proj.{suffix}"] = val elif key.startswith("blocks.") and key.endswith(".gamma_1"): idx = int(key.split(".")[1]) mapped[f"layer.{idx}.layer_scale1.lambda1"] = val elif key.startswith("blocks.") and key.endswith(".gamma_2"): idx = int(key.split(".")[1]) mapped[f"layer.{idx}.layer_scale2.lambda1"] = val elif key == "norm.weight": mapped["norm.weight"] = val elif key == "norm.bias": mapped["norm.bias"] = val # pos_embed is handled by RoPE — skip return mapped # --------------------------------------------------------------------------- # Build the backbone from config + timm weights + LoRA # --------------------------------------------------------------------------- def _build_backbone(device: str, dtype: torch.dtype) -> DINOv3ViT: config = DINOv3ViTConfig( hidden_size=HIDDEN_SIZE, num_hidden_layers=NUM_LAYERS, num_attention_heads=NUM_HEADS, intermediate_size=INTERMEDIATE_SIZE, patch_size=PATCH_SIZE, image_size=IMAGE_SIZE, num_register_tokens=NUM_REGISTER_TOKENS, use_gated_mlp=True, layerscale_value=1e-5, query_bias=False, key_bias=False, value_bias=False, proj_bias=True, mlp_bias=True, ) config._attn_implementation = "eager" backbone = DINOv3ViT(config) # Load timm weights timm_safetensors_path = hf_hub_download( BACKBONE_TIMM_REPO, "model.safetensors" ) timm_sd = load_safetensors(timm_safetensors_path, device="cpu") mapped_sd = _map_timm_to_dinov3(timm_sd) missing, unexpected = backbone.load_state_dict(mapped_sd, strict=False) # embeddings.mask_token won't be in timm weights — that's fine real_missing = [k for k in missing if "mask_token" not in k] if real_missing: print(f"[WARNING] Missing backbone keys after timm load: {real_missing[:10]}") print(f"[INFO] Loaded backbone from timm: {len(mapped_sd)} tensors mapped") # Apply LoRA lora_config = LoraConfig( r=16, lora_alpha=32, lora_dropout=0.05, bias="none", target_modules=[ "patch_embeddings", "q_proj", "k_proj", "v_proj", "o_proj", ], ) backbone = get_peft_model(backbone, lora_config) backbone = backbone.to(device=device, dtype=dtype).eval() return backbone # --------------------------------------------------------------------------- # Build the ViT-Up model from config # --------------------------------------------------------------------------- def _build_vit_up(device: str, dtype: torch.dtype) -> ViTUp: """Instantiate the ViTUp upsampler from the config tree (same as the repo).""" from vit_up.layers.query_encoder import QueryEncoder from vit_up.layers.pos_enc import FourierPositionalEncoding from vit_up.layers.continuous_rope import ContinuousRoPE2D from vit_up.layers.smart_module_list import SmartModuleList from vit_up.layers.mlp import SimpleMLP from vit_up.layers.cross_attention import CrossAttention from vit_up.layers.film import SimpleFiLMV2 dim = HIDDEN_SIZE query_embedding = QueryEncoder( layer_index=0, img_in_size=3584, window_size=0, out_proj_module=None, ) rel_pos_enc = FourierPositionalEncoding(num_bands=16, max_resolution=10.0) q_rope_embeddings = ContinuousRoPE2D(dim=64, base=100.0, scale=2 * math.pi) vit_up_blocks = SmartModuleList( n_blocks=6, block_class_path="vit_up.model.vit_up.ViTUpBlock", block_init_args={ "dim": dim, "dim_h": dim, "transition_mlp": SimpleMLP( dims=[dim, dim * 2, dim], activation="gelu", input_layernorm=True, use_residual=True, ), "cross_attention": CrossAttention( dim=dim, num_heads=NUM_HEADS, cross_attn_window_size=32, qkv_bias=True, attn_dropout=0.0, proj_dropout=0.0, ), "featx": SimpleFiLMV2( input_module=nn.LayerNorm(dim), gamma_beta_mlp=SimpleMLP( dims=[66, dim, dim * 2], activation="gelu", zero_init_last=True, ), post_mlp=SimpleMLP( dims=[dim, dim * 4, dim], activation="gelu", input_layernorm=True, use_residual=False, ), ), "mlp": SimpleMLP( dims=[dim, dim * 4, dim], activation="gelu", ), }, ) decoder_mlp = SmartModuleList( n_blocks=7, block_class_path="vit_up.layers.mlp.SimpleMLP", block_init_args={ "input_layernorm": True, "dims": [dim, dim], }, ) vit_up = ViTUp( layer_indices=LAYER_INDICES, query_embedding=query_embedding, rel_pos_enc=rel_pos_enc, vit_up_blocks=vit_up_blocks, decoder_mlp=decoder_mlp, q_rope_embeddings=q_rope_embeddings, ) vit_up = vit_up.to(device=device, dtype=dtype).eval() return vit_up # --------------------------------------------------------------------------- # Load ViT-Up + LoRA weights from the safetensors checkpoint # --------------------------------------------------------------------------- def _load_vit_up_weights( backbone: nn.Module, vit_up: ViTUp, device: str, ) -> None: """Load the combined LoRA + ViT-Up weights from Krispin/vit-up.""" weights_path = hf_hub_download(VITUP_WEIGHTS_REPO, VITUP_WEIGHTS_FILE) state_dict = load_safetensors(weights_path, device="cpu") backbone_sd: Dict[str, torch.Tensor] = {} vit_up_sd: Dict[str, torch.Tensor] = {} for key, val in state_dict.items(): if key.startswith("backbone."): backbone_sd[key.removeprefix("backbone.")] = val else: vit_up_sd[key] = val # Load backbone LoRA weights missing_b, unexpected_b = backbone.load_state_dict(backbone_sd, strict=False) print(f"[INFO] Loaded backbone LoRA: {len(backbone_sd)} tensors, " f"missing={len(missing_b)}, unexpected={len(unexpected_b)}") # Load ViT-Up weights (with key migration) migrated_vit_up_sd = migrate_vit_up_state_dict_keys(vit_up_sd) missing_v, unexpected_v = vit_up.load_state_dict(migrated_vit_up_sd, strict=False) print(f"[INFO] Loaded ViT-Up: {len(migrated_vit_up_sd)} tensors, " f"missing={len(missing_v)}, unexpected={len(unexpected_v)}") if missing_v: print(f" Missing ViT-Up keys: {missing_v[:10]}") # --------------------------------------------------------------------------- # PCA utilities (from the repo's correspondence.py) # --------------------------------------------------------------------------- def _fit_pca(tokens_nc: torch.Tensor, k: int = 3) -> dict: """Fit a simple PCA on (N, C) feature tokens.""" tokens = tokens_nc.float() mean = tokens.mean(dim=0) centered = tokens - mean _, singular_values, vh = torch.linalg.svd(centered, full_matrices=False) components = vh[:k].T projected = centered @ components color_min = projected.amin(dim=0) color_max = projected.amax(dim=0) flat = torch.isclose(color_max, color_min) color_max = torch.where(flat, color_min + 1.0, color_max) return { "pca_eig": components, "pca_mean": mean, "pca_color_min": color_min, "pca_color_max": color_max, } def _apply_pca_rgb(feats_hwc: torch.Tensor, pca_data: dict) -> torch.Tensor: h, w, c = feats_hwc.shape tokens = feats_hwc.float().reshape(-1, c) mean = pca_data["pca_mean"].to(device=tokens.device, dtype=tokens.dtype) components = pca_data["pca_eig"].to(device=tokens.device, dtype=tokens.dtype) color_min = pca_data["pca_color_min"].to(device=tokens.device, dtype=tokens.dtype) color_max = pca_data["pca_color_max"].to(device=tokens.device, dtype=tokens.dtype) projected = (tokens - mean) @ components rgb = (projected - color_min.view(1, -1)) / (color_max - color_min).view(1, -1).add(1e-8) rgb = rgb.clamp(0.0, 1.0).mul(255.0).to(torch.uint8) return rgb.reshape(h, w, 3) # --------------------------------------------------------------------------- # Image utilities # --------------------------------------------------------------------------- def pad_image_to_square(img: Image.Image) -> Image.Image: w, h = img.size if w == h: return img max_side = max(w, h) if w > h: py = (w - h) // 2 return ImageOps.expand(img, border=(0, py), fill=0) else: px = (h - w) // 2 return ImageOps.expand(img, border=(px, 0), fill=0) def crop_feature_square_to_image_aspect( feat_img: Image.Image, original_size: tuple, ) -> Image.Image: width, height = original_size max_size = max(width, height) px, py = (0, 0) if width > height: py = (width - height) // 2 elif height > width: px = (height - width) // 2 scale_x = feat_img.width / max_size scale_y = feat_img.height / max_size left = int(round(px * scale_x)) top = int(round(py * scale_y)) right = int(round((px + width) * scale_x)) bottom = int(round((py + height) * scale_y)) return feat_img.crop((left, top, right, bottom)) # --------------------------------------------------------------------------- # Build the full model at module scope # --------------------------------------------------------------------------- print("[INFO] Building ViT-Up model...") DEVICE = "cuda" DTYPE = torch.bfloat16 backbone = _build_backbone(DEVICE, DTYPE) vit_up = _build_vit_up(DEVICE, DTYPE) _load_vit_up_weights(backbone, vit_up, DEVICE) backbone = backbone.eval() vit_up = vit_up.eval() print("[INFO] Model ready.") # --------------------------------------------------------------------------- # Inference # --------------------------------------------------------------------------- def _prepare_image(img: Image.Image) -> torch.Tensor: """Pad to square, resize, normalise — return (1, 3, H, W) on device.""" img_square = pad_image_to_square(img.convert("RGB")) transform = T.Compose([ T.ToImage(), T.Resize((IMAGE_SIZE, IMAGE_SIZE), interpolation=T.InterpolationMode.BILINEAR, antialias=True), T.ToDtype(torch.float32, scale=True), T.Normalize(mean=RESNET_MEAN, std=RESNET_STD), ]) return transform(img_square).unsqueeze(0).to(DEVICE) def _compute_query_coords(out_size: int) -> torch.Tensor: coords = torch.linspace(0.5, out_size - 0.5, out_size) / out_size grid_y, grid_x = torch.meshgrid(coords, coords, indexing="ij") return torch.stack((grid_x, grid_y), dim=-1).reshape(1, -1, 2) @spaces.GPU(duration=120) def extract_and_visualize( input_image: Image.Image, output_resolution: int, ) -> tuple[Image.Image, Image.Image, str]: """Extract dense ViT-Up features and visualise them via PCA. Args: input_image: Input PIL image. output_resolution: Output feature map resolution (pixels per side). Returns: Tuple of (pca_visualization, input_resized, info_text). """ if input_image is None: return None, None, "Please provide an input image." out_size = int(output_resolution) orig_w, orig_h = input_image.size # Prepare input pixel_values = _prepare_image(input_image) # Compute cache data (backbone hidden states) with torch.no_grad(), torch.autocast(device_type="cuda", dtype=DTYPE): cache_data = vit_up.compute_cache_data( pixel_values=pixel_values, backbone=backbone, hidden_layer_img_size=IMAGE_SIZE, ) # Query coords for dense output query_coords = _compute_query_coords(out_size).to(DEVICE, dtype=DTYPE) # Extract features chunk_size = 4096 q_chunks = [] for q_start in range(0, query_coords.shape[1], chunk_size): q_end = min(q_start + chunk_size, query_coords.shape[1]) q_chunk = vit_up( pixel_values=None, q_xy_normalized=query_coords[:, q_start:q_end, :], cache_data=cache_data, ) q_chunks.append(q_chunk[-1]) # final layer features = torch.cat(q_chunks, dim=1) # (1, out_size*out_size, D) features_hwc = features[0].reshape(out_size, out_size, -1).float().cpu() # PCA pca_data = _fit_pca(features_hwc.reshape(-1, features_hwc.shape[-1]), k=3) pca_rgb = _apply_pca_rgb(features_hwc, pca_data) pca_img = Image.fromarray(pca_rgb.numpy().astype(np.uint8), mode="RGB") # Crop to original aspect ratio pca_img = crop_feature_square_to_image_aspect(pca_img, (orig_w, orig_h)) # Resize for display display_w, display_h = orig_w, orig_h max_display = 512 if max(display_w, display_h) > max_display: scale = max_display / max(display_w, display_h) display_w = int(display_w * scale) display_h = int(display_h * scale) pca_display = pca_img.resize((display_w, display_h), Image.Resampling.NEAREST) # Also create a resized input for side-by-side comparison input_display = input_image.convert("RGB").resize((display_w, display_h), Image.Resampling.LANCZOS) info = (f"Feature dim: {features_hwc.shape[-1]} | " f"Output resolution: {out_size}x{out_size} | " f"Total query points: {out_size * out_size}") return pca_display, input_display, info # --------------------------------------------------------------------------- # Gradio UI # --------------------------------------------------------------------------- CSS = """ #col-container { max-width: 1100px; margin: 0 auto; } .dark .gradio-container { color: var(--body-text-color); } """ with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo: gr.Markdown("# ViT-Up: Faithful Feature Upsampling for Vision Transformers") gr.Markdown( "Upload an image to extract dense DINOv3 features at arbitrary resolution " "via the ViT-Up feature upsampler. The PCA visualization shows the " "3 principal components of the upsampled feature map as RGB." ) gr.Markdown( "[Paper](https://huggingface.co/papers/2606.14024) | " "[GitHub](https://github.com/krispinwandel/vit-up) | " "[Model Weights](https://huggingface.co/Krispin/vit-up)" ) with gr.Row(): with gr.Column(): input_img = gr.Image(label="Input Image", type="pil") with gr.Accordion("Advanced settings", open=False): out_res = gr.Slider( label="Output resolution (pixels per side)", minimum=28, maximum=224, value=112, step=28, ) run_btn = gr.Button("Extract Features", variant="primary") with gr.Column(): pca_out = gr.Image(label="PCA Feature Visualization") input_display = gr.Image(label="Input (resized)") info_text = gr.Textbox(label="Info", interactive=False) run_btn.click( fn=extract_and_visualize, inputs=[input_img, out_res], outputs=[pca_out, input_display, info_text], api_name="extract_features", ) gr.Examples( examples=[ ["city_with_cars.png", 112], ["fruit_store.png", 112], ], inputs=[input_img, out_res], outputs=[pca_out, input_display, info_text], fn=extract_and_visualize, cache_examples=True, cache_mode="lazy", ) demo.launch(mcp_server=True)