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 embeddings to the children that support it as an extra input.

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

record training numerical indicators.

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

use logging to record log, only work on GPU 0 by judging global_rank

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

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

The full UNet model with attention and embedding. :param in_channel: channels in the input Tensor, for image colorization : Y_channels + X_ch

use tensorboard to record visuals, support 'add_scalar', 'add_scalars', 'add_image', 'add_images', etc. funtion. Also integrated with save r

Any module where forward() takes embeddings as a second argument.