
# easy-unet: Modular UNet Backbone in PyTorch

A **lightweight and flexible PyTorch library** providing a **modular UNet implementation** with advanced attention and normalization blocks. Designed for fast experimentation and development across diverse image processing and deep learning tasks.

## 🚀 Features

- 🧱 **UNet backbone** with residual blocks and flexible channel multipliers  
- 🎯 **Advanced attention modules** including linear and flash attention for improved feature modeling  
- ⚙️ **Configurable architecture** for dropout, channels, and dimensions  
- 🧪 Modular and clean PyTorch codebase suitable for research and production  
- 🔄 Supports easy integration with diffusion models, segmentation, or any custom pipeline  

## 📦 Installation

You can access the [PyPI page](https://pypi.org/project/easy-unet/) or install the package directly.

```bash
pip install easy-unet

```

## 📁 Project Structure

```bash
easy-unet/
├── easy_unet/
│   ├── __init__.py
│   ├── module.py        # All architecture classes and logic
├── pyproject.toml
├── LICENSE
└── README.md

```

## 🚀 Quick Start

### Import and create the model

```python
import torch
from easy_unet import UNet

model = UNet(
    dim=64,
    dim_mults=(1, 2, 4, 8),
    channels=3,
    out_channels=1,
    dropout=0.1
)

x = torch.randn(1, 3, 256, 256)    # sample input
output = model(x)
print(output.shape)  # e.g., torch.Size([1, 3, 256, 256])

```

## ⚙️ Configuration Options

| Argument | Type | Default | Description |
|--|--|--|--|
| `dim` | `int` | `64` | Base number of feature channels |
| `dim_mults` | `tuple` | `(1, 2, 4, 8)` | Channel multipliers per U-Net stage |
| `channels` | `int` | `3` | Number of input channels (e.g., 3 for RGB) |
| `out_channels` | `int` | `1` | Number of output channels (e.g., 1 for binary and 2 or more for multi-class) |
| `dropout` | `float` | `0.0` | Dropout rate for regularization. |

## 🙋‍♂️ Author

Developed by [Mehran Bazrafkan](mailto:mhrn.bzrafkn.dev@gmail.com)

> Created for general-purpose use cases and research requiring flexible UNet architectures in PyTorch.

## ⭐️ Support & Contribute

If you find this project useful, please:

- ⭐️ Star the repo

- 🐛 Report issues

- 📦 Suggest features or improvements

## 🔗 Related Projects

- [diffusion-pytorch-lib · PyPI (by me)](https://pypi.org/project/diffusion-pytorch-lib/)
- [variational-autoencoder-pytorch-lib · PyPI (by me)](https://pypi.org/project/variational-autoencoder-pytorch-lib/)
- [convolutional-autoencoder-pytorch · PyPI (by me)](https://pypi.org/project/convolutional-autoencoder-pytorch/)

## 📜 License

This project is licensed under the terms of the [`MIT LICENSE`](https://github.com/MehranBazrafkan/easy-unet/blob/main/LICENSE).
