Metadata-Version: 2.4
Name: easy-unet
Version: 0.1.0
Summary: A lightweight and flexible PyTorch library providing a modular UNet implementation with advanced attention and normalization blocks for fast image processing and deep learning development.
Author-email: Mehran Bazrafkan <mhrn.bzrafkn.dev@gmail.com>
License-Expression: MIT
Keywords: UNet,PyTorch,deep learning,image segmentation,attention,normalization,machine learning,computer vision
Classifier: Programming Language :: Python :: 3
Classifier: Operating System :: OS Independent
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Intended Audience :: Developers
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=1.13
Requires-Dist: einops
Requires-Dist: numpy
Dynamic: license-file


# 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

```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,
    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/output channels (e.g., 3 for RGB) |
| `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).
