| 449 | |
| 450 | |
| 451 | class ConvMlmLayer(nn.Module): |
| 452 | def __init__( |
| 453 | self, |
| 454 | block_out_channels: int, |
| 455 | in_channels: int, |
| 456 | use_bias: bool, |
| 457 | ln_elementwise_affine: bool, |
| 458 | layer_norm_eps: float, |
| 459 | codebook_size: int, |
| 460 | ): |
| 461 | super().__init__() |
| 462 | self.conv1 = nn.Conv2d(block_out_channels, in_channels, kernel_size=1, bias=use_bias) |
| 463 | self.layer_norm = RMSNorm(in_channels, layer_norm_eps, ln_elementwise_affine) |
| 464 | self.conv2 = nn.Conv2d(in_channels, codebook_size, kernel_size=1, bias=use_bias) |
| 465 | |
| 466 | def forward(self, hidden_states): |
| 467 | hidden_states = self.conv1(hidden_states) |
| 468 | hidden_states = self.layer_norm(hidden_states.permute(0, 2, 3, 1)).permute(0, 3, 1, 2) |
| 469 | logits = self.conv2(hidden_states) |
| 470 | return logits |