change to bigvan

audioldm/bigvgan/__init__.py

@@ -0,0 +1,22 @@
+import torch
+import bigvgan
+from huggingface_hub import hf_hub_download
+
+class BigVGANVocoder:
+    def __init__(self, device='cuda', use_cuda_kernel=False):
+        # Load the pretrained model
+        self.model = bigvgan.BigVGAN.from_pretrained(
+            'nvidia/bigvgan_v2_44khz_128band_512x',
+            use_cuda_kernel=use_cuda_kernel
+        )
+        self.model.remove_weight_norm()
+        self.model.eval().to(device)
+        self.device = device
+        self.h = self.model.h  # This holds config like sampling_rate, etc.
+
+    @torch.no_grad()
+    def infer_waveform(self, mel):
+        # mel shape: [B, n_mels, T], BigVGAN expects mel at model.h.n_mels, typically 128
+        mel = mel.to(self.device)
+        wav_gen = self.model(mel)
+        return wav_gen.squeeze(1)  # Returns [B, T]

audioldm/ldm.py

@@ -234,14 +234,16 @@ class LatentDiffusion(DDPM):
         return self.first_stage_model.decode(z)
 
     def mel_spectrogram_to_waveform(self, mel):
-        # Mel: [bs, 1, t-steps, fbins]
+        # Originally: self.first_stage_model.vocoder(mel)
+        # We'll call BigVGAN here instead.
         if len(mel.size()) == 4:
-            mel = mel.squeeze(1)
-        mel = mel.permute(0, 2, 1)
-        waveform = self.first_stage_model.vocoder(mel)
+            mel = mel.squeeze(1)  # shape: [B, time, n_mels]
+        mel = mel.permute(0, 2, 1)  # BigVGAN expects [B, n_mels, T]
+        waveform = self.vocoder.infer_waveform(mel)  # using BigVGAN vocoder now
         waveform = waveform.cpu().detach().numpy()
         return waveform
 
+
     @torch.no_grad()
     def encode_first_stage(self, x):
         return self.first_stage_model.encode(x)

audioldm/pipeline.py

@@ -10,7 +10,11 @@ from audioldm import LatentDiffusion, seed_everything
 from audioldm.utils import default_audioldm_config, get_duration, get_bit_depth, get_metadata, download_checkpoint
 from audioldm.audio import wav_to_fbank, TacotronSTFT, read_wav_file
 from audioldm.latent_diffusion.ddim import DDIMSampler
+from audioldm.bigvgan import BigVGANVocoder
 from einops import repeat
+from scipy.signal import convolve
+import numpy as np
+
 import os
 
 def make_batch_for_text_to_audio(text, waveform=None, fbank=None, batchsize=1):
@@ -89,6 +93,9 @@ def build_model(
     latent_diffusion = latent_diffusion.to(device)
 
     latent_diffusion.cond_stage_model.embed_mode = "text"
+        # Here is where you add the BigVGAN vocoder initialization
+    
+    latent_diffusion.vocoder = BigVGANVocoder(device='cuda', use_cuda_kernel=False)
     return latent_diffusion
 
 def duration_to_latent_t_size(duration):

audioldm/utils.py

@@ -106,10 +106,10 @@ def default_audioldm_config(model_name="audioldm-s-full"):
             "root": "/mnt/fast/nobackup/users/hl01486/projects/general_audio_generation/AudioLDM-python/config/default/latent_diffusion.yaml",
         },
         "preprocessing": {
-            "audio": {"sampling_rate": 16000, "max_wav_value": 32768},
+            "audio": {"sampling_rate": 44100, "max_wav_value": 32768},
             "stft": {"filter_length": 1024, "hop_length": 160, "win_length": 1024},
             "mel": {
-                "n_mel_channels": 64,
+                "n_mel_channels": 128,
                 "mel_fmin": 0,
                 "mel_fmax": 8000,
                 "freqm": 0,

audioldm/variational_autoencoder/autoencoder.py

@@ -3,7 +3,11 @@ from audioldm.latent_diffusion.ema import *
 from audioldm.variational_autoencoder.modules import Encoder, Decoder
 from audioldm.variational_autoencoder.distributions import DiagonalGaussianDistribution
 
-from audioldm.hifigan.utilities import get_vocoder, vocoder_infer
+#from audioldm.hifigan.utilities import get_vocoder, vocoder_infer
+from audioldm.bigvgan import BigVGANVocoder
+
+# After you create your latent_diffusion instance:
+latent_diffusion.vocoder = BigVGANVocoder(device='cuda', use_cuda_kernel=False)
 
 
 class AutoencoderKL(nn.Module):
@@ -36,7 +40,8 @@ class AutoencoderKL(nn.Module):
         self.quant_conv = torch.nn.Conv2d(2 * ddconfig["z_channels"], 2 * embed_dim, 1)
         self.post_quant_conv = torch.nn.Conv2d(embed_dim, ddconfig["z_channels"], 1)
 
-        self.vocoder = get_vocoder(None, "cpu")
+        self.vocoder = BigVGANVocoder(device=device)
+        #self.vocoder = get_vocoder(None, "cpu")
         self.embed_dim = embed_dim
 
         if monitor is not None:
@@ -65,7 +70,8 @@ class AutoencoderKL(nn.Module):
 
     def decode_to_waveform(self, dec):
         dec = dec.squeeze(1).permute(0, 2, 1)
-        wav_reconstruction = vocoder_infer(dec, self.vocoder)
+        wav_reconstruction = vocoder.infer_waveform(mels)
+        #wav_reconstruction = vocoder_infer(dec, self.vocoder)
         return wav_reconstruction
 
     def forward(self, input, sample_posterior=True):

requirements.txt

@@ -29,4 +29,5 @@ tqdm==4.63.1
 wandb==0.12.14
 ipython==8.12.0
 gradio==4.3.0
-wavio==0.0.7
+wavio==0.0.7
+bigvgan==2.4.1