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

xinference/thirdparty/audiotools/ml/layers/__init__.py

@@ -0,0 +1,2 @@
+from .base import BaseModel
+from .spectral_gate import SpectralGate

xinference/thirdparty/audiotools/ml/layers/base.py

@@ -0,0 +1,328 @@
+import inspect
+import shutil
+import tempfile
+import typing
+from pathlib import Path
+
+import torch
+from torch import nn
+
+
+class BaseModel(nn.Module):
+    """This is a class that adds useful save/load functionality to a
+    ``torch.nn.Module`` object. ``BaseModel`` objects can be saved
+    as ``torch.package`` easily, making them super easy to port between
+    machines without requiring a ton of dependencies. Files can also be
+    saved as just weights, in the standard way.
+
+    >>> class Model(ml.BaseModel):
+    >>>     def __init__(self, arg1: float = 1.0):
+    >>>         super().__init__()
+    >>>         self.arg1 = arg1
+    >>>         self.linear = nn.Linear(1, 1)
+    >>>
+    >>>     def forward(self, x):
+    >>>         return self.linear(x)
+    >>>
+    >>> model1 = Model()
+    >>>
+    >>> with tempfile.NamedTemporaryFile(suffix=".pth") as f:
+    >>>     model1.save(
+    >>>         f.name,
+    >>>     )
+    >>>     model2 = Model.load(f.name)
+    >>>     out2 = seed_and_run(model2, x)
+    >>>     assert torch.allclose(out1, out2)
+    >>>
+    >>>     model1.save(f.name, package=True)
+    >>>     model2 = Model.load(f.name)
+    >>>     model2.save(f.name, package=False)
+    >>>     model3 = Model.load(f.name)
+    >>>     out3 = seed_and_run(model3, x)
+    >>>
+    >>> with tempfile.TemporaryDirectory() as d:
+    >>>     model1.save_to_folder(d, {"data": 1.0})
+    >>>     Model.load_from_folder(d)
+
+    """
+
+    EXTERN = [
+        "audiotools.**",
+        "tqdm",
+        "__main__",
+        "numpy.**",
+        "julius.**",
+        "torchaudio.**",
+        "scipy.**",
+        "einops",
+    ]
+    """Names of libraries that are external to the torch.package saving mechanism.
+    Source code from these libraries will not be packaged into the model. This can
+    be edited by the user of this class by editing ``model.EXTERN``."""
+    INTERN = []
+    """Names of libraries that are internal to the torch.package saving mechanism.
+    Source code from these libraries will be saved alongside the model."""
+
+    def save(
+        self,
+        path: str,
+        metadata: dict = None,
+        package: bool = True,
+        intern: list = [],
+        extern: list = [],
+        mock: list = [],
+    ):
+        """Saves the model, either as a torch package, or just as
+        weights, alongside some specified metadata.
+
+        Parameters
+        ----------
+        path : str
+            Path to save model to.
+        metadata : dict, optional
+            Any metadata to save alongside the model,
+            by default None
+        package : bool, optional
+            Whether to use ``torch.package`` to save the model in
+            a format that is portable, by default True
+        intern : list, optional
+            List of additional libraries that are internal
+            to the model, used with torch.package, by default []
+        extern : list, optional
+            List of additional libraries that are external to
+            the model, used with torch.package, by default []
+        mock : list, optional
+            List of libraries to mock, used with torch.package,
+            by default []
+
+        Returns
+        -------
+        str
+            Path to saved model.
+        """
+        sig = inspect.signature(self.__class__)
+        args = {}
+
+        for key, val in sig.parameters.items():
+            arg_val = val.default
+            if arg_val is not inspect.Parameter.empty:
+                args[key] = arg_val
+
+        # Look up attibutes in self, and if any of them are in args,
+        # overwrite them in args.
+        for attribute in dir(self):
+            if attribute in args:
+                args[attribute] = getattr(self, attribute)
+
+        metadata = {} if metadata is None else metadata
+        metadata["kwargs"] = args
+        if not hasattr(self, "metadata"):
+            self.metadata = {}
+        self.metadata.update(metadata)
+
+        if not package:
+            state_dict = {"state_dict": self.state_dict(), "metadata": metadata}
+            torch.save(state_dict, path)
+        else:
+            self._save_package(path, intern=intern, extern=extern, mock=mock)
+
+        return path
+
+    @property
+    def device(self):
+        """Gets the device the model is on by looking at the device of
+        the first parameter. May not be valid if model is split across
+        multiple devices.
+        """
+        return list(self.parameters())[0].device
+
+    @classmethod
+    def load(
+        cls,
+        location: str,
+        *args,
+        package_name: str = None,
+        strict: bool = False,
+        **kwargs,
+    ):
+        """Load model from a path. Tries first to load as a package, and if
+        that fails, tries to load as weights. The arguments to the class are
+        specified inside the model weights file.
+
+        Parameters
+        ----------
+        location : str
+            Path to file.
+        package_name : str, optional
+            Name of package, by default ``cls.__name__``.
+        strict : bool, optional
+            Ignore unmatched keys, by default False
+        kwargs : dict
+            Additional keyword arguments to the model instantiation, if
+            not loading from package.
+
+        Returns
+        -------
+        BaseModel
+            A model that inherits from BaseModel.
+        """
+        try:
+            model = cls._load_package(location, package_name=package_name)
+        except:
+            model_dict = torch.load(location, "cpu")
+            metadata = model_dict["metadata"]
+            metadata["kwargs"].update(kwargs)
+
+            sig = inspect.signature(cls)
+            class_keys = list(sig.parameters.keys())
+            for k in list(metadata["kwargs"].keys()):
+                if k not in class_keys:
+                    metadata["kwargs"].pop(k)
+
+            model = cls(*args, **metadata["kwargs"])
+            model.load_state_dict(model_dict["state_dict"], strict=strict)
+            model.metadata = metadata
+
+        return model
+
+    def _save_package(self, path, intern=[], extern=[], mock=[], **kwargs):
+        package_name = type(self).__name__
+        resource_name = f"{type(self).__name__}.pth"
+
+        # Below is for loading and re-saving a package.
+        if hasattr(self, "importer"):
+            kwargs["importer"] = (self.importer, torch.package.sys_importer)
+            del self.importer
+
+        # Why do we use a tempfile, you ask?
+        # It's so we can load a packaged model and then re-save
+        # it to the same location. torch.package throws an
+        # error if it's loading and writing to the same
+        # file (this is undocumented).
+        with tempfile.NamedTemporaryFile(suffix=".pth") as f:
+            with torch.package.PackageExporter(f.name, **kwargs) as exp:
+                exp.intern(self.INTERN + intern)
+                exp.mock(mock)
+                exp.extern(self.EXTERN + extern)
+                exp.save_pickle(package_name, resource_name, self)
+
+                if hasattr(self, "metadata"):
+                    exp.save_pickle(
+                        package_name, f"{package_name}.metadata", self.metadata
+                    )
+
+            shutil.copyfile(f.name, path)
+
+        # Must reset the importer back to `self` if it existed
+        # so that you can save the model again!
+        if "importer" in kwargs:
+            self.importer = kwargs["importer"][0]
+        return path
+
+    @classmethod
+    def _load_package(cls, path, package_name=None):
+        package_name = cls.__name__ if package_name is None else package_name
+        resource_name = f"{package_name}.pth"
+
+        imp = torch.package.PackageImporter(path)
+        model = imp.load_pickle(package_name, resource_name, "cpu")
+        try:
+            model.metadata = imp.load_pickle(package_name, f"{package_name}.metadata")
+        except:  # pragma: no cover
+            pass
+        model.importer = imp
+
+        return model
+
+    def save_to_folder(
+        self,
+        folder: typing.Union[str, Path],
+        extra_data: dict = None,
+        package: bool = True,
+    ):
+        """Dumps a model into a folder, as both a package
+        and as weights, as well as anything specified in
+        ``extra_data``. ``extra_data`` is a dictionary of other
+        pickleable files, with the keys being the paths
+        to save them in. The model is saved under a subfolder
+        specified by the name of the class (e.g. ``folder/generator/[package, weights].pth``
+        if the model name was ``Generator``).
+
+        >>> with tempfile.TemporaryDirectory() as d:
+        >>>     extra_data = {
+        >>>         "optimizer.pth": optimizer.state_dict()
+        >>>     }
+        >>>     model.save_to_folder(d, extra_data)
+        >>>     Model.load_from_folder(d)
+
+        Parameters
+        ----------
+        folder : typing.Union[str, Path]
+            _description_
+        extra_data : dict, optional
+            _description_, by default None
+
+        Returns
+        -------
+        str
+            Path to folder
+        """
+        extra_data = {} if extra_data is None else extra_data
+        model_name = type(self).__name__.lower()
+        target_base = Path(f"{folder}/{model_name}/")
+        target_base.mkdir(exist_ok=True, parents=True)
+
+        if package:
+            package_path = target_base / f"package.pth"
+            self.save(package_path)
+
+        weights_path = target_base / f"weights.pth"
+        self.save(weights_path, package=False)
+
+        for path, obj in extra_data.items():
+            torch.save(obj, target_base / path)
+
+        return target_base
+
+    @classmethod
+    def load_from_folder(
+        cls,
+        folder: typing.Union[str, Path],
+        package: bool = True,
+        strict: bool = False,
+        **kwargs,
+    ):
+        """Loads the model from a folder generated by
+        :py:func:`audiotools.ml.layers.base.BaseModel.save_to_folder`.
+        Like that function, this one looks for a subfolder that has
+        the name of the class (e.g. ``folder/generator/[package, weights].pth`` if the
+        model name was ``Generator``).
+
+        Parameters
+        ----------
+        folder : typing.Union[str, Path]
+            _description_
+        package : bool, optional
+            Whether to use ``torch.package`` to load the model,
+            loading the model from ``package.pth``.
+        strict : bool, optional
+            Ignore unmatched keys, by default False
+
+        Returns
+        -------
+        tuple
+            tuple of model and extra data as saved by
+            :py:func:`audiotools.ml.layers.base.BaseModel.save_to_folder`.
+        """
+        folder = Path(folder) / cls.__name__.lower()
+        model_pth = "package.pth" if package else "weights.pth"
+        model_pth = folder / model_pth
+
+        model = cls.load(model_pth, strict=strict)
+        extra_data = {}
+        excluded = ["package.pth", "weights.pth"]
+        files = [x for x in folder.glob("*") if x.is_file() and x.name not in excluded]
+        for f in files:
+            extra_data[f.name] = torch.load(f, **kwargs)
+
+        return model, extra_data

