|
|
|
|
|
|
|
|
|
|
| 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
|
|
|
