lt-tensor 0.0.1a11__py3-none-any.whl → 0.0.1a13__py3-none-any.whl

lt_tensor/__init__.py

@@ -12,6 +12,7 @@ from . import (
     losses,
     processors,
     datasets,
+    torch_commons,
 )
 
 __all__ = [
@@ -26,4 +27,5 @@ __all__ = [
     "losses",
     "processors",
     "datasets",
+    "torch_commons",
 ]

lt_tensor/config_templates.py

@@ -0,0 +1,97 @@
+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] = [
+        "_settings",
+    ]
+
+    def __init__(
+        self,
+        settings: Dict[str, Any] = None,
+        path_name: Optional[Union[str, PathLike]] = None,
+    ):
+        assert is_dict(settings)
+        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):
+        for s_name, setting in self._default_settings.items():
+            if s_name in self._forbidden_list:
+                continue
+            updateDict(self, {s_name: setting})
+
+    def save_config(
+        self,
+        path_name: Union[PathLike, str],
+    ):
+        assert is_pathlike(
+            path_name, True
+        ), f"path_name should be a non-empty string or pathlike object! received instead: {path_name}"
+        self._setup_path_name(path_name)
+        base = {k: y for k, y in self.__dict__.items() if k not in self._forbidden_list}
+        save_json(self.path_name, base, indent=2)
+
+    def to_dict(self):
+        return {k: y for k, y in self.__dict__.items() if k not in self._forbidden_list}
+
+    def set_value(self, var_name: str, value: str) -> None:
+        updateDict(self, {var_name: value})
+
+    def get_value(self, var_name: str) -> Any:
+        return self.__dict__.get(var_name)
+
+    @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/datasets/audio.py

@@ -1,58 +1,140 @@
 __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
 
 
-class WaveMelDataset(Dataset):
-    """Untested!"""
+DEFAULT_DEVICE = torch.tensor([0]).device
+
 
-    data: Union[list[dict[str, Tensor]], Tuple[Tensor, Tensor]] = []
+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
+    max_frame_length: Optional[int] = None
 
     def __init__(
         self,
         audio_processor: AudioProcessor,
-        path: PathLike,
+        dataset_path: PathLike,
         limit_files: 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,
+        n_noises: int = 0,  # TODO: Implement the random noises into the dataset
     ):
         super().__init__()
         assert max_frame_length is None or max_frame_length >= (
             (audio_processor.n_fft // 2) + 1
         )
-        self.post_n_fft = (audio_processor.n_fft // 2) + 1
         self.ap = audio_processor
-        self.files = self.ap.find_audios(path)
+        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
+
+        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
+
+        self.files = self.ap.find_audios(dataset_path, maximum=None)
         if limit_files:
             random.shuffle(self.files)
-            self.files = self.files[:limit_files]
-        self.data = []
+            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)
 
-        for file in self.files:
-            results = self.load_data(file, max_frame_length)
-            self.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_raw: Tensor, audio_mel: Tensor, file: PathLike):
-        return {"mel": audio_mel, "raw": audio_raw, "file": file}
+    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, audio_frames_limit: Optional[int] = None):
-        initial_audio = self.ap.load_audio(file)
-        if not audio_frames_limit or initial_audio.shape[-1] <= audio_frames_limit:
+    def load_data(self, file: PathLike):
+        initial_audio = self.ap.normalize_audio(
+            self.ap.load_audio(
+                file, normalize=self.normalize_waves, alpha=self.alpha_wv
+            )
+        )
+        if initial_audio.shape[-1] < self.post_n_fft:
+            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, file)]
+            return [
+                self._add_dict(initial_audio, audio_mel, audio_pitch, audio_rms, file)
+            ]
         results = []
-        for fragment in torch.split(
-            initial_audio, split_size_or_sections=audio_frames_limit, dim=-1
-        ):
+
+        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:
-                # sometimes the tensor will be too small to be able to pass on mel
+                # Too small
                 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, file))
+            results.append(
+                self._add_dict(fragment, audio_mel, audio_pitch, audio_rms, file)
+            )
         return results
 
     def get_data_loader(
@@ -79,31 +161,58 @@ class WaveMelDataset(Dataset):
             collate_fn=self.collate_fn,
         )
 
-    @staticmethod
-    def collate_fn(batch: Sequence[Dict[str, Tensor]]):
-        mels = []
-        audios = []
-        files = []
+    def collate_fn(self, batch: Sequence[Dict[str, Tensor]]):
+        mel = []
+        wave = []
+        file = []
+        rms = []
+        pitch = []
         for x in batch:
-            mels.append(x["mel"])
-            audios.append(x["raw"])
-            files.append(x["file"])
+            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 mels])
-        max_audio_len = max([a.shape[-1] for a in audios])
+        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_mels = torch.stack(
-            [FT.pad(m, (0, max_mel_len - m.shape[-1])) for m in mels]
+        padded_mel = torch.stack(
+            [FT.pad(m, (0, max_mel_len - m.shape[-1])) for m in mel]
         )  # shape: [B, 80, T_max]
 
-        padded_audios = torch.stack(
-            [FT.pad(a, (0, max_audio_len - a.shape[-1])) for a in audios]
+        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,
+        )
 
-        return padded_mels, padded_audios, files
+    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):
-        return len(self.data)
+        if self.pre_loaded:
+            return len(self.cached_data)
+        return len(self.files)
 
