m_diffae/encoder.py

"""mDiffAE encoder: patchify -> DiCoBlocks -> bottleneck projection."""

from __future__ import annotations

from torch import Tensor, nn

from .dico_block import DiCoBlock
from .norms import ChannelWiseRMSNorm
from .straight_through_encoder import Patchify


class Encoder(nn.Module):
    """Deterministic encoder: Image [B,3,H,W] -> latents [B,bottleneck_dim,h,w].

    Pipeline: Patchify -> RMSNorm -> DiCoBlocks (unconditioned) -> Conv1x1 -> RMSNorm(no affine)
    """

    def __init__(
        self,
        in_channels: int,
        patch_size: int,
        model_dim: int,
        depth: int,
        bottleneck_dim: int,
        mlp_ratio: float,
        depthwise_kernel_size: int,
    ) -> None:
        super().__init__()
        self.patchify = Patchify(in_channels, patch_size, model_dim)
        self.norm_in = ChannelWiseRMSNorm(model_dim, eps=1e-6, affine=True)
        self.blocks = nn.ModuleList(
            [
                DiCoBlock(
                    model_dim,
                    mlp_ratio,
                    depthwise_kernel_size=depthwise_kernel_size,
                    use_external_adaln=False,
                )
                for _ in range(depth)
            ]
        )
        self.to_bottleneck = nn.Conv2d(
            model_dim, bottleneck_dim, kernel_size=1, bias=True
        )
        self.norm_out = ChannelWiseRMSNorm(bottleneck_dim, eps=1e-6, affine=False)

    def forward(self, images: Tensor) -> Tensor:
        """Encode images [B,3,H,W] in [-1,1] to latents [B,bottleneck_dim,h,w]."""
        z = self.patchify(images)
        z = self.norm_in(z)
        for block in self.blocks:
            z = block(z)
        z = self.to_bottleneck(z)
        z = self.norm_out(z)
        return z