| """Standalone loader for clip-vitb-mini-distilled (no repo imports). |
| |
| Usage: |
| from loader import load_model, embed_images |
| model = load_model("student.pt") # or a hf_hub_download path |
| emb = embed_images(model, batch01) # (B,3,H,W) in [0,1] -> (B,512) |
| |
| The student is an 8.66M-parameter ViT (d=240, depth 12, heads 4, patch 16, |
| img 160) distilled from CLIP ViT-B/16 (LAION-2B) pooled image features on |
| COCO-2017 train (118,287 images, ~17 epochs). Outputs live in the teacher's |
| 512-d projection space: compatible with the teacher's TEXT tower for |
| zero-shot/retrieval. See the model card for measured capability and caveats. |
| """ |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| CLIP_MEAN = (0.48145466, 0.4578275, 0.40821073) |
| CLIP_STD = (0.26862954, 0.26130258, 0.27577711) |
|
|
|
|
| class Block(nn.Module): |
| def __init__(self, d, heads): |
| super().__init__() |
| self.n1 = nn.LayerNorm(d) |
| self.qkv = nn.Linear(d, 3 * d) |
| self.proj = nn.Linear(d, d) |
| self.n2 = nn.LayerNorm(d) |
| self.fc1 = nn.Linear(d, 4 * d) |
| self.fc2 = nn.Linear(4 * d, d) |
| self.heads = heads |
|
|
| def forward(self, x): |
| B, N, C = x.shape |
| q, k, v = (self.qkv(self.n1(x)) |
| .reshape(B, N, 3, self.heads, C // self.heads) |
| .permute(2, 0, 3, 1, 4)) |
| a = F.scaled_dot_product_attention(q, k, v) |
| x = x + self.proj(a.transpose(1, 2).reshape(B, N, C)) |
| return x + self.fc2(F.gelu(self.fc1(self.n2(x)))) |
|
|
|
|
| class Student(nn.Module): |
| """CLS-token readout, Linear head to the teacher's 512-d space.""" |
| def __init__(self, out_dim=512, d=240, depth=12, heads=4, patch=16, |
| img=160): |
| super().__init__() |
| self.patch = nn.Conv2d(3, d, patch, patch) |
| self.cls = nn.Parameter(torch.zeros(1, 1, d)) |
| self.pos = nn.Parameter(torch.zeros(1, (img // patch) ** 2 + 1, d)) |
| self.blocks = nn.ModuleList(Block(d, heads) for _ in range(depth)) |
| self.norm = nn.LayerNorm(d) |
| self.head = nn.Linear(d, out_dim) |
|
|
| def forward_features(self, x): |
| x = self.patch(x).flatten(2).transpose(1, 2) |
| x = torch.cat([self.cls.expand(x.shape[0], -1, -1), x], 1) + self.pos |
| for b in self.blocks: |
| x = b(x) |
| return self.norm(x)[:, 0] |
|
|
| def forward(self, x): |
| return self.head(self.forward_features(x)) |
|
|
|
|
| def load_model(path, device=None): |
| device = device or ("cuda" if torch.cuda.is_available() else "cpu") |
| ck = torch.load(path, map_location="cpu", weights_only=True) |
| sd = ck["state_dict"] if "state_dict" in ck else ck |
| model = Student(out_dim=sd["head.weight"].shape[0]) |
| model.load_state_dict(sd, strict=True) |
| return model.to(device).eval() |
|
|
|
|
| def load_rotation(path): |
| """Load a frame rotation shipped beside a consensus-distilled student. |
| The rotation maps the student's consensus frame into the CLIP-B/16 |
| (LAION-2B) deployment frame; it was fitted once, offline, on 2,500 |
| COCO-val pairs (fp64 orthogonal Procrustes) and is FROZEN.""" |
| d = torch.load(path, map_location="cpu", weights_only=True) |
| return d["R"] if isinstance(d, dict) else d |
|
|
|
|
| @torch.no_grad() |
| def embed_images(model, x01, batch=256, rotation=None): |
| """x01: (B,3,H,W) float in [0,1]. Returns L2-normalized (B,512). |
| Pass `rotation` (from load_rotation) with the consensus-distilled |
| champion to place outputs in the LAION text tower's frame — required |
| for zero-shot/retrieval against that tower, and how every headline |
| number on the card was measured.""" |
| dev = next(model.parameters()).device |
| mean = torch.tensor(CLIP_MEAN).view(1, 3, 1, 1).to(dev) |
| std = torch.tensor(CLIP_STD).view(1, 3, 1, 1).to(dev) |
| out = [] |
| for i in range(0, len(x01), batch): |
| x = x01[i:i + batch].to(dev) |
| x = F.interpolate(x, size=(160, 160), mode="bicubic", |
| align_corners=False).clamp(0, 1) |
| z = F.normalize(model((x - mean) / std), dim=-1) |
| if rotation is not None: |
| z = F.normalize(z.double() @ rotation.to(dev).double(), |
| dim=-1).float() |
| out.append(z.cpu()) |
| return torch.cat(out) |
|
|
