| """Minimal model-loading code for the 7 VFM baselines in the paper.""" |
|
|
| from __future__ import annotations |
|
|
| import sys |
| from pathlib import Path |
| from typing import Callable |
|
|
| import torch |
| import torch.nn as nn |
| from torchvision import transforms |
| from transformers import AutoConfig, AutoImageProcessor, AutoModel |
|
|
| ROOT = Path(__file__).resolve().parent |
| WEIGHTS_DIR = ROOT / "weights" |
|
|
| MODEL_SPECS = { |
| "metacliplin": { |
| "paper_name": "MetaCLIP-Linear", |
| "checkpoint": "metacliplin0.pth", |
| "hf_model": "facebook/metaclip-h14-fullcc2.5b", |
| "feature_dim": 1280, |
| "image_size": 224, |
| "pooler_output": True, |
| }, |
| "metaclip2lin": { |
| "paper_name": "MetaCLIP2-Linear", |
| "checkpoint": "metaclip2lin0.pth", |
| "hf_model": "facebook/metaclip-2-worldwide-giant", |
| "feature_dim": 1280, |
| "image_size": 224, |
| "pooler_output": True, |
| }, |
| "sigliplin": { |
| "paper_name": "SigLIP-Linear", |
| "checkpoint": "sigliplin0.pth", |
| "hf_model": "google/siglip-large-patch16-384", |
| "feature_dim": 1024, |
| "image_size": 384, |
| "pooler_output": True, |
| }, |
| "siglip2lin": { |
| "paper_name": "SigLIP2-Linear", |
| "checkpoint": "siglip2lin0.pth", |
| "hf_model": "google/siglip2-giant-opt-patch16-384", |
| "feature_dim": 1536, |
| "image_size": 384, |
| "pooler_output": True, |
| }, |
| "pelin": { |
| "paper_name": "PE-CLIP-Linear", |
| "checkpoint": "pelin0.pth", |
| "feature_dim": 1024, |
| "image_size": 336, |
| "pooler_output": False, |
| }, |
| "dinov2lin": { |
| "paper_name": "DINOv2-Linear", |
| "checkpoint": "dinov2lin0.pth", |
| "feature_dim": 1024, |
| "pooler_output": False, |
| }, |
| "dinov3lin": { |
| "paper_name": "DINOv3-Linear", |
| "checkpoint": "dinov3lin0.pth", |
| "hf_model": "facebook/dinov3-vit7b16-pretrain-lvd1689m", |
| "feature_dim": 4096, |
| "pooler_output": False, |
| }, |
| } |
|
|
| ALIASES = { |
| "MetaCLIP-Linear": "metacliplin", |
| "MetaCLIP2-Linear": "metaclip2lin", |
| "SigLIP-Linear": "sigliplin", |
| "SigLIP2-Linear": "siglip2lin", |
| "PE-CLIP-Linear": "pelin", |
| "DINOv2-Linear": "dinov2lin", |
| "DINOv3-Linear": "dinov3lin", |
| } |
|
|
|
|
| def canonical_model_name(name: str) -> str: |
| if name in MODEL_SPECS: |
| return name |
| if name in ALIASES: |
| return ALIASES[name] |
| raise KeyError(f"Unknown model: {name}") |
|
|
|
|
| def default_checkpoint_path(model_name: str) -> Path: |
| model_name = canonical_model_name(model_name) |
| return WEIGHTS_DIR / MODEL_SPECS[model_name]["checkpoint"] |
|
|
|
|
| def _resolve_device(device: str | torch.device | None = None) -> torch.device: |
| if device is None: |
| return torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| return torch.device(device) |
|
|
|
|
| def _load_checkpoint(checkpoint_path: str | Path) -> dict: |
| checkpoint = torch.load(str(checkpoint_path), map_location="cpu", weights_only=False) |
| if isinstance(checkpoint, dict): |
| for key in ("state_dict", "model", "model_state_dict"): |
| if key in checkpoint and isinstance(checkpoint[key], dict): |
| checkpoint = checkpoint[key] |
| break |
|
|
| normalized = {} |
| for key, value in checkpoint.items(): |
| normalized[key[7:] if key.startswith("module.") else key] = value |
| return normalized |
|
|
|
|
| def _infer_feature_dim(state_dict: dict, default_dim: int) -> int: |
| head_weight = state_dict.get("head.weight") |
| if isinstance(head_weight, torch.Tensor) and head_weight.ndim == 2: |
| return int(head_weight.shape[1]) |
| return default_dim |
|
|
|
|
| def _load_image_processor(model_name: str): |
| try: |
| return AutoImageProcessor.from_pretrained(model_name, local_files_only=True) |
| except Exception: |
| try: |
| return AutoImageProcessor.from_pretrained(model_name) |
| except Exception: |
| return None |
|
|
|
|
| def _load_backbone(model_name: str): |
| try: |
| return AutoModel.from_pretrained(model_name, local_files_only=True) |
| except Exception: |
| config = AutoConfig.from_pretrained(model_name) |
| return AutoModel.from_config(config) |
|
|
|
|
| class _PoolerLinearModel(nn.Module): |
| def __init__(self, backbone: nn.Module, feature_dim: int): |
| super().__init__() |
| self.backbone = backbone |
| self.head = nn.Linear(feature_dim, 2) |
|
|
| def forward(self, x): |
| with torch.no_grad(): |
| outputs = self.backbone(x) |
| features = outputs.pooler_output.float() |
| return self.head(features) |
|
|
|
|
| class _ClsTokenLinearModel(nn.Module): |
| def __init__(self, backbone: nn.Module, feature_dim: int): |
| super().__init__() |
| self.backbone = backbone |
| self.head = nn.Linear(feature_dim, 2) |
|
|
| def forward(self, x): |
| with torch.no_grad(): |
| outputs = self.backbone(x) |
| features = outputs.last_hidden_state[:, 0].float() |
| return self.head(features) |
|
|
|
|
| class _PELinearModel(nn.Module): |
| def __init__(self, backbone: nn.Module, feature_dim: int): |
| super().__init__() |
| self.backbone = backbone |
| self.head = nn.Linear(feature_dim, 2) |
|
|
| def forward(self, x): |
| with torch.no_grad(): |
| features = self.backbone(x) |
| if isinstance(features, torch.Tensor): |
| features = features.float() |
| return self.head(features) |
|
|
|
|
| def _finalize_model(model: nn.Module, state_dict: dict, device=None) -> nn.Module: |
| model.load_state_dict(state_dict, strict=False) |
| model.to(_resolve_device(device)) |
| model.eval() |
| return model |
|
|
|
|
| def _build_clip_transform(image_size: int, image_processor=None): |
| mean = [0.485, 0.456, 0.406] |
| std = [0.229, 0.224, 0.225] |
| if image_processor is not None: |
| mean = getattr(image_processor, "image_mean", mean) |
| std = getattr(image_processor, "image_std", std) |
| return transforms.Compose( |
| [ |
| transforms.Resize(image_size, interpolation=transforms.InterpolationMode.BICUBIC), |
| transforms.CenterCrop(image_size), |
| transforms.ToTensor(), |
| transforms.Normalize(mean=mean, std=std), |
| ] |
| ) |
|
|
|
|
| def _build_dino_transform(): |
| return transforms.Compose( |
| [ |
| transforms.Resize(256, interpolation=transforms.InterpolationMode.BICUBIC), |
| transforms.CenterCrop(224), |
| transforms.ToTensor(), |
| transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), |
| ] |
| ) |
|
|
|
|
| def load_metacliplin(checkpoint_path: str | Path | None = None, device=None): |
| spec = MODEL_SPECS["metacliplin"] |
| checkpoint_path = checkpoint_path or default_checkpoint_path("metacliplin") |
| state_dict = _load_checkpoint(checkpoint_path) |
| feature_dim = _infer_feature_dim(state_dict, spec["feature_dim"]) |
| image_processor = _load_image_processor(spec["hf_model"]) |
| backbone = _load_backbone(spec["hf_model"]) |
| model = _PoolerLinearModel(backbone.vision_model, feature_dim) |
| model = _finalize_model(model, state_dict, device=device) |
| return model, _build_clip_transform(spec["image_size"], image_processor) |
|
|
|
|
| def load_metaclip2lin(checkpoint_path: str | Path | None = None, device=None): |
| spec = MODEL_SPECS["metaclip2lin"] |
| checkpoint_path = checkpoint_path or default_checkpoint_path("metaclip2lin") |
| state_dict = _load_checkpoint(checkpoint_path) |
| feature_dim = _infer_feature_dim(state_dict, spec["feature_dim"]) |
| image_processor = _load_image_processor(spec["hf_model"]) |
| backbone = _load_backbone(spec["hf_model"]) |
| model = _PoolerLinearModel(backbone.vision_model, feature_dim) |
| model = _finalize_model(model, state_dict, device=device) |
| return model, _build_clip_transform(spec["image_size"], image_processor) |
|
|
|
|
| def load_sigliplin(checkpoint_path: str | Path | None = None, device=None): |
| spec = MODEL_SPECS["sigliplin"] |
| checkpoint_path = checkpoint_path or default_checkpoint_path("sigliplin") |
| state_dict = _load_checkpoint(checkpoint_path) |
| feature_dim = _infer_feature_dim(state_dict, spec["feature_dim"]) |
| image_processor = _load_image_processor(spec["hf_model"]) |
| backbone = _load_backbone(spec["hf_model"]) |
| model = _PoolerLinearModel(backbone.vision_model, feature_dim) |
| model = _finalize_model(model, state_dict, device=device) |
| return model, _build_clip_transform(spec["image_size"], image_processor) |
|
|
|
|
| def load_siglip2lin(checkpoint_path: str | Path | None = None, device=None): |
| spec = MODEL_SPECS["siglip2lin"] |
| checkpoint_path = checkpoint_path or default_checkpoint_path("siglip2lin") |
| state_dict = _load_checkpoint(checkpoint_path) |
| feature_dim = _infer_feature_dim(state_dict, spec["feature_dim"]) |
| image_processor = _load_image_processor(spec["hf_model"]) |
| backbone = _load_backbone(spec["hf_model"]) |
| model = _PoolerLinearModel(backbone.vision_model, feature_dim) |
| model = _finalize_model(model, state_dict, device=device) |
| return model, _build_clip_transform(spec["image_size"], image_processor) |
|
|
|
|
| def load_dinov2lin(checkpoint_path: str | Path | None = None, device=None): |
| checkpoint_path = checkpoint_path or default_checkpoint_path("dinov2lin") |
| state_dict = _load_checkpoint(checkpoint_path) |
| feature_dim = _infer_feature_dim(state_dict, MODEL_SPECS["dinov2lin"]["feature_dim"]) |
| if feature_dim == 1536: |
| candidates = ["facebook/dinov2-giant", "facebook/dinov2-large"] |
| elif feature_dim == 1024: |
| candidates = ["facebook/dinov2-large", "facebook/dinov2-base"] |
| elif feature_dim == 768: |
| candidates = ["facebook/dinov2-base", "facebook/dinov2-small"] |
| else: |
| candidates = ["facebook/dinov2-large"] |
|
|
| last_error = None |
| backbone = None |
| for candidate in candidates: |
| try: |
| backbone = _load_backbone(candidate) |
| break |
| except Exception as exc: |
| last_error = exc |
| if backbone is None: |
| raise RuntimeError(f"Failed to load DINOv2 backbone: {last_error}") |
|
|
| model = _ClsTokenLinearModel(backbone, feature_dim) |
| model = _finalize_model(model, state_dict, device=device) |
| return model, _build_dino_transform() |
|
|
|
|
| def load_dinov3lin(checkpoint_path: str | Path | None = None, device=None): |
| checkpoint_path = checkpoint_path or default_checkpoint_path("dinov3lin") |
| state_dict = _load_checkpoint(checkpoint_path) |
| feature_dim = _infer_feature_dim(state_dict, MODEL_SPECS["dinov3lin"]["feature_dim"]) |
| backbone = _load_backbone(MODEL_SPECS["dinov3lin"]["hf_model"]) |
| model = _ClsTokenLinearModel(backbone, feature_dim) |
| model = _finalize_model(model, state_dict, device=device) |
| return model, _build_dino_transform() |
|
|
|
|
| def load_pelin(checkpoint_path: str | Path | None = None, device=None): |
| checkpoint_path = checkpoint_path or default_checkpoint_path("pelin") |
| if str(ROOT) not in sys.path: |
| sys.path.insert(0, str(ROOT)) |
|
|
| import core.vision_encoder.pe as pe |
| import core.vision_encoder.transforms as pe_transforms |
|
|
| state_dict = _load_checkpoint(checkpoint_path) |
| feature_dim = _infer_feature_dim(state_dict, MODEL_SPECS["pelin"]["feature_dim"]) |
| clip_model = pe.CLIP.from_config("PE-Core-L14-336", pretrained=False) |
| model = _PELinearModel(clip_model.visual, feature_dim) |
| model = _finalize_model(model, state_dict, device=device) |
| return model, pe_transforms.get_image_transform(MODEL_SPECS["pelin"]["image_size"]) |
|
|
|
|
| LOADERS: dict[str, Callable] = { |
| "metacliplin": load_metacliplin, |
| "metaclip2lin": load_metaclip2lin, |
| "sigliplin": load_sigliplin, |
| "siglip2lin": load_siglip2lin, |
| "pelin": load_pelin, |
| "dinov2lin": load_dinov2lin, |
| "dinov3lin": load_dinov3lin, |
| } |
|
|
|
|
| def load_model(model_name: str, checkpoint_path: str | Path | None = None, device=None): |
| model_name = canonical_model_name(model_name) |
| return LOADERS[model_name](checkpoint_path=checkpoint_path, device=device) |
|
|
