Residual block with time embedding conditioning. §3.3 — "group normalization throughout... Transformer sinusoidal position embedding into each residual block" [FROM_OFFICIAL_CODE] Structure: GroupNorm -> SiLU -> Conv -> GroupNorm -> SiLU -> Dropout -> Conv + residual
| 88 | # --------------------------------------------------------------------------- |
| 89 | |
| 90 | class ResidualBlock(nn.Module): |
| 91 | """Residual block with time embedding conditioning. |
| 92 | |
| 93 | §3.3 — "group normalization throughout... Transformer sinusoidal position |
| 94 | embedding into each residual block" |
| 95 | |
| 96 | [FROM_OFFICIAL_CODE] Structure: GroupNorm -> SiLU -> Conv -> GroupNorm -> SiLU -> Dropout -> Conv + residual |
| 97 | """ |
| 98 | |
| 99 | def __init__( |
| 100 | self, |
| 101 | in_channels: int, |
| 102 | out_channels: int, |
| 103 | time_embed_dim: int, |
| 104 | dropout: float = 0.0, |
| 105 | num_groups: int = 32, |
| 106 | ): |
| 107 | super().__init__() |
| 108 | self.norm1 = nn.GroupNorm(num_groups, in_channels) |
| 109 | self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1) |
| 110 | |
| 111 | # Time embedding projection |
| 112 | self.time_proj = nn.Linear(time_embed_dim, out_channels) |
| 113 | |
| 114 | self.norm2 = nn.GroupNorm(num_groups, out_channels) |
| 115 | self.dropout = nn.Dropout(dropout) |
| 116 | self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1) |
| 117 | |
| 118 | # Skip connection (1x1 conv if channel count changes) |
| 119 | if in_channels != out_channels: |
| 120 | self.skip = nn.Conv2d(in_channels, out_channels, kernel_size=1) |
| 121 | else: |
| 122 | self.skip = nn.Identity() |
| 123 | |
| 124 | def forward(self, x: torch.Tensor, t_emb: torch.Tensor) -> torch.Tensor: |
| 125 | """ |
| 126 | Args: |
| 127 | x: (batch, in_channels, H, W) |
| 128 | t_emb: (batch, time_embed_dim) |
| 129 | |
| 130 | Returns: |
| 131 | (batch, out_channels, H, W) |
| 132 | """ |
| 133 | h = self.norm1(x) |
| 134 | h = F.silu(h) # (batch, in_channels, H, W) |
| 135 | h = self.conv1(h) # (batch, out_channels, H, W) |
| 136 | |
| 137 | # Add time embedding |
| 138 | t = self.time_proj(F.silu(t_emb)) # (batch, out_channels) |
| 139 | h = h + t.unsqueeze(-1).unsqueeze(-1) # (batch, out_channels, H, W) — broadcast |
| 140 | |
| 141 | h = self.norm2(h) |
| 142 | h = F.silu(h) |
| 143 | h = self.dropout(h) |
| 144 | h = self.conv2(h) # (batch, out_channels, H, W) |
| 145 | |
| 146 | return h + self.skip(x) # residual connection |
| 147 |