Update app.py

app.py

@@ -4,13 +4,27 @@ from PIL import Image
 import torch
 from peft import PeftModel
 import numpy as np
+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]:
+    if not str(filename).endswith("modeling_florence2.py"):
+        return get_imports(filename)
+    imports = get_imports(filename)
+    imports.remove("flash_attn")
+    return imports
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 torch_dtype = torch.float32
 
 # Load the fine-tuned base model
-caption_model = AutoModelForCausalLM.from_pretrained('microsoft/Florence-2-base-ft', trust_remote_code=True, revision='refs/pr/6', torch_dtype=torch_dtype, attn_implementation="eager").to(device)
-model = AutoModelForCausalLM.from_pretrained('byh711/FLODA-deepfake', trust_remote_code=True, torch_dtype=torch_dtype, attn_implementation="eager").to(device)
+with patch("transformers.dynamic_module_utils.get_imports", fixed_get_imports):
+  caption_model = AutoModelForCausalLM.from_pretrained('microsoft/Florence-2-base-ft', trust_remote_code=True, revision='refs/pr/6', torch_dtype=torch_dtype).to(device)
+
+with patch("transformers.dynamic_module_utils.get_imports", fixed_get_imports):
+  model = AutoModelForCausalLM.from_pretrained('byh711/FLODA-deepfake', trust_remote_code=True, torch_dtype=torch_dtype).to(device)
+
 processor = AutoProcessor.from_pretrained('byh711/FLODA-deepfake', trust_remote_code=True)
 model.eval()