lt-tensor 0.0.1a18__py3-none-any.whl → 0.0.1a20__py3-none-any.whl

lt_tensor/__init__.py

@@ -11,7 +11,6 @@ from . import (
     noise_tools,
     losses,
     processors,
-    datasets,
     torch_commons,
 )
 
@@ -26,6 +25,5 @@ __all__ = [
     "noise_tools",
     "losses",
     "processors",
-    "datasets",
     "torch_commons",
 ]

lt_tensor/config_templates.py

@@ -7,75 +7,58 @@ from lt_tensor.misc_utils import updateDict
 
 class ModelConfig(ABC, OrderedDict):
     _default_settings: Dict[str, Any] = {}
-    _forbidden_list: List[str] = ["_default_settings", "_forbidden_list" "path_name"]
+    _forbidden_list: List[str] = [
+        "_default_settings",
+        "_forbidden_list",
+    ]
 
     def __init__(
         self,
-        settings: Dict[str, Any] = {},
-        path_name: Optional[Union[str, PathLike]] = None,
+        **settings,
     ):
-        assert is_dict(settings, False)
         self._default_settings = settings
-        if path_name is not None and is_pathlike(path_name):
-            if not str(path_name).endswith(".json"):
-                self.path_name = str(Path(path_name, "config.json")).replace("\\", "/")
-            else:
-                self.path_name = str(path_name).replace("\\", "/")
-        else:
-            self.path_name = "config.json"
-        self.reset_settings()
-
-    def _setup_path_name(self, path_name: Union[str, PathLike]):
-        if is_file(path_name):
-            self.from_path(path_name)
-            self.path_name = str(path_name).replace("\\", "/")
-        elif is_str(path_name):
-            self.path_name = str(path_name).replace("\\", "/")
-            if not self.path_name.endswith((".json")):
-                self.path_name += ".json"
+        self.set_state_dict(self._default_settings)
 
     def reset_settings(self):
-        dk_keys = self.__dict__.keys()
-        for s_name, setting in self._default_settings.items():
-            if s_name in self._forbidden_list or s_name not in dk_keys:
-                continue
-            updateDict(self, {s_name: setting})
+        raise NotImplementedError("Not implemented")
 
     def save_config(
         self,
-        path_name: Optional[Union[PathLike, str]] = None,
+        path: str,
     ):
-        if not is_pathlike(path_name, True):
-            assert (
-                path_name is None
-            ), f"path_name should be a non-empty string or pathlike object! received instead: {path_name}."
-            path_name = self.path_name
-        else:
-            self._setup_path_name(path_name)
+        base = {
+            k: v
+            for k, v in self.state_dict().items()
+            if isinstance(v, (str, int, float, list, tuple, dict, set, bytes))
+        }
 
-        base = self.get_state_dict()
-        save_json(self.path_name, base, indent=2)
+        save_json(path, base, indent=4)
 
     def set_value(self, var_name: str, value: str) -> None:
-        assert var_name in self.__dict__, "Value not registered!"
         assert var_name not in self._forbidden_list, "Not allowed!"
         updateDict(self, {var_name: value})
+        self.update({var_name: value})
 
     def get_value(self, var_name: str) -> Any:
         return self.__dict__.get(var_name)
 
-    def __getattribute__(self, name):
-        return self.__dict__.get(name)
+    def set_state_dict(self, new_state: dict[str, str]):
+        new_state = {
+            k: y for k, y in new_state.items() if k not in self._forbidden_list
+        }
+        updateDict(self, new_state)
+        self.update(**new_state)
 
-    def get_state_dict(self):
+    def state_dict(self):
         return {k: y for k, y in self.__dict__.items() if k not in self._forbidden_list}
 
     @classmethod
     def from_dict(
-        cls, dictionary: Dict[str, Any], path: Optional[Union[str, PathLike]] = None
+        cls,
+        dictionary: Dict[str, Any],
     ) -> "ModelConfig":
         assert is_dict(dictionary)
-        return ModelConfig(dictionary, path)
+        return ModelConfig(**dictionary)
 
     @classmethod
     def from_path(cls, path_name: PathLike) -> "ModelConfig":
