Update app.py

app.py

@@ -1,4 +1,5 @@
 from transformers import AutoModel, AutoTokenizer
+from copy import deepcopy
 import os
 import ipdb
 import gradio as gr
@@ -10,9 +11,9 @@ import json
 # init the model
 args = {
     'model': 'openllama_peft',
-    'imagebind_ckpt_path': './pretrained_ckpt/imagebind_ckpt/',
+    'imagebind_ckpt_path': 'pretrained_ckpt/imagebind_ckpt',
     'vicuna_ckpt_path': 'openllmplayground/vicuna_7b_v0',
-    'delta_ckpt_path': './pretrained_ckpt/pandagpt_ckpt/pytorch_model.pt',
+    'delta_ckpt_path': 'pretrained_ckpt/pandagpt_ckpt/pytorch_model.pt',
     'stage': 2,
     'max_tgt_len': 128,
     'lora_r': 32,
@@ -22,9 +23,8 @@ args = {
 model = OpenLLAMAPEFTModel(**args)
 delta_ckpt = torch.load(args['delta_ckpt_path'], map_location=torch.device('cpu'))
 model.load_state_dict(delta_ckpt, strict=False)
-model = model.eval().half().cuda()
-print(f'[!] init the model over ...')
-
+model = model.half().cuda().eval()
+print(f'[!] init the 13b model over ...')
 
 """Override Chatbot.postprocess"""
 
@@ -76,6 +76,25 @@ def parse_text(text):
     return text
 
 
+def re_predict(
+    input, 
+    image_path, 
+    audio_path, 
+    video_path, 
+    thermal_path, 
+    chatbot, 
+    max_length, 
+    top_p, 
+    temperature, 
+    history, 
+    modality_cache, 
+):
+    # drop the latest query and answers and generate again
+    q, a = history.pop()
+    chatbot.pop()
+    return predict(q, image_path, audio_path, video_path, thermal_path, chatbot, max_length, top_p, temperature, history, modality_cache)
+
+
 def predict(
     input, 
     image_path, 
@@ -90,11 +109,11 @@ def predict(
     modality_cache, 
 ):
     if image_path is None and audio_path is None and video_path is None and thermal_path is None:
-        return [(input, "There is no image/audio/video provided. Please upload the file to start a conversation.")]
+        return [(input, "图片和音频以及视频为空！请重新上传才能开启对话。")]
     else:
-        print(f'[!] image path: {image_path}\n[!] audio path: {audio_path}\n[!] video path: {video_path}\n[!] thermal pah: {thermal_path}')
-    # prepare the prompt
+        print(f'[!] image path: {image_path}\n[!] audio path: {audio_path}\n[!] video path: {video_path}\n[!] thermal path: {thermal_path}')
 
+    # prepare the prompt
     prompt_text = ''
     for idx, (q, a) in enumerate(history):
         if idx == 0:
@@ -106,17 +125,18 @@ def predict(
     else:
         prompt_text += f' Human: {input}'
 
-    response = model.generate({
-        'prompt': prompt_text,
-        'image_paths': [image_path] if image_path else [],
-        'audio_paths': [audio_path] if audio_path else [],
-        'video_paths': [video_path] if video_path else [],
-        'thermal_paths': [thermal_path] if thermal_path else [],
-        'top_p': top_p,
-        'temperature': temperature,
-        'max_tgt_len': max_length,
-        'modality_embeds': modality_cache
-    })
+    with torch.no_grad():
+        response = model.generate({
+            'prompt': prompt_text,
+            'image_paths': [image_path] if image_path else [],
+            'audio_paths': [audio_path] if audio_path else [],
+            'video_paths': [video_path] if video_path else [],
+            'thermal_paths': [thermal_path] if thermal_path else [],
+            'top_p': top_p,
+            'temperature': temperature,
+            'max_tgt_len': max_length,
+            'modality_embeds': modality_cache
+        })
     chatbot.append((parse_text(input), parse_text(response)))
     history.append((input, response))
     return chatbot, history, modality_cache
@@ -125,70 +145,44 @@ def predict(
 def reset_user_input():
     return gr.update(value='')
 
+def reset_dialog():
+    return [], []
 
 def reset_state():
     return None, None, None, None, [], [], []
 
 
-with gr.Blocks() as demo:
-    gr.HTML("""<h1 align="center">PandaGPT</h1>""")
+with gr.Blocks(scale=4) as demo:
+    gr.HTML("""<h1 align="center">PandaGPT</h1>
+
+    We note that the current online demo uses the 7B version of PandaGPT due to the limitation of computation resource. Better results should be expected when switching to the 13B version of PandaGPT. For more details on how to run 13B PandaGPT, please refer to our [main project repository](https://github.com/yxuansu/PandaGPT).
+            """)
 
     with gr.Row(scale=4):
-        with gr.Column(scale=2):
+        with gr.Column(scale=1):
             image_path = gr.Image(type="filepath", label="Image", value=None)
-
-            gr.Examples(
-                [ 
-                    os.path.join(os.path.dirname(__file__), "/assets/images/bird_image.jpg"),
-                    os.path.join(os.path.dirname(__file__), "/assets/images/dog_image.jpg"),
-                    os.path.join(os.path.dirname(__file__), "/assets/images/car_image.jpg"),
-                ],
-                image_path
-            )
-        with gr.Column(scale=2):
+        with gr.Column(scale=1):
             audio_path = gr.Audio(type="filepath", label="Audio", value=None)
-            gr.Examples(
-                [ 
-                    os.path.join(os.path.dirname(__file__), "/assets/audios/bird_audio.wav"),
-                    os.path.join(os.path.dirname(__file__), "/assets/audios/dog_audio.wav"),
-                    os.path.join(os.path.dirname(__file__), "/assets/audios/car_audio.wav"),
-                ],
-                audio_path
-            )
-    with gr.Row(scale=4):
-        with gr.Column(scale=2):
+        with gr.Column(scale=1):
             video_path = gr.Video(type='file', label="Video")
-
-            gr.Examples(
-                [ 
-                    os.path.join(os.path.dirname(__file__), "/assets/videos/world.mp4"),
-                    os.path.join(os.path.dirname(__file__), "/assets/videos/a.mp4"),
-                ],
-                video_path
-            )
-        with gr.Column(scale=2):
+        with gr.Column(scale=1):
             thermal_path = gr.Image(type="filepath", label="Thermal Image", value=None)
 
-            gr.Examples(
-                [ 
-                    os.path.join(os.path.dirname(__file__), "/assets/thermals/190662.jpg"),
-                    os.path.join(os.path.dirname(__file__), "/assets/thermals/210009.jpg"),
-                ],
-                thermal_path
-            )
-
-    chatbot = gr.Chatbot()
+    chatbot = gr.Chatbot().style(height=300)
     with gr.Row():
         with gr.Column(scale=4):
             with gr.Column(scale=12):
                 user_input = gr.Textbox(show_label=False, placeholder="Input...", lines=10).style(container=False)
             with gr.Column(min_width=32, scale=1):
-                submitBtn = gr.Button("Submit", variant="primary")
+                with gr.Row(scale=1):
+                    submitBtn = gr.Button("Submit", variant="primary")
+                with gr.Row(scale=1):
+                    resubmitBtn = gr.Button("Resubmit", variant="primary")
         with gr.Column(scale=1):
             emptyBtn = gr.Button("Clear History")
-            max_length = gr.Slider(0, 512, value=128, step=1.0, label="Maximum length", interactive=True)
-            top_p = gr.Slider(0, 1, value=0.4, step=0.01, label="Top P", interactive=True)
-            temperature = gr.Slider(0, 1, value=0.8, step=0.01, label="Temperature", interactive=True)
+            max_length = gr.Slider(0, 400, value=256, step=1.0, label="Maximum length", interactive=True)
+            top_p = gr.Slider(0, 1, value=0.01, step=0.01, label="Top P", interactive=True)
+            temperature = gr.Slider(0, 1, value=1.0, step=0.01, label="Temperature", interactive=True)
 
     history = gr.State([])
     modality_cache = gr.State([])
@@ -214,6 +208,28 @@ with gr.Blocks() as demo:
         show_progress=True
     )
 
+    resubmitBtn.click(
+        re_predict, [
+            user_input, 
+            image_path, 
+            audio_path, 
+            video_path, 
+            thermal_path, 
+            chatbot, 
+            max_length, 
+            top_p, 
+            temperature, 
+            history, 
+            modality_cache,
+        ], [
+            chatbot, 
+            history,
+            modality_cache
+        ],
+        show_progress=True
+    )
+
+
     submitBtn.click(reset_user_input, [], [user_input])
     emptyBtn.click(reset_state, outputs=[
         image_path,
@@ -225,4 +241,4 @@ with gr.Blocks() as demo:
         modality_cache
     ], show_progress=True)
 
-demo.queue().launch(share=False, inbrowser=True, server_name='0.0.0.0', server_port=24000)
+demo.launch(enable_queue=True)