:zap: Fix version

app.py

@@ -1,7 +1,6 @@
 import gradio as gr
 import torch
 import numpy as np
-import spaces
 from ola_vlm.constants import DEFAULT_IMAGE_TOKEN
 
 from ola_vlm.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN
@@ -9,8 +8,7 @@ from ola_vlm.conversation import conv_templates, SeparatorStyle
 from ola_vlm.model.builder import load_pretrained_model
 from ola_vlm.mm_utils import tokenizer_image_token, get_model_name_from_path, process_images
 
-from diffusers import StableUnCLIPImg2ImgPipeline
-from diffusers import DPMSolverMultistepScheduler
+from diffusers import StableUnCLIPImg2ImgPipeline, DPMSolverMultistepScheduler
 from transformers import OneFormerProcessor
 from ola_vlm.model.aux_heads.oneformer_head import OneFormerHead
 from ola_vlm.ola_utils import visualize_oneformer_masks_on_image, oneformer_prepare_panoptic_instance_prediction
@@ -150,10 +148,9 @@ our_chatbot = None
 
 pipe = StableUnCLIPImg2ImgPipeline.from_pretrained(f"stabilityai/stable-diffusion-2-1-unclip", torch_dtype=torch.float16, variant="fp16")
 pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
-pipe = pipe.to("cuda")
 
 oneformer_processor = OneFormerProcessor.from_pretrained("shi-labs/oneformer_coco_swin_large")
-oneformer = OneFormerHead.from_pretrained("shi-labs/oneformer_coco_swin_large").to("cuda")
+oneformer = OneFormerHead.from_pretrained("shi-labs/oneformer_coco_swin_large")
 
 gen_layer_indices = model.config.image_gen["img_layer_indices"].split("-")
 seg_layer_indices = model.config.image_seg["seg_layer_indices"].split("-")
@@ -181,6 +178,7 @@ def add_text(state, imagebox, textbox, image_process_mode):
     yield (state, state.to_gradio_chatbot(), "", None) + (disable_btn, disable_btn, disable_btn, enable_btn, enable_btn)
 
 def get_gen_images(out):
+    pipe = pipe.to("cuda")
     img_embeds = out.image_embs
     if len(img_embeds) == 0:
         return None
@@ -213,6 +211,7 @@ def get_depth_images(out, org_size):
     return grid_image
 
 def get_seg_images(out, image):
+    oneformer = oneformer.to("cuda")
     seg_embs = out.seg_embs
     
     if len(seg_embs) == 0:
@@ -252,7 +251,7 @@ def regenerate(state, image_process_mode):
 # @spaces.GPU
 # def get_interm_outs(state):
     
-
+import spaces
 @spaces.GPU
 def generate(state, temperature, top_p, max_output_tokens, is_inter=False):
     if is_inter:

demo.py

