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

lt_tensor/model_zoo/discriminator.py

@@ -1,314 +0,0 @@
-from lt_tensor.torch_commons import *
-import torch.nn.functional as F
-from lt_tensor.model_base import Model
-from lt_utils.common import *
-from einops import rearrange
-import torchaudio
-
-
-def get_padding(ks, d):
-    return int((ks * d - d) / 2)
-
-
-class DiscriminatorP(Model):
-    def __init__(self, period, kernel_size=5, stride=3, use_spectral_norm=False):
-        super().__init__()
-        self.period = period
-        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)))
-        self.activation = nn.LeakyReLU(0.1)
-
-    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(Model):
-    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(Model):
-    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 = 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 MultiScaleDiscriminator(Model):
-    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
-
-
-class MultiResolutionDiscriminator(Model):
-    """Source: https://github.com/gemelo-ai/vocos/blob/main/vocos/discriminators.py"""
-
-    def __init__(
-        self,
-        fft_sizes: Tuple[int, ...] = (2048, 1024, 512),
-        num_embeddings: Optional[int] = None,
-    ):
-        """
-
-        Args:
-            fft_sizes (tuple[int]): Tuple of window lengths for FFT. Defaults to (2048, 1024, 512).
-            num_embeddings (int, optional): Number of embeddings. None means non-conditional discriminator.
-                Defaults to None.
-        """
-
-        super().__init__()
-        self.discriminators = nn.ModuleList(
-            [
-                DiscriminatorR(window_length=w, num_embeddings=num_embeddings)
-                for w in fft_sizes
-            ]
-        )
-
-    def forward(
-        self, y: torch.Tensor, y_hat: torch.Tensor, bandwidth_id: torch.Tensor = None
-    ) -> Tuple[
-        List[torch.Tensor],
-        List[torch.Tensor],
-        List[List[torch.Tensor]],
-        List[List[torch.Tensor]],
-    ]:
-        y_d_rs = []
-        y_d_gs = []
-        fmap_rs = []
-        fmap_gs = []
-
-        for d in self.discriminators:
-            y_d_r, fmap_r = d(x=y, cond_embedding_id=bandwidth_id)
-            y_d_g, fmap_g = d(x=y_hat, cond_embedding_id=bandwidth_id)
-            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 DiscriminatorR(Model):
-    def __init__(
-        self,
-        window_length: int,
-        num_embeddings: Optional[int] = None,
-        channels: int = 32,
-        hop_factor: float = 0.25,
-        bands: Tuple[Tuple[float, float], ...] = (
-            (0.0, 0.1),
-            (0.1, 0.25),
-            (0.25, 0.5),
-            (0.5, 0.75),
-            (0.75, 1.0),
-        ),
-    ):
-        super().__init__()
-        self.window_length = window_length
-        self.hop_factor = hop_factor
-        self.spec_fn = torchaudio.transforms.Spectrogram(
-            n_fft=window_length,
-            hop_length=int(window_length * hop_factor),
-            win_length=window_length,
-            power=None,
-        )
-        n_fft = window_length // 2 + 1
-        bands = [(int(b[0] * n_fft), int(b[1] * n_fft)) for b in bands]
-        self.bands = bands
-        convs = lambda: nn.ModuleList(
-            [
-                weight_norm(nn.Conv2d(2, channels, (3, 9), (1, 1), padding=(1, 4))),
-                weight_norm(
-                    nn.Conv2d(channels, channels, (3, 9), (1, 2), padding=(1, 4))
-                ),
-                weight_norm(
-                    nn.Conv2d(channels, channels, (3, 9), (1, 2), padding=(1, 4))
-                ),
-                weight_norm(
-                    nn.Conv2d(channels, channels, (3, 9), (1, 2), padding=(1, 4))
-                ),
-                weight_norm(
-                    nn.Conv2d(channels, channels, (3, 3), (1, 1), padding=(1, 1))
-                ),
-            ]
-        )
-        self.band_convs = nn.ModuleList([convs() for _ in range(len(self.bands))])
-
-        if num_embeddings is not None:
-            self.emb = torch.nn.Embedding(
-                num_embeddings=num_embeddings, embedding_dim=channels
-            )
-            torch.nn.init.zeros_(self.emb.weight)
-
-        self.conv_post = weight_norm(
-            nn.Conv2d(channels, 1, (3, 3), (1, 1), padding=(1, 1))
-        )
-
-    def spectrogram(self, x):
-        # Remove DC offset
-        x = x - x.mean(dim=-1, keepdims=True)
-        # Peak normalize the volume of input audio
-        x = 0.8 * x / (x.abs().max(dim=-1, keepdim=True)[0] + 1e-9)
-        x = self.spec_fn(x)
-        x = torch.view_as_real(x)
-        x = rearrange(x, "b f t c -> b c t f")
-        # Split into bands
-        x_bands = [x[..., b[0] : b[1]] for b in self.bands]
-        return x_bands
-
-    def forward(self, x: torch.Tensor, cond_embedding_id: torch.Tensor = None):
-        x_bands = self.spectrogram(x)
-        fmap = []
-        x = []
-        for band, stack in zip(x_bands, self.band_convs):
-            for i, layer in enumerate(stack):
-                band = layer(band)
-                band = torch.nn.functional.leaky_relu(band, 0.1)
-                if i > 0:
-                    fmap.append(band)
-            x.append(band)
-        x = torch.cat(x, dim=-1)
-        if cond_embedding_id is not None:
-            emb = self.emb(cond_embedding_id)
-            h = (emb.view(1, -1, 1, 1) * x).sum(dim=1, keepdims=True)
-        else:
-            h = 0
-        x = self.conv_post(x)
-        fmap.append(x)
-        x += h
-
-        return x, fmap

