Initial RadJEPA encoder release

config.json

@@ -0,0 +1,6 @@
+{
+  "model_type": "radjepa",
+  "image_size": 224,
+  "patch_size": 14,
+  "embed_dim": 768
+}

jepa_encoder.pth.tar

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:afce5c46e600354b58033a53f88ecdc0da4a09308c5d0062f142465090e4e2aa
+size 1633156351

modeling_radjepa.py

@@ -0,0 +1,76 @@
+import torch
+import timm
+from timm.layers import PatchEmbed
+from transformers import PreTrainedModel, PretrainedConfig
+
+class RadJEPAConfig(PretrainedConfig):
+    model_type = "radjepa"
+
+    def __init__(
+        self,
+        image_size=224,
+        patch_size=14,
+        embed_dim=768,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.image_size = image_size
+        self.patch_size = patch_size
+        self.embed_dim = embed_dim
+
+
+class RadJEPAEncoder(PreTrainedModel):
+    config_class = RadJEPAConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.model = timm.create_model(
+            "vit_base_patch16_224",
+            pretrained=False,
+            num_classes=0
+        )
+
+        self.model.patch_embed = PatchEmbed(
+            img_size=config.image_size,
+            patch_size=config.patch_size,
+            in_chans=3,
+            embed_dim=config.embed_dim,
+        )
+
+        num_patches = self.model.patch_embed.num_patches
+
+        self.model.cls_token = None
+        self.model.num_prefix_tokens = 0
+
+        self.model.pos_embed = torch.nn.Parameter(
+            torch.zeros(1, num_patches, config.embed_dim)
+        )
+        torch.nn.init.trunc_normal_(self.model.pos_embed, std=0.02)
+
+    def forward(self, pixel_values):
+        tokens = self.model.forward_features(pixel_values)
+        return tokens.mean(dim=1)
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs):
+        config = RadJEPAConfig.from_pretrained(pretrained_model_name_or_path)
+        model = cls(config)
+
+        ckpt_path = f"{pretrained_model_name_or_path}/jepa_encoder.pth.tar"
+        ckpt = torch.load(ckpt_path, map_location="cpu")
+
+        if "encoder" in ckpt:
+            state_dict = ckpt["encoder"]
+        elif "state_dict" in ckpt and "encoder" in ckpt["state_dict"]:
+            state_dict = ckpt["state_dict"]["encoder"]
+        else:
+            raise RuntimeError("Encoder weights not found")
+
+        state_dict = {
+            k.replace("module.", "").replace("encoder.", ""): v
+            for k, v in state_dict.items()
+        }
+
+        model.model.load_state_dict(state_dict, strict=True)
+        return model