lt-tensor 0.0.1a18__tar.gz → 0.0.1a20__tar.gz

{lt_tensor-0.0.1a18 → lt_tensor-0.0.1a20}/PKG-INFO

@@ -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 → lt_tensor-0.0.1a20}/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-0.0.1a20/lt_tensor/config_templates.py

@@ -0,0 +1,91 @@
+from lt_utils.common import *
+from lt_utils.file_ops import load_json, save_json, FileScan
+from lt_utils.misc_utils import log_traceback, get_current_time
+from lt_utils.type_utils import is_pathlike, is_file, is_dir, is_dict, is_str
+from lt_tensor.misc_utils import updateDict
+
+
+class ModelConfig(ABC, OrderedDict):
+    _default_settings: Dict[str, Any] = {}
+    _forbidden_list: List[str] = [
+        "_default_settings",
+        "_forbidden_list",
+    ]
+
+    def __init__(
+        self,
+        **settings,
+    ):
+        self._default_settings = settings
+        self.set_state_dict(self._default_settings)
+
+    def reset_settings(self):
+        raise NotImplementedError("Not implemented")
+
+    def save_config(
+        self,
+        path: str,
+    ):
+        base = {
+            k: v
+            for k, v in self.state_dict().items()
+            if isinstance(v, (str, int, float, list, tuple, dict, set, bytes))
+        }
+
+        save_json(path, base, indent=4)
+
+    def set_value(self, var_name: str, value: str) -> None:
+        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 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 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],
+    ) -> "ModelConfig":
+        assert is_dict(dictionary)
+        return ModelConfig(**dictionary)
+
+    @classmethod
+    def from_path(cls, path_name: PathLike) -> "ModelConfig":
+        assert is_file(path_name) or is_dir(path_name)
+        settings = {}
+
+        if is_file(path_name):
+            settings.update(load_json(path_name, {}, errors="ignore"))
+        else:
+            files = FileScan.files(
+                path_name,
+                [
+                    "*_config.json",
+                    "config_*.json",
+                    "*_config.json",
+                    "cfg_*.json",
+                    "*_cfg.json",
+                    "cfg.json",
+                    "config.json",
+                    "settings.json",
+                    "settings_*.json",
+                    "*_settings.json",
+                ],
+            )
+            assert files, "No config file found in the provided directory!"
+            settings.update(load_json(files[-1], {}, errors="ignore"))
+        settings.pop("path", None)
+        settings.pop("path_name", None)
+
+        return ModelConfig(**settings)

{lt_tensor-0.0.1a18 → lt_tensor-0.0.1a20}/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-0.0.1a18 → lt_tensor-0.0.1a20}/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-0.0.1a18 → lt_tensor-0.0.1a20}/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 → lt_tensor-0.0.1a20}/lt_tensor.egg-info/PKG-INFO

@@ -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 → lt_tensor-0.0.1a20}/lt_tensor.egg-info/SOURCES.txt

@@ -17,8 +17,6 @@ lt_tensor.egg-info/SOURCES.txt
 lt_tensor.egg-info/dependency_links.txt
 lt_tensor.egg-info/requires.txt
 lt_tensor.egg-info/top_level.txt
-lt_tensor/datasets/__init__.py
-lt_tensor/datasets/audio.py
 lt_tensor/model_zoo/__init__.py
 lt_tensor/model_zoo/basic.py
 lt_tensor/model_zoo/features.py

{lt_tensor-0.0.1a18 → lt_tensor-0.0.1a20}/lt_tensor.egg-info/requires.txt

@@ -4,7 +4,7 @@ numpy>=1.26.4
 tokenizers
 pyyaml>=6.0.0
 numba>0.60.0
-lt-utils==0.0.2a2
+lt-utils==0.0.2a3
 librosa==0.11.*
 einops
 plotly

{lt_tensor-0.0.1a18 → lt_tensor-0.0.1a20}/setup.py

@@ -4,7 +4,7 @@ with open("README.md", "r", encoding="utf-8") as f:
     long_description = f.read()
 
 setup(
-    version="0.0.1a18",
+    version="0.0.1a20",
     name="lt-tensor",
     description="General utilities for PyTorch and others. Built for general use.",
     long_description=long_description,
@@ -17,7 +17,7 @@ setup(
         "tokenizers",
         "pyyaml>=6.0.0",
         "numba>0.60.0",
-        "lt-utils==0.0.2a2",
+        "lt-utils==0.0.2a3",
         "librosa==0.11.*",
         "einops",
         "plotly",

lt_tensor-0.0.1a18/lt_tensor/config_templates.py

@@ -1,105 +0,0 @@
-from lt_utils.common import *
-from lt_utils.file_ops import load_json, save_json, FileScan
-from lt_utils.misc_utils import log_traceback, get_current_time
-from lt_utils.type_utils import is_pathlike, is_file, is_dir, is_dict, is_str
-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"]
-
-    def __init__(
-        self,
-        settings: Dict[str, Any] = {},
-        path_name: Optional[Union[str, PathLike]] = None,
-    ):
-        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"
-
-    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})
-
-    def save_config(
-        self,
-        path_name: Optional[Union[PathLike, str]] = None,
-    ):
-        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 = self.get_state_dict()
-        save_json(self.path_name, base, indent=2)
-
-    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})
-
-    def get_value(self, var_name: str) -> Any:
-        return self.__dict__.get(var_name)
-
-    def __getattribute__(self, name):
-        return self.__dict__.get(name)
-
-    def get_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
-    ) -> "ModelConfig":
-        assert is_dict(dictionary)
-        return ModelConfig(dictionary, path)
-
-    @classmethod
-    def from_path(cls, path_name: PathLike) -> "ModelConfig":
-        assert is_file(path_name) or is_dir(path_name)
-        settings = {}
-
-        if is_file(path_name):
-            settings.update(load_json(path_name, {}, errors="ignore"))
-        else:
-            files = FileScan.files(
-                path_name,
-                [
-                    "*_config.json",
-                    "config_*.json",
-                    "*_config.json",
-                    "cfg_*.json",
-                    "*_cfg.json",
-                    "cfg.json",
-                    "config.json",
-                    "settings.json",
-                    "settings_*.json",
-                    "*_settings.json",
-                ],
-            )
-            assert files, "No config file found in the provided directory!"
-            settings.update(load_json(files[-1], {}, errors="ignore"))
-        return ModelConfig(settings, path_name)

lt_tensor-0.0.1a18/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)