Update app.py

app.py

@@ -7,7 +7,7 @@ from threading import Thread
 import gradio as gr
 import spaces
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, AutoModelForImageTextToText
 
 DESCRIPTION = """# AlexHung29629/fix_magistral_fc2
 開頭有<think>"""
@@ -46,7 +46,7 @@ CHAT_TEMPLATE="""{%- set default_system_message = "A user will ask you to solve
 
 if torch.cuda.is_available():
     model_id = "mistralai/Mistral-Small-24B-Instruct-2501"
-    model = AutoModelForCausalLM.from_pretrained("AlexHung29629/fix_magistral_fc2", torch_dtype=torch.bfloat16, device_map="auto")
+    model = AutoModelForImageTextToText.from_pretrained("AlexHung29629/add_vision_4", torch_dtype=torch.bfloat16, device_map="auto")
     tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-Small-24B-Instruct-2501")