| from __future__ import annotations |
|
|
| import torch |
| import torch.nn as nn |
|
|
|
|
| class Decoder4FeatureExtractor(nn.Module): |
| """MOSS audio-tokenizer codebook decode plus decoder blocks 0..4.""" |
|
|
| def __init__( |
| self, |
| audio_tokenizer: nn.Module, |
| num_quantizers: int = 32, |
| output_dtype: torch.dtype = torch.float16, |
| ) -> None: |
| super().__init__() |
| quantizer = getattr(audio_tokenizer, "quantizer") |
| self.quantizers = quantizer.quantizers |
| self.output_proj = quantizer.output_proj |
| self.decoder_prefix = nn.ModuleList(list(audio_tokenizer.decoder[:5])) |
| self.rvq_dim = int(quantizer.rvq_dim) |
| self.num_quantizers = int(num_quantizers) |
| self.output_dtype = output_dtype |
|
|
| def forward(self, codes: torch.Tensor, lengths: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: |
| _, batch, frames = codes.shape |
| emb = torch.zeros(batch, self.rvq_dim, frames, device=codes.device, dtype=self.output_dtype) |
| for index, quantizer in enumerate(self.quantizers[: self.num_quantizers]): |
| if self.output_dtype == torch.float16: |
| z_q = quantizer.embed_code(codes[index]).transpose(1, 2).to(self.output_dtype) |
| z_q = quantizer.out_proj(z_q) |
| else: |
| z_q = quantizer.decode_code(codes[index]) |
| emb = emb + z_q.to(self.output_dtype) |
| features = self.output_proj(emb) |
| feature_lengths = lengths |
| for module in self.decoder_prefix: |
| features, feature_lengths = module(features, feature_lengths) |
| return features.to(self.output_dtype), feature_lengths |
|
|
