| 404 | |
| 405 | |
| 406 | class ConvNextBlock(nn.Module): |
| 407 | def __init__( |
| 408 | self, channels, layer_norm_eps, ln_elementwise_affine, use_bias, hidden_dropout, hidden_size, res_ffn_factor=4 |
| 409 | ): |
| 410 | super().__init__() |
| 411 | self.depthwise = nn.Conv2d( |
| 412 | channels, |
| 413 | channels, |
| 414 | kernel_size=3, |
| 415 | padding=1, |
| 416 | groups=channels, |
| 417 | bias=use_bias, |
| 418 | ) |
| 419 | self.norm = RMSNorm(channels, layer_norm_eps, ln_elementwise_affine) |
| 420 | self.channelwise_linear_1 = nn.Linear(channels, int(channels * res_ffn_factor), bias=use_bias) |
| 421 | self.channelwise_act = nn.GELU() |
| 422 | self.channelwise_norm = GlobalResponseNorm(int(channels * res_ffn_factor)) |
| 423 | self.channelwise_linear_2 = nn.Linear(int(channels * res_ffn_factor), channels, bias=use_bias) |
| 424 | self.channelwise_dropout = nn.Dropout(hidden_dropout) |
| 425 | self.cond_embeds_mapper = nn.Linear(hidden_size, channels * 2, use_bias) |
| 426 | |
| 427 | def forward(self, x, cond_embeds): |
| 428 | x_res = x |
| 429 | |
| 430 | x = self.depthwise(x) |
| 431 | |
| 432 | x = x.permute(0, 2, 3, 1) |
| 433 | x = self.norm(x) |
| 434 | |
| 435 | x = self.channelwise_linear_1(x) |
| 436 | x = self.channelwise_act(x) |
| 437 | x = self.channelwise_norm(x) |
| 438 | x = self.channelwise_linear_2(x) |
| 439 | x = self.channelwise_dropout(x) |
| 440 | |
| 441 | x = x.permute(0, 3, 1, 2) |
| 442 | |
| 443 | x = x + x_res |
| 444 | |
| 445 | scale, shift = self.cond_embeds_mapper(F.silu(cond_embeds)).chunk(2, dim=1) |
| 446 | x = x * (1 + scale[:, :, None, None]) + shift[:, :, None, None] |
| 447 | |
| 448 | return x |
| 449 | |
| 450 | |
| 451 | class ConvMlmLayer(nn.Module): |