lt-tensor 0.0.1a12__py3-none-any.whl → 0.0.1a14__py3-none-any.whl

lt_tensor/datasets/audio.py

@@ -1,6 +1,7 @@
 __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
@@ -9,64 +10,131 @@ from lt_tensor.misc_utils import log_tensor
 from tqdm import tqdm
 
 
+DEFAULT_DEVICE = torch.tensor([0]).device
+
+
 class WaveMelDataset(Dataset):
-    data: Union[list[dict[str, Tensor]], Tuple[Tensor, Tensor]] = []
+    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: bool = False,
+        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.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
+
         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.ap = audio_processor
-        self.files = self.ap.find_audios(path)
+            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)
 
+    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, max_frame_length, randomize_ranges)
-            self.data.extend(results)
-
-    def _add_dict(self, audio_raw: Tensor, audio_mel: Tensor, file: PathLike):
-        return {"mel": audio_mel, "raw": audio_raw, "file": file}
+            results = self.load_data(file)
+            if not results:
+                continue
+            self.cached_data.extend(results)
 
-    def load_data(
+    def _add_dict(
         self,
+        audio_wave: Tensor,
+        audio_mel: Tensor,
+        pitch: Tensor,
+        rms: Tensor,
         file: PathLike,
-        audio_frames_limit: Optional[int] = None,
-        randomize_ranges: bool = False,
     ):
-        initial_audio = self.ap.rebuild_spectrogram(self.ap.load_audio(file))
-        if not audio_frames_limit or initial_audio.shape[-1] <= audio_frames_limit:
-            if initial_audio.shape[-1] < self.post_n_fft:
-                return []
+        return {
+            "wave": audio_wave,
+            "pitch": pitch,
+            "rms": rms,
+            "mel": audio_mel,
+            "file": file,
+        }
+
+    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 = []
-        if randomize_ranges:
-            frame_limit = random.randint(self.r_range, audio_frames_limit)
+
+        if self.randomize_ranges:
+            frame_limit = random.randint(self.r_range, self.max_frame_length)
         else:
-            frame_limit = audio_frames_limit
-        for fragment in torch.split(
-            initial_audio, split_size_or_sections=frame_limit, dim=-1
-        ):
+            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(
@@ -93,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/misc_utils.py

@@ -23,6 +23,7 @@ __all__ = [
     "get_losses",
     "plot_view",
     "get_weights",
+    "get_activated_conv",
 ]
 
 import re
@@ -39,6 +40,42 @@ 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]],
@@ -203,7 +240,7 @@ class LogTensor:
     stored_items: List[
         Dict[str, Union[str, Number, Tensor, List[Union[Tensor, Number, str]]]]
     ] = []
-    max_stored_items: int = 32
+    max_stored_items: int = 8
 
     def _setup_message(self, title: str, t: Union[Tensor, str, int]):
         try:

lt_tensor/model_zoo/__init__.py

@@ -1,11 +1,20 @@
 __all__ = [
-    "bsc",  # basic
-    "rsd",  # residual
-    "tfrms",  # transformer
-    "pos",  # positional encoders
-    "fsn",  # fusion
-    "gns",  # generators
-    "disc",  # discriminators
-    "istft",  # self-explanatory
+    "basic",  # basic
+    "residual",  # residual
+    "transformer",  # transformer
+    "pos_encoder", 
+    "fusion", 
+    "features", 
+    "discriminator",  
+    "istft",
 ]
-from . import bsc, fsn, gns, istft, pos, rsd, tfrms, disc
+from . import (
+    basic,
+    discriminator,
+    features,
+    fusion,
+    istft,
+    pos_encoder,
+    residual,
+    transformer,
+)

lt_tensor/model_zoo/{bsc.py → basic.py}

@@ -8,11 +8,14 @@ __all__ = [
     "StyleEncoder",
     "PatchEmbed1D",
     "MultiScaleEncoder1D",
+    "UpSampleConv1D",
 ]
-
+import torch.nn.functional as F
 from lt_tensor.torch_commons import *
 from lt_tensor.model_base import Model
 from lt_tensor.transform import get_sinusoidal_embedding
+from lt_utils.common import *
+import math
 
 
 class FeedForward(Model):
@@ -30,6 +33,9 @@ class FeedForward(Model):
             nn.Linear(d_model, ff_dim),
             activation,
             nn.Dropout(dropout),
