Use weights_only for loading (#3427)

Co-authored-by: Manuel Schmid <9307310+mashb1t@users.noreply.github.com>

ldm_patched/ldm/modules/encoders/noise_aug_modules.py

@@ -8,7 +8,7 @@ class CLIPEmbeddingNoiseAugmentation(ImageConcatWithNoiseAugmentation):
         if clip_stats_path is None:
             clip_mean, clip_std = torch.zeros(timestep_dim), torch.ones(timestep_dim)
         else:
-            clip_mean, clip_std = torch.load(clip_stats_path, map_location="cpu")
+            clip_mean, clip_std = torch.load(clip_stats_path, map_location="cpu", weights_only=True)
         self.register_buffer("data_mean", clip_mean[None, :], persistent=False)
         self.register_buffer("data_std", clip_std[None, :], persistent=False)
         self.time_embed = Timestep(timestep_dim)

ldm_patched/modules/sd1_clip.py

@@ -326,7 +326,7 @@ def load_embed(embedding_name, embedding_directory, embedding_size, embed_key=No
                 except:
                     embed_out = safe_load_embed_zip(embed_path)
             else:
-                embed = torch.load(embed_path, map_location="cpu")
+                embed = torch.load(embed_path, map_location="cpu", weights_only=True)
     except Exception as e:
         print(traceback.format_exc())
         print()

ldm_patched/pfn/architecture/face/codeformer.py

@@ -377,15 +377,15 @@ class VQAutoEncoder(nn.Module):
         )
 
         if model_path is not None:
-            chkpt = torch.load(model_path, map_location="cpu")
+            chkpt = torch.load(model_path, map_location="cpu", weights_only=True)
             if "params_ema" in chkpt:
                 self.load_state_dict(
-                    torch.load(model_path, map_location="cpu")["params_ema"]
+                    torch.load(model_path, map_location="cpu", weights_only=True)["params_ema"]
                 )
                 logger.info(f"vqgan is loaded from: {model_path} [params_ema]")
             elif "params" in chkpt:
                 self.load_state_dict(
-                    torch.load(model_path, map_location="cpu")["params"]
+                    torch.load(model_path, map_location="cpu", weights_only=True)["params"]
                 )
                 logger.info(f"vqgan is loaded from: {model_path} [params]")
             else:

ldm_patched/pfn/architecture/face/gfpgan_bilinear_arch.py

@@ -273,8 +273,8 @@ class GFPGANBilinear(nn.Module):
         if decoder_load_path:
             self.stylegan_decoder.load_state_dict(
                 torch.load(
-                    decoder_load_path, map_location=lambda storage, loc: storage
-                )["params_ema"]
+                    decoder_load_path, map_location=lambda storage, loc: storage, 
+                weights_only=True)["params_ema"]
             )
         # fix decoder without updating params
         if fix_decoder:

ldm_patched/pfn/architecture/face/gfpganv1_arch.py

@@ -373,8 +373,8 @@ class GFPGANv1(nn.Module):
         if decoder_load_path:
             self.stylegan_decoder.load_state_dict(
                 torch.load(
-                    decoder_load_path, map_location=lambda storage, loc: storage
-                )["params_ema"]
+                    decoder_load_path, map_location=lambda storage, loc: storage, 
+                weights_only=True)["params_ema"]
             )
         # fix decoder without updating params
         if fix_decoder:

ldm_patched/pfn/architecture/face/gfpganv1_clean_arch.py

@@ -284,8 +284,8 @@ class GFPGANv1Clean(nn.Module):
         if decoder_load_path:
             self.stylegan_decoder.load_state_dict(
                 torch.load(
-                    decoder_load_path, map_location=lambda storage, loc: storage
-                )["params_ema"]
+                    decoder_load_path, map_location=lambda storage, loc: storage, 
+                weights_only=True)["params_ema"]
             )
         # fix decoder without updating params
         if fix_decoder: