init project

app.py

@@ -37,10 +37,12 @@ from modules.mobilesamv2.utils.transforms import ResizeLongestSide
 from modules.pe3r.models import Models
 import torchvision.transforms as tvf
 
+from transformers import AutoTokenizer, AutoModel, AutoProcessor
+
 silent = False
-device = 'cuda' if torch.cuda.is_available() else 'cpu' #'cpu' #
-pe3r = Models(device) #
-print(device)
+# device = 'cuda' if torch.cuda.is_available() else 'cpu' #'cpu' #
+pe3r = Models('cpu') #
+# print(device)
 
 def _convert_scene_output_to_glb(outdir, imgs, pts3d, mask, focals, cams2world, cam_size=0.05,
                                  cam_color=None, as_pointcloud=False,
@@ -245,7 +247,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]
@@ -301,7 +305,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)
+    siglip = AutoModel.from_pretrained("google/siglip-large-patch16-256", device_map=device)
+    siglip_processor = AutoProcessor.from_pretrained("google/siglip-large-patch16-256")
 
     cog_seg_maps = []
     rev_cog_seg_maps = []
@@ -395,10 +402,10 @@ def get_cog_feats(images):
         seg_imgs = np.stack(seg_img_list, axis=0) # b,H,W,3
         seg_imgs = torch.from_numpy(seg_imgs).permute(0,3,1,2) # / 255.0
         
-        inputs = pe3r.siglip_processor(images=seg_imgs, return_tensors="pt")
+        inputs = siglip_processor(images=seg_imgs, return_tensors="pt")
         inputs = {key: value.to(device) for key, value in inputs.items()}
         
-        image_features = pe3r.siglip.get_image_features(**inputs)
+        image_features = siglip.get_image_features(**inputs)
         image_features = image_features / image_features.norm(dim=-1, keepdim=True)
         image_features = image_features.detach().cpu()
 
@@ -438,7 +445,7 @@ def get_cog_feats(images):
     return cog_seg_maps, rev_cog_seg_maps, multi_view_clip_feats
 
 
-@spaces.GPU(duration=120)
+@spaces.GPU(duration=60)
 def get_reconstructed_scene(outdir, filelist, schedule, niter, min_conf_thr,
                             as_pointcloud, mask_sky, clean_depth, transparent_cams, cam_size,
                             scenegraph_type, winsize, refid):
@@ -447,7 +454,9 @@ 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.")
@@ -505,22 +514,24 @@ def get_reconstructed_scene(outdir, filelist, schedule, niter, min_conf_thr,
     outfile = get_3D_model_from_scene(outdir, scene, min_conf_thr, as_pointcloud, mask_sky,
                                       clean_depth, transparent_cams, cam_size)
     
-    scene.to('cpu')
+    # scene.to('cpu')
     torch.cuda.empty_cache()
 
     return scene, outfile 
 
-@spaces.GPU(duration=120)
+# @spaces.GPU(duration=60)
 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'
+    siglip_tokenizer = AutoTokenizer.from_pretrained("google/siglip-large-patch16-256")
+    siglip = AutoModel.from_pretrained("google/siglip-large-patch16-256", device_map=device)
 
     texts = [text]
-    inputs = pe3r.siglip_tokenizer(text=texts, padding="max_length", return_tensors="pt")
+    inputs = siglip_tokenizer(text=texts, padding="max_length", return_tensors="pt")
     inputs = {key: value.to(device) for key, value in inputs.items()}
     with torch.no_grad():
-        text_feats =pe3r.siglip.get_text_features(**inputs)
+        text_feats =siglip.get_text_features(**inputs)
         text_feats = text_feats / text_feats.norm(dim=-1, keepdim=True)
     scene.render_image(text_feats, threshold)
     scene.ori_imgs = scene.rendered_imgs

modules/pe3r/models.py

@@ -47,6 +47,6 @@ class Models:
         self.yolov8 = ObjectAwareModel(YOLO8_CKP)
 
         # -- siglip --
-        self.siglip = AutoModel.from_pretrained("google/siglip-large-patch16-256", device_map=device)
-        self.siglip_tokenizer = AutoTokenizer.from_pretrained("google/siglip-large-patch16-256", device_map=device)
-        self.siglip_processor = AutoProcessor.from_pretrained("google/siglip-large-patch16-256", device_map=device)
+        # self.siglip = AutoModel.from_pretrained("google/siglip-large-patch16-256", device_map=device)
+        # self.siglip_tokenizer = AutoTokenizer.from_pretrained("google/siglip-large-patch16-256")
+        # self.siglip_processor = AutoProcessor.from_pretrained("google/siglip-large-patch16-256")