run dpt on CPU

app.py

@@ -17,6 +17,7 @@ import trimesh
 import random
 import imageio
 from einops import repeat
+from threeDFixer.moge.model.v2 import MoGeModel
 from threeDFixer.pipelines import ThreeDFixerPipeline
 from threeDFixer.datasets.utils import (
     edge_mask_morph_gradient,
@@ -51,36 +52,6 @@ work_space = None
 dpt_pack = None
 generated_object_map = {}
 
-# Prepare models
-## Grounding SAM
-segmenter_id = "facebook/sam-vit-base"
-sam_processor = AutoProcessor.from_pretrained(segmenter_id)
-sam_segmentator = AutoModelForMaskGeneration.from_pretrained(segmenter_id).to(
-    DEVICE, DTYPE
-)
-
-# ############## 3D-Fixer model
-# model_dir = 'HorizonRobotics/3D-Fixer'
-# pipeline = ThreeDFixerPipeline.from_pretrained(
-#     model_dir, compile=False
-# )
-# pipeline.to(device=DEVICE)
-# ############## 3D-Fixer model
-
-# rot = np.array([
-#     [-1.0,  0.0,  0.0, 0.0],
-#     [ 0.0,  0.0,  1.0, 0.0],
-#     [ 0.0,  1.0,  0.0, 0.0],
-#     [ 0.0,  0.0,  0.0, 1.0],
-# ], dtype=np.float32)
-
-# c2w = torch.tensor([
-#     [1.0, 0.0, 0.0, 0.0],
-#     [0.0, 0.0, -1.0, 0.0],
-#     [0.0, 1.0, 0.0, 0.0],
-#     [0.0, 0.0, 0.0, 1.0],
-# ], dtype=torch.float32, device=DEVICE)
-
 save_projected_colored_pcd = lambda pts, pts_color, fpath: trimesh.PointCloud(pts.reshape(-1, 3), pts_color.reshape(-1, 3)).export(fpath)
 
 EXAMPLES = [
@@ -213,7 +184,6 @@ def run_segmentation(
 
     return seg_map_pil
 
-@spaces.GPU
 def run_depth_estimation(
     image_prompts: Any,
     seg_image: Union[str, Image.Image],
@@ -223,9 +193,6 @@ def run_depth_estimation(
 
         rgb_image = rgb_image.resize((1024, 1024), Image.Resampling.LANCZOS)
 
-        # global pipeline
-        # pipeline.cuda()
-
         global dpt_pack
         global work_space
         if work_space is None:
@@ -247,7 +214,7 @@ def run_depth_estimation(
         input_image = torch.tensor(input_image / 255, dtype=torch.float32, device=DEVICE).permute(2, 0, 1)
 
         with torch.no_grad():
-            output = pipeline.models['scene_cond_model'].infer(input_image)
+            output = moge_v2_dpt_model.infer(input_image)
         depth = output['depth']
         intrinsics = output['intrinsics']
 
@@ -877,14 +844,19 @@ if __name__ == '__main__':
     dpt_pack = None
     generated_object_map = {}
 
-    DEVICE = "cpu"
+    DEVICE = "cuda"
 
     # Prepare models
     ## Grounding SAM
     segmenter_id = "facebook/sam-vit-base"
     sam_processor = AutoProcessor.from_pretrained(segmenter_id)
     sam_segmentator = AutoModelForMaskGeneration.from_pretrained(segmenter_id).to(
-        DEVICE, DTYPE if DEVICE == 'cuda' else torch.float32
+        "cpu", torch.float32
+    )
+
+    mogev2_id = 'Ruicheng/moge-2-vitl'
+    moge_v2_dpt_model = MoGeModel.from_pretrained(mogev2_id).to(
+        "cpu", torch.float32
     )
 
     ############## 3D-Fixer model

requirements.txt

@@ -37,11 +37,6 @@ huggingface_hub==0.36.0
 git+https://github.com/EasternJournalist/utils3d.git@9a4eb15e4021b67b12c460c7057d642626897ec8
 git+https://github.com/facebookresearch/sam2.git
 pydantic==2.10.6
-# gradio==4.38.1
-# gradio_client==1.1.0
-# fastapi==0.111.1
-# httpx==0.27.0
 kaolin==0.18.0
 flash-attn==2.8.3+pt2.8.0cu129
 nvdiffrast==0.4.0+253ac4fpt2.8.0cu129
-xformers