| | |
| | |
| | |
| | |
| |
|
| | from .modules.seanet import SEANetEncoder, SEANetDecoder |
| | from .modules.quantization import ResidualVectorQuantizer |
| | import torch.nn as nn |
| | from einops import rearrange |
| | import torch |
| | import numpy as np |
| |
|
| |
|
| | class SpeechTokenizer(nn.Module): |
| | def __init__(self, config): |
| | """ |
| | |
| | Parameters |
| | ---------- |
| | config : json |
| | Model Config. |
| | |
| | """ |
| | super().__init__() |
| | self.encoder = SEANetEncoder( |
| | n_filters=config.get("n_filters"), |
| | dimension=config.get("dimension"), |
| | ratios=config.get("strides"), |
| | lstm=config.get("lstm_layers"), |
| | bidirectional=config.get("bidirectional"), |
| | dilation_base=config.get("dilation_base"), |
| | residual_kernel_size=config.get("residual_kernel_size"), |
| | n_residual_layers=config.get("n_residual_layers"), |
| | activation=config.get("activation"), |
| | ) |
| | self.sample_rate = config.get("sample_rate") |
| | self.n_q = config.get("n_q") |
| | self.downsample_rate = np.prod(config.get("strides")) |
| | if config.get("dimension") != config.get("semantic_dimension"): |
| | self.transform = nn.Linear( |
| | config.get("dimension"), config.get("semantic_dimension") |
| | ) |
| | else: |
| | self.transform = nn.Identity() |
| | self.quantizer = ResidualVectorQuantizer( |
| | dimension=config.get("dimension"), |
| | n_q=config.get("n_q"), |
| | bins=config.get("codebook_size"), |
| | ) |
| | self.decoder = SEANetDecoder( |
| | n_filters=config.get("n_filters"), |
| | dimension=config.get("dimension"), |
| | ratios=config.get("strides"), |
| | lstm=config.get("lstm_layers"), |
| | bidirectional=False, |
| | dilation_base=config.get("dilation_base"), |
| | residual_kernel_size=config.get("residual_kernel_size"), |
| | n_residual_layers=config.get("n_residual_layers"), |
| | activation=config.get("activation"), |
| | ) |
| |
|
| | @classmethod |
| | def load_from_checkpoint(cls, config_path: str, ckpt_path: str): |
| | """ |
| | |
| | Parameters |
| | ---------- |
| | config_path : str |
| | Path of model configuration file. |
| | ckpt_path : str |
| | Path of model checkpoint. |
| | |
| | Returns |
| | ------- |
| | model : SpeechTokenizer |
| | SpeechTokenizer model. |
| | |
| | """ |
| | import json |
| |
|
| | with open(config_path) as f: |
| | cfg = json.load(f) |
| | model = cls(cfg) |
| | params = torch.load(ckpt_path, map_location="cpu") |
| | model.load_state_dict(params) |
| | return model |
| |
|
| | def forward(self, x: torch.tensor, n_q: int = None, layers: list = [0]): |
| | """ |
| | |
| | Parameters |
| | ---------- |
| | x : torch.tensor |
| | Input wavs. Shape: (batch, channels, timesteps). |
| | n_q : int, optional |
| | Number of quantizers in RVQ used to encode. The default is all layers. |
| | layers : list[int], optional |
| | Layers of RVQ should return quantized result. The default is the first layer. |
| | |
| | Returns |
| | ------- |
| | o : torch.tensor |
| | Output wavs. Shape: (batch, channels, timesteps). |
| | commit_loss : torch.tensor |
| | Commitment loss from residual vector quantizers. |
| | feature : torch.tensor |
| | Output of RVQ's first layer. Shape: (batch, timesteps, dimension) |
| | |
| | """ |
| | n_q = n_q if n_q else self.n_q |
| | e = self.encoder(x) |
| | quantized, codes, commit_loss, quantized_list = self.quantizer( |
| | e, n_q=n_q, layers=layers |
| | ) |
| | feature = rearrange(quantized_list[0], "b d t -> b t d") |
| | feature = self.transform(feature) |
| | o = self.decoder(quantized) |
| | return o, commit_loss, feature |
| |
|
| | def forward_feature(self, x: torch.tensor, layers: list = None): |
| | """ |
| | |
| | Parameters |
| | ---------- |
| | x : torch.tensor |
| | Input wavs. Shape should be (batch, channels, timesteps). |
| | layers : list[int], optional |
| | Layers of RVQ should return quantized result. The default is all layers. |
| | |
| | Returns |
| | ------- |
| | quantized_list : list[torch.tensor] |
| | Quantized of required layers. |
| | |
| | """ |
| | e = self.encoder(x) |
| | layers = layers if layers else list(range(self.n_q)) |
| | quantized, codes, commit_loss, quantized_list = self.quantizer(e, layers=layers) |
| | return quantized_list |
| |
|
| | def encode(self, x: torch.tensor, n_q: int = None, st: int = None): |
| | """ |
| | |
| | Parameters |
| | ---------- |
| | x : torch.tensor |
| | Input wavs. Shape: (batch, channels, timesteps). |
| | n_q : int, optional |
| | Number of quantizers in RVQ used to encode. The default is all layers. |
| | st : int, optional |
| | Start quantizer index in RVQ. The default is 0. |
| | |
| | Returns |
| | ------- |
| | codes : torch.tensor |
| | Output indices for each quantizer. Shape: (n_q, batch, timesteps) |
| | |
| | """ |
| | e = self.encoder(x) |
| | if st is None: |
| | st = 0 |
| | n_q = n_q if n_q else self.n_q |
| | codes = self.quantizer.encode(e, n_q=n_q, st=st) |
| | return codes |
| |
|
| | def decode(self, codes: torch.tensor, st: int = 0): |
| | """ |
| | |
| | Parameters |
| | ---------- |
| | codes : torch.tensor |
| | Indices for each quantizer. Shape: (n_q, batch, timesteps). |
| | st : int, optional |
| | Start quantizer index in RVQ. The default is 0. |
| | |
| | Returns |
| | ------- |
| | o : torch.tensor |
| | Reconstruct wavs from codes. Shape: (batch, channels, timesteps) |
| | |
| | """ |
| | quantized = self.quantizer.decode(codes, st=st) |
| | o = self.decoder(quantized) |
| | return o |
| |
|
