support streaming output

app.py

@@ -134,29 +134,33 @@ def create_component(params, comp='Slider'):
 
 @spaces.GPU(duration=120)
 def chat(img, msgs, ctx, params=None, vision_hidden_states=None):
-    default_params = {"num_beams":3, "repetition_penalty": 1.2, "max_new_tokens": 1024}
+    default_params = {"stream": False, "sampling": False, "num_beams":3, "repetition_penalty": 1.2, "max_new_tokens": 1024}
     if params is None:
         params = default_params
     if img is None:
-        return -1, "Error, invalid image, please upload a new image", None, None
-    try:
-        image = img.convert('RGB')
-        answer = model.chat(
-            image=image,
-            msgs=msgs,
-            tokenizer=tokenizer,
-            **params
-        )
-        res = re.sub(r'(<box>.*</box>)', '', answer)
-        res = res.replace('<ref>', '')
-        res = res.replace('</ref>', '')
-        res = res.replace('<box>', '')
-        answer = res.replace('</box>', '')
-        return -1, answer, None, None
-    except Exception as err:
-        print(err)
-        traceback.print_exc()
-        return -1, ERROR_MSG, None, None
+        yield "Error, invalid image, please upload a new image"
+    else:
+        try:
+            image = img.convert('RGB')
+            answer = model.chat(
+                image=image,
+                msgs=msgs,
+                tokenizer=tokenizer,
+                **params
+            )
+            # if params['stream'] is False:
+                # res = re.sub(r'(<box>.*</box>)', '', answer)
+                # res = res.replace('<ref>', '')
+                # res = res.replace('</ref>', '')
+                # res = res.replace('<box>', '')
+                # answer = res.replace('</box>', '')
+            # else:
+            for char in answer:
+                yield char
+        except Exception as err:
+            print(err)
+            traceback.print_exc()
+            yield ERROR_MSG
 
 
 def upload_img(image, _chatbot, _app_session):
@@ -169,46 +173,51 @@ def upload_img(image, _chatbot, _app_session):
     return _chatbot, _app_session
 
 
-def respond(_question, _chat_bot, _app_cfg, params_form, num_beams, repetition_penalty, repetition_penalty_2, top_p, top_k, temperature):
+def respond(_chat_bot, _app_cfg, params_form, num_beams, repetition_penalty, repetition_penalty_2, top_p, top_k, temperature):
+    _question = _chat_bot[-1][0]
+    print('<Question>:', _question)
     if _app_cfg.get('ctx', None) is None:
-        _chat_bot.append((_question, 'Please upload an image to start'))
-        return '', _chat_bot, _app_cfg
-
-    _context = _app_cfg['ctx'].copy()
-    if _context:
-        _context.append({"role": "user", "content": _question})
+        _chat_bot[-1][1] = 'Please upload an image to start'
+        yield (_chat_bot, _app_cfg)
     else:
-        _context = [{"role": "user", "content": _question}] 
-    print('<User>:', _question)
+        _context = _app_cfg['ctx'].copy()
+        if _context:
+            _context.append({"role": "user", "content": _question})
+        else:
+            _context = [{"role": "user", "content": _question}]
+        if params_form == 'Beam Search':
+            params = {
+                'sampling': False,
+                'stream': False,
+                'num_beams': num_beams,
+                'repetition_penalty': repetition_penalty,
+                "max_new_tokens": 896 
+            }
+        else:
+            params = {
+                'sampling': True,
+                'stream': True,
+                'top_p': top_p,
+                'top_k': top_k,
+                'temperature': temperature,
+                'repetition_penalty': repetition_penalty_2,
+                "max_new_tokens": 896 
+            }
+    
+        gen = chat(_app_cfg['img'], _context, None, params)
+        _chat_bot[-1][1] = ""
+        for _char in gen:
+            _chat_bot[-1][1] += _char
+            _context[-1]["content"] += _char
+            yield (_chat_bot, _app_cfg)
 
-    if params_form == 'Beam Search':
-        params = {
-            'sampling': False,
-            'num_beams': num_beams,
-            'repetition_penalty': repetition_penalty,
-            "max_new_tokens": 896 
-        }
-    else:
-        params = {
-            'sampling': True,
-            'top_p': top_p,
-            'top_k': top_k,
-            'temperature': temperature,
-            'repetition_penalty': repetition_penalty_2,
-            "max_new_tokens": 896 
-        }
-    code, _answer, _, sts = chat(_app_cfg['img'], _context, None, params)
-    print('<Assistant>:', _answer)
 
