flash attn fix

app.py

@@ -10,6 +10,8 @@ from sam2.sam2_image_predictor import SAM2ImagePredictor
 import cv2
 import traceback
 import matplotlib.pyplot as plt
+from utils import load_model_without_flash_attn
+
 
 # CUDA optimizations
 torch.autocast(device_type="cuda", dtype=torch.bfloat16).__enter__()
@@ -26,9 +28,19 @@ sam2_model = build_sam2(model_cfg, sam2_checkpoint, device="cuda")
 image_predictor = SAM2ImagePredictor(sam2_model)
 
 model_id = 'microsoft/Florence-2-large'
-florence_model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True, torch_dtype=torch.bfloat16).eval().cuda()
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+def load_florence_model():
+    return AutoModelForCausalLM.from_pretrained(
+        model_id, 
+        trust_remote_code=True, 
+        torch_dtype=torch.float16 if device == "cuda" else torch.float32
+    ).eval().to(device)
+
+florence_model = load_model_without_flash_attn(load_florence_model)
 florence_processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
 
+
 def apply_color_mask(frame, mask, obj_id):
     cmap = plt.get_cmap("tab10")
     color = np.array(cmap(obj_id % 10)[:3])  # Use modulo 10 to cycle through colors

requirements.txt

@@ -8,5 +8,4 @@ opencv-python
 matplotlib
 einops
 timm
-pytest
-flash_attn
+pytest

utils.py

@@ -0,0 +1,17 @@
+import os
+from unittest.mock import patch
+from transformers.dynamic_module_utils import get_imports
+
+def fixed_get_imports(filename: str | os.PathLike) -> list[str]:
+    """Workaround for flash_attn import issue."""
+    if not str(filename).endswith(("modeling_phi.py", "configuration_florence2.py")):
+        return get_imports(filename)
+    imports = get_imports(filename)
+    if "flash_attn" in imports:
+        imports.remove("flash_attn")
+    return imports
+
+def load_model_without_flash_attn(model_loader):
+    """Load a model using the flash_attn workaround."""
+    with patch("transformers.dynamic_module_utils.get_imports", fixed_get_imports):
+        return model_loader()