init project

app.py

@@ -38,8 +38,7 @@ from modules.pe3r.models import Models
 import torchvision.transforms as tvf
 
 silent = False
-device = 'cpu' # 'cuda' if torch.cuda.is_available() else 
-pe3r = Models(device) #
+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,6 +243,8 @@ def slerp_multiple(vectors, t_values):
 @torch.no_grad
 def get_mask_from_img_sam1(mobilesamv2, yolov8, sam1_image, yolov8_image, original_size, input_size, transform):
 
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
     sam_mask=[]
     img_area = original_size[0] * original_size[1]
 
@@ -297,7 +298,9 @@ def get_mask_from_img_sam1(mobilesamv2, yolov8, sam1_image, yolov8_image, origin
 
 @torch.no_grad
 def get_cog_feats(images):
+
     device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
     cog_seg_maps = []
     rev_cog_seg_maps = []
     inference_state = pe3r.sam2.init_state(images=images.sam2_images, video_height=images.sam2_video_size[0], video_width=images.sam2_video_size[1])
@@ -441,6 +444,9 @@ def get_reconstructed_scene(outdir, filelist, schedule, niter, min_conf_thr,
     from a list of images, run dust3r inference, global aligner.
     then run get_3D_model_from_scene
     """
+
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
     if len(filelist) < 2:
         raise gradio.Error("Please input at least 2 images.")
 
@@ -504,6 +510,8 @@ 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'
+
     texts = [text]
     inputs = pe3r.siglip_tokenizer(text=texts, padding="max_length", return_tensors="pt")
     inputs = {key: value.to(device) for key, value in inputs.items()}