multimodal_chat.py

import gradio as gr
from fastapi import FastAPI
import os
from PIL import Image as pil_image
import json
from conversations import Conversation, conv_templates
from sglang import RuntimeEndpoint
from datetime import datetime
import sglang as sgl
import hashlib
import argparse
import PIL

from theme_dropdown import create_theme_dropdown  # noqa: F401

dropdown, js = create_theme_dropdown()

title_markdown = """
# [LLaVA-NeXT: Stronger LLMs Supercharge Multimodal Capabilities in the Wild](https://llava-vl.github.io/blog/2024-05-10-llava-next-stronger-llms/)
"""

sub_title_markdown = """
> This is a simple demo webpage showcasing the versatile capabilities of our model. We are operating on 8 * A100-40GB GPUs with SGLang. Please note that during the serving process, there may be congestion leading to delayed responses. If you encounter any issues with the webpage, kindly refresh it.
"""

block_css = """
#buttons button {
    min-width: min(120px,100%);
}
"""

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 model [License](https://github.com/facebookresearch/llama/blob/main/MODEL_CARD.md) of LLaMA, [Terms of Use](https://openai.com/policies/terms-of-use) of the data generated by OpenAI, and [Privacy Practices](https://chrome.google.com/webstore/detail/sharegpt-share-your-chatg/daiacboceoaocpibfodeljbdfacokfjb) of ShareGPT. Please contact us if you find any potential violation.
"""

bibtext = """
## Citation
```
@misc{li2024llavanext-strong,
    title={LLaVA-NeXT: Stronger LLMs Supercharge Multimodal Capabilities in the Wild},
    url={https://llava-vl.github.io/blog/2024-05-10-llava-next-stronger-llms/},
    author={Li, Bo and Zhang, Kaichen and Zhang, Hao and Guo, Dong and Zhang, Renrui and Li, Feng and Zhang, Yuanhan and Liu, Ziwei and Li, Chunyuan},
    month={May},
    year={2024}
}
```
"""

################## BACKEND ##################
os.environ["GRADIO_EXAMPLES_CACHE"] = (
    "/mnt/bn/vl-research/workspace/boli01/projects/demos/cache"
)
os.environ["GRADIO_TEMP_DIR"] = (
    "/mnt/bn/vl-research/workspace/boli01/projects/demos/cache"
)
multimodal_folder_path = (
    "/mnt/bn/vl-research/workspace/boli01/projects/demos/cache/user_logs/medias"
)

if not os.path.exists(multimodal_folder_path):
    os.makedirs(multimodal_folder_path)


def generate_file_hash(file_path):
    sha256_hash = hashlib.sha256()

    with open(file_path, "rb") as f:
        # Read and update hash in chunks of 4K
        for byte_block in iter(lambda: f.read(4096), b""):
            sha256_hash.update(byte_block)

    return sha256_hash.hexdigest()[:6]


@sgl.function
def image_qa(s, image_path, question):
    s += sgl.user(sgl.image(image_path) + question)
    s += sgl.assistant(sgl.gen("answer"))


def single(path, prompt, temperature, max_new_tokens):
    state = image_qa.run(
        image_path=path, question=prompt, max_new_tokens=max_new_tokens, temperature=temperature
    )
    print(state["answer"], "\n")
    return state["answer"]


def print_like_dislike(x: gr.LikeData):
    print(x.index, x.value, x.liked)


def add_message(history, message):
    for x in message["files"]:
        history.append(((x,), None))
    if message["text"] is not None:
        history.append((message["text"], None))
    return history, gr.MultimodalTextbox(value=None, interactive=False)


def process_image_and_prompt(image_path, prompt, temperature=0, max_new_tokens=8192):
    start_time = datetime.now()
    formated_time = start_time.strftime("%Y-%m-%d-%H-%M-%S")

    if not os.path.exists(image_path):
        return "Image is not correctly uploaded and processed. Please try again."

    print(f"Processing Image: {image_path}")
    try:
        if (
            ".png" in image_path.lower()
            or ".jpg" in image_path.lower()
            or ".jpeg" in image_path.lower()
            or ".webp" in image_path.lower()
            or ".bmp" in image_path.lower()
            or ".gif" in image_path.lower()
        ):
            response = single(image_path, prompt, temperature, max_new_tokens)
        else:
            response = (
                "Image format is not supported. Please upload a valid image file."
            )
    except Exception as e:
        print(e)
        return "Image is not correctly uploaded and processed. Please try again."

    hashed_value = generate_file_hash(image_path)
    collected_json_path = os.path.join(
        multimodal_folder_path, f"{formated_time}_{hashed_value}.json"
    )
    collected_user_logs = {}
    collected_user_logs["image_path"] = image_path
    collected_user_logs["user_questions"] = prompt
    collected_user_logs["model_response"] = response

    with open(collected_json_path, "w") as f:
        f.write(json.dumps(collected_user_logs))

    print(f"################# {collected_json_path} #################")
    print(f"Image Path: {image_path}")
    print(f"User Question: {prompt}")
    print(f"Response: {response}")
    print(f"######################### END ############################")
    return response


def bot(history, temperature=0.2, max_new_tokens=8192):
    try:
        if len(history) > 2:
            history = history[-2:]
        response = process_image_and_prompt(history[-2][0][0], history[-1][0], temperature, max_new_tokens)
    except Exception as e:
        print(e)
        response = "Image is not correctly uploaded and processed. Please try again."

    try:
        history[-1][1] = ""
        for character in response:
            history[-1][1] += character
            yield history
    except Exception as e:
        print(e)
        yield history


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--sglang_port", default="10001", help="SGlang port")
    parser.add_argument("--model_name", default="LLaVA-NeXT-110B", help="Model name")
    parser.add_argument("--temperature", default="0", help="Temperature")
    parser.add_argument("--max_new_tokens", default="8192", help="Max new tokens")
    args = parser.parse_args()
    with gr.Blocks(
        theme="finlaymacklon/smooth_slate",
        title="LLaVA-NeXT: Multimodal Chat",
        css=".message-wrap.svelte-1lcyrx4>div.svelte-1lcyrx4  img {min-width: 50px}") as demo:
        print(f"SGlang port: {args.sglang_port}")
        runtime = RuntimeEndpoint(f"http://localhost:{args.sglang_port}")
        sgl.set_default_backend(runtime)
        gr.Markdown(title_markdown)

        # model_selector = gr.Dropdown(
        #     choices=models,
        #     value=models[0] if len(models) > 0 else "",
        #     interactive=True,
        #     show_label=False,
        #     container=False)

        chatbot = gr.Chatbot(
            [],
            label=f"Model: {args.model_name}",
            elem_id="chatbot",
            bubble_full_width=False,
            height=600,
            avatar_images=(
                (os.path.join(os.path.dirname(__file__), "./assets/user_logo.png")),
                (
                    os.path.join(
                        os.path.dirname(__file__), "./assets/assistant_logo.png"
                    )
                ),
            ),
        )

        chat_input = gr.MultimodalTextbox(
            interactive=True,
            file_types=["image"],
            placeholder="Enter message or upload file...",
            show_label=False,
            max_lines=10000,
        )

        chat_msg = chat_input.submit(
            add_message, [chatbot, chat_input], [chatbot, chat_input]
        )
        bot_msg = chat_msg.then(bot, inputs=[chatbot], outputs=[chatbot], api_name="bot_response")
        bot_msg.then(lambda: gr.MultimodalTextbox(interactive=True), None, [chat_input])

        chatbot.like(print_like_dislike, None, None)

        with gr.Row():
            gr.ClearButton(chatbot, chat_input, chat_msg, bot_msg)
            submit_btn = gr.Button("Send", chat_msg)
            submit_btn.click(
                add_message, [chatbot, chat_input], [chatbot, chat_input]
            ).then(bot, chatbot, chatbot, api_name="bot_response").then(
                lambda: gr.MultimodalTextbox(interactive=True), None, [chat_input]
            )

        gr.Examples(
            examples=[
                {
                    "files": [
                        "./image_examples/user_example_01.jpg",
                    ],
                    "text": "Explain this diagram.",
                },
                {
                    "files": [
                        "./image_examples/user_example_03.jpg",
                    ],
                    "text": "What characters are used in this captcha? Write them sequentially",
                },
                {
                    "files": [
                        "./image_examples/user_example_04.jpg",
                    ],
                    "text": "What is the latex code for this image?",
                },
                {
                    "files": [
                        "./image_examples/user_example_06.jpg",
                    ],
                    "text": "Write the content of this table in a Notion format?",
                },
                {
                    "files": [
                        "./image_examples/user_example_07.jpg",
                    ],
                    "text": "这个是什么猫？它在干啥？",
                },
                {
                    "files": [
                        "./image_examples/user_example_05.jpg",
                    ],
                    "text": "この猫の目の大きさは、どのような理由で他の猫と比べて特に大きく見えますか？",
                },
                {
                    "files": [
                        "./image_examples/user_example_09.jpg",
                    ],
                    "text": "请针对于这幅画写一首中文古诗。",
                },
            ],
            inputs=[chat_input],
        )
        gr.Markdown(sub_title_markdown)
        gr.Markdown(bibtext)
        gr.Markdown(tos_markdown)
        gr.Markdown(learn_more_markdown)

    demo.queue(max_size=128)
    demo.launch(max_threads=8)