Method __init__

stable_diffusion/stable_diffusion/unet.py:273–401

(self, config: UNetConfig)

    """The conditional 2D UNet model that actually performs the denoising."""

    def __init__(self, config: UNetConfig):
        super().__init__()

        self.conv_in = nn.Conv2d(
            config.in_channels,
            config.block_out_channels[0],
            config.conv_in_kernel,
            padding=(config.conv_in_kernel - 1) // 2,
        )

        self.timesteps = nn.SinusoidalPositionalEncoding(
            config.block_out_channels[0],
            max_freq=1,
            min_freq=math.exp(
                -math.log(10000) + 2 * math.log(10000) / config.block_out_channels[0]
            ),
            scale=1.0,
            cos_first=True,
            full_turns=False,
        )
        self.time_embedding = TimestepEmbedding(
            config.block_out_channels[0],
            config.block_out_channels[0] * 4,
        )

        if config.addition_embed_type == "text_time":
            self.add_time_proj = nn.SinusoidalPositionalEncoding(
                config.addition_time_embed_dim,
                max_freq=1,
                min_freq=math.exp(
                    -math.log(10000)
                    + 2 * math.log(10000) / config.addition_time_embed_dim
                ),
                scale=1.0,
                cos_first=True,
                full_turns=False,
            )
            self.add_embedding = TimestepEmbedding(
                config.projection_class_embeddings_input_dim,
                config.block_out_channels[0] * 4,
            )

        # Make the downsampling blocks
        block_channels = [config.block_out_channels[0]] + list(
            config.block_out_channels
        )
        self.down_blocks = [
            UNetBlock2D(
                in_channels=in_channels,
                out_channels=out_channels,
                temb_channels=config.block_out_channels[0] * 4,
                num_layers=config.layers_per_block[i],
                transformer_layers_per_block=config.transformer_layers_per_block[i],
                num_attention_heads=config.num_attention_heads[i],
                cross_attention_dim=config.cross_attention_dim[i],
                resnet_groups=config.norm_num_groups,
                add_downsample=(i < len(config.block_out_channels) - 1),
                add_upsample=False,