dit-ml 0.2.6

DiT (Diffusion Transformer)

dit_ml

The goal of dit_ml is to create a python repository to create the base model for diffusion transformer (https://arxiv.org/abs/2212.09748)

Also we incorporate RoPe embedding (for equivariance) taking inspiration from https://arxiv.org/pdf/2104.09864, https://arxiv.org/pdf/2403.13298 (2D mixed rope)

Installation

You can install dit_ml using pip:

pip install dit_ml

Usage

Here's a basic example of how to use dit_ml:

from dit_ml.dit import DiT

model = DiT(
    num_patches=input_size*input_size, # if 2d with flatten size
    hidden_size=hidden_size,
    depth=depth,
    num_heads=num_heads,
    learn_sigma=learn_sigma
)

dummy_x = torch.randn(batch_size, input_size * input_size, hidden_size)

dummy_c = torch.randn(batch_size, hidden_size) # Dummy conditioning vector

output = model(dummy_x, dummy_c) # of shape (batch_size, input_size * input_size, hidden_size)

Development

To set up the development environment:

1. Clone the repository:

git clone https://github.com/Forbu/dit_ml.git
cd dit_ml

2. Install dependencies using uv:

uv sync

3. Run tests:

uv run pytest

Contributing

Contributions are welcome! Please see the LICENSE for details.

License

This project is licensed under the Apache 2.0 - see the LICENSE file for details.