Method __init__

models/light_decoder.py:34–110

(
        self,
        in_channels: int = 4,
        out_channels: int = 3,
        block_out_channels=(64, 128, 256, 256),
        layers_per_block: int = 2,
        norm_num_groups: int = 32,
        act_fn: str = "silu",
        mid_block_add_attention: bool = True,
    )

    """

    def __init__(
        self,
        in_channels: int = 4,
        out_channels: int = 3,
        block_out_channels=(64, 128, 256, 256),
        layers_per_block: int = 2,
        norm_num_groups: int = 32,
        act_fn: str = "silu",
        mid_block_add_attention: bool = True,
    ):
        super().__init__()

        self.in_channels = in_channels
        self.out_channels = out_channels
        self.block_out_channels = block_out_channels
        self.layers_per_block = layers_per_block

        # Reversed block_out_channels for decoder (goes from deepest to shallowest)
        reversed_block_out_channels = list(reversed(block_out_channels))

        # conv_in: latent_channels -> deepest channel dim
        self.conv_in = nn.Conv2d(
            in_channels,
            reversed_block_out_channels[0],
            kernel_size=3,
            stride=1,
            padding=1,
        )

        # Mid block with attention
        # attention_head_dim should match in_channels (single-head attention),
        # consistent with diffusers Decoder which uses attention_head_dim=block_out_channels[-1]
        self.mid_block = UNetMidBlock2D(
            in_channels=reversed_block_out_channels[0],
            temb_channels=None,
            dropout=0.0,
            num_layers=1,
            resnet_eps=1e-6,
            resnet_act_fn=act_fn,
            resnet_groups=norm_num_groups,
            add_attention=mid_block_add_attention,
            attention_head_dim=reversed_block_out_channels[0],
        )

        # Up blocks
        self.up_blocks = nn.ModuleList([])
        output_channel = reversed_block_out_channels[0]
        for i, up_block_out_channel in enumerate(reversed_block_out_channels):
            prev_output_channel = output_channel
            output_channel = up_block_out_channel
            is_final_block = i == len(reversed_block_out_channels) - 1

            up_block = UpDecoderBlock2D(
                in_channels=prev_output_channel,
                out_channels=output_channel,
                num_layers=layers_per_block + 1,  # diffusers convention: layers_per_block + 1 resnets
                resnet_eps=1e-6,
                resnet_act_fn=act_fn,