acids-transforms 1.0.0__py3-none-any.whl

acids_transforms/__init__.py

@@ -0,0 +1,3 @@
+from .utils import *
+from .transforms import *
+

acids_transforms/transforms/__init__.py

@@ -0,0 +1,9 @@
+from .base import *
+from .raw import *
+from .stft import *
+from .dgt import *
+from .norm import *
+from .spectral_repr import *
+from .mel import *
+from .misc import *
+from .oadd import *

acids_transforms/transforms/base.py

@@ -0,0 +1,232 @@
+import torch.nn as nn
+import functools
+import torch
+from typing import Union
+
+
+class NotInvertibleError(Exception):
+    pass
+
+
+InversionEnumType = Union[str, None]
+
+
+class AudioTransform(nn.Module):
+    # python-side descriptors, not compiled by torch.jit.script (e.g. tracable returns a global flag)
+    __jit_unused_properties__ = ["invertible", "scriptable", "tracable", "needs_scaling", "ratio"]
+    invertible = True
+    scriptable = False
+    tracable = False
+    needs_scaling = False
+
+    def __init__(self, sr=44100):
+        super(AudioTransform, self).__init__()
+        self.sr = sr
+
+    def __repr__(self):
+        return "AudioTransform()"
+
+    def __add__(self, transform):
+        if isinstance(transform, ComposeAudioTransform):
+            return ComposeAudioTransform(transforms=[self] + transform.transforms)
+        elif isinstance(transform, AudioTransform):
+            return ComposeAudioTransform(transforms=[self, transform])
+        else:
+            raise TypeError(
+                'AudioTransform cannot be added to type: %s' % type(transform))
+
+    @torch.jit.export
+    def scale_data(self, x: torch.Tensor) -> None:
+        pass
+
+    @torch.jit.export
+    def forward(self, x):
+        return x
+
+    def get_inversion_modes(self):
+        return None
+
+    @torch.jit.export
+    def invert(self, x: torch.Tensor) -> torch.Tensor:
+        return x
+
+    @torch.jit.export
+    def forward_with_time(self, x: torch.Tensor, time: torch.Tensor):
+        return self.forward(x), time
+
+    def realtime(self):
+        return self
+
+    @property
+    def ratio(self):
+        return 1
+
+    def test_forward(self, x: torch.Tensor, time: torch.Tensor = None):
+        if time is None:
+            return self.forward(x)
+        else:
+            return self.forward_with_time(x, time)
+
+    def test_inversion(self, x: torch.Tensor):
+        if not self.invertible:
+            raise NotImplementedError
+        x_transformed = self.forward(x)
+        x_inv = self.invert(x_transformed)
+        return {'inverted': x_inv}
+
+    @classmethod
+    def test_scripted_transform(cls, transform, batch_size=(2,2), invert=True):
+        x = torch.zeros(*batch_size, 44100)
+        time = torch.zeros(*batch_size)
+        x_t = transform.forward(x)
+        x_t, time_t = transform.forward_with_time(x, time)
+        if invert:
+            x_inv = transform.invert(x_t)
+
+
+class ComposeAudioTransform(AudioTransform):
+
+    @property
+    def invertible(self):
+        for t in self.transforms:
+            if not t.invertible:
+                return False
+        return True
+
+    @property
+    def needs_scaling(self):
+        for t in self.transforms:
+            if t.needs_scaling:
+                return True
+        return False
+
+    @property
+    def scriptable(self):
+        for t in self.transforms:
+            if not t.scriptable:
+                return False
+        return True
+
+    @property
+    def tracable(self):
+        for t in self.transforms:
+            if not t.tracable:
+                return False
+        return True
+
+    def __getitem__(self, item):
+        return self.transforms[item]
+
+    def __init__(self, transforms=[], sr=44100):
+        super().__init__()
+        self.transforms = nn.ModuleList(transforms)
+
+    def __repr__(self) -> str:
+        return "ComposeAudioTransform(%s)" % [t.__repr__()+"\n" for t in self.transforms]
+
+    def __add__(self, itm):
+        if not isinstance(itm, AudioTransform):
+            raise TypeError(
+                "ComposeAudioTransform can only be added to other AudioTransforms")
+        if isinstance(itm, ComposeAudioTransform):
+            return ComposeAudioTransform(list(self.transforms) + list(itm.transforms))
+        else:
+            return ComposeAudioTransform(list(self.transforms) + [itm])
+
+    def __radd__(self, other):
+        if not isinstance(other, AudioTransform):
+            raise TypeError(
+                "ComposeAudioTransform can only be added to other AudioTransforms")
+        if isinstance(other, ComposeAudioTransform):
+            return ComposeAudioTransform(list(other.transforms) + list(self.transforms))
+        else:
+            return ComposeAudioTransform([other] + list(self.transforms))
+
+    def realtime(self):
+        return ComposeAudioTransform(transforms=[t.realtime() for t in self.transforms], sr=self.sr)
+
+    @property
+    def ratio(self):
+        ratio = 1
+        for t in self.transforms:
+            ratio = ratio * t.ratio
+        return ratio
+
+    @torch.jit.export
+    def scale_data(self, x):
+        for t in self.transforms:
+            t.scale_data(x)
+            x = t(x)
+
+    @torch.jit.export
+    def forward(self, x: torch.Tensor):
+        for t in self.transforms:
+            x = t(x)
+        return x
+
+    @torch.jit.export
+    def forward_with_time(self, x: torch.Tensor, time: torch.Tensor):
+        for t in self.transforms:
+            x, time = t.forward_with_time(x, time)
+        return x, time
+
+    @torch.jit.export
+    def invert(self, x, inversion_mode: InversionEnumType = None):
+        for t in self.transforms[::-1]:
+            x = t.invert(x, inversion_mode=inversion_mode) 
+        return x
+
+    def get_inversion_modes(self, idx):
+        return type(self.transforms[idx]).get_inversion_modes()
+
+    def test_inversion(self, x: torch.Tensor):
+        if not self.invertible:
+            raise NotImplementedError
+        x_transformed = x
+        for t in self.transforms:
+            x_transformed = t.invert(x_transformed)
+        x_inv = x_transformed
+        for t in reversed(self.transforms):
+            x_inv = t.invert(x_inv)
+        return x_inv
+
+
+def apply_transform_to_list(transform, data, time=None, **kwargs):
+    if time is None:
+        outs = [transform(data[i], **kwargs) for i in range(len(data))]
+        return outs
+    else:
+        outs = [transform(data[i], time=time[i], **kwargs)
+                for i in range(len(data))]
+        return [o[0] for o in outs], [o[1] for o in outs]
+
+
+def apply_invert_transform_to_list(transform, data, time=None, **kwargs):
+    if time is None:
+        outs = [transform.invert(data[i], **kwargs) for i in range(len(data))]
+        return outs
+    else:
+        outs = [transform.invert(data[i], time=time[i], **kwargs)
+                for i in range(len(data))]
+        return [o[0] for o in outs], [o[1] for o in outs]
+
+
+_CanTrace = False
+make_fx = None
+try: 
+    from torch.fx.experimental.proxy_tensor import make_fx as tft_make_fx
+
+    def make_fx(obj, *args, _fn="forward", **kwargs):
+        if _fn == "forward": 
+            return tft_make_fx(obj, *args, **kwargs)
+        else:
+            # @functools.wraps(getattr(obj, _fn))
+            def _closure(*args, **kwargs):
+                return getattr(obj, _fn)(*args, **kwargs)
+            return tft_make_fx(_closure, *args, **kwargs)
+
+    # tracing relies on symbolic make_fx, custom operators (torch >= 2.4, see utils.tracing) and fake fft kernels
+    # (aten._fft_r2c, torch >= 2.7) : transforms are not tracable before torch 2.7.
+    _CanTrace = tuple(int(v) for v in torch.__version__.split("+")[0].split(".")[:2]) >= (2, 7)
+except ModuleNotFoundError:
+    pass