Update app.py

app.py

@@ -1,5 +1,5 @@
 import gradio as gr
-from transformers import LlavaNextProcessor, LlavaNextForConditionalGeneration, TextIteratorStreamer
+from transformers import AutoModelForCausalLM, AutoTokenizer
 from threading import Thread
 import re
 import time 
@@ -7,9 +7,15 @@ from PIL import Image
 import torch
 import spaces
 
-processor = LlavaNextProcessor.from_pretrained("llava-hf/llava-v1.6-mistral-7b-hf")
+tokenizer = AutoTokenizer.from_pretrained(
+    'qnguyen3/nanoLLaVA',
+    trust_remote_code=True)
 
-model = LlavaNextForConditionalGeneration.from_pretrained("llava-hf/llava-v1.6-mistral-7b-hf", torch_dtype=torch.float16, low_cpu_mem_usage=True)
+model = AutoModelForCausalLM.from_pretrained(
+    'qnguyen3/nanoLLaVA',
+    torch_dtype=torch.float16,
+    device_map='auto',
+    trust_remote_code=True)
 model.to("cuda:0")
 
 @spaces.GPU
@@ -18,27 +24,46 @@ def bot_streaming(message, history):
     if message["files"]:
       image = message["files"][-1]["path"]
     else:
-      for hist in history:
+      for i, hist in enumerate(history):
         if type(hist[0])==tuple:
           image = hist[0][0]
+          image_turn = i
             
-    if len(history) > 0 and image:
-        chat_history.append({"role": "user", "content": f'<image>\n{message['text']}'})
-        for human, assistant in history[1:]:
+    if len(history) > 0 and image is not None:
+        chat_history.append({"role": "user", "content": f'<image>\n{history[1][0]}'})
+        chat_history.append({"role": "assistant", "content": history[1][1] })
+        for human, assistant in history[2:]:
             chat_history.append({"role": "user", "content": human })
             chat_history.append({"role": "assistant", "content": assistant })
-            
-    if image is None:
-        gr.Error("You need to upload an image for LLaVA to work.")
+        chat_history.append({"role": "user", "content": message['text']})
+    elif len(history) > 0 and image is None:
+        for human, assistant in history:
+            chat_history.append({"role": "user", "content": human })
+            chat_history.append({"role": "assistant", "content": assistant })
+        chat_history.append({"role": "user", "content": message['text']})
+    elif len(history) == 0 and image is not None:
+        chat_history.append({"role": "user", "content": f'<image>\n{message['text']}'})
+    elif len(history) == 0 and image is None:
+        chat_history.append({"role": "user", "content": message['text'] })
+
+    # if image is None:
+    #     gr.Error("You need to upload an image for LLaVA to work.")
     prompt=f"[INST] <image>\n{message['text']} [/INST]"
     image = Image.open(image).convert("RGB")
-    inputs = processor(prompt, image, return_tensors="pt").to("cuda:0")
-    streamer = TextIteratorStreamer(processor, **{"skip_special_tokens": True})
-    generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=100)
+    text = tokenizer.apply_chat_template(
+    messages,
+    tokenize=False,
+    add_generation_prompt=True)
+    text_chunks = [tokenizer(chunk).input_ids for chunk in text.split('<image>')]
+    input_ids = torch.tensor(text_chunks[0] + [-200] + text_chunks[1], dtype=torch.long).unsqueeze(0)
+    streamer = TextIteratorStreamer(input_ids, **{"skip_special_tokens": True})
+    image = Image.open(image)
+    image_tensor = model.process_images([image], model.config).to(dtype=model.dtype)
+    generation_kwargs = dict(inputs, images=image_tensor, streamer=streamer, max_new_tokens=100)
     generated_text = ""
     thread = Thread(target=model.generate, kwargs=generation_kwargs)
     thread.start()
-    text_prompt =f"[INST]  \n{message['text']} [/INST]"
+    text_prompt =f"<|im_start|>user\n{message['text']}<|im_end|>"
     
     buffer = ""
     for new_text in streamer: