Add XMag AutoModel release

README.md

@@ -1,3 +1,63 @@
 ---
-license: cc-by-nc-nd-4.0
+library_name: transformers
+pipeline_tag: image-feature-extraction
+tags:
+- pathology
+- histopathology
+- foundation-model
+- distillation
+- dinov2
+- xmag
 ---
+
+# XMag
+
+XMag is a pathology image encoder distilled from a high-magnification foundation model to a low-magnification student model.
+
+- Teacher: frozen UNIv2
+- Student: DINOv2 ViT-B/14
+- Input: RGB pathology patch, `224 x 224`
+- Output: CLS embedding and patch embeddings
+- Training objective: global and local cosine feature distillation from high-magnification teacher features
+
+## Usage
+
+```python
+import torch
+from transformers import AutoModel
+
+model = AutoModel.from_pretrained("AI4PATH/XMAG", trust_remote_code=True)
+model.eval()
+
+pixel_values = torch.randn(1, 3, 224, 224)
+
+with torch.no_grad():
+    outputs = model(pixel_values)
+
+cls_embedding = outputs["cls_embedding"]          # (1, 768)
+patch_embeddings = outputs["patch_embeddings"]    # (1, 256, 768)
+```
+
+`pixel_values` should be RGB tensors with shape `(B, 3, 224, 224)`, typically after `torchvision.transforms.ToTensor()`.
+The model applies ImageNet mean/std normalization internally.
+
+## Model Details
+
+The student sees a low-magnification `224 x 224` view of the tissue region. During training, the frozen UNIv2 teacher sees the corresponding high-magnification `896 x 896` region split into `4 x 4` subpatches. The student is trained to match both:
+
+- A global teacher representation, computed by averaging the 16 teacher local features.
+- Local teacher representations, aligned to pooled student patch-token blocks.
+
+This release contains the student EMA backbone weights only. Projection heads used during distillation are not included.
+
+## Requirements
+
+```bash
+pip install torch torchvision transformers huggingface_hub
+```
+
+The model code uses `torch.hub` to instantiate the DINOv2 ViT-B/14 backbone. The first load may need internet access to fetch the DINOv2 hub code, unless it is already cached.
+
+## Citation
+
+If you use this model, please cite the related XMag project and the original teacher/student foundation models where appropriate.

config.json

@@ -0,0 +1,25 @@
+{
+  "model_type": "xmag",
+  "architectures": [
+    "XMagModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_xmag.XMagConfig",
+    "AutoModel": "modeling_xmag.XMagModel"
+  },
+  "backbone_name": "dinov2_vitb14",
+  "embed_dim": 768,
+  "image_size": 224,
+  "image_mean": [
+    0.485,
+    0.456,
+    0.406
+  ],
+  "image_std": [
+    0.229,
+    0.224,
+    0.225
+  ],
+  "torch_dtype": "float32",
+  "transformers_version": "4.50.3"
+}

configuration_xmag.py

@@ -0,0 +1,21 @@
+from transformers import PretrainedConfig
+
+
+class XMagConfig(PretrainedConfig):
+    model_type = "xmag"
+
+    def __init__(
+        self,
+        backbone_name="dinov2_vitb14",
+        embed_dim=768,
+        image_size=224,
+        image_mean=(0.485, 0.456, 0.406),
+        image_std=(0.229, 0.224, 0.225),
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.backbone_name = backbone_name
+        self.embed_dim = embed_dim
+        self.image_size = image_size
+        self.image_mean = list(image_mean)
+        self.image_std = list(image_std)

modeling_xmag.py

@@ -0,0 +1,37 @@
+import torch
+import torch.nn as nn
+from transformers import PreTrainedModel
+
+from .configuration_xmag import XMagConfig
+
+
+class XMagModel(PreTrainedModel):
+    config_class = XMagConfig
+    base_model_prefix = "model"
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = torch.hub.load("facebookresearch/dinov2", config.backbone_name)
+
+        mean = torch.tensor(config.image_mean).view(1, 3, 1, 1)
+        std = torch.tensor(config.image_std).view(1, 3, 1, 1)
+        self.register_buffer("image_mean", mean, persistent=False)
+        self.register_buffer("image_std", std, persistent=False)
+
+    def forward(self, pixel_values):
+        if pixel_values.ndim != 4:
+            raise ValueError("pixel_values must have shape (B, 3, H, W).")
+
+        _, channels, height, width = pixel_values.shape
+        if channels != 3:
+            raise ValueError("pixel_values must have 3 RGB channels.")
+        if height != self.config.image_size or width != self.config.image_size:
+            raise ValueError(f"pixel_values must be {self.config.image_size}x{self.config.image_size}.")
+
+        pixel_values = (pixel_values - self.image_mean) / self.image_std
+        outputs = self.model.forward_features(pixel_values)
+
+        return {
+            "cls_embedding": outputs["x_norm_clstoken"],
+            "patch_embeddings": outputs["x_norm_patchtokens"],
+        }

pytorch_model.bin

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

requirements.txt

@@ -0,0 +1,4 @@
+torch
+torchvision
+transformers
+huggingface_hub

upload_to_hub.py

@@ -0,0 +1,28 @@
+import argparse
+import os
+
+from huggingface_hub import login, upload_folder
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Upload the XMag Hugging Face release folder.")
+    parser.add_argument("--repo-id", default="AI4PATH/XMAG", help="Target Hugging Face model repo.")
+    parser.add_argument("--folder", default=".", help="Release folder to upload.")
+    parser.add_argument("--commit-message", default="Add XMag AutoModel release")
+    args = parser.parse_args()
+
+    if os.environ.get("HF_TOKEN") is None:
+        print("warning: HF_TOKEN is not set. Falling back to cached Hugging Face credentials.")
+
+    login()
+    result = upload_folder(
+        folder_path=args.folder,
+        repo_id=args.repo_id,
+        repo_type="model",
+        commit_message=args.commit_message,
+    )
+    print(result)
+
+
+if __name__ == "__main__":
+    main()