xinference 1.9.1__py3-none-any.whl → 1.10.1__py3-none-any.whl

xinference/thirdparty/audiotools/core/ffmpeg.py

@@ -0,0 +1,211 @@
+import json
+import shlex
+import subprocess
+import tempfile
+from pathlib import Path
+from typing import Tuple
+
+import ffmpy
+import numpy as np
+import torch
+
+
+def r128stats(filepath: str, quiet: bool):
+    """Takes a path to an audio file, returns a dict with the loudness
+    stats computed by the ffmpeg ebur128 filter.
+
+    Parameters
+    ----------
+    filepath : str
+        Path to compute loudness stats on.
+    quiet : bool
+        Whether to show FFMPEG output during computation.
+
+    Returns
+    -------
+    dict
+        Dictionary containing loudness stats.
+    """
+    ffargs = [
+        "ffmpeg",
+        "-nostats",
+        "-i",
+        filepath,
+        "-filter_complex",
+        "ebur128",
+        "-f",
+        "null",
+        "-",
+    ]
+    if quiet:
+        ffargs += ["-hide_banner"]
+    proc = subprocess.Popen(ffargs, stderr=subprocess.PIPE, universal_newlines=True)
+    stats = proc.communicate()[1]
+    summary_index = stats.rfind("Summary:")
+
+    summary_list = stats[summary_index:].split()
+    i_lufs = float(summary_list[summary_list.index("I:") + 1])
+    i_thresh = float(summary_list[summary_list.index("I:") + 4])
+    lra = float(summary_list[summary_list.index("LRA:") + 1])
+    lra_thresh = float(summary_list[summary_list.index("LRA:") + 4])
+    lra_low = float(summary_list[summary_list.index("low:") + 1])
+    lra_high = float(summary_list[summary_list.index("high:") + 1])
+    stats_dict = {
+        "I": i_lufs,
+        "I Threshold": i_thresh,
+        "LRA": lra,
+        "LRA Threshold": lra_thresh,
+        "LRA Low": lra_low,
+        "LRA High": lra_high,
+    }
+
+    return stats_dict
+
+
+def ffprobe_offset_and_codec(path: str) -> Tuple[float, str]:
+    """Given a path to a file, returns the start time offset and codec of
+    the first audio stream.
+    """
+    ff = ffmpy.FFprobe(
+        inputs={path: None},
+        global_options="-show_entries format=start_time:stream=duration,start_time,codec_type,codec_name,start_pts,time_base -of json -v quiet",
+    )
+    streams = json.loads(ff.run(stdout=subprocess.PIPE)[0])["streams"]
+    seconds_offset = 0.0
+    codec = None
+
+    # Get the offset and codec of the first audio stream we find
+    # and return its start time, if it has one.
+    for stream in streams:
+        if stream["codec_type"] == "audio":
+            seconds_offset = stream.get("start_time", 0.0)
+            codec = stream.get("codec_name")
+            break
+    return float(seconds_offset), codec
+
+
+class FFMPEGMixin:
+    _loudness = None
+
+    def ffmpeg_loudness(self, quiet: bool = True):
+        """Computes loudness of audio file using FFMPEG.
+
+        Parameters
+        ----------
+        quiet : bool, optional
+            Whether to show FFMPEG output during computation,
+            by default True
+
+        Returns
+        -------
+        torch.Tensor
+            Loudness of every item in the batch, computed via
+            FFMPEG.
+        """
+        loudness = []
+
+        with tempfile.NamedTemporaryFile(suffix=".wav") as f:
+            for i in range(self.batch_size):
+                self[i].write(f.name)
+                loudness_stats = r128stats(f.name, quiet=quiet)
+                loudness.append(loudness_stats["I"])
+
+        self._loudness = torch.from_numpy(np.array(loudness)).float()
+        return self.loudness()
+
+    def ffmpeg_resample(self, sample_rate: int, quiet: bool = True):
+        """Resamples AudioSignal using FFMPEG. More memory-efficient
+        than using julius.resample for long audio files.
+
+        Parameters
+        ----------
+        sample_rate : int
+            Sample rate to resample to.
+        quiet : bool, optional
+            Whether to show FFMPEG output during computation,
+            by default True
+
+        Returns
+        -------
+        AudioSignal
+            Resampled AudioSignal.
+        """
+        from audiotools import AudioSignal
+
+        if sample_rate == self.sample_rate:
+            return self
+
+        with tempfile.NamedTemporaryFile(suffix=".wav") as f:
+            self.write(f.name)
+            f_out = f.name.replace("wav", "rs.wav")
+            command = f"ffmpeg -i {f.name} -ar {sample_rate} {f_out}"
+            if quiet:
+                command += " -hide_banner -loglevel error"
+            subprocess.check_call(shlex.split(command))
+            resampled = AudioSignal(f_out)
+            Path.unlink(Path(f_out))
+        return resampled
+
+    @classmethod
+    def load_from_file_with_ffmpeg(cls, audio_path: str, quiet: bool = True, **kwargs):
+        """Loads AudioSignal object after decoding it to a wav file using FFMPEG.
+        Useful for loading audio that isn't covered by librosa's loading mechanism. Also
+        useful for loading mp3 files, without any offset.
+
+        Parameters
+        ----------
+        audio_path : str
+            Path to load AudioSignal from.
+        quiet : bool, optional
+            Whether to show FFMPEG output during computation,
+            by default True
+
+        Returns
+        -------
+        AudioSignal
+            AudioSignal loaded from file with FFMPEG.
+        """
+        audio_path = str(audio_path)
+        with tempfile.TemporaryDirectory() as d:
+            wav_file = str(Path(d) / "extracted.wav")
+            padded_wav = str(Path(d) / "padded.wav")
+
+            global_options = "-y"
+            if quiet:
+                global_options += " -loglevel error"
+
+            ff = ffmpy.FFmpeg(
+                inputs={audio_path: None},
+                # For inputs that are m4a (and others?), the input audio can
+                # have samples that don't match the sample rate. This aresample
+                # option forces ffmpeg to read timing information in the source
+                # file instead of assuming constant sample rate.
+                #
+                # This fixes an issue where an input m4a file might be a
+                # different length than the output wav file
+                outputs={wav_file: "-af aresample=async=1000"},
+                global_options=global_options,
+            )
+            ff.run()
+
+            # We pad the file using the start time offset in case it's an audio
+            # stream starting at some offset in a video container.
+            pad, codec = ffprobe_offset_and_codec(audio_path)
+
+            # For mp3s, don't pad files with discrepancies less than 0.027s -
+            # it's likely due to codec latency. The amount of latency introduced
+            # by mp3 is 1152, which is 0.0261 44khz. So we set the threshold
+            # here slightly above that.
+            # Source: https://lame.sourceforge.io/tech-FAQ.txt.
+            if codec == "mp3" and pad < 0.027:
+                pad = 0.0
+            ff = ffmpy.FFmpeg(
+                inputs={wav_file: None},
+                outputs={padded_wav: f"-af 'adelay={pad*1000}:all=true'"},
+                global_options=global_options,
+            )
+            ff.run()
+
+            signal = cls(padded_wav, **kwargs)
+
+        return signal

