Cleanup and update README.

README.md

@@ -16,7 +16,7 @@ pipeline_tag: image-segmentation
 This repository contains an implementation of U-Net [[1]](#references). [unet.py](./unet.py) implements the class UNet. The implementation has been tested with PyTorch 2.7.1 and CUDA 12.6.
 ![](./images/unet_diagram.png)
 
-You can also load the U-Net from PyTorch Hub.
+You can load the U-Net from PyTorch Hub.
 ```python
 import torch
 
@@ -32,6 +32,14 @@ model = torch.hub.load('sm00thix/unet', 'unet', pretrained=False, in_channels=3,
 # model = torch.hub.load('sm00thix/unet', 'unet_transconv', **kwargs) # Convenience function equivalent to torch.hub.load('sm00thix/unet', 'unet', bilinear=False, **kwargs)
 ```
 
+You can also clone this repository to access the U-Net directly.
+```python
+import torch
+from unet import UNet
+
+model = UNet(in_channels=3, out_channels=1, pad=True, bilinear=True, normalization=None)
+```
+
 ## Options
 The UNet class provides the following options for customization.
 

config.json

@@ -1,9 +0,0 @@
-{
-  "model_type": "unet",
-  "architectures": ["UNet"],
-  "in_channels": 3,
-  "out_channels": 1,
-  "pad": true,
-  "bilinear": true,
-  "normalization": null
-}

requirements.txt

@@ -1,2 +1 @@
 torch >= 2.7.1
-transformers >= 4.55.2

unet_config.py

@@ -1,21 +0,0 @@
-from transformers import PretrainedConfig
-
-
-class UNetConfig(PretrainedConfig):
-    model_type = "unet"
-
-    def __init__(
-        self,
-        in_channels=3,
-        out_channels=1,
-        pad=True,
-        bilinear=True,
-        normalization=None,
-        **kwargs,
-    ):
-        super().__init__(**kwargs)
-        self.in_channels = in_channels
-        self.out_channels = out_channels
-        self.pad = pad
-        self.bilinear = bilinear
-        self.normalization = normalization

unet_hf.py

@@ -1,20 +0,0 @@
-from transformers import PreTrainedModel
-from .unet import UNet
-from .unet_config import UNetConfig
-
-
-class UNetModel(PreTrainedModel):
-    config_class = UNetConfig
-
-    def __init__(self, config: UNetConfig):
-        super().__init__(config)
-        self.model = UNet(
-            in_channels=config.in_channels,
-            out_channels=config.out_channels,
-            pad=config.pad,
-            bilinear=config.bilinear,
-            normalization=config.normalization,
-        )
-
-    def forward(self, x):
-        return self.model(x)