image_app.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

title_markdown = """
# 🌋 LLaVA-NeXT: Stronger LLMs Supercharge Multimodal Capabilities in the Wild
"""

sub_title_markdown = """
📚 [[LLaVA-NeXT-2024/01](https://llava-vl.github.io/blog/2024-01-30-llava-1-6/) | [[LLaVA-NeXT-Video](https://llava-vl.github.io/blog/2024-01-30-llava-next/)] | [[LLaVA-NeXT-2024/05]
"""

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{zhang2024llavanextvideo,
  title={LLaVA-NeXT: A Strong Zero-shot Video Understanding Model},
  url={https://llava-vl.github.io/blog/2024-04-30-llava-next-video/},
  author={Zhang, Yuanhan and Li, Bo and Liu, haotian and Lee, Yong jae and Gui, Liangke and Fu, Di and Feng, Jiashi and Liu, Ziwei and Li, Chunyuan},
  month={April},
  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):
    state = image_qa.run(
        image_path=path, question=prompt, max_new_tokens=8192
    )
    print(state["answer"], "\n")
    return state["answer"]

def process_image_and_prompt(image_path, prompt):
    start_time = datetime.now()
    formated_time = start_time.strftime("%Y-%m-%d-%H-%M-%S")
    
    if not os.path.exists(image_path):
        return "Video is not correctly uploaded and processed. Please try again."

    print(f"Processing video: {image_path}")
    if ".png" in image_path or ".jpg" in image_path or ".jpeg" in image_path or ".webp" in image_path or ".bmp" in image_path or ".gif" in image_path:
        response = single(image_path, prompt)
    else:
        return "Video 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"Video Path: {image_path}")
    print(f"User Question: {prompt}")
    print(f"Response: {response}")
    return response

def save_image_to_local(image):
    start_time = datetime.now()
    formated_time = start_time.strftime("%Y-%m-%d-%H-%M-%S")
    image_path = os.path.join(multimodal_folder_path, f"{formated_time}.png")
    image.save(image_path)
    return image_path

def generate(image, textbox_in, first_run, state, state_, images_tensor):
    print("Generating")
    print(image)
    flag = 1
    if not textbox_in:
        if len(state_.messages) > 0:
            textbox_in = state_.messages[-1][1]
            state_.messages.pop(-1)
            flag = 0
        else:
            return "Please enter instruction"

    image = image if image else None
    if type(image) is pil_image.Image:
        image_path = save_image_to_local(image)
    # assert not (os.path.exists(image1) and os.path.exists(video))

    if type(state) is not Conversation:
        state = conv_templates["vicuna_v1"].copy()
        state_ = conv_templates["vicuna_v1"].copy()
        images_tensor = [[], []]
        
    first_run = False if len(state.messages) > 0 else True
    state.append_message(state.roles[0], textbox_in)

    text_en_out = process_image_and_prompt(image_path, textbox_in)
    state_.messages.append((state_.roles[0], textbox_in))
    state_.messages.append((state_.roles[1], text_en_out))

    # text_en_out = text_en_out.split("#")[0]
    textbox_out = text_en_out

    # show_images = ""
    # if os.path.exists(image1):
    #     filename = save_image_to_local(image1)
    #     show_images += f'<img src="./file={filename}" style="display: inline-block;width: 250px;max-height: 400px;">'
    # if os.path.exists(video):
    #     filename = save_video_to_local(video)
    #     show_images += f'<video controls playsinline width="500" style="display: inline-block;"  src="./file={filename}"></video>'

    # if flag:
    #     state.append_message(state.roles[0], textbox_in + "\n" + show_images)
    state.append_message(state.roles[1], textbox_out)
    return (
        state,
        state_,
        state.to_gradio_chatbot(),
        False,
        gr.update(value=None, interactive=True),
        images_tensor,
        gr.update(value=image, interactive=True),
    )


def regenerate(state, state_):
    state.messages.pop(-1)
    state_.messages.pop(-1)
    if len(state.messages) > 0:
        return state, state_, state.to_gradio_chatbot(), False
    return (state, state_, state.to_gradio_chatbot(), True)


def clear_history(state, state_):
    state = conv_templates["vicuna_v1"].copy()
    state_ = conv_templates["vicuna_v1"].copy()
    return (
        gr.update(value=None, interactive=True),
        gr.update(value=None, interactive=True),
        True,
        state,
        state_,
        state.to_gradio_chatbot(),
        [[], []],
    )

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--sglang_port", default="30000", help="SGlang port")
    args = parser.parse_args()

    models = ["LLaVA-NeXT-72B", "LLaVA-NeXT-110B"]
    runtime = RuntimeEndpoint(f"http://localhost:{args.sglang_port}")
    sgl.set_default_backend(runtime)
    textbox = gr.Textbox(
        show_label=False, placeholder="Enter text and press ENTER", container=False, max_lines=100
    )
    with gr.Blocks(
        title="LLaVA-NeXT", theme=gr.themes.Default(), css=block_css, fill_height=True
    ) as demo:
        gr.Markdown(title_markdown)
        gr.Markdown(sub_title_markdown)
        
        state = gr.State()
        state_ = gr.State()
        first_run = gr.State()
        images_tensor = gr.State()

        with gr.Row(equal_height=True):
            with gr.Column(scale=3):
                with gr.Row(elem_id="model_selector_row"):
                    model_selector = gr.Dropdown(
                        choices=models,
                        value=models[0] if len(models) > 0 else "",
                        interactive=True,
                        show_label=False,
                        container=False)
                    
                with gr.Row():
                    image = gr.Image(label="Input Image", type="pil", height=500)
                    cur_dir = os.path.dirname(os.path.abspath(__file__))
                with gr.Row():
                    gr.Examples(
                        examples=[
                            [
                                f"{cur_dir}/image_examples/llava_next.jpg",
                                "What's the image about?",
                            ],
                        ],
                        inputs=[image, textbox],
                    )

            with gr.Column(scale=7):
                chatbot = gr.Chatbot(label="LLaVA-NeXT", bubble_full_width=False, height=800)
                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", interactive=True
                        )
                with gr.Row(elem_id="buttons") as button_row:
                    upvote_btn = gr.Button(value="👍  Upvote", interactive=True)
                    downvote_btn = gr.Button(value="👎  Downvote", interactive=True)
                    flag_btn = gr.Button(value="⚠️  Flag", interactive=True)
                    clear_btn = gr.Button(value="🗑️  Clear history", interactive=True)

        
        gr.Markdown(tos_markdown)
        gr.Markdown(learn_more_markdown)
        gr.Markdown(bibtext)

        submit_btn.click(
            generate,
            [image, textbox, first_run, state, state_, images_tensor],
            [state, state_, chatbot, first_run, textbox, images_tensor, image],
        )

        clear_btn.click(
            clear_history,
            [state, state_],
            [image, textbox, first_run, state, state_, chatbot, images_tensor],
        )

    _, url, _ = demo.launch(server_port=7860)