update

app.py

@@ -1,101 +1,313 @@
 import sys
 import os
-import spaces
 import argparse
 import time
 import subprocess
-import torch
+import spaces
 import cumo.serve.gradio_web_server as gws
 
-#os.system("export BUILD_WITH_CUDA=True")
-#os.system("pip install --upgrade pip")
-#os.system("pip install flash-attn --no-build-isolation")
+import datetime
+import json
+
+import gradio as gr
+import requests
+from PIL import Image
+
+from cumo.conversation import (default_conversation, conv_templates, SeparatorStyle)
+from cumo.constants import LOGDIR
+from cumo.utils import (build_logger, server_error_msg, violates_moderation, moderation_msg)
+import hashlib
+
+import torch
+import io
+from cumo.constants import WORKER_HEART_BEAT_INTERVAL
+from cumo.utils import (build_logger, server_error_msg,
+    pretty_print_semaphore)
+from cumo.model.builder import load_pretrained_model
+from cumo.mm_utils import process_images, load_image_from_base64, tokenizer_image_token
+from cumo.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN
+from transformers import TextIteratorStreamer
+from threading import Thread
 
 # Execute the pip install command with additional options
-#subprocess.check_call([sys.executable, '-m', 'pip', 'install', 'flash-attn', '--no-build-isolation', '-U'])
-
-def start_controller():
-    print("Starting the controller")
-    controller_command = [
-        sys.executable,
-        "-m",
-        "cumo.serve.controller",
-        "--host",
-        "0.0.0.0",
-        "--port",
-        "10000",
-    ]
-    print(controller_command)
-    return subprocess.Popen(controller_command)
-
-def start_worker(model_path: str, bits=16):
-    print(f"Starting the model worker for the model {model_path}")
-    model_name = model_path.strip("/").split("/")[-1]
-    assert bits in [4, 8, 16], "It can be only loaded with 16-bit, 8-bit, and 4-bit."
-    if bits != 16:
-        model_name += f"-{bits}bit"
-    worker_command = [
-        sys.executable,
-        "-m",
-        "cumo.serve.model_worker",
-        "--host",
-        "0.0.0.0",
-        "--controller",
-        "http://localhost:10000",
-        "--model-path",
-        model_path,
-        "--model-name",
-        model_name,
-    ]
-    if bits != 16:
-        worker_command += [f"--load-{bits}bit"]
-    print(worker_command)
-    return subprocess.Popen(worker_command)
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--host", type=str, default="0.0.0.0")
-    parser.add_argument("--port", type=int)
-    parser.add_argument("--model-path", type=str, default="checkpoints/CuMo-mistral-7b")
-    parser.add_argument("--model-base", type=str, default=None)
-    parser.add_argument("--controller-url", type=str, default="http://localhost:10000")
-    parser.add_argument("--concurrency-count", type=int, default=5)
-    parser.add_argument("--bits", type=int, default=16)
-    parser.add_argument("--model-list-mode", type=str, default="reload", choices=["once", "reload"])
-    parser.add_argument("--share", action="store_true")
-    parser.add_argument("--moderate", action="store_true")
-    parser.add_argument("--embed", action="store_true")
-    gws.args = parser.parse_args()
-    gws.models = []
-
-    print(f"args: {gws.args}")
-    model_path = gws.args.model_path
-    bits = gws.args.bits
-    concurrency_count = int(os.getenv("concurrency_count", 5))
-    #device = "cuda" if torch.cuda.is_available() else "cpu"
-    controller_proc = start_controller()
-    worker_proc = start_worker(model_path, bits=bits)
-
-    # Wait for worker and controller to start
-    time.sleep(10)
-    exit_status = 0
-    try:
-        demo = gws.build_demo(embed_mode=False, concurrency_count=concurrency_count)
-        demo.queue(
-            status_update_rate=10,
-            api_open=False
-        ).launch(
-            server_name=gws.args.host,
-            server_port=gws.args.port,
-            share=gws.args.share
-        )
-
-    except Exception as e:
-        print(e)
-        exit_status = 1
-    finally:
-        worker_proc.kill()
-        controller_proc.kill()
-
-        sys.exit(exit_status)
+#subprocess.check_call([sys.executable, '-m', 'pip', 'install', 'flash-attn', '--no-build-isolation', '-U']
+
+headers = {"User-Agent": "CuMo"}
+
+no_change_btn = gr.Button()
+enable_btn = gr.Button(interactive=True)
+disable_btn = gr.Button(interactive=False)
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model_path = './checkpoints/CuMo-mistral-7b'
+model_base = 'mistralai/Mistral-7B-Instruct-v0.2'
+model_name = 'CuMo-mistral-7b'
+conv_mode = 'mistral_instruct_system'
+load_8bit = False
+load_4bit = False
+tokenizer, model, image_processor, context_len = load_pretrained_model(model_path, model_base, model_name, load_8bit, load_4bit, device=device, use_flash_attn=False)
+model.config.training = False
+             
+def upvote_last_response(state):
+    return ("",) + (disable_btn,) * 3
+
+
+def downvote_last_response(state):
+    return ("",) + (disable_btn,) * 3
+
+
+def flag_last_response(state):
+    return ("",) + (disable_btn,) * 3
+
+def clear_history():
+    state = default_conversation.copy()
+    return (state, state.to_gradio_chatbot(), "", None) + (disable_btn,) * 5
+
+def add_text(state, imagebox, textbox, image_process_mode):
+    if state is None:
+        state = conv_templates[conv_mode].copy()
+
+    if imagebox is not None:
+        textbox = DEFAULT_IMAGE_TOKEN + '\n' + textbox
+        image = Image.open(imagebox).convert('RGB')
+
+    if imagebox is not None:
+        textbox = (textbox, image, image_process_mode)
+
+    state.append_message(state.roles[0], textbox)
+    state.append_message(state.roles[1], None)
+
+    yield (state, state.to_gradio_chatbot(), "", None) + (disable_btn, disable_btn, disable_btn, enable_btn, enable_btn)
+
+def delete_text(state, image_process_mode):
+    state.messages[-1][-1] = None
+    prev_human_msg = state.messages[-2]
+    if type(prev_human_msg[1]) in (tuple, list):
+        prev_human_msg[1] = (*prev_human_msg[1][:2], image_process_mode)
+    yield (state, state.to_gradio_chatbot(), "", None) + (disable_btn, disable_btn, disable_btn, enable_btn, enable_btn)
+
+def regenerate(state, image_process_mode):
+    state.messages[-1][-1] = None
+    prev_human_msg = state.messages[-2]
+    if type(prev_human_msg[1]) in (tuple, list):
+        prev_human_msg[1] = (*prev_human_msg[1][:2], image_process_mode)
+    state.skip_next = False
+    return (state, state.to_gradio_chatbot(), "", None) + (disable_btn,) * 5
+
+@spaces.GPU
+def generate(state, imagebox, textbox, image_process_mode, temperature, top_p, max_output_tokens):
+    prompt = state.get_prompt()
+    images = state.get_images(return_pil=True)
+    #prompt, image_args = process_image(prompt, images)
+
+    ori_prompt = prompt
+    num_image_tokens = 0
+
+    if images is not None and len(images) > 0:
+        if len(images) > 0:
+            if len(images) != prompt.count(DEFAULT_IMAGE_TOKEN):
+                raise ValueError("Number of images does not match number of <image> tokens in prompt")
+            
+            #images = [load_image_from_base64(image) for image in images]
+            image_sizes = [image.size for image in images]
+            images = process_images(images, image_processor, model.config)
+
+            if type(images) is list:
+                images = [image.to(model.device, dtype=torch.float16) for image in images]
+            else:
+                images = images.to(model.device, dtype=torch.float16)
+
+            replace_token = DEFAULT_IMAGE_TOKEN
+            if getattr(model.config, 'mm_use_im_start_end', False):
+                replace_token = DEFAULT_IM_START_TOKEN + replace_token + DEFAULT_IM_END_TOKEN
+            prompt = prompt.replace(DEFAULT_IMAGE_TOKEN, replace_token)
+
+            num_image_tokens = prompt.count(replace_token) * model.get_vision_tower().num_patches
+        else:
+            images = None
+            image_sizes = None
+        image_args = {"images": images, "image_sizes": image_sizes}
+    else:
+        images = None
+        image_args = {}
+
+    max_context_length = getattr(model.config, 'max_position_embeddings', 2048)
+    max_new_tokens = 512
+    do_sample = True if temperature > 0.001 else False
+    stop_str = state.sep if state.sep_style in [SeparatorStyle.SINGLE, SeparatorStyle.MPT] else state.sep2
+
+    input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).to(model.device)
+    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True, timeout=15)
+
+    max_new_tokens = min(max_new_tokens, max_context_length - input_ids.shape[-1] - num_image_tokens)
+
+    if max_new_tokens < 1:
+        yield json.dumps({"text": ori_prompt + "Exceeds max token length. Please start a new conversation, thanks.", "error_code": 0}).encode() + b"\0"
+        return
+
+    thread = Thread(target=model.generate, kwargs=dict(
+        inputs=input_ids,
+        do_sample=do_sample,
+        temperature=temperature,
+        top_p=top_p,
+        max_new_tokens=max_new_tokens,
+        streamer=streamer,
+        use_cache=True,
+        pad_token_id=tokenizer.eos_token_id,
+        **image_args
+    ))
+    thread.start()
+    generated_text = ''
+    for new_text in streamer:
+        generated_text += new_text
+        if generated_text.endswith(stop_str):
+            generated_text = generated_text[:-len(stop_str)]
+        state.messages[-1][-1] = generated_text
+        yield (state, state.to_gradio_chatbot(), "", None) + (disable_btn, disable_btn, disable_btn, enable_btn, enable_btn)
+    
+    yield (state, state.to_gradio_chatbot(), "", None) + (enable_btn,) * 5
+    
+    torch.cuda.empty_cache()
+
+title_markdown = ("""
+# CuMo: Scaling Multimodal LLM with Co-Upcycled Mixture-of-Experts
+[[Project Page](https://chrisjuniorli.github.io/project/CuMo/)] [[Code](https://github.com/SHI-Labs/CuMo)] [[Model](https://huggingface.co/shi-labs/CuMo-mistral-7b)] | 📚 [[Arxiv](https://arxiv.org/pdf/2405.05949)]]
+""")
+
+tos_markdown = ("""
+### Terms of use
+By using this service, users are required to agree to the following terms:
+The service is a research preview intended for non-commercial use only. It only provides limited safety measures and may generate offensive content. It must not be used for any illegal, harmful, violent, racist, or sexual purposes. The service may collect user dialogue data for future research.
+Please click the "Flag" button if you get any inappropriate answer! We will collect those to keep improving our moderator.
+For an optimal experience, please use desktop computers for this demo, as mobile devices may compromise its quality.
+""")
+
+
+learn_more_markdown = ("""
+### License
+The service is a research preview intended for non-commercial use only, subject to the. Please contact us if you find any potential violation.
+""")
+
+block_css = """
+
+#buttons button {
+    min-width: min(120px,100%);
+}
+
+"""
+
+textbox = gr.Textbox(show_label=False, placeholder="Enter text and press ENTER", container=False)
+with gr.Blocks(title="CuMo", theme=gr.themes.Default(), css=block_css) as demo:
+    state = gr.State()
+
+    gr.Markdown(title_markdown)
+
+    with gr.Row():
+        with gr.Column(scale=3):
+            imagebox = gr.Image(label="Input Image", type="filepath")
+            image_process_mode = gr.Radio(
+                ["Crop", "Resize", "Pad", "Default"],
+                value="Default",
+                label="Preprocess for non-square image", visible=False)
+
+    
+            #cur_dir = os.path.dirname(os.path.abspath(__file__))
+            cur_dir = './cumo/serve'
+            gr.Examples(examples=[
+                [f"{cur_dir}/examples/aveger.jpg", "Can you introduce this movie based on the poster?"],
+                [f"{cur_dir}/examples/fridge.webp", "Can you describe what groceries are presented in this fridge?"],
+                [f"{cur_dir}/examples/su7_4.jpg", "What car is it in this image?"],
+                [f"{cur_dir}/examples/nvidia.jpeg", "Can you tell me what happened in this image?"],
+                [f"{cur_dir}/examples/animal.webp", "What animals are in this image?"],
+                [f"{cur_dir}/examples/noodle.png", "How many calories estimated in this bowl?"],
+                [f"{cur_dir}/examples/disney.jpeg", "How many characters in this image?"],
+                [f"{cur_dir}/examples/reka_6.jpeg", "What colour is my hat (im sitting on the bear)?"],
+            ], inputs=[imagebox, textbox], cache_examples=False)
+
+            with gr.Accordion("Parameters", open=False) as parameter_row:
+                temperature = gr.Slider(minimum=0.0, maximum=1.0, value=0.2, step=0.1, interactive=True, label="Temperature",)
+                top_p = gr.Slider(minimum=0.0, maximum=1.0, value=0.7, step=0.1, interactive=True, label="Top P",)
+                max_output_tokens = gr.Slider(minimum=0, maximum=1024, value=512, step=64, interactive=True, label="Max output tokens",)
+
+        with gr.Column(scale=8):
+            chatbot = gr.Chatbot(
+                elem_id="chatbot",
+                label="CuMo Chatbot",
+                height=650,
+                layout="panel",
+            )
+            with gr.Row():
+                with gr.Column(scale=8):
+                    textbox.render()
+                with gr.Column(scale=1, min_width=50):
+                    submit_btn = gr.Button(value="Send", variant="primary")
+            with gr.Row(elem_id="buttons") as button_row:
+                upvote_btn = gr.Button(value="👍  Upvote", interactive=False)
+                downvote_btn = gr.Button(value="👎  Downvote", interactive=False)
+                flag_btn = gr.Button(value="⚠️  Flag", interactive=False)
+                #stop_btn = gr.Button(value="⏹️  Stop Generation", interactive=False)
+                regenerate_btn = gr.Button(value="🔄  Regenerate", interactive=False)
+                clear_btn = gr.Button(value="🗑️  Clear", interactive=False)
+
+    gr.Markdown(tos_markdown)
+    gr.Markdown(learn_more_markdown)
+    url_params = gr.JSON(visible=False)
+
+    # Register listeners
+    btn_list = [upvote_btn, downvote_btn, flag_btn, regenerate_btn, clear_btn]
+    upvote_btn.click(
+        upvote_last_response,
+        [state],
+        [textbox, upvote_btn, downvote_btn, flag_btn]
+    )
+    downvote_btn.click(
+        downvote_last_response,
+        [state],
+        [textbox, upvote_btn, downvote_btn, flag_btn]
+    )
+    flag_btn.click(
+        flag_last_response,
+        [state],
+        [textbox, upvote_btn, downvote_btn, flag_btn]
+    )
+
+    clear_btn.click(
+        clear_history,
+        None,
+        [state, chatbot, textbox, imagebox] + btn_list,
+        queue=False
+    )
+
+    regenerate_btn.click(
+        delete_text,
+        [state, image_process_mode],
+        [state, chatbot, textbox, imagebox] + btn_list,
+    ).then(
+        generate,
+        [state, imagebox, textbox, image_process_mode, temperature, top_p, max_output_tokens],
+        [state, chatbot, textbox, imagebox] + btn_list,
+    )
+    textbox.submit(
+        add_text,
+        [state, imagebox, textbox, image_process_mode],
+        [state, chatbot, textbox, imagebox] + btn_list,
+    ).then(
+        generate,
+        [state, imagebox, textbox, image_process_mode, temperature, top_p, max_output_tokens],
+        [state, chatbot, textbox, imagebox] + btn_list,
+    )
+
+    submit_btn.click(
+        add_text,
+        [state, imagebox, textbox, image_process_mode],
+        [state, chatbot, textbox, imagebox] + btn_list,
+    ).then(
+        generate,
+        [state, imagebox, textbox, image_process_mode, temperature, top_p, max_output_tokens],
+        [state, chatbot, textbox, imagebox] + btn_list,
+    )
+
+demo.queue(
+    status_update_rate=10,
+    api_open=False
+).launch()