Update app.py

app.py

@@ -1,5 +1,5 @@
 import torch
-from transformers import AutoTokenizer, AutoModel 
+from transformers import AutoModelForCausalLM, AutoTokenizer
 from PIL import Image 
 import numpy as np 
 import os 
@@ -8,18 +8,16 @@ import gradio as gr
 # Load the model and tokenizer 
 model_path = "ByteDance/Sa2VA-4B"
  
-model = AutoModel.from_pretrained(
+model = AutoModelForCausalLM.from_pretrained(
     model_path,
-    torch_dtype = torch.bfloat16,
-    low_cpu_mem_usage = True,
-    use_flash_attn = True,
-    trust_remote_code = True
+    torch_dtype="auto",
+    device_map="auto",
+    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):