src/models/diffusion/unet.py

from __future__ import annotations

import torch
from torch import nn
import torch.nn.functional as F

from src.models.diffusion.timestep import TimestepEmbedding
from src.models.diffusion.blocks import (
    DownBlock,
    MiddleBlock,
    UpBlock,
    normalization,
)


class LatentDiffusionUNet(nn.Module):
    """
    Lightweight text-conditioned U-Net for latent diffusion.

    Input:
        z_t:
            noisy latent [B, in_channels, H, W]

        timesteps:
            diffusion timestep [B]

        context:
            CLIP token embeddings [B, seq_len, context_dim]

    Output:
        prediction [B, out_channels, H, W]
    """

    def __init__(
        self,
        in_channels: int = 8,
        out_channels: int = 8,
        latent_size: int = 32,
        base_channels: int = 128,
        channel_multipliers: list[int] | tuple[int, ...] = (1, 2, 3),
        num_res_blocks: int = 2,
        attention_resolutions: list[int] | tuple[int, ...] = (16, 8),
        context_dim: int = 768,
        num_heads: int = 4,
        head_dim: int = 32,
        transformer_depth: int = 1,
        dropout: float = 0.0,
        time_embedding_dim: int | None = None,
        use_middle_attention: bool = True,
    ):
        super().__init__()

        self.in_channels = in_channels
        self.out_channels = out_channels
        self.latent_size = latent_size
        self.base_channels = base_channels
        self.channel_multipliers = list(channel_multipliers)
        self.num_res_blocks = num_res_blocks
        self.attention_resolutions = set(attention_resolutions)
        self.context_dim = context_dim
        self.num_heads = num_heads
        self.head_dim = head_dim
        self.transformer_depth = transformer_depth
        self.dropout = dropout

        if time_embedding_dim is None:
            time_embedding_dim = base_channels * 4

        self.time_embedding_dim = time_embedding_dim

        self.time_embed = TimestepEmbedding(
            embedding_dim=base_channels,
            time_embed_dim=time_embedding_dim,
        )

        self.conv_in = nn.Conv2d(
            in_channels,
            base_channels,
            kernel_size=3,
            stride=1,
            padding=1,
        )

        channels_per_level = [
            base_channels * multiplier
            for multiplier in self.channel_multipliers
        ]

        # -------------------------
        # Down path
        # -------------------------
        self.down_blocks = nn.ModuleList()

        current_channels = base_channels
        current_resolution = latent_size

        self.down_block_channels: list[int] = []
        self.down_block_resolutions: list[int] = []

        for level, out_channels_level in enumerate(channels_per_level):
            is_last = level == len(channels_per_level) - 1

            use_attention = current_resolution in self.attention_resolutions

            block = DownBlock(
                in_channels=current_channels,
                out_channels=out_channels_level,
                time_embed_dim=time_embedding_dim,
                num_res_blocks=num_res_blocks,
                context_dim=context_dim,
                num_heads=num_heads,
                head_dim=head_dim,
                transformer_depth=transformer_depth,
                dropout=dropout,
                use_attention=use_attention,
                add_downsample=not is_last,
            )

            self.down_blocks.append(block)

            self.down_block_channels.append(out_channels_level)
            self.down_block_resolutions.append(current_resolution)

            current_channels = out_channels_level

            if not is_last:
                current_resolution //= 2

        # -------------------------
        # Middle
        # -------------------------
        self.middle = MiddleBlock(
            channels=current_channels,
            time_embed_dim=time_embedding_dim,
            context_dim=context_dim,
            num_heads=num_heads,
            head_dim=head_dim,
            transformer_depth=transformer_depth,
            dropout=dropout,
            use_attention=use_middle_attention,
        )

        # -------------------------
        # Up path
        # -------------------------
        self.up_blocks = nn.ModuleList()

        reversed_channels = list(reversed(self.down_block_channels))
        reversed_resolutions = list(reversed(self.down_block_resolutions))

        for level, (skip_channels, resolution) in enumerate(
            zip(reversed_channels, reversed_resolutions)
        ):
            is_last = level == len(reversed_channels) - 1

            out_channels_level = skip_channels
            use_attention = resolution in self.attention_resolutions

            block = UpBlock(
                in_channels=current_channels,
                skip_channels=skip_channels,
                out_channels=out_channels_level,
                time_embed_dim=time_embedding_dim,
                num_res_blocks=num_res_blocks,
                context_dim=context_dim,
                num_heads=num_heads,
                head_dim=head_dim,
                transformer_depth=transformer_depth,
                dropout=dropout,
                use_attention=use_attention,
                add_upsample=not is_last,
            )

            self.up_blocks.append(block)

            current_channels = out_channels_level

        self.norm_out = normalization(current_channels)

        self.conv_out = nn.Conv2d(
            current_channels,
            out_channels,
            kernel_size=3,
            stride=1,
            padding=1,
        )

        # Stable start: initially predicts near zero.
        nn.init.zeros_(self.conv_out.weight)
        nn.init.zeros_(self.conv_out.bias)

    def forward(
        self,
        z_t: torch.Tensor,
        timesteps: torch.Tensor,
        context: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """
        Args:
            z_t:
                Noisy latent [B, C, H, W]

            timesteps:
                Diffusion timesteps [B]

            context:
                CLIP token embeddings [B, seq_len, context_dim]

            attention_mask:
                CLIP token mask [B, seq_len], optional

        Returns:
            model_output:
                [B, out_channels, H, W]
        """
        if z_t.ndim != 4:
            raise ValueError(f"z_t must be [B, C, H, W], got {z_t.shape}")

        if timesteps.ndim != 1:
            raise ValueError(f"timesteps must be [B], got {timesteps.shape}")

        time_emb = self.time_embed(timesteps)

        x = self.conv_in(z_t)

        skips: list[torch.Tensor] = []

        for block in self.down_blocks:
            x, block_skips = block(
                x=x,
                time_emb=time_emb,
                context=context,
                attention_mask=attention_mask,
            )

            skips.extend(block_skips)

        x = self.middle(
            x=x,
            time_emb=time_emb,
            context=context,
            attention_mask=attention_mask,
        )

        for block in self.up_blocks:
            x = block(
                x=x,
                skips=skips,
                time_emb=time_emb,
                context=context,
                attention_mask=attention_mask,
            )

        if len(skips) != 0:
            raise RuntimeError(
                f"Unused skip connections remain: {len(skips)}. "
                "Check U-Net block construction."
            )

        x = self.norm_out(x)
        x = F.silu(x)
        x = self.conv_out(x)

        return x


def count_parameters(
    model: nn.Module,
    trainable_only: bool = True,
) -> int:
    if trainable_only:
        return sum(p.numel() for p in model.parameters() if p.requires_grad)

    return sum(p.numel() for p in model.parameters())


def build_latent_diffusion_unet_from_config(cfg: dict) -> LatentDiffusionUNet:
    """
    Build U-Net from config dictionary.

    Expects:

        cfg["model"]
    """
    m = cfg["model"]

    return LatentDiffusionUNet(
        in_channels=int(m["in_channels"]),
        out_channels=int(m["out_channels"]),
        latent_size=int(m.get("latent_size", 32)),
        base_channels=int(m["base_channels"]),
        channel_multipliers=tuple(m["channel_multipliers"]),
        num_res_blocks=int(m["num_res_blocks"]),
        attention_resolutions=tuple(m.get("attention_resolutions", [16, 8])),
        context_dim=int(m["context_dim"]),
        num_heads=int(m["num_heads"]),
        head_dim=int(m["head_dim"]),
        transformer_depth=int(m.get("transformer_depth", 1)),
        dropout=float(m.get("dropout", 0.0)),
        time_embedding_dim=m.get("time_embedding_dim", None),
        use_middle_attention=bool(m.get("use_middle_attention", True)),
    )