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

lt_tensor/model_zoo/losses/discriminators.py

@@ -726,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

lt_tensor/processors/audio.py

@@ -73,7 +73,7 @@ class AudioProcessorConfig(ModelConfig):
     def post_process(self):
         self.n_stft = self.n_fft // 2 + 1
         # some functions needs this to be a non-zero or not None value.
-        self.f_min = max(self.f_min, (self.sample_rate / (self.n_fft - 1))  * 2)
+        self.f_min = max(self.f_min, (self.sample_rate / (self.n_fft - 1)) * 2)
         self.default_f_max = min(
             default(self.f_max, self.sample_rate // 2), self.sample_rate // 2
         )
@@ -105,7 +105,6 @@ class AudioProcessor(Model):
             f_min=self.cfg.f_min,
             f_max=self.cfg.f_max,
             mel_scale=self.cfg.mel_scale,
-            onesided=self.cfg.onesided,
             normalized=self.cfg.normalized,
         )
         self._mel_rscale = torchaudio.transforms.InverseMelScale(
@@ -122,17 +121,15 @@ class AudioProcessor(Model):
             (torch.hann_window(self.cfg.win_length) if window is None else window),
         )
 
-
-    
     def compute_mel(
         self,
         wave: Tensor,
-        raw_mel_only: bool = False,
         eps: float = 1e-5,
+        raw_mel_only: bool = False,
         *,
         _recall: bool = False,
     ) -> Tensor:
-        """Returns: [B, M, T]"""
+        """Returns: (M, T) or (B, M, T) if batched"""
         try:
             mel_tensor = self._mel_spec(wave.to(self.device))  # [M, T]
             if not raw_mel_only:
@@ -203,13 +200,16 @@ class AudioProcessor(Model):
             rms_ = []
             for i in range(B):
                 _t = _comp_rms_helper(i, audio, mel)
-                _r = librosa.feature.rms(**_t, **rms_kwargs)[
-                    0
-                ]
+                _r = librosa.feature.rms(**_t, **rms_kwargs)[0]
                 rms_.append(_r)
             return self.from_numpy_batch(rms_, default_device, default_dtype).squeeze()
 
-    def pitch_shift(self, audio: torch.Tensor, sample_rate: Optional[int] = None, n_steps: float = 2.0):
+    def pitch_shift(
+        self,
+        audio: torch.Tensor,
+        sample_rate: Optional[int] = None,
+        n_steps: float = 2.0,
+    ):
         """
         Shifts the pitch of an audio tensor by `n_steps` semitones.
 
@@ -225,21 +225,25 @@ class AudioProcessor(Model):
         src_dtype = audio.dtype
         audio = audio.squeeze()
         sample_rate = default(sample_rate, self.cfg.sample_rate)
+
         def _shift_one(wav):
             wav_np = self.to_numpy_safe(wav)
-            shifted_np = librosa.effects.pitch_shift(wav_np, sr=sample_rate, n_steps=n_steps)
+            shifted_np = librosa.effects.pitch_shift(
+                wav_np, sr=sample_rate, n_steps=n_steps
+            )
             return torch.from_numpy(shifted_np)
 
         if audio.ndim == 1:
             return _shift_one(audio).to(device=src_device, dtype=src_dtype)
-        return torch.stack([_shift_one(a) for a in audio]).to(device=src_device, dtype=src_dtype)
-
+        return torch.stack([_shift_one(a) for a in audio]).to(
+            device=src_device, dtype=src_dtype
+        )
 
     @staticmethod
-    def calc_pitch_fmin(sr:int, frame_length:float):
+    def calc_pitch_fmin(sr: int, frame_length: float):
         """For pitch f_min"""
         return (sr / (frame_length - 1)) * 2
-    
+
     def compute_pitch(
         self,
         audio: Tensor,
@@ -261,8 +265,10 @@ class AudioProcessor(Model):
             B = 1
         sr = default(sr, self.cfg.sample_rate)
         frame_length = default(frame_length, self.cfg.n_fft)
-        fmin = max(default(fmin, self.cfg.f_min), self.calc_pitch_fmin(sr, frame_length))
-        fmax = min(max(default(fmax, self.cfg.default_f_max), fmin+1), sr // 2)
+        fmin = max(
+            default(fmin, self.cfg.f_min), self.calc_pitch_fmin(sr, frame_length)
+        )
+        fmax = min(max(default(fmax, self.cfg.default_f_max), fmin + 1), sr // 2)
         hop_length = default(hop_length, self.cfg.hop_length)
         center = default(center, self.cfg.center)
         yn_kwargs = dict(
@@ -361,7 +367,7 @@ class AudioProcessor(Model):
         The modes available for upsampling are: `nearest`, `linear` (3D-only),
         `bilinear`, `bicubic` (4D-only), `trilinear` (5D-only)
         """
