Load model from hub

app.py

@@ -25,19 +25,9 @@ import gradio as gr
 
 import rembg
 
-def main():
+from huggingface_hub import hf_hub_download
 
-    # ============= model loading ==========
-    def load_model(device):
-        experiment_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 
-                                       "model_file", "objaverse")
-        # load cfg
-        training_cfg = OmegaConf.load(os.path.join(experiment_path, ".hydra", "config.yaml"))
-        # load model
-        model = GaussianSplatPredictor(training_cfg)
-        ckpt_loaded = torch.load(os.path.join(experiment_path, "model_latest.pth"), map_location=device)
-        model.load_state_dict(ckpt_loaded["model_state_dict"])
-        return model, training_cfg
+def main():
 
     if torch.cuda.is_available():
         device = "cuda:0"
@@ -45,7 +35,20 @@ def main():
         device = "cpu"
     torch.cuda.set_device(device)
 
-    model, model_cfg = load_model(device)
+    model_cfg = OmegaConf.load(
+        os.path.join(
+            os.path.dirname(os.path.abspath(__file__)), 
+                    "config.yaml"
+                    ))
+
+    model_path = hf_hub_download(repo_id="szymanowiczs/splatter-image-multi-category-v1", 
+                                 filename="model_latest.pth")
+
+    
+    model = GaussianSplatPredictor(model_cfg)
+
+    ckpt_loaded = torch.load(model_path, map_location=device)
+    model.load_state_dict(ckpt_loaded["model_state_dict"])
     model.to(device)
 
     # ============= image preprocessing =============

config.yaml

@@ -0,0 +1,66 @@
+wandb:
+  project: gs_pred
+cam_embd:
+  embedding: null
+  encode_embedding: null
+  dimension: 0
+  method: null
+general:
+  device: 0
+  random_seed: 0
+  num_devices: 2
+  mixed_precision: true
+data:
+  training_resolution: 128
+  fov: 49.134342641202636
+  subset: -1
+  input_images: 1
+  znear: 0.8
+  zfar: 3.2
+  category: objaverse
+  white_background: true
+  origin_distances: false
+opt:
+  iterations: 50001
+  base_lr: 6.34584421e-05
+  batch_size: 16
+  betas:
+  - 0.9
+  - 0.999
+  loss: l2
+  imgs_per_obj: 4
+  ema:
+    use: true
+    update_every: 10
+    update_after_step: 100
+    beta: 0.9999
+  lambda_lpips: 0.33814373
+  start_lpips_after: 0
+  step_lr_at: -1
+model:
+  max_sh_degree: 1
+  inverted_x: false
+  inverted_y: true
+  name: SingleUNet
+  opacity_scale: 1.0
+  opacity_bias: -2.0
+  scale_scale: 0.01
+  scale_bias: 0.02
+  xyz_scale: 0.1
+  xyz_bias: 0.0
+  depth_scale: 1.0
+  depth_bias: 0.0
+  network_without_offset: false
+  network_with_offset: true
+  attention_resolutions:
+  - 16
+  cross_view_attention: true
+  isotropic: false
+  base_dim: 128
+  num_blocks: 4
+logging:
+  ckpt_iterations: 1000
+  val_log: 10000
+  loss_log: 10
+  loop_log: 10000
+  render_log: 10000