Vision Encoders for VLMs › Code Implementation: Processing Images with Vision Encoders

import numpy as np
import torch
import torch.nn as nn
from PIL import Image


## Create synthetic image batch for demonstration
def create_synthetic_images(batch_size=2, size=224):
    """Create synthetic images with geometric patterns."""
    images = []
    for i in range(batch_size):
        # Create pattern with different frequency components
        x = np.linspace(-1, 1, size)
        y = np.linspace(-1, 1, size)
        xx, yy = np.meshgrid(x, y)

        # Different pattern per image
        pattern = np.sin(4 * np.pi * (xx + yy * (i + 1)))
        pattern = (pattern + 1) / 2  # Normalize to [0, 1]

        # Convert to RGB
        img_array = np.stack([pattern, pattern * 0.8, pattern * 0.6], axis=-1)
        img_array = (img_array * 255).astype(np.uint8)
        images.append(Image.fromarray(img_array))

    return images


## Preprocessing
def preprocess_images(images, size=224):
    """Convert PIL images to normalized tensors."""
    tensors = []
    for img in images:
        img = img.resize((size, size))
        arr = np.array(img).astype(np.float32) / 255.0
        # Normalize with ImageNet stats
        mean = np.array([0.485, 0.456, 0.406])
        std = np.array([0.229, 0.224, 0.225])
        arr = (arr - mean) / std
        tensor = torch.from_numpy(arr).permute(2, 0, 1).float()
        tensors.append(tensor)

    return torch.stack(tensors)


# Simulated components for demonstration
class PatchEmbedding(nn.Module):
    """Converts images to patch embeddings."""

    def __init__(
        self, img_size=224, patch_size=16, in_channels=3, embed_dim=768
    ):
        super().__init__()
        self.patch_size = patch_size
        self.num_patches = (img_size // patch_size) ** 2
        self.proj = nn.Conv2d(
            in_channels, embed_dim, kernel_size=patch_size, stride=patch_size
        )

    def forward(self, x):
        # x: (batch, channels, height, width)
        x = self.proj(x)  # (batch, embed_dim, H/P, W/P)
        x = x.flatten(2).transpose(1, 2)  # (batch, num_patches, embed_dim)
        return x


class VisionEncoder(nn.Module):
    """Simplified Vision Transformer encoder for VLM."""

    def __init__(
        self,
        img_size=224,
        patch_size=16,
        embed_dim=768,
        num_layers=12,
        num_heads=12,
    ):
        super().__init__()
        self.patch_embed = PatchEmbedding(img_size, patch_size, 3, embed_dim)
        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))

        # Position embeddings
        num_patches = self.patch_embed.num_patches
        self.pos_embed = nn.Parameter(
            torch.zeros(1, num_patches + 1, embed_dim)
        )

        # Transformer layers
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=embed_dim,
            nhead=num_heads,
            dim_feedforward=embed_dim * 4,
            batch_first=True,
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers)

    def forward(self, x):
        batch_size = x.shape[0]
        x = self.patch_embed(x)

        # Add CLS token
        cls_tokens = self.cls_token.expand(batch_size, -1, -1)
        x = torch.cat([cls_tokens, x], dim=1)

        # Add position embeddings
        x = x + self.pos_embed

        # Transform
        x = self.transformer(x)
        return x  # (batch, num_patches+1, embed_dim)

def interpolate_pos_embed(pos_embed, new_num_patches):
    """
    Interpolate position embeddings for different resolution.

    Args:
        pos_embed: Original position embeddings (1, N+1, D)
        new_num_patches: New number of patches (N')
    Returns:
        Interpolated position embeddings (1, N'+1, D)
    """
    # Separate CLS token
    cls_embed = pos_embed[:, :1]
    patch_embed = pos_embed[:, 1:]

    # Get original grid size (assuming square)
    orig_size = int(np.sqrt(patch_embed.shape[1]))
    new_size = int(np.sqrt(new_num_patches))

    # Reshape to 2D grid
    patch_embed = patch_embed.reshape(1, orig_size, orig_size, -1)
    patch_embed = patch_embed.permute(0, 3, 1, 2)  # (1, D, H, W)

    # Interpolate
    patch_embed = torch.nn.functional.interpolate(
        patch_embed,
        size=(new_size, new_size),
        mode="bilinear",
        align_corners=False,
    )

    # Reshape back
    patch_embed = patch_embed.permute(0, 2, 3, 1).reshape(
        1, new_size * new_size, -1
    )

    # Concatenate CLS token back
    return torch.cat([cls_embed, patch_embed], dim=1)


# Simulate resolution increase
orig_patches = 196  # 14x14 grid at 224px with 16x16 patches
new_patches = 784  # 28x28 grid at 448px with 16x16 patches

# Original position embeddings
pos_embed_224 = torch.randn(1, orig_patches + 1, 768)

# Interpolate for higher resolution
pos_embed_448 = interpolate_pos_embed(pos_embed_224, new_patches)

def compute_similarities(
    image_embeds, text_embeds, method="siglip", temperature=0.01
):
    """
    Compute image-text similarities using different methods.

    Args:
        image_embeds: (batch, dim) image embeddings
        text_embeds: (batch, dim) text embeddings
        method: 'clip' (softmax) or 'siglip' (sigmoid)
        temperature: scaling factor
    """
    # Normalize embeddings
    image_embeds = nn.functional.normalize(image_embeds, dim=-1)
    text_embeds = nn.functional.normalize(text_embeds, dim=-1)

    # Compute similarity matrix
    logits = torch.matmul(image_embeds, text_embeds.t()) / temperature

    if method == "clip":
        # CLIP uses softmax across batch (symmetric loss)
        probs = nn.functional.softmax(logits, dim=-1)
        loss = -torch.log(torch.diag(probs)).mean()
        return logits, probs, loss
    else:
        # SigLIP uses independent sigmoids
        # Create target matrix (identity for matching pairs)
        batch_size = image_embeds.shape[0]
        targets = torch.eye(batch_size, device=image_embeds.device)

        # Binary cross-entropy for each pair independently
        probs = torch.sigmoid(logits)
        loss = -(
            targets * torch.log(probs + 1e-6)
            + (1 - targets) * torch.log(1 - probs + 1e-6)
        ).mean()
        return logits, probs, loss


# Simulate batch of embeddings
torch.manual_seed(42)
batch_size = 4
embed_dim = 512

# Create partially aligned embeddings with substantial noise so that
# off-diagonal similarities are non-negligible -- this makes the
# CLIP vs SigLIP contrast visible in the rendered heatmaps.
image_embeds = torch.randn(batch_size, embed_dim)
text_embeds = image_embeds + torch.randn(batch_size, embed_dim) * 1.2