Update app.py

app.py

@@ -1,6 +1,7 @@
 import gradio as gr
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, StoppingCriteria
+from .modeling_llava_qwen2 import LlavaQwen2ForCausalLM
 from threading import Thread
 import re
 import time 
@@ -16,7 +17,7 @@ tokenizer = AutoTokenizer.from_pretrained(
     'qnguyen3/nanoLLaVA',
     trust_remote_code=True)
 
-model = AutoModelForCausalLM.from_pretrained(
+model = LlavaQwen2ForCausalLM.from_pretrained(
     'qnguyen3/nanoLLaVA',
     torch_dtype=torch.float16,
     trust_remote_code=True)