code/bundle/models.py

"""Vision encoder zoo matching paper Table 12 (public checkpoints)."""

from __future__ import annotations

from dataclasses import dataclass
from typing import Callable

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


IMAGENET_TF = transforms.Compose(
    [
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ]
)


@dataclass
class Encoder:
    name: str
    family: str
    encode: Callable[[list[Image.Image]], np.ndarray]
    encode_tensor: Callable[[torch.Tensor], torch.Tensor]
    device: torch.device


def _pil_batch(images: list[Image.Image], tf, device: torch.device) -> torch.Tensor:
    return torch.stack([tf(im.convert("RGB")) for im in images]).to(device)


def _np(x: torch.Tensor) -> np.ndarray:
    return x.detach().float().cpu().numpy()


def load_open_clip(model_name: str, pretrained: str, device: torch.device, display: str, family: str) -> Encoder:
    import open_clip

    model, _, preprocess = open_clip.create_model_and_transforms(model_name, pretrained=pretrained)
    model = model.to(device).eval()
    for p in model.parameters():
        p.requires_grad_(False)

    def encode_tensor(x: torch.Tensor) -> torch.Tensor:
        # x already normalized ImageNet-style may differ from open_clip preprocess;
        # for JER noise probe we pass ImageNet-normalized tensors; for binding we use preprocess.
        feats = model.encode_image(x)
        return feats / feats.norm(dim=-1, keepdim=True).clamp_min(1e-12)

    def encode(images: list[Image.Image]) -> np.ndarray:
        with torch.no_grad():
            batch = torch.stack([preprocess(im.convert("RGB")) for im in images]).to(device)
            return _np(encode_tensor(batch))

    # For JER, rebuild a tensor path that applies open_clip normalization from [0,1]-ish noise is awkward.
    # Use model's visual forward with open_clip preprocess mean/std baked via a wrapper.
    mean = torch.tensor(preprocess.transforms[-1].mean, device=device).view(1, 3, 1, 1)
    std = torch.tensor(preprocess.transforms[-1].std, device=device).view(1, 3, 1, 1)

    def encode_tensor_from_imagenet_norm(x: torch.Tensor) -> torch.Tensor:
        # Convert ImageNet-normalized x back-ish is hard; instead assume x is already open_clip-normalized
        # when called from JER after we switch preprocess. We'll feed open_clip-normalized tensors from evaluate.
        feats = model.encode_image(x)
        return feats

    return Encoder(display, family, encode, encode_tensor_from_imagenet_norm, device)


def load_timm(model_name: str, device: torch.device, display: str, family: str) -> Encoder:
    import timm

    model = timm.create_model(model_name, pretrained=True, num_classes=0)
    model = model.to(device).eval()
    for p in model.parameters():
        p.requires_grad_(False)

    def encode_tensor(x: torch.Tensor) -> torch.Tensor:
        return model(x)

    def encode(images: list[Image.Image]) -> np.ndarray:
        with torch.no_grad():
            return _np(encode_tensor(_pil_batch(images, IMAGENET_TF, device)))

    return Encoder(display, family, encode, encode_tensor, device)


def load_torch_hub_resnet_features(
    repo: str,
    entry: str,
    device: torch.device,
    display: str,
    family: str,
    feature_layer: str = "avgpool",
) -> Encoder:
    model = torch.hub.load(repo, entry, trust_repo=True)
    model = model.to(device).eval()
    for p in model.parameters():
        p.requires_grad_(False)

    # Barlow twins hub returns a model that outputs 1000-d projection for barlowtwins;
    # VICReg/SwAV typically expose forward to features.
    def encode_tensor(x: torch.Tensor) -> torch.Tensor:
        out = model(x)
        if isinstance(out, (tuple, list)):
            out = out[0]
        if out.ndim > 2:
            out = out.flatten(1)
        return out

    def encode(images: list[Image.Image]) -> np.ndarray:
        with torch.no_grad():
            return _np(encode_tensor(_pil_batch(images, IMAGENET_TF, device)))

    return Encoder(display, family, encode, encode_tensor, device)


def load_dinov2(name: str, device: torch.device, display: str) -> Encoder:
    model = torch.hub.load("facebookresearch/dinov2", name, trust_repo=True)
    model = model.to(device).eval()
    for p in model.parameters():
        p.requires_grad_(False)

    def encode_tensor(x: torch.Tensor) -> torch.Tensor:
        return model(x)

    def encode(images: list[Image.Image]) -> np.ndarray:
        with torch.no_grad():
            return _np(encode_tensor(_pil_batch(images, IMAGENET_TF, device)))

    return Encoder(display, "Self-Distill", encode, encode_tensor, device)


def load_dino(name: str, device: torch.device, display: str) -> Encoder:
    model = torch.hub.load("facebookresearch/dino:main", name, trust_repo=True)
    model = model.to(device).eval()
    for p in model.parameters():
        p.requires_grad_(False)

    def encode_tensor(x: torch.Tensor) -> torch.Tensor:
        return model(x)

    def encode(images: list[Image.Image]) -> np.ndarray:
        with torch.no_grad():
            return _np(encode_tensor(_pil_batch(images, IMAGENET_TF, device)))

    return Encoder(display, "Self-Distill", encode, encode_tensor, device)


