diffusion_models.models.openai_unet.ResBlock

class diffusion_models.models.openai_unet.ResBlock(*args: Any, **kwargs: Any)

Bases: TimestepBlock

A residual block that can optionally change the number of channels.

Parameters:

• channels – the number of input channels.

• emb_channels – the number of timestep embedding channels.

• dropout – the rate of dropout.

• out_channels – if specified, the number of out channels.

• use_conv – if True and out_channels is specified, use a spatial convolution instead of a smaller 1x1 convolution to change the channels in the skip connection.

• dims – determines if the signal is 1D, 2D, or 3D.

• use_checkpoint – if True, use gradient checkpointing on this module.

__init__(channels, emb_channels, dropout, out_channels=None, use_conv=False, use_scale_shift_norm=False, dims=2, use_checkpoint=False)

Methods

__init__(channels, emb_channels, dropout[, ...])
forward(x, emb)	Apply the block to a Tensor, conditioned on a timestep embedding.

forward(x, emb)

Apply the block to a Tensor, conditioned on a timestep embedding.

Parameters:

• x – an [N x C x …] Tensor of features.

• emb – an [N x emb_channels] Tensor of timestep embeddings.

Returns:

an [N x C x …] Tensor of outputs.