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

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

Channel-wise self-attention block.

An attention block that allows spatial positions to attend to each other. Originally ported from here, but adapted to the N-d case. http

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

Maintains (exponential) moving average of a set of parameters.

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

Object that stores parameters like a dictionary. This allows perturbation classes to be only partially instantiated and then fed vari

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

Based on code from https://github.com/yaodongyu/TRADES

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

The 1931 CIE XYZ color space (assuming input is in sRGB). Warning: may have values outside [0, 1] range. Should only be used in the proc

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

Utilities for training and sampling diffusion models. Ported directly from here, and then adapted over time to further experimentation.

modified from: https://pytorch.org/docs/stable/_modules/torchvision/datasets/folder.html#ImageFolder uses cached directory listing if availab

An empty corrector that does nothing.

An empty predictor that does nothing.

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

Puts the color at each pixel in the image through the same transformation. Parameters: - lp_style: number or 'inf' - lp_bound: max

Takes a list of perturbations and composes them. A norm needs to be specified here to describe the perturbations.

A diffusion process which can skip steps in a base diffusion process. :param use_timesteps: a collection (sequence or set) of timesteps from

A UNetModel that performs super-resolution. Expects an extra kwarg `low_res` to condition on a low-resolution image.

Dumps key/value pairs into TensorBoard's numeric format.

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

Wrapper for adversarial attacks. Is helpful for when subsidiary methods are needed.

Wrapper to store an adversarial attack object as well as some extra parameters for how to use it in training

Skeleton class to hold adversarial perturbations FOR A SINGLE MINIBATCH. For general input-agnostic adversarial perturbations, see the

Wrapper for training of a NN with adversarial examples cooked in

Affine transformation -- just has 6 parameters per example: 4 for 2d rotation, and 1 for translation in each direction

The ancestral sampling predictor. Currently only supports VE/VP SDEs.

The original annealed Langevin dynamics predictor in NCSN/NCSNv2. We include this corrector only for completeness. It was not directly used in our

Converts from RGB to approximately the HSV cone using a much smoother transformation.

Channel-wise self-attention block. Modified from DDPM.

Computes and stores the average and current value

Converts to the 1976 CIE L*u*v* color space.