r""" A Resnet block that use normalization layer that incorporate conditioning information. Parameters: in_channels (`int`): The number of channels in the input. out_channels (`int`, *optional*, default to be `None`): The number of output channels for the first c
| 42 | |
| 43 | |
| 44 | class ResnetBlockCondNorm2D(nn.Module): |
| 45 | r""" |
| 46 | A Resnet block that use normalization layer that incorporate conditioning information. |
| 47 | |
| 48 | Parameters: |
| 49 | in_channels (`int`): The number of channels in the input. |
| 50 | out_channels (`int`, *optional*, default to be `None`): |
| 51 | The number of output channels for the first conv2d layer. If None, same as `in_channels`. |
| 52 | dropout (`float`, *optional*, defaults to `0.0`): The dropout probability to use. |
| 53 | temb_channels (`int`, *optional*, default to `512`): the number of channels in timestep embedding. |
| 54 | groups (`int`, *optional*, default to `32`): The number of groups to use for the first normalization layer. |
| 55 | groups_out (`int`, *optional*, default to None): |
| 56 | The number of groups to use for the second normalization layer. if set to None, same as `groups`. |
| 57 | eps (`float`, *optional*, defaults to `1e-6`): The epsilon to use for the normalization. |
| 58 | non_linearity (`str`, *optional*, default to `"swish"`): the activation function to use. |
| 59 | time_embedding_norm (`str`, *optional*, default to `"ada_group"` ): |
| 60 | The normalization layer for time embedding `temb`. Currently only support "ada_group" or "spatial". |
| 61 | kernel (`torch.Tensor`, optional, default to None): FIR filter, see |
| 62 | [`~models.resnet.FirUpsample2D`] and [`~models.resnet.FirDownsample2D`]. |
| 63 | output_scale_factor (`float`, *optional*, default to be `1.0`): the scale factor to use for the output. |
| 64 | use_in_shortcut (`bool`, *optional*, default to `True`): |
| 65 | If `True`, add a 1x1 nn.conv2d layer for skip-connection. |
| 66 | up (`bool`, *optional*, default to `False`): If `True`, add an upsample layer. |
| 67 | down (`bool`, *optional*, default to `False`): If `True`, add a downsample layer. |
| 68 | conv_shortcut_bias (`bool`, *optional*, default to `True`): If `True`, adds a learnable bias to the |
| 69 | `conv_shortcut` output. |
| 70 | conv_2d_out_channels (`int`, *optional*, default to `None`): the number of channels in the output. |
| 71 | If None, same as `out_channels`. |
| 72 | """ |
| 73 | |
| 74 | def __init__( |
| 75 | self, |
| 76 | *, |
| 77 | in_channels: int, |
| 78 | out_channels: Optional[int] = None, |
| 79 | conv_shortcut: bool = False, |
| 80 | dropout: float = 0.0, |
| 81 | temb_channels: int = 512, |
| 82 | groups: int = 32, |
| 83 | groups_out: Optional[int] = None, |
| 84 | eps: float = 1e-6, |
| 85 | non_linearity: str = "swish", |
| 86 | time_embedding_norm: str = "ada_group", # ada_group, spatial |
| 87 | output_scale_factor: float = 1.0, |
| 88 | use_in_shortcut: Optional[bool] = None, |
| 89 | up: bool = False, |
| 90 | down: bool = False, |
| 91 | conv_shortcut_bias: bool = True, |
| 92 | conv_2d_out_channels: Optional[int] = None, |
| 93 | ): |
| 94 | super().__init__() |
| 95 | self.in_channels = in_channels |
| 96 | out_channels = in_channels if out_channels is None else out_channels |
| 97 | self.out_channels = out_channels |
| 98 | self.use_conv_shortcut = conv_shortcut |
| 99 | self.up = up |
| 100 | self.down = down |
| 101 | self.output_scale_factor = output_scale_factor |