fix

SUPIR/util.py

@@ -6,7 +6,7 @@ from PIL import Image
 from torch.nn.functional import interpolate
 from omegaconf import OmegaConf
 from sgm.util import instantiate_from_config
-
+from huggingface_hub import hf_hub_download
 
 def get_state_dict(d):
     return d.get('state_dict', d)
@@ -30,30 +30,44 @@ def create_model(config_path):
     print(f'Loaded model config from [{config_path}]')
     return model
 
+def resolve_ckpt_path(path_or_hub):
+    if os.path.exists(path_or_hub):
+        return path_or_hub  # local path
+    if "/" in path_or_hub and path_or_hub.endswith(".ckpt"):
+        # Assume format: repo_id/path/to/file.ckpt
+        parts = path_or_hub.split("/")
+        repo_id = "/".join(parts[:2])
+        filename = "/".join(parts[2:])
+        return hf_hub_download(repo_id=repo_id, filename=filename)
+    return path_or_hub  # fallback
 
 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.SDXL_CKPT is not None:
-        model.load_state_dict(load_state_dict(config.SDXL_CKPT), strict=False)
-    if config.SUPIR_CKPT is not None:
-        model.load_state_dict(load_state_dict(config.SUPIR_CKPT), strict=False)
+
+    if config.get("SDXL_CKPT") is not None:
+        path = resolve_ckpt_path(config.SDXL_CKPT)
+        model.load_state_dict(torch.load(path, map_location='cpu'), strict=False)
+
+    if config.get("SUPIR_CKPT") is not None:
+        path = resolve_ckpt_path(config.SUPIR_CKPT)
+        model.load_state_dict(torch.load(path, map_location='cpu'), strict=False)
+
     if SUPIR_sign is not None:
         assert SUPIR_sign in ['F', 'Q']
-        if SUPIR_sign == 'F':
-            model.load_state_dict(load_state_dict(config.SUPIR_CKPT_F), strict=False)
-        elif SUPIR_sign == 'Q':
-            model.load_state_dict(load_state_dict(config.SUPIR_CKPT_Q), strict=False)
+        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 load_default_setting:
-        default_setting = config.default_setting
-        return model, default_setting
+        return model, config.default_setting
     return model
 
 def load_QF_ckpt(config_path):
     config = OmegaConf.load(config_path)
-    ckpt_F = torch.load(config.SUPIR_CKPT_F, map_location='cpu')
-    ckpt_Q = torch.load(config.SUPIR_CKPT_Q, map_location='cpu')
+    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')
     return ckpt_Q, ckpt_F