latentfix

audioldm/hifigan/utilities.py

@@ -64,13 +64,13 @@ def get_param_num(model):
     return num_param
 
 
-def get_vocoder(config, device):
-    config = hifigan.AttrDict(HIFIGAN_16K_64)
-    vocoder = hifigan.Generator(config)
-    vocoder.eval()
-    vocoder.remove_weight_norm()
-    vocoder.to(device)
-    return vocoder
+# def get_vocoder(config, device):
+#     config = hifigan.AttrDict(HIFIGAN_16K_64)
+#     vocoder = hifigan.Generator(config)
+#     vocoder.eval()
+#     vocoder.remove_weight_norm()
+#     vocoder.to(device)
+#     return vocoder
 
 
 def vocoder_infer(mels, vocoder, lengths=None):

audioldm/pipeline.py

@@ -87,7 +87,11 @@ def build_model(
     resume_from_checkpoint = ckpt_path
 
     checkpoint = torch.load(resume_from_checkpoint, map_location=device)
-    latent_diffusion.load_state_dict(checkpoint["state_dict"])
+    state_dict = checkpoint["state_dict"]
+    # Filter out vocoder keys
+    filtered_state_dict = {k: v for k, v in state_dict.items() if not k.startswith("vocoder.")}
+    latent_diffusion.load_state_dict(filtered_state_dict, strict=False)
+    
 
     latent_diffusion.eval()
     latent_diffusion = latent_diffusion.to(device)