| 192 | self.initialize_weights() |
| 193 | |
| 194 | def initialize_weights(self): |
| 195 | def _basic_init(module): |
| 196 | if isinstance(module, nn.Linear): |
| 197 | torch.nn.init.xavier_uniform_(module.weight) |
| 198 | if module.bias is not None: |
| 199 | nn.init.constant_(module.bias, 0) |
| 200 | self.apply(_basic_init) |
| 201 | |
| 202 | # Initialize timestep embedding MLP |
| 203 | nn.init.normal_(self.time_embed.mlp[0].weight, std=0.02) |
| 204 | nn.init.normal_(self.time_embed.mlp[2].weight, std=0.02) |
| 205 | |
| 206 | # Zero-out adaLN modulation layers |
| 207 | for block in self.res_blocks: |
| 208 | nn.init.constant_(block.adaLN_modulation[-1].weight, 0) |
| 209 | nn.init.constant_(block.adaLN_modulation[-1].bias, 0) |
| 210 | |
| 211 | # Zero-out output layers |
| 212 | nn.init.constant_(self.final_layer.adaLN_modulation[-1].weight, 0) |
| 213 | nn.init.constant_(self.final_layer.adaLN_modulation[-1].bias, 0) |
| 214 | nn.init.constant_(self.final_layer.linear.weight, 0) |
| 215 | nn.init.constant_(self.final_layer.linear.bias, 0) |
| 216 | |
| 217 | def forward(self, x, t, c): |
| 218 | """ |