| 363 | |
| 364 | class CondInjection(nn.Module): |
| 365 | def __init__(self, fea_dim, cond_dim, hidden_dim, groups=32) -> None: |
| 366 | super().__init__() |
| 367 | self.body = nn.Sequential( |
| 368 | nn.Conv2d(cond_dim, hidden_dim * 4, 3, padding=1, bias=False), |
| 369 | nn.GroupNorm(groups, hidden_dim * 4), |
| 370 | nn.SiLU(), |
| 371 | nn.Conv2d(hidden_dim * 4, hidden_dim * 2, 1, bias=True), |
| 372 | ) |
| 373 | self.x_conv = nn.Conv2d(fea_dim, hidden_dim, 1, bias=True) |
| 374 | nn.init.zeros_(self.body[-1].weight) |
| 375 | nn.init.zeros_(self.body[-1].bias) |
| 376 | |
| 377 | def forward(self, x, cond): |
| 378 | cond = self.body(cond) |