add model

model/__init__.py

+from .chat import chat
+from .infer_util import *
+from .blip2 import BlipImageEvalProcessor

model/blip2.py

+import torch
+import torch.nn as nn
+
+from sat.model import ViTModel, BaseModel
+from sat.model import BaseMixin
+from sat import AutoModel
+from copy import deepcopy
+from torchvision import transforms
+from torchvision.transforms.functional import InterpolationMode
+
+class LNFinalyMixin(BaseMixin):
+    def __init__(self, hidden_size):
+        super().__init__()
+        self.ln_vision = nn.LayerNorm(hidden_size)
+
+    def final_forward(self, logits, **kw_args):
+        return self.ln_vision(logits)
+
+class EVAViT(ViTModel):
+    def __init__(self, args, transformer=None, parallel_output=True, **kwargs):
+        super().__init__(args, transformer=transformer, parallel_output=parallel_output, **kwargs)
+        self.del_mixin("cls")
+        self.add_mixin("cls", LNFinalyMixin(args.hidden_size))
+    
+    def forward(self, image):
+        batch_size = image.size(0)
+        input_ids = torch.zeros(batch_size, 1, dtype=torch.long, device=image.device)
+        attention_mask = torch.tensor([[1.]], dtype=image.dtype, device=image.device)
+        return super().forward(input_ids=input_ids, position_ids=None, attention_mask=attention_mask, image=image)
+
+class QFormer(BaseModel):
+    def __init__(self, args, transformer=None, parallel_output=True, **kwargs):
+        super().__init__(args, transformer=transformer, parallel_output=parallel_output, activation_func=nn.functional.gelu, **kwargs)
+        self.transformer.position_embeddings = None
+    
+    def final_forward(self, logits, **kw_args):
+        return logits
+
+    def position_embedding_forward(self, position_ids, **kw_args):
+        return None
+    
+    def forward(self, encoder_outputs):
+        batch_size = encoder_outputs.size(0)
+        input_ids = torch.arange(32, dtype=torch.long, device=encoder_outputs.device).unsqueeze(0).expand(batch_size, -1)
+        attention_mask = torch.tensor([[1.]], dtype=encoder_outputs.dtype, device=encoder_outputs.device)
+        cross_attention_mask = torch.tensor([[1.]], dtype=encoder_outputs.dtype, device=encoder_outputs.device)
+        return super().forward(input_ids=input_ids, position_ids=None, attention_mask=attention_mask, encoder_outputs=encoder_outputs, cross_attention_mask=cross_attention_mask)
+
+
+class BLIP2(torch.nn.Module):
+    def __init__(self, eva_args, qformer_args, vit=None, qformer=None, **kwargs):
+        super().__init__()
+        if vit is not None:
+            self.vit = vit
+        else:
+            self.vit = EVAViT(EVAViT.get_args(**eva_args))
+        if qformer is not None:
+            self.qformer = qformer
+        else:
+            self.qformer = QFormer(QFormer.get_args(**qformer_args))
+        
+        self.glm_proj = nn.Linear(768, 4096).to(self.qformer.parameters().__next__().device).to(self.qformer.parameters().__next__().dtype)
+
+    def forward(self, image, **kwargs):
+        enc = self.vit(image)[0]
+        out = self.qformer(enc)[0]
+        return self.glm_proj(out)
+    
+class BlipImageBaseProcessor():
+    def __init__(self, mean=None, std=None):
+        if mean is None:
+            mean = (0.48145466, 0.4578275, 0.40821073)
+        if std is None:
+            std = (0.26862954, 0.26130258, 0.27577711)
+
+        self.normalize = transforms.Normalize(mean, std)
+
+class BlipImageEvalProcessor(BlipImageBaseProcessor):
+    def __init__(self, image_size=384, mean=None, std=None):
+        super().__init__(mean=mean, std=std)
+
+        self.transform = transforms.Compose(
+            [
+                transforms.Resize(
+                    (image_size, image_size), interpolation=InterpolationMode.BICUBIC
+                ),
+                transforms.ToTensor(),
+                self.normalize,
+            ]
+        )
+
+    def __call__(self, item):
+        return self.transform(item)

