init project

app.py

@@ -21,8 +21,6 @@ import torch
 #         builtin_print(*args, **kwargs)
 #     builtins.print = print_with_timestamp
 
-device = 'cuda' if torch.cuda.is_available() else 'cpu'
-
 def get_args_parser():
     parser = argparse.ArgumentParser()
     parser_url = parser.add_mutually_exclusive_group()
@@ -53,7 +51,7 @@ if __name__ == '__main__':
     else:
         server_name = '0.0.0.0' if args.local_network else '127.0.0.1'
 
-    pe3r = Models(device=device)
+    pe3r = Models()
 
     with tempfile.TemporaryDirectory(suffix='pe3r_gradio_demo') as tmpdirname:
         if not args.silent:

modules/pe3r/demo.py

@@ -10,6 +10,7 @@ import math
 import gradio
 import os
 import torch
+import spaces
 import numpy as np
 import functools
 import trimesh
@@ -548,7 +549,9 @@ def set_scenegraph_options(inputfiles, winsize, refid, scenegraph_type):
     return winsize, refid
 
 
-def main_demo(tmpdirname, pe3r, device, server_name, server_port, silent=False):
+@spaces.GPU
+def main_demo(tmpdirname, pe3r, server_name, server_port, silent=False):
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
 #    scene, outfile, imgs = get_reconstructed_scene(
 #         outdir=tmpdirname,  pe3r=pe3r, device=device, silent=silent,
 #         filelist=['/home/hujie/pe3r/datasets/mipnerf360_ov/bonsai/black_chair/images/DSCF5590.png',

modules/pe3r/models.py

@@ -10,9 +10,13 @@ from modules.mobilesamv2.promt_mobilesamv2 import ObjectAwareModel
 from modules.mobilesamv2 import sam_model_registry
 
 from sam2.sam2_video_predictor import SAM2VideoPredictor
+import spaces
+import torch
 
 class Models:
-    def __init__(self, device):
+    @spaces.GPU
+    def __init__(self, device='cpu'):
+        device = 'cuda' if torch.cuda.is_available() else 'cpu'
         # -- mast3r --
         # MAST3R_CKP = './checkpoints/MASt3R_ViTLarge_BaseDecoder_512_catmlpdpt_metric.pth'
         MAST3R_CKP = 'naver/MASt3R_ViTLarge_BaseDecoder_512_catmlpdpt_metric'