# Core suite prioritizing models needed for Claims 1–4. Full 26 when --suite full.
CORE_SPECS = [
    ("barlow_twins", "Var-Decorr"),
    ("vicreg", "Var-Decorr"),
    ("swav", "Clustering"),
    ("dinov2_vits14", "Self-Distill"),
    ("dinov2_vitb14", "Self-Distill"),
    ("dino_vits16", "Self-Distill"),
    ("mae_vit_base", "Masked"),
    ("clip_vit_b16", "Vision-Lang"),
    ("clip_vit_l14", "Vision-Lang"),
    ("clip_vit_b32", "Vision-Lang"),
    ("siglip_vit_b16", "Vision-Lang"),
    ("convnext_base", "Supervised"),
    ("vit_base_sup", "Supervised"),
]

FULL_EXTRA = [
    ("dinov2_vitl14", "Self-Distill"),
    ("dinov2_vitg14", "Self-Distill"),
    ("dino_vitb16", "Self-Distill"),
    ("mae_vit_large", "Masked"),
    ("beit_base", "Masked"),
    ("beitv2_base", "Masked"),
    ("siglip_so400m", "Vision-Lang"),
    ("eva_clip_b16", "Vision-Lang"),
    ("eva_clip_l14", "Vision-Lang"),
    ("convnext_large", "Supervised"),
    ("vit_large_sup", "Supervised"),
]


def load_encoder(key: str, device: torch.device | str = "cuda") -> Encoder:
    device = torch.device(device if torch.cuda.is_available() or str(device) == "cpu" else "cpu")
    if key == "barlow_twins":
        return load_torch_hub_resnet_features(
            "facebookresearch/barlowtwins:main", "resnet50", device, "Barlow Twins ResNet-50", "Var-Decorr"
        )
    if key == "vicreg":
        return load_torch_hub_resnet_features(
            "facebookresearch/vicreg:main", "resnet50", device, "VICReg ResNet-50", "Var-Decorr"
        )
    if key == "swav":
        return load_torch_hub_resnet_features(
            "facebookresearch/swav:main", "resnet50", device, "SwAV ResNet-50", "Clustering"
        )
    if key == "dinov2_vits14":
        return load_dinov2("dinov2_vits14", device, "DINOv2 ViT-S/14")
    if key == "dinov2_vitb14":
        return load_dinov2("dinov2_vitb14", device, "DINOv2 ViT-B/14")
    if key == "dinov2_vitl14":
        return load_dinov2("dinov2_vitl14", device, "DINOv2 ViT-L/14")
    if key == "dinov2_vitg14":
        return load_dinov2("dinov2_vitg14", device, "DINOv2 ViT-g/14")
    if key == "dino_vits16":
        return load_dino("dino_vits16", device, "DINO ViT-S/16")
    if key == "dino_vitb16":
        return load_dino("dino_vitb16", device, "DINO ViT-B/16")
    if key == "mae_vit_base":
        return load_timm("vit_base_patch16_224.mae", device, "MAE ViT-B/16", "Masked")
    if key == "mae_vit_large":
        return load_timm("vit_large_patch16_224.mae", device, "MAE ViT-L/16", "Masked")
    if key == "beit_base":
        return load_timm("beit_base_patch16_224", device, "BEiT ViT-B/16", "Masked")
    if key == "beitv2_base":
        return load_timm("beitv2_base_patch16_224", device, "BEiTv2 ViT-B/16", "Masked")
    if key == "clip_vit_b16":
        return load_open_clip("ViT-B-16", "openai", device, "CLIP ViT-B/16", "Vision-Lang")
    if key == "clip_vit_l14":
        return load_open_clip("ViT-L-14", "openai", device, "CLIP ViT-L/14", "Vision-Lang")
    if key == "clip_vit_b32":
        return load_open_clip("ViT-B-32", "openai", device, "CLIP ViT-B/32", "Vision-Lang")
    if key == "siglip_vit_b16":
        return load_open_clip("ViT-B-16-SigLIP", "webli", device, "SigLIP ViT-B/16", "Vision-Lang")
    if key == "siglip_so400m":
        return load_open_clip("ViT-SO400M-14-SigLIP", "webli", device, "SigLIP SoViT-400M", "Vision-Lang")
    if key == "eva_clip_b16":
        return load_open_clip("EVA02-B-16", "merged2b_s8b_b131k", device, "EVA-CLIP ViT-B/16", "Vision-Lang")
    if key == "eva_clip_l14":
        return load_open_clip("EVA02-L-14", "merged2b_s4b_b131k", device, "EVA-CLIP ViT-L/14", "Vision-Lang")
    if key == "convnext_base":
        return load_timm("convnext_base.fb_in22k_ft_in1k", device, "ConvNeXt Base", "Supervised")
    if key == "convnext_large":
        return load_timm("convnext_large.fb_in22k_ft_in1k", device, "ConvNeXt Large", "Supervised")
    if key == "vit_base_sup":
        return load_timm(
            "vit_base_patch16_224.augreg_in21k_ft_in1k", device, "ViT-B/16 Supervised", "Supervised"
        )
    if key == "vit_large_sup":
        return load_timm(
            "vit_large_patch16_224.augreg_in21k_ft_in1k", device, "ViT-L/16 Supervised", "Supervised"
        )
    raise KeyError(key)


def suite_keys(suite: str) -> list[str]:
    if suite == "smoke":
        return ["clip_vit_b32", "dinov2_vits14"]
    if suite == "claim4":
        return ["barlow_twins", "vicreg", "clip_vit_b16", "clip_vit_l14"]
    if suite == "core":
        return [k for k, _ in CORE_SPECS]
    if suite == "full":
        return [k for k, _ in CORE_SPECS] + [k for k, _ in FULL_EXTRA]
    raise ValueError(suite)