From a09e249ead8150beeca31510e485e1d7788e4a80 Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Thu, 27 Jun 2024 03:28:45 +0300
Subject: [PATCH] support safetensors

---
 nodes.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/nodes.py b/nodes.py
index 63c141a..c16ba43 100644
--- a/nodes.py
+++ b/nodes.py
@@ -10,6 +10,7 @@ from .model.ccsr_stage1 import ControlLDM
 from .utils.common import instantiate_from_config, load_state_dict
 
 import comfy.model_management
+import comfy.utils
 import folder_paths
 from nodes import ImageScaleBy
 from nodes import ImageScale
@@ -75,7 +76,7 @@ class CCSR_Upscale:
             config = OmegaConf.load(config_path)
             self.model = instantiate_from_config(config)
             
-            load_state_dict(self.model, torch.load(ccsr_model, map_location="cpu"), strict=True)
+            load_state_dict(self.model, comfy.utils.load_torch_file(ccsr_model), strict=True)
             # reload preprocess model if specified
 
             self.model.freeze()