@@ -102,4 +85,7 @@ class ModelConfig(ABC, OrderedDict):
             )
             assert files, "No config file found in the provided directory!"
             settings.update(load_json(files[-1], {}, errors="ignore"))
-        return ModelConfig(settings, path_name)
+        settings.pop("path", None)
+        settings.pop("path_name", None)
+
+        return ModelConfig(**settings)

lt_tensor/model_zoo/audio_models/diffwave/__init__.py

@@ -1,9 +1,8 @@
 __all__ = ["DiffWave", "DiffWaveConfig", "SpectrogramUpsample", "DiffusionEmbedding"]
 
 import numpy as np
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
+from lt_tensor.torch_commons import *
+from torch.nn import functional as F
 from lt_tensor.config_templates import ModelConfig
 from lt_tensor.torch_commons import *
 from lt_tensor.model_base import Model
@@ -12,16 +11,9 @@ from lt_utils.common import *
 
 
 class DiffWaveConfig(ModelConfig):
-    # Training params
-    batch_size = 16
-    learning_rate = 2e-4
-    max_grad_norm = None
-    # Data params
-    sample_rate = 24000
+    # Model params
     n_mels = 80
-    n_fft = 1024
     hop_samples = 256
-    # Model params
     residual_layers = 30
     residual_channels = 64
     dilation_cycle_length = 10
@@ -35,10 +27,36 @@ class DiffWaveConfig(ModelConfig):
 
     def __init__(
         self,
-        settings: Dict[str, Any] = {},
-        path_name: Optional[Union[str, PathLike]] = None,
+        n_mels=80,
+        hop_samples=256,
+        residual_layers=30,
+        residual_channels=64,
+        dilation_cycle_length=10,
+        unconditional=False,
+        noise_schedule: list[int] = np.linspace(1e-4, 0.05, 50).tolist(),
+        interpolate_cond=False,
+        interpolation_mode: Literal[
+            "nearest",
+            "linear",
+            "bilinear",
+            "bicubic",
+            "trilinear",
+            "area",
+            "nearest-exact",
+        ] = "nearest",
     ):
-        super().__init__(settings, path_name)
+        settings = {
+            "n_mels": n_mels,
+            "hop_samples": hop_samples,
+            "residual_layers": residual_layers,
+            "dilation_cycle_length": dilation_cycle_length,
+            "residual_channels": residual_channels,
+            "unconditional": unconditional,
+            "noise_schedule": noise_schedule,
+            "interpolate": interpolate_cond,
+            "interpolation_mode": interpolation_mode,
+        }
+        super().__init__(**settings)
 
 
 def Conv1d(*args, **kwargs):

lt_tensor/model_zoo/audio_models/hifigan/__init__.py

@@ -4,10 +4,6 @@ from lt_tensor.torch_commons import *
 from lt_tensor.model_zoo.residual import ConvNets
 from torch.nn import functional as F
 
-import torch
-import torch.nn.functional as F
-import torch.nn as nn
-
 
 def get_padding(kernel_size, dilation=1):
     return int((kernel_size * dilation - dilation) / 2)
@@ -21,7 +17,7 @@ class HifiganConfig(ModelConfig):
     in_channels: int = 80
     upsample_rates: List[Union[int, List[int]]] = [8, 8]
     upsample_kernel_sizes: List[Union[int, List[int]]] = [16, 16]
