Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,24 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+demo_images/04133.png filter=lfs diff=lfs merge=lfs -text
+demo_images/04444.png filter=lfs diff=lfs merge=lfs -text
+demo_images/7292.JPG filter=lfs diff=lfs merge=lfs -text
+demo_images/MicrosoftTeams-image.png filter=lfs diff=lfs merge=lfs -text
+demo_images/church_183.png filter=lfs diff=lfs merge=lfs -text
+demo_images/train_2956_0001.png filter=lfs diff=lfs merge=lfs -text
+docs/geochat_supp.pdf filter=lfs diff=lfs merge=lfs -text
+geochat/serve/examples/11760.jpg filter=lfs diff=lfs merge=lfs -text
+geochat/serve/examples/11765.jpg filter=lfs diff=lfs merge=lfs -text
+images/architecture.png filter=lfs diff=lfs merge=lfs -text
+images/dataset.png filter=lfs diff=lfs merge=lfs -text
+images/examples.png filter=lfs diff=lfs merge=lfs -text
+images/grounded.jpg filter=lfs diff=lfs merge=lfs -text
+images/iden.jpg filter=lfs diff=lfs merge=lfs -text
+images/logo_geochat.png filter=lfs diff=lfs merge=lfs -text
+images/overview2.png filter=lfs diff=lfs merge=lfs -text
+images/ref1.jpg filter=lfs diff=lfs merge=lfs -text
+images/ref_2.jpg filter=lfs diff=lfs merge=lfs -text
+images/scene.jpg filter=lfs diff=lfs merge=lfs -text
+images/teaser.png filter=lfs diff=lfs merge=lfs -text
+images/vqa.jpg filter=lfs diff=lfs merge=lfs -text

.ipynb_checkpoints/geochat_demo-checkpoint.py

@@ -0,0 +1,707 @@
+import argparse
+import os
+import random
+from collections import defaultdict
+
+import cv2
+import re
+import math
+import numpy as np
+from PIL import Image
+import torch
+import html
+import gradio as gr
+
+import torchvision.transforms as T
+import torch.backends.cudnn as cudnn
+
+from geochat.conversation import conv_templates, Chat
+from geochat.model.builder import load_pretrained_model
+from geochat.mm_utils import  get_model_name_from_path
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Demo")
+    # parser = argparse.ArgumentParser()
+    parser.add_argument("--model-path", type=str, default="facebook/opt-350m")
+    parser.add_argument("--model-base", type=str, default=None)
+    parser.add_argument("--gpu-id", type=str,default=0)
+    parser.add_argument("--device", type=str, default="cuda")
+    parser.add_argument("--conv-mode", type=str, default=None)
+    parser.add_argument("--max-new-tokens", type=int, default=300)
+    parser.add_argument("--load-8bit", action="store_true")
+    parser.add_argument("--load-4bit", action="store_true")
+    parser.add_argument("--debug", action="store_true")
+    parser.add_argument("--image-aspect-ratio", type=str, default='pad')
+    # args = parser.parse_args()
+    args = parser.parse_args()
+    return args
+
+
+random.seed(42)
+np.random.seed(42)
+torch.manual_seed(42)
+
+cudnn.benchmark = False
+cudnn.deterministic = True
+
+print('Initializing Chat')
+args = parse_args()
+# cfg = Config(args)
+
+model_name = get_model_name_from_path(args.model_path)
+print(model_name)
+tokenizer, model, image_processor, context_len = load_pretrained_model(args.model_path, args.model_base, model_name, args.load_8bit, args.load_4bit, device=args.device)
+
+device = 'cuda:{}'.format(args.gpu_id)
+
+# model_config = cfg.model_cfg
+# model_config.device_8bit = args.gpu_id
+# model_cls = registry.get_model_class(model_config.arch)
+# model = model_cls.from_config(model_config).to(device)
+bounding_box_size = 100
+
+# vis_processor_cfg = cfg.datasets_cfg.cc_sbu_align.vis_processor.train
+# vis_processor = registry.get_processor_class(vis_processor_cfg.name).from_config(vis_processor_cfg)
+
+model = model.eval()
+
+CONV_VISION = conv_templates['llava_v1'].copy()
+
+def bbox_and_angle_to_polygon(x1, y1, x2, y2, a):
+    # Calculate center coordinates
+    x_ctr = (x1 + x2) / 2
+    y_ctr = (y1 + y2) / 2
+    
+    # Calculate width and height
+    w = abs(x2 - x1)
+    h = abs(y2 - y1)
+    
+    # Calculate the angle in radians
+    angle_rad = math.radians(a)
+    
+    # Calculate coordinates of the four corners of the rotated bounding box
+    cos_a = math.cos(angle_rad)
+    sin_a = math.sin(angle_rad)
+    
+    x1_rot = cos_a * (-w / 2) - sin_a * (-h / 2) + x_ctr
+    y1_rot = sin_a * (-w / 2) + cos_a * (-h / 2) + y_ctr
+    
+    x2_rot = cos_a * (w / 2) - sin_a * (-h / 2) + x_ctr
+    y2_rot = sin_a * (w / 2) + cos_a * (-h / 2) + y_ctr
+    
+    x3_rot = cos_a * (w / 2) - sin_a * (h / 2) + x_ctr
+    y3_rot = sin_a * (w / 2) + cos_a * (h / 2) + y_ctr
+    
+    x4_rot = cos_a * (-w / 2) - sin_a * (h / 2) + x_ctr
+    y4_rot = sin_a * (-w / 2) + cos_a * (h / 2) + y_ctr
+    
+    # Return the polygon coordinates
+    polygon_coords = np.array((x1_rot, y1_rot, x2_rot, y2_rot, x3_rot, y3_rot, x4_rot, y4_rot))
+    
+    return polygon_coords
+
+def rotate_bbox(top_right, bottom_left, angle_degrees):
+    # Convert angle to radians
+    angle_radians = np.radians(angle_degrees)
+
+    # Calculate the center of the rectangle
+    center = ((top_right[0] + bottom_left[0]) / 2, (top_right[1] + bottom_left[1]) / 2)
+
+    # Calculate the width and height of the rectangle
+    width = top_right[0] - bottom_left[0]
+    height = top_right[1] - bottom_left[1]
+
+    # Create a rotation matrix
+    rotation_matrix = cv2.getRotationMatrix2D(center, angle_degrees, 1)
+
+    # Create an array of the rectangle corners
+    rectangle_points = np.array([[bottom_left[0], bottom_left[1]],
+                                 [top_right[0], bottom_left[1]],
+                                 [top_right[0], top_right[1]],
+                                 [bottom_left[0], top_right[1]]], dtype=np.float32)
+
+    # Rotate the rectangle points
+    rotated_rectangle = cv2.transform(np.array([rectangle_points]), rotation_matrix)[0]
+
+    return rotated_rectangle
+def extract_substrings(string):
+    # first check if there is no-finished bracket
+    index = string.rfind('}')
+    if index != -1:
+        string = string[:index + 1]
+
+    pattern = r'<p>(.*?)\}(?!<)'
+    matches = re.findall(pattern, string)
+    substrings = [match for match in matches]
+
+    return substrings
+
+
+def is_overlapping(rect1, rect2):
+    x1, y1, x2, y2 = rect1
+    x3, y3, x4, y4 = rect2
+    return not (x2 < x3 or x1 > x4 or y2 < y3 or y1 > y4)
+
+
+def computeIoU(bbox1, bbox2):
+    x1, y1, x2, y2 = bbox1
+    x3, y3, x4, y4 = bbox2
+    intersection_x1 = max(x1, x3)
+    intersection_y1 = max(y1, y3)
+    intersection_x2 = min(x2, x4)
+    intersection_y2 = min(y2, y4)
+    intersection_area = max(0, intersection_x2 - intersection_x1 + 1) * max(0, intersection_y2 - intersection_y1 + 1)
+    bbox1_area = (x2 - x1 + 1) * (y2 - y1 + 1)
+    bbox2_area = (x4 - x3 + 1) * (y4 - y3 + 1)
+    union_area = bbox1_area + bbox2_area - intersection_area
+    iou = intersection_area / union_area
+    return iou
+
+
+def save_tmp_img(visual_img):
+    file_name = "".join([str(random.randint(0, 9)) for _ in range(5)]) + ".jpg"
+    file_path = "/tmp/gradio" + file_name
+    visual_img.save(file_path)
+    return file_path
+
+
+def mask2bbox(mask):
+    if mask is None:
+        return ''
+    mask = mask.resize([100, 100], resample=Image.NEAREST)
+    mask = np.array(mask)[:, :, 0]
+
+    rows = np.any(mask, axis=1)
+    cols = np.any(mask, axis=0)
+
+    if rows.sum():
+        # Get the top, bottom, left, and right boundaries
+        rmin, rmax = np.where(rows)[0][[0, -1]]
+        cmin, cmax = np.where(cols)[0][[0, -1]]
+        bbox = '{{<{}><{}><{}><{}>}}'.format(cmin, rmin, cmax, rmax)
+    else:
+        bbox = ''
+
+    return bbox
+
+
+def escape_markdown(text):
+    # List of Markdown special characters that need to be escaped
+    md_chars = ['<', '>']
+
+    # Escape each special character
+    for char in md_chars:
+        text = text.replace(char, '\\' + char)
+
+    return text
+
+
+def reverse_escape(text):
+    md_chars = ['\\<', '\\>']
+
+    for char in md_chars:
+        text = text.replace(char, char[1:])
+
+    return text
+
+
+colors = [
+    (255, 0, 0),
+    (0, 255, 0),
+    (0, 0, 255),
+    (210, 210, 0),
+    (255, 0, 255),
+    (0, 255, 255),
+    (114, 128, 250),
+    (0, 165, 255),
+    (0, 128, 0),
+    (144, 238, 144),
+    (238, 238, 175),
+    (255, 191, 0),
+    (0, 128, 0),
+    (226, 43, 138),
+    (255, 0, 255),
+    (0, 215, 255),
+]
+
+color_map = {
+    f"{color_id}": f"#{hex(color[2])[2:].zfill(2)}{hex(color[1])[2:].zfill(2)}{hex(color[0])[2:].zfill(2)}" for
+    color_id, color in enumerate(colors)
+}
+
+used_colors = colors
+
+
+def visualize_all_bbox_together(image, generation):
+    if image is None:
+        return None, ''
+
+    generation = html.unescape(generation)
+
+    image_width, image_height = image.size
+    image = image.resize([500, int(500 / image_width * image_height)])
+    image_width, image_height = image.size
+
+    string_list = extract_substrings(generation)
+    if string_list:  # it is grounding or detection
+        mode = 'all'
+        entities = defaultdict(list)
+        i = 0
+        j = 0
+        for string in string_list:
+            try:
+                obj, string = string.split('</p>')
+            except ValueError:
+                print('wrong string: ', string)
+                continue
+            if "}{" in string:
+                string=string.replace("}{","}<delim>{")
+            bbox_list = string.split('<delim>')
+            flag = False
+            for bbox_string in bbox_list:
+                integers = re.findall(r'-?\d+', bbox_string)
+                if len(integers)==4:
+                    angle=0
+                else:
+                    angle=integers[4]
+                integers=integers[:-1]
+                
+                if len(integers) == 4:
+                    x0, y0, x1, y1 = int(integers[0]), int(integers[1]), int(integers[2]), int(integers[3])
+                    left = x0 / bounding_box_size * image_width
+                    bottom = y0 / bounding_box_size * image_height
+                    right = x1 / bounding_box_size * image_width
+                    top = y1 / bounding_box_size * image_height
+
+                    entities[obj].append([left, bottom, right, top,angle])
+
+                    j += 1
+                    flag = True
+            if flag:
+                i += 1
+    else:
+        integers = re.findall(r'-?\d+', generation)
+        # if len(integers)==4:
+        angle=0
+        # else:
+            # angle=integers[4]
+        integers=integers[:-1]
+        if len(integers) == 4:  # it is refer
+            mode = 'single'
+
+            entities = list()
+            x0, y0, x1, y1 = int(integers[0]), int(integers[1]), int(integers[2]), int(integers[3])
+            left = x0 / bounding_box_size * image_width
+            bottom = y0 / bounding_box_size * image_height
+            right = x1 / bounding_box_size * image_width
+            top = y1 / bounding_box_size * image_height
+            entities.append([left, bottom, right, top,angle])
+        else:
+            # don't detect any valid bbox to visualize
+            return None, ''
+
+    if len(entities) == 0:
+        return None, ''
+
+    if isinstance(image, Image.Image):
+        image_h = image.height
+        image_w = image.width
+        image = np.array(image)
+
+    elif isinstance(image, str):
+        if os.path.exists(image):
+            pil_img = Image.open(image).convert("RGB")
+            image = np.array(pil_img)[:, :, [2, 1, 0]]
+            image_h = pil_img.height
+            image_w = pil_img.width
+        else:
+            raise ValueError(f"invaild image path, {image}")
+    elif isinstance(image, torch.Tensor):
+
+        image_tensor = image.cpu()
+        reverse_norm_mean = torch.tensor([0.48145466, 0.4578275, 0.40821073])[:, None, None]
+        reverse_norm_std = torch.tensor([0.26862954, 0.26130258, 0.27577711])[:, None, None]
+        image_tensor = image_tensor * reverse_norm_std + reverse_norm_mean
+        pil_img = T.ToPILImage()(image_tensor)
+        image_h = pil_img.height
+        image_w = pil_img.width
+        image = np.array(pil_img)[:, :, [2, 1, 0]]
+    else:
+        raise ValueError(f"invalid image format, {type(image)} for {image}")
+
+    indices = list(range(len(entities)))
+
+    new_image = image.copy()
+
+    previous_bboxes = []
+    # size of text
+    text_size = 0.4
+    # thickness of text
+    text_line = 1  # int(max(1 * min(image_h, image_w) / 512, 1))
+    box_line = 2
+    (c_width, text_height), _ = cv2.getTextSize("F", cv2.FONT_HERSHEY_COMPLEX, text_size, text_line)
+    base_height = int(text_height * 0.675)
+    text_offset_original = text_height - base_height
+    text_spaces = 2
+
+    # num_bboxes = sum(len(x[-1]) for x in entities)
+    used_colors = colors  # random.sample(colors, k=num_bboxes)
+
+    color_id = -1
+    for entity_idx, entity_name in enumerate(entities):
+        if mode == 'single' or mode == 'identify':
+            bboxes = entity_name
+            bboxes = [bboxes]
+        else:
+            bboxes = entities[entity_name]
+        color_id += 1
+        for bbox_id, (x1_norm, y1_norm, x2_norm, y2_norm,angle) in enumerate(bboxes):
+            skip_flag = False
+            orig_x1, orig_y1, orig_x2, orig_y2,angle = int(x1_norm), int(y1_norm), int(x2_norm), int(y2_norm), int(angle)
+
+            color = used_colors[entity_idx % len(used_colors)] # tuple(np.random.randint(0, 255, size=3).tolist())
+            top_right=(orig_x1,orig_y1)
+            bottom_left=(orig_x2,orig_y2)
+            angle=angle
+            rotated_bbox = rotate_bbox(top_right, bottom_left, angle)
+            new_image=cv2.polylines(new_image, [rotated_bbox.astype(np.int32)], isClosed=True,thickness=2, color=color)
+
+            # new_image = cv2.rectangle(new_image, (orig_x1, orig_y1), (orig_x2, orig_y2), color, box_line)
+
+            if mode == 'all':
+                l_o, r_o = box_line // 2 + box_line % 2, box_line // 2 + box_line % 2 + 1
+
+                x1 = orig_x1 - l_o
+                y1 = orig_y1 - l_o
+
+                if y1 < text_height + text_offset_original + 2 * text_spaces:
+                    y1 = orig_y1 + r_o + text_height + text_offset_original + 2 * text_spaces
+                    x1 = orig_x1 + r_o
+
+                # add text background
+                (text_width, text_height), _ = cv2.getTextSize(f"  {entity_name}", cv2.FONT_HERSHEY_COMPLEX, text_size,
+                                                               text_line)
+                text_bg_x1, text_bg_y1, text_bg_x2, text_bg_y2 = x1, y1 - (
+                            text_height + text_offset_original + 2 * text_spaces), x1 + text_width, y1
+
+                for prev_bbox in previous_bboxes:
+                    if computeIoU((text_bg_x1, text_bg_y1, text_bg_x2, text_bg_y2), prev_bbox['bbox']) > 0.95 and \
+                            prev_bbox['phrase'] == entity_name:
+                        skip_flag = True
+                        break
+                    while is_overlapping((text_bg_x1, text_bg_y1, text_bg_x2, text_bg_y2), prev_bbox['bbox']):
+                        text_bg_y1 += (text_height + text_offset_original + 2 * text_spaces)
+                        text_bg_y2 += (text_height + text_offset_original + 2 * text_spaces)
+                        y1 += (text_height + text_offset_original + 2 * text_spaces)
+
+                        if text_bg_y2 >= image_h:
+                            text_bg_y1 = max(0, image_h - (text_height + text_offset_original + 2 * text_spaces))
+                            text_bg_y2 = image_h
+                            y1 = image_h
+                            break
+                if not skip_flag:
+                    alpha = 0.5
+                    for i in range(text_bg_y1, text_bg_y2):
+                        for j in range(text_bg_x1, text_bg_x2):
+                            if i < image_h and j < image_w:
+                                if j < text_bg_x1 + 1.35 * c_width:
+                                    # original color
+                                    bg_color = color
+                                else:
+                                    # white
+                                    bg_color = [255, 255, 255]
+                                new_image[i, j] = (alpha * new_image[i, j] + (1 - alpha) * np.array(bg_color)).astype(
+                                    np.uint8)
+
+                    cv2.putText(
+                        new_image, f"  {entity_name}", (x1, y1 - text_offset_original - 1 * text_spaces),
+                        cv2.FONT_HERSHEY_COMPLEX, text_size, (0, 0, 0), text_line, cv2.LINE_AA
+                    )
+
+                    previous_bboxes.append(
+                        {'bbox': (text_bg_x1, text_bg_y1, text_bg_x2, text_bg_y2), 'phrase': entity_name})
+
+    if mode == 'all':
+        def color_iterator(colors):
+            while True:
+                for color in colors:
+                    yield color
+
+        color_gen = color_iterator(colors)
+
+        # Add colors to phrases and remove <p></p>
+        def colored_phrases(match):
+            phrase = match.group(1)
+            color = next(color_gen)
+            return f'<span style="color:rgb{color}">{phrase}</span>'
+
+        generation = re.sub(r'{<\d+><\d+><\d+><\d+>}|<delim>', '', generation)
+        generation_colored = re.sub(r'<p>(.*?)</p>', colored_phrases, generation)
+    else:
+        generation_colored = ''
+
+    pil_image = Image.fromarray(new_image)
+    return pil_image, generation_colored
+
+
+def gradio_reset(chat_state, img_list):
+    if chat_state is not None:
+        chat_state.messages = []
+    if img_list is not None:
+        img_list = []
+    return None, gr.update(value=None, interactive=True), gr.update(placeholder='Upload your image and chat',
+                                                                    interactive=True), chat_state, img_list
+
+
+def image_upload_trigger(upload_flag, replace_flag, img_list):
+    # set the upload flag to true when receive a new image.
+    # if there is an old image (and old conversation), set the replace flag to true to reset the conv later.
+    upload_flag = 1
+    if img_list:
+        replace_flag = 1
+    return upload_flag, replace_flag
+
+
+def example_trigger(text_input, image, upload_flag, replace_flag, img_list):
+    # set the upload flag to true when receive a new image.
+    # if there is an old image (and old conversation), set the replace flag to true to reset the conv later.
+    upload_flag = 1
+    if img_list or replace_flag == 1:
+        replace_flag = 1
+
+    return upload_flag, replace_flag
+
+
+def gradio_ask(user_message, chatbot, chat_state, gr_img, img_list, upload_flag, replace_flag):
+    if len(user_message) == 0:
+        text_box_show = 'Input should not be empty!'
+    else:
+        text_box_show = ''
+
+    if isinstance(gr_img, dict):
+        gr_img, mask = gr_img['image'], gr_img['mask']
+    else:
+        mask = None
+
+    if '[identify]' in user_message:
+        # check if user provide bbox in the text input
+        integers = re.findall(r'-?\d+', user_message)
+        if len(integers) != 4:  # no bbox in text
+            bbox = mask2bbox(mask)
+            user_message = user_message + bbox
+
+    if chat_state is None:
+        chat_state = CONV_VISION.copy()
+
+    if upload_flag:
+        if replace_flag:
+            chat_state = CONV_VISION.copy()  # new image, reset everything
+            replace_flag = 0
+            chatbot = []
+        img_list = []
+        llm_message = chat.upload_img(gr_img, chat_state, img_list)
+        upload_flag = 0
+
+    chat.ask(user_message, chat_state)
+
+    chatbot = chatbot + [[user_message, None]]
+
+    if '[identify]' in user_message:
+        visual_img, _ = visualize_all_bbox_together(gr_img, user_message)
+        if visual_img is not None:
+            file_path = save_tmp_img(visual_img)
+            chatbot = chatbot + [[(file_path,), None]]
+
+    return text_box_show, chatbot, chat_state, img_list, upload_flag, replace_flag
+
+
+# def gradio_answer(chatbot, chat_state, img_list, temperature):
+#     llm_message = chat.answer(conv=chat_state,
+#                               img_list=img_list,
+#                               temperature=temperature,
+#                               max_new_tokens=500,
+#                               max_length=2000)[0]
+#     chatbot[-1][1] = llm_message
+#     return chatbot, chat_state
+
+
+def gradio_stream_answer(chatbot, chat_state, img_list, temperature):
+    if len(img_list) > 0:
+        if not isinstance(img_list[0], torch.Tensor):
+            chat.encode_img(img_list)
+    streamer = chat.stream_answer(conv=chat_state,
+                                  img_list=img_list,
+                                  temperature=temperature,
+                                  max_new_tokens=500,
+                                  max_length=2000)
+    # chatbot[-1][1] = output
+    # chat_state.messages[-1][1] = '</s>'
+    
+    output = ''
+    for new_output in streamer:
+        # print(new_output)
+        output=output+new_output
+    print(output)
+    # if "{" in output:
+    #     chatbot[-1][1]="Grounding and referring expression is still under work."
+    # else:
+    output = escape_markdown(output)
+        # output += escapped
+    chatbot[-1][1] = output
+    yield chatbot, chat_state
+    chat_state.messages[-1][1] = '</s>'
+    return chatbot, chat_state
+
+
+def gradio_visualize(chatbot, gr_img):
+    if isinstance(gr_img, dict):
+        gr_img, mask = gr_img['image'], gr_img['mask']
+
+    unescaped = reverse_escape(chatbot[-1][1])
+    visual_img, generation_color = visualize_all_bbox_together(gr_img, unescaped)
+    if visual_img is not None:
+        if len(generation_color):
+            chatbot[-1][1] = generation_color
+        file_path = save_tmp_img(visual_img)
+        chatbot = chatbot + [[None, (file_path,)]]
+
+    return chatbot
+
+
+def gradio_taskselect(idx):
+    prompt_list = [
+        '',
+        'Classify the image in the following classes: ',
+        '[identify] what is this ',
+    ]
+    instruct_list = [
+        '**Hint:** Type in whatever you want',
+        '**Hint:** Type in the classes you want the model to classify in',
+        '**Hint:** Draw a bounding box on the uploaded image then send the command. Click the "clear" botton on the top right of the image before redraw',
+    ]
+    return prompt_list[idx], instruct_list[idx]
+
+
+
+
+chat = Chat(model, image_processor,tokenizer, device=device)
+
+
+title = """<h1 align="center">GeoChat Demo</h1>"""
+description = 'Welcome to Our GeoChat Chatbot Demo!'
+article = """<div style="display: flex;"><p style="display: inline-block;"><a href='https://mbzuai-oryx.github.io/GeoChat'><img src='https://img.shields.io/badge/Project-Page-Green'></a></p><p style="display: inline-block;"><a href='https://arxiv.org/abs/2311.15826'><img src='https://img.shields.io/badge/Paper-PDF-red'></a></p><p style="display: inline-block;"><a href='https://github.com/mbzuai-oryx/GeoChat/tree/main'><img src='https://img.shields.io/badge/GitHub-Repo-blue'></a></p><p style="display: inline-block;"><a href='https://youtu.be/KOKtkkKpNDk?feature=shared'><img src='https://img.shields.io/badge/YouTube-Video-red'></a></p></div>"""
+# article = """<p><a href='https://minigpt-v2.github.io'><img src='https://img.shields.io/badge/Project-Page-Green'></a></p>"""
+
+introduction = '''
+1. Identify: Draw the bounding box on the uploaded image window and CLICK **Send** to generate the bounding box. (CLICK "clear" button before re-drawing next time).
+2. No Tag: Input whatever you want and CLICK **Send** without any tagging
+
+You can also simply chat in free form!
+'''
+
+
+text_input = gr.Textbox(placeholder='Upload your image and chat', interactive=True, show_label=False, container=False,
+                        scale=12)
+with gr.Blocks() as demo:
+    gr.Markdown(title)
+    # gr.Markdown(description)
+    gr.Markdown(article)
+
+    with gr.Row():
+        with gr.Column(scale=0.5):
+            image = gr.Image(type="pil", tool='sketch', brush_radius=20)
+
+            temperature = gr.Slider(
+                minimum=0.1,
+                maximum=1.5,
+                value=0.6,
+                step=0.1,
+                interactive=True,
+                label="Temperature",
+            )
+
+            clear = gr.Button("Restart")
+
+            gr.Markdown(introduction)
+
+        with gr.Column():
+            chat_state = gr.State(value=None)
+            img_list = gr.State(value=[])
+            chatbot = gr.Chatbot(label='GeoChat')
+
+            dataset = gr.Dataset(
+                components=[gr.Textbox(visible=False)],
+                samples=[['No Tag'], ['Scene Classification'],['Identify']],
+                type="index",
+                label='Task Shortcuts',
+            )
+            task_inst = gr.Markdown('**Hint:** Upload your image and chat')
+            with gr.Row():
+                text_input.render()
+                send = gr.Button("Send", variant='primary', size='sm', scale=1)
+
+    upload_flag = gr.State(value=0)
+    replace_flag = gr.State(value=0)
+    image.upload(image_upload_trigger, [upload_flag, replace_flag, img_list], [upload_flag, replace_flag])
+
+    with gr.Row():
+        with gr.Column():
+            gr.Examples(examples=[
+                ["demo_images/train_2956_0001.png", "Where are the airplanes located and what is their type?", upload_flag, replace_flag,
+                 img_list],
+                ["demo_images/7292.JPG", "How many buildings are flooded?", upload_flag,
+                 replace_flag, img_list],
+            ], inputs=[image, text_input, upload_flag, replace_flag, img_list], fn=example_trigger,
+                outputs=[upload_flag, replace_flag])
+        with gr.Column():
+            gr.Examples(examples=[
+                ["demo_images/church_183.png", "Classify the image in the following classes: Church, Beach, Dense Residential, Storage Tanks.",
+                 upload_flag, replace_flag, img_list],
+                ["demo_images/04444.png", "[identify] what is this {<8><26><22><37>}", upload_flag,
+                 replace_flag, img_list],
+            ], inputs=[image, text_input, upload_flag, replace_flag, img_list], fn=example_trigger,
+                outputs=[upload_flag, replace_flag])
+
+    dataset.click(
+        gradio_taskselect,
+        inputs=[dataset],
+        outputs=[text_input, task_inst],
+        show_progress="hidden",
+        postprocess=False,
+        queue=False,
+    )
+
+    text_input.submit(
+        gradio_ask,
+        [text_input, chatbot, chat_state, image, img_list, upload_flag, replace_flag],
+        [text_input, chatbot, chat_state, img_list, upload_flag, replace_flag], queue=False
+    ).success(
+        gradio_stream_answer,
+        [chatbot, chat_state, img_list, temperature],
+        [chatbot, chat_state]
+    ).success(
+        gradio_visualize,
+        [chatbot, image],
+        [chatbot],
+        queue=False,
+    )
+
+    send.click(
+        gradio_ask,
+        [text_input, chatbot, chat_state, image, img_list, upload_flag, replace_flag],
+        [text_input, chatbot, chat_state, img_list, upload_flag, replace_flag], queue=False
+    ).success(
+        gradio_stream_answer,
+        [chatbot, chat_state, img_list, temperature],
+        [chatbot, chat_state]
+    ).success(
+        gradio_visualize,
+        [chatbot, image],
+        [chatbot],
+        queue=False,
+    )
+
+    clear.click(gradio_reset, [chat_state, img_list], [chatbot, image, text_input, chat_state, img_list], queue=False)
+
+demo.launch(share=True, enable_queue=True,server_name='0.0.0.0')

