Spaces:
Running on Zero
Running on Zero
| # -*- coding: utf-8 -*- | |
| # MIT License | |
| # Copyright (c) 2025 Yiwei Guo | |
| import torch | |
| from typing import Union | |
| from lscodec.utils import crop_seq | |
| from lscodec.models.quantization.core_vq import GroupVectorQuantization | |
| from lscodec.models.fairseq_modules.fp32_group_norm import Fp32GroupNorm | |
| from lscodec.models.feature_extractor import LSCodecEncoder | |
| from lscodec.models.prompt_prenet import ConvPromptPrenet | |
| import logging | |
| class QuantizerProjection(torch.nn.Module): | |
| def __init__(self, dim, groups): | |
| super().__init__() | |
| self.projection = torch.nn.Sequential( | |
| torch.nn.Conv1d(dim, dim, kernel_size=1, groups=groups, bias=False), | |
| Fp32GroupNorm(groups, dim), | |
| ) | |
| def forward(self, x): | |
| return self.projection(x) | |
| class LSCodecEncoderQuantizer(torch.nn.Module): | |
| def __init__(self, frontend: LSCodecEncoder, | |
| quantizer: GroupVectorQuantization, | |
| dropout_features: float, | |
| mean_only: bool=False, | |
| length_mismatch_tolerance: int = 3): | |
| super(LSCodecEncoderQuantizer, self).__init__() | |
| self.frontend = frontend | |
| self.quantizer = quantizer | |
| self.dropout_feats = torch.nn.Dropout(p=dropout_features) | |
| self.length_mismatch_tolerance = length_mismatch_tolerance | |
| self.mean_only = mean_only | |
| def sample_from_posterior(self, means, logvars): | |
| # z = mean + N(0,1)*exp(0.5*logvars) | |
| return means + torch.randn_like(means) * torch.exp(0.5*logvars) | |
| def forward(self, waveform: torch.Tensor, prompt, mask=None, prompt_mask=None, perform_sample=True): | |
| features = self.frontend.forward(waveform, mask, self.length_mismatch_tolerance) # [B, D, L] | |
| B, D, L = features.shape | |
| if mask is not None: | |
| features = features.masked_fill(~mask.unsqueeze(1), 0) | |
| features = features.transpose(1, 2) # [B, L, D] | |
| # get mean and logvars from features | |
| means = features[..., :D//2] | |
| if self.mean_only: | |
| logvars = torch.zeros_like(means) | |
| else: | |
| logvars = features[..., D//2:] | |
| if self.quantizer is not None: | |
| if mask is not None: | |
| features, embed_index, commitment_loss = self.quantizer.forward(means[mask]) # NOTE: quantize on mean only. | |
| # NOTE: the input is already a mask-selected tensor [something, D] | |
| # now features is a [something, D] tensor. We need to convert it back | |
| L = mask.shape[-1] | |
| vqvec = torch.zeros(B, L, D//2).float().to(features.device) | |
| vqvec[mask.unsqueeze(-1).expand(-1, -1, D//2)] = features.view(-1) | |
| else: | |
| features, embed_index, commitment_loss = self.quantizer.forward(means.reshape(B*L, D//2)) | |
| vqvec = features.reshape(B, L, D//2) | |
| vqvec = vqvec.transpose(1, 2) # still [B, D, L] for compatibility. | |
| vqvec = self.dropout_feats(vqvec) | |
| if mask is not None and vqvec.shape[-1] != mask.shape[-1]: | |
| if abs(vqvec.shape[-1] - mask.shape[-1]) <= self.length_mismatch_tolerance: | |
| min_len = min(vqvec.shape[-1], mask.shape[-1]) | |
| vqvec = vqvec[..., :min_len] | |
| mask = mask[..., :min_len] | |
| else: | |
| raise RuntimeError(f"After feature extractor, length mismatch: " | |
| f"extracted {vqvec.shape[-1]} frames and original mel is {mask.shape[-1]} frames") | |
| else: | |
| if perform_sample: | |
| features = self.sample_from_posterior(means, logvars) | |
| else: | |
| features = means | |
| vqvec = features.transpose(1, 2) # NOTE: a little dirty hack. This is not actually "VQ" vector. | |
| commitment_loss = None | |
| return commitment_loss, mask, means, logvars | |
| def encode(self, waveform: torch.Tensor, mask=None): | |
| features = self.frontend.forward(waveform, mask, self.length_mismatch_tolerance) # [B, D, L] | |
| B, D, L = features.shape | |
| features = features.transpose(1, 2) # [B, L, D] | |
| # get mean and logvars from features | |
| means = features[..., :D//2] | |
| logvars = features[..., D//2:] | |
| if self.quantizer is not None: | |
| features, embed_index, commitment_loss = self.quantizer.forward(means.reshape(B*L, D//2)) | |
| else: | |
| embed_index = None | |
| if embed_index is not None and len(embed_index.shape) == 1: | |
| embed_index = embed_index.unsqueeze(-1) # [L, 1] | |
| return means, logvars, embed_index | |