-    def __getitem__(self, index):
-        return self.data[index]
+    def __getitem__(self, index: int):
+        return self.get_item(index)

lt_tensor/losses.py

@@ -11,7 +11,7 @@ __all__ = [
 ]
 import math
 import random
-from .torch_commons import *
+from lt_tensor.torch_commons import *
 from lt_utils.common import *
 import torch.nn.functional as F
 

lt_tensor/math_ops.py

@@ -12,7 +12,7 @@ __all__ = [
     "phase",
 ]
 
-from .torch_commons import *
+from lt_tensor.torch_commons import *
 
 
 def sin_tensor(x: Tensor, freq: float = 1.0) -> Tensor:

lt_tensor/misc_utils.py

@@ -21,6 +21,9 @@ __all__ = [
     "Masking",
     "LogTensor",
     "get_losses",
+    "plot_view",
+    "get_weights",
+    "get_activated_conv",
 ]
 
 import re
@@ -28,13 +31,116 @@ import gc
 import sys
 import random
 import numpy as np
-from lt_utils.type_utils import is_str
-from .torch_commons import *
+import warnings
+from lt_utils.type_utils import is_str, is_dir, is_file, is_pathlike, is_path_valid
+from lt_utils.file_ops import FileScan, find_files, path_to_str, load_json, load_yaml
+from lt_tensor.torch_commons import *
 from lt_utils.misc_utils import cache_wrapper
 from lt_utils.common import *
 from lt_utils.misc_utils import ff_list
 import torch.nn.functional as F
 
+CONV_MAP = {
+    "conv1d": nn.Conv1d,
+    "conv1d": nn.Conv2d,
+    "conv1d": nn.Conv3d,
+    "convtranspose1d": nn.ConvTranspose1d,
+    "convtranspose2d": nn.ConvTranspose2d,
+    "convtranspose3d": nn.ConvTranspose3d,
+}
+
+
+def get_activated_conv(
+    in_channels: int,
+    out_channels: int,
+    kernel_size: int = 1,
+    stride: int = 1,
+    padding: Union[int, str] = 0,
+    groups: int = 1,
+    conv_type: Literal[
+        "Conv1d",
+        "Conv2d",
+        "Conv3d",
+        "ConvTranspose1d",
+        "ConvTranspose2d",
+        "ConvTranspose3d",
+    ] = "Conv1d",
+    activation: nn.Module = nn.LeakyReLU(0.2),
+    norm_fn: Callable[[nn.Module], nn.Module] = lambda x: x,
+):
+    assert conv_type.lower() in CONV_MAP, f"Invalid conv type: {conv_type}."
+    return nn.Sequential(
+        activation,
+        norm_fn(
+            nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, groups)
+        ),
+    )
+
+
+def plot_view(
+    data: Dict[str, List[Any]],
+    title: str = "Loss",
+    max_amount: int = 0,
+    xaxis_title="Step/Epoch",
+    yaxis_title="Loss",
+    template="plotly_dark",
+):
+    try:
+        import plotly.graph_objs as go
+    except ModuleNotFoundError:
+        warnings.warn(
+            "No installation of plotly was found. To use it use 'pip install plotly' and restart this application!"
+        )
+        return
+    fig = go.Figure()
+    for mode, values in data.items():
+        if values:
+            items = values if not max_amount > 0 else values[-max_amount:]
+            fig.add_trace(go.Scatter(y=items, name=mode.capitalize()))
+    fig.update_layout(
+        title=title,
+        xaxis_title=xaxis_title,
+        yaxis_title=yaxis_title,
+        template=template,
+    )
+    return fig
+
+
+def get_weights(directory: Union[str, PathLike]):
+    is_path_valid(directory, validate=True)  # raises validation if its invalid path
+    directory = Path(directory)
+    if is_file(directory):
+        if directory.name.endswith((".pt", ".ckpt", ".pth")):
+            return directory
+        directory = directory.parent
+    res = sorted(find_files(directory, ["*.pt", "*.ckpt", "*.pth"]))
+    return res[-1] if res else None
+
+
+def get_config(directory: Union[str, PathLike], default: Optional[Any] = None):
+    # raises validation if its invalid path only when default is None otherwise it returns the defaults.
+    if not is_path_valid(directory, validate=default is None):
+        return default
+    directory = Path(directory)
+    if is_file(directory):
+        if directory.name.endswith((".json", ".yaml", ".yml")):
+            if directory.name.endswith(".json"):
+                return load_json(directory, default)
+            return load_yaml(directory, default)
+        directory = directory.parent
+    res = sorted(find_files(directory, ["*.pt", "*.ckpt", "*.pth"]))
+    if res:
+        res = res[-1]
+        if Path(res).name.endswith(".json"):
+            return load_json(directory, default)
+        return load_yaml(directory, default)
+    return default
+
+
+def updateDict(self, dct: dict[str, Any]):
+    for k, v in dct.items():
+        setattr(self, k, v)
+
 
 def soft_restore(tensor, epsilon=1e-6):
     return torch.where(tensor == 0, torch.full_like(tensor, epsilon), tensor)