.ipynb_checkpoints/pyproject-checkpoint.toml

@@ -0,0 +1,39 @@
+[build-system]
+requires = ["setuptools>=61.0"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "geochat"
+version = "1.1.1"
+description = "Grounded VLM for Remote Sensing"
+readme = "README.md"
+requires-python = ">=3.8"
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "License :: OSI Approved :: Apache Software License",
+]
+dependencies = [
+    "einops", "fastapi", "gradio==3.35.2", "markdown2[all]", "numpy",
+    "requests", "sentencepiece", "tokenizers>=0.12.1",
+    "torch==2.0.1", "torchvision==0.15.2", "uvicorn", "wandb",
+    "shortuuid", "httpx==0.24.0",
+    #"deepspeed==0.9.5",
+    "peft==0.4.0",
+    "transformers==4.31.0",
+    "accelerate==0.21.0",
+    "bitsandbytes==0.41.0",
+    "scikit-learn==1.2.2",
+    "sentencepiece==0.1.99",
+    "einops==0.6.1", "einops-exts==0.0.4", "timm==0.6.13",
+    "gradio_client==0.2.9"
+]
+
+[project.urls]
+"Homepage" = "https://github.com/mbzuai-oryx/GeoChat"
+"Bug Tracker" = "https://github.com/mbzuai-oryx/GeoChat/issues"
+
+[tool.setuptools.packages.find]
+exclude = ["assets*", "benchmark*", "docs", "dist*", "playground*", "scripts*", "tests*"]
+
+[tool.wheel]
+exclude = ["assets*", "benchmark*", "docs", "dist*", "playground*", "scripts*", "tests*"]

README.md

@@ -1,12 +1,231 @@
 ---
-title: Csu
-emoji: 🏢
-colorFrom: gray
-colorTo: yellow
+title: csu
+app_file: geochat_demo.py
 sdk: gradio
-sdk_version: 5.15.0
-app_file: app.py
-pinned: false
+sdk_version: 3.35.2
 ---