+            nn.Linear(ff_dim, ff_dim),
+            activation,
+            nn.Dropout(dropout),
             nn.Linear(ff_dim, d_model),
             normalizer,
         )
@@ -38,6 +44,24 @@ class FeedForward(Model):
         return self.net(x)
 
 
+class UpSampleConv1D(nn.Module):
+    def __init__(self, upsample: bool = False, dim_in: int = 0, dim_out: int = 0):
+        super().__init__()
+        if upsample:
+            self.upsample = lambda x: F.interpolate(x, scale_factor=2, mode="nearest")
+        else:
+            self.upsample = nn.Identity()
+
+        if dim_in == dim_out:
+            self.learned = nn.Identity()
+        else:
+            self.learned = weight_norm(nn.Conv1d(dim_in, dim_out, 1, 1, 0, bias=False))
+
+    def forward(self, x):
+        x = self.upsample(x)
+        return self.learned(x)
+
+
 class MLP(Model):
     def __init__(
         self,
@@ -161,7 +185,6 @@ class StyleEncoder(Model):
         self.linear = nn.Linear(hidden, out_dim)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        # x: [B, Mels, T]
         x = self.net(x).squeeze(-1)  # [B, hidden]
         return self.linear(x)  # [B, out_dim]
 
@@ -230,3 +253,96 @@ class AudioClassifier(Model):
 
     def forward(self, x):
         return self.model(x)
+
+
+class LoRALinearLayer(nn.Module):
+
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        rank: int = 4,
+        alpha: float = 1.0,
+    ):
+        super().__init__()
+        self.down = nn.Linear(in_features, rank, bias=False)
+        self.up = nn.Linear(rank, out_features, bias=False)
+        self.alpha = alpha
+        self.rank = rank
+        self.out_features = out_features
+        self.in_features = in_features
+        self.ah = self.alpha / self.rank
+        self._down_dt = self.down.weight.dtype
+
+        nn.init.normal_(self.down.weight, std=1 / rank)
+        nn.init.zeros_(self.up.weight)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        orig_dtype = hidden_states.dtype
+        down_hidden_states = self.down(hidden_states.to(self._down_dt))
+        up_hidden_states = self.up(down_hidden_states) * self.ah
+        return up_hidden_states.to(orig_dtype)
+
+
+class LoRAConv1DLayer(nn.Module):
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        kernel_size: Union[int, Tuple[int, ...]] = 1,
+        rank: int = 4,
+        alpha: float = 1.0,
+    ):
+        super().__init__()
+        self.down = nn.Conv1d(
+            in_features, rank, kernel_size, padding=kernel_size // 2, bias=False
+        )
+        self.up = nn.Conv1d(
+            rank, out_features, kernel_size, padding=kernel_size // 2, bias=False
+        )
+        self.ah = alpha / rank
+        self._down_dt = self.down.weight.dtype
+        nn.init.kaiming_uniform_(self.down.weight, a=math.sqrt(5))
+        nn.init.zeros_(self.up.weight)
+
+    def forward(self, inputs: Tensor) -> Tensor:
+        orig_dtype = inputs.dtype
+        down_hidden_states = self.down(inputs.to(self._down_dt))
+        up_hidden_states = self.up(down_hidden_states) * self.ah
+        return up_hidden_states.to(orig_dtype)
+
+
+class LoRAConv2DLayer(nn.Module):
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        kernel_size: Union[int, Tuple[int, ...]] = (1, 1),
+        rank: int = 4,
+        alpha: float = 1.0,
+    ):
+        super().__init__()
+        self.down = nn.Conv2d(
+            in_features,
+            rank,
+            kernel_size,
+            padding="same",
+            bias=False,
+        )
+        self.up = nn.Conv2d(
+            rank,
+            out_features,
+            kernel_size,
+            padding="same",
+            bias=False,
+        )
+        self.ah = alpha / rank
+
+        nn.init.kaiming_normal_(self.down.weight, a=0.2)
+        nn.init.zeros_(self.up.weight)
+
+    def forward(self, inputs: Tensor) -> Tensor:
+        orig_dtype = inputs.dtype
+        down_hidden_states = self.down(inputs.to(self._down_dt))
+        up_hidden_states = self.up(down_hidden_states) * self.ah
+        return up_hidden_states.to(orig_dtype)