| 548 | |
| 549 | # UpsampleConv(50->100Hz) + VocosBackbone + ISTFTHead |
| 550 | class AcousticDecoder(nn.Module): |
| 551 | def __init__( |
| 552 | self, |
| 553 | # Transformer |
| 554 | embed_dim: int, |
| 555 | num_layers: int, |
| 556 | num_heads: int, |
| 557 | dropout: float = 0.0, |
| 558 | # iSTFT |
| 559 | hop_length: int = 240, |
| 560 | # Causal |
| 561 | causal: bool = False, |
| 562 | ): |
| 563 | super().__init__() |
| 564 | self.embed_dim = embed_dim |
| 565 | self.num_layers = num_layers |
| 566 | self.num_heads = num_heads |
| 567 | self.hop_length = hop_length |
| 568 | self.causal = causal |
| 569 | |
| 570 | # Output upsample |
| 571 | self.upsample_conv = nn.Sequential( |
| 572 | nn.ConvTranspose1d( |
| 573 | embed_dim, |
| 574 | embed_dim, |
| 575 | kernel_size=3, |
| 576 | stride=2, |
| 577 | padding=0, # Do not fill input side |
| 578 | output_padding=0, # Can be adjusted to precisely control length |
| 579 | ), |
| 580 | nn.GELU(), |
| 581 | nn.ConvTranspose1d( |
| 582 | embed_dim, |
| 583 | embed_dim, |
| 584 | kernel_size=3, |
| 585 | stride=1, |
| 586 | padding=0, # Do not fill input side |
| 587 | ), |
| 588 | nn.GELU(), |
| 589 | ) |
| 590 | self.backbone = ( |
| 591 | CausalVocosBackbone(embed_dim, num_layers, num_heads, dropout) |
| 592 | if causal |
| 593 | else VocosBackbone(embed_dim, num_layers, num_heads, dropout) |
| 594 | ) |
| 595 | self.isift = ISTFTHead(embed_dim, hop_length * 4, hop_length, padding="same") |
| 596 | # Init weights |
| 597 | self.apply(self._init_weights) |
| 598 | |
| 599 | def _init_weights(self, m): |
| 600 | if isinstance(m, nn.Conv1d): |
| 601 | nn.init.trunc_normal_(m.weight, std=0.02) |
| 602 | nn.init.constant_(m.bias, 0) |
| 603 | |
| 604 | def forward(self, x: torch.Tensor, x_lens: torch.Tensor): |
| 605 | """ |
| 606 | Args: |
| 607 | x: shape (b, t, c) |