+# GeoChat <img src="images/logo_geochat.png" height="40">: Grounded Large Vision-Language Model for Remote Sensing [CVPR-2024]
+<p align="center">
+    <img src="https://i.imgur.com/waxVImv.png" alt="Oryx Video-ChatGPT">
+</p>
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+#### [Kartik Kuckreja](https://www.linkedin.com/in/kartik-kuckreja-930531221/)\*, [Muhammad Sohail Danish](https://www.linkedin.com/in/muhammad-sohail-danish/)\*, [Muzammal Naseer](https://muzammal-naseer.com/), [Abhijit Das](https://sites.google.com/site/dasabhijit2048/home), [Salman Khan](https://salman-h-khan.github.io/) and [Fahad Khan](https://sites.google.com/view/fahadkhans/home)
+\* Equally contributing first authors
+
+#### **Mohamed bin Zayed University of AI, Birla Institute of Technology & Science, Australian National University, Linkoping University**
+
+[![Website](https://img.shields.io/badge/Project-Website-87CEEB)](https://mbzuai-oryx.github.io/GeoChat)
+[![paper](https://img.shields.io/badge/arXiv-Paper-<COLOR>.svg)](https://arxiv.org/abs/2311.15826)
+[![video](https://img.shields.io/badge/Video-Presentation-F9D371)](https://youtu.be/KOKtkkKpNDk)
+
+---
+
+## 📢 Latest Updates
+- Supplementary material for the accepted paper is available here: [Supplementary](https://github.com/mbzuai-oryx/GeoChat/blob/main/docs/geochat_supp.pdf).
+- **Feb-28-24**: We open source the code, model, dataset, and evaluation scripts. 
+- **Feb-27-24**: GeoChat has been accepted to **CVPR-24** 🎉. 
+- **Nov-28-23**: GeoChat paper is released [arxiv link](https://arxiv.org/abs/2311.15826). 🔥🔥
+---
+
+
+
+## <img src="images/logo_geochat.png" height="40">Overview
+
+GeoChat is the first grounded Large Vision Language Model, specifically tailored to Remote Sensing(RS) scenarios. Unlike general-domain models, GeoChat excels in handling high-resolution RS imagery, employing region-level reasoning for comprehensive scene interpretation. Leveraging a newly created RS multimodal dataset, GeoChat is fine-tuned using the LLaVA-1.5 architecture. This results in robust zero-shot performance across various RS tasks, including image and region captioning, visual question answering, scene classification, visually grounded conversations, and referring object detection.
+
+---
+## Contents
+- [Install](#install)
+- [Model Zoo](https://github.com/mbzuai-oryx/GeoChat/blob/main/docs/MODEL_ZOO.md)
+- [Dataset](https://huggingface.co/datasets/MBZUAI/GeoChat_Instruct/blob/main/GeoChat_Instruct.json)
+- [Train](#train)
+- [Evaluation](#evaluation)
+
+## Install
+
+1. Clone this repository and navigate to GeoChat folder
+```bash
+git clone https://github.com/mbzuai-oryx/GeoChat.git
+cd GeoChat
+```
+
+2. Install Package
+```Shell
+conda create -n geochat python=3.10 -y
+conda activate geochat
+pip install --upgrade pip  # enable PEP 660 support
+pip install -e .
+```
+
+3. Install additional packages for training cases
+```
+pip install ninja
+pip install flash-attn --no-build-isolation
+```
+
+### Upgrade to latest code base
+
+```Shell
+git pull
+pip uninstall transformers
+pip install -e .
+```
+
+## GeoChat Weights and Demo
+Please check out our [Model Zoo](https://github.com/mbzuai-oryx/GeoChat/blob/main/docs/MODEL_ZOO.md) for all public GeoChat checkpoints, and check [LoRA.md](https://github.com/mbzuai-oryx/GeoChat/blob/main/docs/LoRA.md) for instructions on how to run the demo and training.
+
+## Train
+
+GeoChat training consists of visual instruction tuning using GeoChat_Instruct Dataset: 318k Vicuna-generated multimodal instruction-following data, finetuned over the pretrained weights of LlaVA-v1.5.
+
+We train GeoChat on 3 A100 GPUs with 40GB memory. To train on fewer GPUs, you can reduce the `per_device_train_batch_size` and increase the `gradient_accumulation_steps` accordingly. Always keep the global batch size the same: `per_device_train_batch_size` x `gradient_accumulation_steps` x `num_gpus`.
+
+### Hyperparameters
+We use a similar set of hyperparameters as Vicuna in finetuning.  Both hyperparameters used in pretraining and finetuning are provided below.
+
+| Hyperparameter | Global Batch Size | Learning rate | Epochs | Max length | Weight decay |
+| --- | ---: | ---: | ---: | ---: | ---: |
+| GeoChat-7B | 144 | 2e-5 | 1 | 2048 | 0 |
+
+### Pretrain (feature alignment)
+
+We use the pretrained projector from LLaVAv1.5, which is trained on 558K subset of the LAION-CC-SBU dataset with BLIP captions. It takes around 3.5 hours for LLaVA-v1.5-7B.
+
+- `--mm_projector_type mlp2x_gelu`: the two-layer MLP vision-language connector.
+- `--vision_tower openai/clip-vit-large-patch14-336`: CLIP ViT-L/14 336px.
+
+### Visual Instruction Tuning
+
+1. Prepare data
+
+Please download the annotation of the final mixture of our instruction tuning data [GeoChat_Instruct.json](https://huggingface.co/datasets/MBZUAI/GeoChat_Instruct/blob/main/GeoChat_Instruct.json), and download the split image zips from the [hugging face](https://huggingface.co/datasets/MBZUAI/GeoChat_Instruct). Save the multiple image zips in a single folder and run the following command to merge them:
+```Shell
+cat images_parta* > images.zip
+```
+Unzip the images.zip file to a folder and give the folder's path in [finetune_lora.sh](https://github.com/mbzuai-oryx/GeoChat/blob/main/scripts/finetune_lora.sh).
+
+2. Start training!
+
+Visual instruction tuning takes more time due to the increased resolution of CLIP to 504X504. It takes around ~25 hours to finetune GeoChat-7B on 3x A100 (40G).
+
+Training script with DeepSpeed ZeRO-3: [`finetune_lora.sh`](https://github.com/mbzuai-oryx/GeoChat/blob/main/scripts/finetune_lora.sh).
+
+Options to note:
+
+- `--mm_projector_type mlp2x_gelu`: the two-layer MLP vision-language connector.
+- `--vision_tower openai/clip-vit-large-patch14-336`: CLIP ViT-L/14 336px.
+- `--image_aspect_ratio pad`: this pads the non-square images to square, instead of cropping them; it slightly reduces hallucination.
+- `--group_by_modality_length True`: this should only be used when your instruction tuning dataset contains both language (e.g. ShareGPT) and multimodal (e.g. LLaVA-Instruct).
+- 
+## Evaluation
+
+We evaluate GeoChat on a diverse set of 7 benchmarks. To ensure the reproducibility, we evaluate the models with greedy decoding. We do not evaluate using beam search to make the inference process consistent with the chat demo of real-time outputs.
+See [Evaluation.md](https://github.com/mbzuai-oryx/GeoChat/blob/main/docs/Evaluation.md).
+
+## 🏆 Contributions
+
+- **RS multimodal instruction following dataset.** We present a novel data generation pipeline, to leverage existing object detection dataset to create short descriptions of the images, followed by using Vicuna-v1.5 to create conversations using the generated text alone. Further, we add visual question-answering and scene classification abilities 
+ using their corresponding datasets. This results in a total of 318k instruction pairs for RS domain.
+- **GeoChat.** Leveraging our dataset, we finetune LLaVA-1.5 to create the remote sensing-domain vision-language model - GeoChat. Our LoRA fine-tuning is efficient and avoids forgetting the necessary context embedded in fully-tuned LLaVA model, whose MLP projection is trained to align images into the word embedding space of the LLM (Vicuna-v1.5). This allows GeoChat to retain the conversation and instruction following abilities of LLaVA and extend its domain-knowledge to remote sensing tasks.  
+
+- **Evaluation Benchmark.** We also address the lack of evaluation benchmarks to assess the capability of existing VLMs on remote-sensing conversations. To this end, we setup evaluation protocols for conversation grounding in RS, as well as a setup a suite of tasks to allow comparisons with future efforts in this direction. We show various supervised as well as  zero-shot evaluations for different remote sensing tasks, including image captioning, visual question answering and scene classification to demonstrate the generalisability of GeoChat conversational VLM.
+
+---
+## 👁️💬 GeoChat : Grounded Large Vision-Language Model for Remote Sensing
+
+GeoChat can accomplish multiple tasks for remote-sensing (RS) image comprehension in a unified framework. Given suitable task tokens and user queries, the model can generate visually grounded responses (text with corresponding object locations - shown on top), visual question answering on images and regions (top left and bottom right, respectively) as well as scene classification (top right) and normal natural language conversations (bottom). This makes it the first RS VLM with grounding capability. 
+
+<p align="center">
+  <img src="images/overview2.png" alt="GeoChat Overview">
+</p>
+
+---
+
+## 🛰️ GeoChat : Architecture
+
+An overview of GeoChat - the first grounded large vision-language model for remote sensing. Given an image input together with a user query, a visual backbone is first used to encode patch-level tokens at a higher resolution via interpolating positional encodings. A multi-layer perceptron (MLP) is used to adapt vision-tokens to language space suitable for input to a Large Language Model (Vicuna 1.5). Besides visual inputs, region locations can also be input to the model together with task-specific prompts that specify the desired task required by the user. Given this context, the LLM can generate natural language responses interleaved with corresponding object locations. GeoChat can perform multiple tasks as shown on top e.g., scene classification, image/region captioning, VQA and grounded conversations.
+
+<p align="center">
+  <img src="images/architecture.png" alt="GeoChat Architectural">
+</p>
+
+---
+
+## 🔍 RS Multimodal Instruction Dataset
+
+Types of annotations available in the GeoChat instruction-set. For a given RS image, we obtain object attribute and relationship information, referring expressions and region captions along with their corresponding region annotations (shown over the image). This structured information is used to create the rich instruction-set with a total of 318k image-instruction pairs.
+
+<p align="center">
+  <img src="images/dataset.png" alt="Dataset Annotation Pipeline">
+</p>
+
+
+
+## 🤖 Qualitative results of GeoChat
+
+Qualitative results of GeoChat. (<em>left-right</em>) Results are shown on grounding, referring object detection, and disaster/damage detection. The user can provide task-specific tokens (e.g., <strong>[grounding]</strong>) to shape model responses according to the desired behavior. The model can generate textual responses (<em>right</em>), only visual grounding (<em>center</em>) and both text and object groundings interleaved together (<em>left</em>). The model can also specify object types, object counts, object attributes and object relationships.
+<p align="center">
+  <img src="images/examples.png" alt="Results_GCG">
+</p>
+
+---
+
+## 🤖 Visual Question Answering
+Qualitative examples for Visual Question Answering tasks. GeoChat is able to hold multi-turn conversations, based on various types of questions, including presence, count, complex comparisons and so on. It is able to detect objects and hold conversations against low resolution images as well.
+<p align="center">
+  <img src="images/vqa.jpg" alt="Visual Question Answering">
+</p>
+
+---
+
+## 🤖 Scene Classification
+Qualitative examples for scene classification. We give the model all the classes from the dataset and ask to choose only one.
+<p align="center">
+  <img src="images/scene.jpg" alt="Visual Question Answering">
+</p>
+
+---
+
+## 🤖 Grounded Description
+When asked to describe the image with the special token '[grounding]', GeoChat outputs both the description of the image as well as the bounding boxes for all the objects detected.
+<p align="center">
+  <img src="images/grounded.jpg" alt="Grounded Description">
+</p>
+
+---
+
+## 🤖 Referring Expression
+When asked about an object as a referred expression, GeoChat is able to locate it and draw rotated bounding boxes around it correspondingly.
+<p align="center">
+  <img src="images/ref1.jpg" alt="Referring Expression">
+</p>
+<p align="center">
+  <img src="images/ref_2.jpg" alt="Referring Expression">
+</p>
+
+---
+
+## 🤖 Region Caption
+Qualitative examples for region-based captioning. Given a bounding box, GeoChat is able to provide brief descriptions about the area or the object covered by the bounding box.
+<p align="center">
+  <img src="images/iden.jpg" alt="Region Caption">
+</p>
+
+---
+
+## 📜 Citation
+```bibtex
+  @article{kuckreja2023geochat,
+          title={GeoChat: Grounded Large Vision-Language Model for Remote Sensing},
+          author={Kuckreja, Kartik and Danish, Muhammad S. and Naseer, Muzammal and Das, Abhijit and Khan, Salman and Khan, Fahad S.},
+          journal={The IEEE/CVF Conference on Computer Vision and Pattern Recognition},
+          year={2024}
+  }
+```
+## 🙏 Acknowledgement
+We are thankful to LLaVA and Vicuna for releasing their models and code as open-source contributions.
+
+---
+[<img src="images/IVAL_logo.png" width="200" height="100">](https://www.ival-mbzuai.com)
+[<img src="images/Oryx_logo.png" width="100" height="100">](https://github.com/mbzuai-oryx)
+[<img src="images/MBZUAI_logo.png" width="360" height="85">](https://mbzuai.ac.ae)

docs/Customize_Component.md

@@ -0,0 +1,20 @@
+# Customize Components in GeoChat
+
+This is an initial guide on how to replace the LLMs, visual encoders, etc. with your choice of components.
+
+## LLM
+
+It is quite simple to swap out LLaMA to any other LLMs.  You can refer to our implementation of [`GeoChat_llama.py`](https://github.com/mbzuai-oryx/GeoChat/blob/main/geochat/model/language_model/geochat_llama.py) for an example of how to replace the LLM.
+
+Although it may seem that it still needs ~100 lines of code, most of them are copied from the original `llama.py` from HF.  The only part that is different is to insert some lines for processing the multimodal inputs.
+
+In `forward` function, you can see that we call `self.prepare_inputs_labels_for_multimodal` to process the multimodal inputs.  This function is defined in `GeoChatMetaForCausalLM` and you just need to insert it into the `forward` function of your LLM.
+
+In `prepare_inputs_for_generation` function, you can see that we add `images` to the `model_inputs`.  This is because we need to pass the images to the LLM during generation.
+
+These are basically all the changes you need to make to replace the LLM.
+
+## Visual Encoder
+
+You can check out [`clip_encoder.py`](https://github.com/haotian-liu/LLaVA/blob/main/llava/model/multimodal_encoder/clip_encoder.py) on how we implement the CLIP visual encoder.
+

docs/Data.md

@@ -0,0 +1,24 @@
+##  Finetuning Data
+We use GeoChat-Instruct to finetune our model. The instruction following dataset is present in GeoChat_Instruct.json and the images are present in the [huggingface repo](https://huggingface.co/datasets/MBZUAI/GeoChat_Instruct). The images are split into multiple files. Download the separate files in the same folder and run the following script to merge them. 
+
+```Shell
+cat images_parta* > images.zip
+```
+
+Unzip the images in a folder and provide the folder path in training and evaluation scripts.
+
+| Data file name | Size |
+| --- | ---: |
+| [GeoChat_Instruct](https://huggingface.co/datasets/MBZUAI/GeoChat_Instruct/blob/main/GeoChat_Instruct.json) | 263 MB |
+
+## Pretraining Dataset
+We use the same pretraining dataset as of LlaVA-v1.5.
+The pretraining dataset used in this release is a subset of CC-3M dataset, filtered with a more balanced concept coverage distribution.  Please see [here](https://huggingface.co/datasets/liuhaotian/LLaVA-CC3M-Pretrain-595K) for a detailed description of the dataset structure and how to download the images.
+
+If you already have CC-3M dataset on your disk, the image names follow this format: `GCC_train_000000000.jpg`.  You may edit the `image` field correspondingly if necessary.
+
+| Data | Chat File | Meta Data | Size |
+| --- |  --- |  --- | ---: |
+| CC-3M Concept-balanced 595K | [chat.json](https://huggingface.co/datasets/liuhaotian/LLaVA-CC3M-Pretrain-595K/blob/main/chat.json) | [metadata.json](https://huggingface.co/datasets/liuhaotian/LLaVA-CC3M-Pretrain-595K/blob/main/metadata.json) | 211 MB
+| LAION/CC/SBU BLIP-Caption Concept-balanced 558K | [blip_laion_cc_sbu_558k.json](https://huggingface.co/datasets/liuhaotian/LLaVA-Pretrain/blob/main/blip_laion_cc_sbu_558k.json) | [metadata.json](#) | 181 MB
+

docs/Evaluation.md

@@ -0,0 +1,54 @@
+# Evaluation
+
+We evaluate GeoChat on a variety of tasks, including scene classification, region captioning, visual grounding, grounding description and VQA.
+Converted files in the input format for GeoChat are available at [GeoChat-Bench](https://huggingface.co/datasets/MBZUAI/GeoChat-Bench/tree/main)
+
+
+Below we provide a general guideline for evaluating datasets.
+
+1. LRBEN/HRBEN.
+Images and ground truth for evaluation need to be downloaded from the following sources: [LRBEN](https://zenodo.org/records/6344334), [HRBEN](https://zenodo.org/records/6344367)
+Give the path to the extracted image folder in the evaluation script. We add the following text after each question during our evaluation.
+```
+<question>
+Answer the question using a single word or phrase.
+```
+```Shell
+python geochat/eval/batch_geochat_vqa.py \
+    --model-path /path/to/model \
+    --question-file path/to/jsonl/file \
+    --answer-file path/to/output/jsonl/file \
+    --image_folder path/to/image/folder/
+```
+2. Scene Classification.
+Download the images from the following sources, [UCmerced](http://weegee.vision.ucmerced.edu/datasets/landuse.html), [AID](https://drive.google.com/drive/folders/1-1D9DrYYWMGuuxx-qcvIIOV1oUkAVf-M). We add the following text after each question during our evaluation.
+```
+<question>
+Classify the image from the following classes. Answer in one word or a short phrase.
+```
+```Shell
+python geochat/eval/batch_geochat_scene.py \
+    --model-path /path/to/model \
+    --question-file path/to/jsonl/file \
+    --answer-file path/to/output/jsonl/file \
+    --image_folder path/to/image/folder/
+```
+
+3. Region-Captioning/Visual grounding.
+
+The evaluation images are present in the image.zip folder in [GeoChat_Instruct](https://huggingface.co/datasets/MBZUAI/GeoChat_Instruct/blob/main/images.zip). 
+```Shell
+python geochat/eval/batch_geochat_grounding.py \
+    --model-path /path/to/model \
+    --question-file path/to/jsonl/file \
+    --answer-file path/to/output/jsonl/file \
+    --image_folder path/to/image/folder/
+```
+
+```Shell
+python geochat/eval/batch_geochat_referring.py \
+    --model-path /path/to/model \
+    --question-file path/to/jsonl/file \
+    --answer-file path/to/output/jsonl/file \
+    --image_folder path/to/image/folder/
+```

docs/LoRA.md

@@ -0,0 +1,24 @@
+
+## Demo (Web UI)
+You need GeoChat-7B to run the demo locally. Download the model from [GeoChat-7B](https://huggingface.co/MBZUAI/geochat-7B). After loading the model, run this command by giving the model path to launch the gradio demo.
+#### Launch the demo
+```Shell
+python geochat_demo.py --model-path /path/to/model 
+```
+
+## Training
+
+Please see sample training scripts for [LoRA](https://github.com/mbzuai-oryx/GeoChat/blob/main/scripts/finetune_lora.sh)
+
+We provide sample DeepSpeed configs, [`zero3.json`](https://github.com/haotian-liu/LLaVA/blob/main/scripts/zero3.json) is more like PyTorch FSDP, and [`zero3_offload.json`](https://github.com/haotian-liu/LLaVA/blob/main/scripts/zero3_offload.json) can further save memory consumption by offloading parameters to CPU. `zero3.json` is usually faster than `zero3_offload.json` but requires more GPU memory, therefore, we recommend trying `zero3.json` first, and if you run out of GPU memory, try `zero3_offload.json`. You can also tweak the `per_device_train_batch_size` and `gradient_accumulation_steps` in the config to save memory, and just to make sure that `per_device_train_batch_size` and `gradient_accumulation_steps` remains the same.
+
+If you are having issues with ZeRO-3 configs, and there are enough VRAM, you may try [`zero2.json`](https://github.com/haotian-liu/LLaVA/blob/main/scripts/zero2.json). This consumes slightly more memory than ZeRO-3, and behaves more similar to PyTorch FSDP, while still supporting parameter-efficient tuning.
+
+## Create Merged Checkpoints
+
+```Shell
+python scripts/merge_lora_weights.py \
+    --model-path /path/to/lora_model \
+    --model-base /path/to/base_model \
+    --save-model-path /path/to/merge_model
+```

docs/MODEL_ZOO.md

@@ -0,0 +1,18 @@
+# Model Zoo
+
+| Base LLM | Vision Encoder | Pretrain Data | Pretraining schedule | Finetuning Data | Finetuning schedule | Download |
+|----------|----------------|---------------|----------------------|-----------------|--------------------|------------------
+| Vicuna-13B-v1.3 | CLIP-L-336px(extended to 504) | LCS-558K | 1e | Geochat_Instruct | proj-1e, lora-1e | [LoRA-Merged](https://huggingface.co/MBZUAI/geochat-7B) |
+
+## Projector weights
+We use the projector from LlaVA-1.5 for initialization. [Link](https://huggingface.co/liuhaotian/llava-v1.5-7b-lora)
+
+**NOTE**: When you use our pretrained projector for visual instruction tuning, it is very important to **use the same base LLM and vision encoder** as the one we used for pretraining the projector. Otherwise, the performance will be very bad.
+
+When using these projector weights to instruction tune your LMM, please make sure that these options are correctly set as follows,
+
+```Shell
+--mm_use_im_start_end False
+--mm_use_im_patch_token False
+```
+

docs/geochat_supp.pdf

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:dc9b5d3df4af5c06e59fc2258332e00b779c55d441405cb5a7fd7997d29b63fe
+size 4839915

geochat.egg-info/PKG-INFO

@@ -0,0 +1,260 @@
+Metadata-Version: 2.2
+Name: geochat
+Version: 1.1.1
+Summary: Grounded VLM for Remote Sensing
+Project-URL: Homepage, https://github.com/mbzuai-oryx/GeoChat
+Project-URL: Bug Tracker, https://github.com/mbzuai-oryx/GeoChat/issues
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: Apache Software License
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+Requires-Dist: einops
+Requires-Dist: fastapi
+Requires-Dist: gradio==3.35.2
+Requires-Dist: markdown2[all]
+Requires-Dist: numpy
+Requires-Dist: requests
+Requires-Dist: sentencepiece
+Requires-Dist: tokenizers>=0.12.1
+Requires-Dist: torch==2.0.1
+Requires-Dist: torchvision==0.15.2
+Requires-Dist: uvicorn
+Requires-Dist: wandb
+Requires-Dist: shortuuid
+Requires-Dist: httpx==0.24.0
+Requires-Dist: peft==0.4.0
+Requires-Dist: transformers==4.31.0
+Requires-Dist: accelerate==0.21.0
+Requires-Dist: bitsandbytes==0.41.0
+Requires-Dist: scikit-learn==1.2.2
+Requires-Dist: sentencepiece==0.1.99
+Requires-Dist: einops==0.6.1
+Requires-Dist: einops-exts==0.0.4
+Requires-Dist: timm==0.6.13
+Requires-Dist: gradio_client==0.2.9
+
+# GeoChat <img src="images/logo_geochat.png" height="40">: Grounded Large Vision-Language Model for Remote Sensing [CVPR-2024]
+<p align="center">
+    <img src="https://i.imgur.com/waxVImv.png" alt="Oryx Video-ChatGPT">
+</p>
+
+#### [Kartik Kuckreja](https://www.linkedin.com/in/kartik-kuckreja-930531221/)\*, [Muhammad Sohail Danish](https://www.linkedin.com/in/muhammad-sohail-danish/)\*, [Muzammal Naseer](https://muzammal-naseer.com/), [Abhijit Das](https://sites.google.com/site/dasabhijit2048/home), [Salman Khan](https://salman-h-khan.github.io/) and [Fahad Khan](https://sites.google.com/view/fahadkhans/home)
+\* Equally contributing first authors
+
+#### **Mohamed bin Zayed University of AI, Birla Institute of Technology & Science, Australian National University, Linkoping University**
+
+[![Website](https://img.shields.io/badge/Project-Website-87CEEB)](https://mbzuai-oryx.github.io/GeoChat)
+[![paper](https://img.shields.io/badge/arXiv-Paper-<COLOR>.svg)](https://arxiv.org/abs/2311.15826)
+[![video](https://img.shields.io/badge/Video-Presentation-F9D371)](https://youtu.be/KOKtkkKpNDk)
+
+---
+
+## 📢 Latest Updates
+- Supplementary material for the accepted paper is available here: [Supplementary](https://github.com/mbzuai-oryx/GeoChat/blob/main/docs/geochat_supp.pdf).
+- **Feb-28-24**: We open source the code, model, dataset, and evaluation scripts. 
+- **Feb-27-24**: GeoChat has been accepted to **CVPR-24** 🎉. 
+- **Nov-28-23**: GeoChat paper is released [arxiv link](https://arxiv.org/abs/2311.15826). 🔥🔥
+---
+
+
+
+## <img src="images/logo_geochat.png" height="40">Overview
+
+GeoChat is the first grounded Large Vision Language Model, specifically tailored to Remote Sensing(RS) scenarios. Unlike general-domain models, GeoChat excels in handling high-resolution RS imagery, employing region-level reasoning for comprehensive scene interpretation. Leveraging a newly created RS multimodal dataset, GeoChat is fine-tuned using the LLaVA-1.5 architecture. This results in robust zero-shot performance across various RS tasks, including image and region captioning, visual question answering, scene classification, visually grounded conversations, and referring object detection.
+
+---
+## Contents
+- [Install](#install)
+- [Model Zoo](https://github.com/mbzuai-oryx/GeoChat/blob/main/docs/MODEL_ZOO.md)
+- [Dataset](https://huggingface.co/datasets/MBZUAI/GeoChat_Instruct/blob/main/GeoChat_Instruct.json)
+- [Train](#train)
+- [Evaluation](#evaluation)
+
+## Install
+
+1. Clone this repository and navigate to GeoChat folder
+```bash
+git clone https://github.com/mbzuai-oryx/GeoChat.git
+cd GeoChat
+```
+
+2. Install Package
+```Shell
+conda create -n geochat python=3.10 -y
+conda activate geochat
+pip install --upgrade pip  # enable PEP 660 support
+pip install -e .
+```
+
+3. Install additional packages for training cases
+```
+pip install ninja
+pip install flash-attn --no-build-isolation
+```
+
+### Upgrade to latest code base
+
+```Shell
+git pull
+pip uninstall transformers
+pip install -e .
+```
+
+## GeoChat Weights and Demo
+Please check out our [Model Zoo](https://github.com/mbzuai-oryx/GeoChat/blob/main/docs/MODEL_ZOO.md) for all public GeoChat checkpoints, and check [LoRA.md](https://github.com/mbzuai-oryx/GeoChat/blob/main/docs/LoRA.md) for instructions on how to run the demo and training.
+
+## Train
+
+GeoChat training consists of visual instruction tuning using GeoChat_Instruct Dataset: 318k Vicuna-generated multimodal instruction-following data, finetuned over the pretrained weights of LlaVA-v1.5.
+
+We train GeoChat on 3 A100 GPUs with 40GB memory. To train on fewer GPUs, you can reduce the `per_device_train_batch_size` and increase the `gradient_accumulation_steps` accordingly. Always keep the global batch size the same: `per_device_train_batch_size` x `gradient_accumulation_steps` x `num_gpus`.
+
+### Hyperparameters
+We use a similar set of hyperparameters as Vicuna in finetuning.  Both hyperparameters used in pretraining and finetuning are provided below.
+
+| Hyperparameter | Global Batch Size | Learning rate | Epochs | Max length | Weight decay |
+| --- | ---: | ---: | ---: | ---: | ---: |
+| GeoChat-7B | 144 | 2e-5 | 1 | 2048 | 0 |
+
+### Pretrain (feature alignment)
+
+We use the pretrained projector from LLaVAv1.5, which is trained on 558K subset of the LAION-CC-SBU dataset with BLIP captions. It takes around 3.5 hours for LLaVA-v1.5-7B.
+
+- `--mm_projector_type mlp2x_gelu`: the two-layer MLP vision-language connector.
+- `--vision_tower openai/clip-vit-large-patch14-336`: CLIP ViT-L/14 336px.
+
+### Visual Instruction Tuning
+
+1. Prepare data
+
+Please download the annotation of the final mixture of our instruction tuning data [GeoChat_Instruct.json](https://huggingface.co/datasets/MBZUAI/GeoChat_Instruct/blob/main/GeoChat_Instruct.json), and download the split image zips from the [hugging face](https://huggingface.co/datasets/MBZUAI/GeoChat_Instruct). Save the multiple image zips in a single folder and run the following command to merge them:
+```Shell
+cat images_parta* > images.zip
+```
+Unzip the images.zip file to a folder and give the folder's path in [finetune_lora.sh](https://github.com/mbzuai-oryx/GeoChat/blob/main/scripts/finetune_lora.sh).
+
+2. Start training!
+
+Visual instruction tuning takes more time due to the increased resolution of CLIP to 504X504. It takes around ~25 hours to finetune GeoChat-7B on 3x A100 (40G).
+
+Training script with DeepSpeed ZeRO-3: [`finetune_lora.sh`](https://github.com/mbzuai-oryx/GeoChat/blob/main/scripts/finetune_lora.sh).
+
+Options to note:
+
+- `--mm_projector_type mlp2x_gelu`: the two-layer MLP vision-language connector.
+- `--vision_tower openai/clip-vit-large-patch14-336`: CLIP ViT-L/14 336px.
+- `--image_aspect_ratio pad`: this pads the non-square images to square, instead of cropping them; it slightly reduces hallucination.
+- `--group_by_modality_length True`: this should only be used when your instruction tuning dataset contains both language (e.g. ShareGPT) and multimodal (e.g. LLaVA-Instruct).
+- 
+## Evaluation
+
+We evaluate GeoChat on a diverse set of 7 benchmarks. To ensure the reproducibility, we evaluate the models with greedy decoding. We do not evaluate using beam search to make the inference process consistent with the chat demo of real-time outputs.
+See [Evaluation.md](https://github.com/mbzuai-oryx/GeoChat/blob/main/docs/Evaluation.md).
+
+## 🏆 Contributions
+
+- **RS multimodal instruction following dataset.** We present a novel data generation pipeline, to leverage existing object detection dataset to create short descriptions of the images, followed by using Vicuna-v1.5 to create conversations using the generated text alone. Further, we add visual question-answering and scene classification abilities 
+ using their corresponding datasets. This results in a total of 318k instruction pairs for RS domain.
+- **GeoChat.** Leveraging our dataset, we finetune LLaVA-1.5 to create the remote sensing-domain vision-language model - GeoChat. Our LoRA fine-tuning is efficient and avoids forgetting the necessary context embedded in fully-tuned LLaVA model, whose MLP projection is trained to align images into the word embedding space of the LLM (Vicuna-v1.5). This allows GeoChat to retain the conversation and instruction following abilities of LLaVA and extend its domain-knowledge to remote sensing tasks.  
+
+- **Evaluation Benchmark.** We also address the lack of evaluation benchmarks to assess the capability of existing VLMs on remote-sensing conversations. To this end, we setup evaluation protocols for conversation grounding in RS, as well as a setup a suite of tasks to allow comparisons with future efforts in this direction. We show various supervised as well as  zero-shot evaluations for different remote sensing tasks, including image captioning, visual question answering and scene classification to demonstrate the generalisability of GeoChat conversational VLM.
+
+---
+## 👁️💬 GeoChat : Grounded Large Vision-Language Model for Remote Sensing
+
+GeoChat can accomplish multiple tasks for remote-sensing (RS) image comprehension in a unified framework. Given suitable task tokens and user queries, the model can generate visually grounded responses (text with corresponding object locations - shown on top), visual question answering on images and regions (top left and bottom right, respectively) as well as scene classification (top right) and normal natural language conversations (bottom). This makes it the first RS VLM with grounding capability. 
+
+<p align="center">
+  <img src="images/overview2.png" alt="GeoChat Overview">
+</p>
+
+---
+
+## 🛰️ GeoChat : Architecture
+
+An overview of GeoChat - the first grounded large vision-language model for remote sensing. Given an image input together with a user query, a visual backbone is first used to encode patch-level tokens at a higher resolution via interpolating positional encodings. A multi-layer perceptron (MLP) is used to adapt vision-tokens to language space suitable for input to a Large Language Model (Vicuna 1.5). Besides visual inputs, region locations can also be input to the model together with task-specific prompts that specify the desired task required by the user. Given this context, the LLM can generate natural language responses interleaved with corresponding object locations. GeoChat can perform multiple tasks as shown on top e.g., scene classification, image/region captioning, VQA and grounded conversations.
+
+<p align="center">
+  <img src="images/architecture.png" alt="GeoChat Architectural">
+</p>
+
+---
+
+## 🔍 RS Multimodal Instruction Dataset
+
+Types of annotations available in the GeoChat instruction-set. For a given RS image, we obtain object attribute and relationship information, referring expressions and region captions along with their corresponding region annotations (shown over the image). This structured information is used to create the rich instruction-set with a total of 318k image-instruction pairs.
+
+<p align="center">
+  <img src="images/dataset.png" alt="Dataset Annotation Pipeline">
+</p>
+
+
+
+## 🤖 Qualitative results of GeoChat
+
+Qualitative results of GeoChat. (<em>left-right</em>) Results are shown on grounding, referring object detection, and disaster/damage detection. The user can provide task-specific tokens (e.g., <strong>[grounding]</strong>) to shape model responses according to the desired behavior. The model can generate textual responses (<em>right</em>), only visual grounding (<em>center</em>) and both text and object groundings interleaved together (<em>left</em>). The model can also specify object types, object counts, object attributes and object relationships.
+<p align="center">
+  <img src="images/examples.png" alt="Results_GCG">
+</p>
+
+---
+
+## 🤖 Visual Question Answering
+Qualitative examples for Visual Question Answering tasks. GeoChat is able to hold multi-turn conversations, based on various types of questions, including presence, count, complex comparisons and so on. It is able to detect objects and hold conversations against low resolution images as well.
+<p align="center">
+  <img src="images/vqa.jpg" alt="Visual Question Answering">
+</p>
+
+---
+
+## 🤖 Scene Classification
+Qualitative examples for scene classification. We give the model all the classes from the dataset and ask to choose only one.
+<p align="center">
+  <img src="images/scene.jpg" alt="Visual Question Answering">
+</p>
+
+---
+
+## 🤖 Grounded Description
+When asked to describe the image with the special token '[grounding]', GeoChat outputs both the description of the image as well as the bounding boxes for all the objects detected.
+<p align="center">
+  <img src="images/grounded.jpg" alt="Grounded Description">
+</p>
+
+---
+
+## 🤖 Referring Expression
+When asked about an object as a referred expression, GeoChat is able to locate it and draw rotated bounding boxes around it correspondingly.
+<p align="center">
+  <img src="images/ref1.jpg" alt="Referring Expression">
+</p>
+<p align="center">
+  <img src="images/ref_2.jpg" alt="Referring Expression">
+</p>
+
+---
+
+## 🤖 Region Caption
+Qualitative examples for region-based captioning. Given a bounding box, GeoChat is able to provide brief descriptions about the area or the object covered by the bounding box.
+<p align="center">
+  <img src="images/iden.jpg" alt="Region Caption">
+</p>
+
+---
+
+## 📜 Citation
+```bibtex
+  @article{kuckreja2023geochat,
+          title={GeoChat: Grounded Large Vision-Language Model for Remote Sensing},
+          author={Kuckreja, Kartik and Danish, Muhammad S. and Naseer, Muzammal and Das, Abhijit and Khan, Salman and Khan, Fahad S.},
+          journal={The IEEE/CVF Conference on Computer Vision and Pattern Recognition},
+          year={2024}
+  }
+```
+## 🙏 Acknowledgement
+We are thankful to LLaVA and Vicuna for releasing their models and code as open-source contributions.
+
+---
+[<img src="images/IVAL_logo.png" width="200" height="100">](https://www.ival-mbzuai.com)
+[<img src="images/Oryx_logo.png" width="100" height="100">](https://github.com/mbzuai-oryx)
+[<img src="images/MBZUAI_logo.png" width="360" height="85">](https://mbzuai.ac.ae)

geochat.egg-info/SOURCES.txt

@@ -0,0 +1,51 @@
+README.md
+pyproject.toml
+geochat/__init__.py
+geochat/constants.py
+geochat/conversation.py
+geochat/mm_utils.py
+geochat/utils.py
+geochat.egg-info/PKG-INFO
+geochat.egg-info/SOURCES.txt
+geochat.egg-info/dependency_links.txt
+geochat.egg-info/requires.txt
+geochat.egg-info/top_level.txt
+geochat/eval/batch_geochat_grounding.py
+geochat/eval/batch_geochat_referring.py
+geochat/eval/batch_geochat_scene.py
+geochat/eval/batch_geochat_vqa.py
+geochat/model/__init__.py
+geochat/model/apply_delta.py
+geochat/model/builder.py
+geochat/model/consolidate.py
+geochat/model/geochat_arch.py
+geochat/model/make_delta.py
+geochat/model/utils.py
+geochat/model/language_model/geochat_llama.py
+geochat/model/language_model/geochat_mpt.py
+geochat/model/language_model/mpt/adapt_tokenizer.py
+geochat/model/language_model/mpt/attention.py
+geochat/model/language_model/mpt/blocks.py
+geochat/model/language_model/mpt/configuration_mpt.py
+geochat/model/language_model/mpt/custom_embedding.py
+geochat/model/language_model/mpt/flash_attn_triton.py
+geochat/model/language_model/mpt/hf_prefixlm_converter.py
+geochat/model/language_model/mpt/meta_init_context.py
+geochat/model/language_model/mpt/modeling_mpt.py
+geochat/model/language_model/mpt/norm.py
+geochat/model/language_model/mpt/param_init_fns.py
+geochat/model/multimodal_encoder/builder.py
+geochat/model/multimodal_encoder/clip_encoder.py
+geochat/model/multimodal_projector/builder.py
+geochat/serve/__init__.py
+geochat/serve/cli.py
+geochat/serve/controller.py
+geochat/serve/gradio_trial.py
+geochat/serve/gradio_web_server.py
+geochat/serve/model_worker.py
+geochat/serve/register_worker.py
+geochat/serve/test_message.py
+geochat/train/geochat_trainer.py
+geochat/train/llama_flash_attn_monkey_patch.py
+geochat/train/train.py
+geochat/train/train_mem.py

geochat.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

geochat.egg-info/requires.txt

@@ -0,0 +1,24 @@
+einops
+fastapi
+gradio==3.35.2
+markdown2[all]
+numpy
+requests
+sentencepiece
+tokenizers>=0.12.1
+torch==2.0.1
+torchvision==0.15.2
+uvicorn
+wandb
+shortuuid
+httpx==0.24.0
+peft==0.4.0
+transformers==4.31.0
+accelerate==0.21.0
+bitsandbytes==0.41.0
+scikit-learn==1.2.2
+sentencepiece==0.1.99
+einops==0.6.1
+einops-exts==0.0.4
+timm==0.6.13
+gradio_client==0.2.9

geochat.egg-info/top_level.txt

@@ -0,0 +1,3 @@
+demo_images
+geochat
+images

geochat/__init__.py

@@ -0,0 +1 @@
+from .model import GeoChatLlamaForCausalLM

geochat/constants.py

@@ -0,0 +1,12 @@
+CONTROLLER_HEART_BEAT_EXPIRATION = 30
+WORKER_HEART_BEAT_INTERVAL = 15
+
+LOGDIR = "."
+
+# Model Constants
+IGNORE_INDEX = -100
+IMAGE_TOKEN_INDEX = -200
+DEFAULT_IMAGE_TOKEN = "<image>"
+DEFAULT_IMAGE_PATCH_TOKEN = "<im_patch>"
+DEFAULT_IM_START_TOKEN = "<im_start>"
+DEFAULT_IM_END_TOKEN = "<im_end>"

geochat/conversation.py

@@ -0,0 +1,520 @@
+import dataclasses
+from enum import auto, Enum
+from typing import List, Tuple
+from PIL import Image
+from threading import Thread
+
+from geochat.constants import IMAGE_TOKEN_INDEX, 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 geochat.utils import disable_torch_init
+from geochat.mm_utils import process_images_demo, tokenizer_image_token, get_model_name_from_path, KeywordsStoppingCriteria
+from transformers import StoppingCriteria, StoppingCriteriaList, TextIteratorStreamer,TextStreamer
+import torch
+import dataclasses
+from enum import auto, Enum
+from typing import List, Tuple, Any
+
+
+class SeparatorStyle(Enum):
+    """Different separator style."""
+    SINGLE = auto()
+    TWO = auto()
+    MPT = auto()
+    PLAIN = auto()
+    LLAMA_2 = auto()
+
+
+@dataclasses.dataclass
+class Conversation:
+    """A class that keeps all conversation history."""
+    system: str
+    roles: List[str]
+    messages: List[List[str]]
+    offset: int
+    sep_style: SeparatorStyle = SeparatorStyle.SINGLE
+    sep: str = "###"
+    sep2: str = None
+    version: str = "Unknown"
+
+    skip_next: bool = False
+
+    def get_prompt(self):
+        messages = self.messages
+        if len(messages) > 0 and type(messages[0][1]) is tuple:
+            messages = self.messages.copy()
+            init_role, init_msg = messages[0].copy()
+            init_msg = init_msg[0].replace("<image>", "").strip()
+            if 'mmtag' in self.version:
+                messages[0] = (init_role, init_msg)
+                messages.insert(0, (self.roles[0], "<Image><image></Image>"))
+                messages.insert(1, (self.roles[1], "Received."))
+            else:
+                messages[0] = (init_role, "<image>\n" + init_msg)
+
+        if self.sep_style == SeparatorStyle.SINGLE:
+            ret = self.system + self.sep
+            for role, message in messages:
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    ret += role + ": " + message + self.sep
+                else:
+                    ret += role + ":"
+        elif self.sep_style == SeparatorStyle.TWO:
+            seps = [self.sep, self.sep2]
+            ret = self.system + seps[0]
+            for i, (role, message) in enumerate(messages):
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    ret += role + ": " + message + seps[i % 2]
+                else:
+                    ret += role + ":"
+        elif self.sep_style == SeparatorStyle.MPT:
+            ret = self.system + self.sep
+            for role, message in messages:
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    ret += role + message + self.sep
+                else:
+                    ret += role
+        elif self.sep_style == SeparatorStyle.LLAMA_2:
+            wrap_sys = lambda msg: f"<<SYS>>\n{msg}\n<</SYS>>\n\n"
+            wrap_inst = lambda msg: f"[INST] {msg} [/INST]"
+            ret = ""
+
+            for i, (role, message) in enumerate(messages):
+                if i == 0:
+                    assert message, "first message should not be none"
+                    assert role == self.roles[0], "first message should come from user"
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    if i == 0: message = wrap_sys(self.system) + message
+                    if i % 2 == 0:
+                        message = wrap_inst(message)
+                        ret += self.sep + message
+                    else:
+                        ret += " " + message + " " + self.sep2
+                else:
+                    ret += ""
+            ret = ret.lstrip(self.sep)
+        elif self.sep_style == SeparatorStyle.PLAIN:
+            seps = [self.sep, self.sep2]
+            ret = self.system
+            for i, (role, message) in enumerate(messages):
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    ret += message + seps[i % 2]
+                else:
+                    ret += ""
+        else:
+            raise ValueError(f"Invalid style: {self.sep_style}")
+
+        return ret
+
+    def append_message(self, role, message):
+        self.messages.append([role, message])
+
+    def get_images(self, return_pil=False):
+        images = []
+        for i, (role, msg) in enumerate(self.messages[self.offset:]):
+            if i % 2 == 0:
+                if type(msg) is tuple:
+                    import base64
+                    from io import BytesIO
+                    from PIL import Image
+                    msg, image, image_process_mode = msg
+                    if image_process_mode == "Pad":
+                        def expand2square(pil_img, background_color=(122, 116, 104)):
+                            width, height = pil_img.size
+                            if width == height:
+                                return pil_img
+                            elif width > height:
+                                result = Image.new(pil_img.mode, (width, width), background_color)
+                                result.paste(pil_img, (0, (width - height) // 2))
+                                return result
+                            else:
+                                result = Image.new(pil_img.mode, (height, height), background_color)
+                                result.paste(pil_img, ((height - width) // 2, 0))
+                                return result
+                        image = expand2square(image)
+                    elif image_process_mode in ["Default", "Crop"]:
+                        pass
+                    elif image_process_mode == "Resize":
+                        image = image.resize((336, 336))
+                    else:
+                        raise ValueError(f"Invalid image_process_mode: {image_process_mode}")
+                    max_hw, min_hw = max(image.size), min(image.size)
+                    aspect_ratio = max_hw / min_hw
+                    max_len, min_len = 800, 400
+                    shortest_edge = int(min(max_len / aspect_ratio, min_len, min_hw))
+                    longest_edge = int(shortest_edge * aspect_ratio)
+                    W, H = image.size
+                    if longest_edge != max(image.size):
+                        if H > W:
+                            H, W = longest_edge, shortest_edge
+                        else:
+                            H, W = shortest_edge, longest_edge
+                        image = image.resize((W, H))
+                    if return_pil:
+                        images.append(image)
+                    else:
+                        buffered = BytesIO()
+                        image.save(buffered, format="PNG")
+                        img_b64_str = base64.b64encode(buffered.getvalue()).decode()
+                        images.append(img_b64_str)
+        return images
+
+    def to_gradio_chatbot(self):
+        ret = []
+        for i, (role, msg) in enumerate(self.messages[self.offset:]):
+            if i % 2 == 0:
+                if type(msg) is tuple:
+                    import base64
+                    from io import BytesIO
+                    msg, image, image_process_mode = msg
+                    max_hw, min_hw = max(image.size), min(image.size)
+                    aspect_ratio = max_hw / min_hw
+                    max_len, min_len = 800, 400
+                    shortest_edge = int(min(max_len / aspect_ratio, min_len, min_hw))
+                    longest_edge = int(shortest_edge * aspect_ratio)
+                    W, H = image.size
+                    if H > W:
+                        H, W = longest_edge, shortest_edge
+                    else:
+                        H, W = shortest_edge, longest_edge
+                    image = image.resize((W, H))
+                    buffered = BytesIO()
+                    image.save(buffered, format="JPEG")
+                    img_b64_str = base64.b64encode(buffered.getvalue()).decode()
+                    img_str = f'<img src="data:image/png;base64,{img_b64_str}" alt="user upload image" />'
+                    msg = img_str + msg.replace('<image>', '').strip()
+                    ret.append([msg, None])
+                else:
+                    ret.append([msg, None])
+            else:
+                ret[-1][-1] = msg
+        return ret
+
+    def copy(self):
+        return Conversation(
+            system=self.system,
+            roles=self.roles,
+            messages=[[x, y] for x, y in self.messages],
+            offset=self.offset,
+            sep_style=self.sep_style,
+            sep=self.sep,
+            sep2=self.sep2,
+            version=self.version)
+
+    def dict(self):
+        if len(self.get_images()) > 0:
+            return {
+                "system": self.system,
+                "roles": self.roles,
+                "messages": [[x, y[0] if type(y) is tuple else y] for x, y in self.messages],
+                "offset": self.offset,
+                "sep": self.sep,
+                "sep2": self.sep2,
+            }
+        return {
+            "system": self.system,
+            "roles": self.roles,
+            "messages": self.messages,
+            "offset": self.offset,
+            "sep": self.sep,
+            "sep2": self.sep2,
+        }
+
+
+conv_vicuna_v0 = Conversation(
+    system="A chat between a curious human and an artificial intelligence assistant. "
+           "The assistant gives helpful, detailed, and polite answers to the human's questions.",
+    roles=("Human", "Assistant"),
+    messages=(
+        ("Human", "What are the key differences between renewable and non-renewable energy sources?"),
+        ("Assistant",
+            "Renewable energy sources are those that can be replenished naturally in a relatively "
+            "short amount of time, such as solar, wind, hydro, geothermal, and biomass. "
+            "Non-renewable energy sources, on the other hand, are finite and will eventually be "
+            "depleted, such as coal, oil, and natural gas. Here are some key differences between "
+            "renewable and non-renewable energy sources:\n"
+            "1. Availability: Renewable energy sources are virtually inexhaustible, while non-renewable "
+            "energy sources are finite and will eventually run out.\n"
+            "2. Environmental impact: Renewable energy sources have a much lower environmental impact "
+            "than non-renewable sources, which can lead to air and water pollution, greenhouse gas emissions, "
+            "and other negative effects.\n"
+            "3. Cost: Renewable energy sources can be more expensive to initially set up, but they typically "
+            "have lower operational costs than non-renewable sources.\n"
+            "4. Reliability: Renewable energy sources are often more reliable and can be used in more remote "
+            "locations than non-renewable sources.\n"
+            "5. Flexibility: Renewable energy sources are often more flexible and can be adapted to different "
+            "situations and needs, while non-renewable sources are more rigid and inflexible.\n"
+            "6. Sustainability: Renewable energy sources are more sustainable over the long term, while "
+            "non-renewable sources are not, and their depletion can lead to economic and social instability.\n")
+    ),
+    offset=2,
+    sep_style=SeparatorStyle.SINGLE,
+    sep="###",
+)
+
+conv_vicuna_v1 = Conversation(
+    system="A chat between a curious user and an artificial intelligence assistant. "
+    "The assistant gives helpful, detailed, and polite answers to the user's questions.",
+    roles=("USER", "ASSISTANT"),
+    version="v1",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.TWO,
+    sep=" ",
+    sep2="</s>",
+)
+
+conv_llama_2 = Conversation(
+    system="""You are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe.  Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature.
+
+If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.""",
+    roles=("USER", "ASSISTANT"),
+    version="llama_v2",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.LLAMA_2,
+    sep="<s>",
+    sep2="</s>",
+)
+
+conv_llava_llama_2 = Conversation(
+    system="You are a helpful language and vision assistant. "
+           "You are able to understand the visual content that the user provides, "
+           "and assist the user with a variety of tasks using natural language.",
+    roles=("USER", "ASSISTANT"),
+    version="llama_v2",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.LLAMA_2,
+    sep="<s>",
+    sep2="</s>",
+)
+
+conv_mpt = Conversation(
+    system="""<|im_start|>system
+A conversation between a user and an LLM-based AI assistant. The assistant gives helpful and honest answers.""",
+    roles=("<|im_start|>user\n", "<|im_start|>assistant\n"),
+    version="mpt",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.MPT,
+    sep="<|im_end|>",
+)
+
+conv_llava_plain = Conversation(
+    system="",
+    roles=("", ""),
+    messages=(
+    ),
+    offset=0,
+    sep_style=SeparatorStyle.PLAIN,
+    sep="\n",
+)
+
+conv_llava_v0 = Conversation(
+    system="A chat between a curious human and an artificial intelligence assistant. "
+           "The assistant gives helpful, detailed, and polite answers to the human's questions.",
+    roles=("Human", "Assistant"),
+    messages=(
+    ),
+    offset=0,
+    sep_style=SeparatorStyle.SINGLE,
+    sep="###",
+)
+
+conv_llava_v0_mmtag = Conversation(
+    system="A chat between a curious user and an artificial intelligence assistant. "
+           "The assistant is able to understand the visual content that the user provides, and assist the user with a variety of tasks using natural language."
+           "The visual content will be provided with the following format: <Image>visual content</Image>.",
+    roles=("Human", "Assistant"),
+    messages=(
+    ),
+    offset=0,
+    sep_style=SeparatorStyle.SINGLE,
+    sep="###",
+    version="v0_mmtag",
+)
+
+conv_llava_v1 = Conversation(
+    system="A chat between a curious human and an artificial intelligence assistant. "
+           "The assistant gives helpful, detailed, and polite answers to the human's questions.",
+    roles=("USER", "ASSISTANT"),
+    version="v1",
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.TWO,
+    sep=" ",
+    sep2="</s>",
+)
+
+conv_llava_v1_mmtag = Conversation(
+    system="A chat between a curious user and an artificial intelligence assistant. "
+           "The assistant is able to understand the visual content that the user provides, and assist the user with a variety of tasks using natural language."
+           "The visual content will be provided with the following format: <Image>visual content</Image>.",
+    roles=("USER", "ASSISTANT"),
+    messages=(),
+    offset=0,
+    sep_style=SeparatorStyle.TWO,
+    sep=" ",
+    sep2="</s>",
+    version="v1_mmtag",
+)
+
+default_conversation = conv_vicuna_v0
+conv_templates = {
+    "default": conv_vicuna_v0,
+    "v0": conv_vicuna_v0,
+    "v1": conv_vicuna_v1,
+    "vicuna_v1": conv_vicuna_v1,
+    "llama_2": conv_llama_2,
+
+    "plain": conv_llava_plain,
+    "v0_plain": conv_llava_plain,
+    "llava_v0": conv_llava_v0,
+    "v0_mmtag": conv_llava_v0_mmtag,
+    "llava_v1": conv_llava_v1,
+    "v1_mmtag": conv_llava_v1_mmtag,
+    "llava_llama_2": conv_llava_llama_2,
+
+    "mpt": conv_mpt,
+}
+
+class Chat:
+    def __init__(self, model, image_processor,tokenizer, device='cuda:0', stopping_criteria=None):
+        self.device = device
+        self.model = model
+        self.vis_processor = image_processor
+        self.tokenizer=tokenizer
+
+        # if stopping_criteria is not None:
+        #     self.stopping_criteria = stopping_criteria
+        # else:
+        #     stop_words_ids = [torch.tensor([2]).to(self.device)]
+        #     self.stopping_criteria = StoppingCriteriaList([StoppingCriteriaSub(stops=stop_words_ids)])
+
+    def ask(self, text, conv):
+        # import pdb;pdb.set_trace()
+        if len(conv.messages) > 0 and conv.messages[-1][0] == conv.roles[0] \
+                and conv.messages[-1][1][-9:] == '<image>\n':  # last message is image.
+            conv.messages[-1][1] = ' '.join([conv.messages[-1][1], text])
+        else:
+            conv.append_message(conv.roles[0], text)
+
+    def answer_prepare(self, conv, img_list, max_new_tokens=300, num_beams=1, min_length=1, top_p=0.9,
+                       repetition_penalty=1.05, length_penalty=1, temperature=1.0, max_length=2000):
+        conv.append_message(conv.roles[1], None)
+        prompt = conv.get_prompt()
+        # prompt='A chat between a curious human and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the human\'s questions. USER: <image>\n hello ASSISTANT:'
+        text_input_ids = tokenizer_image_token(prompt, self.tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).to(device=self.device)
+
+        stop_str = conv.sep if conv.sep_style != SeparatorStyle.TWO else conv.sep2
+        keywords = [stop_str]
+        stopping_criteria = KeywordsStoppingCriteria(keywords, self.tokenizer, text_input_ids)
+        current_max_len = text_input_ids.shape[1] + max_new_tokens
+        if current_max_len - max_length > 0:
+            print('Warning: The number of tokens in current conversation exceeds the max length. '
+                  'The model will not see the contexts outside the range.')
+        begin_idx = max(0, current_max_len - max_length)
+        embs = text_input_ids[:, begin_idx:]
+
+        generation_kwargs = dict(
+            input_ids=embs,
+            images=img_list[0],
+            max_new_tokens=max_new_tokens,
+            stopping_criteria=[stopping_criteria],
+            num_beams=num_beams,
+            do_sample=True,
+            min_length=min_length,
+            top_p=top_p,
+            use_cache=True,
+            repetition_penalty=repetition_penalty,
+            length_penalty=length_penalty,
+            temperature=float(temperature),
+        )
+        return generation_kwargs
+
+    # def answer(self, conv, img_list, **kargs):
+    #     generation_dict = self.answer_prepare(conv, img_list, **kargs)
+    #     output_token = self.model_generate(**generation_dict)[0]
+    #     output_text = self.model.llama_tokenizer.decode(output_token, skip_special_tokens=True)
+
+    #     output_text = output_text.split('###')[0]  # remove the stop sign '###'
+    #     output_text = output_text.split('Assistant:')[-1].strip()
+
+    #     conv.messages[-1][1] = output_text
+    #     return output_text, output_token.cpu().numpy()
+
+    def stream_answer(self, conv, img_list, **kargs):
+        generation_kwargs = self.answer_prepare(conv, img_list, **kargs)
+        
+        streamer = TextIteratorStreamer(self.tokenizer,skip_prompt=True, skip_special_tokens=True)
+        generation_kwargs['streamer'] = streamer
+        # import pdb;pdb.set_trace()
+        # output_ids=self.model.generate(*generation_kwargs)
+        output=self.model_generate(kwargs=generation_kwargs)
+        # thread = Thread(target=self.model_generate, kwargs=generation_kwargs)
+        # thread.start()
+        return streamer
+
+    def model_generate(self, *args, **kwargs):
+        # for 8 bit and 16 bit compatibility
+        with torch.inference_mode():
+            output = self.model.generate(kwargs['kwargs']['input_ids'],
+                                         images=kwargs['kwargs']['images'],
+                                         do_sample=False,
+                                         temperature=kwargs['kwargs']['temperature'],
+                                         max_new_tokens=kwargs['kwargs']['max_new_tokens'],
+                                         streamer=kwargs['kwargs']['streamer'],
+                                         use_cache=kwargs['kwargs']['use_cache'],
+                                         stopping_criteria=kwargs['kwargs']['stopping_criteria'])
+            # import pdb;pdb.set_trace()
+            # print(output)
+            outputs = self.tokenizer.decode(output[0,kwargs['kwargs']['input_ids'].shape[1]:]).strip()
+            # print(outputs)
+        return output
+
+    def encode_img(self, img_list):
+        
+        image = img_list[0]
+        # image='/share/data/drive_3/kartik/LLaVA/output_images/output.jpg'
+        img_list.pop(0)
+        if isinstance(image, str):  # is a image path
+            raw_image = Image.open(image).convert('RGB')
+            image = process_images_demo([raw_image], self.vis_processor)
+            # print("raw")
+            # image = self.vis_processor(raw_image).unsqueeze(0).to(self.device)
+        elif isinstance(image, Image.Image):
+            raw_image = image
+            image = process_images_demo([raw_image], self.vis_processor )
+            image=image.to(device=self.device,dtype=torch.float16)
+            # print("Image")
+            # image = self.vis_processor(raw_image).unsqueeze(0).to(self.device)
+        elif isinstance(image, torch.Tensor):
+            if len(image.shape) == 3:
+                image = image.unsqueeze(0)
+        image = image.to(self.device)
+
+        # image_emb, _ = self.model.encode_img(image)
+        img_list.append(image)
+
+    def upload_img(self, image, conv, img_list):
+        conv.append_message(conv.roles[0], DEFAULT_IMAGE_TOKEN+'\n')
+        img_list.append(image)
+        msg = "Received."
+
+        return msg
+
+
+
+# if __name__ == "__main__":
+#     print(default_conversation.get_prompt())

geochat/eval/batch_geochat_grounding.py

@@ -0,0 +1,138 @@
+import argparse
+import torch
+import os
+import json
+from tqdm import tqdm
+import shortuuid
+
+from geochat.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN
+from geochat.conversation import conv_templates, SeparatorStyle
+from geochat.model.builder import load_pretrained_model
+from geochat.utils import disable_torch_init
+from geochat.mm_utils import tokenizer_image_token, get_model_name_from_path, KeywordsStoppingCriteria
+
+from PIL import Image
+import math
+def split_list(lst, n):
+    """Split a list into n (roughly) equal-sized chunks"""
+    chunk_size = math.ceil(len(lst) / n)  # integer division
+    return [lst[i:i+chunk_size] for i in range(0, len(lst), chunk_size)]
+
+
+def get_chunk(lst, n, k):
+    chunks = split_list(lst, n)
+    return chunks[k]
+
+
+def eval_model(args):
+    # Model
+    disable_torch_init()
+    model_path = os.path.expanduser(args.model_path)
+    model_name = get_model_name_from_path(model_path)
+    tokenizer, model, image_processor, context_len = load_pretrained_model(args.model_path, args.model_base, model_name)
+    import pdb;pdb.set_trace()
+    # print(model)
+    questions=[]
+    questions = [json.loads(q) for q in open(os.path.expanduser(args.question_file), "r")]
+
+        
+    questions = get_chunk(questions, args.num_chunks, args.chunk_idx)
+    answers_file = os.path.expanduser(args.answers_file)
+    os.makedirs(os.path.dirname(answers_file), exist_ok=True)
+    
+    ans_file = open(answers_file, "w")
+    
+    for i in tqdm(range(0,len(questions),args.batch_size)):
+        input_batch=[]
+        input_image_batch=[]
+        count=i
+        image_folder=[]     
+        batch_end = min(i + args.batch_size, len(questions))
+
+             
+        for j in range(i,batch_end):
+            image_file=questions[j]['image_id']+'.png'
+
+            if questions[j]['type']=='ref':
+                qs="[refer] Give me the location of <p> " + qs+" </p>"
+            else:
+                qs="[grounding]" + qs
+            
+            if model.config.mm_use_im_start_end:
+                qs = DEFAULT_IM_START_TOKEN + DEFAULT_IMAGE_TOKEN + DEFAULT_IM_END_TOKEN + '\n' + qs
+            else:
+                qs = DEFAULT_IMAGE_TOKEN + '\n' + qs
+
+            conv = conv_templates[args.conv_mode].copy()
+            conv.append_message(conv.roles[0], qs)
+            conv.append_message(conv.roles[1], None)
+            prompt = conv.get_prompt()
+
+            input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).cuda()
+            input_batch.append(input_ids)
+
+            image = Image.open(os.path.join(args.image_folder, image_file))
+
+            image_folder.append(image)
+
+            stop_str = conv.sep if conv.sep_style != SeparatorStyle.TWO else conv.sep2
+            keywords = [stop_str]
+            stopping_criteria = KeywordsStoppingCriteria(keywords, tokenizer, input_ids)
+
+        max_length = max(tensor.size(1) for tensor in input_batch)
+
+        final_input_list = [torch.cat((torch.zeros((1,max_length - tensor.size(1)), dtype=tensor.dtype,device=tensor.get_device()), tensor),dim=1) for tensor in input_batch]
+        final_input_tensors=torch.cat(final_input_list,dim=0)
+        image_tensor_batch = image_processor.preprocess(image_folder,crop_size ={'height': 504, 'width': 504},size = {'shortest_edge': 504}, return_tensors='pt')['pixel_values']
+
+        with torch.inference_mode():
+            output_ids = model.generate( final_input_tensors, images=image_tensor_batch.half().cuda(), do_sample=False , temperature=args.temperature, top_p=args.top_p, num_beams=1, max_new_tokens=256,length_penalty=2.0, use_cache=True)
+
+        input_token_len = final_input_tensors.shape[1]
+        n_diff_input_output = (final_input_tensors != output_ids[:, :input_token_len]).sum().item()
+        if n_diff_input_output > 0:
+            print(f'[Warning] {n_diff_input_output} output_ids are not the same as the input_ids')
+        outputs = tokenizer.batch_decode(output_ids[:, input_token_len:], skip_special_tokens=True)
+        for k in range(0,len(final_input_list)):
+            output = outputs[k].strip()
+            if output.endswith(stop_str):
+                output = output[:-len(stop_str)]
+            output = output.strip()
+
+            ans_id = shortuuid.uuid()
+            
+            ans_file.write(json.dumps({
+
+                                    "question_id": questions[count]["question_id"],
+                                    "image_id": questions[count]["image_id"],
+                                    "answer": output,
+                                    "ground_truth": questions[count]['ground_truth'],
+                                    "question":questions[count]['question'],
+                                    "type": questions[count]['type'],
+                                    "dataset": questions[count]['dataset'],
+                                    "obj_ids": questions[count]['obj_ids'],
+                                    "size_group": questions[count]['size_group'],                            
+                                    
+                                    }) + "\n")
+            count=count+1
+            ans_file.flush()
+    ans_file.close()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model-path", type=str, default="facebook/opt-350m")
+    parser.add_argument("--model-base", type=str, default=None)
+    parser.add_argument("--image-folder", type=str, default="")
+    parser.add_argument("--question-file", type=str, default="tables/question.jsonl")
+    parser.add_argument("--answers-file", type=str, default="answer.jsonl")
+    parser.add_argument("--conv-mode", type=str, default="llava_v1")
+    parser.add_argument("--num-chunks", type=int, default=1)
+    parser.add_argument("--chunk-idx", type=int, default=0)
+    parser.add_argument("--temperature", type=float, default=0.2)
+    parser.add_argument("--top_p", type=float, default=None)
+    parser.add_argument("--num_beams", type=int, default=1)
+    parser.add_argument("--batch_size",type=int, default=1)
+    args = parser.parse_args()
+
+    eval_model(args)

geochat/eval/batch_geochat_referring.py

@@ -0,0 +1,132 @@
+import argparse
+import torch
+import os
+import json
+from tqdm import tqdm
+import shortuuid
+
+from geochat.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN
+from geochat.conversation import conv_templates, SeparatorStyle
+from geochat.model.builder import load_pretrained_model
+from geochat.utils import disable_torch_init
+from geochat.mm_utils import tokenizer_image_token, get_model_name_from_path, KeywordsStoppingCriteria
+
+from PIL import Image
+import math
+def split_list(lst, n):
+    """Split a list into n (roughly) equal-sized chunks"""
+    chunk_size = math.ceil(len(lst) / n)  # integer division
+    return [lst[i:i+chunk_size] for i in range(0, len(lst), chunk_size)]
+
+
+def get_chunk(lst, n, k):
+    chunks = split_list(lst, n)
+    return chunks[k]
+
+
+def eval_model(args):
+    # Model
+    disable_torch_init()
+    model_path = os.path.expanduser(args.model_path)
+    model_name = get_model_name_from_path(model_path)
+    tokenizer, model, image_processor, context_len = load_pretrained_model(args.model_path, args.model_base, model_name)
+    # print(model)
+    questions=[]
+    questions = [json.loads(q) for q in open(os.path.expanduser(args.question_file), "r")]
+
+        
+    questions = get_chunk(questions, args.num_chunks, args.chunk_idx)
+    answers_file = os.path.expanduser(args.answers_file)
+    os.makedirs(os.path.dirname(answers_file), exist_ok=True)
+    
+    ans_file = open(answers_file, "w")
+    
+    for i in tqdm(range(0,len(questions),args.batch_size)):
+        input_batch=[]
+        input_image_batch=[]
+        count=i
+        image_folder=[]     
+        batch_end = min(i + args.batch_size, len(questions))
+
+             
+        for j in range(i,batch_end):
+            image_file=questions[j]['image_id']+'.png'
+            qs="[identify] What is the object present at " + questions[j]['question']
+            
+            if model.config.mm_use_im_start_end:
+                qs = DEFAULT_IM_START_TOKEN + DEFAULT_IMAGE_TOKEN + DEFAULT_IM_END_TOKEN + '\n' + qs
+            else:
+                qs = DEFAULT_IMAGE_TOKEN + '\n' + qs
+
+            conv = conv_templates[args.conv_mode].copy()
+            conv.append_message(conv.roles[0], qs)
+            conv.append_message(conv.roles[1], None)
+            prompt = conv.get_prompt()
+
+            input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).cuda()
+            input_batch.append(input_ids)
+
+            image = Image.open(os.path.join(args.image_folder, image_file))
+
+            image_folder.append(image)
+
+            stop_str = conv.sep if conv.sep_style != SeparatorStyle.TWO else conv.sep2
+            keywords = [stop_str]
+            stopping_criteria = KeywordsStoppingCriteria(keywords, tokenizer, input_ids)
+
+        max_length = max(tensor.size(1) for tensor in input_batch)
+
+        final_input_list = [torch.cat((torch.zeros((1,max_length - tensor.size(1)), dtype=tensor.dtype,device=tensor.get_device()), tensor),dim=1) for tensor in input_batch]
+        final_input_tensors=torch.cat(final_input_list,dim=0)
+        image_tensor_batch = image_processor.preprocess(image_folder,crop_size ={'height': 504, 'width': 504},size = {'shortest_edge': 504}, return_tensors='pt')['pixel_values']
+
+        with torch.inference_mode():
+            output_ids = model.generate( final_input_tensors, images=image_tensor_batch.half().cuda(), do_sample=False , temperature=args.temperature, top_p=args.top_p, num_beams=1, max_new_tokens=256,length_penalty=2.0, use_cache=True)
+
+        input_token_len = final_input_tensors.shape[1]
+        n_diff_input_output = (final_input_tensors != output_ids[:, :input_token_len]).sum().item()
+        if n_diff_input_output > 0:
+            print(f'[Warning] {n_diff_input_output} output_ids are not the same as the input_ids')
+        outputs = tokenizer.batch_decode(output_ids[:, input_token_len:], skip_special_tokens=True)
+        for k in range(0,len(final_input_list)):
+            output = outputs[k].strip()
+            if output.endswith(stop_str):
+                output = output[:-len(stop_str)]
+            output = output.strip()
+
+            ans_id = shortuuid.uuid()
+            
+            ans_file.write(json.dumps({
+                                    "question_id": questions[count]["question_id"],
+                                    "image_id": questions[count]["image_id"],
+                                    "answer": output,
+                                    "ground_truth": questions[count]['ground_truth'],
+                                    "question":questions[count]['question'],
+                                    "type": questions[count]['type'],
+                                    "dataset": questions[count]['dataset'],
+                                    "obj_ids": questions[count]['obj_ids'],
+                                    "size_group": questions[count]['size_group'],                            
+                                    
+                                    }) + "\n")
+            count=count+1
+            ans_file.flush()
+    ans_file.close()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model-path", type=str, default="facebook/opt-350m")
+    parser.add_argument("--model-base", type=str, default=None)
+    parser.add_argument("--image-folder", type=str, default="")
+    parser.add_argument("--question-file", type=str, default="tables/question.jsonl")
+    parser.add_argument("--answers-file", type=str, default="answer.jsonl")
+    parser.add_argument("--conv-mode", type=str, default="llava_v1")
+    parser.add_argument("--num-chunks", type=int, default=1)
+    parser.add_argument("--chunk-idx", type=int, default=0)
+    parser.add_argument("--temperature", type=float, default=0.2)
+    parser.add_argument("--top_p", type=float, default=None)
+    parser.add_argument("--num_beams", type=int, default=1)
+    parser.add_argument("--batch_size",type=int, default=1)
+    args = parser.parse_args()
+
+    eval_model(args)

geochat/eval/batch_geochat_scene.py

@@ -0,0 +1,139 @@
+import argparse
+import torch
+import os
+import json
+from tqdm import tqdm
+import shortuuid
+
+from geochat.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN
+from geochat.conversation import conv_templates, SeparatorStyle
+from geochat.model.builder import load_pretrained_model
+from geochat.utils import disable_torch_init
+from geochat.mm_utils import tokenizer_image_token, get_model_name_from_path, KeywordsStoppingCriteria
+
+from PIL import Image
+import math
+
+def evaluation_metrics(data_path):
+        
+    base = [json.loads(q) for q in open(data_path, "r")]
+    correct=0
+    incorrect=0
+    for answers in tqdm(base):
+        gt=answers['question_id'].split('/')[0].lower()
+        answer=answers['answer'].replace(' ','').lower().replace('.','')
+        if gt==answer:
+                correct=correct+1
+        else:
+                incorrect=incorrect+1
+        # else:
+        #     continue
+    print('correct:',correct)
+    print('incorrect:',incorrect)
+    print('Total:',correct+incorrect)
+    print('Acc:',(correct/(correct+incorrect)))
+
+
+            
+
+def eval_model(args):
+    # Model
+    disable_torch_init()
+    model_path = os.path.expanduser(args.model_path)
+    model_name = get_model_name_from_path(model_path)
+    tokenizer, model, image_processor, context_len = load_pretrained_model(args.model_path, args.model_base, model_name)
+    # print(model)
+    questions=[]
+    questions = [json.loads(q) for q in open(os.path.expanduser(args.question_file), "r")]
+
+    questions = get_chunk(questions, args.num_chunks, args.chunk_idx)
+    answers_file = os.path.expanduser(args.answers_file)
+    os.makedirs(os.path.dirname(answers_file), exist_ok=True)
+    
+    ans_file = open(answers_file, "w")
+    
+    for i in tqdm(range(0,len(questions),args.batch_size)):
+        input_batch=[]
+        input_image_batch=[]
+        count=i
+        image_folder=[]     
+        batch_end = min(i + args.batch_size, len(questions))
+
+             
+        for j in range(i,batch_end):
+            image_file=questions[j]['image']
+            qs=questions[j]['text']
+            
+            if model.config.mm_use_im_start_end:
+                qs = DEFAULT_IM_START_TOKEN + DEFAULT_IMAGE_TOKEN + DEFAULT_IM_END_TOKEN + '\n' + qs
+            else:
+                qs = DEFAULT_IMAGE_TOKEN + '\n' + qs
+
+            conv = conv_templates[args.conv_mode].copy()
+            conv.append_message(conv.roles[0], qs)
+            conv.append_message(conv.roles[1], None)
+            prompt = conv.get_prompt()
+
+            input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).cuda()
+            input_batch.append(input_ids)
+
+            image = Image.open(os.path.join(args.image_folder, image_file))
+
+            image_folder.append(image)
+
+            stop_str = conv.sep if conv.sep_style != SeparatorStyle.TWO else conv.sep2
+            keywords = [stop_str]
+            stopping_criteria = KeywordsStoppingCriteria(keywords, tokenizer, input_ids)
+
+        max_length = max(tensor.size(1) for tensor in input_batch)
+
+        final_input_list = [torch.cat((torch.zeros((1,max_length - tensor.size(1)), dtype=tensor.dtype,device=tensor.get_device()), tensor),dim=1) for tensor in input_batch]
+        final_input_tensors=torch.cat(final_input_list,dim=0)
+        image_tensor_batch = image_processor.preprocess(image_folder,crop_size ={'height': 504, 'width': 504},size = {'shortest_edge': 504}, return_tensors='pt')['pixel_values']
+
+        with torch.inference_mode():
+            output_ids = model.generate( final_input_tensors, images=image_tensor_batch.half().cuda(), do_sample=False , temperature=args.temperature, top_p=args.top_p, num_beams=1, max_new_tokens=256,length_penalty=2.0, use_cache=True)
+
+        input_token_len = final_input_tensors.shape[1]
+        n_diff_input_output = (final_input_tensors != output_ids[:, :input_token_len]).sum().item()
+        if n_diff_input_output > 0:
+            print(f'[Warning] {n_diff_input_output} output_ids are not the same as the input_ids')
+        outputs = tokenizer.batch_decode(output_ids[:, input_token_len:], skip_special_tokens=True)
+        for k in range(0,len(final_input_list)):
+            output = outputs[k].strip()
+            if output.endswith(stop_str):
+                output = output[:-len(stop_str)]
+            output = output.strip()
+
+            ans_id = shortuuid.uuid()
+            
+            ans_file.write(json.dumps({
+
+                                    "question_id": questions[count]["question_id"],
+                                    "image_id": questions[count]["image"],
+                                    "answer": output,
+                                    "ground_truth": questions[count]['ground_truth']
+                                    }) + "\n")
+            count=count+1
+            ans_file.flush()
+    ans_file.close()
+    evaluation_metrics(answers_file)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model-path", type=str, default="facebook/opt-350m")
+    parser.add_argument("--model-base", type=str, default=None)
+    parser.add_argument("--image-folder", type=str, default="")
+    parser.add_argument("--question-file", type=str, default="tables/question.jsonl")
+    parser.add_argument("--answers-file", type=str, default="answer.jsonl")
+    parser.add_argument("--conv-mode", type=str, default="llava_v1")
+    parser.add_argument("--num-chunks", type=int, default=1)
+    parser.add_argument("--chunk-idx", type=int, default=0)
+    parser.add_argument("--temperature", type=float, default=0.2)
+    parser.add_argument("--top_p", type=float, default=None)
+    parser.add_argument("--num_beams", type=int, default=1)
+    parser.add_argument("--batch_size",type=int, default=1)
+    args = parser.parse_args()
+
+    eval_model(args)

geochat/eval/batch_geochat_vqa.py

@@ -0,0 +1,125 @@
+import argparse
+import torch
+import os
+import json
+from tqdm import tqdm
+import shortuuid
+
+from geochat.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN
+from geochat.conversation import conv_templates, SeparatorStyle
+from geochat.model.builder import load_pretrained_model
+from geochat.utils import disable_torch_init
+from geochat.mm_utils import tokenizer_image_token, get_model_name_from_path, KeywordsStoppingCriteria
+
+from PIL import Image
+import math
+def split_list(lst, n):
+    """Split a list into n (roughly) equal-sized chunks"""
+    chunk_size = math.ceil(len(lst) / n)  # integer division
+    return [lst[i:i+chunk_size] for i in range(0, len(lst), chunk_size)]
+
+
+def get_chunk(lst, n, k):
+    chunks = split_list(lst, n)
+    return chunks[k]
+
+
+def eval_model(args):
+    # Model
+    disable_torch_init()
+    model_path = os.path.expanduser(args.model_path)
+    model_name = get_model_name_from_path(model_path)
+    tokenizer, model, image_processor, context_len = load_pretrained_model(args.model_path, args.model_base, model_name)
+    
+    questions=[]
+    questions = [json.loads(q) for q in open(os.path.expanduser(args.question_file), "r")]
+
+        
+    questions = get_chunk(questions, args.num_chunks, args.chunk_idx)
+    answers_file = os.path.expanduser(args.answers_file)
+    os.makedirs(os.path.dirname(answers_file), exist_ok=True)
+    
+    ans_file = open(answers_file, "w")
+    
+    for i in tqdm(range(0,len(questions),args.batch_size)):
+        input_batch=[]
+        input_image_batch=[]
+        count=i
+        image_folder=[]     
+        batch_end = min(i + args.batch_size, len(questions))
+
+             
+        for j in range(i,batch_end):
+            image_file=questions[j]['image']
+            qs=questions[j]['text']
+            
+            if model.config.mm_use_im_start_end:
+                qs = DEFAULT_IM_START_TOKEN + DEFAULT_IMAGE_TOKEN + DEFAULT_IM_END_TOKEN + '\n' + qs
+            else:
+                qs = DEFAULT_IMAGE_TOKEN + '\n' + qs
+
+            conv = conv_templates[args.conv_mode].copy()
+            conv.append_message(conv.roles[0], qs)
+            conv.append_message(conv.roles[1], None)
+            prompt = conv.get_prompt()
+
+            input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).cuda()
+            input_batch.append(input_ids)
+
+            image = Image.open(os.path.join(args.image_folder, image_file))
+
+            image_folder.append(image)
+
+            stop_str = conv.sep if conv.sep_style != SeparatorStyle.TWO else conv.sep2
+            keywords = [stop_str]
+            stopping_criteria = KeywordsStoppingCriteria(keywords, tokenizer, input_ids)
+
+        max_length = max(tensor.size(1) for tensor in input_batch)
+
+        final_input_list = [torch.cat((torch.zeros((1,max_length - tensor.size(1)), dtype=tensor.dtype,device=tensor.get_device()), tensor),dim=1) for tensor in input_batch]
+        final_input_tensors=torch.cat(final_input_list,dim=0)
+        image_tensor_batch = image_processor.preprocess(image_folder,crop_size ={'height': 504, 'width': 504},size = {'shortest_edge': 504}, return_tensors='pt')['pixel_values']
+
+        with torch.inference_mode():
+            output_ids = model.generate( final_input_tensors, images=image_tensor_batch.half().cuda(), do_sample=False , temperature=args.temperature, top_p=args.top_p, num_beams=1, max_new_tokens=256,length_penalty=2.0, use_cache=True)
+
+        input_token_len = final_input_tensors.shape[1]
+        n_diff_input_output = (final_input_tensors != output_ids[:, :input_token_len]).sum().item()
+        if n_diff_input_output > 0:
+            print(f'[Warning] {n_diff_input_output} output_ids are not the same as the input_ids')
+        outputs = tokenizer.batch_decode(output_ids[:, input_token_len:], skip_special_tokens=True)
+        for k in range(0,len(final_input_list)):
+            output = outputs[k].strip()
+            if output.endswith(stop_str):
+                output = output[:-len(stop_str)]
+            output = output.strip()
+
+            ans_id = shortuuid.uuid()
+            
+            ans_file.write(json.dumps({
+                                    "question_id": questions[count]["question_id"],
+                                    "image_id": questions[count]["image"],
+                                    "answer": output,
+                                    }) + "\n")
+            count=count+1
+            ans_file.flush()
+    ans_file.close()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model-path", type=str, default="facebook/opt-350m")
+    parser.add_argument("--model-base", type=str, default=None)
+    parser.add_argument("--image-folder", type=str, default="")
+    parser.add_argument("--question-file", type=str, default="tables/question.jsonl")
+    parser.add_argument("--answers-file", type=str, default="answer.jsonl")
+    parser.add_argument("--conv-mode", type=str, default="llava_v1")
+    parser.add_argument("--num-chunks", type=int, default=1)
+    parser.add_argument("--chunk-idx", type=int, default=0)
+    parser.add_argument("--temperature", type=float, default=0.2)
+    parser.add_argument("--top_p", type=float, default=None)
+    parser.add_argument("--num_beams", type=int, default=1)
+    parser.add_argument("--batch_size",type=int, default=1)
+    args = parser.parse_args()
+
+    eval_model(args)

geochat/mm_utils.py

@@ -0,0 +1,121 @@
+from PIL import Image
+from io import BytesIO
+import base64
+
+import torch
+from transformers import StoppingCriteria
+from geochat.constants import IMAGE_TOKEN_INDEX
+import numpy as np
+
+def load_image_from_base64(image):
+    return Image.open(BytesIO(base64.b64decode(image)))
+
+
+def expand2square(pil_img, background_color):
+    width, height = pil_img.size
+    if width == height:
+        return pil_img
+    elif width > height:
+        result = Image.new(pil_img.mode, (width, width), background_color)
+        result.paste(pil_img, (0, (width - height) // 2))
+        return result
+    else:
+        result = Image.new(pil_img.mode, (height, height), background_color)
+        result.paste(pil_img, ((height - width) // 2, 0))
+        return result
+
+
+def process_images(images, image_processor, model_cfg):
+    image_aspect_ratio = getattr(model_cfg, "image_aspect_ratio", None)
+    new_images = []
+    if image_aspect_ratio == 'pad':
+        for image in images:
+            image = expand2square(image, tuple(int(x*255) for x in image_processor.image_mean))
+            image = image_processor.preprocess(image,crop_size ={'height': 504, 'width': 504},size = {'shortest_edge': 504},return_tensors='pt')['pixel_values'][0]
+            # image = image_processor.preprocess(image,return_tensors='pt')['pixel_values'][0]
+
+            new_images.append(image)
+    else:
+        return image_processor(images, return_tensors='pt')['pixel_values']
+    if all(x.shape == new_images[0].shape for x in new_images):
+        new_images = torch.stack(new_images, dim=0)
+    return new_images
+
+def process_images_demo(images, image_processor):
+    new_images = []
+    # image_aspect_ratio = 'pad'
+    for image in images:
+        image = expand2square(image, tuple(int(x*255) for x in image_processor.image_mean))
+        image = image_processor.preprocess(image,crop_size ={'height': 504, 'width': 504},size = {'shortest_edge': 504},return_tensors='pt')['pixel_values'][0]
+        # image = image_processor.preprocess(image,return_tensors='pt')['pixel_values'][0]
+
+        new_images.append(image)
+
+    if all(x.shape == new_images[0].shape for x in new_images):
+        new_images = torch.stack(new_images, dim=0)
+    return new_images
+
+def tokenizer_image_token(prompt, tokenizer, image_token_index=IMAGE_TOKEN_INDEX, return_tensors=None):
+    prompt_chunks = [tokenizer(chunk).input_ids for chunk in prompt.split('<image>')]
+
+    def insert_separator(X, sep):
+        return [ele for sublist in zip(X, [sep]*len(X)) for ele in sublist][:-1]
+
+    input_ids = []
+    offset = 0
+    if len(prompt_chunks) > 0 and len(prompt_chunks[0]) > 0 and prompt_chunks[0][0] == tokenizer.bos_token_id:
+        offset = 1
+        input_ids.append(prompt_chunks[0][0])
+
+    for x in insert_separator(prompt_chunks, [image_token_index] * (offset + 1)):
+        input_ids.extend(x[offset:])
+
+    if return_tensors is not None:
+        if return_tensors == 'pt':
+            return torch.tensor(input_ids, dtype=torch.long)
+        raise ValueError(f'Unsupported tensor type: {return_tensors}')
+    return input_ids
+
+
+def get_model_name_from_path(model_path):
+    model_path = model_path.strip("/")
+    model_paths = model_path.split("/")
+    if model_paths[-1].startswith('checkpoint-'):
+        return model_paths[-2] + "_" + model_paths[-1]
+    else:
+        return model_paths[-1]
+
+
+
+
+class KeywordsStoppingCriteria(StoppingCriteria):
+    def __init__(self, keywords, tokenizer, input_ids):
+        self.keywords = keywords
+        self.keyword_ids = []
+        self.max_keyword_len = 0
+        for keyword in keywords:
+            cur_keyword_ids = tokenizer(keyword).input_ids
+            if len(cur_keyword_ids) > 1 and cur_keyword_ids[0] == tokenizer.bos_token_id:
+                cur_keyword_ids = cur_keyword_ids[1:]
+            if len(cur_keyword_ids) > self.max_keyword_len:
+                self.max_keyword_len = len(cur_keyword_ids)
+            self.keyword_ids.append(torch.tensor(cur_keyword_ids))
+        self.tokenizer = tokenizer
+        self.start_len = input_ids.shape[1]
+
+    def __call__(self, output_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
+        # assert output_ids.shape[0] == 1, "Only support batch size 1 (yet)"  # TODO
+        offset = min(output_ids.shape[1] - self.start_len, self.max_keyword_len)
+        self.keyword_ids = [keyword_id.to(output_ids.device) for keyword_id in self.keyword_ids]
+        for keyword_id in self.keyword_ids:
+            if (output_ids[0, -keyword_id.shape[0]:] == keyword_id).all():
+                return True
+        outputs = self.tokenizer.batch_decode(output_ids[:, -offset:], skip_special_tokens=True)[0]
+        flag=False
+        for output in outputs:
+                
+            for keyword in self.keywords:
+                if keyword in output:
+                    flag=True
+                    return flag
+        return flag

geochat/model/.ipynb_checkpoints/__init__-checkpoint.py

@@ -0,0 +1,2 @@
+from .language_model.geochat_llama import GeoChatLlamaForCausalLM, GeoChatConfig
+from .language_model.geochat_mpt import GeoChatMPTForCausalLM, GeoChatMPTConfig

geochat/model/.ipynb_checkpoints/builder-checkpoint.py

@@ -0,0 +1,149 @@
+#    Copyright 2023 Haotian Liu
+#
+#    Licensed under the Apache License, Version 2.0 (the "License");
+#    you may not use this file except in compliance with the License.
+#    You may obtain a copy of the License at
+#
+#        http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS,
+#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#    See the License for the specific language governing permissions and
+#    limitations under the License.
+
+
+import os
+import warnings
+import shutil
+
+from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig, BitsAndBytesConfig
+import torch
+from geochat.model import *
+from geochat.constants import DEFAULT_IMAGE_PATCH_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN
+
+
+def load_pretrained_model(model_path, model_base, model_name, load_8bit=False, load_4bit=False, device_map="auto", device="cuda"):
+    kwargs = {"device_map": device_map}
+
+    if load_8bit:
+        kwargs['load_in_8bit'] = True
+    elif load_4bit:
+        kwargs['load_in_4bit'] = True
+        kwargs['quantization_config'] = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_compute_dtype=torch.float16,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type='nf4'
+        )
+    else:
+        kwargs['torch_dtype'] = torch.float16
+
+    if 'geochat' in model_name.lower():
+        # Load LLaVA model
+        if 'lora' in model_name.lower() and model_base is None:
+            warnings.warn('There is `lora` in model name but no `model_base` is provided. If you are loading a LoRA model, please provide the `model_base` argument. Detailed instruction: https://github.com/haotian-liu/LLaVA#launch-a-model-worker-lora-weights-unmerged.')
+        if 'lora' in model_name.lower() and model_base is not None:
+            lora_cfg_pretrained = AutoConfig.from_pretrained(model_path)
+            tokenizer = AutoTokenizer.from_pretrained(model_base, use_fast=False)
+            print('Loading Geochat from base model...')
+            model = GeoChatLlamaForCausalLM.from_pretrained(model_base, low_cpu_mem_usage=True, config=lora_cfg_pretrained, **kwargs)
+            token_num, tokem_dim = model.lm_head.out_features, model.lm_head.in_features
+            if model.lm_head.weight.shape[0] != token_num:
+                model.lm_head.weight = torch.nn.Parameter(torch.empty(token_num, tokem_dim, device=model.device, dtype=model.dtype))
+                model.model.embed_tokens.weight = torch.nn.Parameter(torch.empty(token_num, tokem_dim, device=model.device, dtype=model.dtype))
+
+            print('Loading additional GeoChat weights...')
+            if os.path.exists(os.path.join(model_path, 'non_lora_trainables.bin')):
+                non_lora_trainables = torch.load(os.path.join(model_path, 'non_lora_trainables.bin'), map_location='cpu')
+            else:
+                # this is probably from HF Hub
+                from huggingface_hub import hf_hub_download
+                def load_from_hf(repo_id, filename, subfolder=None):
+                    cache_file = hf_hub_download(
+                        repo_id=repo_id,
+                        filename=filename,
+                        subfolder=subfolder)
+                    return torch.load(cache_file, map_location='cpu')
+                non_lora_trainables = load_from_hf(model_path, 'non_lora_trainables.bin')
+            non_lora_trainables = {(k[11:] if k.startswith('base_model.') else k): v for k, v in non_lora_trainables.items()}
+            if any(k.startswith('model.model.') for k in non_lora_trainables):
+                non_lora_trainables = {(k[6:] if k.startswith('model.') else k): v for k, v in non_lora_trainables.items()}
+            model.load_state_dict(non_lora_trainables, strict=False)
+
+            from peft import PeftModel
+            print('Loading LoRA weights...')
+            model = PeftModel.from_pretrained(model, model_path)
+            print('Merging LoRA weights...')
+            model = model.merge_and_unload()
+            print('Model is loaded...')
+        elif model_base is not None:
+            # this may be mm projector only
+            print('Loading GeoChat from base model...')
+            if 'mpt' in model_name.lower():
+                if not os.path.isfile(os.path.join(model_path, 'configuration_mpt.py')):
+                    shutil.copyfile(os.path.join(model_base, 'configuration_mpt.py'), os.path.join(model_path, 'configuration_mpt.py'))
+                tokenizer = AutoTokenizer.from_pretrained(model_base, use_fast=True)
+                cfg_pretrained = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
+                model = GeoChatMPTForCausalLM.from_pretrained(model_base, low_cpu_mem_usage=True, config=cfg_pretrained, **kwargs)
+            else:
+                tokenizer = AutoTokenizer.from_pretrained(model_base, use_fast=False)
+                cfg_pretrained = AutoConfig.from_pretrained(model_path)
+                model = GeoChatLlamaForCausalLM.from_pretrained(model_base, low_cpu_mem_usage=True, config=cfg_pretrained, **kwargs)
+
+            mm_projector_weights = torch.load(os.path.join(model_path, 'mm_projector.bin'), map_location='cpu')
+            mm_projector_weights = {k: v.to(torch.float16) for k, v in mm_projector_weights.items()}
+            model.load_state_dict(mm_projector_weights, strict=False)
+        else:
+            if 'mpt' in model_name.lower():
+                tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=True)
+                model = GeoChatMPTForCausalLM.from_pretrained(model_path, low_cpu_mem_usage=True, **kwargs)
+            else:
+                print("Loading GeoChat......")
+                tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)
+                model = GeoChatLlamaForCausalLM.from_pretrained(model_path, low_cpu_mem_usage=True, **kwargs)
+    else:
+        # Load language model
+        if model_base is not None:
+            # PEFT model
+            from peft import PeftModel
+            tokenizer = AutoTokenizer.from_pretrained(model_base, use_fast=False)
+            model = AutoModelForCausalLM.from_pretrained(model_base, torch_dtype=torch.float16, low_cpu_mem_usage=True, device_map="auto")
+            print(f"Loading LoRA weights from {model_path}")
+            model = PeftModel.from_pretrained(model, model_path)
+            print(f"Merging weights")
+            model = model.merge_and_unload()
+            print('Convert to FP16...')
+            model.to(torch.float16)
+        else:
+            use_fast = False
+            if 'mpt' in model_name.lower():
+                tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=True)
+                model = AutoModelForCausalLM.from_pretrained(model_path, low_cpu_mem_usage=True, trust_remote_code=True, **kwargs)
+            else:
+                tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)
+                model = AutoModelForCausalLM.from_pretrained(model_path, low_cpu_mem_usage=True, **kwargs)
+
+    image_processor = None
+
+    if 'geochat' in model_name.lower():
+        mm_use_im_start_end = getattr(model.config, "mm_use_im_start_end", False)
+        mm_use_im_patch_token = getattr(model.config, "mm_use_im_patch_token", True)
+        if mm_use_im_patch_token:
+            tokenizer.add_tokens([DEFAULT_IMAGE_PATCH_TOKEN], special_tokens=True)
+        if mm_use_im_start_end:
+            tokenizer.add_tokens([DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN], special_tokens=True)
+        model.resize_token_embeddings(len(tokenizer))
+
+        vision_tower = model.get_vision_tower()
+        if not vision_tower.is_loaded:
+            vision_tower.load_model()
+        vision_tower.to(device=device, dtype=torch.float16)
+        image_processor = vision_tower.image_processor
+
+    if hasattr(model.config, "max_sequence_length"):
+        context_len = model.config.max_sequence_length
+    else:
+        context_len = 2048
+
+    return tokenizer, model, image_processor, context_len

geochat/model/__init__.py

@@ -0,0 +1,2 @@
+from .language_model.geochat_llama import GeoChatLlamaForCausalLM, GeoChatConfig
+from .language_model.geochat_mpt import GeoChatMPTForCausalLM, GeoChatMPTConfig

geochat/model/apply_delta.py

@@ -0,0 +1,48 @@
+"""
+Usage:
+python3 -m fastchat.model.apply_delta --base ~/model_weights/llama-7b --target ~/model_weights/vicuna-7b --delta lmsys/vicuna-7b-delta
+"""
+import argparse
+
+import torch
+from tqdm import tqdm
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from geochat import GeoChatLlamaForCausalLM
+
+
+def apply_delta(base_model_path, target_model_path, delta_path):
+    print("Loading base model")
+    base = AutoModelForCausalLM.from_pretrained(
+        base_model_path, torch_dtype=torch.float16, low_cpu_mem_usage=True)
+
+    print("Loading delta")
+    delta = LlavaLlamaForCausalLM.from_pretrained(delta_path, torch_dtype=torch.float16, low_cpu_mem_usage=True)
+    delta_tokenizer = AutoTokenizer.from_pretrained(delta_path)
+
+    print("Applying delta")
+    for name, param in tqdm(delta.state_dict().items(), desc="Applying delta"):
+        if name not in base.state_dict():
+            assert name in ['model.mm_projector.weight', 'model.mm_projector.bias'], f'{name} not in base model'
+            continue
+        if param.data.shape == base.state_dict()[name].shape:
+            param.data += base.state_dict()[name]
+        else:
+            assert name in ['model.embed_tokens.weight', 'lm_head.weight'], \
+                f'{name} dimension mismatch: {param.data.shape} vs {base.state_dict()[name].shape}'
+            bparam = base.state_dict()[name]
+            param.data[:bparam.shape[0], :bparam.shape[1]] += bparam
+
+    print("Saving target model")
+    delta.save_pretrained(target_model_path)
+    delta_tokenizer.save_pretrained(target_model_path)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--base-model-path", type=str, required=True)
+    parser.add_argument("--target-model-path", type=str, required=True)
+    parser.add_argument("--delta-path", type=str, required=True)
+
+    args = parser.parse_args()
+
+    apply_delta(args.base_model_path, args.target_model_path, args.delta_path)

geochat/model/builder.py

@@ -0,0 +1,149 @@
+#    Copyright 2023 Haotian Liu
+#
+#    Licensed under the Apache License, Version 2.0 (the "License");
+#    you may not use this file except in compliance with the License.
+#    You may obtain a copy of the License at
+#
+#        http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS,
+#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#    See the License for the specific language governing permissions and
+#    limitations under the License.
+
+
+import os
+import warnings
+import shutil
+
+from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig, BitsAndBytesConfig
+import torch
+from geochat.model import *
+from geochat.constants import DEFAULT_IMAGE_PATCH_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN
+
+
+def load_pretrained_model(model_path, model_base, model_name, load_8bit=False, load_4bit=False, device_map="auto", device="cuda"):
+    kwargs = {"device_map": device_map}
+
+    if load_8bit:
+        kwargs['load_in_8bit'] = True
+    elif load_4bit:
+        kwargs['load_in_4bit'] = True
+        kwargs['quantization_config'] = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_compute_dtype=torch.float16,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type='nf4'
+        )
+    else:
+        kwargs['torch_dtype'] = torch.float16
+
+    if 'geochat' in model_name.lower():
+        # Load LLaVA model
+        if 'lora' in model_name.lower() and model_base is None:
+            warnings.warn('There is `lora` in model name but no `model_base` is provided. If you are loading a LoRA model, please provide the `model_base` argument. Detailed instruction: https://github.com/haotian-liu/LLaVA#launch-a-model-worker-lora-weights-unmerged.')
+        if 'lora' in model_name.lower() and model_base is not None:
+            lora_cfg_pretrained = AutoConfig.from_pretrained(model_path)
+            tokenizer = AutoTokenizer.from_pretrained(model_base, use_fast=False)
+            print('Loading Geochat from base model...')
+            model = GeoChatLlamaForCausalLM.from_pretrained(model_base, low_cpu_mem_usage=True, config=lora_cfg_pretrained, **kwargs)
+            token_num, tokem_dim = model.lm_head.out_features, model.lm_head.in_features
+            if model.lm_head.weight.shape[0] != token_num:
+                model.lm_head.weight = torch.nn.Parameter(torch.empty(token_num, tokem_dim, device=model.device, dtype=model.dtype))
+                model.model.embed_tokens.weight = torch.nn.Parameter(torch.empty(token_num, tokem_dim, device=model.device, dtype=model.dtype))
+
+            print('Loading additional GeoChat weights...')
+            if os.path.exists(os.path.join(model_path, 'non_lora_trainables.bin')):
+                non_lora_trainables = torch.load(os.path.join(model_path, 'non_lora_trainables.bin'), map_location='cpu')
+            else:
+                # this is probably from HF Hub
+                from huggingface_hub import hf_hub_download
+                def load_from_hf(repo_id, filename, subfolder=None):
+                    cache_file = hf_hub_download(
+                        repo_id=repo_id,
+                        filename=filename,
+                        subfolder=subfolder)
+                    return torch.load(cache_file, map_location='cpu')
+                non_lora_trainables = load_from_hf(model_path, 'non_lora_trainables.bin')
+            non_lora_trainables = {(k[11:] if k.startswith('base_model.') else k): v for k, v in non_lora_trainables.items()}
+            if any(k.startswith('model.model.') for k in non_lora_trainables):
+                non_lora_trainables = {(k[6:] if k.startswith('model.') else k): v for k, v in non_lora_trainables.items()}
+            model.load_state_dict(non_lora_trainables, strict=False)
+
+            from peft import PeftModel
+            print('Loading LoRA weights...')
+            model = PeftModel.from_pretrained(model, model_path)
+            print('Merging LoRA weights...')
+            model = model.merge_and_unload()
+            print('Model is loaded...')
+        elif model_base is not None:
+            # this may be mm projector only
+            print('Loading GeoChat from base model...')
+            if 'mpt' in model_name.lower():
+                if not os.path.isfile(os.path.join(model_path, 'configuration_mpt.py')):
+                    shutil.copyfile(os.path.join(model_base, 'configuration_mpt.py'), os.path.join(model_path, 'configuration_mpt.py'))
+                tokenizer = AutoTokenizer.from_pretrained(model_base, use_fast=True)
+                cfg_pretrained = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
+                model = GeoChatMPTForCausalLM.from_pretrained(model_base, low_cpu_mem_usage=True, config=cfg_pretrained, **kwargs)
+            else:
+                tokenizer = AutoTokenizer.from_pretrained(model_base, use_fast=False)
+                cfg_pretrained = AutoConfig.from_pretrained(model_path)
+                model = GeoChatLlamaForCausalLM.from_pretrained(model_base, low_cpu_mem_usage=True, config=cfg_pretrained, **kwargs)
+
+            mm_projector_weights = torch.load(os.path.join(model_path, 'mm_projector.bin'), map_location='cpu')
+            mm_projector_weights = {k: v.to(torch.float16) for k, v in mm_projector_weights.items()}
+            model.load_state_dict(mm_projector_weights, strict=False)
+        else:
+            if 'mpt' in model_name.lower():
+                tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=True)
+                model = GeoChatMPTForCausalLM.from_pretrained(model_path, low_cpu_mem_usage=True, **kwargs)
+            else:
+                print("Loading GeoChat......")
+                tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)
+                model = GeoChatLlamaForCausalLM.from_pretrained(model_path, low_cpu_mem_usage=True, **kwargs)
+    else:
+        # Load language model
+        if model_base is not None:
+            # PEFT model
+            from peft import PeftModel
+            tokenizer = AutoTokenizer.from_pretrained(model_base, use_fast=False)
+            model = AutoModelForCausalLM.from_pretrained(model_base, torch_dtype=torch.float16, low_cpu_mem_usage=True, device_map="auto")
+            print(f"Loading LoRA weights from {model_path}")
+            model = PeftModel.from_pretrained(model, model_path)
+            print(f"Merging weights")
+            model = model.merge_and_unload()
+            print('Convert to FP16...')
+            model.to(torch.float16)
+        else:
+            use_fast = False
+            if 'mpt' in model_name.lower():
+                tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=True)
+                model = AutoModelForCausalLM.from_pretrained(model_path, low_cpu_mem_usage=True, trust_remote_code=True, **kwargs)
+            else:
+                tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)
+                model = AutoModelForCausalLM.from_pretrained(model_path, low_cpu_mem_usage=True, **kwargs)
+
+    image_processor = None
+
+    if 'geochat' in model_name.lower():
+        mm_use_im_start_end = getattr(model.config, "mm_use_im_start_end", False)
+        mm_use_im_patch_token = getattr(model.config, "mm_use_im_patch_token", True)
+        if mm_use_im_patch_token:
+            tokenizer.add_tokens([DEFAULT_IMAGE_PATCH_TOKEN], special_tokens=True)
+        if mm_use_im_start_end:
+            tokenizer.add_tokens([DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN], special_tokens=True)
+        model.resize_token_embeddings(len(tokenizer))
+
+        vision_tower = model.get_vision_tower()
+        if not vision_tower.is_loaded:
+            vision_tower.load_model()
+        vision_tower.to(device=device, dtype=torch.float16)
+        image_processor = vision_tower.image_processor
+
+    if hasattr(model.config, "max_sequence_length"):
+        context_len = model.config.max_sequence_length
+    else:
+        context_len = 2048
+
+    return tokenizer, model, image_processor, context_len

geochat/model/consolidate.py

@@ -0,0 +1,29 @@
+"""
+Usage:
+python3 -m llava.model.consolidate --src ~/model_weights/llava-7b --dst ~/model_weights/llava-7b_consolidate
+"""
+import argparse
+
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from geochat.model import *
+from geochat.model.utils import auto_upgrade
+
+
+def consolidate_ckpt(src_path, dst_path):
+    print("Loading model")
+    auto_upgrade(src_path)
+    src_model = AutoModelForCausalLM.from_pretrained(src_path, torch_dtype=torch.float16, low_cpu_mem_usage=True)
+    src_tokenizer = AutoTokenizer.from_pretrained(src_path, use_fast=False)
+    src_model.save_pretrained(dst_path)
+    src_tokenizer.save_pretrained(dst_path)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--src", type=str, required=True)
+    parser.add_argument("--dst", type=str, required=True)
+
+    args = parser.parse_args()
+
+    consolidate_ckpt(args.src, args.dst)

geochat/model/geochat_arch.py

@@ -0,0 +1,262 @@
+#    Copyright 2023 Haotian Liu
+#
+#    Licensed under the Apache License, Version 2.0 (the "License");
+#    you may not use this file except in compliance with the License.
+#    You may obtain a copy of the License at
+#
+#        http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS,
+#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#    See the License for the specific language governing permissions and
+#    limitations under the License.
+
+
+from abc import ABC, abstractmethod
+
+import torch
+import torch.nn as nn
+
+from .multimodal_encoder.builder import build_vision_tower
+from .multimodal_projector.builder import build_vision_projector
+
+from geochat.constants import IGNORE_INDEX, IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_PATCH_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN
+
+
+class GeoChatMetaModel:
+
+    def __init__(self, config):
+        super(GeoChatMetaModel, self).__init__(config)
+
+        if hasattr(config, "mm_vision_tower"):
+            self.vision_tower = build_vision_tower(config, delay_load=True)
+            self.mm_projector = build_vision_projector(config)
+
+    def get_vision_tower(self):
+        vision_tower = getattr(self, 'vision_tower', None)
+        if type(vision_tower) is list:
+            vision_tower = vision_tower[0]
+        return vision_tower
+
+    def initialize_vision_modules(self, model_args, fsdp=None):
+        vision_tower = model_args.vision_tower
+        mm_vision_select_layer = model_args.mm_vision_select_layer
+        mm_vision_select_feature = model_args.mm_vision_select_feature
+        pretrain_mm_mlp_adapter = model_args.pretrain_mm_mlp_adapter
+
+        self.config.mm_vision_tower = vision_tower
+
+        if self.get_vision_tower() is None:
+            vision_tower = build_vision_tower(model_args)
+
+            if fsdp is not None and len(fsdp) > 0:
+                self.vision_tower = [vision_tower]
+            else:
+                self.vision_tower = vision_tower
+        else:
+            if fsdp is not None and len(fsdp) > 0:
+                vision_tower = self.vision_tower[0]
+            else:
+                vision_tower = self.vision_tower
+            vision_tower.load_model()
+
+        self.config.use_mm_proj = True
+        self.config.mm_projector_type = getattr(model_args, 'mm_projector_type', 'linear')
+        self.config.mm_hidden_size = vision_tower.hidden_size
+        self.config.mm_vision_select_layer = mm_vision_select_layer
+        self.config.mm_vision_select_feature = mm_vision_select_feature
+
+        if getattr(self, 'mm_projector', None) is None:
+            self.mm_projector = build_vision_projector(self.config)
+            # print(mm_projector)
+            
+
+        if pretrain_mm_mlp_adapter is not None:
+            mm_projector_weights = torch.load(pretrain_mm_mlp_adapter, map_location='cpu')
+
+            def get_w(weights, keyword):
+                return {k.split(keyword + '.')[1]: v for k, v in weights.items() if keyword in k}
+
+            self.mm_projector.load_state_dict(get_w(mm_projector_weights, 'mm_projector'))
+            
+        
+
+
+class GeoChatMetaForCausalLM(ABC):
+
+    @abstractmethod
+    def get_model(self):
+        pass
+
+    def get_vision_tower(self):
+        return self.get_model().get_vision_tower()
+
+    def encode_images(self, images):
+        image_features = self.get_model().get_vision_tower()(images)
+        image_features = self.get_model().mm_projector(image_features)
+        return image_features
+
+    def prepare_inputs_labels_for_multimodal(
+        self, input_ids, attention_mask, past_key_values, labels, images
+    ):
+        vision_tower = self.get_vision_tower()
+        if vision_tower is None or images is None or input_ids.shape[1] == 1:
+            if past_key_values is not None and vision_tower is not None and images is not None and input_ids.shape[1] == 1:
+                attention_mask = torch.ones((attention_mask.shape[0], past_key_values[-1][-1].shape[-2] + 1), dtype=attention_mask.dtype, device=attention_mask.device)
+            return input_ids, attention_mask, past_key_values, None, labels
+
+        if type(images) is list or images.ndim == 5:
+            concat_images = torch.cat([image for image in images], dim=0)
+            image_features = self.encode_images(concat_images)
+            split_sizes = [image.shape[0] for image in images]
+            image_features = torch.split(image_features, split_sizes, dim=0)
+            image_features = [x.flatten(0, 1) for x in image_features]
+        else:
+            image_features = self.encode_images(images)
+
+        new_input_embeds = []
+        new_labels = [] if labels is not None else None
+        cur_image_idx = 0
+        for batch_idx, cur_input_ids in enumerate(input_ids):
+            if (cur_input_ids == IMAGE_TOKEN_INDEX).sum() == 0:
+                # multimodal LLM, but the current sample is not multimodal
+                # FIXME: this is a hacky fix, for deepspeed zero3 to work
+                half_len = cur_input_ids.shape[0] // 2
+                cur_image_features = image_features[cur_image_idx]
+                cur_input_embeds_1 = self.get_model().embed_tokens(cur_input_ids[:half_len])
+                cur_input_embeds_2 = self.get_model().embed_tokens(cur_input_ids[half_len:])
+                cur_input_embeds = torch.cat([cur_input_embeds_1, cur_image_features[0:0], cur_input_embeds_2], dim=0)
+                new_input_embeds.append(cur_input_embeds)
+                if labels is not None:
+                    new_labels.append(labels[batch_idx])
+                cur_image_idx += 1
+                continue
+            image_token_indices = torch.where(cur_input_ids == IMAGE_TOKEN_INDEX)[0]
+            cur_new_input_embeds = []
+            if labels is not None:
+                cur_labels = labels[batch_idx]
+                cur_new_labels = []
+                assert cur_labels.shape == cur_input_ids.shape
+            while image_token_indices.numel() > 0:
+                cur_image_features = image_features[cur_image_idx]
+                image_token_start = image_token_indices[0]
+                if getattr(self.config, 'tune_mm_mlp_adapter', False) and getattr(self.config, 'mm_use_im_start_end', False):
+                    cur_new_input_embeds.append(self.get_model().embed_tokens(cur_input_ids[:image_token_start-1]).detach())
+                    cur_new_input_embeds.append(self.get_model().embed_tokens(cur_input_ids[image_token_start-1:image_token_start]))
+                    cur_new_input_embeds.append(cur_image_features)
+                    cur_new_input_embeds.append(self.get_model().embed_tokens(cur_input_ids[image_token_start+1:image_token_start+2]))
+                    if labels is not None:
+                        cur_new_labels.append(cur_labels[:image_token_start])
+                        cur_new_labels.append(torch.full((cur_image_features.shape[0],), IGNORE_INDEX, device=labels.device, dtype=labels.dtype))
+                        cur_new_labels.append(cur_labels[image_token_start:image_token_start+1])
+                        cur_labels = cur_labels[image_token_start+2:]
+                else:
+                    cur_new_input_embeds.append(self.get_model().embed_tokens(cur_input_ids[:image_token_start]))
+                    cur_new_input_embeds.append(cur_image_features)
+                    if labels is not None:
+                        cur_new_labels.append(cur_labels[:image_token_start])
+                        cur_new_labels.append(torch.full((cur_image_features.shape[0],), IGNORE_INDEX, device=labels.device, dtype=labels.dtype))
+                        cur_labels = cur_labels[image_token_start+1:]
+                cur_image_idx += 1
+                if getattr(self.config, 'tune_mm_mlp_adapter', False) and getattr(self.config, 'mm_use_im_start_end', False):
+                    cur_input_ids = cur_input_ids[image_token_start+2:]
+                else:
+                    cur_input_ids = cur_input_ids[image_token_start+1:]
+                image_token_indices = torch.where(cur_input_ids == IMAGE_TOKEN_INDEX)[0]
+            if cur_input_ids.numel() > 0:
+                if getattr(self.config, 'tune_mm_mlp_adapter', False) and getattr(self.config, 'mm_use_im_start_end', False):
+                    cur_new_input_embeds.append(self.get_model().embed_tokens(cur_input_ids).detach())
+                else:
+                    cur_new_input_embeds.append(self.get_model().embed_tokens(cur_input_ids))
+                if labels is not None:
+                    cur_new_labels.append(cur_labels)
+            cur_new_input_embeds = [x.to(device=self.device) for x in cur_new_input_embeds]
+            cur_new_input_embeds = torch.cat(cur_new_input_embeds, dim=0)
+            new_input_embeds.append(cur_new_input_embeds)
+            if labels is not None:
+                cur_new_labels = torch.cat(cur_new_labels, dim=0)
+                new_labels.append(cur_new_labels)
+
+        if any(x.shape != new_input_embeds[0].shape for x in new_input_embeds):
+            max_len = max(x.shape[0] for x in new_input_embeds)
+
+            new_input_embeds_align = []
+            for cur_new_embed in new_input_embeds:
+                cur_new_embed = torch.cat((cur_new_embed, torch.zeros((max_len - cur_new_embed.shape[0], cur_new_embed.shape[1]), dtype=cur_new_embed.dtype, device=cur_new_embed.device)), dim=0)
+                new_input_embeds_align.append(cur_new_embed)
+            new_input_embeds = torch.stack(new_input_embeds_align, dim=0)
+
+            if labels is not None:
+                new_labels_align = []
+                _new_labels = new_labels
+                for cur_new_label in new_labels:
+                    cur_new_label = torch.cat((cur_new_label, torch.full((max_len - cur_new_label.shape[0],), IGNORE_INDEX, dtype=cur_new_label.dtype, device=cur_new_label.device)), dim=0)
+                    new_labels_align.append(cur_new_label)
+                new_labels = torch.stack(new_labels_align, dim=0)
+
+            if attention_mask is not None:
+                new_attention_mask = []
+                for cur_attention_mask, cur_new_labels, cur_new_labels_align in zip(attention_mask, _new_labels, new_labels):
+                    new_attn_mask_pad_left = torch.full((cur_new_labels.shape[0] - labels.shape[1],), True, dtype=attention_mask.dtype, device=attention_mask.device)
+                    new_attn_mask_pad_right = torch.full((cur_new_labels_align.shape[0] - cur_new_labels.shape[0],), False, dtype=attention_mask.dtype, device=attention_mask.device)
+                    cur_new_attention_mask = torch.cat((new_attn_mask_pad_left, cur_attention_mask, new_attn_mask_pad_right), dim=0)
+                    new_attention_mask.append(cur_new_attention_mask)
+                attention_mask = torch.stack(new_attention_mask, dim=0)
+                assert attention_mask.shape == new_labels.shape
+        else:
+            new_input_embeds = torch.stack(new_input_embeds, dim=0)
+            if labels is not None:
+                new_labels  = torch.stack(new_labels, dim=0)
+
+            if attention_mask is not None:
+                new_attn_mask_pad_left = torch.full((attention_mask.shape[0], new_input_embeds.shape[1] - input_ids.shape[1]), True, dtype=attention_mask.dtype, device=attention_mask.device)
+                attention_mask = torch.cat((new_attn_mask_pad_left, attention_mask), dim=1)
+                assert attention_mask.shape == new_input_embeds.shape[:2]
+
+        return None, attention_mask, past_key_values, new_input_embeds, new_labels
+
+    def initialize_vision_tokenizer(self, model_args, tokenizer):
+        if model_args.mm_use_im_patch_token:
+            tokenizer.add_tokens([DEFAULT_IMAGE_PATCH_TOKEN], special_tokens=True)
+            self.resize_token_embeddings(len(tokenizer))
+
+        if model_args.mm_use_im_start_end:
+            num_new_tokens = tokenizer.add_tokens([DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN], special_tokens=True)
+            self.resize_token_embeddings(len(tokenizer))
+
+            if num_new_tokens > 0:
+                input_embeddings = self.get_input_embeddings().weight.data
+                output_embeddings = self.get_output_embeddings().weight.data
+
+                input_embeddings_avg = input_embeddings[:-num_new_tokens].mean(
+                    dim=0, keepdim=True)
+                output_embeddings_avg = output_embeddings[:-num_new_tokens].mean(
+                    dim=0, keepdim=True)
+
+                input_embeddings[-num_new_tokens:] = input_embeddings_avg
+                output_embeddings[-num_new_tokens:] = output_embeddings_avg
+
+            if model_args.tune_mm_mlp_adapter:
+                for p in self.get_input_embeddings().parameters():
+                    p.requires_grad = True
+                for p in self.get_output_embeddings().parameters():
+                    p.requires_grad = False
+
+            if model_args.pretrain_mm_mlp_adapter:
+                mm_projector_weights = torch.load(model_args.pretrain_mm_mlp_adapter, map_location='cpu')
+                print(mm_projector_weights)
+                embed_tokens_weight = mm_projector_weights['model.embed_tokens.weight']
+                assert num_new_tokens == 2
+                if input_embeddings.shape == embed_tokens_weight.shape:
+                    input_embeddings[-num_new_tokens:] = embed_tokens_weight[-num_new_tokens:]
+                elif embed_tokens_weight.shape[0] == num_new_tokens:
+                    input_embeddings[-num_new_tokens:] = embed_tokens_weight
+                else:
+                    raise ValueError(f"Unexpected embed_tokens_weight shape. Pretrained: {embed_tokens_weight.shape}. Current: {input_embeddings.shape}. Numer of new tokens: {num_new_tokens}.")
+        elif model_args.mm_use_im_patch_token:
+            if model_args.tune_mm_mlp_adapter:
+                for p in self.get_input_embeddings().parameters():
+                    p.requires_grad = False
+                for p in self.get_output_embeddings().parameters():
+                    p.requires_grad = False

geochat/model/language_model/.ipynb_checkpoints/geochat_llama-checkpoint.py

@@ -0,0 +1,140 @@
+#    Copyright 2023 Haotian Liu
+#
+#    Licensed under the Apache License, Version 2.0 (the "License");
+#    you may not use this file except in compliance with the License.
+#    You may obtain a copy of the License at
+#
+#        http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS,
+#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#    See the License for the specific language governing permissions and
+#    limitations under the License.
+
+
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+from torch.nn import CrossEntropyLoss
+
+from transformers import AutoConfig, AutoModelForCausalLM, \
+                         LlamaConfig, LlamaModel, LlamaForCausalLM
+
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+from ..geochat_arch import GeoChatMetaModel, GeoChatMetaForCausalLM
+
+
+class GeoChatConfig(LlamaConfig):
+    model_type = "geochat"
+
+
+class GeoChatLlamaModel(GeoChatMetaModel, LlamaModel):
+    config_class = GeoChatConfig
+
+    def __init__(self, config: LlamaConfig):
+        super(GeoChatLlamaModel, self).__init__(config)
+
+
+class GeoChatLlamaForCausalLM(LlamaForCausalLM, GeoChatMetaForCausalLM):
+    config_class = GeoChatConfig
+
+    def __init__(self, config):
+        super(LlamaForCausalLM, self).__init__(config)
+        self.model = GeoChatLlamaModel(config)
+
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_model(self):
+        return self.model
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        images: Optional[torch.FloatTensor] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        input_ids, attention_mask, past_key_values, inputs_embeds, labels = self.prepare_inputs_labels_for_multimodal(input_ids, attention_mask, past_key_values, labels, images)
+
+        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict
+        )
+
+        hidden_states = outputs[0]
+        logits = self.lm_head(hidden_states)
+
+        loss = None
+        if labels is not None:
+            # Shift so that tokens < n predict n
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            # Flatten the tokens
+            loss_fct = CrossEntropyLoss()
+            shift_logits = shift_logits.view(-1, self.config.vocab_size)
+            shift_labels = shift_labels.view(-1)
+            # Enable model/pipeline parallelism
+            shift_labels = shift_labels.to(shift_logits.device)
+            loss = loss_fct(shift_logits, shift_labels)
+
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return (loss,) + output if loss is not None else output
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+    def prepare_inputs_for_generation(
+        self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs
+    ):
+        if past_key_values:
+            input_ids = input_ids[:, -1:]
+
+        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
+        if inputs_embeds is not None and past_key_values is None:
+            model_inputs = {"inputs_embeds": inputs_embeds}
+        else:
+            model_inputs = {"input_ids": input_ids}
+
+        model_inputs.update(
+            {
+                "past_key_values": past_key_values,
+                "use_cache": kwargs.get("use_cache"),
+                "attention_mask": attention_mask,
+                "images": kwargs.get("images", None),
+            }
+        )
+        return model_inputs
+
+AutoConfig.register("geochat", GeoChatConfig)
+AutoModelForCausalLM.register(GeoChatConfig, GeoChatLlamaForCausalLM)

geochat/model/language_model/geochat_llama.py

@@ -0,0 +1,140 @@
+#    Copyright 2023 Haotian Liu
+#
+#    Licensed under the Apache License, Version 2.0 (the "License");
+#    you may not use this file except in compliance with the License.
+#    You may obtain a copy of the License at
+#
+#        http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS,
+#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#    See the License for the specific language governing permissions and
+#    limitations under the License.
+
+
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+from torch.nn import CrossEntropyLoss
+
+from transformers import AutoConfig, AutoModelForCausalLM, \
+                         LlamaConfig, LlamaModel, LlamaForCausalLM
+
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+from ..geochat_arch import GeoChatMetaModel, GeoChatMetaForCausalLM
+
+
+class GeoChatConfig(LlamaConfig):
+    model_type = "geochat"
+
+
+class GeoChatLlamaModel(GeoChatMetaModel, LlamaModel):
+    config_class = GeoChatConfig
+
+    def __init__(self, config: LlamaConfig):
+        super(GeoChatLlamaModel, self).__init__(config)
+
+
+class GeoChatLlamaForCausalLM(LlamaForCausalLM, GeoChatMetaForCausalLM):
+    config_class = GeoChatConfig
+
+    def __init__(self, config):
+        super(LlamaForCausalLM, self).__init__(config)
+        self.model = GeoChatLlamaModel(config)
+
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_model(self):
+        return self.model
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        images: Optional[torch.FloatTensor] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        input_ids, attention_mask, past_key_values, inputs_embeds, labels = self.prepare_inputs_labels_for_multimodal(input_ids, attention_mask, past_key_values, labels, images)
+
+        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict
+        )
+
+        hidden_states = outputs[0]
+        logits = self.lm_head(hidden_states)
+
+        loss = None
+        if labels is not None:
+            # Shift so that tokens < n predict n
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            # Flatten the tokens
+            loss_fct = CrossEntropyLoss()
+            shift_logits = shift_logits.view(-1, self.config.vocab_size)
+            shift_labels = shift_labels.view(-1)
+            # Enable model/pipeline parallelism
+            shift_labels = shift_labels.to(shift_logits.device)
+            loss = loss_fct(shift_logits, shift_labels)
+
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return (loss,) + output if loss is not None else output
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+    def prepare_inputs_for_generation(
+        self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs
+    ):
+        if past_key_values:
+            input_ids = input_ids[:, -1:]
+
+        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
+        if inputs_embeds is not None and past_key_values is None:
+            model_inputs = {"inputs_embeds": inputs_embeds}
+        else:
+            model_inputs = {"input_ids": input_ids}
+
+        model_inputs.update(
+            {
+                "past_key_values": past_key_values,
+                "use_cache": kwargs.get("use_cache"),
+                "attention_mask": attention_mask,
+                "images": kwargs.get("images", None),
+            }
+        )
+        return model_inputs
+
+AutoConfig.register("geochat", GeoChatConfig)
+AutoModelForCausalLM.register(GeoChatConfig, GeoChatLlamaForCausalLM)

geochat/model/language_model/geochat_mpt.py

@@ -0,0 +1,113 @@
+#    Copyright 2023 Haotian Liu
+#
+#    Licensed under the Apache License, Version 2.0 (the "License");
+#    you may not use this file except in compliance with the License.
+#    You may obtain a copy of the License at
+#
+#        http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS,
+#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#    See the License for the specific language governing permissions and
+#    limitations under the License.
+
+
+from typing import List, Optional, Tuple
+import warnings
+
+import torch
+import torch.nn.functional as F
+import math
+
+from transformers import AutoConfig, AutoModelForCausalLM
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+from .mpt.modeling_mpt import MPTConfig, MPTForCausalLM, MPTModel
+from geochat.model.geochat_arch import GeoChatMetaModel, GeoChatMetaForCausalLM
+
+
+class GeoChatMPTConfig(MPTConfig):
+    model_type = "geochat_mpt"
+
+
+class GeoChatMPTModel(GeoChatMetaModel, MPTModel):
+    config_class = GeoChatMPTConfig
+
+    def __init__(self, config: MPTConfig):
+        config.hidden_size = config.d_model
+        super(GeoChatMPTModel, self).__init__(config)
+    
+    def embed_tokens(self, x):
+        return self.wte(x)
+
+
+class GeoChatMPTForCausalLM(MPTForCausalLM, GeoChatMetaForCausalLM):
+    config_class = GeoChatMPTConfig
+    supports_gradient_checkpointing = True
+
+    def __init__(self, config):
+        super(MPTForCausalLM, self).__init__(config)
+
+        if not config.tie_word_embeddings:
+            raise ValueError('MPTForCausalLM only supports tied word embeddings')
+        self.transformer = GeoChatMPTModel(config)
+        self.logit_scale = None
+        if config.logit_scale is not None:
+            logit_scale = config.logit_scale
+            if isinstance(logit_scale, str):
+                if logit_scale == 'inv_sqrt_d_model':
+                    logit_scale = 1 / math.sqrt(config.d_model)
+                else:
+                    raise ValueError(f"logit_scale={logit_scale!r} is not recognized as an option; use numeric value or 'inv_sqrt_d_model'.")
+            self.logit_scale = logit_scale
+
+    def get_model(self):
+        return self.transformer
+
+    def _set_gradient_checkpointing(self, module, value=False):
+        if isinstance(module, GeoChatMPTModel):
+            module.gradient_checkpointing = value
+
+    def forward(self, input_ids: torch.LongTensor, past_key_values: Optional[List[Tuple[torch.FloatTensor]]]=None, attention_mask: Optional[torch.ByteTensor]=None, prefix_mask: Optional[torch.ByteTensor]=None, sequence_id: Optional[torch.LongTensor]=None, labels: Optional[torch.LongTensor]=None, return_dict: Optional[bool]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, use_cache: Optional[bool]=None, images=None):
+        return_dict = return_dict if return_dict is not None else self.config.return_dict
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+
+        input_ids, attention_mask, past_key_values, inputs_embeds, labels = self.prepare_inputs_labels_for_multimodal(input_ids, attention_mask, past_key_values, labels, images)
+        outputs = self.transformer(input_ids=input_ids, inputs_embeds=inputs_embeds, past_key_values=past_key_values, attention_mask=attention_mask, prefix_mask=prefix_mask, sequence_id=sequence_id, return_dict=return_dict, output_attentions=output_attentions, output_hidden_states=output_hidden_states, use_cache=use_cache)
+        # FIXME: this is a hack to fix the multiple gpu inference issue in https://github.com/haotian-liu/LLaVA/issues/338
+        logits = F.linear(outputs.last_hidden_state.to(self.transformer.wte.weight.device), self.transformer.wte.weight)
+        if self.logit_scale is not None:
+            if self.logit_scale == 0:
+                warnings.warn(f'Multiplying logits by self.logit_scale={self.logit_scale!r}. This will produce uniform (uninformative) outputs.')
+            logits *= self.logit_scale
+        loss = None
+        if labels is not None:
+            labels = torch.roll(labels, shifts=-1)
+            labels[:, -1] = -100
+            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), labels.to(logits.device).view(-1))
+        return CausalLMOutputWithPast(loss=loss, logits=logits, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states)
+
+    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, inputs_embeds=None, **kwargs):
+        if inputs_embeds is not None:
+            raise NotImplementedError('inputs_embeds is not implemented for MPT yet')
+        attention_mask = kwargs['attention_mask'].bool()
+        if attention_mask[:, -1].sum() != attention_mask.shape[0]:
+            raise NotImplementedError('MPT does not support generation with right padding.')
+        if self.transformer.attn_uses_sequence_id and self.training:
+            sequence_id = torch.zeros_like(input_ids[:1])
+        else:
+            sequence_id = None
+        if past_key_values is not None:
+            input_ids = input_ids[:, -1].unsqueeze(-1)
+        if self.transformer.prefix_lm:
+            prefix_mask = torch.ones_like(attention_mask)
+            if kwargs.get('use_cache') == False:
+                raise NotImplementedError('MPT with prefix_lm=True does not support use_cache=False.')
+        else:
+            prefix_mask = None
+        return {'input_ids': input_ids, 'attention_mask': attention_mask, 'prefix_mask': prefix_mask, 'sequence_id': sequence_id, 'past_key_values': past_key_values, 'use_cache': kwargs.get('use_cache', True), "images": kwargs.get("images", None)}
+
+
+AutoConfig.register("geochat_mpt", GeoChatMPTConfig)
+AutoModelForCausalLM.register(GeoChatMPTConfig, GeoChatMPTForCausalLM)