(
self, channels, layer_norm_eps, ln_elementwise_affine, use_bias, hidden_dropout, hidden_size, res_ffn_factor=4
)
| 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 |
nothing calls this directly
no test coverage detected