app.py

# import gradio as gr
# import os
# import torch
# import spaces

# from llava import conversation as conversation_lib
# from llava.constants import IMAGE_TOKEN_IDX, DEFAULT_IMAGE_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN
# from llava.conversation import conv_templates, SeparatorStyle
# from llava.model.builder import load_pretrained_model
# from llava.utils import disable_torch_init
# from llava.mm_utils import tokenizer_image_token, get_model_name_from_path, process_images

# from PIL import Image
# import argparse


# from transformers import TextIteratorStreamer
# from threading import Thread

# import subprocess
# # Install flash attention, skipping CUDA build if necessary
# subprocess.run(
#     "pip install flash-attn --no-build-isolation",
#     env={"FLASH_ATTENTION_SKIP_CUDA_BUILD": "TRUE"},
#     shell=True,
# )

# # os.environ['GRADIO_TEMP_DIR'] = './gradio_tmp'
# no_change_btn = gr.Button()
# enable_btn = gr.Button(interactive=True)
# disable_btn = gr.Button(interactive=False)

# argparser = argparse.ArgumentParser()
# argparser.add_argument("--model-path", default="umd-vt-nyu/clip-evaclip-und-gen-sft", type=str)
# argparser.add_argument("--model-base", type=str, default=None)
# argparser.add_argument("--num-gpus", type=int, default=1)
# argparser.add_argument("--conv-mode", type=str, default="llama3")
# argparser.add_argument("--temperature", type=float, default=0.2)
# argparser.add_argument("--max-new-tokens", type=int, default=64)
# argparser.add_argument("--num_frames", type=int, default=16)
# argparser.add_argument("--load-8bit", action="store_true")
# argparser.add_argument("--load-4bit", action="store_true")
# argparser.add_argument("--debug", action="store_true")

# args = argparser.parse_args()
# model_path = args.model_path
# conv_mode = args.conv_mode
# filt_invalid="cut"
# model_name = get_model_name_from_path(args.model_path)
# model_name = 'clip-evaclip-und-gen-sft'
# model_kwargs = {
#     "use_cache": False,
#     "trust_remote_code": True,
#     "torch_dtype": torch.bfloat16,
#     "attn_implementation": "sdpa"
# }
# tokenizer, model, image_processor, context_len = load_pretrained_model(args.model_path, args.model_base, model_name, device_map="cuda:0", **model_kwargs)
# our_chatbot = None

# 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 =conv_templates[conv_mode].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)
#         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_IDX, 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()

# txt = gr.Textbox(
#     scale=4,
#     show_label=False,
#     placeholder="Enter text and press enter.",
#     container=False,
# )


# title_markdown = ("""
# # Florence-llama
# [[Code](TBD)] [[Model](TBD)] | 📚 [[Arxiv](TBD)]]
# """)

# # title_markdown = ("""
# # # Florence-llama
# # """)

# 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="llava", 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__))
#             # gr.Examples(examples=[
#             #     [f"{cur_dir}/assets/health-insurance.png", "Under which circumstances do I need to be enrolled in mandatory health insurance if I am an international student?"],
#             #     [f"{cur_dir}/assets/leasing-apartment.png", "I don't have any 3rd party renter's insurance now. Do I need to get one for myself?"],
#             #     [f"{cur_dir}/assets/nvidia.jpeg", "Who is the person in the middle?"],
#             #     [f"{cur_dir}/assets/animal-compare.png", "Are these two pictures showing the same kind of animal?"],
#             #     [f"{cur_dir}/assets/georgia-tech.jpeg", "Where is this photo taken?"]
#             # ], inputs=[imagebox, textbox], cache_examples=False)


            
            
#             gr.Examples(examples=[
#                 [f"{cur_dir}/assets/animal-compare.png", "Provide a detailed description of the given image."]
#             ], 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="llava 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()









import gradio as gr
import os
import torch
import argparse
from transformers import TextIteratorStreamer
from threading import Thread
from PIL import Image
from llava import conversation as conversation_lib
from llava.constants import *
from llava.conversation import conv_templates, SeparatorStyle
from llava.model.builder import load_pretrained_model
from llava.utils import disable_torch_init
from llava.mm_utils import tokenizer_image_token, get_model_name_from_path, process_images
from diffusers import DiffusionPipeline

# Define paths and configurations
# diffusion_path = "/export/jchen169/hub/models--BAAI--Emu2-Gen/snapshots/a41a2dcd777a68225dddc72c7213b064ee06f4a0"

argparser = argparse.ArgumentParser()
argparser.add_argument("--model-path", default="umd-vt-nyu/clip-evaclip-und-gen-sft-3v", type=str)
argparser.add_argument("--conv-mode", type=str, default="llama3")
argparser.add_argument("--temperature", type=float, default=0.2)
argparser.add_argument("--max-new-tokens", type=int, default=64)
argparser.add_argument("--num_frames", type=int, default=16)
argparser.add_argument("--load-8bit", action="store_true")
argparser.add_argument("--load-4bit", action="store_true")
argparser.add_argument("--debug", action="store_true")
args = argparser.parse_args()

# Load LLaVA model
disable_torch_init()
model_name = get_model_name_from_path(args.model_path)
tokenizer, model, image_processor, context_len = load_pretrained_model(args.model_path, None, model_name)
our_chatbot = None

# Load Diffusion model for image generation
pipe = DiffusionPipeline.from_pretrained(
    'BAAI/Emu2-Gen',
    custom_pipeline="pipeline_llava_gen",
    torch_dtype=torch.bfloat16,
    use_safetensors=True,
    variant="bf16",
    multimodal_encoder=model,
    tokenizer=tokenizer,
)
pipe.vae.to("cuda:0")
pipe.unet.to("cuda:0")
pipe.safety_checker.to("cuda:0")

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 = conv_templates[conv_mode].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 generate(state, imagebox, textbox, image_process_mode, temperature, top_p, max_output_tokens):
    prompt = state.get_prompt()
    images = state.get_images(return_pil=True)
    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")
            
            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)
        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_IDX, 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:
        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()

def add_template(prompt):
   conv = conv_templates['llama3'].copy()
   conv.append_message(conv.roles[0], prompt[0])
   conv.append_message(conv.roles[1], None)
   prompt = conv.get_prompt()
   return [prompt]


def generate_image(prompt):
    prompt = add_template(prompt)
    gen_img = pipe(prompt, guidance_scale=3.0)
    return gen_img.image

# Interface setup
with gr.Blocks(title="LLaVA Chatbot with Image Generation") as demo:
    state = gr.State()
    gr.Markdown("# LLaVA Chatbot with Image Generation")

    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)
            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(label="LLaVA Chatbot", height=650, layout="panel")
            textbox = gr.Textbox(show_label=False, placeholder="Enter text and press ENTER", container=False)
            submit_btn = gr.Button(value="Send", variant="primary")

            with gr.Row() as button_row:
                clear_btn = gr.Button(value="🗑️  Clear", interactive=False)
    
    # Define actions
    submit_btn.click(
        lambda state, imagebox, textbox, image_process_mode, temperature, top_p, max_output_tokens: (
            generate_image([textbox]) if "generate image" in textbox.lower() else add_text(
                state, imagebox, textbox, image_process_mode)),
        [state, imagebox, textbox, image_process_mode, temperature, top_p, max_output_tokens],
        [state, chatbot, textbox, imagebox]
    )
    
    demo.queue(status_update_rate=10, api_open=False).launch()