Add AutoModel.from_pretrained support (unet3plus+efficientnet)

README.md

@@ -106,42 +106,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/unet3plus-efficientnet-kvasir-seg")
-sys.path.insert(0, repo_dir)
-
-# 2. Import and instantiate the model
-from models import create_model
-model = create_model("unet3plus", backbone="efficientnet")
-
-# 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/unet3plus-efficientnet-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,10 +1,14 @@
 {
-  "architecture": "unet3plus",
+  "model_type": "unet3plus",
   "backbone": "efficientnet",
   "img_size": 256,
   "in_channels": 3,
   "out_channels": 1,
   "inter_ch": 64,
+  "auto_map": {
+    "AutoConfig": "configuration_unet3plus.UNet3PlusConfig",
+    "AutoModel": "modeling_unet3plus.UNet3PlusForSegmentation"
+  },
   "training": {
     "loss": "dice_focal",
     "focal_gamma": 1.1217038657488427,

configuration_unet3plus.py

@@ -0,0 +1,21 @@
+from transformers import PretrainedConfig
+
+
+class UNet3PlusConfig(PretrainedConfig):
+    model_type = "unet3plus"
+
+    def __init__(
+        self,
+        backbone: str = "efficientnet",
+        in_channels: int = 3,
+        out_channels: int = 1,
+        inter_ch: int = 64,
+        img_size: int = 256,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.backbone = backbone
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.inter_ch = inter_ch
+        self.img_size = img_size

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:94528930d89b8b147886031f2dbec583407cc3662b59909454c7721e513da72a
-size 52240388
+oid sha256:4b20d1c13554c4da7ea0a35318b90fda300f220a023224abdac5d4e658d780a1
+size 52243740

modeling_unet3plus.py

@@ -0,0 +1,48 @@
+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_unet3plus import UNet3PlusConfig
+from models import create_model
+
+
+class UNet3PlusForSegmentation(PreTrainedModel):
+    """UNet 3+ segmentation model with EfficientNet-B0 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 = UNet3PlusConfig
+    _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: UNet3PlusConfig):
+        super().__init__(config)
+        self.model = create_model(
+            architecture="unet3plus",
+            backbone=config.backbone,
+            in_channels=config.in_channels,
+            out_channels=config.out_channels,
+            inter_ch=config.inter_ch,
+        )
+
+    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}