lt-tensor 0.0.1a34__py3-none-any.whl → 0.0.1a36__py3-none-any.whl

lt_tensor/model_zoo/losses/discriminators.py

@@ -7,8 +7,6 @@ from lt_tensor.model_base import Model
 from lt_tensor.model_zoo.convs import ConvNets
 from torch.nn import functional as F
 from torchaudio import transforms as T
-from lt_tensor.processors import AudioProcessor, AudioProcessorConfig
-
 
 MULTI_DISC_OUT_TYPE: TypeAlias = Tuple[
     List[Tensor],
@@ -19,9 +17,11 @@ MULTI_DISC_OUT_TYPE: TypeAlias = Tuple[
 
 
 class MultiDiscriminatorWrapper(Model):
-    def __init__(self, list_discriminator: List["_MultiDiscriminatorT"]):
+    def __init__(
+        self, list_discriminator: Union[List["_MultiDiscriminatorT"], nn.ModuleList]
+    ):
         """Setup example:
-        model_d = MultiDiscriminatorStep(
+        model_d = MultiDiscriminatorWrapper(
             [
                 MultiEnvelopeDiscriminator(),
                 MultiBandDiscriminator(),
@@ -31,7 +31,12 @@ class MultiDiscriminatorWrapper(Model):
         )
         """
         super().__init__()
-        self.disc: Sequence[_MultiDiscriminatorT] = nn.ModuleList(list_discriminator)
+
+        self.disc: Sequence[_MultiDiscriminatorT] = (
+            nn.ModuleList(list_discriminator)
+            if isinstance(list_discriminator, (list, tuple, set))
+            else list_discriminator
+        )
         self.total = len(self.disc)
 
     def forward(
@@ -96,7 +101,6 @@ class _MultiDiscriminatorT(ConvNets):
     def forward(self, y: Tensor, y_hat: Tensor) -> MULTI_DISC_OUT_TYPE:
         pass
 
-    # for type hinting
     def __call__(self, *args, **kwds) -> MULTI_DISC_OUT_TYPE:
         return super().__call__(*args, **kwds)
 
@@ -159,7 +163,7 @@ class DiscriminatorP(ConvNets):
     def __init__(
         self,
         period: List[int],
-        discriminator_channel_mult: Number = 1,
+        discriminator_channel_multi: Number = 1,
         kernel_size: int = 5,
         stride: int = 3,
         use_spectral_norm: bool = False,
@@ -167,7 +171,7 @@ class DiscriminatorP(ConvNets):
         super().__init__()
         self.period = period
         norm_f = weight_norm if not use_spectral_norm else spectral_norm
-        dsc = lambda x: int(x * discriminator_channel_mult)
+        dsc = lambda x: int(x * discriminator_channel_multi)
         self.convs = nn.ModuleList(
             [
                 norm_f(
@@ -242,19 +246,18 @@ class DiscriminatorP(ConvNets):
 class MultiPeriodDiscriminator(_MultiDiscriminatorT):
     def __init__(
         self,
-        discriminator_channel_mult: Number = 1,
+        discriminator_channel_multi: Number = 1,
         mpd_reshapes: list[int] = [2, 3, 5, 7, 11],
         use_spectral_norm: bool = False,
     ):
         super().__init__()
         self.mpd_reshapes = mpd_reshapes
-        print(f"mpd_reshapes: {self.mpd_reshapes}")
         self.discriminators = nn.ModuleList(
             [
                 DiscriminatorP(
                     rs,
                     use_spectral_norm=use_spectral_norm,
-                    discriminator_channel_mult=discriminator_channel_mult,
+                    discriminator_channel_multi=discriminator_channel_multi,
                 )
                 for rs in self.mpd_reshapes
             ]
@@ -276,6 +279,79 @@ class MultiPeriodDiscriminator(_MultiDiscriminatorT):
         return y_d_rs, y_d_gs, fmap_rs, fmap_gs
 
 
+class DiscriminatorS(ConvNets):
+    def __init__(
+        self,
+        use_spectral_norm=False,
+        discriminator_channel_multi: Number = 1,
+    ):
+        super().__init__()
+        norm_f = weight_norm if use_spectral_norm == False else spectral_norm
+        dsc = lambda x: int(x * discriminator_channel_multi)
+        self.convs = nn.ModuleList(
+            [
+                norm_f(nn.Conv1d(1, dsc(128), 15, 1, padding=7)),
+                norm_f(nn.Conv1d(dsc(128), dsc(128), 41, 2, groups=4, padding=20)),
+                norm_f(nn.Conv1d(dsc(128), dsc(256), 41, 2, groups=16, padding=20)),
+                norm_f(nn.Conv1d(dsc(256), dsc(512), 41, 4, groups=16, padding=20)),
+                norm_f(nn.Conv1d(dsc(512), dsc(1024), 41, 4, groups=16, padding=20)),
+                norm_f(nn.Conv1d(dsc(1024), dsc(1024), 41, 1, groups=16, padding=20)),
+                norm_f(nn.Conv1d(dsc(1024), dsc(1024), 5, 1, padding=2)),
+            ]
+        )
+        self.conv_post = norm_f(nn.Conv1d(dsc(1024), 1, 3, 1, padding=1))
+        self.activation = nn.LeakyReLU(0.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)
+        return x.flatten(1, -1), fmap
+
+
+class MultiScaleDiscriminator(ConvNets):
+    def __init__(
+        self,
+        discriminator_channel_multi: Number = 1,
+    ):
+        super().__init__()
+        self.discriminators = nn.ModuleList(
+            [
+                DiscriminatorS(
+                    use_spectral_norm=True,
+                    discriminator_channel_multi=discriminator_channel_multi,
+                ),
+                DiscriminatorS(discriminator_channel_multi=discriminator_channel_multi),
+                DiscriminatorS(discriminator_channel_multi=discriminator_channel_multi),
+            ]
+        )
+        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 EnvelopeExtractor(Model):
     """Extracts the amplitude envelope of the audio signal."""
 
@@ -297,21 +373,35 @@ class EnvelopeExtractor(Model):
 
 
 class DiscriminatorEnvelope(ConvNets):
-    def __init__(self, use_spectral_norm=False):
+    def __init__(
+        self,
+        use_spectral_norm=False,
+        discriminator_channel_multi: Number = 1,
+        kernel_size: int = 101,
+    ):
         super().__init__()
         norm_f = weight_norm if not use_spectral_norm else spectral_norm
-        self.extractor = EnvelopeExtractor(kernel_size=101)
+        self.extractor = EnvelopeExtractor(kernel_size=kernel_size)
+        dsc = lambda x: int(x * discriminator_channel_multi)
         self.convs = nn.ModuleList(
             [
-                norm_f(nn.Conv1d(1, 64, 15, stride=1, padding=7)),
-                norm_f(nn.Conv1d(64, 128, 41, stride=2, groups=4, padding=20)),
-                norm_f(nn.Conv1d(128, 256, 41, stride=2, groups=16, padding=20)),
-                norm_f(nn.Conv1d(256, 512, 41, stride=4, groups=16, padding=20)),
-                norm_f(nn.Conv1d(512, 512, 41, stride=4, groups=16, padding=20)),
-                norm_f(nn.Conv1d(512, 512, 5, stride=1, padding=2)),
+                norm_f(nn.Conv1d(1, dsc(64), 15, stride=1, padding=7)),
+                norm_f(
+                    nn.Conv1d(dsc(64), dsc(128), 41, stride=2, groups=4, padding=20)
+                ),
+                norm_f(
+                    nn.Conv1d(dsc(128), dsc(256), 41, stride=2, groups=16, padding=20)
+                ),
+                norm_f(
+                    nn.Conv1d(dsc(256), dsc(512), 41, stride=4, groups=16, padding=20)
+                ),
+                norm_f(
+                    nn.Conv1d(dsc(512), dsc(512), 41, stride=4, groups=16, padding=20)
+                ),
+                norm_f(nn.Conv1d(dsc(512), dsc(512), 5, stride=1, padding=2)),
             ]
         )
-        self.conv_post = norm_f(nn.Conv1d(512, 1, 3, stride=1, padding=1))
+        self.conv_post = norm_f(nn.Conv1d(dsc(512), 1, 3, stride=1, padding=1))
         self.activation = nn.LeakyReLU(0.1)
 
     def forward(self, x):
@@ -327,11 +417,17 @@ class DiscriminatorEnvelope(ConvNets):
 
 
 class MultiEnvelopeDiscriminator(_MultiDiscriminatorT):
-    def __init__(self, use_spectral_norm: bool = False):
+    def __init__(
+        self,
+        use_spectral_norm: bool = False,
+        discriminator_channel_multi: Number = 1,
+    ):
         super().__init__()
         self.discriminators = nn.ModuleList(
             [
-                DiscriminatorEnvelope(use_spectral_norm),  # raw envelope
+                DiscriminatorEnvelope(
+                    use_spectral_norm, discriminator_channel_multi
+                ),  # raw envelope
                 DiscriminatorEnvelope(use_spectral_norm),  # downsampled once
                 DiscriminatorEnvelope(use_spectral_norm),  # downsampled twice
             ]
@@ -431,7 +527,7 @@ class DiscriminatorB(ConvNets):
         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)
+                band = F.leaky_relu(band, 0.1)
                 if i > 0:
                     fmap.append(band)
             x.append(band)
@@ -452,11 +548,21 @@ class MultiBandDiscriminator(_MultiDiscriminatorT):
     def __init__(
         self,
         mbd_fft_sizes: list[int] = [2048, 1024, 512],
+        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.fft_sizes = mbd_fft_sizes
+        kwargs_disc = dict(channels=channels, hop_factor=hop_factor, bands=bands)
         self.discriminators = nn.ModuleList(
-            [DiscriminatorB(window_length=w) for w in self.fft_sizes]
+            [DiscriminatorB(window_length=w, **kwargs_disc) for w in self.fft_sizes]
         )
 
     def forward(self, y: Tensor, y_hat: Tensor) -> MULTI_DISC_OUT_TYPE:
@@ -483,7 +589,7 @@ class DiscriminatorR(ConvNets):
         self,
         resolution: List[int],
         use_spectral_norm: bool = False,
-        discriminator_channel_mult: int = 1,
+        discriminator_channel_multi: Number = 1,
     ):
         super().__init__()
 
@@ -501,13 +607,13 @@ class DiscriminatorR(ConvNets):
             [
                 norm_f(
                     nn.Conv2d(
-                        1, int(32 * discriminator_channel_mult), (3, 9), padding=(1, 4)
+                        1, int(32 * discriminator_channel_multi), (3, 9), padding=(1, 4)
                     )
                 ),
                 norm_f(
                     nn.Conv2d(
-                        int(32 * discriminator_channel_mult),
-                        int(32 * discriminator_channel_mult),
+                        int(32 * discriminator_channel_multi),
+                        int(32 * discriminator_channel_multi),
                         (3, 9),
                         stride=(1, 2),
                         padding=(1, 4),
@@ -515,8 +621,8 @@ class DiscriminatorR(ConvNets):
                 ),
                 norm_f(
                     nn.Conv2d(
-                        int(32 * discriminator_channel_mult),
-                        int(32 * discriminator_channel_mult),
+                        int(32 * discriminator_channel_multi),
+                        int(32 * discriminator_channel_multi),
                         (3, 9),
                         stride=(1, 2),
                         padding=(1, 4),
@@ -524,8 +630,8 @@ class DiscriminatorR(ConvNets):
                 ),
                 norm_f(
                     nn.Conv2d(
-                        int(32 * discriminator_channel_mult),
-                        int(32 * discriminator_channel_mult),
+                        int(32 * discriminator_channel_multi),
+                        int(32 * discriminator_channel_multi),
                         (3, 9),
                         stride=(1, 2),
                         padding=(1, 4),
@@ -533,8 +639,8 @@ class DiscriminatorR(ConvNets):
                 ),
                 norm_f(
                     nn.Conv2d(
-                        int(32 * discriminator_channel_mult),
-                        int(32 * discriminator_channel_mult),
+                        int(32 * discriminator_channel_multi),
+                        int(32 * discriminator_channel_multi),
                         (3, 3),
                         padding=(1, 1),
                     )
@@ -542,7 +648,7 @@ class DiscriminatorR(ConvNets):
             ]
         )
         self.conv_post = norm_f(
-            nn.Conv2d(int(32 * discriminator_channel_mult), 1, (3, 3), padding=(1, 1))
+            nn.Conv2d(int(32 * discriminator_channel_multi), 1, (3, 3), padding=(1, 1))
         )
 
     def forward(self, x: Tensor) -> Tuple[Tensor, List[Tensor]]:
@@ -586,7 +692,7 @@ class MultiResolutionDiscriminator(_MultiDiscriminatorT):
     def __init__(
         self,
         use_spectral_norm: bool = False,
-        discriminator_channel_mult: int = 1,
+        discriminator_channel_multi: Number = 1,
         resolutions: List[List[int]] = [
             [1024, 120, 600],
             [2048, 240, 1200],
@@ -601,7 +707,7 @@ class MultiResolutionDiscriminator(_MultiDiscriminatorT):
         self.discriminators = nn.ModuleList(
             [
                 DiscriminatorR(
-                    resolution, use_spectral_norm, discriminator_channel_mult
+                    resolution, use_spectral_norm, discriminator_channel_multi
                 )
                 for resolution in self.resolutions
             ]
@@ -620,3 +726,235 @@ class MultiResolutionDiscriminator(_MultiDiscriminatorT):
             y_d_gs.append(y_d_g)
             fmap_gs.append(fmap_g)
         return y_d_rs, y_d_gs, fmap_rs, fmap_gs
+
+
+class DiscriminatorCQT(ConvNets):
+    """Adapted from https://github.com/open-mmlab/Amphion/blob/main/models/vocoders/gan/discriminator/mssbcqtd.py under the MIT license."""
+
+    def __init__(
+        self,
+        hop_length: int,
+        n_octaves: int,
+        bins_per_octave: int,
+        sampling_rate: int,
+        cqtd_filters: int = 128,
+        cqtd_max_filters: int = 1024,
+        cqtd_filters_scale: int = 1,
+        cqtd_dilations: list = [1, 2, 4],
+        cqtd_in_channels: int = 1,
+        cqtd_out_channels: int = 1,
+        cqtd_normalize_volume: bool = False,
+    ):
+        super().__init__()
+        self.filters = cqtd_filters
+        self.max_filters = cqtd_max_filters
+        self.filters_scale = cqtd_filters_scale
+        self.kernel_size = (3, 9)
+        self.dilations = cqtd_dilations
+        self.stride = (1, 2)
+
+        self.fs = sampling_rate
+        self.in_channels = cqtd_in_channels
+        self.out_channels = cqtd_out_channels
+        self.hop_length = hop_length
+        self.n_octaves = n_octaves
+        self.bins_per_octave = bins_per_octave
+
+        # Lazy-load
+        from lt_tensor.model_zoo.losses.CQT.transforms import CQT2010v2
+
+        self.cqt_transform = CQT2010v2(
+            sr=self.fs * 2,
+            hop_length=self.hop_length,
+            n_bins=self.bins_per_octave * self.n_octaves,
+            bins_per_octave=self.bins_per_octave,
+            output_format="Complex",
+            pad_mode="constant",
+        )
+
+        self.conv_pres = nn.ModuleList()
+        for _ in range(self.n_octaves):
+            self.conv_pres.append(
+                nn.Conv2d(
+                    self.in_channels * 2,
+                    self.in_channels * 2,
+                    kernel_size=self.kernel_size,
+                    padding=self.get_2d_padding(self.kernel_size),
+                )
+            )
+
+        self.convs = nn.ModuleList()
+
+        self.convs.append(
+            nn.Conv2d(
+                self.in_channels * 2,
+                self.filters,
+                kernel_size=self.kernel_size,
+                padding=self.get_2d_padding(self.kernel_size),
+            )
+        )
+
+        in_chs = min(self.filters_scale * self.filters, self.max_filters)
+        for i, dilation in enumerate(self.dilations):
+            out_chs = min(
+                (self.filters_scale ** (i + 1)) * self.filters, self.max_filters
+            )
+            self.convs.append(
+                weight_norm(
+                    nn.Conv2d(
+                        in_chs,
+                        out_chs,
+                        kernel_size=self.kernel_size,
+                        stride=self.stride,
+                        dilation=(dilation, 1),
+                        padding=self.get_2d_padding(self.kernel_size, (dilation, 1)),
+                    )
+                )
+            )
+            in_chs = out_chs
+        out_chs = min(
+            (self.filters_scale ** (len(self.dilations) + 1)) * self.filters,
+            self.max_filters,
+        )
+        self.convs.append(
+            weight_norm(
+                nn.Conv2d(
+                    in_chs,
+                    out_chs,
+                    kernel_size=(self.kernel_size[0], self.kernel_size[0]),
+                    padding=self.get_2d_padding(
+                        (self.kernel_size[0], self.kernel_size[0])
+                    ),
+                )
+            )
+        )
+
+        self.conv_post = weight_norm(
+            nn.Conv2d(
+                out_chs,
+                self.out_channels,
+                kernel_size=(self.kernel_size[0], self.kernel_size[0]),
+                padding=self.get_2d_padding((self.kernel_size[0], self.kernel_size[0])),
+            )
+        )
+
+        self.activation = torch.nn.LeakyReLU(negative_slope=0.1)
+        self.resample = T.Resample(orig_freq=self.fs, new_freq=self.fs * 2)
+
+        self.cqtd_normalize_volume = cqtd_normalize_volume
+        if self.cqtd_normalize_volume:
+            print(
+                f"[INFO] cqtd_normalize_volume set to True. Will apply DC offset removal & peak volume normalization in CQTD!"
+            )
+
+    def get_2d_padding(
+        self,
+        kernel_size: Tuple[int, int],
+        dilation: Tuple[int, int] = (1, 1),
+    ):
+        return (
+            ((kernel_size[0] - 1) * dilation[0]) // 2,
+            ((kernel_size[1] - 1) * dilation[1]) // 2,
+        )
+
+    def forward(self, x: torch.tensor) -> Tuple[torch.Tensor, List[torch.Tensor]]:
+        fmap = []
+
+        if self.cqtd_normalize_volume:
+            # 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.resample(x)
+
+        z = self.cqt_transform(x)
+
+        z_amplitude = z[:, :, :, 0].unsqueeze(1)
+        z_phase = z[:, :, :, 1].unsqueeze(1)
+
+        z = torch.cat([z_amplitude, z_phase], dim=1)
+        z = torch.permute(z, (0, 1, 3, 2))  # [B, C, W, T] -> [B, C, T, W]
+
+        latent_z = []
+        for i in range(self.n_octaves):
+            latent_z.append(
+                self.conv_pres[i](
+                    z[
+                        :,
+                        :,
+                        :,
+                        i * self.bins_per_octave : (i + 1) * self.bins_per_octave,
+                    ]
+                )
+            )
+        latent_z = torch.cat(latent_z, dim=-1)
+
+        for i, l in enumerate(self.convs):
+            latent_z = l(latent_z)
+
+            latent_z = self.activation(latent_z)
+            fmap.append(latent_z)
+
+        latent_z = self.conv_post(latent_z)
+
+        return latent_z, fmap
+
+
+class MultiScaleSubbandCQTDiscriminator(_MultiDiscriminatorT):
+    def __init__(
+        self,
+        sampling_rate: int,
+        cqtd_filters: int = 128,
+        cqtd_max_filters: int = 1024,
+        cqtd_filters_scale: Number = 1,
+        cqtd_dilations: list = [1, 2, 4],
+        cqtd_hop_lengths: list = [512, 256, 256],
+        cqtd_n_octaves: list = [9, 9, 9],
+        cqtd_bins_per_octaves: list = [24, 36, 48],
+        cqtd_in_channels: int = 1,
+        cqtd_out_channels: int = 1,
+        cqtd_normalize_volume: bool = False,
+    ):
+        super().__init__()
+
+        self.discriminators = nn.ModuleList(
+            [
+                DiscriminatorCQT(
+                    hop_length=cqtd_hop_lengths[i],
+                    n_octaves=cqtd_n_octaves[i],
+                    bins_per_octave=cqtd_bins_per_octaves[i],
+                    sampling_rate=sampling_rate,
+                    cqtd_filters=cqtd_filters,
+                    cqtd_max_filters=cqtd_max_filters,
+                    cqtd_filters_scale=cqtd_filters_scale,
+                    cqtd_dilations=cqtd_dilations,
+                    cqtd_in_channels=cqtd_in_channels,
+                    cqtd_out_channels=cqtd_out_channels,
+                    cqtd_normalize_volume=cqtd_normalize_volume,
+                )
+                for i in range(len(cqtd_hop_lengths))
+            ]
+        )
+
+    def forward(self, y: torch.Tensor, y_hat: torch.Tensor) -> 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 disc in self.discriminators:
+            y_d_r, fmap_r = disc(y)
+            y_d_g, fmap_g = disc(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