Add AutoModel.from_pretrained support (attention_unet+convnext)

README.md

@@ -90,42 +90,33 @@ This model uses a custom PyTorch architecture. The model code is included in the
 ### Installation
 
 ```bash
-pip install torch torchvision timm safetensors huggingface_hub
+pip install torch torchvision timm transformers
 ```
 
 ### Inference
 
 ```python
-import sys
 import torch
-from huggingface_hub import snapshot_download
-from safetensors.torch import load_file
+from transformers import AutoModel
 from torchvision.transforms import functional as TF
 from PIL import Image
 
-# 1. Download the full repo (weights + model code) — handles LFS automatically
-repo_dir = snapshot_download("andreribeiro87/attention-unet-convnext-kvasir-seg")
-sys.path.insert(0, repo_dir)
-
-# 2. Import and instantiate the model
-from models import create_model
-model = create_model("attention_unet", backbone="convnext")
-
-# 3. Load weights
-state_dict = load_file(f"{repo_dir}/model.safetensors")
-model.load_state_dict(state_dict)
+# Load model — downloads weights + code automatically
+model = AutoModel.from_pretrained(
+    "andreribeiro87/attention-unet-convnext-kvasir-seg",
+    trust_remote_code=True,
+)
 model.eval()
 
-# 4. Preprocess an image
+# Preprocess
 image = Image.open("your_colonoscopy_image.jpg").convert("RGB")
 x = TF.to_tensor(TF.resize(image, [256, 256])).unsqueeze(0)  # (1, 3, 256, 256)
 
-# 5. Predict
+# Predict
 with torch.no_grad():
-    logit = model(x)                          # (1, 1, 256, 256)
-    mask = (logit.sigmoid() > 0.5).squeeze()  # bool tensor (256, 256)
+    outputs = model(pixel_values=x)
+    mask = (outputs["logits"].sigmoid() > 0.5).squeeze()  # bool (256, 256)
 
-# 6. Convert to PIL
 pred_mask = TF.to_pil_image(mask.float())
 ```
 

config.json

@@ -1,9 +1,13 @@
 {
-  "architecture": "attention_unet",
+  "model_type": "attention_unet",
   "backbone": "convnext",
   "img_size": 256,
   "in_channels": 3,
   "out_channels": 1,
+  "auto_map": {
+    "AutoConfig": "configuration_attention_unet.AttentionUNetConfig",
+    "AutoModel": "modeling_attention_unet.AttentionUNetForSegmentation"
+  },
   "training": {
     "loss": "bce_dice",
     "lr": 0.001,

configuration_attention_unet.py

@@ -0,0 +1,19 @@
+from transformers import PretrainedConfig
+
+
+class AttentionUNetConfig(PretrainedConfig):
+    model_type = "attention_unet"
+
+    def __init__(
+        self,
+        backbone: str = "convnext",
+        in_channels: int = 3,
+        out_channels: int = 1,
+        img_size: int = 256,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.backbone = backbone
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.img_size = img_size

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:b0c29289995adfc04f79dc1c5a0f1c0e0f85b3b19f11f0fab4ee0dd9699a2a5d
-size 139209476
+oid sha256:29d062c1486094cb2acbd057b6fa40f8883c48fe60db26f6db5332c2bd3a4ed4
+size 139211100

modeling_attention_unet.py

@@ -0,0 +1,47 @@
+import os
+import sys
+
+_here = os.path.dirname(os.path.abspath(__file__))
+if _here not in sys.path:
+    sys.path.insert(0, _here)
+
+import torch
+import torch.nn.functional as F
+from transformers import PreTrainedModel
+
+from configuration_attention_unet import AttentionUNetConfig
+from models import create_model
+
+
+class AttentionUNetForSegmentation(PreTrainedModel):
+    """Attention U-Net segmentation model with ConvNeXt-Tiny backbone.
+
+    Returns a dict with key "logits" (raw sigmoid input, shape B×1×H×W).
+    Pass pixel_values as a float32 tensor normalised to [0, 1], shape B×3×H×W.
+    """
+    config_class = AttentionUNetConfig
+    _tied_weights_keys = None
+    # Class-level fallback so transformers v5 finalization never triggers
+    # nn.Module.__getattr__ for this attribute (instance attr set in __init__ takes priority)
+    all_tied_weights_keys = {}
+
+    def __init__(self, config: AttentionUNetConfig):
+        super().__init__(config)
+        self.model = create_model(
+            architecture="attention_unet",
+            backbone=config.backbone,
+            in_channels=config.in_channels,
+            out_channels=config.out_channels,
+        )
+
+    def forward(
+        self,
+        pixel_values: torch.Tensor,
+        labels: torch.Tensor | None = None,
+        **kwargs,
+    ):
+        logits = self.model(pixel_values)
+        loss = None
+        if labels is not None:
+            loss = F.binary_cross_entropy_with_logits(logits, labels.float())
+        return {"loss": loss, "logits": logits}