A generic multi-modal dataset loader where the samples are arranged in this way: :: root/modality_a/class_x/xxx.ext root/modality_a/c

4M model. Args: encoder_embeddings: Dict of encoder embedding modules. decoder_embeddings: Dict of decoder embedding modules.

Modified 4M model, adapted to behave as a simple RGB-only ViT. Args: img_size (int): Input image size. patch_size (int): Patch si

UViT model = Conditional UNet with Transformer bottleneck blocks and optionalpatching. See https://arxiv.org/abs/2301.11093 for more details.

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

Subclass torch's LayerNorm to handle fp16.

A single linear layer which does a 1x1 conv

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

Sampler that wraps a trained 4M model for generation use cases. Implements standard autoregressive, MaskGIT, and ROAR generation schemes with chai

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

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

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

Implements SwiGLU and other gated feed-forward layers from Noam Shazeer's paper: https://arxiv.org/abs/2002.05202

Model Exponential Moving Average (DEPRECATED) Keep a moving average of everything in the model state_dict (parameters and buffers). This ver

Pipeline for conditional image generation. This model inherits from `DiffusionPipeline`. Check the superclass documentation for the generic metho

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

Transformer to map latent features back to images / feature maps. Args: out_channels: Number of output channels. patch_size:

Transformer to map images / feature maps to latent features. Args: in_channels: Number of input channels. patch_size: Patch s

r""" ConvNeXt Block. There are two equivalent implementations: (1) DwConv -> LayerNorm (channels_first) -> 1x1 Conv -> GELU -> 1x1 Conv; all in (N

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

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

Track a series of values and provide access to smoothed values over a window or the global series average.

Iterates through each of `model`'s modules and matches modules using unix pattern matching based on `hook_fn_locs` and registers `hook_fn` to the

AugMix Transform Adapted and improved from impl here: https://github.com/google-research/augmix/blob/master/imagenet.py From paper: 'AugMix:

Output class for the scheduler's step function output. Args: prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, h

Output class for the scheduler's step function output. Args: prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, h

DiVAE ("Diffusion VQ-VAE") model = simple encoder + diffusion decoder with a discrete bottleneck, inspired by https://arxiv.org/abs/2206.00386.

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

Empty dataset

CNN Feature Map Embedding Extract feature map from CNN, flatten, project to embedding dim.

Embedding module for spatial inputs, like images or feature maps. Creates tokens from patches over the image. This adapter / embedding differ

Low-Rank Adaptation Wrapper for linear layers. See https://arxiv.org/abs/2106.09685 Args: linear: nn.Linear layer to wrap

A ResNet class that is similar to torchvision's but contains the following changes: - There are now 3 "stem" convolutions as opposed to 1, wi