Enable safe loading for upscale models.

2 years ago · 587f89fe5a
parent acff543d66
commit 587f89fe5a
1 changed files with 1 additions and 1 deletions
--- a/comfy_extras/nodes_upscale_model.py
+++ b/comfy_extras/nodes_upscale_model.py
@ -17,7 +17,7 @@ class UpscaleModelLoader:

    def load_model(self, model_name):
        model_path = folder_paths.get_full_path("upscale_models", model_name)
-        sd = comfy.utils.load_torch_file(model_path)
+        sd = comfy.utils.load_torch_file(model_path, safe_load=True)
        out = model_loading.load_state_dict(sd).eval()
        return (out, )