Move preprocessing outside XMag model

README.md

@@ -24,12 +24,21 @@ XMag is a pathology image encoder distilled from a high-magnification foundation
 
 ```python
 import torch
+from PIL import Image
+from torchvision import transforms
 from transformers import AutoModel
 
+eval_transform = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
+])
+
 model = AutoModel.from_pretrained("AI4PATH/XMAG", trust_remote_code=True)
 model.eval()
 
-pixel_values = torch.randn(1, 3, 224, 224)
+image = Image.open("patch.png").convert("RGB")
+pixel_values = eval_transform(image).unsqueeze(0)
 
 with torch.no_grad():
     outputs = model(pixel_values)
@@ -38,8 +47,8 @@ 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.
+`pixel_values` should be normalized RGB tensors with shape `(B, 3, 224, 224)`.
+The model does not apply preprocessing internally. Resize, `ToTensor()`, and ImageNet mean/std normalization should be done before calling the model.
 
 ## Model Details
 

config.json

@@ -20,6 +20,7 @@
     0.224,
     0.225
   ],
+  "expects_normalized_input": true,
   "torch_dtype": "float32",
   "transformers_version": "4.50.3"
 }

modeling_xmag.py

@@ -1,5 +1,4 @@
 import torch
-import torch.nn as nn
 from transformers import PreTrainedModel
 
 from .configuration_xmag import XMagConfig
@@ -13,11 +12,6 @@ class XMagModel(PreTrainedModel):
         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).")
@@ -28,7 +22,6 @@ class XMagModel(PreTrainedModel):
         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 {

upload_to_hub.py

@@ -11,10 +11,13 @@ def main():
     parser.add_argument("--commit-message", default="Add XMag AutoModel release")
     args = parser.parse_args()
 
-    if os.environ.get("HF_TOKEN") is None:
+    token = os.environ.get("HF_TOKEN")
+    if token is None:
         print("warning: HF_TOKEN is not set. Falling back to cached Hugging Face credentials.")
+        login()
+    else:
+        login(token=token)
 
-    login()
     result = upload_folder(
         folder_path=args.folder,
         repo_id=args.repo_id,