Update app.py

app.py

@@ -1,6 +1,5 @@
 #!/usr/bin/env python
 # encoding: utf-8
-
 import spaces
 import gradio as gr
 from PIL import Image
@@ -8,13 +7,15 @@ import traceback
 import re
 import torch
 import argparse
-import logging
 from transformers import AutoModel, AutoTokenizer
-from huggingface_hub import hf_hub_download
 
-# Set up logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
+# README, How to run demo on different devices
+
+# For Nvidia GPUs.
+# python web_demo_2.5.py --device cuda
+
+# For Mac with MPS (Apple silicon or AMD GPUs).
+# PYTORCH_ENABLE_MPS_FALLBACK=1 python web_demo_2.5.py --device mps
 
 # Argparser
 parser = argparse.ArgumentParser(description='demo')
@@ -25,44 +26,30 @@ assert device in ['cuda', 'mps']
 
 # 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)
+tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+model.eval()
 
-def download_model_files(repo_id, filenames):
-    for filename in filenames:
-        try:
-            file_path = hf_hub_download(repo_id=repo_id, filename=filename, resume_download=True)
-            logger.info(f"Downloaded {filename} successfully.")
-        except Exception as e:
-            logger.error(f"Error downloading {filename}: {e}")
-            raise
-
-model_files = ["configuration_minicpm.py", "resampler.py", "modeling_minicpmv.py"]
-download_model_files(model_path, model_files)
 
-try:
-    if 'int4' in model_path:
-        if device == 'mps':
-            logger.error('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)
-    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
-    model.eval()
-except Exception as e:
-    logger.error(f"Error loading model or tokenizer: {e}")
-    raise
 
 ERROR_MSG = "Error, please retry"
 model_name = 'MiniCPM-Llama3-V 2.5'
 
 form_radio = {
     'choices': ['Beam Search', 'Sampling'],
+    #'value': 'Beam Search',
     'value': 'Sampling',
     'interactive': True,
     'label': 'Decode Type'
 }
-
+# Beam Form
 num_beams_slider = {
     'minimum': 0,
     'maximum': 5,
@@ -71,7 +58,6 @@ num_beams_slider = {
     'interactive': True,
     'label': 'Num Beams'
 }
-
 repetition_penalty_slider = {
     'minimum': 0,
     'maximum': 3,
@@ -80,7 +66,6 @@ repetition_penalty_slider = {
     'interactive': True,
     'label': 'Repetition Penalty'
 }
-
 repetition_penalty_slider2 = {
     'minimum': 0,
     'maximum': 3,
@@ -89,14 +74,13 @@ repetition_penalty_slider2 = {
     'interactive': True,
     'label': 'Repetition Penalty'
 }
-
 max_new_tokens_slider = {
     'minimum': 1,
     'maximum': 4096,
     'value': 1024,
     'step': 1,
     'interactive': True,
-    'label': 'Max New Tokens'
+    'label': 'Max New Tokens'    
 }
 
 top_p_slider = {
@@ -105,27 +89,26 @@ top_p_slider = {
     'value': 0.8,
     'step': 0.05,
     'interactive': True,
-    'label': 'Top P'
+    'label': 'Top P'    
 }
-
 top_k_slider = {
     'minimum': 0,
     'maximum': 200,
     'value': 100,
     'step': 1,
     'interactive': True,
-    'label': 'Top K'
+    'label': 'Top K'    
 }
-
 temperature_slider = {
     'minimum': 0,
     'maximum': 2,
     'value': 0.7,
     'step': 0.05,
     'interactive': True,
-    'label': 'Temperature'
+    'label': 'Temperature'    
 }
 
+
 def create_component(params, comp='Slider'):
     if comp == 'Slider':
         return gr.Slider(
@@ -151,7 +134,7 @@ 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 = {"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:
@@ -165,24 +148,34 @@ def chat(img, msgs, ctx, params=None, vision_hidden_states=None):
                 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:
-            logger.error(f"Error during chat: {err}")
+            print(err)
             traceback.print_exc()
             yield ERROR_MSG
 
+
 def upload_img(image, _chatbot, _app_session):
     image = Image.fromarray(image)
-    _app_session['sts'] = None
-    _app_session['ctx'] = []
-    _app_session['img'] = image
+
+    _app_session['sts']=None
+    _app_session['ctx']=[]
+    _app_session['img']=image 
     _chatbot.append(('', 'Image uploaded successfully, you can talk to me now'))
     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]
-    logger.info(f'<Question>: {_question}')
+    print('<Question>:', _question)
     if _app_cfg.get('ctx', None) is None:
         _chat_bot[-1][1] = 'Please upload an image to start'
         yield (_chat_bot, _app_cfg)
@@ -198,7 +191,7 @@ def respond(_chat_bot, _app_cfg, params_form, num_beams, repetition_penalty, rep
                 'stream': False,
                 'num_beams': num_beams,
                 'repetition_penalty': repetition_penalty,
-                "max_new_tokens": 896
+                "max_new_tokens": 896 
             }
         else:
             params = {
@@ -208,9 +201,9 @@ def respond(_chat_bot, _app_cfg, params_form, num_beams, repetition_penalty, rep
                 'top_k': top_k,
                 'temperature': temperature,
                 'repetition_penalty': repetition_penalty_2,
-                "max_new_tokens": 896
+                "max_new_tokens": 896 
             }
-
+    
         gen = chat(_app_cfg['img'], _context, None, params)
         _chat_bot[-1][1] = ""
         for _char in gen:
@@ -218,10 +211,12 @@ def respond(_chat_bot, _app_cfg, params_form, num_beams, repetition_penalty, rep
             _context[-1]["content"] += _char
             yield (_chat_bot, _app_cfg)
 
+
 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):
     if len(_chat_bot) <= 1:
         _chat_bot.append(('Regenerate', 'No question for regeneration.'))
@@ -233,6 +228,8 @@ def regenerate_button_clicked(_question, _chat_bot, _app_cfg):
         _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)
+
 
 def clear_button_clicked(_question, _chat_bot, _app_cfg, _bt_pic):
     _chat_bot.clear()
@@ -241,6 +238,7 @@ def clear_button_clicked(_question, _chat_bot, _app_cfg, _bt_pic):
     _app_cfg['img'] = None
     _bt_pic = None
     return '', _chat_bot, _app_cfg, _bt_pic
+    
 
 with gr.Blocks() as demo:
     with gr.Row():
@@ -257,11 +255,11 @@ with gr.Blocks() as demo:
             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})
+            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],
@@ -269,7 +267,8 @@ with gr.Blocks() as demo:
                 queue=False
             )
             txt_message.submit(
-                request,
+                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
@@ -288,8 +287,9 @@ with gr.Blocks() as demo:
                 [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])
+            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 the demo
+# 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()