Ecosyste.ms: Awesome
An open API service indexing awesome lists of open source software.
https://github.com/dirmeier/diffusion-transformer
A diffusion transformer implementation in Flax
https://github.com/dirmeier/diffusion-transformer
diffusion-transformer dit flax jax python
Last synced: about 1 month ago
JSON representation
A diffusion transformer implementation in Flax
- Host: GitHub
- URL: https://github.com/dirmeier/diffusion-transformer
- Owner: dirmeier
- License: apache-2.0
- Created: 2024-08-23T13:09:58.000Z (4 months ago)
- Default Branch: main
- Last Pushed: 2024-09-09T07:05:11.000Z (4 months ago)
- Last Synced: 2024-10-13T00:13:04.661Z (3 months ago)
- Topics: diffusion-transformer, dit, flax, jax, python
- Language: Python
- Homepage:
- Size: 95.7 KB
- Stars: 1
- Watchers: 2
- Forks: 0
- Open Issues: 0
-
Metadata Files:
- Readme: README.md
- License: LICENSE
Awesome Lists containing this project
README
# DiT
[![status](http://www.repostatus.org/badges/latest/concept.svg)](http://www.repostatus.org/#concept)
[![ci](https://github.com/dirmeier/diffusion-transformer/actions/workflows/ci.yaml/badge.svg)](https://github.com/dirmeier/diffusion-transformer/actions/workflows/ci.yaml)> A diffusion transformer implementation in Flax
## About
This repository implements the diffusion transformer (DiT) architecture
which has been proposed in [Scalable Diffusion Models with Transformers](https://arxiv.org/abs/2212.09748)
in Flax. We test the architecture using the EDM parameterization introduced in [Elucidating the Design Space of Diffusion-Based Generative Models](https://arxiv.org/abs/2206.00364).> [!NOTE]
> The architecture does not follow the [original implementation](https://github.com/facebookresearch/DiT) exactly.
> For instance, we don't use label conditioning and just use a learnable positional encoding for the patches.
> The rest is fairly similar.## Example usage
The `experiments` folder contains a use case that trains an EDM (diffusion model) on MNIST-SDF.
To train the model, just execute:```bash
cd experiments/mnist_sdf
python main.py
--config=config.py
--workdir=
(--usewand)
```Below are some samples drawn from the EDM using a DiT-B after training 100 epochs.
In my experiments, the UNet still works better, but that might just be how the hyperparameters have been chosen.
## Installation
To install the latest GitHub , just call the following on the command line:
```bash
pip install git+https://github.com/dirmeier/diffusion-transformer@
```## Author
Simon Dirmeier sfyrbnd @ pm me