xinference/thirdparty/audiotools/ml/layers/spectral_gate.py

@@ -0,0 +1,127 @@
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from ...core import AudioSignal
+from ...core import STFTParams
+from ...core import util
+
+
+class SpectralGate(nn.Module):
+    """Spectral gating algorithm for noise reduction,
+    as in Audacity/Ocenaudio. The steps are as follows:
+
+    1.  An FFT is calculated over the noise audio clip
+    2.  Statistics are calculated over FFT of the the noise
+        (in frequency)
+    3.  A threshold is calculated based upon the statistics
+        of the noise (and the desired sensitivity of the algorithm)
+    4.  An FFT is calculated over the signal
+    5.  A mask is determined by comparing the signal FFT to the
+        threshold
+    6.  The mask is smoothed with a filter over frequency and time
+    7.  The mask is appled to the FFT of the signal, and is inverted
+
+    Implementation inspired by Tim Sainburg's noisereduce:
+
+    https://timsainburg.com/noise-reduction-python.html
+
+    Parameters
+    ----------
+    n_freq : int, optional
+        Number of frequency bins to smooth by, by default 3
+    n_time : int, optional
+        Number of time bins to smooth by, by default 5
+    """
+
+    def __init__(self, n_freq: int = 3, n_time: int = 5):
+        super().__init__()
+
+        smoothing_filter = torch.outer(
+            torch.cat(
+                [
+                    torch.linspace(0, 1, n_freq + 2)[:-1],
+                    torch.linspace(1, 0, n_freq + 2),
+                ]
+            )[..., 1:-1],
+            torch.cat(
+                [
+                    torch.linspace(0, 1, n_time + 2)[:-1],
+                    torch.linspace(1, 0, n_time + 2),
+                ]
+            )[..., 1:-1],
+        )
+        smoothing_filter = smoothing_filter / smoothing_filter.sum()
+        smoothing_filter = smoothing_filter.unsqueeze(0).unsqueeze(0)
+        self.register_buffer("smoothing_filter", smoothing_filter)
+
+    def forward(
+        self,
+        audio_signal: AudioSignal,
+        nz_signal: AudioSignal,
+        denoise_amount: float = 1.0,
+        n_std: float = 3.0,
+        win_length: int = 2048,
+        hop_length: int = 512,
+    ):
+        """Perform noise reduction.
+
+        Parameters
+        ----------
+        audio_signal : AudioSignal
+            Audio signal that noise will be removed from.
+        nz_signal : AudioSignal, optional
+            Noise signal to compute noise statistics from.
+        denoise_amount : float, optional
+            Amount to denoise by, by default 1.0
+        n_std : float, optional
+            Number of standard deviations above which to consider
+            noise, by default 3.0
+        win_length : int, optional
+            Length of window for STFT, by default 2048
+        hop_length : int, optional
+            Hop length for STFT, by default 512
+
+        Returns
+        -------
+        AudioSignal
+            Denoised audio signal.
+        """
+        stft_params = STFTParams(win_length, hop_length, "sqrt_hann")
+
+        audio_signal = audio_signal.clone()
+        audio_signal.stft_data = None
+        audio_signal.stft_params = stft_params
+
+        nz_signal = nz_signal.clone()
+        nz_signal.stft_params = stft_params
+
+        nz_stft_db = 20 * nz_signal.magnitude.clamp(1e-4).log10()
+        nz_freq_mean = nz_stft_db.mean(keepdim=True, dim=-1)
+        nz_freq_std = nz_stft_db.std(keepdim=True, dim=-1)
+
+        nz_thresh = nz_freq_mean + nz_freq_std * n_std
+
+        stft_db = 20 * audio_signal.magnitude.clamp(1e-4).log10()
+        nb, nac, nf, nt = stft_db.shape
+        db_thresh = nz_thresh.expand(nb, nac, -1, nt)
+
+        stft_mask = (stft_db < db_thresh).float()
+        shape = stft_mask.shape
+
+        stft_mask = stft_mask.reshape(nb * nac, 1, nf, nt)
+        pad_tuple = (
+            self.smoothing_filter.shape[-2] // 2,
+            self.smoothing_filter.shape[-1] // 2,
+        )
+        stft_mask = F.conv2d(stft_mask, self.smoothing_filter, padding=pad_tuple)
+        stft_mask = stft_mask.reshape(*shape)
+        stft_mask *= util.ensure_tensor(denoise_amount, ndim=stft_mask.ndim).to(
+            audio_signal.device
+        )
+        stft_mask = 1 - stft_mask
+
+        audio_signal.stft_data *= stft_mask
+        audio_signal.istft()
+
+        return audio_signal

