Upload folder using huggingface_hub

.idea/workspace.xml

@@ -54,7 +54,7 @@
       <option name="number" value="Default" />
       <option name="presentableId" value="Default" />
       <updated>1760056680813</updated>
-      <workItem from="1760056681869" duration="2244000" />
+      <workItem from="1760056681869" duration="2999000" />
     </task>
     <servers />
   </component>

app.py

@@ -19,7 +19,7 @@ from scripts.camera.visualization.visualize_batch import make_perspective_figure
 
 from huggingface_hub import snapshot_download
 import os
-local_path = snapshot_download(
+_ = snapshot_download(
     repo_id="KangLiao/Puffin",
     repo_type="model",       
     #filename="Puffin-Base.pth",         
@@ -27,7 +27,14 @@ local_path = snapshot_download(
     local_dir_use_symlinks=False,     
     revision="main",                  
 )
-
+_ = snapshot_download(
+    repo_id="wusize/Puffin",
+    repo_type="model",
+    #filename="Puffin-Base.pth",
+    local_dir="checkpoints/",
+    local_dir_use_symlinks=False,
+    revision="main",
+)
 
 NUM = r"[+-]?(?:\d+(?:\.\d+)?|\.\d+)(?:[eE][+-]?\d+)?"
 CAM_PATTERN = re.compile(r"(?:camera parameters.*?:|roll.*?:)\s*("+NUM+r")\s*,\s*("+NUM+r")\s*,\s*("+NUM+r")", re.IGNORECASE|re.DOTALL)
@@ -46,7 +53,8 @@ config = Config.fromfile(config)
 model = BUILDER.build(config.model).cuda().bfloat16().eval()
 checkpoint_path = "checkpoints/Puffin-Base.pth"
 checkpoint = torch.load(checkpoint_path)
-info = model.load_state_dict(checkpoint, strict=False)
+_ = model.load_state_dict(checkpoint, strict=False)
+_ = model.vae.load_state_dict(torch.load('checkpoints/vae.pth'), strict=True)
 
 def fig_to_image(fig):
     buf = io.BytesIO()