@@ -1,486 +0,0 @@
-import gradio as gr
-import os
-import torch
-import numpy as np
-
-from ola_vlm.constants import DEFAULT_IMAGE_TOKEN
-
-from ola_vlm.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN
-from ola_vlm.conversation import conv_templates, SeparatorStyle
-from ola_vlm.model.builder import load_pretrained_model
-from ola_vlm.mm_utils import tokenizer_image_token, get_model_name_from_path, process_images
-
-from diffusers import StableUnCLIPImg2ImgPipeline
-from diffusers import DPMSolverMultistepScheduler
-from transformers import OneFormerProcessor
-from ola_vlm.model.aux_heads.oneformer_head import OneFormerHead
-from ola_vlm.ola_utils import visualize_oneformer_masks_on_image, oneformer_prepare_panoptic_instance_prediction
-import matplotlib
-from PIL import Image, ImageDraw, ImageFont
-import argparse
-import math
-
-from transformers import TextIteratorStreamer
-from threading import Thread
-
-def make_grid(pil_images, layer_indices=None):
-    new_images = []
-    new_captions = []
-    
-    # Resize images and prepare captions
-    for i, pil_image in enumerate(pil_images):
-        pil_image = pil_image.resize((256, 256))
-        new_images.append(pil_image)
-        if layer_indices is not None:
-            new_captions.append(f"Layer: {layer_indices[i]}")
-        else:
-            new_captions.append(f"Layer: {i+1}")
-    
-    images = new_images
-    captions = new_captions
-
-    width, height = images[0].size
-    font_size = 18
-
-    # Calculate the number of rows and columns for the grid
-    images_per_row = min(len(images), 4)  # Max 4 images per row
-    row_count = math.ceil(len(images) / images_per_row)
-    total_width = width * images_per_row
-    total_height = height * row_count
-
-    # Create a new blank image
-    new_image = Image.new("RGB", (total_width, total_height), "white")
-    draw = ImageDraw.Draw(new_image)
-
-    # Load a default font
-    try:
-        font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", font_size)
-    except:
-        font = ImageFont.load_default()
-
-    # Place images and captions in the grid
-    for i, (image, caption) in enumerate(zip(images, captions)):
-        row = i // images_per_row
-        col = i % images_per_row
-        x_offset = col * width
-        y_offset = row * height
-        
-        # Paste the image
-        new_image.paste(image, (x_offset, y_offset))
-        
-        # Calculate text and background positions
-        text_width, text_height = draw.textsize(caption, font=font)
-        text_position = (x_offset + 10, y_offset + height - text_height - 10)
-        background_position = (
-            text_position[0] - 5,
-            text_position[1] - 5,
-            text_position[0] + text_width + 5,
-            text_position[1] + text_height + 5,
-        )
-
-        # Draw background rectangle and text
-        draw.rectangle(background_position, fill="white", outline="black")
-        draw.text(text_position, caption, fill="black", font=font)
-    
-    return new_image
-
-def reload_from_ckpt(model_path, model, cache_dir=None):
-    import os
-    from safetensors import safe_open
-    from huggingface_hub import hf_hub_download, list_repo_files
-
-    state_dict = {}
-
-    # Check if the path is a local directory or HF Hub model
-    if os.path.isdir(model_path):
-        # Local directory: Load safetensors files
-        safetensors_paths = [os.path.join(model_path, f) for f in os.listdir(model_path) if f.endswith('.safetensors')]
-    else:
-        # HF Hub: Get list of safetensors files and download them
-        repo_files = list_repo_files(model_path)
-        safetensors_paths = [
-            hf_hub_download(model_path, file_name, cache_dir=cache_dir)
-            for file_name in repo_files if file_name.endswith('.safetensors')
-        ]
-
-    # Load safetensors files into the state_dict
-    for path in safetensors_paths:
-        with safe_open(path, framework="pt", device="cpu") as f:
-            for key in f.keys():
-                state_dict[key] = f.get_tensor(key)
-
-    # Load the state dict into the model
-    model.load_state_dict(state_dict, strict=False)
-    return model
-
-# 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("--server_name", default="0.0.0.0", type=str)
-argparser.add_argument("--port", default="6324", type=str)
-argparser.add_argument("--model-path", default="shi-labs/pretrain_dsg_OLA-VLM-CLIP-ViT-Llama3-8b", 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="llava_llama_3")
-argparser.add_argument("--temperature", type=float, default=0.2)
-argparser.add_argument("--max-new-tokens", type=int, default=512)
-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)
-tokenizer, model, image_processor, context_len = load_pretrained_model(args.model_path, args.model_base, model_name, args.load_8bit, args.load_4bit)
-model = reload_from_ckpt("shi-labs/OLA-VLM-CLIP-ViT-Llama3-8b", model)
-our_chatbot = None
-
-pipe = StableUnCLIPImg2ImgPipeline.from_pretrained(f"stabilityai/stable-diffusion-2-1-unclip", torch_dtype=torch.float16, variant="fp16")
-pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
-pipe = pipe.to("cuda")
-
-oneformer_processor = OneFormerProcessor.from_pretrained("shi-labs/oneformer_coco_swin_large")
-oneformer = OneFormerHead.from_pretrained("shi-labs/oneformer_coco_swin_large").to("cuda")
-
-gen_layer_indices = model.config.image_gen["img_layer_indices"].split("-")
-seg_layer_indices = model.config.image_seg["seg_layer_indices"].split("-")
-depth_layer_indices = model.config.image_depth["depth_layer_indices"].split("-")
-
-
-def clear_history():
-    state =conv_templates[conv_mode].copy()
-    return (state, state.to_gradio_chatbot(), "", None, None, None, 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 get_gen_images(out):
-    img_embeds = out.image_embs
-    if len(img_embeds) == 0:
-        return None
-    images = []
-    for img_embed in img_embeds:
-        gen_image = pipe(image_embeds=img_embed.squeeze(1),
-                num_inference_steps=25,
-            ).images[0]
-        images.append(gen_image)
-    grid_image = make_grid(images, gen_layer_indices)
-    return grid_image
-
-def get_depth_images(out, org_size):
-    depth_preds = out.depth_preds
-
-    if len(depth_preds) == 0:
-        return None
-    depths = []
-
-    for i, depth_pred in enumerate(depth_preds):
-        depth = (depth_pred - depth_pred.min()) / (depth_pred.max() - depth_pred.min()) * 255.0
-        depth = depth.squeeze(0).cpu().numpy()
-        depth = depth.astype(np.uint8)
-        cmap = matplotlib.colormaps.get_cmap('Spectral_r')
-        depth = (cmap(depth)[:, :, :3] * 255).astype(np.uint8)       
-        depth = Image.fromarray(depth)
-        depth = depth.resize(org_size)
-        depths.append(depth)
-    grid_image = make_grid(depths, depth_layer_indices)
-    return grid_image
-
-def get_seg_images(out, image):
-    seg_embs = out.seg_embs
-    
-    if len(seg_embs) == 0:
-        return None
-    
-    seg_preds = []
-    inputs = oneformer_processor(image, ["semantic"], return_tensors="pt")
-    inputs["pixel_values"] = inputs["pixel_values"].to(out.logits.device, out.logits.dtype)
-    inputs["task_inputs"] = inputs["task_inputs"].to(out.logits.device, out.logits.dtype)
-    backbone_features = oneformer.get_backbone_feats(**inputs)
-    for i, seg_emb in enumerate(seg_embs):
-        pred = oneformer.get_masks(**inputs, backbone_last_feature=seg_emb.float(), all_backbone_features=backbone_features)
-        pred = oneformer_processor.post_process_panoptic_segmentation(
-                                pred, target_sizes=[image.size[::-1]]
-                            )[0]
-        pred_msk, pred_cls = oneformer_prepare_panoptic_instance_prediction(**pred, oneformer=oneformer)
-        pred = visualize_oneformer_masks_on_image(image, pred_msk, pred_cls)
-        seg_preds.append(pred)
-    grid_image = make_grid(seg_preds, seg_layer_indices)
-    return grid_image
-
-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
-
-def get_interm_outs(state):
-    prompt = state.get_prompt()
-    images = state.get_images(return_pil=True)
-    #prompt, image_args = process_image(prompt, images)
-
-    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]
-            inp_images = process_images(images, image_processor, model.config)
-
-            if type(inp_images) is list:
-                inp_images = [image.to(model.device, dtype=torch.float16) for image in images]
-            else:
-                inp_images = inp_images.to(model.device, dtype=torch.float16)
-        else:
-            inp_images = None
-            image_sizes = None
-        image_args = {"images": inp_images, "image_sizes": image_sizes}
-    else:
-        inp_images = None
-        image_args = {}
-
-    input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).to(model.device)
-
-    interm_outs = model.get_visual_interpretations(
-                input_ids,
-                **image_args
-         )
-    
-    depth_outs = get_depth_images(interm_outs, image_sizes[0]) 
-    seg_outs =  get_seg_images(interm_outs, images[0])
-    gen_outs = get_gen_images(interm_outs)
-
-    return depth_outs, seg_outs, gen_outs
-
-# @spaces.GPU
-def generate(state, 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 = max_output_tokens
-    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:
-        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 = "<h1 style='margin-bottom: -10px; text-align: center'>OLA-VLM: Optimizing Language Model Representations for Enhanced Visual Quality and Alignment</h1>"
-description = "<p style='font-size: 16px; margin: 5px; font-weight: w300; text-align: center'> <a href='https://praeclarumjj3.github.io/' style='text-decoration:none' target='_blank'>Jitesh Jain</a> &nbsp;&nbsp <a href='https://zyang-ur.github.io/' style='text-decoration:none' target='_blank'>Zhengyuan Yang</a> &nbsp;&nbsp <a href='https://www.humphreyshi.com/home' style='text-decoration:none' target='_blank'>Humphrey Shi<sup>*</sup></a> &nbsp;&nbsp <a href='https://www.humphreyshi.com/home' style='text-decoration:none' target='_blank'>Jianfeng Gao<sup>*</sup></a> &nbsp;&nbsp <a href='https://jwyang.github.io/' style='text-decoration:none' target='_blank'>Jianwei Yang<sup>*</sup></a></p>" \
-            + "<p style='font-size: 12px; margin: 5px; font-weight: w300; text-align: center'><sup>*</sup>Equal Advising</p>" \
-            + "<p style='font-size: 16px; margin: 5px; font-weight: w600; text-align: center'> <a href='https://praeclarumjj3.github.io/ola_vlm/' target='_blank'>Project Page</a> | <a href='https://youtu.be/' target='_blank'>Video</a> | <a href='https://arxiv.org/abs/' target='_blank'>ArXiv</a> | <a href='https://github.com/SHI-Labs/OLA-VLM' target='_blank'>Github</a></p>"
-
-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.
-""")
-
-
-learn_more_markdown = ("""
-### License
-The service is a research preview intended for non-commercial use only, subject to the [License](https://huggingface.co/lmsys/vicuna-7b-v1.5) of Vicuna-v1.5, [License](https://github.com/haotian-liu/LLaVA/blob/main/LICENSE) of LLaVA, [Terms of Use](https://cocodataset.org/#termsofuse) of the COCO dataset, [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.
-""")
-
-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="OLA-VLM", theme=gr.themes.Default(), css=block_css) as demo:
-    state = gr.State()
-
-    gr.Markdown(title)
-    gr.Markdown(description)
-
-    with gr.Row():
-        with gr.Column(scale=4):
-            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)
-
-            # with gr.Accordion("Parameters", open=False) as parameter_row:
-            with gr.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="OLA-VLM",
-                height=300,
-                layout="panel",
-            )
-            textbox.render()
-            with gr.Row(elem_id="buttons") as button_row:
-                upvote_btn = gr.Button(value="👍  Upvote", interactive=False, visible=False)
-                downvote_btn = gr.Button(value="👎  Downvote", interactive=False, visible=False)
-                flag_btn = gr.Button(value="⚠️  Flag", interactive=False, visible=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)
-                submit_btn = gr.Button(value="Send", variant="primary")
-
-    with gr.Accordion("Representations from selected layers of the LLM (expects only a single image input)", open=False) as interm_out:
-        inter_vis_btn = gr.Button(value="✨ Visualize")
-        with gr.Row():
-            depth_box = gr.Image(label="depth", type="pil", visible=True)
-            seg_box = gr.Image(label="seg", type="pil", visible=True)
-            gen_box = gr.Image(label="gen", type="pil", visible=True)
-    
-    gr.Examples(examples=[
-            [f"assets/cars.jpg", "Which car is in front: the blue or the brown one?"],
-            [f"assets/pb.jpg", "Where is the bulding located with respect to the man?"],
-        ], inputs=[imagebox, textbox], cache_examples=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]
-
-    inter_vis_btn.click(
-        get_interm_outs,
-        [state],
-        [depth_box, seg_box, gen_box],
-    )
-
-    clear_btn.click(
-        clear_history,
-        None,
-        [state, chatbot, textbox, imagebox, depth_box, gen_box, seg_box] + btn_list,
-        queue=False
-    )
-
-    regenerate_btn.click(
-        delete_text,
-        [state, image_process_mode],
-        [state, chatbot, textbox, imagebox] + btn_list,
-    ).then(
-        generate,
-        [state, 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, 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, temperature, top_p, max_output_tokens],
-        [state, chatbot, textbox, imagebox] + btn_list,
-    )
-
-demo.queue(
-    status_update_rate=10,
-    api_open=False
-).launch(share=True)
-demo.queue()