fix

SUPIR/util.py

@@ -8,6 +8,8 @@ from omegaconf import OmegaConf
 from sgm.util import instantiate_from_config
 from huggingface_hub import hf_hub_download
 
+from safetensors.torch import load_file as load_safetensors
+
 def get_state_dict(d):
     return d.get('state_dict', d)
 
@@ -30,12 +32,14 @@ def create_model(config_path):
     print(f'Loaded model config from [{config_path}]')
     return model
 
+
 def resolve_ckpt_path(path_or_hub):
+    path_or_hub = path_or_hub.strip()
+
     if os.path.exists(path_or_hub):
-        return path_or_hub  # Local path
+        return path_or_hub
 
-    # Handle Hugging Face Hub repo/file pattern
-    if "/" in path_or_hub and path_or_hub.endswith((".ckpt", ".safetensors", ".bin")):
+    if "/" in path_or_hub and path_or_hub.endswith((".ckpt", ".pt", ".bin", ".safetensors")):
         parts = path_or_hub.split("/")
         repo_id = "/".join(parts[:2])
         filename = "/".join(parts[2:])
@@ -43,24 +47,31 @@ def resolve_ckpt_path(path_or_hub):
 
     raise FileNotFoundError(f"Could not resolve checkpoint path: {path_or_hub}")
 
+def load_checkpoint(path):
+    if path.endswith(".safetensors"):
+        return load_safetensors(path, device='cpu')
+    else:
+        return torch.load(path, map_location='cpu')
+
 def create_SUPIR_model(config_path, SUPIR_sign=None, load_default_setting=False):
     config = OmegaConf.load(config_path)
     model = instantiate_from_config(config.model).cpu()
     print(f'Loaded model config from [{config_path}]')
 
-    if config.get("SDXL_CKPT") is not None:
+    if config.get("SDXL_CKPT"):
         path = resolve_ckpt_path(config.SDXL_CKPT)
-        model.load_state_dict(torch.load(path, map_location='cpu'), strict=False)
+        model.load_state_dict(load_checkpoint(path), strict=False)
 
-    if config.get("SUPIR_CKPT") is not None:
+    if config.get("SUPIR_CKPT"):
         path = resolve_ckpt_path(config.SUPIR_CKPT)
-        model.load_state_dict(torch.load(path, map_location='cpu'), strict=False)
+        model.load_state_dict(load_checkpoint(path), strict=False)
 
     if SUPIR_sign is not None:
-        assert SUPIR_sign in ['F', 'Q']
+        assert SUPIR_sign in ['F', 'Q'], "SUPIR_sign must be 'F' or 'Q'"
         key = f"SUPIR_CKPT_{SUPIR_sign}"
-        path = resolve_ckpt_path(config[key])
-        model.load_state_dict(torch.load(path, map_location='cpu'), strict=False)
+        if config.get(key):
+            path = resolve_ckpt_path(config[key])
+            model.load_state_dict(load_checkpoint(path), strict=False)
 
     if load_default_setting:
         return model, config.default_setting
@@ -68,11 +79,10 @@ def create_SUPIR_model(config_path, SUPIR_sign=None, load_default_setting=False)
 
 def load_QF_ckpt(config_path):
     config = OmegaConf.load(config_path)
-    ckpt_F = torch.load(resolve_ckpt_path(config.SUPIR_CKPT_F), map_location='cpu')
-    ckpt_Q = torch.load(resolve_ckpt_path(config.SUPIR_CKPT_Q), map_location='cpu')
+    ckpt_F = load_checkpoint(resolve_ckpt_path(config.SUPIR_CKPT_F))
+    ckpt_Q = load_checkpoint(resolve_ckpt_path(config.SUPIR_CKPT_Q))
     return ckpt_Q, ckpt_F
 
-
 def PIL2Tensor(img, upsacle=1, min_size=1024, fix_resize=None):
     '''
     PIL.Image -> Tensor[C, H, W], RGB, [-1, 1]