crowsonkb/k-diffusion
A PyTorch implementation of diffusion-based generative models from the Karras et al. (2022) paper with transformer-based enhancements.

This library implements the Elucidating the Design Space of Diffusion-Based Generative Models paper by Karras et al. (2022) for PyTorch. It includes enhancements such as improved sampling algorithms and transformer-based diffusion models like image_transformer_v2, which combines ideas from Hourglass Transformer and DiT. The project supports training on image datasets like Oxford Flowers and requires custom CUDA kernels (NATTEN, FlashAttention-2) for optimal performance.