soulxpodcast/utils/audio.py

import torch
import numpy as np
from librosa.filters import mel as librosa_mel_fn
from scipy.io.wavfile import read

MAX_WAV_VALUE = 32768.0


def load_wav(full_path):
    sampling_rate, data = read(full_path)
    return data, sampling_rate


def dynamic_range_compression(x, C=1, clip_val=1e-5):
    return np.log(np.clip(x, a_min=clip_val, a_max=None) * C)


def dynamic_range_decompression(x, C=1):
    return np.exp(x) / C


def dynamic_range_compression_torch(x, C=1, clip_val=1e-5):
    return torch.log(torch.clamp(x, min=clip_val) * C)


def dynamic_range_decompression_torch(x, C=1):
    return torch.exp(x) / C


def spectral_normalize_torch(magnitudes):
    output = dynamic_range_compression_torch(magnitudes)
    return output


def spectral_de_normalize_torch(magnitudes):
    output = dynamic_range_decompression_torch(magnitudes)
    return output


mel_basis = {}
hann_window = {}


def mel_spectrogram(y, n_fft=1920, num_mels=80, sampling_rate=24000, hop_size=480,
                    win_size=1920, fmin=0, fmax=8000, center=False):
    global mel_basis, hann_window
    if f"{str(fmax)}_{str(y.device)}" not in mel_basis:
        mel = librosa_mel_fn(sr=sampling_rate, n_fft=n_fft, n_mels=num_mels, fmin=fmin, fmax=fmax)
        mel_basis[str(fmax) + "_" + str(y.device)] = torch.from_numpy(mel).float().to(y.device)
        hann_window[str(y.device)] = torch.hann_window(win_size).to(y.device)

    y = torch.nn.functional.pad(
        y.unsqueeze(1), (int((n_fft - hop_size) / 2), int((n_fft - hop_size) / 2)), mode="reflect"
    )
    y = y.squeeze(1)

    spec = torch.view_as_real(
        torch.stft(
            y,
            n_fft,
            hop_length=hop_size,
            win_length=win_size,
            window=hann_window[str(y.device)],
            center=center,
            pad_mode="reflect",
            normalized=False,
            onesided=True,
            return_complex=True,
        )
    )

    spec = torch.sqrt(spec.pow(2).sum(-1) + (1e-9))

    spec = torch.matmul(mel_basis[str(fmax) + "_" + str(y.device)], spec)
    spec = spectral_normalize_torch(spec)

    return spec


def audio_volume_normalize(audio: torch.Tensor, coeff=0.1):
    """
    Normalize the volume of an audio signal.

    Parameters:
        audio (torch tensor): Input audio signal array.
        coeff (float): Target coefficient for normalization, default is 0.1.

    Returns:
        torch tensor: The volume-normalized audio signal.
    """

    device = audio.device
    audio = audio.cpu().numpy()
    temp = np.sort(np.abs(audio))


    if temp[-1] < 0.1:
        scaling_factor = max(
            temp[-1], 1e-3
        )
        audio = audio / scaling_factor * 0.1


    temp = temp[temp > 0.01]
    L = temp.shape[0]


    if L <= 10:
        return audio


    volume = np.mean(temp[int(0.9 * L) : int(0.99 * L)])


    audio = audio * np.clip(coeff / volume, a_min=0.1, a_max=10)


    max_value = np.max(np.abs(audio))
    if max_value > 1:
        audio = audio / max_value

    audio = torch.from_numpy(audio).to(device)
    return audio