easy-unet 0.1.1

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.

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

pip install easy-unet

📁 Project Structure

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

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

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)
• variational-autoencoder-pytorch-lib · PyPI (by me)
• convolutional-autoencoder-pytorch · PyPI (by me)

📜 License

This project is licensed under the terms of the MIT LICENSE.