-    upsample_initial_channel: int = (512,)
+    upsample_initial_channel: int = 512
     resblock_kernel_sizes: List[Union[int, List[int]]] = [3, 7, 11]
     resblock_dilation_sizes: List[Union[int, List[int]]] = [
         [1, 3, 5],
@@ -34,10 +30,32 @@ class HifiganConfig(ModelConfig):
 
     def __init__(
         self,
-        settings: Dict[str, Any] = {},
-        path_name: Optional[Union[str, PathLike]] = None,
+        in_channels: int = 80,
+        upsample_rates: List[Union[int, List[int]]] = [8, 8],
+        upsample_kernel_sizes: List[Union[int, List[int]]] = [16, 16],
+        upsample_initial_channel: int = 512,
+        resblock_kernel_sizes: List[Union[int, List[int]]] = [3, 7, 11],
+        resblock_dilation_sizes: List[Union[int, List[int]]] = [
+            [1, 3, 5],
+            [1, 3, 5],
+            [1, 3, 5],
+        ],
+        activation: nn.Module = nn.LeakyReLU(0.1),
+        resblock: int = 0,
+        *args,
+        **kwargs,
     ):
-        super().__init__(settings, path_name)
+        settings = {
+            "in_channels": in_channels,
+            "upsample_rates": upsample_rates,
+            "upsample_kernel_sizes": upsample_kernel_sizes,
+            "upsample_initial_channel": upsample_initial_channel,
+            "resblock_kernel_sizes": resblock_kernel_sizes,
+            "resblock_dilation_sizes": resblock_dilation_sizes,
+            "activation": activation,
+            "resblock": resblock,
+        }
+        super().__init__(**settings)
 
 
 class ResBlock1(ConvNets):
@@ -177,10 +195,10 @@ class HifiganGenerator(ConvNets):
         self.conv_pre = weight_norm(
             nn.Conv1d(cfg.in_channels, cfg.upsample_initial_channel, 7, 1, padding=3)
         )
-        resblock = ResBlock1 if resblock == 0 else ResBlock2
+        resblock = ResBlock1 if cfg.resblock == 0 else ResBlock2
         self.activation = cfg.activation
         self.ups = nn.ModuleList()
-        for i, (u, k) in enumerate(zip(cfg.psample_rates, cfg.upsample_kernel_sizes)):
+        for i, (u, k) in enumerate(zip(cfg.upsample_rates, cfg.upsample_kernel_sizes)):
             self.ups.append(
                 weight_norm(
                     nn.ConvTranspose1d(

lt_tensor/model_zoo/audio_models/istft/__init__.py

@@ -11,7 +11,7 @@ class iSTFTNetConfig(ModelConfig):
     in_channels: int = 80
     upsample_rates: List[Union[int, List[int]]] = [8, 8]
     upsample_kernel_sizes: List[Union[int, List[int]]] = [16, 16]
-    upsample_initial_channel: int = (512,)
+    upsample_initial_channel: int = 512
     resblock_kernel_sizes: List[Union[int, List[int]]] = [3, 7, 11]
     resblock_dilation_sizes: List[Union[int, List[int]]] = [
         [1, 3, 5],
@@ -26,10 +26,36 @@ class iSTFTNetConfig(ModelConfig):
 
     def __init__(
         self,
-        settings: Dict[str, Any] = {},
-        path_name: Optional[Union[str, PathLike]] = None,
+        in_channels: int = 80,
+        upsample_rates: List[Union[int, List[int]]] = [8, 8],
+        upsample_kernel_sizes: List[Union[int, List[int]]] = [16, 16],
+        upsample_initial_channel: int = 512,
+        resblock_kernel_sizes: List[Union[int, List[int]]] = [3, 7, 11],
+        resblock_dilation_sizes: List[Union[int, List[int]]] = [
+            [1, 3, 5],
+            [1, 3, 5],
+            [1, 3, 5],
+        ],
+        activation: nn.Module = nn.LeakyReLU(0.1),
+        resblock: int = 0,
+        gen_istft_n_fft: int = 16,
+        sampling_rate: Number = 24000,
+        *args,
+        **kwargs,
     ):
-        super().__init__(settings, path_name)
+        settings = {
+            "in_channels": in_channels,
+            "upsample_rates": upsample_rates,
+            "upsample_kernel_sizes": upsample_kernel_sizes,
+            "upsample_initial_channel": upsample_initial_channel,
+            "resblock_kernel_sizes": resblock_kernel_sizes,
+            "resblock_dilation_sizes": resblock_dilation_sizes,
+            "activation": activation,
+            "resblock": resblock,
+            "gen_istft_n_fft": gen_istft_n_fft,
+            "sampling_rate": sampling_rate,
+        }
+        super().__init__(**settings)
 
 
 def get_padding(ks, d):
@@ -169,7 +195,7 @@ class iSTFTNetGenerator(ConvNets):
         self.conv_pre = weight_norm(
             nn.Conv1d(cfg.in_channels, cfg.upsample_initial_channel, 7, 1, padding=3)
         )
-        resblock = ResBlock1 if resblock == 0 else ResBlock2
+        resblock = ResBlock1 if cfg.resblock == 0 else ResBlock2
 
         self.ups = nn.ModuleList()
         for i, (u, k) in enumerate(zip(cfg.upsample_rates, cfg.upsample_kernel_sizes)):

{lt_tensor-0.0.1a18.dist-info → lt_tensor-0.0.1a20.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lt-tensor
-Version: 0.0.1a18
+Version: 0.0.1a20
 Summary: General utilities for PyTorch and others. Built for general use.
 Home-page: https://github.com/gr1336/lt-tensor/
 Author: gr1336
@@ -17,7 +17,7 @@ Requires-Dist: numpy>=1.26.4
 Requires-Dist: tokenizers
 Requires-Dist: pyyaml>=6.0.0
 Requires-Dist: numba>0.60.0
-Requires-Dist: lt-utils==0.0.2a2
+Requires-Dist: lt-utils==0.0.2a3
 Requires-Dist: librosa==0.11.*
 Requires-Dist: einops
 Requires-Dist: plotly

{lt_tensor-0.0.1a18.dist-info → lt_tensor-0.0.1a20.dist-info}/RECORD

@@ -1,5 +1,5 @@
-lt_tensor/__init__.py,sha256=XxNCGcVL-haJyMpifr-GRaamo32R6jmqe3iOuS4ecfs,469
-lt_tensor/config_templates.py,sha256=xWZhktYVlkwvJVreqyACpWo-lJ5htG9vTZyqZ6OexzA,3899
+lt_tensor/__init__.py,sha256=8FTxpJ6td2bMr_GqzW2tCV6Tr5CelbQle8N5JRWtx8M,439
+lt_tensor/config_templates.py,sha256=RP7EFVRj6mRUj6xDLe7FMXgN5TIo8_o9h1Kb8epdmfo,2825
 lt_tensor/losses.py,sha256=zvkCOnE5XpF3v6ymivRIdqPTsMM5zc94ZMom7YDi3zM,4946
 lt_tensor/lr_schedulers.py,sha256=LSZzqrOOLzSthD8k-W4cYPJt0vCjmHkiJkLr5e3yRTE,3659
 lt_tensor/math_ops.py,sha256=TkD4WQG42KsQ9Fg7FXOjf8f-ixtW0apf2XjaooecVx4,2257
@@ -9,8 +9,6 @@ lt_tensor/monotonic_align.py,sha256=LhBd8p1xdBzg6jQrQX1j7b4PNeYGwIqM24zcU-pHOLE,
 lt_tensor/noise_tools.py,sha256=wFeAsHhLhSlEc5XU5LbFKaXoHeVxrWjiMeljjGdIKyM,11363
 lt_tensor/torch_commons.py,sha256=fntsEU8lhBQo0ebonI1iXBkMbWMN3HpBsG13EWlP5s8,718
 lt_tensor/transform.py,sha256=dZm8T_ov0blHMQu6nGiehsdG1VSB7bZBUVmTkT-PBdc,13257
-lt_tensor/datasets/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-lt_tensor/datasets/audio.py,sha256=5Wvz1BJ7xXkLYpVLLw9RY3X3RgMdPPeGiN0-MmJDQy0,8045
 lt_tensor/model_zoo/__init__.py,sha256=ltVTvmOlbOCfDc5Trvg0-Ta_Ujgkw0UVF9V5rqHx-RI,378
 lt_tensor/model_zoo/basic.py,sha256=pI8HyiHK-cmWcEEaVY_EduUJOjZW6HOtXvJd8Rbhq30,15452
 lt_tensor/model_zoo/features.py,sha256=DO8dlE0kmPKTNC1Xkv9wKegOOYkQa_rkxM4hhcNwJWA,15655
@@ -19,13 +17,13 @@ lt_tensor/model_zoo/pos_encoder.py,sha256=3d1EYLinCU9UAy-WuEWeYMGhMqaGknCiQ5qEmh
 lt_tensor/model_zoo/residual.py,sha256=i5V4ju7DB3WesKBVm6KH_LyPoKGDUOyo2Usfs-PyP58,9394
 lt_tensor/model_zoo/transformer.py,sha256=HUFoFFh7EQJErxdd9XIxhssdjvNVx2tNGDJOTUfwG2A,4301
 lt_tensor/model_zoo/audio_models/__init__.py,sha256=MoG9YjxLyvscq_6njK1ljGBletK9iedBXt66bplzW-s,83
-lt_tensor/model_zoo/audio_models/diffwave/__init__.py,sha256=R14hY-nCbCO-T3ox9f4MXCPgQQogFUKAJ2WtntLz09w,7393
-lt_tensor/model_zoo/audio_models/hifigan/__init__.py,sha256=6ZGYyNiTMGHnOjGU0gq_TSM8Y9LtYlP3neGwa01Ghyk,13135
-lt_tensor/model_zoo/audio_models/istft/__init__.py,sha256=noi4GLGZQ_qg5H-ipe5d7j8rvt4Hic_sXiME-TE-B2c,13783
+lt_tensor/model_zoo/audio_models/diffwave/__init__.py,sha256=OUyh421xRCcxOMi_Ek6Ak3-FPe1k6WTDQ-6gd6OjaCU,8091
+lt_tensor/model_zoo/audio_models/hifigan/__init__.py,sha256=JNebaYO3nsyyqpYCCOyL13zY2uxLY3NOCeNynF6-96k,13940
+lt_tensor/model_zoo/audio_models/istft/__init__.py,sha256=JdFChpPhURaI2qb9mDV6vzDcZN757FBGGtgzN3vxtJ0,14821
 lt_tensor/processors/__init__.py,sha256=4b9MxAJolXiJfSm20ZEspQTDm1tgLazwlPWA_jB1yLM,63
 lt_tensor/processors/audio.py,sha256=SMqNSl4Den-x1awTCQ8-TcR-0jPiv5lDaUpU93SRRaw,14749
-lt_tensor-0.0.1a18.dist-info/licenses/LICENSE,sha256=HUnu_iSPpnDfZS_PINhO3AoVizJD1A2vee8WX7D7uXo,11358
-lt_tensor-0.0.1a18.dist-info/METADATA,sha256=fgRzOiw5tMmkaEY9HrGEKNL2v9mN5JVbf9r-bf18Am0,1033
-lt_tensor-0.0.1a18.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-lt_tensor-0.0.1a18.dist-info/top_level.txt,sha256=35FuhFeXnUyvHWdbVHGPh0hS8euofafnJ_GJAVSF4Kk,10
-lt_tensor-0.0.1a18.dist-info/RECORD,,
+lt_tensor-0.0.1a20.dist-info/licenses/LICENSE,sha256=HUnu_iSPpnDfZS_PINhO3AoVizJD1A2vee8WX7D7uXo,11358
+lt_tensor-0.0.1a20.dist-info/METADATA,sha256=y-nfT5fMD5c8JgD299060Aij5Ugz2xHlDYYNoPXtAy8,1033
+lt_tensor-0.0.1a20.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+lt_tensor-0.0.1a20.dist-info/top_level.txt,sha256=35FuhFeXnUyvHWdbVHGPh0hS8euofafnJ_GJAVSF4Kk,10
+lt_tensor-0.0.1a20.dist-info/RECORD,,

lt_tensor/datasets/audio.py

@@ -1,235 +0,0 @@
-__all__ = ["WaveMelDataset"]
-from lt_tensor.torch_commons import *
-from lt_utils.common import *
-from lt_utils.misc_utils import default
-import random
-from torch.utils.data import Dataset, DataLoader, Sampler
-from lt_tensor.processors import AudioProcessor
-import torch.nn.functional as FT
-from lt_tensor.misc_utils import log_tensor
-from tqdm import tqdm
-
-
-DEFAULT_DEVICE = torch.tensor([0]).device
-
-
-class WaveMelDataset(Dataset):
-    cached_data: Union[list[dict[str, Tensor]], Tuple[Tensor, Tensor]] = []
-    loaded_files: Dict[str, List[Dict[str, Tensor]]] = {}
-    normalize_waves: bool = False
-    randomize_ranges: bool = False
-    alpha_wv: float = 1.0
-    limit_files: Optional[int] = None
-    min_frame_length: Optional[int] = None
-    max_frame_length: Optional[int] = None
-
-    def __init__(
-        self,
-        audio_processor: AudioProcessor,
-        dataset_path: PathLike,
-        limit_files: Optional[int] = None,
-        min_frame_length: Optional[int] = None,
-        max_frame_length: Optional[int] = None,
-        randomize_ranges: Optional[bool] = None,
-        pre_load: bool = False,
-        normalize_waves: Optional[bool] = None,
-        alpha_wv: Optional[float] = None,
-        lib_norm: bool = True,
-    ):
-        super().__init__()
-        assert max_frame_length is None or max_frame_length >= (
-            (audio_processor.n_fft // 2) + 1
-        )
-        self.ap = audio_processor
-        self.dataset_path = dataset_path
-        if limit_files:
-            self.limit_files = limit_files
-        if normalize_waves is not None:
-            self.normalize_waves = normalize_waves
-        if alpha_wv is not None:
-            self.alpha_wv = alpha_wv
-        if pre_load is not None:
-            self.pre_loaded = pre_load
-        if randomize_ranges is not None:
-            self.randomize_ranges = randomize_ranges
-
-        self.post_n_fft = (audio_processor.n_fft // 2) + 1
-        self.lib_norm = lib_norm
-        if max_frame_length is not None:
-            max_frame_length = max(self.post_n_fft + 1, max_frame_length)
-            self.r_range = max(self.post_n_fft + 1, max_frame_length // 3)
-            self.max_frame_length = max_frame_length
-        if min_frame_length is not None:
-            self.min_frame_length = max(
-                self.post_n_fft + 1, min(min_frame_length, max_frame_length)
-            )
-
-        self.files = self.ap.find_audios(dataset_path, maximum=None)
-        if limit_files:
-            random.shuffle(self.files)
-            self.files = self.files[-self.limit_files :]
-        if pre_load:
-            for file in tqdm(self.files, "Loading files"):
-                results = self.load_data(file)
-                if not results:
-                    continue
-                self.cached_data.extend(results)
-
-    def renew_dataset(self, new_path: Optional[PathLike] = None):
-        new_path = default(new_path, self.dataset_path)
-        self.files = self.ap.find_audios(new_path, maximum=None)
-        random.shuffle(self.files)
-        for file in tqdm(self.files, "Loading files"):
-            results = self.load_data(file)
-            if not results:
-                continue
-            self.cached_data.extend(results)
-
-    def _add_dict(
-        self,
-        audio_wave: Tensor,
-        audio_mel: Tensor,
-        pitch: Tensor,
-        rms: Tensor,
-        file: PathLike,
-    ):
-        return {
-            "wave": audio_wave,
-            "pitch": pitch,
-            "rms": rms,
-            "mel": audio_mel,
-            "file": file,
-        }
-
-    def load_data(self, file: PathLike):
-        initial_audio = self.ap.load_audio(
-            file, normalize=self.lib_norm, alpha=self.alpha_wv
-        )
-        if self.normalize_waves:
-            initial_audio = self.ap.normalize_audio(initial_audio)
-        if initial_audio.shape[-1] < self.post_n_fft:
-            return None
-
-        if self.min_frame_length is not None:
-            if self.min_frame_length > initial_audio.shape[-1]:
-                return None
-        if (
-            not self.max_frame_length
-            or initial_audio.shape[-1] <= self.max_frame_length
-        ):
-
-            audio_rms = self.ap.compute_rms(initial_audio)
-            audio_pitch = self.ap.compute_pitch(initial_audio)
-            audio_mel = self.ap.compute_mel(initial_audio, add_base=True)
-
-            return [
-                self._add_dict(initial_audio, audio_mel, audio_pitch, audio_rms, file)
-            ]
-        results = []
-
-        if self.randomize_ranges:
-            frame_limit = random.randint(self.r_range, self.max_frame_length)
-        else:
-            frame_limit = self.max_frame_length
-
-        fragments = list(
-            torch.split(initial_audio, split_size_or_sections=frame_limit, dim=-1)
-        )
-        random.shuffle(fragments)
-        for fragment in fragments:
-            if fragment.shape[-1] < self.post_n_fft:
-                # Too small
-                continue
-            if (
-                self.min_frame_length is not None
-                and self.min_frame_length > fragment.shape[-1]
-            ):
-                continue
-
-            audio_rms = self.ap.compute_rms(fragment)
-            audio_pitch = self.ap.compute_pitch(fragment)
-            audio_mel = self.ap.compute_mel(fragment, add_base=True)
-            results.append(
-                self._add_dict(fragment, audio_mel, audio_pitch, audio_rms, file)
-            )
-        return results
-
-    def get_data_loader(
-        self,
-        batch_size: int = 1,
-        shuffle: Optional[bool] = None,
-        sampler: Optional[Union[Sampler, Iterable]] = None,
-        batch_sampler: Optional[Union[Sampler[list], Iterable[list]]] = None,
-        num_workers: int = 0,
-        pin_memory: bool = False,
-        drop_last: bool = False,
-        timeout: float = 0,
-    ):
-        return DataLoader(
-            self,
-            batch_size=batch_size,
-            shuffle=shuffle,
-            sampler=sampler,
-            batch_sampler=batch_sampler,
-            num_workers=num_workers,
-            pin_memory=pin_memory,
-            drop_last=drop_last,
-            timeout=timeout,
-            collate_fn=self.collate_fn,
-        )
-
-    def collate_fn(self, batch: Sequence[Dict[str, Tensor]]):
-        mel = []
-        wave = []
-        file = []
-        rms = []
-        pitch = []
-        for x in batch:
-            mel.append(x["mel"])
-            wave.append(x["wave"])
-            file.append(x["file"])
-            rms.append(x["rms"])
-            pitch.append(x["pitch"])
-        # Find max time in mel (dim -1), and max audio length
-        max_mel_len = max([m.shape[-1] for m in mel])
-        max_audio_len = max([a.shape[-1] for a in wave])
-        max_pitch_len = max([a.shape[-1] for a in pitch])
-        max_rms_len = max([a.shape[-1] for a in rms])
-
-        padded_mel = torch.stack(
-            [FT.pad(m, (0, max_mel_len - m.shape[-1])) for m in mel]
-        )  # shape: [B, 80, T_max]
-
-        padded_wave = torch.stack(
-            [FT.pad(a, (0, max_audio_len - a.shape[-1])) for a in wave]
-        )  # shape: [B, L_max]
-
-        padded_pitch = torch.stack(
-            [FT.pad(a, (0, max_pitch_len - a.shape[-1])) for a in pitch]
-        )  # shape: [B, L_max]
-        padded_rms = torch.stack(
-            [FT.pad(a, (0, max_rms_len - a.shape[-1])) for a in rms]
-        )  # shape: [B, L_max]
-        return dict(
-            mel=padded_mel,
-            wave=padded_wave,
-            pitch=padded_pitch,
-            rms=padded_rms,
-            file=file,
-        )
-
-    def get_item(self, idx: int):
-        if self.pre_loaded:
-            return self.cached_data[idx]
-        file = self.files[idx]
-        if file not in self.loaded_files:
-            self.loaded_files[file] = self.load_data(file)
-        return random.choice(self.loaded_files[file])
-
-    def __len__(self):
-        if self.pre_loaded:
-            return len(self.cached_data)
-        return len(self.files)
-
-    def __getitem__(self, index: int):
-        return self.get_item(index)