Update diffusionsfm/inference/load_model.py

diffusionsfm/inference/load_model.py

@@ -27,15 +27,17 @@ def load_model(
         device (str): Device to load the model on.
         custom_keys (dict): Dictionary of custom keys to override in the config.
     """
-    if checkpoint is None:
-        checkpoint_path = sorted(glob(osp.join(output_dir, "checkpoints", "*.pth")))[-1]
-    else:
-        if isinstance(checkpoint, int):
-            checkpoint_name = f"ckpt_{checkpoint:08d}.pth"
-        else:
-            checkpoint_name = checkpoint
-        checkpoint_path = osp.join(output_dir, "checkpoints", checkpoint_name)
-    print("Loading checkpoint", osp.basename(checkpoint_path))
+    # if checkpoint is None:
+    #     checkpoint_path = sorted(glob(osp.join(output_dir, "checkpoints", "*.pth")))[-1]
+    # else:
+    #     if isinstance(checkpoint, int):
+    #         checkpoint_name = f"ckpt_{checkpoint:08d}.pth"
+    #     else:
+    #         checkpoint_name = checkpoint
+    #     checkpoint_path = osp.join(output_dir, "checkpoints", checkpoint_name)
+    _URL = "https://huggingface.co/qitaoz/DiffusionSfM/resolve/main/ckpt_00800000.pth"
+    data = torch.hub.load_state_dict_from_url(_URL)
+    print("Loading checkpoint", _URL)
 
     cfg = OmegaConf.load(osp.join(output_dir, "hydra", "config.yaml"))
     if custom_keys is not None:
@@ -78,7 +80,7 @@ def load_model(
             cond_depth_mask=cfg.model.get("cond_depth_mask", False),
         ).to(device)
 
-    data = torch.load(checkpoint_path)
+    # data = torch.load(checkpoint_path)
     state_dict = {}
     for k, v in data["state_dict"].items():
         include = True