Update app.py

app.py

@@ -1,5 +1,6 @@
 #!/usr/bin/env python
 # encoding: utf-8
+import timm
 import spaces
 import gradio as gr
 from PIL import Image
@@ -10,37 +11,39 @@ import argparse
 from transformers import AutoModel, AutoTokenizer
 
 # README, How to run demo on different devices
+# For Nvidia GPUs support BF16 (like A100, H100, RTX3090)
+# python web_demo.py --device cuda --dtype bf16
 
-# For Nvidia GPUs.
-# python web_demo_2.5.py --device cuda
+# For Nvidia GPUs do NOT support BF16 (like V100, T4, RTX2080)
+# python web_demo.py --device cuda --dtype fp16
 
 # For Mac with MPS (Apple silicon or AMD GPUs).
-# PYTORCH_ENABLE_MPS_FALLBACK=1 python web_demo_2.5.py --device mps
+# PYTORCH_ENABLE_MPS_FALLBACK=1 python web_demo.py --device mps --dtype fp16
 
 # Argparser
 parser = argparse.ArgumentParser(description='demo')
 parser.add_argument('--device', type=str, default='cuda', help='cuda or mps')
+parser.add_argument('--dtype', type=str, default='bf16', help='bf16 or fp16')
 args = parser.parse_args()
 device = args.device
 assert device in ['cuda', 'mps']
+if args.dtype == 'bf16':
+    dtype = torch.bfloat16
+else:
+    dtype = torch.float16
 
 # Load model
-model_path = 'openbmb/MiniCPM-Llama3-V-2_5'
-if 'int4' in model_path:
-    if device == 'mps':
-        print('Error: running int4 model with bitsandbytes on Mac is not supported right now.')
-        exit()
-    model = AutoModel.from_pretrained(model_path, trust_remote_code=True)
-else:
-    model = AutoModel.from_pretrained(model_path, trust_remote_code=True).to(dtype=torch.float16)
-    model = model.to(device=device)
+model_path = 'openbmb/MiniCPM-V-2'
+model = AutoModel.from_pretrained(model_path, trust_remote_code=True).to(dtype=torch.bfloat16)
 tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+
+model = model.to(device=device, dtype=dtype)
 model.eval()
 
 
 
 ERROR_MSG = "Error, please retry"
-model_name = 'MiniCPM-Llama3-V 2.5'
+model_name = 'MiniCPM-V 2.0'
 
 form_radio = {
     'choices': ['Beam Search', 'Sampling'],
@@ -134,33 +137,30 @@ def create_component(params, comp='Slider'):
 
 @spaces.GPU(duration=120)
 def chat(img, msgs, ctx, params=None, vision_hidden_states=None):
-    default_params = {"stream": False, "sampling": False, "num_beams":3, "repetition_penalty": 1.2, "max_new_tokens": 1024}
+    default_params = {"num_beams":3, "repetition_penalty": 1.2, "max_new_tokens": 1024}
     if params is None:
         params = default_params
     if img is 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
+        return -1, "Error, invalid image, please upload a new image", None, None
+    try:
+        image = img.convert('RGB')
+        answer, context, _ = model.chat(
+            image=image,
+            msgs=msgs,
+            context=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>', '')
+        return -1, answer, None, None
+    except Exception as err:
+        print(err)
+        traceback.print_exc()
+        return -1, ERROR_MSG, None, None
 
 
 def upload_img(image, _chatbot, _app_session):
@@ -173,51 +173,46 @@ def upload_img(image, _chatbot, _app_session):
     return _chatbot, _app_session
 
 
-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)
+def respond(_question, _chat_bot, _app_cfg, params_form, num_beams, repetition_penalty, repetition_penalty_2, top_p, top_k, temperature):
     if _app_cfg.get('ctx', None) is None:
-        _chat_bot[-1][1] = 'Please upload an image to start'
-        yield (_chat_bot, _app_cfg)
+        _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})
     else:
-        _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)
+        _context = [{"role": "user", "content": _question}] 
+    print('<User>:', _question)
 
+    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)
 
-def request(_question, _chat_bot, _app_cfg):
-    _chat_bot.append((_question, None))
+    _context.append({"role": "assistant", "content": _answer}) 
+    _chat_bot.append((_question, _answer))
+    if code == 0:
+        _app_cfg['ctx']=_context
+        _app_cfg['sts']=sts
     return '', _chat_bot, _app_cfg
 
 
-def regenerate_button_clicked(_question, _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):
     if len(_chat_bot) <= 1:
         _chat_bot.append(('Regenerate', 'No question for regeneration.'))
         return '', _chat_bot, _app_cfg
@@ -227,18 +222,9 @@ def regenerate_button_clicked(_question, _chat_bot, _app_cfg):
         _question = _chat_bot[-1][0]
         _chat_bot = _chat_bot[:-1]
         _app_cfg['ctx'] = _app_cfg['ctx'][:-2]
-    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)
+    return respond(_question, _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():
@@ -253,43 +239,27 @@ 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")
             
-            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(
-                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]
+                [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.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]
             )
             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()
+demo.launch()
+
+
+