Improve Hugging Face integration

vggsfm/models/vggsfm.py

@@ -6,15 +6,30 @@
 
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
 
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Dict
 
 from hydra.utils import instantiate
 
-
-class VGGSfM(nn.Module):
-    def __init__(self, TRACK: Dict, CAMERA: Dict, TRIANGULAE: Dict, cfg=None):
+from omegaconf import OmegaConf
+from omegaconf.dictconfig import DictConfig
+
+from huggingface_hub import PyTorchModelHubMixin
+
+
+class VGGSfM(nn.Module,
+             PyTorchModelHubMixin,
+             repo_url="https://github.com/facebookresearch/vggsfm",
+             pipeline_tag="image-to-3d",
+             license="cc-by-nc-sa-4.0",
+             coders={
+                DictConfig : (
+                    lambda x: OmegaConf.to_container(x, resolve=True),  # Encoder: how to convert a `DictConfig` to a valid jsonable value?
+                    lambda data: OmegaConf.create(data),  # Decoder: how to reconstruct a `DictConfig` from a dictionary?
+                ),
+            }
+    ):
+    def __init__(self, TRACK: DictConfig, CAMERA: DictConfig, TRIANGULAE: DictConfig, cfg: DictConfig = None):
         """
         Initializes a VGGSfM model
 
@@ -33,19 +48,3 @@ def __init__(self, TRACK: Dict, CAMERA: Dict, TRIANGULAE: Dict, cfg=None):
 
         # models.Triangulator
         self.triangulator = instantiate(TRIANGULAE, _recursive_=False, cfg=cfg)
-
-    def from_pretrained(self, model_name):
-        try:
-            from huggingface_hub import hf_hub_download
-
-            ckpt_path = hf_hub_download(
-                repo_id="facebook/VGGSfM", filename=model_name + ".bin"
-            )
-            checkpoint = torch.load(ckpt_path)
-        except:
-            # In case the model is not hosted on huggingface
-            # or the user cannot import huggingface_hub correctly
-            _VGGSFM_URL = "https://huggingface.co/facebook/VGGSfM/resolve/main/vggsfm_v2_0_0.bin"
-            checkpoint = torch.hub.load_state_dict_from_url(_VGGSFM_URL)
-
-        self.load_state_dict(checkpoint, strict=True)

vggsfm/runners/runner.py

@@ -112,11 +112,13 @@ def build_vggsfm_model(self):
                 self.cfg.MODEL, _recursive_=False, cfg=self.cfg
             )
 
-        if self.cfg.auto_download_ckpt:
-            vggsfm.from_pretrained(self.cfg.model_name)
-        else:
-            checkpoint = torch.load(self.cfg.resume_ckpt)
-            vggsfm.load_state_dict(checkpoint, strict=True)
+        # if self.cfg.auto_download_ckpt:
+        #     vggsfm.from_pretrained(self.cfg.model_name)
+        # else:
+        from huggingface_hub import hf_hub_download
+        filepath = hf_hub_download(repo_id="facebook/VGGSfM", filename="vggsfm_v2_0_0.bin", repo_type="model")
+        checkpoint = torch.load(filepath, map_location="cpu")
+        vggsfm.load_state_dict(checkpoint, strict=True)
         self.vggsfm_model = vggsfm.to(self.device).eval()
         print("VGGSfM built successfully")