Update app.py

app.py

@@ -8,19 +8,31 @@ import gradio as gr
 # Load the model and tokenizer 
 model_path = "ByteDance/Sa2VA-4B"
 
-model = AutoModel.from_pretrained(
-    model_path,
-    torch_dtype = torch.bfloat16,
-    low_cpu_mem_usage = True,
-    use_flash_attn = True,
-    #trust_remote_code = True
-).eval().cuda()
+from transformers.dynamic_module_utils import get_imports
 
-tokenizer = AutoTokenizer.from_pretrained(
-    model_path,
-    #trust_remote_code = True,
-    use_fast = False
-)
+def fixed_get_imports(filename: str | os.PathLike) -> list[str]:
+    if not str(filename).endswith("/modeling_phi3.py"):
+        return get_imports(filename)
+    imports = get_imports(filename)
+    imports.remove("flash_attn")
+    return imports
+
+
+with patch("transformers.dynamic_module_utils.get_imports", fixed_get_imports):
+    
+    model = AutoModel.from_pretrained(
+        model_path,
+        torch_dtype = torch.bfloat16,
+        low_cpu_mem_usage = True,
+        use_flash_attn = True,
+        trust_remote_code = True
+    ).eval().cuda()
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_path,
+        trust_remote_code = True,
+        use_fast = False
+    )
 
 def image_vision(image_input_path, prompt):
     image_path = image_input_path