Update app.py

app.py

@@ -1,9 +1,7 @@
 import os
-from unittest.mock import patch
 import spaces
 import gradio as gr
 from transformers import AutoProcessor, AutoModelForCausalLM
-from transformers.dynamic_module_utils import get_imports
 import torch
 from PIL import Image, ImageDraw
 import random
@@ -14,24 +12,19 @@ import cv2
 import io
 import uuid
 
-def workaround_fixed_get_imports(filename: str | os.PathLike) -> list[str]:
-    if not str(filename).endswith("/modeling_florence2.py"):
-        return get_imports(filename)
-    imports = get_imports(filename)
-    imports.remove("flash_attn")
-    return imports
+import subprocess
+subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
 
 def load_model():
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    with patch("transformers.dynamic_module_utils.get_imports", workaround_fixed_get_imports):
-        model = AutoModelForCausalLM.from_pretrained(
-            "microsoft/Florence-2-large-ft", 
-            trust_remote_code=True
-        ).to(device).eval()
-        processor = AutoProcessor.from_pretrained(
-            "microsoft/Florence-2-large-ft", 
-            trust_remote_code=True
-        )
+    model = AutoModelForCausalLM.from_pretrained(
+        "microsoft/Florence-2-large-ft", 
+        trust_remote_code=True
+    ).to(device).eval()
+    processor = AutoProcessor.from_pretrained(
+        "microsoft/Florence-2-large-ft", 
+        trust_remote_code=True
+    )
     return model, processor, device
 
 model = None