k-diffusion

k-diffusion

一个用于 PyTorch 的 Elucidating the Design Space of Diffusion-Based Generative Models (Karras 等人, 2022) 的实现,具有增强功能和额外特性,例如改进的采样算法和基于 Transformer 的扩散模型。

• NATTEN 用于层级低层的稀疏（邻域）注意力。该模型类型有一个移位窗口注意力版本，不需要自定义 CUDA 内核，但性能较差，且训练和推理速度较慢。

• FlashAttention-2 用于全局注意力。如果未安装，它将回退到纯 PyTorch。

此外，您应确保您的 PyTorch 安装能够使用 torch.compile()。如果 torch.compile() 不可用，它将回退到 eager 模式，但训练速度会更慢，且内存占用更高。

用法

演示

要在不安装自定义CUDA内核的情况下，在Oxford Flowers上训练一个256x256的RGB模型，请安装Hugging Face Datasets：

pip install datasets

然后运行：

python train.py --config configs/config_oxford_flowers_shifted_window.json --name flowers_demo_001 --evaluate-n 0 --batch-size 32 --sample-n 36 --mixed-precision bf16

如果内存不足，请尝试添加--checkpointing或减小批次大小。如果您使用的是较旧的GPU（Ampere之前），请省略--mixed-precision bf16以在FP32模式下训练。不建议在FP16模式下训练。

如果您已安装并正常使用NATTEN（推荐），可以通过指定--config configs/config_oxford_flowers.json，使用邻域注意力替代移位窗口注意力进行训练。

续