Update caption_models.py

caption_models.py

@@ -11,13 +11,45 @@ import torch.nn as nn
 
 subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, check=True, shell=True)
 
+import sys
+import importlib.util
+from unittest.mock import MagicMock
+
+# Create a fake spec object
+class FakeFlashAttnSpec:
+    name = 'flash_attn'
+    loader = None
+    origin = None
+    submodule_search_locations = []
+    
+fake_spec = FakeFlashAttnSpec()
+
+# Create mock modules with proper __spec__ attributes
+flash_attn_mock = MagicMock()
+flash_attn_mock.__spec__ = fake_spec
+flash_attn_mock.__version__ = "0.0.0"  # Force version check to fail
+
+sys.modules['flash_attn'] = flash_attn_mock
+sys.modules['flash_attn.flash_attn_interface'] = MagicMock()
+sys.modules['flash_attn.bert_padding'] = MagicMock()
+
+# Patch find_spec to return our fake spec
+_original_find_spec = importlib.util.find_spec
+
+def _patched_find_spec(name, package=None):
+    if name == 'flash_attn' or name.startswith('flash_attn.'):
+        return fake_spec
+    return _original_find_spec(name, package)
+
+importlib.util.find_spec = _patched_find_spec
+
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 HF_TOKEN = os.getenv("HUGGINGFACE_TOKEN", None)
 
 # Initialize Florence model
-florence_model = AutoModelForCausalLM.from_pretrained('microsoft/Florence-2-large', trust_remote_code=True).to(device).eval()
-florence_processor = AutoProcessor.from_pretrained('microsoft/Florence-2-large', trust_remote_code=True)
+florence_model = AutoModelForCausalLM.from_pretrained('microsoft/Florence-2-large', trust_remote_code=True, attn_implementation="sdpa", revision="00d2f1570b00c6dea5df998f5635db96840436bc").to(device).eval()
+florence_processor = AutoProcessor.from_pretrained('microsoft/Florence-2-large', trust_remote_code=True, attn_implementation="sdpa", revision="00d2f1570b00c6dea5df998f5635db96840436bc")
 
 # Initialize Qwen2-VL-2B model
 qwen_model = Qwen2VLForConditionalGeneration.from_pretrained("Qwen/Qwen2-VL-2B-Instruct", trust_remote_code=True, torch_dtype="auto").to(device).eval()