Update app.py

app.py

@@ -1,12 +1,12 @@
 import spaces
 import os
 import re
-import time
 import gradio as gr
 import torch
 from transformers import AutoModelForCausalLM
 from transformers import TextIteratorStreamer
 from threading import Thread
+from PIL import Image
 
 model_name = 'AIDC-AI/Ovis1.6-Gemma2-9B'
 
@@ -21,26 +21,28 @@ streamer = TextIteratorStreamer(text_tokenizer, skip_prompt=True, skip_special_t
 image_placeholder = '<image>'
 cur_dir = os.path.dirname(os.path.abspath(__file__))
 
-def submit_chat(chatbot, text_input):
-    response = ''
-    chatbot.append((text_input, response))
-    return chatbot ,''
 
 @spaces.GPU
-def ovis_chat(chatbot, image_input):
+def ovis_chat(message, history):
+    try :
+        image_input = Image.open(message["files"][0]).convert("RGB")
+    except : 
+        image_input = None
     # preprocess inputs
     conversations = []
     response = ""
-    text_input = chatbot[-1][0]
-    for query, response in chatbot[:-1]:
-        conversations.append({
-            "from": "human",
-            "value": query
-        })
-        conversations.append({
-            "from": "gpt",
-            "value": response
-        })
+    text_input = message["text"]
+    for msg in enumerate(history):
+        if msg["role"] == "user" and "size" not in msg.keys():
+            conversations.append({
+                "from": "human",
+                "value": msg["text"]
+            })
+        elif msg["role"] == "assistant":
+            conversations.append({
+                "from": "gpt",
+                "value": response
+            })
     text_input = text_input.replace(image_placeholder, '')
     conversations.append({
         "from": "human",
@@ -79,19 +81,8 @@ def ovis_chat(chatbot, image_input):
     thread.start()
     for new_text in streamer:
         response += new_text
-        chatbot[-1][1] = response
-        yield chatbot
+        yield response
     thread.join()
-    # debug
-    print('*'*60)
-    print('*'*60)
-    print('OVIS_CONV_START')
-    for i, (request, answer) in enumerate(chatbot[:-1], 1):
-        print(f'Q{i}:\n {request}')
-        print(f'A{i}:\n {answer}')
-    print('New_Q:\n', text_input)
-    print('New_A:\n', response)
-    print('OVIS_CONV_END')
 
 def clear_chat():
     return [], None, ""
@@ -138,29 +129,5 @@ latex_delimiters_set = [{
         "display": True
     }]
 
-text_input = gr.Textbox(label="prompt", placeholder="Enter your text here...", lines=1, container=False)
-with gr.Blocks(title=model_name.split('/')[-1], theme=gr.themes.Ocean()) as demo:
-    gr.HTML(html)
-    with gr.Row():
-        with gr.Column(scale=3):
-            image_input = gr.Image(label="image", height=350, type="pil")
-            gr.Examples(
-                examples=[
-                    [f"{cur_dir}/examples/case0.png", "Find the area of the shaded region."],
-                    [f"{cur_dir}/examples/case1.png", "explain this model to me."],
-                    [f"{cur_dir}/examples/case2.png", "What is net profit margin as a percentage of total revenue?"],
-                ],
-                inputs=[image_input, text_input]
-            )
-        with gr.Column(scale=7):
-            chatbot = gr.Chatbot(label="Ovis", layout="panel", height=600, show_copy_button=True, latex_delimiters=latex_delimiters_set)
-            text_input.render()
-            with gr.Row():
-                send_btn = gr.Button("Send", variant="primary")
-                clear_btn = gr.Button("Clear", variant="secondary")
-
-    send_click_event = send_btn.click(submit_chat, [chatbot, text_input], [chatbot, text_input]).then(ovis_chat,[chatbot, image_input],chatbot)
-    submit_event = text_input.submit(submit_chat, [chatbot, text_input], [chatbot, text_input]).then(ovis_chat,[chatbot, image_input],chatbot)
-    clear_btn.click(clear_chat, outputs=[chatbot, image_input, text_input])
-
-demo.launch()
+demo = gr.ChatInterface(fn=ovis_chat, type="messages", textbox=gr.MultimodalTextbox(),multimodal=True)
+demo.launch(debug=True)