Package: smalldiffusion
smalldiffusionA lightweight diffusion library for training and sampling from diffusion
models. It is built for easy experimentation when training new models and
developing new samplers, supporting minimal toy models to state-of-the-art
pretrained models. Thecore of this libraryfor diffusion
training and sampling is implemented in less than 100 lines of very readable
pytorch code. To install frompypi:pip install smalldiffusionToy modelsTo train and sample from theSwissrolltoy dataset in 10 lines of code (seeexamples/toyexample.ipynbfor a detailed
guide):fromtorch.utils.dataimportDataLoaderfromsmalldiffusionimportSwissroll,TimeInputMLP,ScheduleLogLinear,training_loop,samplesdataset=Swissroll(np.pi/2,5*np.pi,100)loader=DataLoader(dataset,batch_size=2048)model=TimeInputMLP(hidden_dims=(16,128,128,128,128,16))schedule=ScheduleLogLinear(N=200,sigma_min=0.005,sigma_max=10)trainer=training_loop(loader,model,schedule,epochs=15000)losses=[ns.loss.item()fornsintrainer]*xt,x0=samples(model,schedule.sample_sigmas(20),gam=2)Results on various toy datasets:U-Net modelsThe same code can be used to trainU-Net-based models. To
train a model on the FashionMNIST dataset and generate a batch of samples (after
first runningaccelerate config):accelerate launch examples/fashion_mnist.pyWith the provided default parameters and training on a single GPU for around 2
hours, the model can achieve aFID
scoreof
around 12-13, producing the following generated outputs:StableDiffusionsmalldiffusion's sampler works with any pretrained diffusion model, and supports
DDPM, DDIM as well as accelerated sampling algorithms. Inexamples/diffusers_wrapper.py, we provide a
simple wrapper for any pretrainedhuggingface
diffuserslatent diffusion model,
enabling sampling from pretrained models with only a few lines of code:fromdiffusers_wrapperimportModelLatentDiffusionfromsmalldiffusionimportScheduleLDM,samplesschedule=ScheduleLDM(1000)model=ModelLatentDiffusion('stabilityai/stable-diffusion-2-1-base')model.set_text_condition('An astronaut riding a horse')*xts,x0=samples(model,schedule.sample_sigmas(50))decoded=model.decode_latents(x0)It is easy to experiment with different sampler parameters and sampling
schedules, as demonstrated inexamples/stablediffusion.py. A
few examples on tweaking the parametergam:How to useThe core of smalldiffusion depends on the interaction betweendata,modelandscheduleobjects. Here we give a specification of these objects. For a
detailed introduction to diffusion models and the notation used in the code, see
theaccompanying tutorial.DataFor training diffusion models, smalldiffusion supports pytorchDatasetsandDataLoaders.
The training code expects the iterates from aDataLoaderobject to be batches
of data, without labels. To remove labels from existing datasets, extract the
data with the providedMappedDatasetwrapper before constructing aDataLoader.Two toy datasets,SwissrollandDatasaurusDozen,
are provided.ModelAll model objects should be a subclass oftorch.nn.Module. Models should have:A parameterinput_dims, a tuple containing the dimensions of the input to
the model (not including batch-size).A methodrand_input(batchsize)which takes in a batch-size and returns an
i.i.d. standard normal random input with shape[batchsize, *input_dims]. This method can be inherited from the providedModelMixinclass when theinput_dimsparameter is set.Models are called with two arguments:xis a batch of data of batch-sizeBand shape[B, *model.input_dims].sigmais either a singleton or a batch.Ifsigma.shape == [], the same value will be used for eachx.Otherwisesigma.shape == [B, 1, ..., 1], andx[i]will be paired withsigma[i].Models should return a predicted noise value with the same shape asx.ScheduleAScheduleobject determines the rate at which the noise levelsigmaincreases during the diffusion process. It is constructed by simply passing in a
tensor of increasingsigmavalues.Scheduleobjects have the methodssample_sigmas(steps)which subsamples the schedule for sampling.sample_batch(batchsize)which generates batch ofsigmavalues selected
uniformly at random, for use in training.Three schedules are provided:ScheduleLogLinearis a simple schedule which works well on small
datasets and toy models.ScheduleDDPMis commonly used in pixel-space image diffusion models.ScheduleLDMis commonly used in latent diffusion models,
e.g. StableDiffusion.The following plot shows these three schedules with default parameters.TrainingThetraining_loopgenerator function provides a simple training loop for
training a diffusion model , givenloader,modelandscheduleobjects
described above. It yields a namespace with the local variables, for easy
evaluation during training. For example, to print out the loss every iteration:for ns in training_loop(loader, model, schedule):
    print(ns.loss.item())Multi-GPU training and sampling is also supported viaaccelerate.SamplingTo sample from a diffusion model, thesamplesgenerator function takes in amodeland a decreasing list ofsigmasto use during sampling. This list is
usually created by calling thesample_sigmas(steps)method of aScheduleobject. The generator will yield a sequence ofxts produced during
sampling. The sampling loop generalizes most commonly-used samplers:For DDPM[Ho et. al. ], usegam=1, mu=0.5.For DDIM[Song et. al. ], usegam=1, mu=0.For accelerated sampling[Permenter and Yuan], usegam=2.For more details on how these sampling algorithms can be simplified, generalized
and implemented in only 5 lines of code, see Appendix A of[Permenter and
Yuan].