lt-tensor 0.0.1a12__tar.gz → 0.0.1a13__tar.gz

{lt_tensor-0.0.1a12 → lt_tensor-0.0.1a13}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lt-tensor
-Version: 0.0.1a12
+Version: 0.0.1a13
 Summary: General utilities for PyTorch and others. Built for general use.
 Home-page: https://github.com/gr1336/lt-tensor/
 Author: gr1336
@@ -11,15 +11,17 @@ Classifier: Topic :: Software Development :: Libraries
 Classifier: Topic :: Utilities
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: torch>=2.2.0
-Requires-Dist: torchaudio>=2.2.0
+Requires-Dist: torch>=2.7.0
+Requires-Dist: torchaudio>=2.7.0
 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.2a1
-Requires-Dist: librosa>=0.11.0
+Requires-Dist: librosa==0.11.*
+Requires-Dist: einops
 Requires-Dist: plotly
+Requires-Dist: scipy
 Dynamic: author
 Dynamic: classifier
 Dynamic: description

lt_tensor-0.0.1a13/lt_tensor/datasets/audio.py

@@ -0,0 +1,218 @@
+__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
+    max_frame_length: Optional[int] = None
+
+    def __init__(
+        self,
+        audio_processor: AudioProcessor,
+        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.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.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[-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.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, 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
+            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)

{lt_tensor-0.0.1a12 → lt_tensor-0.0.1a13}/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]],

lt_tensor-0.0.1a13/lt_tensor/model_zoo/__init__.py

@@ -0,0 +1,20 @@
+__all__ = [
+    "basic",  # basic
+    "residual",  # residual
+    "transformer",  # transformer
+    "pos_encoder", 
+    "fusion", 
+    "features", 
+    "discriminator",  
+    "istft",
+]
+from . import (
+    basic,
+    discriminator,
+    features,
+    fusion,
+    istft,
+    pos_encoder,
+    residual,
+    transformer,
+)

lt_tensor-0.0.1a12/lt_tensor/model_zoo/bsc.py → lt_tensor-0.0.1a13/lt_tensor/model_zoo/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)