init project

app.py

@@ -38,8 +38,8 @@ from modules.pe3r.models import Models
 import torchvision.transforms as tvf
 
 silent = False
-device = 'cuda' if torch.cuda.is_available() else 'cpu'
-pe3r = Models(device) #
+# device = 'cuda' if torch.cuda.is_available() else 'cpu'
+pe3r = Models('cpu') #
 
 def _convert_scene_output_to_glb(outdir, imgs, pts3d, mask, focals, cams2world, cam_size=0.05,
                                  cam_color=None, as_pointcloud=False,
@@ -244,7 +244,9 @@ def slerp_multiple(vectors, t_values):
 @torch.no_grad
 def get_mask_from_img_sam1(sam1_image, yolov8_image, original_size, input_size, transform):
 
-    # device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    pe3r.yolov8.to(device)
+    pe3r.mobilesamv2.to(device)
 
     sam_mask=[]
     img_area = original_size[0] * original_size[1]
@@ -300,7 +302,10 @@ def get_mask_from_img_sam1(sam1_image, yolov8_image, original_size, input_size,
 @torch.no_grad
 def get_cog_feats(images):
 
-    # device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    pe3r.sam2.to(device)
+    pe3r.siglip_processor.to(device)
+    pe3r.siglip.to(device)
 
     cog_seg_maps = []
     rev_cog_seg_maps = []
@@ -446,7 +451,8 @@ def get_reconstructed_scene(outdir, filelist, schedule, niter, min_conf_thr,
     then run get_3D_model_from_scene
     """
 
-    # device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    pe3r.mast3r.to(device)
 
     if len(filelist) < 2:
         raise gradio.Error("Please input at least 2 images.")
@@ -511,7 +517,9 @@ def get_reconstructed_scene(outdir, filelist, schedule, niter, min_conf_thr,
 def get_3D_object_from_scene(outdir, text, threshold, scene, min_conf_thr, as_pointcloud, 
                  mask_sky, clean_depth, transparent_cams, cam_size):
     
-    # device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    pe3r.siglip_tokenizer.to(device)
+
 
     texts = [text]
     inputs = pe3r.siglip_tokenizer(text=texts, padding="max_length", return_tensors="pt")