(
self, channels, layer_norm_eps, ln_elementwise_affine, use_bias, hidden_dropout, hidden_size, res_ffn_factor=4
)
| 342 | |
| 343 | class ConvNextBlock(nn.Module): |
| 344 | def __init__( |
| 345 | self, channels, layer_norm_eps, ln_elementwise_affine, use_bias, hidden_dropout, hidden_size, res_ffn_factor=4 |
| 346 | ): |
| 347 | super().__init__() |
| 348 | self.depthwise = nn.Conv2d( |
| 349 | channels, |
| 350 | channels, |
| 351 | kernel_size=3, |
| 352 | padding=1, |
| 353 | groups=channels, |
| 354 | bias=use_bias, |
| 355 | ) |
| 356 | self.norm = RMSNorm(channels, layer_norm_eps, ln_elementwise_affine) |
| 357 | self.channelwise_linear_1 = nn.Linear(channels, int(channels * res_ffn_factor), bias=use_bias) |
| 358 | self.channelwise_act = nn.GELU() |
| 359 | self.channelwise_norm = GlobalResponseNorm(int(channels * res_ffn_factor)) |
| 360 | self.channelwise_linear_2 = nn.Linear(int(channels * res_ffn_factor), channels, bias=use_bias) |
| 361 | self.channelwise_dropout = nn.Dropout(hidden_dropout) |
| 362 | self.cond_embeds_mapper = nn.Linear(hidden_size, channels * 2, use_bias) |
| 363 | |
| 364 | def forward(self, x, cond_embeds): |
| 365 | x_res = x |
nothing calls this directly
no test coverage detected