Causal 3d convolusion.

r""" A custom 3D causal convolution layer with feature caching support. This layer extends the standard Conv3D layer by ensuring causality in

r""" Norm layer adaptive layer norm single (adaLN-single). As proposed in PixArt-Alpha (see: https://arxiv.org/abs/2310.00426; Section 2.3).

Causal self-attention with a single head.

r""" Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0).

Shape of audio in VAE latent space: (batch, channels, frames, mel_bins). mel_bins is the number of frequency bins from the mel-spectrogram en

Per-pixel (per-location) RMS normalization layer. For each element along the chosen dimension, this layer normalizes the tensor by the ro

r""" A custom RMS normalization layer. Args: dim (int): The number of dimensions to normalize over. channel_first (bool, opti

N-dimensional pixel shuffle operation for upsampling tensors. Args: dims (int): Number of dimensions to apply pixel shuffle to.

r""" A custom residual block module. Args: in_dim (int): Number of input channels. out_dim (int): Number of output channels.

r""" A Resnet block. Parameters: in_channels (`int`): The number of channels in the input. out_channels (`int`, *optional*, d

For PixArt-Alpha. Reference: https://github.com/PixArt-alpha/PixArt-alpha/blob/0f55e922376d8b797edd44d25d0e7464b260dcab/diffusion/model/n

r""" A Resnet block. Parameters: in_channels (`int`): The number of channels in the input. out_channels (`int`, *optional*, de

A 3D UNet mid-block [`UNetMidBlock3D`] with multiple residual blocks. Args: in_channels (`int`): The number of input channels.

Intervals which a single dimension of the latent space is split into. Each interval is defined by its start, end, left ramp, and right ramp. T

A 2D downsampling layer with an optional convolution. Parameters: channels (`int`): number of channels in the inputs and outp

Embeddings1DConnector applies a 1D transformer-based processing to sequential embeddings (e.g., for video, audio, or other modalities). It su

Flux 2 uses a SwiGLU-style activation in the transformer feedforward sub-blocks, but with the linear projection layer fused into the first li

Input data for a single modality (video or audio) in the transformer. Bundles the latent tokens, timestep embeddings, positional information,

Per-channel statistics for normalizing and denormalizing the latent representation. This statics is computed over the entire dataset and stor

Per-channel statistics for normalizing and denormalizing the latent representation. This statics is computed over the entire dataset and stor

Projects caption embeddings. Also handles dropout for classifier-free guidance. Adapted from https://github.com/PixArt-alpha/PixArt-alpha/blo

r""" Causal self-attention with a single head. Args: dim (int): The number of channels in the input tensor.

Middle block for WanVAE encoder and decoder. Args: dim (int): Number of input/output channels. dropout (float): Dropout rate

r""" A custom resampling module for 2D and 3D data. Args: dim (int): The number of input/output channels. mode (str): The res

r""" Perform upsampling while ensuring the output tensor has the same data type as the input. Args: x (torch.Tensor): Input tensor to

Residual block with two convolutional layers, group normalization, and SiLU activation. Args: channels (int): Number of input and out

Configuration for dividing each frame into spatial tiles with optional overlap. Args: tile_size_in_pixels (int): Size of each tile in pixe

Configuration for dividing a video into temporal tiles (chunks of frames) with optional overlap. Args: tile_size_in_frames (int): Number o

A 2D UNet mid-block [`UNetMidBlock2D`] with multiple residual blocks and optional attention blocks. Args: in_channels (`int`): The n

A 2D upsampling layer with an optional convolution. Parameters: channels (`int`): number of channels in the inputs and output

Shape of the tensor representing video in VAE latent space. The latent representation is a 5D tensor with dimensions ordered as (batch, c

XLMRobertaModel with no pooler and no LM head.