add online demo

.gitignore

@@ -0,0 +1,2 @@
+*.pth
+*.bin

README.md

@@ -9,4 +9,33 @@ app_file: app.py
 pinned: false
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+<img src="https://github.com/chencn2020/SEAGULL/raw/main/imgs/Logo/logo.png" alt="SEAGULL" style="height: auto; width: 100%;">
+
+<div style="display: flex; justify-content: center; gap: 10px; flex-wrap: wrap; width: 100%;">
+  <a href=""><img src="https://huggingface.co/datasets/huggingface/badges/raw/main/open-in-hf-spaces-sm-dark.svg" alt="Open in Spaces" style="max-width: 100%; height: auto;"></a>
+  <a href="https://arxiv.org/abs/2411.10161"><img src="https://img.shields.io/badge/Arxiv-2411:10161-red" style="max-width: 100%; height: auto;"></a>
+  <a href="https://hits.seeyoufarm.com"><img src="https://hits.seeyoufarm.com/api/count/incr/badge.svg?url=https%3A%2F%2Fhuggingface.co%2Fdatasets%2FZevin2023%2FSEAGULL-100w&count_bg=%2379C83D&title_bg=%23555555&icon=&icon_color=%23E7E7E7&title=Visitors&edge_flat=false" style="max-width: 100%; height: auto;"></a>
+  <a href='https://github.com/chencn2020/SEAGULL/'><img src='https://img.shields.io/github/stars/chencn2020/Seagull.svg?style=social' style="max-width: 100%; height: auto;"></a>
+</div>
+
+## Acknowledgement 💌
+<div id="Acknowledgement"></div>
+- [Osprey](https://github.com/CircleRadon/Osprey) and [LLaVA-v1.5](https://github.com/haotian-liu/LLaVA): We build this repostory based on them.
+- [RAISE](http://loki.disi.unitn.it/RAISE/): The Dist. images in SEAGULL-100w are constructed based on this dataset.
+- [SAM](https://segment-anything.com/) and [SEEM](https://github.com/UX-Decoder/Segment-Everything-Everywhere-All-At-Once): The mask-based ROIs are generated using these two awesome works. And SAM are used to get the segmentation result in the demo.
+- [TOPIQ](https://github.com/chaofengc/IQA-PyTorch): The quality scores and importance scores for ROIs are generated using this great FR-IQA.
+
+
+## Citation 🖊️
+If our work is useful to your research, we will be grateful for you to cite our paper:
+```
+@misc{chen2024seagull,
+      title={SEAGULL: No-reference Image Quality Assessment for Regions of Interest via Vision-Language Instruction Tuning}, 
+      author={Zewen Chen and Juan Wang and Wen Wang and Sunhan Xu and Hang Xiong and Yun Zeng and Jian Guo and Shuxun Wang and Chunfeng Yuan and Bing Li and Weiming Hu},
+      year={2024},
+      eprint={2411.10161},
+      archivePrefix={arXiv},
+      primaryClass={cs.CV},
+      url={https://arxiv.org/abs/2411.10161}, 
+}
+```

app.py

@@ -0,0 +1,62 @@
+import argparse
+from demo.UI import Main_ui
+
+if __name__ == '__main__':
+    import subprocess 
+    import sys
+    def run_command(command):
+        subprocess.check_call([sys.executable, '-m'] + command.split(), shell=False)
+
+    # Install the package in editable mode
+    run_command("pip install -e .")
+    
+    # Install NVM (Node Version Manager)
+    run_command("curl -o- https://raw.githubusercontent.com/nvm-sh/nvm/v0.39.3/install.sh | bash")
+
+    # Source the appropriate shell configuration file
+    run_command("source ~/.bashrc")  # You can change to ~/.zshrc based on your shell
+
+    # Install Node.js version 18.16.0
+    run_command("nvm install v18.16.0")
+
+    # Install pnpm (package manager)
+    run_command("curl -fsSL https://get.pnpm.io/install.sh | sh -")
+
+    # Source the shell configuration file again (for pnpm)
+    run_command("source ~/.bashrc")  # You can change to ~/.zshrc based on your shell
+
+    # Verify if pnpm was installed correctly
+    run_command("pnpm --version")
+
+    # Clone the Gradio BBox repository
+    run_command("git clone https://github.com/chencn2020/gradio-bbox.git")
+
+    # Change into the cloned repository directory
+    run_command("cd gradio-bbox")
+
+    # Build frontend
+    run_command("bash scripts/build_frontend.sh")
+
+    
+
+    # Change back to the previous directory
+    run_command("cd ..")
+
+    # Install the package again in editable mode
+    run_command("pip install -e .")
+
+    # Install Segment Anything repository from GitHub
+    run_command("pip install git+https://github.com/facebookresearch/segment-anything.git")
+
+    # Download the model checkpoint
+    run_command("curl -o ./checkpoints/sam_vit_b_01ec64.pth https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth")
+
+ 
+
+    parser = argparse.ArgumentParser(description='SEAGULL', formatter_class=argparse.RawTextHelpFormatter)
+    parser.add_argument('--model', help='path to seagull model', default='Zevin2023/SEAGULL-7B')
+    parser.add_argument('--example_path', help='path to examples', default='./imgs/Examples')
+    args = parser.parse_args()
+    
+    demo = Main_ui(args).load_demo()
+    demo.launch(server_port=7530)

demo/UI.py

@@ -0,0 +1,143 @@
+import os
+import gradio as gr
+from demo.sam_inference import SAM_Inference
+from demo.seagull_inference import Seagull
+from demo.mask_utils import ImageSketcher
+
+class Main_ui():
+    def __init__(self, args) -> None: 
+        self.args = args
+        self.seagull = Seagull(model_path=args.model)
+        
+        self.example_list = self.load_example()
+        self.sam = SAM_Inference()
+        # self.sam_predictor = get_sam_predictor()
+        # self.mask_generator = get_mask_generator()
+    
+    def load_example(self):
+        examples = []
+        for file in sorted(os.listdir(self.args.example_path)):
+            examples.append([os.path.join(self.args.example_path, file)])
+        return examples
+            
+    def load_demo(self):
+        with gr.Blocks() as demo:
+            preprocessed_img = gr.State(value=None)
+            binary_mask = gr.State(value=None)
+            
+            with gr.Row():
+                gr.Markdown("""
+                            <img src="https://github.com/chencn2020/SEAGULL/raw/main/imgs/Logo/logo.png" alt="SEAGULL" style="height: auto; width: 100%; margin-bottom: 3%;">
+                            
+                            ## 🔔 Usage
+                            
+                            Firstly, you need to upload an image and choose the analyse types **(quality score, importance score and distortion analysis)**. 
+                            
+                            Then you can click **(points)** or pull a frame **(bbox)** on the image to indicate the region of interest (ROIs). 
+                            
+                            After that, this demo process the following steps:
+                            
+                            > 1. SAM extracts the mask-based ROIs based on your clicked points or frame.
+                            
+                            > 2. Based on the uploaded image and mask-based ROIs, SEAGULL analyses the quality of the ROIs.
+                        
+                            """)
+
+            with gr.TabItem("Mask-based ROIs (Points)"):
+                with gr.Row():
+                    input_image_ponit = gr.Image(type="numpy", label='Input image', height=512) # input image
+                    output_mask_ponit = gr.Image(label='Mask-based ROI', height=512) # output binary mask
+                
+                with gr.Row():
+                    output_mask_point_on_img = gr.Image(label='Mask on image', height=512) # mask on image for better view
+                        
+                    with gr.Column():
+                        radio_point = gr.Radio(label='Analysis type', choices=['Quality Score', 'Importance Score', 'Distortion Analysis'], value='Quality Score')
+                        output_text_point = gr.Textbox(label='Analysis Results')
+                        point_seg_button = gr.Button('Analysis')
+                        
+                        point_example = gr.Dataset(label='Examples', components=[input_image_ponit], samples=self.example_list)
+
+            with gr.TabItem("Mask-based ROIs (BBox)"):
+                with gr.Row():
+                    input_image_BBOX = ImageSketcher(type="numpy", label='Input image', height=512)
+                    output_mask_BBOX = gr.Image(label='Mask-based ROI', height=512)
+                
+                with gr.Row():
+                    output_BBOX_mask_on_img = gr.Image(label='Mask on image', height=512)
+                        
+                    with gr.Column():
+                        radio_BBOX = gr.Radio(label='Analysis type', choices=['Quality Score', 'Importance Score', 'Distortion Analysis'], value='Quality Score')
+                        output_text_BBOX = gr.Textbox(label='ROI Quality Analysis')
+                        box_seg_button = gr.Button('Generate mask and analysis')
+                        box_analyse_button = gr.Button('Analysis')
+                        
+                        BBOX_example = gr.Dataset(label='Examples', components=[input_image_BBOX], samples=self.example_list)
+
+            # click point
+            input_image_ponit.upload(
+                self.seagull.init_image,
+                [input_image_ponit],
+                [preprocessed_img, input_image_ponit, input_image_BBOX]
+            )
+
+            point_example.click(
+                self.seagull.init_image,
+                [point_example],
+                [preprocessed_img, input_image_ponit, input_image_BBOX]
+            )
+            
+            # after clicking on the image
+            input_image_ponit.select(
+                self.sam.img_select_point,
+                [preprocessed_img],
+                [input_image_ponit, output_mask_ponit, output_mask_point_on_img, binary_mask]
+            ).then(
+                self.seagull.seagull_predict,
+                [preprocessed_img, binary_mask, radio_point],
+                [output_text_point]
+            )
+            
+            point_seg_button.click(
+                self.seagull.seagull_predict,
+                [preprocessed_img, binary_mask, radio_point],
+                [output_text_point]
+            )
+
+            # draw frame
+            input_image_BBOX.upload(
+                self.seagull.init_image,
+                [input_image_BBOX],
+                [preprocessed_img, input_image_ponit, input_image_BBOX]
+            )
+
+            BBOX_example.click(
+                self.seagull.init_image,
+                [BBOX_example],
+                [preprocessed_img, input_image_ponit, input_image_BBOX]
+            )
+
+            # after drawing a frame on the image
+            input_image_BBOX.select(
+                self.sam.gen_box_seg,
+                [input_image_BBOX],
+                [output_mask_BBOX, output_BBOX_mask_on_img, binary_mask]
+            )
+            
+            box_seg_button.click(
+                self.sam.gen_box_seg,
+                [input_image_BBOX],
+                [output_mask_BBOX, output_BBOX_mask_on_img, binary_mask]
+            ).then(
+                self.seagull.seagull_predict,
+                [preprocessed_img, binary_mask, radio_BBOX],
+                [output_text_BBOX]
+            )
+            
+            box_analyse_button.click(
+                self.seagull.seagull_predict,
+                [preprocessed_img, binary_mask, radio_BBOX],
+                [output_text_BBOX]
+            )
+
+        return demo

demo/mask_utils.py

@@ -0,0 +1,144 @@
+import cv2
+from PIL import Image
+import numpy as np
+import torch
+import gradio as gr
+
+class ImageSketcher(gr.Image):
+    """
+    Code is from https://github.com/jshilong/GPT4RoI/blob/7c157b5f33914f21cfbc804fb301d3ce06324193/gpt4roi/app.py#L365
+
+    Fix the bug of gradio.Image that cannot upload with tool == 'sketch'.
+    """
+
+    is_template = True  # Magic to make this work with gradio.Block, don't remove unless you know what you're doing.
+
+    def __init__(self, **kwargs):
+        super().__init__(tool='boxes', **kwargs)
+
+    def preprocess(self, x):
+        if x is None:
+            return x
+        if self.tool == 'boxes' and self.source in ['upload', 'webcam']:
+            if isinstance(x, str):
+                x = {'image': x, 'boxes': []}
+            else:
+                assert isinstance(x, dict)
+                assert isinstance(x['image'], str)
+                assert isinstance(x['boxes'], list)
+        x = super().preprocess(x)
+        return x
+    
+def process_mask_to_show(mask):
+    '''
+        Process the mask to show on the gradio.Image
+    '''
+    mask = np.array(mask > 0.1, dtype=np.uint8) * 255
+    mask_stacked = np.stack([mask] * 3, axis=-1)
+
+    return mask_stacked
+
+def img_add_masks(img_, colored_mask, mask, linewidth=2):
+    if type(img_) is np.ndarray:
+        img = Image.fromarray(img_, mode='RGB').convert('RGBA')
+    else:
+        img = img_.copy()
+    h, w = img.height, img.width
+    # contour
+    temp = np.zeros((h, w, 1))
+    contours, _ = cv2.findContours(mask.astype(np.uint8), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
+    cv2.drawContours(temp, contours, -1, (255, 255, 255), linewidth)
+    color = np.array([1, 1, 1, 1])
+    contour_mask = temp * color.reshape(1, 1, -1)
+
+    overlay_inner = Image.fromarray(colored_mask.astype(np.uint8), 'RGBA')
+    img.paste(overlay_inner, (0, 0), overlay_inner)
+
+    overlay_contour = Image.fromarray(contour_mask.astype(np.uint8), 'RGBA')
+    img.paste(overlay_contour, (0, 0), overlay_contour)
+    return img
+
+def gen_colored_masks(
+        annotation,
+        random_color=False,
+):
+    """
+    Code is largely based on https://github.com/CASIA-IVA-Lab/FastSAM/blob/4d153e909f0ad9c8ecd7632566e5a24e21cf0071/utils/tools_gradio.py#L130
+    """
+    device = annotation.device
+    mask_sum = annotation.shape[0]
+    height = annotation.shape[1]
+    weight = annotation.shape[2]
+    areas = torch.sum(annotation, dim=(1, 2))
+    sorted_indices = torch.argsort(areas, descending=False)
+    annotation = annotation[sorted_indices]
+
+    index = (annotation != 0).to(torch.long).argmax(dim=0)
+    if random_color:
+        color = torch.rand((mask_sum, 1, 1, 3)).to(device)
+    else:
+        color = torch.ones((mask_sum, 1, 1, 3)).to(device) * torch.tensor(
+            [30 / 255, 144 / 255, 255 / 255]
+        ).to(device)
+    transparency = torch.ones((mask_sum, 1, 1, 1)).to(device) * 0.6
+    visual = torch.cat([color, transparency], dim=-1)
+    mask_image = torch.unsqueeze(annotation, -1) * visual
+
+    mask = torch.zeros((height, weight, 4)).to(device)
+    h_indices, w_indices = torch.meshgrid(torch.arange(height), torch.arange(weight))
+    indices = (index[h_indices, w_indices], h_indices, w_indices, slice(None))
+
+    mask[h_indices, w_indices, :] = mask_image[indices]
+    mask_cpu = mask.cpu().numpy()
+
+    return mask_cpu, sorted_indices
+
+def mask_foreground(mask, trans=0.6, random_color=True):
+    if random_color:
+        color = np.concatenate([np.random.random(3) * 255, np.array([trans * 255])], axis=0)
+    else:
+        color = np.array([30, 144, 255, trans * 255])
+    h, w = mask.shape[-2:]
+    mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
+
+    return mask_image
+
+
+def mask_background(mask, trans=0.5):
+    h, w = mask.shape[-2:]
+    mask_image = (1 - mask.reshape(h, w, 1)) * np.array([0, 0, 0, trans * 255])
+
+    return mask_image
+
+
+def mask_select_point(all_masks, output_mask_2_raw, mask_order, evt: gr.SelectData):
+    h, w = output_mask_2_raw.height, output_mask_2_raw.width
+    pointed_mask = None
+    for i in range(len(mask_order)):
+        idx = mask_order[i]
+        msk = all_masks[idx]
+        if msk[evt.index[1], evt.index[0]] == 1:
+            pointed_mask = msk.copy()
+            break
+
+    if pointed_mask is not None:
+        contours, hierarchy = cv2.findContours(pointed_mask.astype("uint8"), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
+        ret = output_mask_2_raw.copy()
+
+        temp = np.zeros((h, w, 1))
+        contours, _ = cv2.findContours(msk.astype(np.uint8), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
+        cv2.drawContours(temp, contours, -1, (255, 255, 255), 3)
+        color = np.array([1, 1, 1, 1])
+        contour_mask = temp * color.reshape(1, 1, -1)
+
+        colored_mask = mask_background(pointed_mask)
+
+        overlay_inner = Image.fromarray(colored_mask.astype(np.uint8), 'RGBA')
+        ret.paste(overlay_inner, (0, 0), overlay_inner)
+
+        overlay_contour = Image.fromarray(contour_mask.astype(np.uint8), 'RGBA')
+        ret.paste(overlay_contour, (0, 0), overlay_contour)
+
+        return ret, pointed_mask
+    else:
+        return output_mask_2_raw, None

demo/sam_inference.py

@@ -0,0 +1,102 @@
+import gc
+
+import numpy as np
+import torch
+from segment_anything import SamPredictor, sam_model_registry, SamAutomaticMaskGenerator
+import gradio as gr
+import cv2
+from demo.mask_utils import *
+
+class SAM_Inference:
+    def __init__(self, model_type='vit_b', device='cuda') -> None:
+        models = {
+            'vit_b': './checkpoints/sam_vit_b_01ec64.pth',
+            'vit_l': './checkpoints/sam_vit_l_0b3195.pth',
+            'vit_h': './checkpoints/sam_vit_h_4b8939.pth'
+        }
+
+        sam = sam_model_registry[model_type](checkpoint=models[model_type])
+        sam = sam.to(device)
+
+        self.predictor = SamPredictor(sam)
+        self.mask_generator = SamAutomaticMaskGenerator(model=sam)
+
+    def img_select_point(self, original_img: np.ndarray, evt: gr.SelectData):
+        img = original_img.copy()
+        sel_pix = [(evt.index, 1)]  # append the foreground_point
+
+        masks = self.run_inference(original_img, sel_pix)
+        for point, label in sel_pix:
+            cv2.circle(img, point, 5, (240, 240, 240), -1, 0)
+            cv2.circle(img, point, 5, (30, 144, 255), 2, 0)
+
+        mask = masks[0][0]
+        colored_mask = mask_foreground(mask)
+        res = img_add_masks(original_img, colored_mask, mask)
+        return img, process_mask_to_show(mask), res, mask
+
+    def gen_box_seg(self, inp):
+        if inp is None:
+            raise gr.Error("Please upload an image first!")
+        image = inp['image']
+        if len(inp['boxes']) == 0:
+            raise gr.Error("Please clear the raw boxes and draw a box first!")
+        boxes = inp['boxes'][-1]
+
+        input_box = np.array([boxes[0], boxes[1], boxes[2], boxes[3]]).astype(int)
+
+        masks = self.predict_box(image, input_box)
+
+        mask = masks[0][0]
+        colored_mask = mask_foreground(mask)
+        res = img_add_masks(image, colored_mask, mask)
+
+        return process_mask_to_show(mask), res, mask
+    
+    def run_inference(self, input_x, selected_points):
+        if len(selected_points) == 0:
+            return []
+
+        self.predictor.set_image(input_x)
+
+        points = torch.Tensor(
+            [p for p, _ in selected_points]
+        ).to(self.predictor.device).unsqueeze(0)
+
+        labels = torch.Tensor(
+            [int(l) for _, l in selected_points]
+        ).to(self.predictor.device).unsqueeze(0)
+
+        transformed_points = self.predictor.transform.apply_coords_torch(
+            points, input_x.shape[:2])
+
+        # predict segmentation according to the boxes
+        masks, scores, logits = self.predictor.predict_torch(
+            point_coords=transformed_points,
+            point_labels=labels,
+            multimask_output=False,
+        )
+        masks = masks.cpu().detach().numpy()
+
+        gc.collect()
+        torch.cuda.empty_cache()
+
+        return masks
+
+    def predict_box(self, input_x, input_box):
+        self.predictor.set_image(input_x)
+
+        input_boxes = torch.tensor(input_box[None, :], device=self.predictor.device)
+        transformed_boxes = self.predictor.transform.apply_boxes_torch(input_boxes, input_x.shape[:2])
+
+        masks, _, _ = self.predictor.predict_torch(
+            point_coords=None,
+            point_labels=None,
+            boxes=transformed_boxes,
+            multimask_output=False
+        )
+        masks = masks.cpu().detach().numpy()
+
+        gc.collect()
+        torch.cuda.empty_cache()
+        return masks

demo/seagull_inference.py

@@ -0,0 +1,163 @@
+import torch
+from seagull.utils import disable_torch_init
+from transformers import AutoTokenizer, CLIPImageProcessor
+from seagull.model.language_model.seagull_llama import SeagullLlamaForCausalLM
+from seagull.mm_utils import tokenizer_image_token
+from seagull.conversation import conv_templates, SeparatorStyle
+from seagull.constants import IMAGE_TOKEN_INDEX
+from seagull.train.train import DataArguments
+
+from functools import partial
+import os
+import numpy as np
+import cv2
+from typing import List
+from PIL import Image
+
+class Seagull():
+    def __init__(self, model_path, device='cuda'):
+        disable_torch_init()
+        model_path = os.path.expanduser(model_path)
+        self.tokenizer = AutoTokenizer.from_pretrained(model_path, model_max_length=2048, padding_side="right", use_fast=True)
+        self.model = SeagullLlamaForCausalLM.from_pretrained(model_path, torch_dtype=torch.bfloat16,).to(device)
+        self.tokenizer.pad_token = self.tokenizer.unk_token
+
+        self.image_processor = CLIPImageProcessor(do_resize=True, size={"shortest_edge":512}, resample=3,  do_center_crop=True, crop_size={"height": 512, "width": 512},
+                                                    do_rescale=True, rescale_factor=0.00392156862745098, do_normalize=True, image_mean=[0.48145466, 0.4578275, 0.40821073],
+                                                    image_std=[0.26862954, 0.26130258, 0.27577711], do_convert_rgb=True, )
+        
+        spi_tokens = ['<global>', '<local>']
+        self.tokenizer.add_tokens(spi_tokens, special_tokens=True)
+        
+        for m in self.model.modules():
+            m.tokenizer = self.tokenizer
+
+        vision_tower = self.model.get_vision_tower()
+        if not vision_tower.is_loaded:
+            vision_tower.load_model()
+        vision_tower.to(dtype=torch.float16, device=device)
+
+        begin_str = "<image>\nThis provides an overview of the image.\n Please answer the following questions about the provided region. Note: Distortions include: blur, colorfulness, compression, contrast exposure and noise.\n Here is the region <global><local>. "
+        
+        instruction = {
+            'distortion analysis': 'Provide the distortion type of this region.',
+            'quality score': 'Analyze the quality of this region.',
+            'importance score': 'Consider the impact of this region on the overall image quality. Analyze its importance to the overall image quality.'
+        }
+        
+        self.ids_input = {}
+        for ins_type, ins in instruction.items():
+            conv = conv_templates['seagull_v1'].copy()
+            qs = begin_str + ins
+            conv.append_message(conv.roles[0], qs)
+            conv.append_message(conv.roles[1], None)
+            prompt = conv.get_prompt()
+            self.ids_input[ins_type] = tokenizer_image_token(prompt, self.tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).to(self.model.device)
+
+        self.stop_str = conv.sep if conv.sep_style != SeparatorStyle.TWO else conv.sep2
+        
+    def init_image(self, img):
+        if isinstance(img, dict):
+            img = img['image']
+        elif isinstance(img, List):
+            img = cv2.imread(img[0])
+            img = img[:, :, ::-1]
+        h_, w_ = img.shape[:2]
+        if h_ > 512:
+            ratio = 512 / h_
+            new_h, new_w = int(h_ * ratio), int(w_ * ratio)
+            preprocessed_img = cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_LINEAR)
+        else:
+            preprocessed_img = img.copy()
+
+        return (preprocessed_img, preprocessed_img, preprocessed_img)
+
+    def preprocess(self, img):
+        image = self.image_processor.preprocess(img,
+                                do_center_crop=False,
+                                return_tensors='pt')['pixel_values'][0]
+
+        image = torch.nn.functional.interpolate(image.unsqueeze(0),
+                                                size=(512, 512),
+                                                mode='bilinear',
+                                                align_corners=False).squeeze(0)
+        
+        return image
+    
+    def seagull_predict(self, img, mask, instruct_type):
+        image = self.preprocess(img)
+        
+        mask = np.array(mask, dtype=np.int)
+        ys, xs = np.where(mask > 0)
+        if len(xs) > 0 and len(ys) > 0:
+            # Find the minimal bounding rectangle for the entire mask
+            x_min, x_max = np.min(xs), np.max(xs)
+            y_min, y_max = np.min(ys), np.max(ys)
+            w1 = x_max - x_min
+            h1 = y_max - y_min
+            
+            bounding_box = (x_min, y_min, w1, h1)
+        else:
+            bounding_box = None
+            
+        mask = cv2.resize(mask, (512, 512), interpolation=cv2.INTER_NEAREST)
+        mask = np.array(mask > 0.1, dtype=np.uint8)
+        masks = torch.Tensor(mask).unsqueeze(0).to(self.model.device)
+        
+        input_ids = self.ids_input[instruct_type.lower()]
+        
+        x1, y1, w1, h1 = list(map(int, bounding_box))  # x y w h
+        cropped_img = img[y1:y1 + h1, x1:x1 + w1]
+        cropped_img = Image.fromarray(cropped_img)
+        cropped_img = self.preprocess(cropped_img)
+            
+        with torch.inference_mode():
+
+            self.model.orig_forward = self.model.forward
+            self.model.forward = partial(self.model.orig_forward,
+                                        img_metas=[None],
+                                        masks=[masks.half()],
+                                        cropped_img=cropped_img.unsqueeze(0)
+                                        )
+            output_ids = self.model.generate(
+                input_ids,
+                images=image.unsqueeze(0).half().to(self.model.device),
+                do_sample=False,
+                temperature=1,
+                max_new_tokens=2048,
+                use_cache=True,
+                num_beams=1,
+                top_k = 0, # 不进行topk
+                top_p = 1, # 累计概率为
+                )
+
+            self.model.forward = self.model.orig_forward
+
+        input_token_len = input_ids.shape[1]
+        n_diff_input_output = (
+            input_ids != 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 = self.tokenizer.batch_decode(output_ids[:, input_token_len:],
+                                            skip_special_tokens=True)[0]
+    
+        outputs = outputs.strip()
+        if outputs.endswith(self.stop_str):
+            outputs = outputs[:-len(self.stop_str)]
+        outputs = outputs.strip()
+        if ':' in outputs:
+            outputs = outputs.split(':')[1]
+
+        outputs_list = outputs.split('.')
+        outputs_list_final = []
+        outputs_str = ''
+        for output in outputs_list:
+            if output not in outputs_list_final:
+                if output=='':
+                    continue
+                outputs_list_final.append(output)
+                outputs_str+=output+'.'
+            else:
+                break
+        return outputs_str

seagull/__init__.py

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

seagull/builder.py

@@ -0,0 +1,171 @@
+#    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 seagull.model import *
+from seagull.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", use_flash_attn=False, **kwargs):
+    kwargs = {"device_map": device_map, **kwargs}
+
+    if device != "cuda":
+        kwargs['device_map'] = {"": device}
+
+    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 use_flash_attn:
+        kwargs['attn_implementation'] = 'flash_attention_2'
+
+    if 'seagull' in model_name.lower() or True:
+        # 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 or True:
+            from seagull.model.language_model.seagull_llama import SeagullConfig
+            lora_cfg_pretrained = SeagullConfig.from_pretrained(model_path)
+            tokenizer = AutoTokenizer.from_pretrained(model_base, use_fast=False)
+            print('Loading LLaVA from base model...')
+            model = SeagullLlamaForCausalLM.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 LLaVA 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')
+            
+            for k, v in non_lora_trainables.items():
+                print(k)
+            print('print non lora')
+            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 LLaVA 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 = SeagullMptForCausalLM.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 = SeagullLlamaForCausalLM.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 = SeagullMptForCausalLM.from_pretrained(model_path, low_cpu_mem_usage=True, **kwargs)
+            elif 'mistral' in model_name.lower():
+                tokenizer = AutoTokenizer.from_pretrained(model_path)
+                model = SeagullMistralForCausalLM.from_pretrained(
+                    model_path,
+                    low_cpu_mem_usage=True,
+                    **kwargs
+                )
+            else:
+                tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)
+                model = SeagullLlamaForCausalLM.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, low_cpu_mem_usage=True, **kwargs)
+            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 'seagull' in model_name.lower() or True:
+        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(device_map=device_map)
+        if device_map != 'auto':
+            vision_tower.to(device=device_map, 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

seagull/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>"

seagull/conversation.py

@@ -0,0 +1,381 @@
+import dataclasses
+from enum import auto, Enum
+from typing import List, Tuple
+
+
+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_seagull_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_seagull_plain = Conversation(
+    system="",
+    roles=("", ""),
+    messages=(
+    ),
+    offset=0,
+    sep_style=SeparatorStyle.PLAIN,
+    sep="\n",
+)
+
+conv_seagull_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_seagull_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_seagull_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_seagull_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_seagull_plain,
+    "v0_plain": conv_seagull_plain,
+    "seagull_v0": conv_seagull_v0,
+    "v0_mmtag": conv_seagull_v0_mmtag,
+    "seagull_v1": conv_seagull_v1,
+    "v1_mmtag": conv_seagull_v1_mmtag,
+    "seagull_llama_2": conv_seagull_llama_2,
+
+    "mpt": conv_mpt,
+}
+
+
+if __name__ == "__main__":
+    print(default_conversation.get_prompt())

seagull/mm_utils.py

@@ -0,0 +1,95 @@
+from PIL import Image
+from io import BytesIO
+import base64
+
+import torch
+from transformers import StoppingCriteria
+from seagull.constants import IMAGE_TOKEN_INDEX
+
+
+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, 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 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 = []
+        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:]
+            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, 3)
+        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:
+                return True
+        outputs = self.tokenizer.batch_decode(output_ids[:, -offset:], skip_special_tokens=True)[0]
+        for keyword in self.keywords:
+            if keyword in outputs:
+                return True
+        return False

seagull/model/__init__.py

@@ -0,0 +1 @@
+from .language_model.seagull_llama import SeagullLlamaForCausalLM, SeagullConfig

seagull/model/consolidate.py

@@ -0,0 +1,26 @@
+
+import argparse
+
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from seagull.model import *
+from seagull.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)

seagull/model/language_model/seagull_llama.py

@@ -0,0 +1,128 @@
+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 ..seagull_arch import SeagullMetaModel, SeagullMetaForCausalLM
+from ..layer import MaskExtractor
+
+class SeagullConfig(LlamaConfig):
+    model_type = "seagull"
+
+class SeagullLlamaModel(SeagullMetaModel, LlamaModel):
+    config_class = SeagullConfig
+
+    def __init__(self, config: LlamaConfig):
+        super(SeagullLlamaModel, self).__init__(config)
+
+class SeagullLlamaForCausalLM(LlamaForCausalLM, SeagullMetaForCausalLM):
+    config_class = SeagullConfig
+
+    def __init__(self, config):
+        super(LlamaForCausalLM, self).__init__(config)
+        self.model = SeagullLlamaModel(config)
+
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.mask_extractor = MaskExtractor()
+
+        self.post_init()
+
+    def get_model(self):
+        return self.model
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        img_metas = None,
+        masks = 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,
+        preprocessed_img_dict = None,
+        return_dict: Optional[bool] = None,
+        cropped_img: Optional[torch.FloatTensor] = 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, masks, attention_mask, past_key_values, labels, images, preprocessed_img_dict=preprocessed_img_dict, cropped_img=cropped_img)
+
+        if inputs_embeds is not None:
+            inputs_embeds = inputs_embeds.bfloat16()
+  
+        self.model = self.model.bfloat16()
+        
+        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]
+        self.lm_head = self.lm_head.to(hidden_states.dtype)
+        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("seagull", SeagullConfig)
+AutoModelForCausalLM.register(SeagullConfig, SeagullLlamaForCausalLM)

