mmagic.models.editors.ddpm.res_blocks

Module Contents

Classes

ResnetBlock2D	resnet block support down sample and up sample.
Upsample2D	An upsampling layer with an optional convolution.
Downsample2D	A downsampling layer with an optional convolution.

class mmagic.models.editors.ddpm.res_blocks.ResnetBlock2D(in_channels, out_channels=None, conv_shortcut=False, dropout=0.0, temb_channels=512, groups=32, groups_out=None, pre_norm=True, eps=1e-06, non_linearity='silu', time_embedding_norm='default', output_scale_factor=1.0, use_in_shortcut=None, up=False, down=False)

Bases: torch.nn.Module

resnet block support down sample and up sample.

Parameters

• in_channels (int) – input channels.

• out_channels (int) – output channels.

• conv_shortcut (bool) – whether to use conv shortcut.

• dropout (float) – dropout rate.

• temb_channels (int) – time embedding channels.

• groups (int) – conv groups.

• groups_out (int) – conv out groups.

• pre_norm (bool) – whether to norm before conv. Todo: remove.

• eps (float) – eps for groupnorm.

• non_linearity (str) – non linearity type.

• time_embedding_norm (str) – time embedding norm type.

• output_scale_factor (float) – factor to scale input and output.

• use_in_shortcut (bool) – whether to use conv in shortcut.

• up (bool) – whether to upsample.

• down (bool) – whether to downsample.

forward(input_tensor, temb)

forward with hidden states and time embeddings.

class mmagic.models.editors.ddpm.res_blocks.Upsample2D(channels, use_conv=False, use_conv_transpose=False, out_channels=None, name='conv')

Bases: torch.nn.Module

An upsampling layer with an optional convolution.

Parameters

• channels (int) – channels in the inputs and outputs.

• use_conv (bool) – a bool determining if a convolution is applied.

• use_conv_transpose (bool) – whether to use conv transpose.

• out_channels (int) – output channels.

forward(hidden_states, output_size=None)

forward with hidden states.

class mmagic.models.editors.ddpm.res_blocks.Downsample2D(channels, use_conv=False, out_channels=None, padding=1, name='conv')

Bases: torch.nn.Module

A downsampling layer with an optional convolution.

Parameters

• channels (int) – channels in the inputs and outputs.

• use_conv (bool) – a bool determining if a convolution is applied.

• out_channels (int) – output channels

• padding (int) – padding num

forward(hidden_states)

forward with hidden states.