lt_tensor/model_zoo/istft/__init__.py

@@ -1,5 +0,0 @@
-from .generator import iSTFTGenerator
-from . import trainer
-
-
-__all__ = ["iSTFTGenerator", "trainer"]

lt_tensor/model_zoo/istft/generator.py

@@ -1,94 +0,0 @@
-__all__ = ["iSTFTGenerator"]
-from lt_utils.common import *
-from lt_tensor.torch_commons import *
-from lt_tensor.model_zoo.residual import ConvNets, ResBlocks1D, ResBlock1D, ResBlock1D2
-
-
-class iSTFTGenerator(ConvNets):
-    def __init__(
-        self,
-        in_channels: int = 80,
-        upsample_rates: List[Union[int, List[int]]] = [8, 8],
-        upsample_kernel_sizes: List[Union[int, List[int]]] = [16, 16],
-        upsample_initial_channel: int = 512,
-        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],
-        ],
-        n_fft: int = 16,
-        activation: nn.Module = nn.LeakyReLU(0.1),
-        hop_length: int = 256,
-        residual_cls: Union[ResBlock1D, ResBlock1D2] = ResBlock1D
-    ):
-        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)
-        )
-        self.blocks = nn.ModuleList()
-        self.activation = activation
-        for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
-            self.blocks.append(
-                self._make_blocks(
-                    (i, k, u),
-                    upsample_initial_channel,
-                    resblock_kernel_sizes,
-                    resblock_dilation_sizes,
-                    residual_cls
-                )
-            )
-
-        ch = upsample_initial_channel // (2 ** (i + 1))
-        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 = nn.ReflectionPad1d((1, 0))
-
-    def _make_blocks(
-        self,
-        state: Tuple[int, int, int],
-        upsample_initial_channel: int,
-        resblock_kernel_sizes: List[Union[int, List[int]]],
-        resblock_dilation_sizes: List[int | List[int]],
-        residual: nn.Module
-    ):
-        i, k, u = state
-        channels = upsample_initial_channel // (2 ** (i + 1))
-        return nn.ModuleDict(
-            dict(
-                up=nn.Sequential(
-                    self.activation,
-                    weight_norm(
-                        nn.ConvTranspose1d(
-                            upsample_initial_channel // (2**i),
-                            channels,
-                            k,
-                            u,
-                            padding=(k - u) // 2,
-                        )
-                    ).apply(self.init_weights),
-                ),
-                residual=ResBlocks1D(
-                    channels,
-                    resblock_kernel_sizes,
-                    resblock_dilation_sizes,
-                    self.activation,
-                    residual
-                ),
-            )
-        )
-
-    def forward(self, x):
-        x = self.conv_pre(x)
-        for block in self.blocks:
-            x = block["up"](x)
-            x = block["residual"](x)
-
-        x = self.conv_post(self.activation(self.reflection_pad(x)))
-        spec = torch.exp(x[:, : self.post_n_fft, :])
-        phase = torch.sin(x[:, self.post_n_fft :, :])
-        return spec, phase

lt_tensor/model_zoo/istft/sg.py

@@ -1,142 +0,0 @@
-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