lt-tensor 0.0.1a22__py3-none-any.whl → 0.0.1a27__py3-none-any.whl

lt_tensor/processors/audio.py

@@ -1,10 +1,11 @@
-__all__ = ["AudioProcessor"]
+__all__ = ["AudioProcessor", "AudioProcessorConfig"]
 from lt_tensor.torch_commons import *
 from lt_utils.common import *
 import librosa
 import torchaudio
 import numpy as np
 from lt_tensor.model_base import Model
+from lt_utils.misc_utils import default
 from lt_utils.type_utils import is_file, is_array
 from lt_utils.file_ops import FileScan, get_file_name, path_to_str
 from torchaudio.functional import detect_pitch_frequency
@@ -12,8 +13,27 @@ import torch.nn.functional as F
 
 DEFAULT_DEVICE = torch.tensor([0]).device
 
+from lt_tensor.config_templates import ModelConfig
+
+
+class AudioProcessorConfig(ModelConfig):
+    sample_rate: int = 24000
+    n_mels: int = 80
+    n_fft: int = 1024
+    win_length: int = 1024
+    hop_length: int = 256
+    f_min: float = 0
+    f_max: float = 8000.0
+    center: bool = True
+    mel_scale: Literal["htk" "slaney"] = "htk"
+    std: int = 4
+    mean: int = -4
+    n_iter: int = 32
+    window: Optional[Tensor] = None
+    normalized: bool = False
+    onesided: Optional[bool] = None
+    n_stft: int = None
 