model/chat.py

+# -*- encoding: utf-8 -*-
+'''
+@File    :   chat.py
+@Time    :   2023/05/08 19:10:08
+@Author  :   Ming Ding 
+@Contact :   dm18@mails.tsinghua.edu.cn
+'''
+
+import os
+import sys
+import re
+from functools import partial
+from typing import Optional, Tuple, Union, List, Callable, Dict, Any
+import requests
+from PIL import Image
+from io import BytesIO
+
+import torch
+from sat.generation.autoregressive_sampling import filling_sequence, BaseStrategy
+
+from .blip2 import BlipImageEvalProcessor
+
+def get_masks_and_position_ids_glm(seq, mask_position, context_length):
+    '''GLM model, different from GPT.
+    Args:
+        seq: torch.IntTensor, [seq_len]
+        mask_position: int, the position of the masked place.
+        context_length: int, the length of context.
+    Returns:
+        tokens: torch.IntTensor, [1, seq_len]
+        attention_mask: torch.FloatTensor, [1, seq_len, seq_len]
+        position_ids: torch.IntTensor, [2, seq_len]
+    '''
+    tokens = seq.unsqueeze(0)
+
+    attention_mask = torch.ones((1, len(seq), len(seq)), device=tokens.device)
+    attention_mask.tril_()
+    attention_mask[..., :context_length] = 1
+    attention_mask.unsqueeze_(1)
+
+    # 2D position ids
+    position_ids = torch.zeros(2, len(seq), device=tokens.device, dtype=torch.long)
+    torch.arange(0, context_length, out=position_ids[0, :context_length])
+    position_ids[0, context_length:] = mask_position
+    torch.arange(1, len(seq) - context_length + 1, out=position_ids[1, context_length:])
+
+    position_ids = position_ids.unsqueeze(0)
+    return tokens, attention_mask, position_ids
+
+def process_response(response):
+    response = response.strip()
+    response = response.replace("[[训练时间]]", "2023年")
+    punkts = [
+        [",", "，"],
+        ["!", "！"],
+        [":", "："],
+        [";", "；"],
+        ["\?", "？"],
+    ]
+    for item in punkts:
+        response = re.sub(r"([\u4e00-\u9fff])%s" % item[0], r"\1%s" % item[1], response)
+        response = re.sub(r"%s([\u4e00-\u9fff])" % item[0], r"%s\1" % item[1], response)
+    return response
+
+def process_image(text, image=None):
+    '''Process image in text.
+    Args:
+        text: str, text.
+        image: Optional, image path / url / PIL image.
+    '''
+    image_position = text.rfind("<img>") + 5
+    # extract path from <img></img> using re
+    image_path = re.findall(r"<img>(.*?)</img>", text)
+    image_path = image_path[-1] if image_path[-1] else None
+    if image_path is not None:
+        assert image is None, "image and image_path cannot be both not None."
+        text = text.replace(image_path, "")
+        image_path = image_path.strip()
+        # url
+        if image_path.startswith("http"):
+            response = requests.get(image_path, timeout=10)
+            image = Image.open(BytesIO(response.content))
+        # local path
+        else:
+            image = Image.open(image_path)
+    if image is not None and isinstance(image, Image.Image):
+        processor = BlipImageEvalProcessor(224)
+        image = processor(image.convert('RGB'))
+        image = image.unsqueeze(0)
+    return text, image_position, image
+
+
+def chat(image_path, model, tokenizer, 
+        query: str, history: List[Tuple[str, str]] = None, image: Image = None,
+        max_length: int = 1024, top_p=0.7, top_k=30, temperature=0.95, repetition_penalty=1.2,
+        invalid_slices=[], english=False
+        ):
+    if not history:
+        history = []
+    if image_path:
+        prompt = "<img>{}</img>".format(image_path if image_path else "")
+    else:
+        prompt = "<img></img>"
+    if english:
+        for i, (old_query, response) in enumerate(history):
+            prompt += "Q:{}\nA:{}\n".format(old_query, response)
+        prompt += "Q:{}\nA:".format(query)
+    else:
+        for i, (old_query, response) in enumerate(history):
+            prompt += "问：{}\n答：{}\n".format(old_query, response)
+        prompt += "问：{}\n答：".format(query)
+    # ---------------
+    # tokenizer, this is an example of huggingface tokenizer.
+    # input str, output['input_ids'] = tensor([[tokenized str, gmask, sop]])
+    prompt, image_position, torch_image = process_image(prompt, image=image)
+    if torch_image is not None:
+        torch_image = torch_image.to(next(model.parameters()).dtype).to(next(model.parameters()).device)
+    if image_position < 5: # no image
+        inputs = tokenizer([prompt], return_tensors="pt").to(model.parameters().__next__().device)['input_ids'][0]
+        pre_image = 0
+    else:
+        input0 = tokenizer.encode(prompt[:image_position], add_special_tokens=False)
+        input1 = [tokenizer.pad_token_id] * model.image_length
+        input2 = tokenizer.encode(prompt[image_position:], add_special_tokens=False)
+        inputs = sum([input0, input1, input2], [])
+        inputs = torch.tensor(tokenizer.build_inputs_with_special_tokens(inputs)).to(model.parameters().__next__().device)
+        pre_image = len(input0)
+    # ---------------
+    # Next, we manually set the format to keep flexibility.
+    mask_position = len(inputs) - 2
+    context_length = len(inputs) - 1 # all before sop
+    get_func = partial(get_masks_and_position_ids_glm, mask_position=mask_position, context_length=context_length)
+    seq = torch.cat(
+        [inputs, torch.tensor([-1]*(max_length-len(inputs)), device=inputs.device)], dim=0
+    )
+    # ---------------
+    # from sat.generation.sampling_strategies import BeamSearchStrategy
+    # strategy = BeamSearchStrategy(num_beams, length_penalty=1., prefer_min_length=5, end_tokens=[tokenizer.eos_token_id], consider_end=True, no_repeat_ngram_size=5, stop_n_iter_unchanged=30, temperature=temperature, top_p=top_p, top_k=60, repetition_penalty=1.1)
+    strategy = BaseStrategy(temperature=temperature, top_p=top_p, top_k=top_k, end_tokens=[tokenizer.eos_token_id],
+                            invalid_slices=invalid_slices, repetition_penalty=repetition_penalty)
+    output = filling_sequence(
+        model, seq,
+        batch_size=1,
+        get_masks_and_position_ids=get_func,
+        strategy=strategy,
+        pre_image=pre_image,
+        image=torch_image,
+    )[0] # drop memory
+    
+    # ---------------
+    # port from inference_glm.py, more general than chat mode
+    # clip -1s and fill back generated things into seq
+    if type(output) is not list:
+        output_list = output.tolist()
+    else:
+        output_list = output
+    for i in range(len(output_list)):
+        output = output_list[i]
+        if type(output) is not list:
+            output = output.tolist()
+        try:
+            unfinished = output.index(-1)
+        except ValueError:
+            unfinished = len(output)
+        if output[unfinished - 1] == tokenizer.eos_token_id:
+            unfinished -= 1
+        bog = output.index(tokenizer.bos_token_id)
+        output_list[i] = output[:mask_position] + output[bog + 1:unfinished] + output[mask_position + 1:bog]
+    # ---------------
+
+    response = tokenizer.decode(output_list[0])
+    sep = 'A:' if english else '答：'
+    response = process_response(response).split(sep)[-1].strip()
+    history = history + [(query, response)]
+    return response, history, torch_image

