10.1 Vision Transformers & Image Encoders¶

import torch
import torch.nn as nn
import einops  # pip install einops

# ---------------------------------------------------------------------------
# patchify: split an image tensor into a sequence of patch vectors
# ---------------------------------------------------------------------------
def patchify(images: torch.Tensor, patch_size: int = 16) -> torch.Tensor:
    """
    Args:
        images: (B, C, H, W) float tensor, pre-normalized
        patch_size: side length of each square patch (P)
    Returns:
        patches: (B, N, P*P*C) where N = (H/P) * (W/P)
    """
    B, C, H, W = images.shape
    assert H % patch_size == 0 and W % patch_size == 0, \
        f"Image size ({H},{W}) must be divisible by patch_size={patch_size}"

    # Rearrange from (B, C, H, W) → (B, N, P*P*C)
    # 'h1 h2 w1 w2' means height is split into h1 groups of h2=P rows, same for width
    patches = einops.rearrange(
        images,
        'b c (h1 h2) (w1 w2) -> b (h1 w1) (h2 w2 c)',
        h2=patch_size, w2=patch_size
    )
    return patches  # (B, N, P*P*C)


# ---------------------------------------------------------------------------
# PatchEmbed: learnable linear projection of patches into D-dim vectors
# ---------------------------------------------------------------------------
class PatchEmbed(nn.Module):
    """
    Mathematically the same operation as a Conv2d with kernel_size=stride=
    patch_size, but implemented as a matrix multiply for clarity. Real
    implementations (timm, HF `transformers`, open_clip) all use
    `nn.Conv2d(in_chans, embed_dim, kernel_size=P, stride=P)` because cuDNN
    fuses the gather-and-project into one pass.

    Caveat if you ever port weights: Conv2d flattens each patch in (C, P, P)
    order, whereas the einops pattern below flattens in (P, P, C) order, so the
    two weight matrices differ by a permutation of the input columns. Either
    ordering trains fine — but you cannot load one checkpoint into the other
    without permuting.
    """
    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768):
        super().__init__()
        self.img_size = img_size
        self.patch_size = patch_size
        self.num_patches = (img_size // patch_size) ** 2   # 196 for 224/16
        self.patch_dim = patch_size * patch_size * in_chans  # 768 = 16*16*3

        # A single weight matrix E ∈ R^{D × (P²C)}
        self.proj = nn.Linear(self.patch_dim, embed_dim, bias=True)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        x: (B, C, H, W) pre-normalized image
        returns: (B, N, D) patch embeddings
        """
        patches = patchify(x, self.patch_size)  # (B, N, P²C)
        return self.proj(patches)               # (B, N, D)


# ---------------------------------------------------------------------------
# Quick sanity check
# ---------------------------------------------------------------------------
if __name__ == "__main__":
    B, C, H, W = 4, 3, 224, 224
    images = torch.randn(B, C, H, W)

    embed = PatchEmbed(img_size=224, patch_size=16, in_chans=3, embed_dim=768)
    out = embed(images)
    print(f"Input shape:  {images.shape}")   # (4, 3, 224, 224)
    print(f"Output shape: {out.shape}")       # (4, 196, 768)
    print(f"Params in PatchEmbed: {sum(p.numel() for p in embed.parameters()):,}")
    # Expected: 768 * 768 + 768 = 590,592  (weight + bias)

import torch
import torch.nn as nn
import math


class MultiHeadSelfAttention(nn.Module):
    """Standard multi-head self-attention — see chapter 2.4 for the full treatment."""
    def __init__(self, embed_dim: int, num_heads: int, dropout: float = 0.0):
        super().__init__()
        assert embed_dim % num_heads == 0
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads
        self.scale = math.sqrt(self.head_dim)  # 1/√d_k scaling factor

        # Fused QKV projection: one matrix, split afterward
        self.qkv = nn.Linear(embed_dim, 3 * embed_dim, bias=True)
        self.out_proj = nn.Linear(embed_dim, embed_dim, bias=True)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        x: (B, N, D)  where N = num_patches+1 (includes CLS)
        returns: (B, N, D)
        """
        B, N, D = x.shape
        H = self.num_heads

        # Project and split into queries, keys, values
        qkv = self.qkv(x)                         # (B, N, 3D)
        q, k, v = qkv.chunk(3, dim=-1)            # each (B, N, D)

        # Reshape to (B, H, N, head_dim) for multi-head computation
        def reshape(t):
            return t.view(B, N, H, self.head_dim).transpose(1, 2)
        q, k, v = reshape(q), reshape(k), reshape(v)  # each (B, H, N, head_dim)

        # Scaled dot-product attention, written out explicitly so you can see
        # the (B, H, N, N) matrix. In production, replace these four lines with
        #   out = F.scaled_dot_product_attention(q, k, v, dropout_p=...)
        # which dispatches to the FlashAttention kernel — same math, no N×N
        # matrix ever materialized. That swap is what makes 1024-patch (512px)
        # and 4096-patch (1024px) ViTs affordable.
        attn = (q @ k.transpose(-2, -1)) / self.scale   # (B, H, N, N)
        attn = attn.softmax(dim=-1)
        attn = self.dropout(attn)

        # Aggregate values and reshape back
        out = (attn @ v)                          # (B, H, N, head_dim)
        out = out.transpose(1, 2).reshape(B, N, D)
        return self.out_proj(out)                 # (B, N, D)


class ViTBlock(nn.Module):
    """
    One Transformer block in ViT, using Pre-LayerNorm (LN before attention).
    Pre-LN improves training stability vs the original Post-LN used by BERT.
    See: chapter 2.6 for a detailed discussion of norm placement.
    """
    def __init__(self, embed_dim: int, num_heads: int, mlp_ratio: float = 4.0,
                 dropout: float = 0.0):
        super().__init__()
        self.norm1 = nn.LayerNorm(embed_dim)
        self.attn = MultiHeadSelfAttention(embed_dim, num_heads, dropout)
        self.norm2 = nn.LayerNorm(embed_dim)
        mlp_hidden = int(embed_dim * mlp_ratio)  # typically 4× hidden dim = 3072

        # Two-layer MLP with GELU activation — standard ViT uses GELU
        self.mlp = nn.Sequential(
            nn.Linear(embed_dim, mlp_hidden),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(mlp_hidden, embed_dim),
            nn.Dropout(dropout),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Pre-LN residual: x = x + Attn(LN(x))
        x = x + self.attn(self.norm1(x))
        # Pre-LN residual: x = x + MLP(LN(x))
        x = x + self.mlp(self.norm2(x))
        return x


class VisionTransformer(nn.Module):
    """
    Minimal ViT implementation matching the ViT-B/16 configuration.
    For the full model (ViT-L, ViT-H, ViT-G), only embed_dim, num_heads,
    and depth change — everything else stays the same.
    """
    def __init__(
        self,
        img_size: int = 224,
        patch_size: int = 16,
        in_chans: int = 3,
        num_classes: int = 1000,
        embed_dim: int = 768,
        depth: int = 12,
        num_heads: int = 12,
        mlp_ratio: float = 4.0,
        dropout: float = 0.0,
    ):
        super().__init__()
        self.patch_embed = PatchEmbed(img_size, patch_size, in_chans, embed_dim)
        num_patches = self.patch_embed.num_patches  # 196

        # Learnable CLS token (prepended to the patch sequence)
        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))

        # Learnable positional embeddings for N+1 positions (patches + CLS)
        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))

        self.pos_drop = nn.Dropout(dropout)

        # Stack of Transformer blocks
        self.blocks = nn.ModuleList([
            ViTBlock(embed_dim, num_heads, mlp_ratio, dropout)
            for _ in range(depth)
        ])
        self.norm = nn.LayerNorm(embed_dim)

        # Classification head: linear layer on top of the CLS token
        self.head = nn.Linear(embed_dim, num_classes)

        # Initialize weights following the ViT paper
        nn.init.trunc_normal_(self.pos_embed, std=0.02)
        nn.init.trunc_normal_(self.cls_token, std=0.02)
        self.apply(self._init_weights)

    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            nn.init.trunc_normal_(m.weight, std=0.02)
            if m.bias is not None:
                nn.init.zeros_(m.bias)
        elif isinstance(m, nn.LayerNorm):
            nn.init.zeros_(m.bias)
            nn.init.ones_(m.weight)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        x: (B, 3, H, W) pre-normalized images
        returns: (B, num_classes) logits
        """
        B = x.shape[0]
        x = self.patch_embed(x)   # (B, N, D)

        # Expand and prepend CLS token
        cls = self.cls_token.expand(B, -1, -1)  # (B, 1, D)
        x = torch.cat([cls, x], dim=1)           # (B, N+1, D)

        # Add positional embeddings and apply dropout
        x = x + self.pos_embed                   # (B, N+1, D)
        x = self.pos_drop(x)

        # Apply Transformer blocks
        for block in self.blocks:
            x = block(x)                         # (B, N+1, D)

        x = self.norm(x)                         # (B, N+1, D)

        # Extract CLS token output and classify
        cls_out = x[:, 0]                        # (B, D)
        return self.head(cls_out)                # (B, num_classes)


# Verify shapes
if __name__ == "__main__":
    model = VisionTransformer()  # ViT-B/16 defaults
    imgs = torch.randn(2, 3, 224, 224)
    logits = model(imgs)
    print(f"Logits shape: {logits.shape}")   # (2, 1000)
    total = sum(p.numel() for p in model.parameters())
    print(f"Total parameters: {total:,}")    # ~86M

import torch
import torch.nn as nn


class DropPath(nn.Module):
    """Stochastic depth: zero the *whole* residual branch for some samples."""
    def __init__(self, p: float = 0.0):
        super().__init__()
        self.p = p

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self.p == 0.0 or not self.training:
            return x
        keep = 1.0 - self.p
        # One Bernoulli draw per sample, broadcast over (N, D)
        shape = (x.shape[0],) + (1,) * (x.ndim - 1)
        mask = x.new_empty(shape).bernoulli_(keep)
        return x * mask / keep          # rescale so E[output] is unchanged


class LayerScale(nn.Module):
    """Per-channel learnable gain on a residual branch (CaiT). Init near zero."""
    def __init__(self, dim: int, init_value: float = 1e-5):
        super().__init__()
        self.gamma = nn.Parameter(init_value * torch.ones(dim))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x * self.gamma


class RegularizedViTBlock(nn.Module):
    """ViTBlock + LayerScale + DropPath — the block deep ViTs actually use."""
    def __init__(self, embed_dim: int, num_heads: int, mlp_ratio: float = 4.0,
                 drop_path: float = 0.0, ls_init: float = 1e-5):
        super().__init__()
        self.norm1 = nn.LayerNorm(embed_dim)
        self.attn = MultiHeadSelfAttention(embed_dim, num_heads)
        self.ls1, self.dp1 = LayerScale(embed_dim, ls_init), DropPath(drop_path)

        self.norm2 = nn.LayerNorm(embed_dim)
        hidden = int(embed_dim * mlp_ratio)
        self.mlp = nn.Sequential(nn.Linear(embed_dim, hidden), nn.GELU(),
                                 nn.Linear(hidden, embed_dim))
        self.ls2, self.dp2 = LayerScale(embed_dim, ls_init), DropPath(drop_path)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # residual = x + DropPath(LayerScale(Branch(LN(x))))
        x = x + self.dp1(self.ls1(self.attn(self.norm1(x))))
        x = x + self.dp2(self.ls2(self.mlp(self.norm2(x))))
        return x


if __name__ == "__main__":
    depth, max_dp = 12, 0.1
    # Linearly ramped drop-path rate: shallow blocks survive, deep blocks drop
    dpr = [max_dp * i / (depth - 1) for i in range(depth)]
    blocks = nn.ModuleList([RegularizedViTBlock(768, 12, drop_path=p) for p in dpr])

    x = torch.randn(2, 197, 768)
    for blk in blocks:
        x = blk(x)
    print(x.shape, f"drop-path ramp: {dpr[0]:.3f} → {dpr[-1]:.3f}")

import torch
import torch.nn.functional as F

def interpolate_pos_embed(pos_embed: torch.Tensor,
                          new_h: int, new_w: int,
                          patch_size: int = 16) -> torch.Tensor:
    """
    Interpolate positional embeddings from training resolution to a new resolution.

    Args:
        pos_embed: (1, N+1, D) learned position embeddings (includes CLS at index 0)
        new_h, new_w: new image height and width in pixels
        patch_size: patch size in pixels
    Returns:
        (1, new_N+1, D) interpolated positional embeddings
    """
    # Separate CLS token from patch position embeddings
    cls_pos = pos_embed[:, :1, :]           # (1, 1, D)
    patch_pos = pos_embed[:, 1:, :]         # (1, N, D)

    # Infer original grid size from N = h_orig * w_orig
    N_orig = patch_pos.shape[1]
    h_orig = w_orig = int(N_orig ** 0.5)    # e.g., 14 for 196 patches
    D = patch_pos.shape[-1]

    # Reshape to 2D grid, interpolate
    patch_pos = patch_pos.reshape(1, h_orig, w_orig, D)
    patch_pos = patch_pos.permute(0, 3, 1, 2)  # (1, D, h_orig, w_orig)

    new_h_patches = new_h // patch_size
    new_w_patches = new_w // patch_size

    patch_pos = F.interpolate(
        patch_pos,
        size=(new_h_patches, new_w_patches),
        mode='bicubic',
        align_corners=False
    )  # (1, D, new_h_patches, new_w_patches)

    patch_pos = patch_pos.permute(0, 2, 3, 1)  # (1, new_h_patches, new_w_patches, D)
    patch_pos = patch_pos.reshape(1, -1, D)     # (1, new_N, D)

    return torch.cat([cls_pos, patch_pos], dim=1)  # (1, new_N+1, D)

import torch

def build_2d_rope(h: int, w: int, head_dim: int, theta: float = 100.0):
    """
    Axial 2D RoPE cache for an h x w patch grid.

    Half of head_dim rotates by the ROW index, the other half by the COLUMN
    index, so a patch's phase encodes (row, col). theta is much smaller than
    the 10000 used for text because patch grids are short in each axis.

    Returns cos, sin of shape (h*w, head_dim), in rotate_half convention.
    """
    assert head_dim % 4 == 0, "head_dim must be divisible by 4 for axial 2D RoPE"
    d_axis = head_dim // 2                       # dims allocated to each axis
    inv_freq = 1.0 / (theta ** (torch.arange(0, d_axis, 2).float() / d_axis))

    ys, xs = torch.meshgrid(torch.arange(h).float(),
                            torch.arange(w).float(), indexing="ij")
    ang_y = ys.reshape(-1, 1) * inv_freq         # (h*w, d_axis/2)
    ang_x = xs.reshape(-1, 1) * inv_freq         # (h*w, d_axis/2)
    ang = torch.cat([ang_y, ang_x], dim=-1)      # (h*w, head_dim/2)
    ang = torch.cat([ang, ang], dim=-1)          # duplicate for rotate_half
    return ang.cos(), ang.sin()


def rotate_half(x: torch.Tensor) -> torch.Tensor:
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat([-x2, x1], dim=-1)


def apply_2d_rope(q: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor):
    """q: (B, H, N, head_dim) — PATCH tokens only; CLS/registers are left alone."""
    return q * cos + rotate_half(q) * sin


if __name__ == "__main__":
    head_dim = 64
    for (gh, gw) in [(14, 14), (32, 24)]:     # 224x224 and 512x384 at P=16
        cos, sin = build_2d_rope(gh, gw, head_dim)
        q = torch.randn(2, 12, gh * gw, head_dim)
        out = apply_2d_rope(q, cos, sin)
        # A rotation preserves norms — a good invariant to assert in tests
        assert torch.allclose(out.norm(dim=-1), q.norm(dim=-1), atol=1e-5)
        print(f"grid {gh}x{gw}: {tuple(out.shape)}, same weights, no interpolation")

import torch
import torch.nn.functional as F


def clip_contrastive_loss(image_embeds: torch.Tensor,
                          text_embeds: torch.Tensor,
                          logit_scale: torch.Tensor) -> torch.Tensor:
    """
    Compute the symmetric CLIP contrastive loss (InfoNCE in both directions).

    Args:
        image_embeds: (N, D) L2-normalized image embeddings
        text_embeds:  (N, D) L2-normalized text embeddings
        logit_scale:  scalar parameter, stored as log(τ) and exp()'d for stability
    Returns:
        scalar loss
    """
    # Exponentiate the log-temperature; clamp to prevent collapse or explosion
    scale = logit_scale.exp().clamp(max=100.0)

    # Compute (N, N) similarity matrix: S[i,j] = τ * <image_i, text_j>
    logits_per_image = scale * image_embeds @ text_embeds.t()  # (N, N)
    logits_per_text  = logits_per_image.t()                   # (N, N)

    # Ground-truth: diagonal is the positive pair
    N = image_embeds.shape[0]
    labels = torch.arange(N, device=image_embeds.device)

    # Cross-entropy in both directions, then average
    loss_i = F.cross_entropy(logits_per_image, labels)
    loss_t = F.cross_entropy(logits_per_text,  labels)
    return (loss_i + loss_t) / 2.0


def zero_shot_classify(image_embed: torch.Tensor,
                       class_text_embeds: torch.Tensor) -> torch.Tensor:
    """
    Args:
        image_embed:       (D,)   L2-normalized embedding of a single image
        class_text_embeds: (C, D) L2-normalized embeddings for C class prompts
    Returns:
        (C,) softmax probabilities
    """
    # Cosine similarities (temperature=1 here; CLIP uses learned τ)
    sims = image_embed @ class_text_embeds.t()   # (C,)
    return sims.softmax(dim=-1)

class VisionTransformerWithRegisters(nn.Module):
    """
    ViT with register tokens, following Darcet et al. (2023).
    Register tokens are appended after patch tokens and discarded at output.
    """
    def __init__(
        self,
        img_size=224, patch_size=14, in_chans=3, embed_dim=1024,
        depth=24, num_heads=16, mlp_ratio=4.0,
        num_registers=4,    # key new parameter
        num_classes=0,      # 0 = return features, not logits
    ):
        super().__init__()
        self.patch_embed = PatchEmbed(img_size, patch_size, in_chans, embed_dim)
        num_patches = self.patch_embed.num_patches

        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))

        # Register tokens — no positional embedding added to them
        self.num_registers = num_registers
        self.register_tokens = nn.Parameter(
            torch.zeros(1, num_registers, embed_dim)
        )

        self.blocks = nn.ModuleList([
            ViTBlock(embed_dim, num_heads, mlp_ratio)
            for _ in range(depth)
        ])
        self.norm = nn.LayerNorm(embed_dim)

        nn.init.trunc_normal_(self.pos_embed, std=0.02)
        nn.init.trunc_normal_(self.cls_token, std=0.02)
        nn.init.trunc_normal_(self.register_tokens, std=0.02)

    def forward(self, x: torch.Tensor):
        B = x.shape[0]
        patches = self.patch_embed(x)                     # (B, N, D)

        # Build sequence: [CLS | patches] with positional embeddings
        cls = self.cls_token.expand(B, -1, -1)
        x = torch.cat([cls, patches], dim=1)              # (B, N+1, D)
        x = x + self.pos_embed                            # add pos embeddings

        # Append register tokens (no positional embedding for registers)
        regs = self.register_tokens.expand(B, -1, -1)
        x = torch.cat([x, regs], dim=1)                   # (B, N+1+R, D)

        for block in self.blocks:
            x = block(x)

        x = self.norm(x)

        # Discard register tokens; return CLS and patch tokens
        # Register tokens occupy the last R positions
        x = x[:, :-self.num_registers]                    # (B, N+1, D)
        cls_out = x[:, 0]                                 # (B, D)
        patch_out = x[:, 1:]                              # (B, N, D)

        return cls_out, patch_out   # global feature, spatial features

import torchvision.transforms.v2 as T
from PIL import Image
import torch

# Standard ViT/CLIP preprocessing (ImageNet statistics)
IMAGENET_MEAN = (0.485, 0.456, 0.406)
IMAGENET_STD  = (0.229, 0.224, 0.225)

# CLIP uses slightly different mean/std
CLIP_MEAN = (0.48145466, 0.4578275, 0.40821073)
CLIP_STD  = (0.26862954, 0.26130258, 0.27577711)


def build_vit_transform(img_size: int = 224,
                        mean=IMAGENET_MEAN,
                        std=IMAGENET_STD,
                        is_train: bool = False) -> T.Compose:
    """
    Build the standard preprocessing pipeline for a ViT model.

    During training, we add random resized crop and horizontal flip for
    data augmentation. During inference, we use deterministic center crop.
    """
    if is_train:
        return T.Compose([
            T.RandomResizedCrop(img_size, scale=(0.2, 1.0)),  # random crop
            T.RandomHorizontalFlip(),
            T.ToImage(),
            T.ToDtype(torch.float32, scale=True),   # → [0, 1]
            T.Normalize(mean=mean, std=std),         # → ~N(0,1)
        ])
    else:
        return T.Compose([
            T.Resize(int(img_size * 256 / 224)),    # resize short edge
            T.CenterCrop(img_size),                 # deterministic center crop
            T.ToImage(),
            T.ToDtype(torch.float32, scale=True),
            T.Normalize(mean=mean, std=std),
        ])


# Example usage
if __name__ == "__main__":
    transform = build_vit_transform(224, is_train=False)

    # Load a PIL image
    img = Image.new("RGB", (640, 480), color=(128, 64, 32))
    tensor = transform(img)               # (3, 224, 224)
    print(f"Preprocessed tensor shape: {tensor.shape}")    # torch.Size([3, 224, 224])
    print(f"Value range: [{tensor.min():.2f}, {tensor.max():.2f}]")  # ~[-2, 2]

    # Batch and feed to model
    batch = tensor.unsqueeze(0)           # (1, 3, 224, 224)
    model = VisionTransformer()
    model.eval()
    with torch.no_grad():
        logits = model(batch)
    print(f"Logits shape: {logits.shape}")  # (1, 1000)

# pip install timm
import timm
import torch

# `num_classes=0` strips the classification head and returns pooled features.
# Names follow "<arch>.<pretrain_tag>"; run timm.list_models('*dinov2*', pretrained=True)
# to see the exact tags available in your installed version.
model = timm.create_model("vit_base_patch14_reg4_dinov2.lvd142m",
                          pretrained=True, num_classes=0)
model.eval()

# The checkpoint carries its own preprocessing config — never hand-roll it.
cfg = timm.data.resolve_model_data_config(model)
transform = timm.data.create_transform(**cfg, is_training=False)
print(cfg["input_size"], cfg["mean"], cfg["std"])

x = torch.randn(1, *cfg["input_size"])           # stand-in for transform(pil_image)
with torch.no_grad():
    pooled = model(x)                            # (1, D) global feature
    tokens = model.forward_features(x)           # (1, 1 + R + N, D) all tokens
print(pooled.shape, tokens.shape)

# pip install transformers pillow
from transformers import AutoModel, AutoProcessor
from PIL import Image
import torch

ckpt = "google/siglip-so400m-patch14-384"        # see also the facebook/dinov3-* collection
model = AutoModel.from_pretrained(ckpt).eval()
processor = AutoProcessor.from_pretrained(ckpt)  # resize + normalize + tokenize

image = Image.new("RGB", (640, 480), color=(120, 90, 60))
labels = ["a photo of a cat", "a photo of a dog", "a photo of a wall"]
inputs = processor(text=labels, images=image,
                   padding="max_length", return_tensors="pt")

with torch.no_grad():
    out = model(**inputs)

# SigLIP was trained with a per-pair sigmoid, so score each label INDEPENDENTLY
# with sigmoid — using softmax here (the CLIP habit) mis-calibrates the scores.
probs = torch.sigmoid(out.logits_per_image)      # (1, num_labels)
print({l: round(p, 4) for l, p in zip(labels, probs[0].tolist())})

# pip install open_clip_torch
import open_clip
import torch

model, _, preprocess = open_clip.create_model_and_transforms(
    "ViT-B-32", pretrained="laion2b_s34b_b79k")
tokenizer = open_clip.get_tokenizer("ViT-B-32")

with torch.no_grad():
    text = tokenizer(["a diagram", "a dog", "a cat"])
    text_features = model.encode_text(text)
    text_features /= text_features.norm(dim=-1, keepdim=True)  # L2-normalize
print(text_features.shape)                        # (3, 512)