-class AudioProcessor(Model):
     def __init__(
         self,
         sample_rate: int = 24000,
@@ -21,63 +41,95 @@ class AudioProcessor(Model):
         n_fft: int = 1024,
         win_length: Optional[int] = None,
         hop_length: Optional[int] = None,
-        f_min: float = 0,
+        f_min: float = 1,
         f_max: float = 12000.0,
         center: bool = True,
         mel_scale: Literal["htk", "slaney"] = "htk",
         std: int = 4,
         mean: int = -4,
-        n_iter: int = 32,
-        window: Optional[Tensor] = None,
         normalized: bool = False,
         onesided: Optional[bool] = None,
+        *args,
+        **kwargs,
+    ):
+        settings = {
+            "sample_rate": sample_rate,
+            "n_mels": n_mels,
+            "n_fft": n_fft,
+            "win_length": win_length,
+            "hop_length": hop_length,
+            "f_min": f_min,
+            "f_max": f_max,
+            "center": center,
+            "mel_scale": mel_scale,
+            "std": std,
+            "mean": mean,
+            "normalized": normalized,
+            "onesided": onesided,
+        }
+        super().__init__(**settings)
+        self.post_process()
+
+    def post_process(self):
+        self.f_min = max(self.f_min, 1)
+        self.f_max = max(min(self.f_max, self.n_fft // 2), self.f_min + 1)
+        self.n_stft = self.n_fft // 2 + 1
+        self.hop_length = default(self.hop_length, self.n_fft // 4)
+        self.win_length = default(self.win_length, self.n_fft)
+
+
+def _comp_rms_helper(i: int, audio: Tensor, mel: Optional[Tensor]):
+    if mel is None:
+        return {"y": audio[i, :]}
+    return {"y": audio[i, :], "S": mel[i, :, :]}
+
+
+class AudioProcessor(Model):
+    def __init__(
+        self,
+        config: AudioProcessorConfig = AudioProcessorConfig(),
+        window: Optional[Tensor] = None,
     ):
         super().__init__()
-        self.mean = mean
-        self.std = std
-        self.n_mels = n_mels
-        self.n_fft = n_fft
-        self.n_stft = n_fft // 2 + 1
-        self.f_min = f_min
-        self.f_max = max(min(f_max, 12000), self.f_min + 1)
-        self.n_iter = n_iter
-        self.hop_length = hop_length or n_fft // 4
-        self.win_length = win_length or n_fft
-        self.sample_rate = sample_rate
-        self.center = center
+        self.cfg = config
         self._mel_spec = torchaudio.transforms.MelSpectrogram(
-            sample_rate=sample_rate,
-            n_mels=n_mels,
-            n_fft=n_fft,
-            win_length=win_length,
-            hop_length=hop_length,
-            center=center,
-            f_min=f_min,
-            f_max=f_max,
-            mel_scale=mel_scale,
+            sample_rate=self.cfg.sample_rate,
+            n_mels=self.cfg.n_mels,
+            n_fft=self.cfg.n_fft,
+            win_length=self.cfg.win_length,
+            hop_length=self.cfg.hop_length,
+            center=self.cfg.center,
+            f_min=self.cfg.f_min,
+            f_max=self.cfg.f_max,
+            mel_scale=self.cfg.mel_scale,
+            onesided=self.cfg.onesided,
+            normalized=self.cfg.normalized,
         )
+        self.griffin_lm_iters = 32
         self.mel_rscale = torchaudio.transforms.InverseMelScale(
-            n_stft=self.n_stft,
-            n_mels=n_mels,
-            sample_rate=sample_rate,
-            f_min=f_min,
-            f_max=f_max,
-            mel_scale=mel_scale,
+            n_stft=self.cfg.n_stft,
+            n_mels=self.cfg.n_mels,
+            sample_rate=self.cfg.sample_rate,
+            f_min=self.cfg.f_min,
+            f_max=self.cfg.f_max,
+            mel_scale=self.cfg.mel_scale,
         )
         self.giffin_lim = torchaudio.transforms.GriffinLim(
-            n_fft=n_fft,
-            n_iter=n_iter,
-            win_length=win_length,
-            hop_length=hop_length,
+            n_fft=self.cfg.n_fft,
+            win_length=self.cfg.win_length,
+            hop_length=self.cfg.hop_length,
         )
-        self.normalized = normalized
-        self.onesided = onesided
-
         self.register_buffer(
             "window",
-            (torch.hann_window(self.win_length) if window is None else window),
+            (torch.hann_window(self.cfg.win_length) if window is None else window),
         )
 
+    def _apply_device(self):
+        print(f"Audio Processor Device: {self.device.type}")
+        self.giffin_lim.to(device=self.device)
+        self._mel_spec.to(device=self.device)
+        self.mel_rscale.to(device=self.device)
+
     def from_numpy(
         self,
         array: np.ndarray,
@@ -85,8 +137,8 @@ class AudioProcessor(Model):
         dtype: Optional[torch.dtype] = None,
     ):
         converted = torch.from_numpy(array)
-        if not any([device is not None, dtype is not None]):
-            return converted
+        if device is None:
+            device = self.device
         return converted.to(device=device, dtype=dtype)
 
     def from_numpy_batch(
@@ -96,18 +148,32 @@ class AudioProcessor(Model):
         dtype: Optional[torch.dtype] = None,
     ):
         stacked = torch.stack([torch.from_numpy(x) for x in arrays])
-        if not any([device is not None, dtype is not None]):
-            return stacked
+        if device is None:
+            device = self.device
         return stacked.to(device=device, dtype=dtype)
 
-    def to_numpy_safe(self, tensor: Tensor):
+    def to_numpy_safe(self, tensor: Union[Tensor, np.ndarray]):
+        if isinstance(tensor, np.ndarray):
+            return tensor
         return tensor.detach().to(DEFAULT_DEVICE).numpy(force=True)
 
     def compute_rms(
         self,
-        audio: Union[Tensor, np.ndarray],
+        audio: Optional[Union[Tensor, np.ndarray]] = None,
         mel: Optional[Tensor] = None,
+        frame_length: Optional[int] = None,
+        hop_length: Optional[int] = None,
+        center: Optional[int] = None,
     ):
+        assert any([audio is not None, mel is not None])
+        rms_kwargs = dict(
+            frame_length=default(frame_length, self.cfg.n_fft),
+            hop_length=default(hop_length, self.cfg.hop_length),
+            center=default(center, self.cfg.center),
+        )
+
+        if audio is None and mel is not None:
+            return self.from_numpy(librosa.feature.rms(S=mel, **rms_kwargs)[0])
         default_dtype = audio.dtype
         default_device = audio.device
         if audio.ndim > 1:
@@ -123,34 +189,32 @@ class AudioProcessor(Model):
             else:
                 assert B == mel.shape[0], "Batch from mel and audio must be the same!"
             mel = self.to_numpy_safe(mel)
-            gt_mel = lambda idx: mel[idx, :, :]
-
-        else:
-            gt_mel = lambda idx: None
         audio = self.to_numpy_safe(audio)
         if B == 1:
-            _r = librosa.feature.rms(
-                y=audio, frame_length=self.n_fft, hop_length=self.hop_length
-            )[0]
-            rms = self.from_numpy(_r, default_device, default_dtype)
-
+            if mel is None:
+                return self.from_numpy(librosa.feature.rms(y=audio, **rms_kwargs)[0])
+            return self.from_numpy(librosa.feature.rms(y=audio, S=mel, **rms_kwargs)[0])
         else:
             rms_ = []
             for i in range(B):
-                _r = librosa.feature.rms(
-                    y=audio[i, :],
-                    S=gt_mel(i),
-                    frame_length=self.n_fft,
-                    hop_length=self.hop_length,
-                )[0]
+                _r = librosa.feature.rms(_comp_rms_helper(i, audio, mel), **rms_kwargs)[
+                    0
+                ]
                 rms_.append(_r)
-            rms = self.from_numpy_batch(rms_, default_device, default_dtype)
-
-        return rms
+            return self.from_numpy_batch(rms_, default_device, default_dtype)
 
     def compute_pitch(
         self,
         audio: Tensor,
+        *,
+        pad_mode: str = "constant",
+        trough_threshold: float = 0.1,
+        fmin: Optional[float] = None,
+        fmax: Optional[float] = None,
+        sr: Optional[float] = None,
+        frame_length: Optional[int] = None,
+        hop_length: Optional[int] = None,
+        center: Optional[bool] = None,
     ):
         default_dtype = audio.dtype
         default_device = audio.device
@@ -158,18 +222,25 @@ class AudioProcessor(Model):
             B = audio.shape[0]
         else:
             B = 1
-        fmin = max(self.f_min, 80)
+        sr = default(sr, self.cfg.sample_rate)
+        fmin = max(default(fmin, self.cfg.f_min), 65)
+        fmax = min(default(fmax, self.cfg.f_max), sr // 2)
+        frame_length = default(frame_length, self.cfg.n_fft)
+        hop_length = default(hop_length, self.cfg.hop_length)
+        center = default(center, self.cfg.center)
+        yn_kwargs = dict(
+            fmin=fmin,
+            fmax=fmax,
+            frame_length=frame_length,
+            sr=sr,
+            hop_length=hop_length,
+            center=center,
+            trough_threshold=trough_threshold,
+            pad_mode=pad_mode,
+        )
         if B == 1:
             f0 = self.from_numpy(
-                librosa.yin(
-                    self.to_numpy_safe(audio),
-                    fmin=fmin,
-                    fmax=min(self.f_max, self.sample_rate // 2 - 1),
-                    frame_length=self.n_fft,
-                    sr=self.sample_rate,
-                    hop_length=self.hop_length,
-                    center=self.center,
-                ),
+                librosa.yin(self.to_numpy_safe(audio), **yn_kwargs),
                 default_device,
                 default_dtype,
             )
@@ -177,32 +248,34 @@ class AudioProcessor(Model):
         else:
             f0_ = []
             for i in range(B):
-                r = librosa.yin(
-                    self.to_numpy_safe(audio[i, :]),
-                    fmin=fmin,
-                    fmax=self.f_max,
-                    frame_length=self.n_fft,
-                    sr=self.sample_rate,
-                    hop_length=self.hop_length,
-                    center=self.center,
-                )
-                f0_.append(r)
+                f0_.append(librosa.yin(self.to_numpy_safe(audio[i, :]), **yn_kwargs))
             f0 = self.from_numpy_batch(f0_, default_device, default_dtype)
+        return f0.squeeze()
 
-        # librosa.pyin(self.f_min, self.f_max)
-        return f0  # dict(f0=f0, attention_mask=f0 != f_max)
-
-    def compute_pitch_torch(self, audio: Tensor):
+    def compute_pitch_torch(
+        self,
+        audio: Tensor,
+        fmin: Optional[float] = None,
+        fmax: Optional[float] = None,
+        sr: Optional[float] = None,
+        win_length: Optional[Number] = None,
+        frame_length: Optional[Number] = None,
+    ):
+        sr = default(sr, self.sample_rate)
+        fmin = max(default(fmin, self.f_min), 1)
+        fmax = min(default(fmax, self.f_max), sr // 2)
+        win_length = default(win_length, self.win_length)
+        frame_length = default(frame_length, self.n_fft)
         return detect_pitch_frequency(
             audio,
-            sample_rate=self.sample_rate,
-            frame_time=self.n_fft,
-            win_length=self.win_length,
-            freq_low=max(self.f_min, 35),
-            freq_high=self.f_max,
+            sample_rate=sr,
+            frame_time=frame_length,
+            win_length=win_length,
+            freq_low=fmin,
+            freq_high=fmax,
         )
 
-    def interpolate_tensor(
+    def interpolate(
         self,
         tensor: Tensor,
         target_len: int,
@@ -227,6 +300,8 @@ class AudioProcessor(Model):
 
         if tensor.ndim == 2:  # [1, T]
             tensor = tensor.unsqueeze(1)  # [1, 1, T]
+        elif tensor.ndim == 1:
+            tensor = tensor.unsqueeze(0).unsqueeze(0)  # [1, 1, T]
         return F.interpolate(
             tensor,
             size=target_len,
@@ -248,18 +323,21 @@ class AudioProcessor(Model):
         *,
         _recall: bool = False,
     ):
+        if win_length is not None and win_length != self.cfg.win_length:
+            window = torch.hann_window(win_length, device=spec.device)
+        else:
+            window = self.window
+
         try:
             return torch.istft(
                 spec * torch.exp(phase * 1j),
-                n_fft=n_fft or self.n_fft,
-                hop_length=hop_length or self.hop_length,
-                win_length=win_length or self.win_length,
-                window=torch.hann_window(
-                    win_length or self.win_length, device=spec.device
-                ),
-                center=self.center,
-                normalized=self.normalized,
-                onesided=self.onesided,
+                n_fft=n_fft or self.cfg.n_fft,
+                hop_length=hop_length or self.cfg.hop_length,
+                win_length=win_length or self.cfg.win_length,
+                window=window,
+                center=self.cfg.center,
+                normalized=self.cfg.normalized,
+                onesided=self.cfg.onesided,
                 length=length,
                 return_complex=False,
             )
@@ -281,15 +359,15 @@ class AudioProcessor(Model):
         try:
             spectrogram = torch.stft(
                 input=wave,
-                n_fft=self.n_fft,
-                hop_length=self.hop_length,
-                win_length=self.win_length,
+                n_fft=self.cfg.n_fft,
+                hop_length=self.cfg.hop_length,
+                win_length=self.cfg.win_length,
                 window=self.window,
-                center=self.center,
+                center=self.cfg.center,
                 pad_mode="reflect",
-                normalized=self.normalized,
-                onesided=self.onesided,
-                return_complex=True,  # needed for the istft
+                normalized=self.cfg.normalized,
+                onesided=self.cfg.onesided,
+                return_complex=True,
             )
             return torch.istft(
                 spectrogram
@@ -298,14 +376,14 @@ class AudioProcessor(Model):
                     fill_value=1,
                     device=spectrogram.device,
                 ),
-                n_fft=self.n_fft,
-                hop_length=self.hop_length,
-                win_length=self.win_length,
+                n_fft=self.cfg.n_fft,
+                hop_length=self.cfg.hop_length,
+                win_length=self.cfg.win_length,
                 window=self.window,
                 length=length,
-                center=self.center,
-                normalized=self.normalized,
-                onesided=self.onesided,
+                center=self.cfg.center,
+                normalized=self.cfg.normalized,
+                onesided=self.cfg.onesided,
                 return_complex=False,
             )
         except RuntimeError as e:
@@ -317,21 +395,30 @@ class AudioProcessor(Model):
     def compute_mel(
         self,
         wave: Tensor,
-        base: float = 1e-5,
-        add_base: bool = True,
+        raw_mel_only: bool = False,
+        eps: float = 1e-5,
+        *,
+        _recall: bool = False,
     ) -> Tensor:
-        """Returns: [B, M, ML]"""
-        mel_tensor = self._mel_spec(wave.to(self.device))  # [M, ML]
-        if not add_base:
-            return (mel_tensor - self.mean) / self.std
-        return (
-            (torch.log(base + mel_tensor.unsqueeze(0)) - self.mean) / self.std
-        ).squeeze()
+        """Returns: [B, M, T]"""
+        try:
+            mel_tensor = self._mel_spec(wave.to(self.device))  # [M, T]
+            if not raw_mel_only:
+                mel_tensor = (
+                    torch.log(eps + mel_tensor.unsqueeze(0)) - self.cfg.mean
+                ) / self.cfg.std
+            return mel_tensor.squeeze()
+
+        except RuntimeError as e:
+            if not _recall:
+                self._mel_spec.to(self.device)
+                return self.compute_mel(wave, raw_mel_only, eps, _recall=True)
+            raise e
 
     def inverse_mel_spectogram(self, mel: Tensor, n_iter: Optional[int] = None):
-        if isinstance(n_iter, int) and n_iter != self.n_iter:
+        if isinstance(n_iter, int) and n_iter != self.griffin_lm_iters:
             self.giffin_lim.n_iter = n_iter
-            self.n_iter = n_iter
+            self.griffin_lm_iters = n_iter
         return self.giffin_lim.forward(
             self.mel_rscale(mel),
         )
@@ -341,18 +428,59 @@ class AudioProcessor(Model):
         path: PathLike,
         top_db: float = 30,
         normalize: bool = False,
-        alpha: float = 1.0,
+        *,
+        ref: float | Callable[[np.ndarray], Any] = np.max,
+        frame_length: int = 2048,
+        hop_length: int = 512,
+        mono: bool = True,
+        offset: float = 0.0,
+        duration: Optional[float] = None,
+        dtype: Any = np.float32,
+        res_type: str = "soxr_hq",
+        fix: bool = True,
+        scale: bool = False,
+        axis: int = -1,
+        norm: Optional[float] = np.inf,
+        norm_axis: int = 0,
+        norm_threshold: Optional[float] = None,
+        norm_fill: Optional[bool] = None,
     ) -> Tensor:
         is_file(path, True)
-        wave, sr = librosa.load(str(path), sr=self.sample_rate)
-        wave, _ = librosa.effects.trim(wave, top_db=top_db)
-        if sr != self.sample_rate:
-            wave = librosa.resample(wave, orig_sr=sr, target_sr=self.sample_rate)
+        wave, sr = librosa.load(
+            str(path),
+            sr=self.cfg.sample_rate,
+            mono=mono,
+            offset=offset,
+            duration=duration,
+            dtype=dtype,
+            res_type=res_type,
+        )
+        wave, _ = librosa.effects.trim(
+            wave,
+            top_db=top_db,
+            ref=ref,
+            frame_length=frame_length,
+            hop_length=hop_length,
+        )
+        if sr != self.cfg.sample_rate:
+            wave = librosa.resample(
+                wave,
+                orig_sr=sr,
+                target_sr=self.cfg.sample_rate,
+                res_type=res_type,
+                fix=fix,
+                scale=scale,
+                axis=axis,
+            )
         if normalize:
-            wave = librosa.util.normalize(wave)
-        if alpha not in [0.0, 1.0]:
-            wave = wave * alpha
-        return torch.from_numpy(wave).float().unsqueeze(0)
+            wave = librosa.util.normalize(
+                wave,
+                norm=norm,
+                axis=norm_axis,
+                threshold=norm_threshold,
+                fill=norm_fill,
+            )
+        return torch.from_numpy(wave).float().unsqueeze(0).to(self.device)
 
     def find_audios(
         self,
@@ -378,79 +506,18 @@ class AudioProcessor(Model):
             maximum,
         )
 
-    def find_audio_text_pairs(
-        self,
-        path,
-        additional_extensions: List[str] = [],
-        text_file_patterns: List[str] = [".normalized.txt", ".original.txt"],
-    ):
-        is_array(text_file_patterns, True, validate=True)  # Rases if empty or not valid
-        additional_extensions = [
-            x
-            for x in additional_extensions
-            if isinstance(x, str)
-            and "*" in x
-            and not any(list(map(lambda y: y in x), text_file_patterns))
-        ]
-        audio_files = self.find_audios(path, additional_extensions)
-        results = []
-        for audio in audio_files:
-            base_audio_dir = Path(audio).parent
-            audio_name = get_file_name(audio, False)
-            for pattern in text_file_patterns:
-                possible_txt_file = Path(base_audio_dir, audio_name + pattern)
-                if is_file(possible_txt_file):
-                    results.append((audio, path_to_str(possible_txt_file)))
-                    break
-        return results
-
     def stft_loss(
         self,
         signal: Tensor,
         ground: Tensor,
     ):
-        ground = F.interpolate(ground, signal.shape[-1]).to(signal.device)
-        if ground.ndim != signal.ndim:
-            assert ground.ndim in [1, 2, 3]
-            assert signal.ndim in [1, 2, 3]
-            if ground.ndim == 3:
-                ground = ground.squeeze(1)
-            elif ground.ndim == 1:
-                ground = ground.unsqueeze(0)
-            if signal.ndim == 3:
-                signal = signal.squeeze(1)
-            elif signal.ndim == 1:
-                signal = signal.unsqueeze(0)
+        with torch.no_grad():
+            ground = F.interpolate(ground, signal.shape[-1]).to(signal.device)
         return F.l1_loss(signal, ground)
 
-    @staticmethod
-    def plot_spectrogram(spectrogram, ax_):
-        import matplotlib.pylab as plt
-
-        fig, ax = plt.subplots(figsize=(10, 2))
-        im = ax.imshow(spectrogram, aspect="auto", origin="lower", interpolation="none")
-        plt.colorbar(im, ax=ax)
-
-        fig.canvas.draw()
-        plt.close()
-
-        return fig
-
     def forward(
         self,
         *inputs: Union[Tensor, float],
-        ap_task: Literal[
-            "get_mel", "get_loss", "inv_transform", "revert_mel"
-        ] = "get_mel",
         **inputs_kwargs,
     ):
-        if ap_task == "get_mel":
-            return self.compute_mel(*inputs, **inputs_kwargs)
-        elif ap_task == "get_loss":
-            return self.stft_loss(*inputs, **inputs_kwargs)
-        elif ap_task == "inv_transform":
-            return self.inverse_transform(*inputs, **inputs_kwargs)
-        elif ap_task == "revert_mel":
-            return self.inverse_mel_spectogram(*inputs, **inputs_kwargs)
-        else:
-            raise ValueError(f"Invalid task '{ap_task}'")
+        return self.compute_mel(*inputs, **inputs_kwargs)

{lt_tensor-0.0.1a22.dist-info → lt_tensor-0.0.1a27.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lt-tensor
-Version: 0.0.1a22
+Version: 0.0.1a27
 Summary: General utilities for PyTorch and others. Built for general use.
 Home-page: https://github.com/gr1336/lt-tensor/
 Author: gr1336