(self)
| 1009 | self.adaln_modulation_mlp = nn.Sequential(nn.SiLU(), nn.Linear(x_dim, 3 * x_dim, bias=False)) |
| 1010 | |
| 1011 | def reset_parameters(self) -> None: |
| 1012 | self.layer_norm_self_attn.reset_parameters() |
| 1013 | self.layer_norm_cross_attn.reset_parameters() |
| 1014 | self.layer_norm_mlp.reset_parameters() |
| 1015 | |
| 1016 | if self.use_adaln_lora: |
| 1017 | std = 1.0 / math.sqrt(self.x_dim) |
| 1018 | torch.nn.init.trunc_normal_(self.adaln_modulation_self_attn[1].weight, std=std, a=-3 * std, b=3 * std) |
| 1019 | torch.nn.init.trunc_normal_(self.adaln_modulation_cross_attn[1].weight, std=std, a=-3 * std, b=3 * std) |
| 1020 | torch.nn.init.trunc_normal_(self.adaln_modulation_mlp[1].weight, std=std, a=-3 * std, b=3 * std) |
| 1021 | torch.nn.init.zeros_(self.adaln_modulation_self_attn[2].weight) |
| 1022 | torch.nn.init.zeros_(self.adaln_modulation_cross_attn[2].weight) |
| 1023 | torch.nn.init.zeros_(self.adaln_modulation_mlp[2].weight) |
| 1024 | else: |
| 1025 | torch.nn.init.zeros_(self.adaln_modulation_self_attn[1].weight) |
| 1026 | torch.nn.init.zeros_(self.adaln_modulation_cross_attn[1].weight) |
| 1027 | torch.nn.init.zeros_(self.adaln_modulation_mlp[1].weight) |
| 1028 | |
| 1029 | def init_weights(self) -> None: |
| 1030 | self.reset_parameters() |
no outgoing calls
no test coverage detected