init project

app.py

@@ -37,11 +37,9 @@ from modules.mobilesamv2.utils.transforms import ResizeLongestSide
 from modules.pe3r.models import Models
 import torchvision.transforms as tvf
 
-from transformers import AutoTokenizer, AutoModel, AutoProcessor, SamModel
-
 silent = False
-# device = 'cuda' if torch.cuda.is_available() else 'cpu'
-pe3r = Models('cpu') #
+device = 'cpu' #'cuda' if torch.cuda.is_available() else 'cpu'
+pe3r = Models(device) #
 
 def _convert_scene_output_to_glb(outdir, imgs, pts3d, mask, focals, cams2world, cam_size=0.05,
                                  cam_color=None, as_pointcloud=False,
@@ -246,9 +244,7 @@ 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'
-    pe3r.yolov8.to(device)
-    pe3r.mobilesamv2.to(device)
+    # device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
     sam_mask=[]
     img_area = original_size[0] * original_size[1]
@@ -304,14 +300,7 @@ 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'
-    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", device_map=device)
-
-    # pe3r.siglip_processor.to(device)
-    # pe3r.siglip.to(device)
+    # device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
     cog_seg_maps = []
     rev_cog_seg_maps = []
@@ -405,10 +394,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 = siglip_processor(images=seg_imgs, return_tensors="pt")
+        inputs = pe3r.siglip_processor(images=seg_imgs, return_tensors="pt")
         inputs = {key: value.to(device) for key, value in inputs.items()}
         
-        image_features = siglip.get_image_features(**inputs)
+        image_features = pe3r.siglip.get_image_features(**inputs)
         image_features = image_features / image_features.norm(dim=-1, keepdim=True)
         image_features = image_features.detach().cpu()
 
@@ -457,8 +446,7 @@ 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'
-    pe3r.mast3r.to(device)
+    # device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
     if len(filelist) < 2:
         raise gradio.Error("Please input at least 2 images.")
@@ -523,16 +511,13 @@ 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'
-    # pe3r.siglip_tokenizer.to(device)
-    siglip = AutoModel.from_pretrained("google/siglip-large-patch16-256", device_map=device)
-    siglip_tokenizer = AutoTokenizer.from_pretrained("google/siglip-large-patch16-256", device_map=device)
+    # device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
     texts = [text]
-    inputs = siglip_tokenizer(text=texts, padding="max_length", return_tensors="pt")
+    inputs = pe3r.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 =siglip.get_text_features(**inputs)
+        text_feats =pe3r.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", device_map=device)
+        self.siglip_processor = AutoProcessor.from_pretrained("google/siglip-large-patch16-256", device_map=device)