enhancement_VAE/enhancement_vae.py

#!/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")