Implementation of a Diffusion Transformer and Rectified Flow Sampling in Jax
github.com