A residual block that can optionally change the number of channels. :param channels: the number of input channels. :param emb_channels: the number of timestep embedding channels. :param dropout: the rate of dropout. :param out_channels: if specified, the number of out channels.
| 207 | |
| 208 | |
| 209 | class ResBlock(TimestepBlock): |
| 210 | """ |
| 211 | A residual block that can optionally change the number of channels. |
| 212 | :param channels: the number of input channels. |
| 213 | :param emb_channels: the number of timestep embedding channels. |
| 214 | :param dropout: the rate of dropout. |
| 215 | :param out_channels: if specified, the number of out channels. |
| 216 | :param use_conv: if True and out_channels is specified, use a spatial |
| 217 | convolution instead of a smaller 1x1 convolution to change the |
| 218 | channels in the skip connection. |
| 219 | :param dims: determines if the signal is 1D, 2D, or 3D. |
| 220 | :param use_checkpoint: if True, use gradient checkpointing on this module. |
| 221 | :param up: if True, use this block for upsampling. |
| 222 | :param down: if True, use this block for downsampling. |
| 223 | """ |
| 224 | |
| 225 | def __init__( |
| 226 | self, |
| 227 | channels, |
| 228 | emb_channels, |
| 229 | dropout, |
| 230 | out_channels=None, |
| 231 | use_conv=False, |
| 232 | use_scale_shift_norm=False, |
| 233 | dims=2, |
| 234 | use_checkpoint=False, |
| 235 | up=False, |
| 236 | down=False, |
| 237 | kernel_size=3, |
| 238 | exchange_temb_dims=False, |
| 239 | skip_t_emb=False, |
| 240 | ): |
| 241 | super().__init__() |
| 242 | self.channels = channels |
| 243 | self.emb_channels = emb_channels |
| 244 | self.dropout = dropout |
| 245 | self.out_channels = out_channels or channels |
| 246 | self.use_conv = use_conv |
| 247 | self.use_checkpoint = use_checkpoint |
| 248 | self.use_scale_shift_norm = use_scale_shift_norm |
| 249 | self.exchange_temb_dims = exchange_temb_dims |
| 250 | |
| 251 | if isinstance(kernel_size, Iterable): |
| 252 | padding = [k // 2 for k in kernel_size] |
| 253 | else: |
| 254 | padding = kernel_size // 2 |
| 255 | |
| 256 | self.in_layers = nn.Sequential( |
| 257 | normalization(channels), |
| 258 | nn.SiLU(), |
| 259 | conv_nd(dims, channels, self.out_channels, kernel_size, padding=padding), |
| 260 | ) |
| 261 | |
| 262 | self.updown = up or down |
| 263 | |
| 264 | if up: |
| 265 | self.h_upd = Upsample(channels, False, dims) |
| 266 | self.x_upd = Upsample(channels, False, dims) |