-
+        tensor = tensor.squeeze()
         if tensor.ndim == 2:  # [1, T]
             tensor = tensor.unsqueeze(1)  # [1, 1, T]
         elif tensor.ndim == 1:
@@ -384,7 +390,7 @@ class AudioProcessor(Model):
         hop_length: Optional[int] = None,
         win_length: Optional[int] = None,
         length: Optional[int] = None,
-        center: Optional[bool] = None,
+        center: bool = True,
         normalized: Optional[bool] = None,
         onesided: Optional[bool] = None,
         return_complex: bool = False,
@@ -403,7 +409,7 @@ class AudioProcessor(Model):
                 hop_length=default(hop_length, self.cfg.hop_length),
                 win_length=default(win_length, self.cfg.win_length),
                 window=window,
-                center=default(center, self.cfg.center),
+                center=center,
                 normalized=default(normalized, self.cfg.normalized),
                 onesided=default(onesided, self.cfg.onesided),
                 length=length,
@@ -421,44 +427,48 @@ class AudioProcessor(Model):
         self,
         wave: Tensor,
         length: Optional[int] = None,
-        *,
-        _recall: bool = False,
+        center: bool = True,
+        n_fft: Optional[int] = None,
+        hop_length: Optional[int] = None,
+        win_length: Optional[int] = None,
+        normalized: Optional[bool] = None,
+        onesided: Optional[bool] = None,
+        return_complex: bool = False,
     ):
