Update mast3r/demo.py

mast3r/demo.py

@@ -212,13 +212,19 @@ def set_scenegraph_options(inputfiles, win_cyclic, refid, scenegraph_type):
                           maximum=num_files - 1, step=1, visible=scenegraph_type == 'oneref')
     return win_col, winsize, win_cyclic, refid
 
-
+def recon_fun(get_reconstructed_scene, tmpdirname, gradio_delete_cache, model, device,
+                                  silent, image_size):
+    +import spaces
+    functools.partial(get_reconstructed_scene, tmpdirname, gradio_delete_cache, model, device,
+                                  silent, image_size)
+    +@spaces.GPU
+    
 def main_demo(tmpdirname, model, device, image_size, server_name, server_port, silent=False,
               share=False, gradio_delete_cache=False):
     if not silent:
         print('Outputing stuff in', tmpdirname)
 
-    recon_fun = functools.partial(get_reconstructed_scene, tmpdirname, gradio_delete_cache, model, device,
+    recon_fun = recon_fun(get_reconstructed_scene, tmpdirname, gradio_delete_cache, model, device,
                                   silent, image_size)
     model_from_scene_fun = functools.partial(get_3D_model_from_scene, silent)