Method forward

skills/paper2code/worked/ddpm/src/model.py:316–370

Predict noise ε_θ(x_t, t). Args: x: (batch, C, H, W) — noisy image x_t t: (batch,) — integer timesteps Returns: (batch, C, H, W) — predicted noise ε_θ

(self, x: torch.Tensor, t: torch.Tensor)

        nn.init.zeros_(self.output_conv.bias)

    def forward(self, x: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
        """Predict noise ε_θ(x_t, t).

        Args:
            x: (batch, C, H, W) — noisy image x_t
            t: (batch,) — integer timesteps

        Returns:
            (batch, C, H, W) — predicted noise ε_θ
        """
        # Time embedding
        t_emb = self.time_embed(t)  # (batch, time_embed_dim)

        # Initial conv
        h = self.input_conv(x)  # (batch, base_channels, H, W)

        # Downsampling with skip connections
        skips = [h]
        block_idx = 0
        for level in range(len(self.config.channel_mults)):
            for _ in range(self.config.num_res_blocks):
                layers = self.down_blocks[block_idx]
                h = layers[0](h, t_emb)  # ResidualBlock
                if len(layers) > 1:
                    h = layers[1](h)      # AttentionBlock (if present)
                skips.append(h)
                block_idx += 1

            h = self.down_samples[level](h)
            if not isinstance(self.down_samples[level], nn.Identity):
                skips.append(h)

        # Middle
        h = self.mid_block1(h, t_emb)
        h = self.mid_attn(h)
        h = self.mid_block2(h, t_emb)

        # Upsampling with skip connections
        block_idx = 0
        for level in reversed(range(len(self.config.channel_mults))):
            for _ in range(self.config.num_res_blocks + 1):
                skip = skips.pop()
                h = torch.cat([h, skip], dim=1)  # Concatenate skip connection
                layers = self.up_blocks[block_idx]
                h = layers[0](h, t_emb)  # ResidualBlock
                if len(layers) > 1:
                    h = layers[1](h)      # AttentionBlock (if present)
                block_idx += 1

            h = self.up_samples[level - len(self.config.channel_mults)](h) if level > 0 else h

        # Output
        h = self.output_norm(h)
        h = F.silu(h)
        return self.output_conv(h)  # (batch, C, H, W) — predicted noise

    def __repr__(self) -> str:
        total_params = sum(p.numel() for p in self.parameters())