Add caption text decoder (encoder + decoder = full OpenVision2 generative model)

README.md

@@ -0,0 +1,47 @@
+---
+pipeline_tag: image-to-text
+tags:
+- openvision2
+- image-captioning
+- vision-encoder
+library_name: open_clip
+---
+
+# OpenVision2 · L/14 @336 — vision encoder + caption decoder
+
+This repo now contains the **full OpenVision2 generative model**: the **vision encoder**
+(originally released here) **plus the caption text decoder** it was jointly trained with.
+Together they map an image -> a descriptive caption. The encoder files are unchanged;
+only the decoder (and this card) were added.
+
+## Files
+| file | role |
+|---|---|
+| `open_clip_pytorch_model.bin`, `open_clip_config.json` | L/14 vision encoder (open_clip format) |
+| `caption_decoder.safetensors` | caption text decoder weights |
+| `text_decoder_config.json` | decoder architecture config |
+| `modeling_openvision2_decoder.py` | standalone PyTorch decoder + `generate()` |
+| `bert_base_vocab_bos_eos.txt` | tokenizer vocab (`[PAD]=0 [bos]=1 [eos]=2`) |
+| `caption_example.py` | end-to-end image -> caption demo |
+
+## Decoder architecture
+A **concat / prefix-LM** autoregressive transformer (not a CoCa cross-attention decoder):
+ViT patch tokens are linearly projected and **prepended as a bidirectional prefix**, and
+text is generated **causally** while attending to all image tokens.
+
+`12 layers · width 768 · 12 heads · mlp 3072 · vocab 32000 ·
+pre-LN · gelu(tanh) · LayerNorm eps 1e-6 · no positional embedding on the text stream.`
+It consumes the encoder's pre-final-norm patch tokens (open_clip `output_tokens=True`).
+
+## Usage
+Needs the patched `open_clip` providing `create_vision_encoder_and_transforms`
+(https://github.com/UCSC-VLAA/OpenVision), plus `torch`, `safetensors`, `pillow`.
+See `caption_example.py`:
+
+```bash
+python caption_example.py --image your.jpg
+```
+
+## Notes
+- Encoder and decoder are a **matched pair** exported from the same training checkpoint.
+- Captions are LLaVA-style **dense** descriptions (multi-sentence, detailed).

caption_decoder.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0241150fd5f5327e9fa5048a005f8ed82912a29b894b8c2be6151a37226558aa
+size 539991872

caption_example.py

@@ -0,0 +1,72 @@
+r'''
+OpenVision2 ViT-H/14 @224 — image -> caption demo.
+
+Combines the released vision encoder (this repo, open_clip format) with the caption
+text decoder (this repo). Requires the patched open_clip that ships
+`create_vision_encoder_and_transforms` (from the OpenVision repo:
+https://github.com/UCSC-VLAA/OpenVision), plus torch / safetensors / pillow.
+
+    python caption_example.py --image path/to/img.jpg
+'''
+import argparse, json
+import numpy as np, torch
+from PIL import Image
+from huggingface_hub import hf_hub_download
+from open_clip.factory import create_vision_encoder_and_transforms  # patched open_clip
+from modeling_openvision2_decoder import OpenVision2TextDecoder, OpenVision2TextDecoderConfig
+from safetensors.torch import load_file
+
+REPO = "UCSC-VLAA/openvision2-vit-large-patch14-336-vision-only"
+IMAGENET_MEAN = np.array([0.485, 0.456, 0.406]) * 255
+IMAGENET_STD = np.array([0.229, 0.224, 0.225]) * 255
+
+
+def preprocess(path, res=224):
+    im = Image.open(path).convert("RGB")
+    w, h = im.size; s = res / min(w, h)
+    im = im.resize((round(w * s), round(h * s)), Image.BILINEAR)
+    w, h = im.size; l, t = (w - res) // 2, (h - res) // 2
+    im = im.crop((l, t, l + res, t + res))
+    x = (np.asarray(im, np.float32) - IMAGENET_MEAN) / IMAGENET_STD
+    return torch.tensor(x.transpose(2, 0, 1)[None], dtype=torch.float32)  # [1,C,H,W]
+
+
+def load_decoder(cache_dir=None):
+    cfg = json.load(open(hf_hub_download(REPO, "text_decoder_config.json", cache_dir=cache_dir)))
+    dec = OpenVision2TextDecoder(OpenVision2TextDecoderConfig(
+        width=cfg["width"], depth=cfg["depth"], num_heads=cfg["num_heads"], mlp_dim=cfg["mlp_dim"],
+        vocab_size=cfg["vocab_size"], vision_width=cfg["vision_width"]))
+    dec.load_state_dict(load_file(hf_hub_download(REPO, "caption_decoder.safetensors", cache_dir=cache_dir)))
+    dec.eval()
+    vocab = [l.rstrip("\n") for l in open(hf_hub_download(REPO, "bert_base_vocab_bos_eos.txt", cache_dir=cache_dir))]
+    return dec, vocab
+
+
+def detok(ids, vocab):
+    w = []
+    for i in ids:
+        if i == 2: break               # eos
+        if i in (0, 1, 2): continue    # pad / bos / eos
+        tk = vocab[i] if 0 <= i < len(vocab) else "[UNK]"
+        if tk.startswith("##"): w[-1] = (w[-1] + tk[2:]) if w else tk[2:]
+        else: w.append(tk)
+    return " ".join(w)
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--image", required=True)
+    ap.add_argument("--max_len", type=int, default=64)
+    ap.add_argument("--cache_dir", default=None)
+    args = ap.parse_args()
+
+    enc = create_vision_encoder_and_transforms(model_name=f"hf-hub:{REPO}", cache_dir=args.cache_dir).eval()
+    dec, vocab = load_decoder(args.cache_dir)
+    with torch.no_grad():
+        _, tokens = enc(preprocess(args.image))          # [1, N, 1280]
+        ids = dec.generate(tokens, max_len=args.max_len, bos_id=1, eos_id=2)[0].tolist()
+    print(detok(ids, vocab))
+
+
+if __name__ == "__main__":
+    main()

modeling_openvision2_decoder.py

@@ -0,0 +1,155 @@
+"""PyTorch reference implementation of the OpenVision2 caption decoder.
+
+The OpenVision2 generative model is: a ViT vision encoder (released separately as
+the `*-vision-only` open_clip checkpoint) whose patch tokens condition a small
+autoregressive text decoder trained with a captioning loss.
+
+This decoder is a **prefix-LM / concat-fusion** transformer (NOT a CoCa-style
+cross-attention decoder):
+
+    text_embeds  = Embed(text_tokens)                      # no positional embedding
+    image_embeds = image_projection(vit_patch_tokens)      # Linear, no bias
+    x            = concat([image_embeds, text_embeds], 1)  # [B, N_img + L, width]
+    x            = prefix_lm_transformer(x)                 # image=bidirectional prefix,
+                                                            # text=causal, text->image=full,
+                                                            # image-/->text
+    logits       = lm_head( ln_final( x[:, N_img:] ) )     # over text positions
+
+Faithful details that matter for numerical parity with the JAX model:
+  * LayerNorm eps = 1e-6 (flax default), pre-LN blocks.
+  * MLP activation = gelu(tanh approximation).
+  * attention scale = head_dim ** -0.5.
+  * NO positional embedding on the text stream.
+  * vision patch tokens are the pre-final-norm tokens, cls excluded
+    (open_clip vision model with output_tokens=True returns exactly these).
+"""
+from dataclasses import dataclass
+from typing import Optional
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+@dataclass
+class OpenVision2TextDecoderConfig:
+    width: int = 1024          # decoder hidden width (H decoder -> 1024)
+    depth: int = 24            # number of transformer blocks
+    num_heads: int = 16
+    mlp_dim: int = 4096
+    vocab_size: int = 32000
+    vision_width: int = 1280   # ViT token dim feeding image_projection (H ViT -> 1280)
+    layer_norm_eps: float = 1e-6
+    pad_id: int = 0
+    bos_id: int = 1
+    eos_id: int = 2
+
+
+class _Mlp(nn.Module):
+    def __init__(self, width: int, mlp_dim: int):
+        super().__init__()
+        self.c_fc = nn.Linear(width, mlp_dim)
+        self.c_proj = nn.Linear(mlp_dim, width)
+
+    def forward(self, x):
+        return self.c_proj(F.gelu(self.c_fc(x), approximate="tanh"))
+
+
+class _Attention(nn.Module):
+    """MHA with a fused in_proj (q,k,v) and an explicit boolean attend-mask."""
+
+    def __init__(self, width: int, num_heads: int):
+        super().__init__()
+        assert width % num_heads == 0
+        self.num_heads = num_heads
+        self.head_dim = width // num_heads
+        self.scale = self.head_dim ** -0.5
+        self.in_proj_weight = nn.Parameter(torch.empty(3 * width, width))
+        self.in_proj_bias = nn.Parameter(torch.zeros(3 * width))
+        self.out_proj = nn.Linear(width, width)
+
+    def forward(self, x, attend_mask):
+        # attend_mask: [L, L] bool, True = keep, False = -inf
+        B, L, D = x.shape
+        qkv = F.linear(x, self.in_proj_weight, self.in_proj_bias)
+        q, k, v = qkv.chunk(3, dim=-1)
+        H, hd = self.num_heads, self.head_dim
+        q = q.view(B, L, H, hd).transpose(1, 2)   # [B,H,L,hd]
+        k = k.view(B, L, H, hd).transpose(1, 2)
+        v = v.view(B, L, H, hd).transpose(1, 2)
+        attn = (q @ k.transpose(-2, -1)) * self.scale  # [B,H,L,L]
+        attn = attn.masked_fill(~attend_mask[None, None], float("-inf"))
+        attn = attn.softmax(dim=-1)
+        out = attn @ v                                  # [B,H,L,hd]
+        out = out.transpose(1, 2).reshape(B, L, D)
+        return self.out_proj(out)
+
+
+class _Block(nn.Module):
+    def __init__(self, cfg: OpenVision2TextDecoderConfig):
+        super().__init__()
+        self.ln_1 = nn.LayerNorm(cfg.width, eps=cfg.layer_norm_eps)
+        self.attn = _Attention(cfg.width, cfg.num_heads)
+        self.ln_2 = nn.LayerNorm(cfg.width, eps=cfg.layer_norm_eps)
+        self.mlp = _Mlp(cfg.width, cfg.mlp_dim)
+
+    def forward(self, x, attend_mask):
+        x = x + self.attn(self.ln_1(x), attend_mask)
+        x = x + self.mlp(self.ln_2(x))
+        return x
+
+
+class OpenVision2TextDecoder(nn.Module):
+    def __init__(self, cfg: OpenVision2TextDecoderConfig):
+        super().__init__()
+        self.cfg = cfg
+        self.token_embedding = nn.Embedding(cfg.vocab_size, cfg.width)
+        self.image_projection = nn.Linear(cfg.vision_width, cfg.width, bias=False)
+        self.blocks = nn.ModuleList([_Block(cfg) for _ in range(cfg.depth)])
+        self.ln_final = nn.LayerNorm(cfg.width, eps=cfg.layer_norm_eps)
+        self.lm_head = nn.Linear(cfg.width, cfg.vocab_size, bias=False)
+
+    @staticmethod
+    def _prefix_lm_mask(li: int, lt: int, device) -> torch.Tensor:
+        """[l, l] bool attend-mask. image(:li) bidirectional prefix; text(li:)
+        causal + attends all image; image does NOT attend text."""
+        l = li + lt
+        mask = torch.zeros(l, l, dtype=torch.bool, device=device)
+        mask[:li, :li] = True                              # image <-> image
+        mask[li:, :li] = True                              # text  -> image
+        text_causal = torch.tril(torch.ones(lt, lt, dtype=torch.bool, device=device))
+        mask[li:, li:] = text_causal                       # text causal self
+        return mask
+
+    def forward(self, image_tokens: torch.Tensor, text_tokens: torch.Tensor) -> torch.Tensor:
+        """image_tokens: [B, N_img, vision_width]; text_tokens: [B, L] (input side).
+        Returns logits [B, L, vocab] (next-token logits at each text position)."""
+        text_embeds = self.token_embedding(text_tokens)          # [B, L, width]
+        image_embeds = self.image_projection(image_tokens)       # [B, N_img, width]
+        li, lt = image_embeds.shape[1], text_embeds.shape[1]
+        x = torch.cat([image_embeds, text_embeds], dim=1)        # [B, li+lt, width]
+        mask = self._prefix_lm_mask(li, lt, x.device)
+        for blk in self.blocks:
+            x = blk(x, mask)
+        x = x[:, li:]                                            # text positions
+        x = self.ln_final(x)
+        return self.lm_head(x)
+
+    @torch.no_grad()
+    def generate(self, image_tokens: torch.Tensor, max_len: int = 64,
+                 bos_id: Optional[int] = None, eos_id: Optional[int] = None) -> torch.Tensor:
+        """Greedy autoregressive caption generation. image_tokens: [B, N, vision_width]."""
+        bos_id = self.cfg.bos_id if bos_id is None else bos_id
+        eos_id = self.cfg.eos_id if eos_id is None else eos_id
+        B = image_tokens.shape[0]
+        device = image_tokens.device
+        seq = torch.full((B, 1), bos_id, dtype=torch.long, device=device)
+        done = torch.zeros(B, dtype=torch.bool, device=device)
+        for _ in range(max_len):
+            logits = self.forward(image_tokens, seq)             # [B, cur_len, vocab]
+            nxt = logits[:, -1].argmax(dim=-1)                   # [B]
+            nxt = torch.where(done, torch.full_like(nxt, self.cfg.pad_id), nxt)
+            seq = torch.cat([seq, nxt[:, None]], dim=1)
+            done = done | (nxt == eos_id)
+            if bool(done.all()):
+                break
+        return seq

text_decoder_config.json

@@ -0,0 +1,17 @@
+{
+  "architecture": "OpenVision2TextDecoder",
+  "fusion_style": "concat",
+  "width": 768,
+  "depth": 12,
+  "num_heads": 12,
+  "mlp_dim": 3072,
+  "vocab_size": 32000,
+  "vision_width": 1024,
+  "layer_norm_eps": 1e-06,
+  "pad_id": 0,
+  "bos_id": 1,
+  "eos_id": 2,
+  "weights_file": "caption_decoder.safetensors",
+  "modeling_file": "modeling_openvision2_decoder.py",
+  "tokenizer": "bert wordpiece (bert_base_vocab_bos_eos.txt), lowercase, [PAD]=0 [bos]=1 [eos]=2"
+}