/k-diffusion

Karras et al. (2022) diffusion models for PyTorch

Primary LanguagePythonMIT LicenseMIT

k-diffusion

An implementation of Elucidating the Design Space of Diffusion-Based Generative Models (Karras et al., 2022) for PyTorch.

This repo is a work in progress (models may break on later versions, script options may change).

Multi-GPU and multi-node training is supported with Hugging Face Accelerate. You can configure Accelerate by running:

$ accelerate config

on all nodes, then running:

$ accelerate launch train.py --train-set LOCATION_OF_TRAINING_SET --model-config MODEL_CONFIG

on all nodes.