Update app.py

app.py

@@ -14,40 +14,40 @@ model.to("cuda:0")
 
 @spaces.GPU
 def bot_streaming(message, history):
-  print(history)
-  if message["files"]:
-    image = message["files"][-1]["path"]
-  else:
-    # if there's no image uploaded for this turn, look for images in the past turns
-    # kept inside tuples, take the last one
-    for hist in history:
-      if type(hist[0])==tuple:
-        image = hist[0][0]
-
-  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)
-  generated_text = ""
-
-  thread = Thread(target=model.generate, kwargs=generation_kwargs)
-  thread.start()
-
-  text_prompt =f"[INST]  \n{message['text']} [/INST]"
-  
-
-  buffer = ""
-  for new_text in streamer:
-    
-    buffer += new_text
+    chat_history = []
+    if message["files"]:
+      image = message["files"][-1]["path"]
+    else:
+      for hist in history:
+        if type(hist[0])==tuple:
+          image = hist[0][0]
+            
+    if len(history) > 0 and image:
+        chat_history.append({"role": "user", "content": f'<image>\n{message['text']}'})
+        for human, assistant in history[1:]:
+            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.")
+    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)
+    generated_text = ""
+    thread = Thread(target=model.generate, kwargs=generation_kwargs)
+    thread.start()
+    text_prompt =f"[INST]  \n{message['text']} [/INST]"
     
-    generated_text_without_prompt = buffer[len(text_prompt):]
-    time.sleep(0.04)
-    yield generated_text_without_prompt
+    buffer = ""
+    for new_text in streamer:
+      
+      buffer += new_text
+      
+      generated_text_without_prompt = buffer[len(text_prompt):]
+      time.sleep(0.04)
+      yield generated_text_without_prompt
 
 
 demo = gr.ChatInterface(fn=bot_streaming, title="LLaVA NeXT", examples=[{"text": "What is on the flower?", "files":["./bee.jpg"]},