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

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

A class that keeps all conversation history.

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

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

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

Processes an answer similar to Eval AI copied from https://github.com/facebookresearch/mmf/blob/c46b3b3391275b4181567db80943473a8

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

Fake file-like stream object that redirects writes to a logger instance.

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

Collate examples for supervised fine-tuning.

Helper for the face restoration pipeline (base class).

Uses the CLIP transformer encoder for text (from huggingface)

Uses the T5 transformer encoder for text

Dataset for supervised fine-tuning.

r""" Sampler that samples indices in a way that groups together features of the dataset of roughly the same length while keeping a bit of rand

to match AttnBlock usage

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

Defines a PatchGAN discriminator as in Pix2Pix --> see https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/blob/master/models/networks.py

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

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

This is the base class for all autoencoders, including image autoencoders, image autoencoders with discriminators, unCLIP models, etc. Hence,

AutoTokenizer + Adaptation for MOD. A simple wrapper around AutoTokenizer to make instantiating an MOD-adapted tokenizer a bit easier. M

Base class for all image autoencoders that we train, like VQGAN or AutoencoderKL (we also restore them explicitly as special cases for legacy

embeds each dimension independently and concatenates them

The half UNet model with attention and timestep embedding. For usage, see UNet.