streaming output

README.md

@@ -4,7 +4,7 @@ emoji: 💬
 colorFrom: yellow
 colorTo: purple
 sdk: gradio
-sdk_version: 4.22.0
+sdk_version: 4.41.0
 app_file: app.py
 pinned: false
 license: apache-2.0

app.py

@@ -26,10 +26,6 @@ import modelscope_studio as mgr
 # For Mac with MPS (Apple silicon or AMD GPUs).
 # PYTORCH_ENABLE_MPS_FALLBACK=1 python web_demo_2.6.py --device mps
 
-os.system("pip list|grep torch")
-os.system("pip list|grep trans")
-os.system("pip list|grep flash")
-
 # Argparser
 parser = argparse.ArgumentParser(description='demo')
 parser.add_argument('--device', type=str, default='cuda', help='cuda or mps')
@@ -131,7 +127,7 @@ def create_component(params, comp='Slider'):
 
 
 def create_multimodal_input(upload_image_disabled=False, upload_video_disabled=False):
-    return mgr.MultimodalInput(upload_image_button_props={'label': 'Upload Image', 'disabled': upload_image_disabled, 'file_count': 'multiple'}, 
+    return mgr.MultimodalInput(value=None, upload_image_button_props={'label': 'Upload Image', 'disabled': upload_image_disabled, 'file_count': 'multiple'}, 
                                         upload_video_button_props={'label': 'Upload Video', 'disabled': upload_video_disabled, 'file_count': 'single'},
                                         submit_button_props={'label': 'Submit'})
 
@@ -139,6 +135,8 @@ def create_multimodal_input(upload_image_disabled=False, upload_video_disabled=F
 @spaces.GPU(duration=120)
 def chat(img, msgs, ctx, params=None, vision_hidden_states=None):
     try:
+        if msgs[-1]['role'] == 'assistant':
+            msgs = msgs[:-1] # remove last which is added for streaming
         print('msgs:', msgs)
         answer = model.chat(
             image=None,
@@ -146,17 +144,18 @@ def chat(img, msgs, ctx, params=None, vision_hidden_states=None):
             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>', '')
-        print('answer:', answer)
-        return 0, answer, None, None
+        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>', '')
+        for char in answer:
+            yield char
     except Exception as e:
         print(e)
         traceback.print_exc()
-        return -1, ERROR_MSG, None, None
+        yield ERROR_MSG
 
 
 def encode_image(image):
@@ -270,11 +269,7 @@ def count_video_frames(_context):
     return num_frames
 
 
-def respond(_question, _chat_bot, _app_cfg, params_form):
-    print("[respond] question:", _question)
-    _context = _app_cfg['ctx'].copy()
-    _context.append({'role': 'user', 'content': encode_message(_question)})
-
+def request(_question, _chat_bot, _app_cfg):
     images_cnt = _app_cfg['images_cnt']
     videos_cnt = _app_cfg['videos_cnt']
     files_cnts = check_has_videos(_question)
@@ -284,47 +279,67 @@ def respond(_question, _chat_bot, _app_cfg, params_form):
     if files_cnts[1] + videos_cnt + files_cnts[0] + images_cnt <= 0:
         gr.Warning("Please chat with at least one image or video.")
         return _question, _chat_bot, _app_cfg
-
-    if params_form == 'Beam Search':
-        params = {
-            'sampling': False,
-            'num_beams': 3,
-            'repetition_penalty': 1.2,
-            "max_new_tokens": 2048
-        }
-    else:
-        params = {
-            'sampling': True,
-            'top_p': 0.8,
-            'top_k': 100,
-            'temperature': 0.7,
-            'repetition_penalty': 1.05,
-            "max_new_tokens": 2048
-        }
-    params["max_inp_length"] = 4352 # 4096+256
-
-    if files_cnts[1] + videos_cnt > 0:
-        #params["max_inp_length"] = 4352 # 4096+256
-        params["use_image_id"] = False
-        params["max_slice_nums"] = 1 if count_video_frames(_context) > 16 else 2
-
-    code, _answer, _, sts = chat("", _context, None, params)
-
+    _chat_bot.append((_question, None))
     images_cnt += files_cnts[0]
     videos_cnt += files_cnts[1]
-    _context.append({"role": "assistant", "content": [make_text(_answer)]}) 
-    _chat_bot.append((_question, _answer))
-    if code == 0:
-        _app_cfg['ctx']=_context
-        _app_cfg['sts']=sts
     _app_cfg['images_cnt'] = images_cnt
     _app_cfg['videos_cnt'] = videos_cnt
-
     upload_image_disabled = videos_cnt > 0
     upload_video_disabled = videos_cnt > 0 or images_cnt > 0
     return create_multimodal_input(upload_image_disabled, upload_video_disabled), _chat_bot, _app_cfg
 
 
+def respond(_chat_bot, _app_cfg, params_form):
+    if len(_app_cfg) == 0:
+        yield (_chat_bot, _app_cfg)
+    elif _app_cfg['images_cnt'] == 0 and _app_cfg['videos_cnt'] == 0:
+        yield(_chat_bot, _app_cfg)
+    else:
+        _question = _chat_bot[-1][0]
+        _context = _app_cfg['ctx'].copy()
+        _context.append({'role': 'user', 'content': encode_message(_question)})
+
+        videos_cnt = _app_cfg['videos_cnt']
+
+        if params_form == 'Beam Search':
+            params = {
+                'sampling': False,
+                'stream': False,
+                'num_beams': 3,
+                'repetition_penalty': 1.2,
+                "max_new_tokens": 2048
+            }
+        else:
+            params = {
+                'sampling': True,
+                'stream': True,
+                'top_p': 0.8,
+                'top_k': 100,
+                'temperature': 0.7,
+                'repetition_penalty': 1.05,
+                "max_new_tokens": 2048
+            }
+        params["max_inp_length"] = 4352 # 4096+256
+
+        if videos_cnt > 0:
+            #params["max_inp_length"] = 4352 # 4096+256
+            params["use_image_id"] = False
+            params["max_slice_nums"] = 1 if count_video_frames(_context) > 16 else 2
+
+        gen = chat("", _context, None, params)
+
+        _context.append({"role": "assistant", "content": [""]}) 
+        _chat_bot[-1][1] = ""
+
+        for _char in gen:
+            _chat_bot[-1][1] += _char
+            _context[-1]["content"][0] += _char
+            yield (_chat_bot, _app_cfg)
+        
+        _app_cfg['ctx']=_context
+        yield (_chat_bot, _app_cfg)
+
+
 def fewshot_add_demonstration(_image, _user_message, _assistant_message, _chat_bot, _app_cfg):
     ctx = _app_cfg["ctx"]
     message_item = []
@@ -332,6 +347,7 @@ def fewshot_add_demonstration(_image, _user_message, _assistant_message, _chat_b
         image = Image.open(_image).convert("RGB")
         ctx.append({"role": "user", "content": [encode_image(image), make_text(_user_message)]})
         message_item.append({"text": "[mm_media]1[/mm_media]" + _user_message, "files": [_image]})
+        _app_cfg["images_cnt"] += 1
     else:
         if _user_message:
             ctx.append({"role": "user", "content": [make_text(_user_message)]})
@@ -348,65 +364,29 @@ def fewshot_add_demonstration(_image, _user_message, _assistant_message, _chat_b
     return None, "", "", _chat_bot, _app_cfg
 
 
-def fewshot_respond(_image, _user_message, _chat_bot, _app_cfg, params_form):
-    user_message_contents = []
-    _context = _app_cfg["ctx"].copy()
-    images_cnt = _app_cfg["images_cnt"]
-    if _image:
-        image = Image.open(_image).convert("RGB")
-        user_message_contents += [encode_image(image)]
-        images_cnt += 1
-    if _user_message:
-        user_message_contents += [make_text(_user_message)]
-    if user_message_contents:
-        _context.append({"role": "user", "content": user_message_contents})
-
-    if params_form == 'Beam Search':
-        params = {
-            'sampling': False,
-            'num_beams': 3,
-            'repetition_penalty': 1.2,
-            "max_new_tokens": 2048
-        }
-    else:
-        params = {
-            'sampling': True,
-            'top_p': 0.8,
-            'top_k': 100,
-            'temperature': 0.7,
-            'repetition_penalty': 1.05,
-            "max_new_tokens": 2048
-        }
-
-    if images_cnt == 0:
-        gr.Warning("Please chat with at least one image or video.")
-        return _image, _user_message, '', _chat_bot, _app_cfg
-    
-    code, _answer, _, sts = chat("", _context, None, params)
-
-    _context.append({"role": "assistant", "content": [make_text(_answer)]})
-
+def fewshot_request(_image, _user_message, _chat_bot, _app_cfg):
+    if _app_cfg["images_cnt"] == 0 and not _image:
+        gr.Warning("Please chat with at least one image.")
+        return None, '', '', _chat_bot, _app_cfg
     if _image:
         _chat_bot.append([
             {"text": "[mm_media]1[/mm_media]" + _user_message, "files": [_image]},
-            {"text": _answer, "files": []}        
+            ""        
         ])
+        _app_cfg["images_cnt"] += 1
     else:
         _chat_bot.append([
             {"text": _user_message, "files": [_image]},
-            {"text": _answer, "files": []}        
+            ""
         ])
-    if code == 0:
-        _app_cfg['ctx']=_context
-        _app_cfg['sts']=sts
-        _app_cfg['images_cnt'] = images_cnt
+
     return None, '', '', _chat_bot, _app_cfg
 
 
-def regenerate_button_clicked(_question, _image, _user_message, _assistant_message, _chat_bot, _app_cfg, params_form):
+def regenerate_button_clicked(_chat_bot, _app_cfg):
     if len(_chat_bot) <= 1 or not _chat_bot[-1][1]:
         gr.Warning('No question for regeneration.')
-        return '', _image, _user_message, _assistant_message, _chat_bot, _app_cfg
+        return None, None, '', '', _chat_bot, _app_cfg
     if _app_cfg["chat_type"] == "Chat":
         images_cnt = _app_cfg['images_cnt']
         videos_cnt = _app_cfg['videos_cnt']
@@ -418,10 +398,9 @@ def regenerate_button_clicked(_question, _image, _user_message, _assistant_messa
         videos_cnt -= files_cnts[1]
         _app_cfg['images_cnt'] = images_cnt
         _app_cfg['videos_cnt'] = videos_cnt
-        upload_image_disabled = videos_cnt > 0
-        upload_video_disabled = videos_cnt > 0 or images_cnt > 0
-        _question, _chat_bot, _app_cfg = respond(_question, _chat_bot, _app_cfg, params_form)
-        return _question, _image, _user_message, _assistant_message, _chat_bot, _app_cfg
+
+        _question, _chat_bot, _app_cfg = request(_question, _chat_bot, _app_cfg)
+        return _question, None, '', '', _chat_bot, _app_cfg
     else: 
         last_message = _chat_bot[-1][0]
         last_image = None
@@ -430,10 +409,9 @@ def regenerate_button_clicked(_question, _image, _user_message, _assistant_messa
             last_user_message = last_message.text
         if last_message.files:
             last_image = last_message.files[0].file.path
-        _chat_bot = _chat_bot[:-1]
+        _chat_bot[-1][1] = ""
         _app_cfg['ctx'] = _app_cfg['ctx'][:-2]
-        _image, _user_message, _assistant_message, _chat_bot, _app_cfg = fewshot_respond(last_image, last_user_message, _chat_bot, _app_cfg, params_form)
-        return _question, _image, _user_message, _assistant_message, _chat_bot, _app_cfg
+        return _question, None, '', '', _chat_bot, _app_cfg
 
 
 def flushed():
@@ -469,7 +447,6 @@ init_conversation = [
 
 
 css = """
-video { height: auto !important; }
 .example label { font-size: 16px;}
 """
 
@@ -503,9 +480,13 @@ with gr.Blocks(css=css) as demo:
                     chat_tab_label = gr.Textbox(value="Chat", interactive=False, visible=False)
 
                     txt_message.submit(
-                        respond,
-                        [txt_message, chat_bot, app_session, params_form], 
+                        request,
+                        [txt_message, chat_bot, app_session], 
                         [txt_message, chat_bot, app_session]
+                    ).then(
+                        respond,
+                        [chat_bot, app_session, params_form],
+                        [chat_bot, app_session]
                     )
 
                 with gr.Tab("Few Shot") as fewshot_tab:
@@ -525,9 +506,13 @@ with gr.Blocks(css=css) as demo:
                         [image_input, user_message, assistant_message, chat_bot, app_session]
                     )
                     generate_button.click(
-                        fewshot_respond,
-                        [image_input, user_message, chat_bot, app_session, params_form],
+                        fewshot_request,
+                        [image_input, user_message, chat_bot, app_session],
                         [image_input, user_message, assistant_message, chat_bot, app_session]
+                    ).then(
+                        respond,
+                        [chat_bot, app_session, params_form],
+                        [chat_bot, app_session]
                     )
 
                 chat_tab.select(
@@ -556,8 +541,12 @@ with gr.Blocks(css=css) as demo:
                 )
                 regenerate.click(
                     regenerate_button_clicked,
-                    [txt_message, image_input, user_message, assistant_message, chat_bot, app_session, params_form],
+                    [chat_bot, app_session],
                     [txt_message, image_input, user_message, assistant_message, chat_bot, app_session]
+                ).then(
+                    respond,
+                    [chat_bot, app_session, params_form],
+                    [chat_bot, app_session]
                 )
                 clear_button.click(
                     clear,

requirements.txt

@@ -6,6 +6,7 @@ sentencepiece==0.1.99
 https://github.com/Dao-AILab/flash-attention/releases/download/v2.6.2/flash_attn-2.6.2+cu123torch2.1cxx11abiFALSE-cp310-cp310-linux_x86_64.whl
 opencv-python
 decord
-gradio==4.22.0
+#gradio==4.22.0
+gradio==4.41.0
 http://thunlp.oss-cn-qingdao.aliyuncs.com/multi_modal/never_delete/modelscope_studio-0.4.0.9-py3-none-any.whl
 accelerate