xinference/thirdparty/audiotools/core/loudness.py

@@ -0,0 +1,320 @@
+import copy
+
+import julius
+import numpy as np
+import scipy
+import torch
+import torch.nn.functional as F
+import torchaudio
+
+
+class Meter(torch.nn.Module):
+    """Tensorized version of pyloudnorm.Meter. Works with batched audio tensors.
+
+    Parameters
+    ----------
+    rate : int
+        Sample rate of audio.
+    filter_class : str, optional
+        Class of weighting filter used.
+        K-weighting' (default), 'Fenton/Lee 1'
+        'Fenton/Lee 2', 'Dash et al.'
+        by default "K-weighting"
+    block_size : float, optional
+        Gating block size in seconds, by default 0.400
+    zeros : int, optional
+         Number of zeros to use in FIR approximation of
+         IIR filters, by default 512
+    use_fir : bool, optional
+        Whether to use FIR approximation or exact IIR formulation.
+        If computing on GPU, ``use_fir=True`` will be used, as its
+        much faster, by default False
+    """
+
+    def __init__(
+        self,
+        rate: int,
+        filter_class: str = "K-weighting",
+        block_size: float = 0.400,
+        zeros: int = 512,
+        use_fir: bool = False,
+    ):
+        super().__init__()
+
+        self.rate = rate
+        self.filter_class = filter_class
+        self.block_size = block_size
+        self.use_fir = use_fir
+
+        G = torch.from_numpy(np.array([1.0, 1.0, 1.0, 1.41, 1.41]))
+        self.register_buffer("G", G)
+
+        # Compute impulse responses so that filtering is fast via
+        # a convolution at runtime, on GPU, unlike lfilter.
+        impulse = np.zeros((zeros,))
+        impulse[..., 0] = 1.0
+
+        firs = np.zeros((len(self._filters), 1, zeros))
+        passband_gain = torch.zeros(len(self._filters))
+
+        for i, (_, filter_stage) in enumerate(self._filters.items()):
+            firs[i] = scipy.signal.lfilter(filter_stage.b, filter_stage.a, impulse)
+            passband_gain[i] = filter_stage.passband_gain
+
+        firs = torch.from_numpy(firs[..., ::-1].copy()).float()
+
+        self.register_buffer("firs", firs)
+        self.register_buffer("passband_gain", passband_gain)
+
+    def apply_filter_gpu(self, data: torch.Tensor):
+        """Performs FIR approximation of loudness computation.
+
+        Parameters
+        ----------
+        data : torch.Tensor
+            Audio data of shape (nb, nch, nt).
+
+        Returns
+        -------
+        torch.Tensor
+            Filtered audio data.
+        """
+        # Data is of shape (nb, nch, nt)
+        # Reshape to (nb*nch, 1, nt)
+        nb, nt, nch = data.shape
+        data = data.permute(0, 2, 1)
+        data = data.reshape(nb * nch, 1, nt)
+
+        # Apply padding
+        pad_length = self.firs.shape[-1]
+
+        # Apply filtering in sequence
+        for i in range(self.firs.shape[0]):
+            data = F.pad(data, (pad_length, pad_length))
+            data = julius.fftconv.fft_conv1d(data, self.firs[i, None, ...])
+            data = self.passband_gain[i] * data
+            data = data[..., 1 : nt + 1]
+
+        data = data.permute(0, 2, 1)
+        data = data[:, :nt, :]
+        return data
+
+    def apply_filter_cpu(self, data: torch.Tensor):
+        """Performs IIR formulation of loudness computation.
+
+        Parameters
+        ----------
+        data : torch.Tensor
+            Audio data of shape (nb, nch, nt).
+
+        Returns
+        -------
+        torch.Tensor
+            Filtered audio data.
+        """
+        for _, filter_stage in self._filters.items():
+            passband_gain = filter_stage.passband_gain
+
+            a_coeffs = torch.from_numpy(filter_stage.a).float().to(data.device)
+            b_coeffs = torch.from_numpy(filter_stage.b).float().to(data.device)
+
+            _data = data.permute(0, 2, 1)
+            filtered = torchaudio.functional.lfilter(
+                _data, a_coeffs, b_coeffs, clamp=False
+            )
+            data = passband_gain * filtered.permute(0, 2, 1)
+        return data
+
+    def apply_filter(self, data: torch.Tensor):
+        """Applies filter on either CPU or GPU, depending
+        on if the audio is on GPU or is on CPU, or if
+        ``self.use_fir`` is True.
+
+        Parameters
+        ----------
+        data : torch.Tensor
+            Audio data of shape (nb, nch, nt).
+
+        Returns
+        -------
+        torch.Tensor
+            Filtered audio data.
+        """
+        if data.is_cuda or self.use_fir:
+            data = self.apply_filter_gpu(data)
+        else:
+            data = self.apply_filter_cpu(data)
+        return data
+
+    def forward(self, data: torch.Tensor):
+        """Computes integrated loudness of data.
+
+        Parameters
+        ----------
+        data : torch.Tensor
+            Audio data of shape (nb, nch, nt).
+
+        Returns
+        -------
+        torch.Tensor
+            Filtered audio data.
+        """
+        return self.integrated_loudness(data)
+
+    def _unfold(self, input_data):
+        T_g = self.block_size
+        overlap = 0.75  # overlap of 75% of the block duration
+        step = 1.0 - overlap  # step size by percentage
+
+        kernel_size = int(T_g * self.rate)
+        stride = int(T_g * self.rate * step)
+        unfolded = julius.core.unfold(input_data.permute(0, 2, 1), kernel_size, stride)
+        unfolded = unfolded.transpose(-1, -2)
+
+        return unfolded
+
+    def integrated_loudness(self, data: torch.Tensor):
+        """Computes integrated loudness of data.
+
+        Parameters
+        ----------
+        data : torch.Tensor
+            Audio data of shape (nb, nch, nt).
+
+        Returns
+        -------
+        torch.Tensor
+            Filtered audio data.
+        """
+        if not torch.is_tensor(data):
+            data = torch.from_numpy(data).float()
+        else:
+            data = data.float()
+
+        input_data = copy.copy(data)
+        # Data always has a batch and channel dimension.
+        # Is of shape (nb, nt, nch)
+        if input_data.ndim < 2:
+            input_data = input_data.unsqueeze(-1)
+        if input_data.ndim < 3:
+            input_data = input_data.unsqueeze(0)
+
+        nb, nt, nch = input_data.shape
+
+        # Apply frequency weighting filters - account
+        # for the acoustic respose of the head and auditory system
+        input_data = self.apply_filter(input_data)
+
+        G = self.G  # channel gains
+        T_g = self.block_size  # 400 ms gating block standard
+        Gamma_a = -70.0  # -70 LKFS = absolute loudness threshold
+
+        unfolded = self._unfold(input_data)
+
+        z = (1.0 / (T_g * self.rate)) * unfolded.square().sum(2)
+        l = -0.691 + 10.0 * torch.log10((G[None, :nch, None] * z).sum(1, keepdim=True))
+        l = l.expand_as(z)
+
+        # find gating block indices above absolute threshold
+        z_avg_gated = z
+        z_avg_gated[l <= Gamma_a] = 0
+        masked = l > Gamma_a
+        z_avg_gated = z_avg_gated.sum(2) / masked.sum(2)
+
+        # calculate the relative threshold value (see eq. 6)
+        Gamma_r = (
+            -0.691 + 10.0 * torch.log10((z_avg_gated * G[None, :nch]).sum(-1)) - 10.0
+        )
+        Gamma_r = Gamma_r[:, None, None]
+        Gamma_r = Gamma_r.expand(nb, nch, l.shape[-1])
+
+        # find gating block indices above relative and absolute thresholds  (end of eq. 7)
+        z_avg_gated = z
+        z_avg_gated[l <= Gamma_a] = 0
+        z_avg_gated[l <= Gamma_r] = 0
+        masked = (l > Gamma_a) * (l > Gamma_r)
+        z_avg_gated = z_avg_gated.sum(2) / masked.sum(2)
+
+        # # Cannot use nan_to_num (pytorch 1.8 does not come with GCP-supported cuda version)
+        # z_avg_gated = torch.nan_to_num(z_avg_gated)
+        z_avg_gated = torch.where(
+            z_avg_gated.isnan(), torch.zeros_like(z_avg_gated), z_avg_gated
+        )
+        z_avg_gated[z_avg_gated == float("inf")] = float(np.finfo(np.float32).max)
+        z_avg_gated[z_avg_gated == -float("inf")] = float(np.finfo(np.float32).min)
+
+        LUFS = -0.691 + 10.0 * torch.log10((G[None, :nch] * z_avg_gated).sum(1))
+        return LUFS.float()
+
+    @property
+    def filter_class(self):
+        return self._filter_class
+
+    @filter_class.setter
+    def filter_class(self, value):
+        from pyloudnorm import Meter
+
+        meter = Meter(self.rate)
+        meter.filter_class = value
+        self._filter_class = value
+        self._filters = meter._filters
+
+
+class LoudnessMixin:
+    _loudness = None
+    MIN_LOUDNESS = -70
+    """Minimum loudness possible."""
+
+    def loudness(
+        self, filter_class: str = "K-weighting", block_size: float = 0.400, **kwargs
+    ):
+        """Calculates loudness using an implementation of ITU-R BS.1770-4.
+        Allows control over gating block size and frequency weighting filters for
+        additional control. Measure the integrated gated loudness of a signal.
+
+        API is derived from PyLoudnorm, but this implementation is ported to PyTorch
+        and is tensorized across batches. When on GPU, an FIR approximation of the IIR
+        filters is used to compute loudness for speed.
+
+        Uses the weighting filters and block size defined by the meter
+        the integrated loudness is measured based upon the gating algorithm
+        defined in the ITU-R BS.1770-4 specification.
+
+        Parameters
+        ----------
+        filter_class : str, optional
+            Class of weighting filter used.
+            K-weighting' (default), 'Fenton/Lee 1'
+            'Fenton/Lee 2', 'Dash et al.'
+            by default "K-weighting"
+        block_size : float, optional
+            Gating block size in seconds, by default 0.400
+        kwargs : dict, optional
+            Keyword arguments to :py:func:`audiotools.core.loudness.Meter`.
+
+        Returns
+        -------
+        torch.Tensor
+            Loudness of audio data.
+        """
+        if self._loudness is not None:
+            return self._loudness.to(self.device)
+        original_length = self.signal_length
+        if self.signal_duration < 0.5:
+            pad_len = int((0.5 - self.signal_duration) * self.sample_rate)
+            self.zero_pad(0, pad_len)
+
+        # create BS.1770 meter
+        meter = Meter(
+            self.sample_rate, filter_class=filter_class, block_size=block_size, **kwargs
+        )
+        meter = meter.to(self.device)
+        # measure loudness
+        loudness = meter.integrated_loudness(self.audio_data.permute(0, 2, 1))
+        self.truncate_samples(original_length)
+        min_loudness = (
+            torch.ones_like(loudness, device=loudness.device) * self.MIN_LOUDNESS
+        )
+        self._loudness = torch.maximum(loudness, min_loudness)
+
+        return self._loudness.to(self.device)