Class UNet

skills/paper2code/worked/ddpm/src/model.py:219–382

§3.3, Appendix B — U-Net noise prediction network ε_θ(x_t, t). "We use a U-Net backbone similar to an unmasked PixelCNN++ with group normalization throughout, and we add one head of self-attention at the 16×16 feature map resolution." The U-Net takes a noisy image x_t and a times

# ---------------------------------------------------------------------------

class UNet(nn.Module):
    """§3.3, Appendix B — U-Net noise prediction network ε_θ(x_t, t).

    "We use a U-Net backbone similar to an unmasked PixelCNN++ with
     group normalization throughout, and we add one head of self-attention
     at the 16×16 feature map resolution."

    The U-Net takes a noisy image x_t and a timestep t, and predicts the
    noise ε that was added. This is NOT the paper's core contribution —
    it is the backbone model that enables the diffusion process.

    Architecture (for CIFAR-10 32×32):
      Down: 32→32→16→8→4 (with skip connections)
      Middle: bottleneck with attention
      Up: 4→8→16→32→32 (with skip connections from down path)
    """

    def __init__(self, config: UNetConfig):
        super().__init__()
        self.config = config
        ch = config.base_channels

        # Time embedding: sinusoidal -> MLP
        # §3.3 — "Transformer sinusoidal position embedding"
        time_embed_dim = config.time_embed_dim
        self.time_embed = nn.Sequential(
            SinusoidalTimeEmbedding(ch),
            nn.Linear(ch, time_embed_dim),
            nn.SiLU(),
            nn.Linear(time_embed_dim, time_embed_dim),
        )

        # Initial convolution
        self.input_conv = nn.Conv2d(config.image_channels, ch, kernel_size=3, padding=1)

        # Downsampling path
        self.down_blocks = nn.ModuleList()
        self.down_samples = nn.ModuleList()
        channels = [ch]
        current_res = config.image_size
        in_ch = ch

        for level, mult in enumerate(config.channel_mults):
            out_ch = ch * mult
            for _ in range(config.num_res_blocks):
                layers = [ResidualBlock(in_ch, out_ch, time_embed_dim,
                                        config.dropout, config.num_groups)]
                if current_res in config.attention_resolutions:
                    layers.append(AttentionBlock(out_ch, config.num_groups))
                self.down_blocks.append(nn.ModuleList(layers))
                channels.append(out_ch)
                in_ch = out_ch

            if level < len(config.channel_mults) - 1:
                self.down_samples.append(Downsample(out_ch))
                channels.append(out_ch)
                current_res //= 2
            else: