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

lt_tensor/model_zoo/{istft.py → istft/generator.py}

@@ -1,15 +1,55 @@
-from ..torch_commons import *
-from ..model_base import Model
-from .rsd import ResBlocks
-from ..misc_utils import log_tensor
-
+__all__ = ["iSTFTGenerator", "ResBlocks"]
+import gc
+import math
+import itertools
+from lt_utils.common import *
+from lt_tensor.torch_commons import *
+from lt_tensor.model_base import Model
+from lt_tensor.misc_utils import log_tensor
+from lt_tensor.model_zoo.residual import ResBlock1D, ConvNets, get_weight_norm
+from lt_utils.misc_utils import log_traceback
+from lt_tensor.processors import AudioProcessor
+from lt_utils.type_utils import is_dir, is_pathlike
+from lt_tensor.misc_utils import set_seed, clear_cache
+from lt_tensor.model_zoo.discriminator import MultiPeriodDiscriminator, MultiScaleDiscriminator
 import torch.nn.functional as F
+from lt_tensor.config_templates import updateDict, ModelConfig
 
 
-class Generator(Model):
-    """Based on the adaptation made by from Rishikesh
-    A Generator for audio processing, can be usd for tother things."""
+class ResBlocks(ConvNets):
+    def __init__(
+        self,
+        channels: int,
+        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),
+    ):
+        super().__init__()
+        self.num_kernels = len(resblock_kernel_sizes)
+        self.rb = nn.ModuleList()
+        self.activation = activation
+
+        for k, j in zip(resblock_kernel_sizes, resblock_dilation_sizes):
+            self.rb.append(ResBlock1D(channels, k, j, activation))
+
+        self.rb.apply(self.init_weights)
+
+    def forward(self, x: torch.Tensor):
+        xs = None
+        for i, block in enumerate(self.rb):
+            if i == 0:
+                xs = block(x)
+            else:
+                xs += block(x)
+        x = xs / self.num_kernels
+        return self.activation(x)
+
 
+class iSTFTGenerator(ConvNets):
     def __init__(
         self,
         in_channels: int = 80,
@@ -24,10 +64,12 @@ class Generator(Model):
         ],
         n_fft: int = 16,
         activation: nn.Module = nn.LeakyReLU(0.1),
+        hop_length: int = 256,
     ):
         super().__init__()
         self.num_kernels = len(resblock_kernel_sizes)
         self.num_upsamples = len(upsample_rates)
+        self.hop_length = hop_length
         self.conv_pre = weight_norm(
             nn.Conv1d(in_channels, upsample_initial_channel, 7, 1, padding=3)
         )
@@ -47,7 +89,20 @@ class Generator(Model):
         self.post_n_fft = n_fft // 2 + 1
         self.conv_post = weight_norm(nn.Conv1d(ch, n_fft + 2, 7, 1, padding=3))
         self.conv_post.apply(self.init_weights)
-        self.reflection_pad = torch.nn.ReflectionPad1d((1, 0))
+        self.reflection_pad = nn.ReflectionPad1d((1, 0))
+
+        self.phase = nn.Sequential(
+            nn.LeakyReLU(0.2),
+            nn.Conv1d(self.post_n_fft, self.post_n_fft, kernel_size=3, padding=1),
+            nn.LeakyReLU(0.2),
+            nn.Conv1d(self.post_n_fft, self.post_n_fft, kernel_size=3, padding=1),
+        )
+        self.spec = nn.Sequential(
+            nn.LeakyReLU(0.2),
+            nn.Conv1d(self.post_n_fft, self.post_n_fft, kernel_size=3, padding=1),
+            nn.LeakyReLU(0.2),
+            nn.Conv1d(self.post_n_fft, self.post_n_fft, kernel_size=3, padding=1),
+        )
 
     def _make_blocks(
         self,
@@ -70,7 +125,7 @@ class Generator(Model):
                             u,
                             padding=(k - u) // 2,
                         )
