demo

app.py

@@ -1,360 +1,39 @@
-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
+# import spaces
+
 # from llava import conversation as conversation_lib
-# from llava.constants import *
+# 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 diffusers import DiffusionPipeline
 
-# # Define paths and configurations
-# # diffusion_path = "/export/jchen169/hub/models--BAAI--Emu2-Gen/snapshots/a41a2dcd777a68225dddc72c7213b064ee06f4a0"
+# 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-3v", type=str)
+# 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)
@@ -362,48 +41,45 @@ demo.queue(
 # 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()
+# 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)
-# tokenizer, model, image_processor, context_len = load_pretrained_model(args.model_path, None, model_name)
+# 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
 
-# # 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()
+#     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)
 
@@ -412,9 +88,27 @@ demo.queue(
 
 #     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
 
@@ -423,8 +117,11 @@ demo.queue(
 #             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:
@@ -444,9 +141,11 @@ demo.queue(
 
 #     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(
@@ -470,25 +169,51 @@ demo.queue(
 #         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]
+# 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.
+# """)
 
-# def generate_image(prompt):
-#     prompt = add_template(prompt)
-#     gen_img = pipe(prompt, guidance_scale=3.0)
-#     return gen_img.image
+# block_css = """
+# #buttons button {
+#     min-width: min(120px,100%);
+# }
+# """
 
-# # Interface setup
-# with gr.Blocks(title="LLaVA Chatbot with Image Generation") as demo:
+# 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("# LLaVA Chatbot with Image Generation")
+
+#     gr.Markdown(title_markdown)
 
 #     with gr.Row():
 #         with gr.Column(scale=3):
@@ -497,24 +222,299 @@ demo.queue(
 #                 ["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")
+    
+#             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)
 
-#             with gr.Row() as button_row:
+
+            
+            
+#             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)
-    
-#     # Define actions
+
+#     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(
-#         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)),
+#         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]
+#         [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()
+    demo.queue(status_update_rate=10, api_open=False).launch()