"""Differentiable loader for NVIDIA C-RADIOv3-H (agglomerative vision encoder).

C-RADIO distills CLIP + SigLIP2 + DINOv2 + SAM into one tower, so it's a strong,
architecturally-distinct held-out JUDGE: a perturbation that transfers to it is
moving shared cross-architecture directions rather than CLIP idiosyncrasies.

The model takes pixel values in [0,1] and applies its own input conditioner, so
image_feat only resizes to a patch-multiple resolution and returns the pooled
`summary` embedding. Loaded in float32 (proven differentiable; 652M params ~2.6GB).
"""

from __future__ import annotations

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import AutoModel

REPO = "nvidia/C-RADIOv3-H"
RES = 384  # multiple of the 16px patch size


def load_cradio(dtype: torch.dtype = torch.float32) -> nn.Module:
    model = AutoModel.from_pretrained(REPO, trust_remote_code=True).eval().to("cuda")
    for p in model.parameters():
        p.requires_grad_(False)
    return model


def image_feat_cradio(model: nn.Module, x: torch.Tensor) -> torch.Tensor:
    """x: (1,3,H,W) in [0,1] on cuda. Returns (1, D) pooled summary embedding."""
    mdtype = next(model.parameters()).dtype
    px = F.interpolate(x, size=(RES, RES), mode="bicubic", align_corners=False).clamp(0, 1)
    out = model(px.to(mdtype))
    summ = out.summary if hasattr(out, "summary") else (out[0] if isinstance(out, (tuple, list)) else out)
    return summ.float()