Update app.py

app.py

@@ -75,6 +75,28 @@ print('Initialization Finished')
 #             Gradio Setting
 # ========================================
 
+def gradio_ask(user_message, chatbot, chat_state):
+    print("building prompt...")
+    if len(user_message) == 0:
+        return gr.update(interactive=True, placeholder='Input should not be empty!'), chatbot, chat_state
+    chat.ask(user_message, chat_state)
+    chatbot = chatbot + [[user_message, None]]
+    return '', chatbot, chat_state
+
+
+def gradio_answer(chatbot, chat_state, img_list, num_beams, temperature):
+    print("generating answer...")
+    llm_message = chat.answer(conv=chat_state,
+                              img_list=img_list,
+                              num_beams=1,
+                              temperature=temperature,
+                              max_new_tokens=240,
+                              max_length=511)[0]
+    chatbot[-1][1] = llm_message
+    print(chat_state.get_prompt())
+    print(chat_state)
+    return chatbot, chat_state, img_list
+
 def gradio_reset(chat_state, img_list):
     if chat_state is not None:
         chat_state.messages = []
@@ -104,8 +126,16 @@ def upload_imgorvideo(gr_video, gr_img, text_input,chatbot,audio_flag):
         )
         img_list = []
         llm_message = chat.upload_video(gr_video, chat_state, img_list)
-        print(chat_state.messages)
-        return text_input,chat_state, chatbot
+        llm_message = chat.ask(text_input, chat_state)
+        llm_message = chat.answer(conv=chat_state,
+                              img_list=img_list,
+                              num_beams=1,
+                              temperature=1.0,
+                              max_new_tokens=240,
+                              max_length=511)[0]
+        print(llm_message)
+        output = [[llm_message]]
+        return llm_message, output
     elif gr_img is not None:
         print(gr_img)
         chatbot = [((gr_img,), None)]
@@ -125,28 +155,6 @@ def upload_imgorvideo(gr_video, gr_img, text_input,chatbot,audio_flag):
         # img_list = []
         return gr.update(interactive=False), gr.update(interactive=False, placeholder='Currently, only one input is supported'), gr.update(value="Currently, only one input is supported", interactive=False), chat_state, None,chatbot
 
-def gradio_ask(user_message, chatbot, chat_state):
-    print("building prompt...")
-    if len(user_message) == 0:
-        return gr.update(interactive=True, placeholder='Input should not be empty!'), chatbot, chat_state
-    chat.ask(user_message, chat_state)
-    chatbot = chatbot + [[user_message, None]]
-    return '', chatbot, chat_state
-
-
-def gradio_answer(chatbot, chat_state, img_list, num_beams, temperature):
-    print("generating answer...")
-    llm_message = chat.answer(conv=chat_state,
-                              img_list=img_list,
-                              num_beams=1,
-                              temperature=temperature,
-                              max_new_tokens=240,
-                              max_length=511)[0]
-    chatbot[-1][1] = llm_message
-    print(chat_state.get_prompt())
-    print(chat_state)
-    return chatbot, chat_state, img_list
-
 
 title = """
 <h1 align="center"><a href="https://github.com/DAMO-NLP-SG/Video-LLaMA"><img src="https://s1.ax1x.com/2023/05/22/p9oQ0FP.jpg", alt="Video-LLaMA" border="0" style="margin: 0 auto; height: 200px;" /></a> </h1>
@@ -219,14 +227,13 @@ with gr.Blocks() as demo:
             img_list = gr.State()
             chatbot = gr.Chatbot(label='Video-LLaMA')
             text_input = gr.Textbox(label='User', placeholder='Upload your image/video first, or directly click the examples at the bottom of the page.', interactive=False)
+            output = gr.Textbox(label='Output')
             
         
     gr.Markdown(cite_markdown)
     #upload_button.click(upload_imgorvideo, inputs=[video, image, text_input], outputs=[chat_state,chatbot])
     
-    text_input.submit(upload_imgorvideo, inputs=[video, image, text_input], outputs=[text_input,chatbot, chat_state]).then(
-    gradio_ask, [text_input, chatbot, chat_state], [text_input, chatbot, chat_state]).then(gradio_answer, [chatbot, chat_state, img_list, num_beams, temperature], [chatbot, chat_state, img_list]
-    )
+    text_input.submit(upload_imgorvideo, inputs=[video, image, text_input], outputs=[output])
     #clear.click(gradio_reset, [chat_state, img_list], [chatbot, video, image, text_input, upload_button, chat_state, img_list], queue=False)
     
 demo.queue().launch(debug=True)