xinference/thirdparty/audiotools/post.py

@@ -0,0 +1,140 @@
+import tempfile
+import typing
+import zipfile
+from pathlib import Path
+
+import markdown2 as md
+import matplotlib.pyplot as plt
+import torch
+from IPython.display import HTML
+
+
+def audio_table(
+    audio_dict: dict,
+    first_column: str = None,
+    format_fn: typing.Callable = None,
+    **kwargs,
+):  # pragma: no cover
+    """Embeds an audio table into HTML, or as the output cell
+    in a notebook.
+
+    Parameters
+    ----------
+    audio_dict : dict
+        Dictionary of data to embed.
+    first_column : str, optional
+        The label for the first column of the table, by default None
+    format_fn : typing.Callable, optional
+        How to format the data, by default None
+
+    Returns
+    -------
+    str
+        Table as a string
+
+    Examples
+    --------
+
+    >>> audio_dict = {}
+    >>> for i in range(signal_batch.batch_size):
+    >>>     audio_dict[i] = {
+    >>>         "input": signal_batch[i],
+    >>>         "output": output_batch[i]
+    >>>     }
+    >>> audiotools.post.audio_zip(audio_dict)
+
+    """
+    from audiotools import AudioSignal
+
+    output = []
+    columns = None
+
+    def _default_format_fn(label, x, **kwargs):
+        if torch.is_tensor(x):
+            x = x.tolist()
+
+        if x is None:
+            return "."
+        elif isinstance(x, AudioSignal):
+            return x.embed(display=False, return_html=True, **kwargs)
+        else:
+            return str(x)
+
+    if format_fn is None:
+        format_fn = _default_format_fn
+
+    if first_column is None:
+        first_column = "."
+
+    for k, v in audio_dict.items():
+        if not isinstance(v, dict):
+            v = {"Audio": v}
+
+        v_keys = list(v.keys())
+        if columns is None:
+            columns = [first_column] + v_keys
+            output.append(" | ".join(columns))
+
+            layout = "|---" + len(v_keys) * "|:-:"
+            output.append(layout)
+
+        formatted_audio = []
+        for col in columns[1:]:
+            formatted_audio.append(format_fn(col, v[col], **kwargs))
+
+        row = f"| {k} | "
+        row += " | ".join(formatted_audio)
+        output.append(row)
+
+    output = "\n" + "\n".join(output)
+    return output
+
+
+def in_notebook():  # pragma: no cover
+    """Determines if code is running in a notebook.
+
+    Returns
+    -------
+    bool
+        Whether or not this is running in a notebook.
+    """
+    try:
+        from IPython import get_ipython
+
+        if "IPKernelApp" not in get_ipython().config:  # pragma: no cover
+            return False
+    except ImportError:
+        return False
+    except AttributeError:
+        return False
+    return True
+
+
+def disp(obj, **kwargs):  # pragma: no cover
+    """Displays an object, depending on if its in a notebook
+    or not.
+
+    Parameters
+    ----------
+    obj : typing.Any
+        Any object to display.
+
+    """
+    from audiotools import AudioSignal
+
+    IN_NOTEBOOK = in_notebook()
+
+    if isinstance(obj, AudioSignal):
+        audio_elem = obj.embed(display=False, return_html=True)
+        if IN_NOTEBOOK:
+            return HTML(audio_elem)
+        else:
+            print(audio_elem)
+    if isinstance(obj, dict):
+        table = audio_table(obj, **kwargs)
+        if IN_NOTEBOOK:
+            return HTML(md.markdown(table, extras=["tables"]))
+        else:
+            print(table)
+    if isinstance(obj, plt.Figure):
+        plt.show()