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()