Update app.py

app.py

@@ -7,28 +7,14 @@ import gradio as gr
 
 # Load the model and tokenizer 
 model_path = "ByteDance/Sa2VA-4B"
-
-
-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_sa2va_chat.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 = False,
-        use_flash_attn = False,
-        trust_remote_code = True
-    ).eval().cuda()
+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,