Spaces:
Running on Zero
Running on Zero
File size: 2,114 Bytes
f45c5e8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 | """NeMo audio-codec wrapper used at inference time.
``UnfoldedCodecModel`` extends NeMo's ``AudioCodecModel`` with direct decoding
from per-dimension discrete FSQ codes (the format the Gepard model produces),
bypassing mixed-radix composition/decomposition.
"""
import torch
from omegaconf import open_dict
from nemo.collections.tts.models import AudioCodecModel
class UnfoldedCodecModel(AudioCodecModel):
"""AudioCodecModel + decoding from unfolded per-dimension FSQ codes.
Works with any GroupFiniteScalarQuantizer configuration — the number of
groups, dimensions per group, and FSQ levels are read from the model's
vector_quantizer at runtime.
"""
def __init__(self, cfg, trainer=None):
# SLMDiscriminator downloads microsoft/wavlm-base-plus (~360MB) and is
# only used during training — strip it from the config before init.
with open_dict(cfg):
disc = cfg.get("discriminator", None)
if disc is not None and "discriminators" in disc:
disc.discriminators = [
d for d in disc.discriminators if "SLM" not in d._target_
]
super().__init__(cfg, trainer)
def decode_from_codes(self, codes: torch.Tensor, codes_len: torch.Tensor):
"""Decode audio from unfolded per-dimension discrete codes.
Args:
codes: (B, D, T) — per-dimension discrete values, where
D = num_groups * dims_per_group.
codes_len: (B,) — valid frame count per batch element.
Returns:
audio: (B, T_audio) — decoded waveform
audio_len: (B,) — valid audio lengths in samples
"""
num_levels = self.vector_quantizer.fsqs[0].num_levels.squeeze()
scale = (num_levels // 2).float().to(codes.device)
groups = codes.chunk(self.vector_quantizer.num_groups, dim=1)
dequantized = torch.cat(
[(g - scale[None, :, None]) / scale[None, :, None] for g in groups],
dim=1,
)
return self.decode_audio(inputs=dequantized, input_len=codes_len)
|