-    _context.append({"role": "assistant", "content": _answer}) 
-    _chat_bot.append((_question, _answer))
-    if code == 0:
-        _app_cfg['ctx']=_context
-        _app_cfg['sts']=sts
+def request(_question, _chat_bot, _app_cfg):
+    _chat_bot.append((_question, None))
     return '', _chat_bot, _app_cfg
 
 
-def regenerate_button_clicked(_question, _chat_bot, _app_cfg, params_form, num_beams, repetition_penalty, repetition_penalty_2, top_p, top_k, temperature):
+def regenerate_button_clicked(_question, _chat_bot, _app_cfg):
     if len(_chat_bot) <= 1:
         _chat_bot.append(('Regenerate', 'No question for regeneration.'))
         return '', _chat_bot, _app_cfg
@@ -218,9 +227,18 @@ def regenerate_button_clicked(_question, _chat_bot, _app_cfg, params_form, num_b
         _question = _chat_bot[-1][0]
         _chat_bot = _chat_bot[:-1]
         _app_cfg['ctx'] = _app_cfg['ctx'][:-2]
-    return respond(_question, _chat_bot, _app_cfg, params_form, num_beams, repetition_penalty, repetition_penalty_2, top_p, top_k, temperature)
+    return request(_question, _chat_bot, _app_cfg)
+    # return respond(_chat_bot, _app_cfg, params_form, num_beams, repetition_penalty, repetition_penalty_2, top_p, top_k, temperature)
 
 
+def clear_button_clicked(_question, _chat_bot, _app_cfg, _bt_pic):
+    _chat_bot.clear()
+    _app_cfg['sts'] = None
+    _app_cfg['ctx'] = None
+    _app_cfg['img'] = None
+    _bt_pic = None
+    return '', _chat_bot, _app_cfg, _bt_pic
+    
 
 with gr.Blocks() as demo:
     with gr.Row():
@@ -235,25 +253,43 @@ with gr.Blocks() as demo:
                 temperature = create_component(temperature_slider)
                 repetition_penalty_2 = create_component(repetition_penalty_slider2)
             regenerate = create_component({'value': 'Regenerate'}, comp='Button')
+            clear = create_component({'value': 'Clear'}, comp='Button')
         with gr.Column(scale=3, min_width=500):
             app_session = gr.State({'sts':None,'ctx':None,'img':None})
             bt_pic = gr.Image(label="Upload an image to start")
             chat_bot = gr.Chatbot(label=f"Chat with {model_name}")
             txt_message = gr.Textbox(label="Input text")
             
-            regenerate.click(
-                regenerate_button_clicked,
-                [txt_message, chat_bot, app_session, params_form, num_beams, repetition_penalty, repetition_penalty_2, top_p, top_k, temperature],
-                [txt_message, chat_bot, app_session]
+            clear.click(
+                clear_button_clicked,
+                [txt_message, chat_bot, app_session, bt_pic],
+                [txt_message, chat_bot, app_session, bt_pic],
+                queue=False
             )
             txt_message.submit(
-                respond, 
-                [txt_message, chat_bot, app_session, params_form, num_beams, repetition_penalty, repetition_penalty_2, top_p, top_k, temperature], 
-                [txt_message, chat_bot, app_session]
+                request, 
+                #[txt_message, chat_bot, app_session, params_form, num_beams, repetition_penalty, repetition_penalty_2, top_p, top_k, temperature], 
+                [txt_message, chat_bot, app_session],
+                [txt_message, chat_bot, app_session],
+                queue=False
+            ).then(
+                respond,
+                [chat_bot, app_session, params_form, num_beams, repetition_penalty, repetition_penalty_2, top_p, top_k, temperature],
+                [chat_bot, app_session]
+            )
+            regenerate.click(
+                regenerate_button_clicked,
+                [txt_message, chat_bot, app_session],
+                [txt_message, chat_bot, app_session],
+                queue=False
+            ).then(
+                respond,
+                [chat_bot, app_session, params_form, num_beams, repetition_penalty, repetition_penalty_2, top_p, top_k, temperature],
+                [chat_bot, app_session]
             )
             bt_pic.upload(lambda: None, None, chat_bot, queue=False).then(upload_img, inputs=[bt_pic,chat_bot,app_session], outputs=[chat_bot,app_session])
 
 # launch
 #demo.launch(share=False, debug=True, show_api=False, server_port=8080, server_name="0.0.0.0")
+demo.queue()
 demo.launch()
-