update

bigvgan.py

@@ -257,14 +257,18 @@ class BigVGAN(
         return x
 
     def remove_weight_norm(self):
-        print('Removing weight norm...')
-        for l in self.ups:
-            for l_i in l:
-                remove_weight_norm(l_i)
-        for l in self.resblocks:
-            l.remove_weight_norm()
-        remove_weight_norm(self.conv_pre)
-        remove_weight_norm(self.conv_post)
+        try:
+            print('Removing weight norm...')
+            for l in self.ups:
+                for l_i in l:
+                    remove_weight_norm(l_i)
+            for l in self.resblocks:
+                l.remove_weight_norm()
+            remove_weight_norm(self.conv_pre)
+            remove_weight_norm(self.conv_post)
+        except ValueError:
+            print('[INFO] Model already removed weight norm. Skipping!')
+            pass
 
     ##################################################################
     # additional methods for huggingface_hub support
@@ -304,17 +308,21 @@ class BigVGAN(
         ##################################################################
         # download and load hyperparameters (h) used by BigVGAN
         ##################################################################
-        config_file = hf_hub_download(
-            repo_id=model_id,
-            filename='config.json',
-            revision=revision,
-            cache_dir=cache_dir,
-            force_download=force_download,
-            proxies=proxies,
-            resume_download=resume_download,
-            token=token,
-            local_files_only=local_files_only,
-        )
+        if os.path.isdir(model_id):
+            print("Loading config.json from local directory")
+            config_file = os.path.join(model_id, 'config.json')
+        else:
+            config_file = hf_hub_download(
+                repo_id=model_id,
+                filename='config.json',
+                revision=revision,
+                cache_dir=cache_dir,
+                force_download=force_download,
+                proxies=proxies,
+                resume_download=resume_download,
+                token=token,
+                local_files_only=local_files_only,
+            )
         h = load_hparams_from_json(config_file)
 
         ##################################################################
@@ -347,6 +355,12 @@ class BigVGAN(
                 )
         
         checkpoint_dict = torch.load(model_file, map_location=map_location)
-        model.load_state_dict(checkpoint_dict['generator'])
+
+        try:
+            model.load_state_dict(checkpoint_dict['generator'])
+        except RuntimeError:
+            print(f"[INFO] the pretrained checkpoint does not contain weight norm. Loading the checkpoint after removing weight norm!")
+            model.remove_weight_norm()
+            model.load_state_dict(checkpoint_dict['generator'])
 
         return model