Create a 1D, 2D, or 3D average pooling module.

Create a 1D, 2D, or 3D average pooling module.

Create a 1D, 2D, or 3D average pooling module.

Compute the entire variational lower-bound, measured in bits-per-dim, as well as other related quantities. :param model: t

Compute the entire variational lower-bound, measured in bits-per-dim, as well as other related quantities. :param model: t

Evaluate a function without caching intermediate activations, allowing for reduced memory at the expense of extra compute in the backward p

Evaluate a function without caching intermediate activations, allowing for reduced memory at the expense of extra compute in the backward pas

Evaluate a function without caching intermediate activations, allowing for reduced memory at the expense of extra compute in the backward pas

End DDP training.

r""" Copy from torch.nn.utils.clip_grad_norm_ Clips gradient norm of an iterable of parameters. The norm is computed over all gradients

Compute the log-likelihood of a continuous Gaussian distribution. :param x: the targets :param means: the Gaussian mean Tensor. :para

Compute the log-likelihood of a continuous Gaussian distribution. :param x: the targets :param means: the Gaussian mean Tensor. :para

Create a 1D, 2D, or 3D convolution module.

Create a 1D, 2D, or 3D convolution module.

Create a 1D, 2D, or 3D convolution module.

A counter for the `thop` package to count the operations in an attention operation. Meant to be used like: macs, params = tho

A counter for the `thop` package to count the operations in an attention operation. Meant to be used like: macs, params = thop.pr

A counter for the `thop` package to count the operations in an attention operation. Meant to be used like: macs, params = thop.pr

Create a logger that writes to a log file and stdout.

Create a logger that writes to a log file and stdout.

Create a ScheduleSampler from a library of pre-defined samplers. :param name: the name of the sampler. :param diffusion: the diffusion ob

Create a ScheduleSampler from a library of pre-defined samplers. :param name: the name of the sampler. :param diffusion: the diffusion ob

Create a tensorboard that saves losses.

Sample x_{t+1} from the model using DDIM reverse ODE.

Sample x_{t+1} from the model using DDIM reverse ODE.

r""" Disable sliced VAE decoding. If `enable_vae_slicing` was previously invoked, this method will go back to computing decoding in on

r""" Disable tiled VAE decoding. If `enable_vae_tiling` was previously invoked, this method will go back to computing decoding in one

r""" Offloads all models to CPU using accelerate, reducing memory usage with a low impact on performance. Compared to `enable_sequenti

r""" Offloads all models to CPU using accelerate, reducing memory usage with a low impact on performance. Compared to `enable_sequenti

r""" Offloads all models to CPU using accelerate, significantly reducing memory usage. When called, unet, text_encoder, vae and safety

r""" Offloads all models to CPU using accelerate, significantly reducing memory usage. When called, unet, text_encoder, vae and safety

r""" Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to compute decoding in se

r""" Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to compute decoding and e

The [`Transformer2DModel`] forward method. Args: hidden_states (`torch.LongTensor` of shape `(batch size, num latent pix