-                    ),
+                    ).apply(self.init_weights),
                 ),
                 residual=ResBlocks(
                     channels,
@@ -89,20 +144,7 @@ class Generator(Model):
 
         x = self.reflection_pad(x)
         x = self.conv_post(x)
-        spec = torch.exp(x[:, : self.post_n_fft, :])
-        phase = torch.sin(x[:, self.post_n_fft :, :])
+        spec = torch.exp(self.spec(x[:, : self.post_n_fft, :]))
+        phase = torch.sin(self.phase(x[:, self.post_n_fft :, :]))
 
         return spec, phase
-
-    def remove_weight_norm(self):
-        for module in self.modules():
-            try:
-                remove_weight_norm(module)
-            except ValueError:
-                pass  # Not normed, skip
-
-    @staticmethod
-    def init_weights(m, mean=0.0, std=0.01):
-        classname = m.__class__.__name__
-        if "Conv" in classname:
-            m.weight.data.normal_(mean, std)

lt_tensor/model_zoo/istft/sg.py

@@ -0,0 +1,142 @@
+import torch
+import torch.nn as nn
+import math
+from einops import repeat
+
+
+class SineGen(nn.Module):
+    def __init__(
+        self,
+        samp_rate,
+        upsample_scale,
+        harmonic_num=0,
+        sine_amp=0.1,
+        noise_std=0.003,
+        voiced_threshold=0,
+        flag_for_pulse=False,
+    ):
+        super().__init__()
+        self.sampling_rate = samp_rate
+        self.upsample_scale = upsample_scale
+        self.harmonic_num = harmonic_num
+        self.sine_amp = sine_amp
+        self.noise_std = noise_std
+        self.voiced_threshold = voiced_threshold
+        self.flag_for_pulse = flag_for_pulse
+        self.dim = self.harmonic_num + 1  # fundamental + harmonics
+
+    def _f02uv_b(self, f0):
+        return (f0 > self.voiced_threshold).float()  # [B, T]
+
+    def _f02uv(self, f0):
+        return (f0 > self.voiced_threshold).float().unsqueeze(-1)  # -> (B, T, 1)
+
+    @torch.no_grad()
+    def _f02sine(self, f0_values):
+        """
+        f0_values: (B, T, 1)
+        Output: sine waves (B, T * upsample, dim)
+        """
+        B, T, _ = f0_values.size()
+        f0_upsampled = repeat(
+            f0_values, "b t d -> b (t r) d", r=self.upsample_scale
+        )  # (B, T_up, 1)
+
+        # Create harmonics
+        harmonics = (
+            torch.arange(1, self.dim + 1, device=f0_values.device)
+            .float()
+            .view(1, 1, -1)
+        )
+        f0_harm = f0_upsampled * harmonics  # (B, T_up, dim)
+
+        # Convert Hz to radians (2πf/sr), then integrate to get phase
+        rad_values = f0_harm / self.sampling_rate  # normalized freq
+        rad_values = rad_values % 1.0  # remove multiples of 2π
+
+        # Random initial phase for each harmonic (except 0th if pulse mode)
+        if self.flag_for_pulse:
+            rand_ini = torch.zeros((B, 1, self.dim), device=f0_values.device)
+        else:
+            rand_ini = torch.rand((B, 1, self.dim), device=f0_values.device)
+
+        rand_ini = rand_ini * 2 * math.pi
+
+        # Compute cumulative phase
+        rad_values = rad_values * 2 * math.pi
+        phase = torch.cumsum(rad_values, dim=1) + rand_ini  # (B, T_up, dim)
+
+        sine_waves = torch.sin(phase)  # (B, T_up, dim)
+        return sine_waves
+
+    def _forward(self, f0):
+        """
+        f0: (B, T, 1)
+        returns: sine signal with harmonics and noise added
+        """
+        sine_waves = self._f02sine(f0)  # (B, T_up, dim)
+        uv = self._f02uv_b(f0) # (B, T, 1)
+        uv = repeat(uv, "b t d -> b (t r) d", r=self.upsample_scale)  # (B, T_up, 1)
+
+        # voiced sine + unvoiced noise
+        sine_signal = self.sine_amp * sine_waves * uv  # (B, T_up, dim)
+        noise = torch.randn_like(sine_signal) * self.noise_std
+        output = sine_signal + noise * (1.0 - uv)  # noise added only on unvoiced
+
+        return output  # (B, T_up, dim)
+
+    def forward(self, f0):
+        """
+        Args:
+            f0: (B, T) in Hz (before upsampling)
+        Returns:
+            sine_waves: (B, T_up, dim)
+            uv: (B, T_up, 1)
+            noise: (B, T_up, 1)
+        """
+        B, T = f0.shape
+        device = f0.device
+
+        # Get uv mask (before upsampling)
+        uv = self._f02uv(f0)  # (B, T, 1)
+
+        # Expand f0 to include harmonics: (B, T, dim)
+        f0 = f0.unsqueeze(-1)  # (B, T, 1)
+        harmonics = (
+            torch.arange(1, self.dim + 1, device=device).float().view(1, 1, -1)
+        )  # (1, 1, dim)
+        f0_harm = f0 * harmonics  # (B, T, dim)
+
+        # Upsample
+        f0_harm_up = repeat(
+            f0_harm, "b t d -> b (t r) d", r=self.upsample_scale
+        )  # (B, T_up, dim)
+        uv_up = repeat(uv, "b t d -> b (t r) d", r=self.upsample_scale)  # (B, T_up, 1)
+
+        # Convert to radians
+        rad_per_sample = f0_harm_up / self.sampling_rate  # Hz → cycles/sample
+        rad_per_sample = rad_per_sample * 2 * math.pi  # cycles → radians/sample
+
+        # Random phase init for each sample
+        B, T_up, D = rad_per_sample.shape
+        rand_phase = torch.rand(B, D, device=device) * 2 * math.pi  # (B, D)
+
+        # Compute cumulative phase
+        phase = torch.cumsum(rad_per_sample, dim=1) + rand_phase.unsqueeze(
+            1
+        )  # (B, T_up, D)
+
+        # Apply sine
+        sine_waves = torch.sin(phase) * self.sine_amp  # (B, T_up, D)
+
+        # Handle unvoiced: create noise only for fundamental
+        noise = torch.randn(B, T_up, 1, device=device) * self.noise_std
+        if self.flag_for_pulse:
+            # If pulse mode is on, align phase at start of voiced segments
+            # Optional and tricky to implement — may require segmenting uv
+            pass
+
+        # Replace sine by noise for unvoiced (only on fundamental)
+        sine_waves[:, :, 0:1] = sine_waves[:, :, 0:1] * uv_up + noise * (1 - uv_up)
+
+        return sine_waves, uv_up, noise