From 9ecda2e689d72298b4dca39403a85d13e53ea659 Mon Sep 17 00:00:00 2001
From: AI Lab <129358391+1038lab@users.noreply.github.com>
Date: Mon, 15 Dec 2025 23:01:05 -0800
Subject: [PATCH] Add files via upload

---
 AILab_SAM2Segment.py | 26 +++++++++++++++-----------
 1 file changed, 15 insertions(+), 11 deletions(-)

diff --git a/AILab_SAM2Segment.py b/AILab_SAM2Segment.py
index 815e67f..dad934f 100644
--- a/AILab_SAM2Segment.py
+++ b/AILab_SAM2Segment.py
@@ -24,16 +24,6 @@ except ImportError:
     GROUNDINGDINO_AVAILABLE = False
     print("Warning: GroundingDINO not available. Text prompts will use fallback method.")
 
-# Disable PyTorch JIT compilation
-# torch.jit._state._python_cu = None
-# torch._C._jit_set_profiling_mode(False)
-# torch._C._jit_set_profiling_executor(False)
-
-# original_jit_script = torch.jit.script
-# def patched_jit_script(obj):
-#     return obj
-# torch.jit.script = patched_jit_script
-
 current_dir = os.path.dirname(__file__)
 sam2_path = os.path.join(current_dir, "sam2")
 sys.path.insert(0, sam2_path)
@@ -48,6 +38,17 @@ def patched_torch_load(*args, **kwargs):
 
 torch.load = patched_torch_load
 
+from contextlib import contextmanager
+
+@contextmanager
+def _sam2_no_jit():
+    _orig = torch.jit.script
+    torch.jit.script = lambda x, *a, **k: x
+    try:
+        yield
+    finally:
+        torch.jit.script = _orig
+
 from sam2_image_predictor import SAM2ImagePredictor
 from AILab_ImageMaskTools import pil2tensor, tensor2pil
 
@@ -247,7 +248,10 @@ class SAM2Segment:
             sam_model.load_state_dict(state_dict, strict=False)
             sam_model = sam_model.to(dtype).to(sam_device).eval()
             
-            predictor = SAM2ImagePredictor(sam_model)
+            # predictor = SAM2ImagePredictor(sam_model)
+            with _sam2_no_jit():
+                predictor = SAM2ImagePredictor(sam_model)
+                
             self.sam2_model_cache[cache_key] = predictor
         return self.sam2_model_cache[cache_key]