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