softwareweaver/MusicGen

1# Copyright (c) Meta Platforms, Inc. and affiliates.2# All rights reserved.3#4# This source code is licensed under the license found in the5# LICENSE file in the root directory of this source tree.6"""Various utilities for audio convertion (pcm format, sample rate and channels),7and volume normalization."""8import sys9import typing as tp10 11import julius12import torch13import torchaudio14 15 16def convert_audio_channels(wav: torch.Tensor, channels: int = 2) -> torch.Tensor:17    """Convert audio to the given number of channels.18 19    Args:20        wav (torch.Tensor): Audio wave of shape [B, C, T].21        channels (int): Expected number of channels as output.22    Returns:23        torch.Tensor: Downmixed or unchanged audio wave [B, C, T].24    """25    *shape, src_channels, length = wav.shape26    if src_channels == channels:27        pass28    elif channels == 1:29        # Case 1:30        # The caller asked 1-channel audio, and the stream has multiple31        # channels, downmix all channels.32        wav = wav.mean(dim=-2, keepdim=True)33    elif src_channels == 1:34        # Case 2:35        # The caller asked for multiple channels, but the input file has36        # a single channel, replicate the audio over all channels.37        wav = wav.expand(*shape, channels, length)38    elif src_channels >= channels:39        # Case 3:40        # The caller asked for multiple channels, and the input file has41        # more channels than requested. In that case return the first channels.42        wav = wav[..., :channels, :]43    else:44        # Case 4: What is a reasonable choice here?45        raise ValueError('The audio file has less channels than requested but is not mono.')46    return wav47 48 49def convert_audio(wav: torch.Tensor, from_rate: float,50                  to_rate: float, to_channels: int) -> torch.Tensor:51    """Convert audio to new sample rate and number of audio channels."""52    wav = julius.resample_frac(wav, int(from_rate), int(to_rate))53    wav = convert_audio_channels(wav, to_channels)54    return wav55 56 57def normalize_loudness(wav: torch.Tensor, sample_rate: int, loudness_headroom_db: float = 14,58                       loudness_compressor: bool = False, energy_floor: float = 2e-3):59    """Normalize an input signal to a user loudness in dB LKFS.60    Audio loudness is defined according to the ITU-R BS.1770-4 recommendation.61 62    Args:63        wav (torch.Tensor): Input multichannel audio data.64        sample_rate (int): Sample rate.65        loudness_headroom_db (float): Target loudness of the output in dB LUFS.66        loudness_compressor (bool): Uses tanh for soft clipping.67        energy_floor (float): anything below that RMS level will not be rescaled.68    Returns:69        torch.Tensor: Loudness normalized output data.70    """71    energy = wav.pow(2).mean().sqrt().item()72    if energy < energy_floor:73        return wav74    transform = torchaudio.transforms.Loudness(sample_rate)75    input_loudness_db = transform(wav).item()76    # calculate the gain needed to scale to the desired loudness level77    delta_loudness = -loudness_headroom_db - input_loudness_db78    gain = 10.0 ** (delta_loudness / 20.0)79    output = gain * wav80    if loudness_compressor:81        output = torch.tanh(output)82    assert output.isfinite().all(), (input_loudness_db, wav.pow(2).mean().sqrt())83    return output84 85 86def _clip_wav(wav: torch.Tensor, log_clipping: bool = False, stem_name: tp.Optional[str] = None) -> None:87    """Utility function to clip the audio with logging if specified."""88    max_scale = wav.abs().max()89    if log_clipping and max_scale > 1:90        clamp_prob = (wav.abs() > 1).float().mean().item()91        print(f"CLIPPING {stem_name or ''} happening with proba (a bit of clipping is okay):",92              clamp_prob, "maximum scale: ", max_scale.item(), file=sys.stderr)93    wav.clamp_(-1, 1)94 95 96def normalize_audio(wav: torch.Tensor, normalize: bool = True,97                    strategy: str = 'peak', peak_clip_headroom_db: float = 1,98                    rms_headroom_db: float = 18, loudness_headroom_db: float = 14,99                    loudness_compressor: bool = False, log_clipping: bool = False,100                    sample_rate: tp.Optional[int] = None,101                    stem_name: tp.Optional[str] = None) -> torch.Tensor:102    """Normalize the audio according to the prescribed strategy (see after).103 104    Args:105        wav (torch.Tensor): Audio data.106        normalize (bool): if `True` (default), normalizes according to the prescribed107            strategy (see after). If `False`, the strategy is only used in case clipping108            would happen.109        strategy (str): Can be either 'clip', 'peak', or 'rms'. Default is 'peak',110            i.e. audio is normalized by its largest value. RMS normalizes by root-mean-square111            with extra headroom to avoid clipping. 'clip' just clips.112        peak_clip_headroom_db (float): Headroom in dB when doing 'peak' or 'clip' strategy.113        rms_headroom_db (float): Headroom in dB when doing 'rms' strategy. This must be much larger114            than the `peak_clip` one to avoid further clipping.115        loudness_headroom_db (float): Target loudness for loudness normalization.116        loudness_compressor (bool): If True, uses tanh based soft clipping.117        log_clipping (bool): If True, basic logging on stderr when clipping still118            occurs despite strategy (only for 'rms').119        sample_rate (int): Sample rate for the audio data (required for loudness).120        stem_name (str, optional): Stem name for clipping logging.121    Returns:122        torch.Tensor: Normalized audio.123    """124    scale_peak = 10 ** (-peak_clip_headroom_db / 20)125    scale_rms = 10 ** (-rms_headroom_db / 20)126    if strategy == 'peak':127        rescaling = (scale_peak / wav.abs().max())128        if normalize or rescaling < 1:129            wav = wav * rescaling130    elif strategy == 'clip':131        wav = wav.clamp(-scale_peak, scale_peak)132    elif strategy == 'rms':133        mono = wav.mean(dim=0)134        rescaling = scale_rms / mono.pow(2).mean().sqrt()135        if normalize or rescaling < 1:136            wav = wav * rescaling137        _clip_wav(wav, log_clipping=log_clipping, stem_name=stem_name)138    elif strategy == 'loudness':139        assert sample_rate is not None, "Loudness normalization requires sample rate."140        wav = normalize_loudness(wav, sample_rate, loudness_headroom_db, loudness_compressor)141        _clip_wav(wav, log_clipping=log_clipping, stem_name=stem_name)142    else:143        assert wav.abs().max() < 1144        assert strategy == '' or strategy == 'none', f"Unexpected strategy: '{strategy}'"145    return wav146 147 148def f32_pcm(wav: torch.Tensor) -> torch.Tensor:149    """Convert audio to float 32 bits PCM format.150    """151    if wav.dtype.is_floating_point:152        return wav153    elif wav.dtype == torch.int16:154        return wav.float() / 2**15155    elif wav.dtype == torch.int32:156        return wav.float() / 2**31157    raise ValueError(f"Unsupported wav dtype: {wav.dtype}")158 159 160def i16_pcm(wav: torch.Tensor) -> torch.Tensor:161    """Convert audio to int 16 bits PCM format.162 163    ..Warning:: There exist many formula for doing this conversion. None are perfect164    due to the asymmetry of the int16 range. One either have possible clipping, DC offset,165    or inconsistencies with f32_pcm. If the given wav doesn't have enough headroom,166    it is possible that `i16_pcm(f32_pcm)) != Identity`.167    """168    if wav.dtype.is_floating_point:169        assert wav.abs().max() <= 1170        candidate = (wav * 2 ** 15).round()171        if candidate.max() >= 2 ** 15:  # clipping would occur172            candidate = (wav * (2 ** 15 - 1)).round()173        return candidate.short()174    else:175        assert wav.dtype == torch.int16176        return wav177