seagull/model/layer.py

@@ -0,0 +1,250 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Optional, Tuple, Type, Any
+from torch import Tensor
+import math
+import numpy as np
+from einops import rearrange
+
+class MLP(nn.Module):
+
+    def __init__(self, input_dim: int, hidden_dim: int, output_dim: int,
+                 num_layers: int) -> None:
+        super().__init__()
+        self.num_layers = num_layers
+        h = [hidden_dim] * (num_layers - 1)
+        self.layers = nn.ModuleList(
+            nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim]))
+
+    def forward(self, x):
+        for i, layer in enumerate(self.layers):
+            x = F.relu(layer(x)) if i < self.num_layers - 1 else layer(x)
+        return x
+
+class MaskExtractor(nn.Module): # Mask-based Feature Extractor
+    def __init__(self, mask_shape=112, embed_dim=1024, out_dim=4096, num_heads=8, mlp_dim=2048, downsample_rate=2, skip_first_layer_pe=False):
+        super(MaskExtractor, self).__init__()
+        self.mask_shape = mask_shape
+        self.mask_pooling = MaskPooling()
+        self.feat_linear = nn.Linear(embed_dim, out_dim)
+        self.cross_feat_linear = nn.Linear(embed_dim, out_dim)
+        self.mask_linear = MLP(mask_shape*mask_shape, embed_dim, out_dim, 3)
+
+        self.feature_name = ['res2', 'res3', 'res4', 'res5']
+
+        self.cross_att_res = CrossAttention(
+            embedding_dim=embed_dim,
+            num_heads=num_heads,
+            mlp_dim=mlp_dim,
+            douwnsample_rate=downsample_rate,
+            skip_first_layer_pe=skip_first_layer_pe
+        )
+        
+        self.res2 = nn.Linear(192, 1024)
+        self.res3 = nn.Linear(384, 1024)
+        self.res4 = nn.Linear(768, 1024)
+        self.res5 = nn.Linear(1536, 1024)
+
+        self.g_res2 = nn.Linear(16384, 1024) # h * w
+        self.g_res3 = nn.Linear(4096, 1024)
+        self.g_res4 = nn.Linear(1024, 1024)
+        self.g_res5 = nn.Linear(256, 1024)
+
+        self.final_mlp = nn.Linear(2 * out_dim, out_dim)
+        
+        self.global_vit = nn.Sequential(
+            nn.Conv2d(3, 5, 1),
+            nn.GELU(),
+            nn.AvgPool2d(4, 4),
+
+            nn.Conv2d(5, 1, 1),
+            nn.GELU(),
+            nn.AvgPool2d(4, 4),
+        )
+        self.is_first = 0
+        
+        self.sa = Attention(32 * 32, num_heads) # self-attention
+        self.mlp =  MLP(32 * 32, 512, out_dim, 3)
+
+    def cal_globa_local(self, mask_feat_raw, feat_new, res, g_res, cross_attention):
+        mask_feat_flatten = mask_feat_raw.to(device=res.weight.device, dtype=res.weight.dtype) 
+        mask_feat = res(mask_feat_flatten) # (b, q, 1024)
+        
+        feat_new = feat_new.to(device=g_res.weight.device, dtype=g_res.weight.dtype) 
+        all_feat_new = g_res(feat_new) # (b, c, 1024)
+        global_mask = cross_attention(mask_feat, all_feat_new)
+        return mask_feat, global_mask
+                    
+    def forward(self, feats, masks, cropped_img):
+        global_features = []
+        local_features = []
+        num_imgs = len(masks)
+
+        for idx in range(num_imgs):
+            mask = masks[idx].unsqueeze(0).float() #(1, q, h, w)
+            cropped_ = cropped_img[idx] # (q, 3, h, w)
+
+            num_feats = len(self.feature_name)
+            mask_feats = mask.new_zeros(num_feats, mask.shape[1], 1024)
+            global_masks = mask.new_zeros(num_feats, mask.shape[1], 1024)
+
+            for i, name in enumerate(self.feature_name):
+                feat = feats[name][idx].unsqueeze(0) 
+                feat = feat.to(mask.dtype)
+                
+                mask_feat_raw = self.mask_pooling(feat, mask)
+                feat_new = rearrange(feat, 'b c h w -> b c (h w)')
+                
+                mask_feat, global_mask = self.cal_globa_local(mask_feat_raw, feat_new, res=getattr(self, name),  g_res=getattr(self, 'g_{}'.format(name)),  cross_attention=getattr(self,"cross_att_res"))
+
+                mask_feats[i] = mask_feat.squeeze(0) # (q, 1024)
+                global_masks[i] = global_mask.squeeze(0)
+            mask_feats = mask_feats.sum(0) # (1, q, 1024)
+            global_masks = global_masks.sum(0)  # (1, q, 1024)
+            global_masks = global_masks.to(device=self.cross_feat_linear.weight.device, dtype=self.cross_feat_linear.weight.dtype)
+            global_masks_linear = self.cross_feat_linear(global_masks)
+            mask_feats = mask_feats.to(device=self.feat_linear.weight.device, dtype=self.feat_linear.weight.dtype) 
+            mask_feats_linear = self.feat_linear(mask_feats) #(1, q, 4096)
+
+            query_feat = self.final_mlp(torch.cat((global_masks_linear, mask_feats_linear), dim=-1))
+            global_features.append(query_feat) # global
+
+            cropped_ = cropped_.to(device=self.feat_linear.weight.device, dtype=self.feat_linear.weight.dtype) 
+            global_features = self.global_vit(cropped_).to(device=self.feat_linear.weight.device, dtype=self.feat_linear.weight.dtype)  # q, 1, 32, 32
+            global_features = global_features.reshape(-1, 1, 32 * 32) # q, 1, 32 * 32
+            pos_feat = self.mlp(self.sa(global_features, global_features, global_features).squeeze(1))  # q, output
+
+            local_features.append(pos_feat) #(imgs_num, 1, q, 4096) # local
+        
+        return global_features, local_features
+    
+class MaskPooling(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x, mask):
+
+        if not x.shape[-2:] == mask.shape[-2:]:
+            # reshape mask to x
+            mask = F.interpolate(mask, size=x.shape[-2:], mode='bilinear', align_corners=False)
+
+        mask = (mask > 0).to(mask.dtype)
+        denorm = mask.sum(dim=(-1, -2), keepdim=True) + 1e-8
+
+        mask_pooled_x = torch.einsum(
+            "bchw,bqhw->bqc",
+            x,
+            mask / denorm,
+        )
+        return mask_pooled_x
+    
+
+class CrossAttention(nn.Module):
+    def __init__(
+            self, 
+            embedding_dim: int,
+            num_heads: int,
+            mlp_dim: int = 2048,
+            douwnsample_rate: int = 2, 
+            activation: Type[nn.Module] = nn.ReLU,
+            skip_first_layer_pe: bool = False
+        ) -> None:
+        super().__init__()
+        self.embedding_dim = embedding_dim
+        self.num_heads =num_heads
+        self.self_attn = Attention(embedding_dim, num_heads) # self-attention
+        self.skip_first_layer_pe = skip_first_layer_pe
+        self.norm1 = nn.LayerNorm(embedding_dim)
+
+        # cross-attention
+        self.cross_attn = Attention(embedding_dim, num_heads, downsample_rate=douwnsample_rate)
+        self.norm2 = nn.LayerNorm(embedding_dim)
+
+        self.mlp = MLPBlock(embedding_dim, mlp_dim, activation) # MLP
+
+    def forward(self, queries, keys):
+        attn_out = self.self_attn(queries, queries, queries)
+        queries = queries + attn_out
+        queries = self.norm1(queries)
+
+        attn_out = self.cross_attn(q=queries, k=keys, v=keys)
+        queries = attn_out + queries
+        queries = self.norm2(queries)
+
+        # MLP
+        mlp_out = self.mlp(queries)
+        queries = queries + mlp_out
+        return queries
+
+class MLPBlock(nn.Module):
+    def __init__(
+        self,
+        embedding_dim: int,
+        mlp_dim: int,
+        act: Type[nn.Module] = nn.GELU,
+    ) -> None:
+        super().__init__()
+        self.lin1 = nn.Linear(embedding_dim, mlp_dim)
+        self.lin2 = nn.Linear(mlp_dim, embedding_dim)
+        self.act = act()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.lin2(self.act(self.lin1(x)))
+
+class Attention(nn.Module):
+    """
+    An attention layer that allows for downscaling the size of the embedding
+    after projection to queries, keys, and values.
+    """
+
+    def __init__(
+        self,
+        embedding_dim: int,
+        num_heads: int,
+        downsample_rate: int = 1,
+    ) -> None:
+        super().__init__()
+        self.embedding_dim = embedding_dim
+        self.internal_dim = embedding_dim // downsample_rate
+        self.num_heads = num_heads
+        assert self.internal_dim % num_heads == 0, "num_heads must divide embedding_dim."
+
+        self.q_proj = nn.Linear(embedding_dim, self.internal_dim)
+        self.k_proj = nn.Linear(embedding_dim, self.internal_dim)
+        self.v_proj = nn.Linear(embedding_dim, self.internal_dim)
+        self.out_proj = nn.Linear(self.internal_dim, embedding_dim)
+
+    def _separate_heads(self, x: Tensor, num_heads: int) -> Tensor:
+        b, n, c = x.shape
+        x = x.reshape(b, n, num_heads, c // num_heads)
+        return x.transpose(1, 2)  # B x N_heads x N_tokens x C_per_head
+
+    def _recombine_heads(self, x: Tensor) -> Tensor:
+        b, n_heads, n_tokens, c_per_head = x.shape
+        x = x.transpose(1, 2)
+        return x.reshape(b, n_tokens, n_heads * c_per_head)  # B x N_tokens x C
+
+    def forward(self, q: Tensor, k: Tensor, v: Tensor) -> Tensor:
+        # Input projections
+        q = self.q_proj(q)
+        k = self.k_proj(k)
+        v = self.v_proj(v)
+
+        # Separate into heads
+        q = self._separate_heads(q, self.num_heads)
+        k = self._separate_heads(k, self.num_heads)
+        v = self._separate_heads(v, self.num_heads)
+
+        # Attention
+        _, _, _, c_per_head = q.shape
+        attn = q @ k.permute(0, 1, 3, 2)  # B x N_heads x N_tokens x N_tokens
+        attn = attn / math.sqrt(c_per_head)
+        attn = torch.softmax(attn, dim=-1)
+
+        # Get output
+        out = attn @ v
+        out = self._recombine_heads(out)
+        out = self.out_proj(out)
+
+        return out

seagull/model/multimodal_encoder/builder.py

@@ -0,0 +1,7 @@
+import os
+from .clip_encoder import CLIPVisionTower
+
+
+def build_vision_tower(vision_tower_cfg, delay_load=False):
+
+    return CLIPVisionTower(args=vision_tower_cfg)

seagull/model/multimodal_encoder/clip.py

@@ -0,0 +1,40 @@
+import torch
+import torch.nn.functional as F
+import torch.nn as nn
+
+from open_clip.model import _build_vision_tower
+
+
+class CLIP(nn.Module):
+    def __init__(self):
+        super().__init__()
+        model_name = 'convnext_large'
+
+        vision_cfg = {'timm_model_name': model_name, 'timm_model_pretrained': False, 'timm_pool': '', 'timm_proj': 'mlp', 'timm_drop': 0.0, 'timm_drop_path': 0.1, 'image_size': 320}
+        self.visual = _build_vision_tower(embed_dim=768, vision_cfg=vision_cfg, quick_gelu=False)
+
+        self.eval()
+        self.freeze_everything()
+
+    def freeze_everything(self):
+        for param in self.visual.parameters():
+            param.requires_grad = False
+
+    def extract_features(self, x):
+        out = {}
+        x = x.to(self.visual.trunk.stem.state_dict()['1.bias'].dtype)
+        x = self.visual.trunk.stem(x)
+        out['stem'] = x.contiguous() 
+        for i in range(4):
+            x = self.visual.trunk.stages[i](x)
+            out[f'res{i+2}'] = x.contiguous() 
+        
+        x = self.visual.trunk.norm_pre(x)
+        out['clip_vis_dense'] = x.contiguous()
+        return out
+
+    def forward(self, x):
+        self.eval()
+        with torch.no_grad():
+            return self.extract_features(x)
+    

seagull/model/multimodal_encoder/clip_encoder.py

@@ -0,0 +1,59 @@
+import torch
+import torch.nn as nn
+
+from transformers import CLIPImageProcessor
+from .clip import CLIP
+
+class CLIPVisionTower(nn.Module):
+    def __init__(self, args, img_size=512, delay_load=False):
+        super().__init__()
+
+        # test
+        if hasattr(args, 'mm_vision_tower'):
+            self.clip_model = args.mm_vision_tower
+        else: # train
+            self.clip_model = args.vision_tower
+        self.is_loaded = False
+        self.img_size = img_size
+
+        if not delay_load:
+            self.load_model()
+
+    def load_model(self):
+        self.image_processor = CLIPImageProcessor(do_resize=True, size={"shortest_edge":self.img_size}, resample=3,  do_center_crop=True, crop_size={"height": self.img_size, "width": self.img_size},
+                                                  do_rescale=True, rescale_factor=0.00392156862745098, do_normalize=True, image_mean=[0.48145466, 0.4578275, 0.40821073],
+                                                  image_std=[0.26862954, 0.26130258, 0.27577711], do_convert_rgb=True, )
+
+        self.vision_tower = CLIP()
+
+        self.vision_tower.load_state_dict(torch.load(self.clip_model),strict=False)
+        
+        self.is_loaded = True
+
+    @torch.no_grad()
+    def forward(self, images):
+        if type(images) is list:
+            image_features = []
+            image_features_dict = []
+            for image in images:
+                image_feature_dict = self.vision_tower(image.unsqueeze(0))
+                image_features_dict.append(image_feature_dict)
+                image_feature = image_feature_dict['res4']
+                image_feature = image_feature.reshape(*image_feature.shape[:2],-1).permute(0,2,1)
+                image_features.append(image_feature)
+        else:
+            # print(images.device)
+            # print(self.vision_tower.device)
+            image_features_dict = self.vision_tower(images)
+            image_features = image_features_dict['res4']
+            image_features = image_features.reshape(*image_features.shape[:2],-1).permute(0,2,1)
+
+        return image_features, image_features_dict
+
+    @property
+    def dtype(self):
+        return self.vision_tower.dtype
+
+    @property
+    def device(self):
+        return self.vision_tower.device

seagull/model/multimodal_projector/builder.py

@@ -0,0 +1,52 @@
+import torch
+import torch.nn as nn
+import re
+
+
+class IdentityMap(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x, *args, **kwargs):
+        return x
+
+    @property
+    def config(self):
+        return {"mm_projector_type": 'identity'}
+
+
+class SimpleResBlock(nn.Module):
+    def __init__(self, channels):
+        super().__init__()
+        self.pre_norm = nn.LayerNorm(channels)
+
+        self.proj = nn.Sequential(
+            nn.Linear(channels, channels),
+            nn.GELU(),
+            nn.Linear(channels, channels)
+        )
+    def forward(self, x):
+        x = self.pre_norm(x)
+        return x + self.proj(x)
+
+
+def build_vision_projector(config, delay_load=False, **kwargs):
+    mm_hidden_size = getattr(config, 'mm_hidden_size', 768)
+    projector_type = getattr(config, 'mm_projector_type', 'linear')
+
+    if projector_type == 'linear':
+        return nn.Linear(mm_hidden_size, config.hidden_size)
+
+    mlp_gelu_match = re.match(r'^mlp(\d+)x_gelu$', projector_type)
+    if mlp_gelu_match:
+        mlp_depth = int(mlp_gelu_match.group(1))
+        modules = [nn.Linear(mm_hidden_size, config.hidden_size)]
+        for _ in range(1, mlp_depth):
+            modules.append(nn.GELU())
+            modules.append(nn.Linear(config.hidden_size, config.hidden_size))
+        return nn.Sequential(*modules)
+
+    if projector_type == 'identity':
+        return IdentityMap()
+
+    raise ValueError(f'Unknown projector type: {projector_type}')

seagull/model/seagull_arch.py

@@ -0,0 +1,281 @@
+from abc import ABC, abstractmethod
+
+import torch
+
+from .multimodal_encoder.builder import build_vision_tower
+from .multimodal_projector.builder import build_vision_projector
+
+from seagull.constants import IGNORE_INDEX, IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_PATCH_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN
+
+
+class SeagullMetaModel:
+
+    def __init__(self, config):
+        super(SeagullMetaModel, self).__init__(config)
+
+        if hasattr(config, "mm_vision_tower"):
+            self.vision_tower = build_vision_tower(config, delay_load=False)
+            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
+        pretrain_mm_mlp_adapter = model_args.pretrain_mm_mlp_adapter
+
+        if not hasattr(self.config, "mm_vision_tower"):
+            self.config.mm_vision_tower = vision_tower
+
+            vision_tower = build_vision_tower(model_args)
+
+            if fsdp is not None and len(fsdp) > 0:
+                self.vision_tower = [self.vision_tower]
+            else:
+                self.vision_tower = vision_tower
+
+            self.config.use_mm_proj = True
+            self.config.mm_projector_type = getattr(model_args, 'mm_projector_type', 'linear')
+ 
+            self.mm_projector = build_vision_projector(self.config)
+
+        if pretrain_mm_mlp_adapter is not None:
+            print("***********load projector_weights********")
+            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 SeagullMetaForCausalLM(ABC):
+    def __init__(self):
+        super(SeagullMetaForCausalLM, self).__init__()
+
+    @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, image_features_dict = self.get_model().get_vision_tower()(images)
+        self.get_model().mm_projector.to(device=image_features.device, dtype=image_features.dtype)
+        image_features = self.get_model().mm_projector(image_features)
+        return image_features, image_features_dict
+
+    def prepare_inputs_labels_for_multimodal(
+        self, input_ids, masks, attention_mask, past_key_values, labels, images, preprocessed_img_dict=None, cropped_img=None
+    ):
+        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 preprocessed_img_dict is not None:
+            image_features, image_features_dict = images, preprocessed_img_dict
+        else:
+            if type(images) is list or images.ndim == 5:
+                concat_images = torch.cat([image for image in images], dim=0)
+                image_features, image_features_dict = 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).to(concat_images.device) for x in image_features]
+            else:
+                image_features, image_features_dict = self.encode_images(images)
+        
+
+        mask_feats, pos_feats = self.mask_extractor(image_features_dict, masks, cropped_img=cropped_img)
+
+        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):
+                    mask_idx = torch.nonzero(cur_input_ids==self.tokenizer.convert_tokens_to_ids(['<global>'])[0])
+    
+                    _l = 0
+                    for i, idx in enumerate(mask_idx):
+                        cur_new_input_embeds.append(self.get_model().embed_tokens(cur_input_ids[_l:idx[0]]).detach())
+                        ## mask
+                        cur_new_input_embeds.append(mask_feats[batch_idx][i:i+1].detach())
+                        ## pos
+                        cur_new_input_embeds.append(pos_feats[batch_idx][i:i+1].detach())
+                        if labels is not None:
+                            cur_labels[idx[0]:idx[0]+2] = torch.full((2,), IGNORE_INDEX, device=labels.device, dtype=labels.dtype)
+                        _l = idx[0]+2
+                    if _l< len(cur_input_ids):
+                        cur_new_input_embeds.append(self.get_model().embed_tokens(cur_input_ids[_l:]).detach())
+
+                else:
+                    
+                    mask_idx = torch.nonzero(cur_input_ids==self.tokenizer.convert_tokens_to_ids(['<global>'])[0])
+                    assert len(mask_idx) == len(mask_feats[batch_idx]), "mask num not equal to mask feats"
+                   
+                    _l = 0
+                    for i, idx in enumerate(mask_idx):
+                        cur_raw_new_input_embeds = self.get_model().embed_tokens(cur_input_ids[_l:idx[0]])
+                        cur_new_input_embeds.append(cur_raw_new_input_embeds)
+                        ## mask
+                        cur_new_input_embeds.append(mask_feats[batch_idx][i:i+1].to(cur_raw_new_input_embeds.dtype))
+                        ## pos
+                        cur_new_input_embeds.append(pos_feats[batch_idx][i:i+1].to(cur_raw_new_input_embeds.dtype))
+
+                        if labels is not None:
+                            cur_labels[idx[0]:idx[0]+2] = torch.full((2,), IGNORE_INDEX, device=labels.device, dtype=labels.dtype)
+
+                        _l = idx[0]+2
+                    if _l< len(cur_input_ids):
+                        cur_new_input_embeds.append(self.get_model().embed_tokens(cur_input_ids[_l:]))
+
+                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))
+
+        mask_tokens = ['<global>', '<pos>']
+        num_new_tokens = tokenizer.add_tokens(mask_tokens, special_tokens=True)
+
+        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')
+                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
+        
+        for m in self.modules():
+            m.tokenizer = tokenizer