model/infer_util.py

+import os
+from PIL import Image
+from io import BytesIO
+import base64
+import re
+import argparse
+import torch
+from transformers import AutoTokenizer
+from sat.model.mixins import CachedAutoregressiveMixin
+from sat.quantization.kernels import quantize
+import hashlib
+from .visualglm import VisualGLMModel
+
+def get_infer_setting(gpu_device=0, quant=None):
+    os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu_device)
+    args = argparse.Namespace(
+        fp16=True,
+        skip_init=True,
+        device='cuda' if quant is None else 'cpu',
+    )
+    model, args = VisualGLMModel.from_pretrained('visualglm-6b', args)
+    model.add_mixin('auto-regressive', CachedAutoregressiveMixin())
+    assert quant in [None, 4, 8]
+    if quant is not None:
+        quantize(model.transformer, quant)
+    model.eval()
+    model = model.cuda()
+    tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True)
+    return model, tokenizer
+
+def is_chinese(text):
+    zh_pattern = re.compile(u'[\u4e00-\u9fa5]+')
+    return zh_pattern.search(text)
+
+def generate_input(input_text, input_image_prompt, history=[], input_para=None, image_is_encoded=True):
+    if not image_is_encoded:
+        image = input_image_prompt
+    else:
+        decoded_image = base64.b64decode(input_image_prompt)
+        image = Image.open(BytesIO(decoded_image))
+
+    input_data = {'input_query': input_text, 'input_image': image, 'history': history, 'gen_kwargs': input_para}
+    return input_data
+
+
+def process_image(image_encoded):
+    decoded_image = base64.b64decode(image_encoded)
+    image = Image.open(BytesIO(decoded_image))
+    image_hash = hashlib.sha256(image.tobytes()).hexdigest()
+    image_path = f'./examples/{image_hash}.png'
+    if not os.path.isfile(image_path):
+        image.save(image_path)
+    return os.path.abspath(image_path)  
\ No newline at end of file

