Upload folder using huggingface_hub

hf_wrapper/brain_unet_hf/config.json

@@ -0,0 +1,11 @@
+{
+  "architectures": [
+    "UNetTransformerModel"
+  ],
+  "image_size": 256,
+  "in_channels": 1,
+  "model_type": "unet",
+  "out_channels": 3,
+  "torch_dtype": "float32",
+  "transformers_version": "4.53.2"
+}

hf_wrapper/brain_unet_hf/model.safetensors

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

hf_wrapper/convert_save.py

@@ -0,0 +1,15 @@
+import torch
+from model import UNetTransformerModel, UNetConfig
+
+# Load config and model
+config = UNetConfig(in_channels=1, out_channels=3, image_size=256)
+model = UNetTransformerModel(config)
+
+# Load your existing model weights
+model.model.load_state_dict(torch.load("unet_epoch20.pth", map_location="cpu"))
+
+# Save the model and config in HF-compatible format
+model.save_pretrained("brain_unet_hf")
+config.save_pretrained("brain_unet_hf")
+
+print("✅ Saved to brain_unet_hf/")

hf_wrapper/model.py

@@ -0,0 +1,119 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PreTrainedModel, PretrainedConfig  
+
+class DoubleConv(nn.Module):
+    """
+    This is the core building block of the U-Net architecture.
+    Use consecutive convolutional layers
+    Each followed by batch normalization and ReLU activation
+    """
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        """
+        nn.Conv2d:
+        Applies a 2D convolution filter (kernel size 3×3)
+        padding=1 ensures the output spatial size stays the same
+        First conv changes input channels → output channels
+        Second conv keeps it at out_channels
+
+        nn.BatchNorm2d
+        Normalizes activations across the batch and channels
+        Helps stabilize and speed up training
+        Reduces internal covariate shift
+        """
+        self.double_conv = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True)
+        )
+    
+    def forward(self, x):
+        return self.double_conv(x)
+
+class UNet(nn.Module):
+    def __init__(self, in_channels=1, out_channels=3):
+        super().__init__()
+
+        # Encoder
+        self.down1 = DoubleConv(in_channels, 64)
+        self.pool1 = nn.MaxPool2d(2)
+
+        self.down2 = DoubleConv(64, 128)
+        self.pool2 = nn.MaxPool2d(2)
+
+        self.down3 = DoubleConv(128, 256)
+        self.pool3 = nn.MaxPool2d(2)
+
+        self.down4 = DoubleConv(256, 512)
+        self.pool4 = nn.MaxPool2d(2)
+
+        # Bottleneck
+        self.bottleneck = DoubleConv(512, 1024)
+
+        # Decoder
+        self.up4 = nn.ConvTranspose2d(1024, 512, kernel_size=2, stride=2)
+        self.dec4 = DoubleConv(1024, 512)
+
+        self.up3 = nn.ConvTranspose2d(512, 256, kernel_size=2, stride=2)
+        self.dec3 = DoubleConv(512, 256)
+
+        self.up2 = nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2)
+        self.dec2 = DoubleConv(256, 128)
+
+        self.up1 = nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2)
+        self.dec1 = DoubleConv(128, 64)
+
+        # Final output layer
+        self.out_conv = nn.Conv2d(64, out_channels, kernel_size=1)
+
+    def forward(self, x):
+        # Encoder
+        d1 = self.down1(x)
+        d2 = self.down2(self.pool1(d1))
+        d3 = self.down3(self.pool2(d2))
+        d4 = self.down4(self.pool3(d3))
+
+        # Bottleneck
+        bn = self.bottleneck(self.pool4(d4))
+
+        # Decoder
+        up4 = self.up4(bn)
+        dec4 = self.dec4(torch.cat([up4, d4], dim=1))
+
+        up3 = self.up3(dec4)
+        dec3 = self.dec3(torch.cat([up3, d3], dim=1))
+
+        up2 = self.up2(dec3)
+        dec2 = self.dec2(torch.cat([up2, d2], dim=1))
+
+        up1 = self.up1(dec2)
+        dec1 = self.dec1(torch.cat([up1, d1], dim=1))
+
+        # Output
+        return self.out_conv(dec1)
+    
+# --- Config ---
+class UNetConfig(PretrainedConfig):
+    model_type = "unet"
+
+    def __init__(self, in_channels=1, out_channels=3, image_size=256, **kwargs):
+        super().__init__(**kwargs)
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.image_size = image_size
+
+# --- Transformer-compatible Model ---
+class UNetTransformerModel(PreTrainedModel):
+    config_class = UNetConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = UNet(config.in_channels, config.out_channels)
+
+    def forward(self, pixel_values):
+        return self.model(pixel_values)

hf_wrapper/unet_epoch20.pth

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

unet.py

@@ -1,6 +1,7 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+from transformers import PreTrainedModel, PretrainedConfig  
 
 class DoubleConv(nn.Module):
     """
@@ -94,4 +95,25 @@ class UNet(nn.Module):
         dec1 = self.dec1(torch.cat([up1, d1], dim=1))
 
         # Output
-        return self.out_conv(dec1)
+        return self.out_conv(dec1)
+    
+# --- Config ---
+class UNetConfig(PretrainedConfig):
+    model_type = "unet"
+
+    def __init__(self, in_channels=1, out_channels=3, image_size=256, **kwargs):
+        super().__init__(**kwargs)
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.image_size = image_size
+
+# --- Transformer-compatible Model ---
+class UNetTransformerModel(PreTrainedModel):
+    config_class = UNetConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = UNet(config.in_channels, config.out_channels)
+
+    def forward(self, pixel_values):
+        return self.model(pixel_values)