Update app.py

app.py

@@ -1,5 +1,5 @@
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer
 from threading import Thread
 import re
 import time 
@@ -50,15 +50,14 @@ def bot_streaming(message, history):
 
     # 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")
     text = tokenizer.apply_chat_template(
-    messages,
-    tokenize=False,
-    add_generation_prompt=True)
+        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(tokenizer, **{"skip_special_tokens": True})
+    streamer = TextStreamer(tokenizer, **{"skip_special_tokens": True})
     image_tensor = model.process_images([image], model.config).to(dtype=model.dtype)
     generation_kwargs = dict(inputs=input_ids, images=image_tensor, streamer=streamer, max_new_tokens=100)
     generated_text = ""