Update app.py

app.py

@@ -8,12 +8,12 @@ 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_phi3.py"):
+    if not str(filename).endswith("flash_attention.py"):
         return get_imports(filename)
     imports = get_imports(filename)
     imports.remove("flash_attn")
@@ -21,14 +21,14 @@ def fixed_get_imports(filename: str | os.PathLike) -> list[str]:
 
 
 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 = False,
+        use_flash_attn = False,
+        trust_remote_code = True
+    ).eval().cuda()
 
 tokenizer = AutoTokenizer.from_pretrained(
     model_path,