A residual block that can optionally change the number of channels. :param channels: the number of input channels. :param emb_channels: t

A sequential module that passes timestep embeddings to the children that support it as an extra input.

An attention block that allows spatial positions to attend to each other. Originally ported from here, but adapted to the N-d case. https

A downsampling layer with an optional convolution. :param channels: channels in the inputs and outputs. :param use_conv: a bool determini

Transformer block for image-like data. First, project the input (aka embedding) and reshape to b, t, d. Then apply standard transform

An upsampling layer with an optional convolution. :param channels: channels in the inputs and outputs. :param use_conv: a bool determinin

A module which performs QKV attention and splits in a different order.

Implements exponential moving average shadowing for your model. Utilizes an inverse decay schedule to manage longer term training runs.

Uses the CLIP transformer encoder for text (from Hugging Face)

Utilities for training and sampling diffusion models. Original ported from this codebase: https://github.com/hojonathanho/diffusion/blob/

Uses xformers efficient implementation, see https://github.com/MatthieuTPHR/diffusers/blob/d80b531ff8060ec1ea982b65a1b8df70f73aa67c/src/diffu

A module which performs QKV attention. Matches legacy QKVAttention + input/ouput heads shaping

A diffusion process which can skip steps in a base diffusion process. :param use_timesteps: a collection (sequence or set) of timesteps from

SDE solver class

The full UNet model with attention and timestep embedding. :param in_channels: channels in the input Tensor. :param model_channels: base

Adapted from CLIP: https://github.com/openai/CLIP/blob/main/clip/model.py

Uses the OpenCLIP transformer encoder for text

Which type of output the model predicts.

What is used as the model's output variance. The LEARNED_RANGE option has been added to allow the model to predict values between FIXED_S

A distribution over timesteps in the diffusion process, intended to reduce variance of the objective. By default, samplers perform unbias

Any module where forward() takes timestep embeddings as a second argument.

Learned 2x upsampling without padding