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.

Video to Patch Embedding. Args: patch_size (int): Patch token size. Default: (2,4,4). in_chans (int): Number of input video chann

Embeds class labels into vector representations. Also handles label dropout for classifier-free guidance.

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

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

Embeds scalar timesteps into vector representations.

Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).

The final layer of PixArt.

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 single linear layer which does a 1x1 conv

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

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

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

This JPEG algorithm result is slightly different from cv2. DiffJPEG supports batch processing. Args: differentiable(bool): If True,

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

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

Diffusion model with a Transformer backbone.

Uses the OpenCLIP transformer encoder for text

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

Embeds scalar timesteps into vector representations.

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

A wrapper around torchsde.BrownianTree that enables batches of entropy.

Merge patches into image

Splitting image into patches

A noise sampler backed by a torchsde.BrownianTree. Args: x (Tensor): The tensor whose shape, device and dtype to use to generate

Dequantize CbCr channel

JPEG Quantization for CbCr channels Args: rounding(function): rounding function to use

Chroma subsampling on CbCr channels

Upsample chroma layers

Full JPEG compression algorithm Args: rounding(function): rounding function to use