fix: Update app.py

app.py

@@ -3,6 +3,7 @@ import time
 import torch
 import numpy as np
 import gradio as gr
+import urllib.parse
 import tempfile
 import subprocess
 from dust3r.losses import L21
@@ -11,6 +12,7 @@ from spann3r.datasets import Demo
 from torch.utils.data import DataLoader
 import trimesh
 from scipy.spatial.transform import Rotation
+import spaces
 
 # Default values
 DEFAULT_CKPT_PATH = './checkpoints/spann3r.pth'
@@ -45,10 +47,22 @@ def cat_meshes(meshes):
     faces = np.concatenate(faces)
     return dict(vertices=vertices, face_colors=colors, faces=faces)
 
+def load_ckpt(model_path_or_url, verbose=True):
+    if verbose:
+        print('... loading model from', model_path_or_url)
+    is_url = urllib.parse.urlparse(model_path_or_url).scheme in ('http', 'https')
+    
+    if is_url:
+        ckpt = torch.hub.load_state_dict_from_url(model_path_or_url, map_location='cpu', progress=verbose)
+    else:
+        ckpt = torch.load(model_path_or_url, map_location='cpu')
+    return ckpt
+
 def load_model(ckpt_path, device):
     model = Spann3R(dus3r_name=DEFAULT_DUST3R_PATH, 
                     use_feat=False).to(device)
-    model.load_state_dict(torch.load(ckpt_path)['model'])
+    
+    model.load_state_dict(load_ckpt(ckpt_path)['model'])
     model.eval()
     return model
 
@@ -91,14 +105,14 @@ def pts3d_to_trimesh(img, pts3d, valid=None):
     assert len(faces) == len(face_colors)
     return dict(vertices=vertices, face_colors=face_colors, faces=faces)
 
+model = load_model(DEFAULT_CKPT_PATH, DEFAULT_DEVICE)
+
+@spaces.GPU
 @torch.no_grad()
-def reconstruct(video_path, conf_thresh, kf_every, voxel_size=0.05, as_pointcloud=False):
+def reconstruct(video_path, conf_thresh, kf_every, as_pointcloud=False):
     # Extract frames from video
     demo_path = extract_frames(video_path)
     
-    # Load model
-    model = load_model(DEFAULT_CKPT_PATH, DEFAULT_DEVICE)
-    
     # Load dataset
     dataset = Demo(ROOT=demo_path, resolution=224, full_video=True, kf_every=kf_every)
     dataloader = DataLoader(dataset, batch_size=1, shuffle=False, num_workers=0)
@@ -168,16 +182,15 @@ iface = gr.Interface(
     inputs=[
         gr.Video(label="Input Video"),
         gr.Slider(0, 1, value=1e-3, label="Confidence Threshold"),
-        gr.Slider(1, 30, step=1, value=10, label="Keyframe Interval"),
-        gr.Slider(0.001, 0.01, value=0.005, step=0.001, label="Voxel Size for Downsampling"),
+        gr.Slider(1, 30, step=1, value=5, label="Keyframe Interval"),
         gr.Checkbox(label="As Pointcloud", value=False)
     ],
     outputs=[
         gr.Model3D(label="3D Model (GLB)", display_mode="solid"),
         gr.Textbox(label="Status")
     ],
-    title="3D Reconstruction from Video",
+    title="3D Reconstruction with Spatial Memory",
 )
 
 if __name__ == "__main__":
-    iface.launch(share=True)
+    iface.launch()