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

lt_tensor/model_zoo/audio_models/istft/__init__.py

@@ -0,0 +1,409 @@
+__all__ = ["iSTFTGenerator"]
+from lt_utils.common import *
+from lt_tensor.torch_commons import *
+from lt_tensor.model_zoo.residual import ConvNets
+from torch.nn import functional as F
+
+
+def get_padding(ks, d):
+    return int((ks * d - d) / 2)
+
+
+class ResBlock1(ConvNets):
+    def __init__(self, h, channels, kernel_size=3, dilation=(1, 3, 5)):
+        super().__init__()
+        self.h = h
+        self.convs1 = nn.ModuleList(
+            [
+                weight_norm(
+                    nn.Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=dilation[0],
+                        padding=get_padding(kernel_size, dilation[0]),
+                    )
+                ),
+                weight_norm(
+                    nn.Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=dilation[1],
+                        padding=get_padding(kernel_size, dilation[1]),
+                    )
+                ),
+                weight_norm(
+                    nn.Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=dilation[2],
+                        padding=get_padding(kernel_size, dilation[2]),
+                    )
+                ),
+            ]
+        )
+        self.convs1.apply(self.init_weights)
+
+        self.convs2 = nn.ModuleList(
+            [
+                weight_norm(
+                    nn.Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=1,
+                        padding=get_padding(kernel_size, 1),
+                    )
+                ),
+                weight_norm(
+                    nn.Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=1,
+                        padding=get_padding(kernel_size, 1),
+                    )
+                ),
+                weight_norm(
+                    nn.Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=1,
+                        padding=get_padding(kernel_size, 1),
+                    )
+                ),
+            ]
+        )
+        self.activation = nn.LeakyReLU(0.1)
+        self.convs2.apply(self.init_weights)
+
+    def forward(self, x):
+        for c1, c2 in zip(self.convs1, self.convs2):
+            xt = self.activation(x)
+            xt = c1(xt)
+            xt = self.activation(xt)
+            xt = c2(xt)
+            x = xt + x
+        return x
+
+class ResBlock2(ConvNets):
+    def __init__(self, h, channels, kernel_size=3, dilation=(1, 3)):
+        super().__init__()
+        self.h = h
+        self.convs = nn.ModuleList(
+            [
+                weight_norm(
+                    nn.Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=dilation[0],
+                        padding=get_padding(kernel_size, dilation[0]),
+                    )
+                ),
+                weight_norm(
+                    nn.Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=dilation[1],
+                        padding=get_padding(kernel_size, dilation[1]),
+                    )
+                ),
+            ]
+        )
+        self.activation = nn.LeakyReLU(0.1)
+        self.convs.apply(self.init_weights)
+
+    def forward(self, x):
+        for c in self.convs:
+            xt = self.activation(x)
+            xt = c(xt)
+            x = xt + x
+        return x
+
+
+class iSTFTGenerator(ConvNets):
+    def __init__(self, h):
+        super().__init__()
+        self.h = h
+        self.num_kernels = len(h.resblock_kernel_sizes)
+        self.num_upsamples = len(h.upsample_rates)
+        self.conv_pre = weight_norm(
+            nn.Conv1d(80, h.upsample_initial_channel, 7, 1, padding=3)
+        )
+        resblock = ResBlock1 if h.resblock == "1" else ResBlock2
+
+        self.ups = nn.ModuleList()
+        for i, (u, k) in enumerate(zip(h.upsample_rates, h.upsample_kernel_sizes)):
+            if h.sampling_rate % 16000:
+                self.ups.append(
+                    weight_norm(
+                        nn.ConvTranspose1d(
+                            h.upsample_initial_channel // (2**i),
+                            h.upsample_initial_channel // (2 ** (i + 1)),
+                            k,
+                            u,
+                            padding=(k - u) // 2,
+                        )
+                    )
+                )
+            else:
+                self.ups.append(
+                    weight_norm(
+                        nn.ConvTranspose1d(
+                            h.upsample_initial_channel // (2**i),
+                            h.upsample_initial_channel // (2 ** (i + 1)),
+                            k,
+                            u,
+                            padding=(u // 2 + u % 2),
+                            output_padding=u % 2,
+                        )
+                    )
+                )
+
+        self.resblocks = nn.ModuleList()
+        for i in range(len(self.ups)):
+            ch = h.upsample_initial_channel // (2 ** (i + 1))
+            for j, (k, d) in enumerate(
+                zip(h.resblock_kernel_sizes, h.resblock_dilation_sizes)
+            ):
+                self.resblocks.append(resblock(h, ch, k, d))
+
+        self.post_n_fft = h.gen_istft_n_fft
+        self.conv_post = weight_norm(
+            nn.Conv1d(ch, self.post_n_fft + 2, 7, 1, padding=3)
+        )
+        self.ups.apply(self.init_weights)
+        self.conv_post.apply(self.init_weights)
+        self.activation = nn.LeakyReLU(0.1)
+        self.reflection_pad = torch.nn.ReflectionPad1d((1, 0))
+
+    def forward(self, x):
+        x = self.conv_pre(x)
+        for i in range(self.num_upsamples):
+            x = self.activation(x)
+            x = self.ups[i](x)
+            xs = None
+            for j in range(self.num_kernels):
+                if xs is None:
+                    xs = self.resblocks[i * self.num_kernels + j](x)
+                else:
+                    xs += self.resblocks[i * self.num_kernels + j](x)
+            x = xs / self.num_kernels
+        x = self.activation(x)
+        x = self.reflection_pad(x)
+        x = self.conv_post(x)
+        spec = torch.exp(x[:, : self.post_n_fft // 2 + 1, :])
+        phase = torch.sin(x[:, self.post_n_fft // 2 + 1 :, :])
+
+        return spec, phase
+
+
+class DiscriminatorP(ConvNets):
+    def __init__(self, period, kernel_size=5, stride=3, use_spectral_norm=False):
+        super().__init__()
+        self.period = period
+        self.activation = nn.LeakyReLU(0.1)
+        norm_f = weight_norm if use_spectral_norm == False else spectral_norm
+        self.convs = nn.ModuleList(
+            [
+                norm_f(
+                    nn.Conv2d(
+                        1,
+                        32,
+                        (kernel_size, 1),
+                        (stride, 1),
+                        padding=(get_padding(5, 1), 0),
+                    )
+                ),
+                norm_f(
+                    nn.Conv2d(
+                        32,
+                        128,
+                        (kernel_size, 1),
+                        (stride, 1),
+                        padding=(get_padding(5, 1), 0),
+                    )
+                ),
+                norm_f(
+                    nn.Conv2d(
+                        128,
+                        512,
+                        (kernel_size, 1),
+                        (stride, 1),
+                        padding=(get_padding(5, 1), 0),
+                    )
+                ),
+                norm_f(
+                    nn.Conv2d(
+                        512,
+                        1024,
+                        (kernel_size, 1),
+                        (stride, 1),
+                        padding=(get_padding(5, 1), 0),
+                    )
+                ),
+                norm_f(nn.Conv2d(1024, 1024, (kernel_size, 1), 1, padding=(2, 0))),
+            ]
+        )
+        self.conv_post = norm_f(nn.Conv2d(1024, 1, (3, 1), 1, padding=(1, 0)))
+
+    def forward(self, x):
+        fmap = []
+
+        # 1d to 2d
+        b, c, t = x.shape
+        if t % self.period != 0:  # pad first
+            n_pad = self.period - (t % self.period)
+            x = F.pad(x, (0, n_pad), "reflect")
+            t = t + n_pad
+        x = x.view(b, c, t // self.period, self.period)
+
+        for l in self.convs:
+            x = l(x)
+            x = self.activation(x)
+            fmap.append(x)
+        x = self.conv_post(x)
+        fmap.append(x)
+        x = torch.flatten(x, 1, -1)
+
+        return x, fmap
+
+
+class MultiPeriodDiscriminator(ConvNets):
+    def __init__(self):
+        super().__init__()
+        self.discriminators = nn.ModuleList(
+            [
+                DiscriminatorP(2),
+                DiscriminatorP(3),
+                DiscriminatorP(5),
+                DiscriminatorP(7),
+                DiscriminatorP(11),
+            ]
+        )
+
+    def forward(self, y, y_hat):
+        y_d_rs = []
+        y_d_gs = []
+        fmap_rs = []
+        fmap_gs = []
+        for i, d in enumerate(self.discriminators):
+            y_d_r, fmap_r = d(y)
+            y_d_g, fmap_g = d(y_hat)
+            y_d_rs.append(y_d_r)
+            fmap_rs.append(fmap_r)
+            y_d_gs.append(y_d_g)
+            fmap_gs.append(fmap_g)
+
+        return y_d_rs, y_d_gs, fmap_rs, fmap_gs
+
+
+class DiscriminatorS(ConvNets):
+    def __init__(self, use_spectral_norm=False):
+        super().__init__()
+        norm_f = weight_norm if use_spectral_norm == False else spectral_norm
+        self.convs = nn.ModuleList(
+            [
+                norm_f(nn.Conv1d(1, 128, 15, 1, padding=7)),
+                norm_f(nn.Conv1d(128, 128, 41, 2, groups=4, padding=20)),
+                norm_f(nn.Conv1d(128, 256, 41, 2, groups=16, padding=20)),
+                norm_f(nn.Conv1d(256, 512, 41, 4, groups=16, padding=20)),
+                norm_f(nn.Conv1d(512, 1024, 41, 4, groups=16, padding=20)),
+                norm_f(nn.Conv1d(1024, 1024, 41, 1, groups=16, padding=20)),
+                norm_f(nn.Conv1d(1024, 1024, 5, 1, padding=2)),
+            ]
+        )
+        self.activation = nn.LeakyReLU(0.1)
+        self.conv_post = norm_f(nn.Conv1d(1024, 1, 3, 1, padding=1))
+
+    def forward(self, x):
+        fmap = []
+        for l in self.convs:
+            x = self.activation(l(x))
+
+            fmap.append(x)
+        x = self.conv_post(x)
+        fmap.append(x)
+        x = torch.flatten(x, 1, -1)
+
+        return x, fmap
+
+
+class MultiScaleDiscriminator(ConvNets):
+    def __init__(self):
+        super().__init__()
+        self.discriminators = nn.ModuleList(
+            [
+                DiscriminatorS(use_spectral_norm=True),
+                DiscriminatorS(),
+                DiscriminatorS(),
+            ]
+        )
+        self.meanpools = nn.ModuleList(
+            [nn.AvgPool1d(4, 2, padding=2), nn.AvgPool1d(4, 2, padding=2)]
+        )
+
+    def forward(self, y, y_hat):
+        y_d_rs = []
+        y_d_gs = []
+        fmap_rs = []
+        fmap_gs = []
+        for i, d in enumerate(self.discriminators):
+            if i != 0:
+                y = self.meanpools[i - 1](y)
+                y_hat = self.meanpools[i - 1](y_hat)
+            y_d_r, fmap_r = d(y)
+            y_d_g, fmap_g = d(y_hat)
+            y_d_rs.append(y_d_r)
+            fmap_rs.append(fmap_r)
+            y_d_gs.append(y_d_g)
+            fmap_gs.append(fmap_g)
+
+        return y_d_rs, y_d_gs, fmap_rs, fmap_gs
+
+
+def feature_loss(fmap_r, fmap_g):
+    loss = 0
+    for dr, dg in zip(fmap_r, fmap_g):
+        for rl, gl in zip(dr, dg):
+            loss += torch.mean(torch.abs(rl - gl))
+
+    return loss * 2
+
+
+def discriminator_loss(disc_real_outputs, disc_generated_outputs):
+    loss = 0
+    r_losses = []
+    g_losses = []
+    for dr, dg in zip(disc_real_outputs, disc_generated_outputs):
+        r_loss = torch.mean((1 - dr) ** 2)
+        g_loss = torch.mean(dg**2)
+        loss += r_loss + g_loss
+        r_losses.append(r_loss.item())
+        g_losses.append(g_loss.item())
+
+    return loss, r_losses, g_losses
+
+
+def generator_loss(disc_outputs):
+    loss = 0
+    gen_losses = []
+    for dg in disc_outputs:
+        l = torch.mean((1 - dg) ** 2)
+        gen_losses.append(l)
+        loss += l
+
+    return loss, gen_losses

lt_tensor/model_zoo/basic.py

@@ -16,6 +16,7 @@ from lt_tensor.model_base import Model
 from lt_tensor.transform import get_sinusoidal_embedding
 from lt_utils.common import *
 import math
+from einops import repeat
 
 
 class FeedForward(Model):
@@ -346,3 +347,141 @@ class LoRAConv2DLayer(nn.Module):
         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 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