model/visualglm.py

+import torch
+from sat.model.official import ChatGLMModel
+from sat.model.base_model import BaseMixin
+from copy import deepcopy
+import json
+from .blip2 import BLIP2
+
+from sat.resources.urls import MODEL_URLS
+MODEL_URLS['visualglm-6b'] = 'https://cloud.tsinghua.edu.cn/f/348b98dffcc940b6a09d/?dl=1'
+
+class ImageMixin(BaseMixin):
+    def __init__(self, args):
+        super().__init__()
+        self.args = deepcopy(args)
+        self.model = BLIP2(args.eva_args, args.qformer_args)
+
+    def word_embedding_forward(self, input_ids, output_cross_layer, **kw_args):
+        if kw_args["pre_image"] > input_ids.shape[1] or kw_args.get("image", None) is None:
+            return self.transformer.word_embeddings(input_ids)
+        image_emb = self.model(**kw_args)
+        # the image is inserted after 问：<img>, override 32 pads
+        pre_id, pads, post_id = torch.tensor_split(input_ids, [kw_args["pre_image"], kw_args["pre_image"]+self.args.image_length], dim=1)
+        pre_txt_emb = self.transformer.word_embeddings(pre_id)
+        post_txt_emb = self.transformer.word_embeddings(post_id)
+        return torch.cat([pre_txt_emb, image_emb, post_txt_emb], dim=1)
+
+class VisualGLMModel(ChatGLMModel):
+    def __init__(self, args, transformer=None, **kwargs):
+        super().__init__(args, transformer=transformer, **kwargs)
+        self.image_length = args.image_length
+        self.add_mixin("eva", ImageMixin(args))
+
+    @classmethod
+    def add_model_specific_args(cls, parser):
+        group = parser.add_argument_group('VisualGLM', 'VisualGLM Configurations')
+        group.add_argument('--image_length', type=int, default=32)
+        group.add_argument('--eva_args', type=json.loads, default={})
+        group.add_argument('--qformer_args', type=json.loads, default={})
+        return super().add_model_specific_args(parser)
+