requirements

WavLM.py

@@ -48,7 +48,7 @@ class WavLM_wrapper(nn.Module):
             )
         if not os.path.exists(model_path):
             self.download_model(model_path, model_size)
-        checkpoint = torch.load(model_path)
+        checkpoint = torch.load(model_path, weights_only=False)
         cfg = WavLMConfig(checkpoint["cfg"])
         self.cfg = cfg
         self.model = WavLM(cfg)

sgm/models/diffusion.py

@@ -119,7 +119,7 @@ class DiffusionEngine(pl.LightningModule):
                 pattern_to_remove=pattern_to_remove,
             )
             if separate_unet_ckpt is not None:
-                sd = torch.load(separate_unet_ckpt)["state_dict"]
+                sd = torch.load(separate_unet_ckpt, weights_only=False)["state_dict"]
                 if remove_keys_from_unet_weights is not None:
                     for k in list(sd.keys()):
                         for remove_key in remove_keys_from_unet_weights:
@@ -190,7 +190,7 @@ class DiffusionEngine(pl.LightningModule):
 
     def load_bad_model_weights(self, path: str) -> None:
         print(f"Restoring bad model from {path}")
-        state_dict = torch.load(path, map_location="cpu")
+        state_dict = torch.load(path, map_location="cpu", weights_only=False)
         new_dict = {}
         for k, v in state_dict["module"].items():
             if "learned_mask" in k:
@@ -221,13 +221,13 @@ class DiffusionEngine(pl.LightningModule):
     ) -> None:
         print(f"Restoring from {path}")
         if path.endswith("ckpt"):
-            sd = torch.load(path, map_location="cpu")["state_dict"]
+            sd = torch.load(path, map_location="cpu", weights_only=False)["state_dict"]
         elif path.endswith("pt"):
-            sd = torch.load(path, map_location="cpu")["module"]
+            sd = torch.load(path, map_location="cpu", weights_only=False)["module"]
             # Remove leading _forward_module from keys
             sd = {k.replace("_forward_module.", ""): v for k, v in sd.items()}
         elif path.endswith("bin"):
-            sd = torch.load(path, map_location="cpu")
+            sd = torch.load(path, map_location="cpu", weights_only=False)
             # Remove leading _forward_module from keys
             sd = {k.replace("_forward_module.", ""): v for k, v in sd.items()}
         elif path.endswith("safetensors"):