try xformers

app.py

@@ -20,16 +20,16 @@ import trimesh
 import random
 import imageio
 from einops import repeat
-# from threeDFixer.pipelines import ThreeDFixerPipeline
-# from threeDFixer.datasets.utils import (
-#     edge_mask_morph_gradient,
-#     process_scene_image,
-#     process_instance_image,
-#     transform_vertices,
-#     normalize_vertices,
-#     project2ply
-# )
-# from threeDFixer.utils import render_utils, postprocessing_utils
+from threeDFixer.pipelines import ThreeDFixerPipeline
+from threeDFixer.datasets.utils import (
+    edge_mask_morph_gradient,
+    process_scene_image,
+    process_instance_image,
+    transform_vertices,
+    normalize_vertices,
+    project2ply
+)
+from threeDFixer.utils import render_utils, postprocessing_utils
 from transformers import AutoModelForMaskGeneration, AutoProcessor
 from scripts.grounding_sam import plot_segmentation, segment
 import copy
@@ -176,57 +176,43 @@ def cleanup_tmp(tmp_root: str = "./tmp", expire_seconds: int = 3600) -> None:
         except Exception as e:
             print(f"[cleanup_tmp] failed to remove {path}: {e}")
 
-
+# run seg on CPU
 def run_segmentation(
     image_prompts: Any,
     polygon_refinement: bool = True,
 ) -> Image.Image:
-    try:
-        gr.Info('in run_segmentation')
-        rgb_image = image_prompts["image"].convert("RGB")
-        gr.Info('in run_segmentation, read image')
+    rgb_image = image_prompts["image"].convert("RGB")
 
-        global work_space
+    global work_space
 
-        # device = "cuda" if torch.cuda.is_available() else "cpu"
-        device = "cpu"
-        sam_segmentator.to(device=device, dtype=DTYPE if device == 'cuda' else torch.float32)
-        gr.Info('in run_segmentation, move sam')
+    device = "cpu"
+    sam_segmentator.to(device=device, dtype=DTYPE if device == 'cuda' else torch.float32)
 
-        # pre-process the layers and get the xyxy boxes of each layer
-        if len(image_prompts["points"]) == 0:
-            raise gr.Error("No points provided for segmentation. Please add points to the image.")
-        
-        boxes = [
-            [
-                [int(box[0]), int(box[1]), int(box[3]), int(box[4])]
-                for box in image_prompts["points"]
-            ]
+    # pre-process the layers and get the xyxy boxes of each layer
+    if len(image_prompts["points"]) == 0:
+        raise gr.Error("No points provided for segmentation. Please add points to the image.")
+    
+    boxes = [
+        [
+            [int(box[0]), int(box[1]), int(box[3]), int(box[4])]
+            for box in image_prompts["points"]
         ]
+    ]
 
-        gr.Info('in run_segmentation, run info')
-        with torch.no_grad():
-            detections = segment(
-                sam_processor,
-                sam_segmentator,
-                rgb_image,
-                boxes=[boxes],
-                polygon_refinement=polygon_refinement,
-            )
-            seg_map_pil = plot_segmentation(rgb_image, detections)
-
-        torch.cuda.empty_cache()
-        
-        cleanup_tmp(TMP_DIR, expire_seconds=3600)
-
-        work_space = os.path.join(TMP_DIR, f"work_space_{uuid.uuid4()}")
-        os.makedirs(work_space, exist_ok=True)
-        seg_map_pil.save(os.path.join(work_space, 'mask.png'))
-
-    except Exception as e:
-        import traceback
-        traceback.print_exc()
-        raise gr.Error(f"run_segmentation failed: {type(e).__name__}: {e}")
+    with torch.no_grad():
+        detections = segment(
+            sam_processor,
+            sam_segmentator,
+            rgb_image,
+            boxes=[boxes],
+            polygon_refinement=polygon_refinement,
+        )
+        seg_map_pil = plot_segmentation(rgb_image, detections)
+    
+    cleanup_tmp(TMP_DIR, expire_seconds=3600)
+    work_space = os.path.join(TMP_DIR, f"work_space_{uuid.uuid4()}")
+    os.makedirs(work_space, exist_ok=True)
+    seg_map_pil.save(os.path.join(work_space, 'mask.png'))
 
     return seg_map_pil
 
@@ -901,15 +887,15 @@ if __name__ == '__main__':
     segmenter_id = "facebook/sam-vit-base"
     sam_processor = AutoProcessor.from_pretrained(segmenter_id)
     sam_segmentator = AutoModelForMaskGeneration.from_pretrained(segmenter_id).to(
-        DEVICE, DTYPE
+        DEVICE, DTYPE if DEVICE == 'cuda' else torch.float32
     )
 
     ############## 3D-Fixer model
     model_dir = 'HorizonRobotics/3D-Fixer'
-    # pipeline = ThreeDFixerPipeline.from_pretrained(
-    #     model_dir, compile=False
-    # )
-    # pipeline.to(device=DEVICE)
+    pipeline = ThreeDFixerPipeline.from_pretrained(
+        model_dir, compile=False
+    )
+    pipeline.to(device=DEVICE)
     ############## 3D-Fixer model
 
     rot = np.array([

requirements.txt

@@ -43,4 +43,5 @@ pydantic==2.10.6
 # httpx==0.27.0
 kaolin==0.18.0
 flash-attn==2.8.3+pt2.8.0cu129
-nvdiffrast==0.4.0+253ac4fpt2.8.0cu129
+nvdiffrast==0.4.0+253ac4fpt2.8.0cu129
+xformers