Upload folder using huggingface_hub

README.md

@@ -1,13 +1,31 @@
 ---
-title: Uncha Hyperbolic Vlm
-emoji: 📊
-colorFrom: gray
-colorTo: purple
+title: UNCHA Hyperbolic VLM
+emoji: 🔮
+colorFrom: pink
+colorTo: gray
 sdk: gradio
 sdk_version: 6.20.0
-python_version: '3.12'
 app_file: app.py
-pinned: false
+short_description: Hyperbolic vision-language zero-shot classification
+python_version: "3.12"
+startup_duration_timeout: 30m
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# UNCHA: Uncertainty-guided Compositional Hyperbolic Alignment
+
+This Space demonstrates **zero-shot image classification** using the UNCHA model,
+a hyperbolic vision-language model that embeds images and text in Lorentz (hyperbolic)
+space to capture hierarchical part-whole relationships.
+
+## How it works
+
+1. Upload an image (or use an example)
+2. Provide a comma-separated list of candidate labels
+3. The model encodes the image and each label's text prompts into hyperbolic space
+4. Classification scores are computed via the Lorentzian inner product
+
+## References
+
+- [Paper (arXiv)](https://arxiv.org/abs/2603.22042)
+- [GitHub](https://github.com/jeeit17/UNCHA)
+- [Model weights](https://huggingface.co/hayeonkim/uncha)

app.py

@@ -0,0 +1,485 @@
+"""
+UNCHA: Uncertainty-guided Compositional Hyperbolic Alignment
+Zero-shot image classification demo using hyperbolic (Lorentz) embeddings.
+
+Paper: https://arxiv.org/abs/2603.22042
+Code:  https://github.com/jeeit17/UNCHA
+"""
+
+import os
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+
+import spaces  # MUST come before torch / any CUDA-touching import
+
+import gzip
+import html
+import math
+import re
+from collections import OrderedDict
+from pathlib import Path
+
+import numpy as np
+import torch
+from torch import nn
+from torch.nn import functional as F
+import timm
+import gradio as gr
+from PIL import Image
+
+import torchvision.transforms as T
+
+# ---------------------------------------------------------------------------
+# Tokenizer (adapted from UNCHA/CLIP BPE tokenizer)
+# ---------------------------------------------------------------------------
+
+class Tokenizer:
+    """Byte-Pair Encoding tokenizer compatible with CLIP / UNCHA checkpoints."""
+
+    def __init__(self, bpe_path: str | Path | None = None):
+        bs = (
+            list(range(ord("!"), ord("~") + 1))
+            + list(range(ord("\xa1"), ord("\xac") + 1))
+            + list(range(ord("\xae"), ord("\xff") + 1))
+        )
+        self.byte_encoder = {b: chr(b) for b in bs}
+        n = 0
+        for b in range(2**8):
+            if b not in self.byte_encoder:
+                self.byte_encoder[b] = chr(2**8 + n)
+                n += 1
+
+        if bpe_path is None:
+            bpe_path = Path(__file__).resolve().parent / "bpe_simple_vocab_16e6.txt.gz"
+        merges = gzip.open(bpe_path).read().decode("utf-8").split("\n")
+        merges = merges[1 : 49152 - 256 - 2 + 1]
+        merges = [tuple(merge.split()) for merge in merges]
+        vocab = list(self.byte_encoder.values())
+        vocab = vocab + [v + "</w>" for v in vocab]
+        for merge in merges:
+            vocab.append("".join(merge))
+        vocab.extend(["<|startoftext|>", ""])
+
+        self.encoder = dict(zip(vocab, range(len(vocab))))
+        self.bpe_ranks = dict(zip(merges, range(len(merges))))
+        self.cache = {"<|startoftext|>": "<|startoftext|>", "": ""}
+        self.pat = re.compile(
+            r"""<\|startoftext\|>|<\|endoftext\|>|'s|'t|'re|'ve|'m|'ll|'d|[\p{L}]+|[\p{N}]|[^\s\p{L}\p{N}]+""",
+            re.IGNORECASE,
+        )
+
+    def __call__(self, text):
+        import ftfy
+
+        text_list = [text] if isinstance(text, str) else text
+        token_tensors = []
+        for text in text_list:
+            bpe_tokens = []
+            text = ftfy.fix_text(text)
+            text = html.unescape(html.unescape(text))
+            text = re.sub(r"\s+", " ", text)
+            text = text.strip().lower()
+            for token in re.findall(self.pat, text):
+                token = "".join(self.byte_encoder[b] for b in token.encode("utf-8"))
+                bpe_tokens.extend(
+                    self.encoder[bpe_token]
+                    for bpe_token in self.bpe(token).split(" ")
+                )
+            sot = self.encoder["<|startoftext|>"]
+            eot = self.encoder[""]
+            bpe_tokens = [sot, *bpe_tokens, eot]
+            token_tensors.append(torch.IntTensor(bpe_tokens))
+        return token_tensors
+
+    @staticmethod
+    def get_pairs(word):
+        pairs = set()
+        prev_char = word[0]
+        for char in word[1:]:
+            pairs.add((prev_char, char))
+            prev_char = char
+        return pairs
+
+    def bpe(self, token):
+        if token in self.cache:
+            return self.cache[token]
+        word = tuple(token[:-1]) + (token[-1] + "</w>",)
+        pairs = self.get_pairs(word)
+        if not pairs:
+            return token + "</w>"
+        while True:
+            bigram = min(pairs, key=lambda pair: self.bpe_ranks.get(pair, float("inf")))
+            if bigram not in self.bpe_ranks:
+                break
+            first, second = bigram
+            new_word = []
+            i = 0
+            while i < len(word):
+                try:
+                    j = word.index(first, i)
+                    new_word.extend(word[i:j])
+                    i = j
+                except ValueError:
+                    new_word.extend(word[i:])
+                    break
+                if word[i] == first and i < len(word) - 1 and word[i + 1] == second:
+                    new_word.append(first + second)
+                    i += 2
+                else:
+                    new_word.append(word[i])
+                    i += 1
+            new_word = tuple(new_word)
+            word = new_word
+            if len(word) == 1:
+                break
+            else:
+                pairs = self.get_pairs(word)
+        word = " ".join(word)
+        self.cache[token] = word
+        return word
+
+
+# ---------------------------------------------------------------------------
+# Lorentz model hyperbolic operations (adapted from UNCHA/meru)
+# ---------------------------------------------------------------------------
+
+def pairwise_inner(x, y, curv=1.0):
+    x_time = torch.sqrt(1 / curv + torch.sum(x**2, dim=-1, keepdim=True))
+    y_time = torch.sqrt(1 / curv + torch.sum(y**2, dim=-1, keepdim=True))
+    return x @ y.T - x_time @ y_time.T
+
+
+def exp_map0(x, curv=1.0, eps=1e-8):
+    rc_xnorm = curv**0.5 * torch.norm(x, dim=-1, keepdim=True)
+    sinh_input = torch.clamp(rc_xnorm, min=eps, max=math.asinh(2**15))
+    return torch.sinh(sinh_input) * x / torch.clamp(rc_xnorm, min=eps)
+
+
+# ---------------------------------------------------------------------------
+# Text encoder (adapted from UNCHA TransformerTextEncoder)
+# ---------------------------------------------------------------------------
+
+class _TransformerBlock(nn.Module):
+    def __init__(self, d_model, n_head):
+        super().__init__()
+        self.attn = nn.MultiheadAttention(d_model, n_head, batch_first=True)
+        self.ln_1 = nn.LayerNorm(d_model)
+        self.mlp = nn.Sequential(
+            OrderedDict([
+                ("c_fc", nn.Linear(d_model, d_model * 4)),
+                ("gelu", nn.GELU()),
+                ("c_proj", nn.Linear(d_model * 4, d_model)),
+            ])
+        )
+        self.ln_2 = nn.LayerNorm(d_model)
+
+    def forward(self, x, attn_mask=None):
+        lx = self.ln_1(x)
+        ax = self.attn(lx, lx, lx, need_weights=False, attn_mask=attn_mask)[0]
+        x = x + ax
+        x = x + self.mlp(self.ln_2(x))
+        return x
+
+
+class TransformerTextEncoder(nn.Module):
+    def __init__(self, arch="L12_W512", vocab_size=49408, context_length=77):
+        super().__init__()
+        self.vocab_size = vocab_size
+        self.context_length = context_length
+        self.layers = int(re.search(r"L(\d+)", arch).group(1))
+        self.width = int(re.search(r"W(\d+)", arch).group(1))
+        _attn = re.search(r"A(\d+)", arch)
+        self.heads = int(_attn.group(1)) if _attn else self.width // 64
+
+        self.token_embed = nn.Embedding(vocab_size, self.width)
+        self.posit_embed = nn.Parameter(torch.empty(context_length, self.width))
+        _resblocks = [_TransformerBlock(self.width, self.heads) for _ in range(self.layers)]
+        self.resblocks = nn.ModuleList(_resblocks)
+        self.ln_final = nn.LayerNorm(self.width)
+
+        attn_mask = torch.triu(
+            torch.full((context_length, context_length), float("-inf")), diagonal=1
+        )
+        self.register_buffer("attn_mask", attn_mask.bool())
+
+        nn.init.normal_(self.token_embed.weight, std=0.02)
+        nn.init.normal_(self.posit_embed.data, std=0.01)
+        out_proj_std = (2 * self.width * self.layers) ** -0.5
+        for block in self.resblocks:
+            nn.init.normal_(block.attn.in_proj_weight, std=self.width**-0.5)
+            nn.init.normal_(block.attn.out_proj.weight, std=out_proj_std)
+            nn.init.normal_(block.mlp[0].weight, std=(2 * self.width) ** -0.5)
+            nn.init.normal_(block.mlp[2].weight, std=out_proj_std)
+
+
+def build_timm_vit(arch="vit_base_patch16_224", global_pool="token",
+                   use_sincos2d_pos=True):
+    model = timm.create_model(
+        arch, num_classes=0, global_pool=global_pool,
+        class_token=global_pool == "token", norm_layer=nn.LayerNorm,
+    )
+    model.width = model.embed_dim
+    if use_sincos2d_pos:
+        h, w = model.patch_embed.grid_size
+        grid_w = torch.arange(w, dtype=torch.float32)
+        grid_h = torch.arange(h, dtype=torch.float32)
+        grid_w, grid_h = torch.meshgrid(grid_w, grid_h)
+        pos_dim = model.embed_dim // 4
+        omega = torch.arange(pos_dim, dtype=torch.float32) / pos_dim
+        omega = 1.0 / (10000.0**omega)
+        out_w = torch.einsum("m,d->md", [grid_w.flatten(), omega])
+        out_h = torch.einsum("m,d->md", [grid_h.flatten(), omega])
+        pos_emb = torch.cat(
+            [torch.sin(out_w), torch.cos(out_w), torch.sin(out_h), torch.cos(out_h)],
+            dim=1,
+        )[None, :, :]
+        if global_pool == "token":
+            pe_token = torch.zeros([1, 1, model.embed_dim], dtype=torch.float32)
+            pos_emb = torch.cat([pe_token, pos_emb], dim=1)
+        model.pos_embed.data.copy_(pos_emb)
+        model.pos_embed.requires_grad = False
+    return model
+
+
+# ---------------------------------------------------------------------------
+# UNCHA model (inference-only)
+# ---------------------------------------------------------------------------
+
+class UNCHAModel(nn.Module):
+    """Inference-only UNCHA model: hyperbolic image-text alignment."""
+
+    def __init__(self, embed_dim=512, visual_arch="vit_base_patch16_224",
+                 text_arch="L12_W512", vocab_size=49408, context_length=77):
+        super().__init__()
+        self.visual = build_timm_vit(arch=visual_arch)
+        self.textual = TransformerTextEncoder(
+            arch=text_arch, vocab_size=vocab_size, context_length=context_length
+        )
+        self.embed_dim = embed_dim
+        self.visual_proj = nn.Linear(self.visual.width, embed_dim, bias=False)
+        self.textual_proj = nn.Linear(self.textual.width, embed_dim, bias=False)
+        self.logit_scale = nn.Parameter(torch.tensor(1 / 0.07).log())
+        self.curv = nn.Parameter(torch.tensor(1.0).log())
+        self.visual_alpha = nn.Parameter(torch.tensor(embed_dim**-0.5).log())
+        self.textual_alpha = nn.Parameter(torch.tensor(embed_dim**-0.5).log())
+        self.tokenizer = Tokenizer()
+        self.register_buffer("pixel_mean", torch.tensor((0.485, 0.456, 0.406)).view(-1, 1, 1))
+        self.register_buffer("pixel_std", torch.tensor((0.229, 0.224, 0.225)).view(-1, 1, 1))
+
+    @property
+    def device(self):
+        return self.logit_scale.device
+
+    def encode_image(self, images, project=True):
+        images = (images - self.pixel_mean) / self.pixel_std
+        feats = self.visual(images)
+        feats = self.visual_proj(feats)
+        if project:
+            feats = feats * self.visual_alpha.exp()
+            with torch.autocast(self.device.type, dtype=torch.float32):
+                feats = exp_map0(feats, self.curv.exp())
+        return feats
+
+    def encode_text(self, tokens, project=True):
+        context_len = self.textual.context_length
+        batch_size = len(tokens)
+        padded = torch.zeros((batch_size, context_len), dtype=torch.long)
+        for idx, inst in enumerate(tokens):
+            L_ = min(inst.shape[0], context_len)
+            if inst.shape[0] > context_len:
+                inst = inst[:context_len]
+            padded[idx, :L_] = inst[:L_]
+        padded = padded.to(self.device)
+        feats = self.textual(padded)
+        eos = padded.argmax(dim=-1)
+        batch_idx = torch.arange(batch_size, device=self.device)
+        feats = feats[batch_idx, eos]
+        feats = self.textual_proj(feats)
+        if project:
+            feats = feats * self.textual_alpha.exp()
+            with torch.autocast(self.device.type, dtype=torch.float32):
+                feats = exp_map0(feats, self.curv.exp())
+        return feats
+
+
+# ---------------------------------------------------------------------------
+# Image preprocessing (matching the evaluation pipeline)
+# ---------------------------------------------------------------------------
+
+IMAGE_TRANSFORM = T.Compose([
+    T.Resize(224, T.InterpolationMode.BICUBIC),
+    T.CenterCrop(224),
+    T.ToTensor(),
+])
+
+# ---------------------------------------------------------------------------
+# Load model at module scope
+# ---------------------------------------------------------------------------
+
+CHECKPOINT_REPO = "hayeonkim/uncha"
+CHECKPOINT_FILE = "uncha_vit_b.pth"
+
+print("Loading UNCHA model...")
+model = UNCHAModel(
+    embed_dim=512,
+    visual_arch="vit_base_patch16_224",
+    text_arch="L12_W512",
+    vocab_size=49408,
+    context_length=77,
+)
+
+# Download and load checkpoint
+from huggingface_hub import hf_hub_download
+_ckpt_path = hf_hub_download(CHECKPOINT_REPO, CHECKPOINT_FILE)
+_ckpt = torch.load(_ckpt_path, map_location="cpu", weights_only=False)
+_sd = _ckpt["model"]
+
+# The checkpoint may contain min_radius_head keys from the text encoder that
+# our inference model doesn't have — filter them out.
+_model_sd = model.state_dict()
+_filtered_sd = {}
+for k, v in _sd.items():
+    if k in _model_sd:
+        _filtered_sd[k] = v
+    else:
+        print(f"  Skipping checkpoint key not in model: {k}")
+
+_missing, _unexpected = model.load_state_dict(_filtered_sd, strict=False)
+if _missing:
+    print(f"  Missing keys: {_missing}")
+if _unexpected:
+    print(f"  Unexpected keys: {_unexpected}")
+
+model = model.eval().to("cuda")
+print(f"Model loaded. curv={model.curv.exp().item():.4f}, "
+      f"visual_alpha={model.visual_alpha.exp().item():.4f}, "
+      f"textual_alpha={model.textual_alpha.exp().item():.4f}")
+
+# ---------------------------------------------------------------------------
+# Inference
+# ---------------------------------------------------------------------------
+
+PROMPT_TEMPLATES = [
+    "a photo of a {}.",
+    "a blurry photo of a {}.",
+    "a black and white photo of a {}.",
+    "a low contrast photo of a {}.",
+    "a high contrast photo of a {}.",
+    "a bad photo of a {}.",
+    "a good photo of a {}.",
+    "a photo of a small {}.",
+    "a photo of a big {}.",
+    "a photo of the {}.",
+]
+
+@spaces.GPU(duration=60)
+def classify(image: Image.Image, candidate_labels: str) -> dict:
+    """Zero-shot image classification using UNCHA hyperbolic vision-language model.
+
+    Args:
+        image: Input image to classify.
+        candidate_labels: Comma-separated list of candidate class labels.
+
+    Returns:
+        Dictionary mapping each label to its probability score.
+    """
+    if image is None:
+        return {}
+    if image.mode != "RGB":
+        image = image.convert("RGB")
+
+    # Parse labels
+    labels = [l.strip() for l in candidate_labels.split(",") if l.strip()]
+    if not labels:
+        return {}
+
+    # Preprocess image
+    img_tensor = IMAGE_TRANSFORM(image).unsqueeze(0).to(model.device)
+
+    # Encode image into hyperbolic space
+    with torch.inference_mode():
+        img_feats = model.encode_image(img_tensor, project=True)  # (1, D)
+
+    # Encode text prompts for each label (prompt ensemble in tangent space)
+    with torch.inference_mode():
+        all_class_feats = []
+        for label in labels:
+            prompts = [pt.format(label) for pt in PROMPT_TEMPLATES]
+            tokens = model.tokenizer(prompts)
+            text_feats = model.encode_text(tokens, project=False)  # (N_prompts, D)
+            # Ensemble in tangent space, then project to hyperboloid
+            text_feats = text_feats.mean(dim=0)  # (D,)
+            text_feats = text_feats * model.textual_alpha.exp()
+            text_feats = exp_map0(text_feats.unsqueeze(0), model.curv.exp())  # (1, D)
+            all_class_feats.append(text_feats.squeeze(0))
+
+        classifier = torch.stack(all_class_feats, dim=0)  # (num_classes, D)
+
+        # Lorentzian pairwise inner product as classification scores
+        scores = pairwise_inner(img_feats, classifier, model.curv.exp())  # (1, num_classes)
+        scores = scores.squeeze(0)  # (num_classes,)
+
+        # Convert to probabilities via softmax
+        probs = F.softmax(scores * model.logit_scale.exp(), dim=-1)
+
+    # Build label->prob dict, sorted by probability
+    result = {label: float(probs[i]) for i, label in enumerate(labels)}
+    result = dict(sorted(result.items(), key=lambda x: x[1], reverse=True))
+    return result
+
+
+# ---------------------------------------------------------------------------
+# 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:
+    with gr.Column(elem_id="col-container"):
+        gr.Markdown(
+            """
+            # UNCHA: Uncertainty-guided Compositional Hyperbolic Alignment
+
+            Zero-shot image classification using a hyperbolic vision-language model.
+            Upload an image and provide candidate labels — the model will compute
+            similarity scores using Lorentzian (hyperbolic) geometry.
+
+            [Paper](https://arxiv.org/abs/2603.22042) | [Code](https://github.com/jeeit17/UNCHA) | [Model](https://huggingface.co/hayeonkim/uncha)
+            """
+        )
+
+        with gr.Row():
+            with gr.Column(scale=1):
+                image_input = gr.Image(
+                    type="pil", label="Input Image",
+                    sources=["upload", "clipboard"],
+                )
+                labels_input = gr.Textbox(
+                    label="Candidate Labels (comma-separated)",
+                    placeholder="cat, dog, bird, car, person",
+                    value="cat, dog, bird, car, person",
+                )
+                run_btn = gr.Button("Classify", variant="primary")
+
+            with gr.Column(scale=1):
+                output_labels = gr.Label(
+                    label="Classification Scores",
+                    num_top_classes=10,
+                )
+
+    gr.Examples(
+        examples=[
+            ["examples/sample1.jpg", "dog, cat, animal, pet, mammal"],
+            ["examples/sample2.jpg", "building, landscape, nature, city, water"],
+            ["examples/sample3.jpg", "person, food, vehicle, animal, plant"],
+        ],
+        inputs=[image_input, labels_input],
+        outputs=output_labels,
+        fn=classify,
+        cache_examples=True,
+        cache_mode="lazy",
+    )
+
+demo.launch(mcp_server=True)

bpe_simple_vocab_16e6.txt.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:924691ac288e54409236115652ad4aa250f48203de50a9e4722a6ecd48d6804a
+size 1356917

requirements.txt

@@ -0,0 +1,3 @@
+timm
+ftfy
+regex