(self, config: UNetConfig)
| 234 | """ |
| 235 | |
| 236 | def __init__(self, config: UNetConfig): |
| 237 | super().__init__() |
| 238 | self.config = config |
| 239 | ch = config.base_channels |
| 240 | |
| 241 | # Time embedding: sinusoidal -> MLP |
| 242 | # §3.3 — "Transformer sinusoidal position embedding" |
| 243 | time_embed_dim = config.time_embed_dim |
| 244 | self.time_embed = nn.Sequential( |
| 245 | SinusoidalTimeEmbedding(ch), |
| 246 | nn.Linear(ch, time_embed_dim), |
| 247 | nn.SiLU(), |
| 248 | nn.Linear(time_embed_dim, time_embed_dim), |
| 249 | ) |
| 250 | |
| 251 | # Initial convolution |
| 252 | self.input_conv = nn.Conv2d(config.image_channels, ch, kernel_size=3, padding=1) |
| 253 | |
| 254 | # Downsampling path |
| 255 | self.down_blocks = nn.ModuleList() |
| 256 | self.down_samples = nn.ModuleList() |
| 257 | channels = [ch] |
| 258 | current_res = config.image_size |
| 259 | in_ch = ch |
| 260 | |
| 261 | for level, mult in enumerate(config.channel_mults): |
| 262 | out_ch = ch * mult |
| 263 | for _ in range(config.num_res_blocks): |
| 264 | layers = [ResidualBlock(in_ch, out_ch, time_embed_dim, |
| 265 | config.dropout, config.num_groups)] |
| 266 | if current_res in config.attention_resolutions: |
| 267 | layers.append(AttentionBlock(out_ch, config.num_groups)) |
| 268 | self.down_blocks.append(nn.ModuleList(layers)) |
| 269 | channels.append(out_ch) |
| 270 | in_ch = out_ch |
| 271 | |
| 272 | if level < len(config.channel_mults) - 1: |
| 273 | self.down_samples.append(Downsample(out_ch)) |
| 274 | channels.append(out_ch) |
| 275 | current_res //= 2 |
| 276 | else: |
| 277 | self.down_samples.append(nn.Identity()) |
| 278 | |
| 279 | # Middle (bottleneck) |
| 280 | self.mid_block1 = ResidualBlock(in_ch, in_ch, time_embed_dim, |
| 281 | config.dropout, config.num_groups) |
| 282 | self.mid_attn = AttentionBlock(in_ch, config.num_groups) |
| 283 | self.mid_block2 = ResidualBlock(in_ch, in_ch, time_embed_dim, |
| 284 | config.dropout, config.num_groups) |
| 285 | |
| 286 | # Upsampling path |
| 287 | self.up_blocks = nn.ModuleList() |
| 288 | self.up_samples = nn.ModuleList() |
| 289 | |
| 290 | for level in reversed(range(len(config.channel_mults))): |
| 291 | mult = config.channel_mults[level] |
| 292 | out_ch = ch * mult |
| 293 | for i in range(config.num_res_blocks + 1): |
no test coverage detected