-        try:
-            spectrogram = torch.stft(
-                input=wave,
-                n_fft=self.cfg.n_fft,
-                hop_length=self.cfg.hop_length,
-                win_length=self.cfg.win_length,
-                window=self.window,
-                center=self.cfg.center,
-                pad_mode="reflect",
-                normalized=self.cfg.normalized,
-                onesided=self.cfg.onesided,
-                return_complex=True,
-            )
-            return torch.istft(
-                spectrogram
-                * torch.full(
-                    spectrogram.size(),
-                    fill_value=1,
-                    device=spectrogram.device,
-                ),
-                n_fft=self.cfg.n_fft,
-                hop_length=self.cfg.hop_length,
-                win_length=self.cfg.win_length,
-                window=self.window,
-                length=length,
-                center=self.cfg.center,
-                normalized=self.cfg.normalized,
-                onesided=self.cfg.onesided,
-                return_complex=False,
-            )
-        except RuntimeError as e:
-            if not _recall and wave.device != self.window.device:
-                self.window = self.window.to(wave.device)
-                return self.istft_norm(wave, length, _recall=True)
-            raise e
+
+        if win_length is not None and win_length != self.cfg.win_length:
+            window = torch.hann_window(win_length, device=wave.device)
+        else:
+            window = self.window
+        spectrogram = torch.stft(
+            input=wave,
+            n_fft=default(n_fft, self.cfg.n_fft),
+            hop_length=default(hop_length, self.cfg.hop_length),
+            win_length=default(win_length, self.cfg.win_length),
+            window=window,
+            center=center,
+            pad_mode="reflect",
+            normalized=default(normalized, self.cfg.normalized),
+            onesided=default(onesided, self.cfg.onesided),
+            return_complex=True,
+        )
+        return torch.istft(
+            spectrogram
+            * torch.full(
+                spectrogram.size(),
+                fill_value=1,
+                device=spectrogram.device,
+            ),
+            n_fft=default(n_fft, self.cfg.n_fft),
+            hop_length=default(hop_length, self.cfg.hop_length),
+            win_length=default(win_length, self.cfg.win_length),
+            window=self.window,
+            length=length,
+            center=center,
+            normalized=default(normalized, self.cfg.normalized),
+            onesided=default(onesided, self.cfg.onesided),
+            return_complex=return_complex,
+        )
 
     def load_audio(
         self,

{lt_tensor-0.0.1a35.dist-info → lt_tensor-0.0.1a36.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lt-tensor
-Version: 0.0.1a35
+Version: 0.0.1a36
 Summary: General utilities for PyTorch and others. Built for general use.
 Home-page: https://github.com/gr1336/lt-tensor/
 Author: gr1336

{lt_tensor-0.0.1a35.dist-info → lt_tensor-0.0.1a36.dist-info}/RECORD

@@ -1,4 +1,4 @@
-lt_tensor/__init__.py,sha256=4NqhrI_O5q4YQMBpyoLtNUUbBnnbWkO92GE1hxHcrd8,441
+lt_tensor/__init__.py,sha256=nBbiGH1byHU0aTTKKorRj8MIEO2oEMBXl7kt5DOCatU,441
 lt_tensor/config_templates.py,sha256=F9UvL8paAjkSvio890kp8WznpYeI50pYnm9iqQroBxk,2797
 lt_tensor/losses.py,sha256=Heco_WyoC1HkNkcJEircOAzS9umusATHiNAG-FKGyzc,8918
 lt_tensor/lr_schedulers.py,sha256=6_vcfaPHrozfH3wvmNEdKSFYl6iTIijYoHL8vuG-45U,7651
@@ -25,16 +25,19 @@ lt_tensor/model_zoo/activations/alias_free/resample.py,sha256=3iM4fNr9fLNXXMyXvz
 lt_tensor/model_zoo/activations/snake/__init__.py,sha256=AtOAbJuMinxmKkppITGMzRbcbPQaALnl9mCtl1c3x0Q,4356
 lt_tensor/model_zoo/audio_models/__init__.py,sha256=WwiP9MekJreMOfKPWLl24VkRJIpLk6hhL8ch0aKgOss,103
 lt_tensor/model_zoo/audio_models/resblocks.py,sha256=u-foHxaFDUICjxSkpyHXljQYQG9zMxVYaOGqLR_nJ-k,7978
-lt_tensor/model_zoo/audio_models/bigvgan/__init__.py,sha256=Dpt_3JXUToldxQrZx4a1gfI-awsLIVipAXqWm4lzBzM,8495
+lt_tensor/model_zoo/audio_models/bigvgan/__init__.py,sha256=4EZG8Non75dHoDCizMHbMTvPrKwdUlPYGHc7hkfT_nw,8526
 lt_tensor/model_zoo/audio_models/diffwave/__init__.py,sha256=PDuDYN1omD1RoAXcmxH3tEgfAuM3ZHAWzimD6ElMqEQ,9073
-lt_tensor/model_zoo/audio_models/hifigan/__init__.py,sha256=3HVfEreQ4NqYIC9AWEkmL4ePcIbR1kTyH0cBG8u_Jik,6387
+lt_tensor/model_zoo/audio_models/hifigan/__init__.py,sha256=ITSXHg3c0Um1P2HaPaXkQKI7meG5Ne60wTbyyYju3hY,6360
 lt_tensor/model_zoo/audio_models/istft/__init__.py,sha256=blICjLX_z_IFmR3_TCz_dJiSayLYGza9eG6fd9aKyvE,7448
 lt_tensor/model_zoo/losses/__init__.py,sha256=B9RAUxBiOZwooztnij1oLeRwZ7_MjnN3mPoum7saD6s,59
-lt_tensor/model_zoo/losses/discriminators.py,sha256=HBO7jwCsUGsYfSz-JZPZccuYLnto6jfZs3Ve5j51JQE,24247
+lt_tensor/model_zoo/losses/discriminators.py,sha256=o4cicNdOv0jH3ink7jTNeDqOnwmkmRtEj9E7IUIGnEI,31866
+lt_tensor/model_zoo/losses/CQT/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+lt_tensor/model_zoo/losses/CQT/transforms.py,sha256=Vkid0J9dqLnlINfyyUlQf-qB3gOQAgU7W9j7xLOjDFw,13218
+lt_tensor/model_zoo/losses/CQT/utils.py,sha256=twGw6FVD7V5Ksfx_1BUEN3EP1tAS6wo-9LL3VnuHB8c,16751
 lt_tensor/processors/__init__.py,sha256=Pvxhh0KR65zLCgUd53_k5Z0y5JWWcO0ZBXFK9rv0o5w,109
-lt_tensor/processors/audio.py,sha256=HNr1GS-6M2q0Rda4cErf5y2Jlc9f4jD58FvpX2ua9d4,18369
-lt_tensor-0.0.1a35.dist-info/licenses/LICENSE,sha256=TbiyJWLgNqqgqhfCnrGwFIxy7EqGNrIZZcKhHrefcuU,11354
-lt_tensor-0.0.1a35.dist-info/METADATA,sha256=0FrtLNnbU49bKOlyshasXPZOZ90Sok03XkXbtxP4VMI,1062
-lt_tensor-0.0.1a35.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-lt_tensor-0.0.1a35.dist-info/top_level.txt,sha256=35FuhFeXnUyvHWdbVHGPh0hS8euofafnJ_GJAVSF4Kk,10
-lt_tensor-0.0.1a35.dist-info/RECORD,,
+lt_tensor/processors/audio.py,sha256=3YzyEpMwh124rb1KMAly62qweeruF200BnM-vQIbzy0,18645
+lt_tensor-0.0.1a36.dist-info/licenses/LICENSE,sha256=TbiyJWLgNqqgqhfCnrGwFIxy7EqGNrIZZcKhHrefcuU,11354
+lt_tensor-0.0.1a36.dist-info/METADATA,sha256=mTmnoWn8EG48j_VOM3rr_8RLLgaxB5pWZE1tkPdFrac,1062
+lt_tensor-0.0.1a36.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+lt_tensor-0.0.1a36.dist-info/top_level.txt,sha256=35FuhFeXnUyvHWdbVHGPh0hS8euofafnJ_GJAVSF4Kk,10
+lt_tensor-0.0.1a36.dist-info/RECORD,,