#!/usr/bin/env python3 # coding=utf-8 # Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Inference-only Enhancement VAE model/API for the Audex release.""" from __future__ import annotations import json import math from pathlib import Path from typing import Any import numpy as np import torch import torch.nn.functional as F from safetensors.torch import load_file from scipy.io import wavfile from scipy.signal import resample_poly from torch import Tensor, nn from torch.nn.utils import remove_weight_norm from torch.nn.utils.parametrize import remove_parametrizations from torch.nn.utils.parametrizations import weight_norm DEFAULT_CHECKPOINT = "XCodec_RVQ4_mono_causal_fp32.safetensors" DEFAULT_CONFIG = "config.json" def load_json(path: Path) -> dict[str, Any]: with path.open() as f: return json.load(f) @torch.jit.script def snake_beta(x: Tensor, alpha: Tensor, beta: Tensor) -> Tensor: return x + (1.0 / (beta + 1e-9)) * torch.sin(x * alpha).pow(2) class SnakeBeta(nn.Module): def __init__(self, in_features: int, alpha: float = 1.0, alpha_trainable: bool = True) -> None: super().__init__() self.alpha = nn.Parameter(torch.zeros(in_features) * alpha, requires_grad=alpha_trainable) self.beta = nn.Parameter(torch.zeros(in_features) * alpha, requires_grad=alpha_trainable) def forward(self, x: Tensor) -> Tensor: alpha = torch.exp(self.alpha).unsqueeze(0).unsqueeze(-1) beta = torch.exp(self.beta).unsqueeze(0).unsqueeze(-1) return snake_beta(x, alpha, beta) class LayerNorm(nn.Module): def __init__(self, size: int, eps: float = 1e-5) -> None: super().__init__() self.weight = nn.Parameter(torch.ones(size)) self.bias = None self.eps = eps def forward(self, tensor: Tensor) -> Tensor: dtype = tensor.dtype tensor = F.layer_norm(tensor.float(), self.weight.shape, self.weight.float(), self.bias, self.eps) return tensor.to(dtype) def wn_conv1d(*args: Any, **kwargs: Any) -> nn.Conv1d: return weight_norm(nn.Conv1d(*args, **kwargs)) def wn_conv_transpose1d(*args: Any, **kwargs: Any) -> nn.ConvTranspose1d: return weight_norm(nn.ConvTranspose1d(*args, **kwargs)) def pad1d(x: Tensor, paddings: tuple[int, int], mode: str = "zero", value: float = 0.0) -> Tensor: left, right = paddings if mode == "reflect": max_pad = max(left, right) extra_pad = 0 if x.shape[-1] <= max_pad: extra_pad = max_pad - x.shape[-1] + 1 x = F.pad(x, (0, extra_pad)) padded = F.pad(x, paddings, mode, value) return padded[..., : padded.shape[-1] - extra_pad] return F.pad(x, paddings, mode, value) def unpad1d(x: Tensor, paddings: tuple[int, int]) -> Tensor: left, right = paddings end = x.shape[-1] - right return x[..., left:end] def get_extra_padding_for_conv1d( x: Tensor, kernel_size: int, stride: int, padding_total: int = 0, ) -> int: length = x.shape[-1] n_frames = (length - kernel_size + padding_total) / stride + 1 ideal_length = (math.ceil(n_frames) - 1) * stride + (kernel_size - padding_total) return ideal_length - length class NormConv1d(nn.Module): def __init__(self, *args: Any, norm: str = "none", **kwargs: Any) -> None: super().__init__() conv = nn.Conv1d(*args, **kwargs) self.conv = weight_norm(conv) if norm == "weight_norm" else conv self.norm = nn.Identity() def forward(self, x: Tensor) -> Tensor: return self.norm(self.conv(x)) class NormConvTranspose1d(nn.Module): def __init__(self, *args: Any, norm: str = "none", causal: bool = False, **kwargs: Any) -> None: super().__init__() convtr = nn.ConvTranspose1d(*args, **kwargs) self.convtr = weight_norm(convtr) if norm == "weight_norm" else convtr self.norm = nn.Identity() def forward(self, x: Tensor) -> Tensor: return self.norm(self.convtr(x)) class SConv1d(nn.Module): def __init__( self, in_channels: int, out_channels: int, kernel_size: int, stride: int = 1, dilation: int = 1, groups: int = 1, bias: bool = True, causal: bool = False, norm: str = "none", pad_mode: str = "reflect", ) -> None: super().__init__() self.conv = NormConv1d( in_channels, out_channels, kernel_size, stride, dilation=dilation, groups=groups, bias=bias, norm=norm, ) self.causal = causal self.pad_mode = pad_mode def forward(self, x: Tensor) -> Tensor: kernel_size = self.conv.conv.kernel_size[0] stride = self.conv.conv.stride[0] dilation = self.conv.conv.dilation[0] padding_total = (kernel_size - 1) * dilation - (stride - 1) extra_padding = get_extra_padding_for_conv1d(x, kernel_size, stride, padding_total) if self.causal: x = pad1d(x, (padding_total, extra_padding), mode=self.pad_mode) else: right = padding_total // 2 left = padding_total - right x = pad1d(x, (left, right + extra_padding), mode=self.pad_mode) return self.conv(x) class SConvTranspose1d(nn.Module): def __init__( self, in_channels: int, out_channels: int, kernel_size: int, stride: int = 1, causal: bool = False, norm: str = "none", trim_right_ratio: float = 1.0, ) -> None: super().__init__() self.convtr = NormConvTranspose1d( in_channels, out_channels, kernel_size, stride, causal=causal, norm=norm, ) self.causal = causal self.trim_right_ratio = trim_right_ratio def forward(self, x: Tensor) -> Tensor: kernel_size = self.convtr.convtr.kernel_size[0] stride = self.convtr.convtr.stride[0] padding_total = kernel_size - stride y = self.convtr(x) if self.causal: right = math.ceil(padding_total * self.trim_right_ratio) left = padding_total - right else: right = padding_total // 2 left = padding_total - right return unpad1d(y, (left, right)) class TrimPadding(nn.Module): def __init__(self, padding: int) -> None: super().__init__() self.padding = padding def forward(self, x: Tensor) -> Tensor: return x[:, :, : -self.padding] class ConvNeXtBlock(nn.Module): def __init__( self, dim: int, intermediate_dim: int, identity_init: bool = False, use_snake: bool = False, causal: bool = False, ) -> None: super().__init__() pad = (6, 0) if causal else (3, 3) self.dwconv = nn.Sequential(nn.ConstantPad1d(pad, 0), nn.Conv1d(dim, dim, kernel_size=7, groups=dim)) self.norm = LayerNorm(dim) self.pwconv1 = nn.Conv1d(dim, intermediate_dim, 1) self.act = SnakeBeta(intermediate_dim) if use_snake else nn.GELU() self.pwconv2 = nn.Conv1d(intermediate_dim, dim, 1) if identity_init: nn.init.zeros_(self.pwconv2.weight) if self.pwconv2.bias is not None: nn.init.zeros_(self.pwconv2.bias) def forward(self, x: Tensor) -> Tensor: residual = x x = self.dwconv(x) x = self.norm(x.permute(0, 2, 1)).permute(0, 2, 1) x = self.pwconv1(x) x = self.act(x) x = self.pwconv2(x) return residual + x def spectrogram( wav: Tensor, n_fft: int, hop_length: int, win_length: int, ) -> Tensor: left = (n_fft - hop_length) // 2 right = (n_fft - hop_length) - left wav = F.pad(wav, (left, right)).float() return torch.stft( wav, n_fft, hop_length=hop_length, win_length=win_length, window=torch.hann_window(win_length, device=wav.device).to(wav), center=False, normalized=False, onesided=True, return_complex=True, ) class SpectrogramConvNeXtEncoder(nn.Module): def __init__( self, in_channels: int, channels: int, latent_dim: int, c_mults: list[int], strides: list[int], identity_init: bool, n_fft: int, hop_length: int, use_snake: bool, causal: bool, padding_mode: str, num_blocks: int = 2, ) -> None: super().__init__() self.in_channels = in_channels self.n_fft = n_fft self.hop_length = hop_length layers: list[nn.Module] = [ wn_conv1d((n_fft + 2) * in_channels, c_mults[0] * channels, kernel_size=1, bias=False) ] for i, _ in enumerate(c_mults): dim_in = c_mults[i] * channels dim_out = c_mults[i + 1] * channels if i < len(c_mults) - 1 else c_mults[-1] * channels for _ in range(num_blocks): layers.append( ConvNeXtBlock( dim=dim_in, intermediate_dim=dim_in * 4, identity_init=identity_init, use_snake=use_snake, causal=causal, ) ) if causal: layers.append( SConv1d( in_channels=dim_in, out_channels=dim_out, kernel_size=2 * strides[i], stride=strides[i], causal=True, norm="weight_norm", ) ) else: layers.append( wn_conv1d( in_channels=dim_in, out_channels=dim_out, kernel_size=2 * strides[i], stride=strides[i], padding=math.ceil(strides[i] / 2), padding_mode=padding_mode, ) ) layers.append(wn_conv1d(c_mults[-1] * channels, latent_dim, kernel_size=1, bias=False)) self.layers = nn.Sequential(*layers) def forward(self, x: Tensor) -> Tensor: batch, channels, length = x.shape x_spec_in = x.reshape(batch * channels, 1, length) if channels > 1 else x with torch.autocast(device_type=x.device.type, enabled=False): spec = spectrogram( x_spec_in.float().squeeze(1), n_fft=self.n_fft, hop_length=self.hop_length, win_length=self.n_fft, ) real, imag = torch.view_as_real(spec).chunk(2, dim=-1) spec_features = torch.cat([real, imag], dim=1).squeeze(-1) spec_features = spec_features.to(x.dtype) if channels > 1: freq = spec_features.shape[1] spec_features = spec_features.reshape(batch, channels * freq, *spec_features.shape[2:]) return self.layers(spec_features) class ResidualUnit(nn.Module): def __init__( self, in_channels: int, out_channels: int, dilation: int, kernel_size: int = 7, use_snake: bool = False, causal: bool = False, padding_mode: str = "zeros", ) -> None: super().__init__() padding = dilation * (kernel_size - 1) if causal else (dilation * (kernel_size - 1)) // 2 layers: list[nn.Module] = [ SnakeBeta(out_channels) if use_snake else nn.ELU(), wn_conv1d( in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, dilation=dilation, padding=padding, padding_mode=padding_mode, ), SnakeBeta(out_channels) if use_snake else nn.ELU(), wn_conv1d(in_channels=out_channels, out_channels=out_channels, kernel_size=1, padding=0), ] self.layers = nn.Sequential(*layers) self.causal = causal self.padding = padding def forward(self, x: Tensor) -> Tensor: out = self.layers(x) if self.causal: out = out[:, :, : -self.padding] return x + out class DecoderBlock(nn.Module): def __init__( self, in_channels: int, out_channels: int, stride: int, use_snake: bool = False, causal: bool = False, padding_mode: str = "zeros", ) -> None: super().__init__() upsample = ( SConvTranspose1d( in_channels=in_channels, out_channels=out_channels, kernel_size=2 * stride, stride=stride, causal=True, norm="weight_norm", ) if causal else wn_conv_transpose1d( in_channels=in_channels, out_channels=out_channels, kernel_size=2 * stride, stride=stride, padding=math.ceil(stride / 2), output_padding=stride % 2, padding_mode="zeros", ) ) self.layers = nn.Sequential( SnakeBeta(in_channels) if use_snake else nn.ELU(), upsample, ResidualUnit(out_channels, out_channels, dilation=1, use_snake=use_snake, causal=causal, padding_mode=padding_mode), ResidualUnit(out_channels, out_channels, dilation=3, use_snake=use_snake, causal=causal, padding_mode=padding_mode), ResidualUnit(out_channels, out_channels, dilation=9, use_snake=use_snake, causal=causal, padding_mode=padding_mode), ) def forward(self, x: Tensor) -> Tensor: return self.layers(x) class OobleckDecoder(nn.Module): def __init__( self, out_channels: int, channels: int, latent_dim: int, c_mults: list[int], strides: list[int], use_snake: bool, final_tanh: bool, causal: bool, padding_mode: str, ) -> None: super().__init__() c_mults = [1] + c_mults first_padding = 6 if causal else 3 first_conv: nn.Module = wn_conv1d( in_channels=latent_dim, out_channels=c_mults[-1] * channels, kernel_size=7, padding=first_padding, padding_mode=padding_mode, ) if causal: first_conv = nn.Sequential(first_conv, TrimPadding(first_padding)) layers: list[nn.Module] = [first_conv] for i in range(len(c_mults) - 1, 0, -1): layers.append( DecoderBlock( in_channels=c_mults[i] * channels, out_channels=c_mults[i - 1] * channels, stride=strides[i - 1], use_snake=use_snake, causal=causal, padding_mode=padding_mode, ) ) final_padding = 6 if causal else 3 final_conv: nn.Module = wn_conv1d( in_channels=c_mults[0] * channels, out_channels=out_channels, kernel_size=7, padding=final_padding, padding_mode=padding_mode, bias=False, ) if causal: final_conv = nn.Sequential(final_conv, TrimPadding(final_padding)) layers += [SnakeBeta(c_mults[0] * channels) if use_snake else nn.ELU(), final_conv, nn.Tanh() if final_tanh else nn.Identity()] self.layers = nn.Sequential(*layers) def forward(self, x: Tensor) -> Tensor: return self.layers(x) def vae_sample(mean: Tensor, scale: Tensor, deterministic: bool) -> Tensor: if deterministic: return mean stdev = F.softplus(scale) + 1e-4 return torch.randn_like(mean) * stdev + mean class VAEDistillationBottleneck(nn.Module): def __init__(self, latent_dim: int, distillation_proj_dim: int) -> None: super().__init__() self.distillation_proj = wn_conv1d(latent_dim, distillation_proj_dim, kernel_size=1, bias=False) def encode(self, x: Tensor, deterministic: bool = False) -> Tensor: mean, scale = x.chunk(2, dim=1) return vae_sample(mean, scale, deterministic=deterministic) def decode(self, x: Tensor) -> Tensor: return x class AudioAutoencoder(nn.Module): def __init__( self, encoder: nn.Module, decoder: nn.Module, bottleneck: VAEDistillationBottleneck, downsampling_ratio: int, sample_rate: int, ) -> None: super().__init__() self.encoder = encoder self.decoder = decoder self.bottleneck = bottleneck self.downsampling_ratio = downsampling_ratio self.sample_rate = sample_rate def encode_audio(self, audio: Tensor, deterministic: bool = False) -> Tensor: latents = self.encoder(audio) return self.bottleneck.encode(latents, deterministic=deterministic) def decode(self, latents: Tensor) -> Tensor: return self.decoder(self.bottleneck.decode(latents)) def create_model_from_config(config: dict[str, Any]) -> AudioAutoencoder: model_config = config["model"] encoder_config = model_config["encoder"]["config"] decoder_config = model_config["decoder"]["config"] bottleneck_config = model_config["bottleneck"]["config"] return AudioAutoencoder( encoder=SpectrogramConvNeXtEncoder(**encoder_config), decoder=OobleckDecoder(**decoder_config), bottleneck=VAEDistillationBottleneck(**bottleneck_config), downsampling_ratio=model_config["downsampling_ratio"], sample_rate=config["sample_rate"], ) def load_checkpoint_state(path: Path) -> dict[str, Tensor]: if path.suffix == ".safetensors": return load_file(path, device="cpu") checkpoint = torch.load(path, map_location="cpu", weights_only=True) return checkpoint["state_dict"] def remove_weight_norm_from_model(model: nn.Module) -> None: for module in model.modules(): if hasattr(module, "parametrizations"): try: remove_parametrizations(module, "weight") continue except ValueError: pass try: remove_weight_norm(module) except ValueError: pass def load_model( checkpoint_path: Path, config_path: Path, device: torch.device, ) -> AudioAutoencoder: config = load_json(config_path) model = create_model_from_config(config) state_dict = load_checkpoint_state(checkpoint_path) model.load_state_dict(state_dict, strict=True) remove_weight_norm_from_model(model) model.eval().requires_grad_(False) return model.to(device) def wav_to_float32(audio: np.ndarray) -> np.ndarray: if np.issubdtype(audio.dtype, np.floating): return audio.astype(np.float32) if audio.dtype == np.uint8: return (audio.astype(np.float32) - 128.0) / 128.0 info = np.iinfo(audio.dtype) return audio.astype(np.float32) / max(abs(info.min), info.max) def read_audio_mono(path: Path) -> tuple[np.ndarray, int]: sample_rate, audio = wavfile.read(path) audio = wav_to_float32(audio) if audio.ndim == 2: audio = audio.mean(axis=1) return audio.astype(np.float32), sample_rate def write_audio_mono(path: Path, audio: np.ndarray, sample_rate: int) -> None: path.parent.mkdir(parents=True, exist_ok=True) wavfile.write(path, sample_rate, audio.astype(np.float32)) def resample_audio(audio: np.ndarray, source_sr: int, target_sr: int) -> np.ndarray: if source_sr == target_sr: return audio.astype(np.float32) gcd = math.gcd(source_sr, target_sr) up = target_sr // gcd down = source_sr // gcd return resample_poly(audio, up, down).astype(np.float32) def pad_to_multiple(audio: Tensor, multiple: int) -> Tensor: pad = (multiple - (audio.shape[-1] % multiple)) % multiple if pad == 0: return audio return F.pad(audio, (0, pad)) def enhance_file( model: AudioAutoencoder, input_path: Path, output_path: Path, deterministic: bool, ) -> None: audio, sample_rate = read_audio_mono(input_path) audio_48k = resample_audio(audio, sample_rate, model.sample_rate) original_len = audio_48k.shape[0] device = next(model.parameters()).device audio_tensor = torch.from_numpy(audio_48k).to(device).view(1, 1, -1) audio_tensor = pad_to_multiple(audio_tensor, model.downsampling_ratio) with torch.inference_mode(): latents = model.encode_audio(audio_tensor, deterministic=deterministic) enhanced = model.decode(latents)[..., :original_len] enhanced_np = enhanced.squeeze(0).squeeze(0).float().clamp(-1.0, 1.0).cpu().numpy() write_audio_mono(output_path, enhanced_np, model.sample_rate) def iter_input_files(path: Path) -> list[Path]: if path.is_file(): return [path] if path.suffix.lower() == ".wav" else [] return sorted(p for p in path.iterdir() if p.is_file() and p.suffix.lower() == ".wav")