chat_demo_minimonkey.py

# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
#
# 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 argparse
import paddle
import paddle.vision.transforms as T
from PIL import Image
from paddlemix.models.internvl2.internlm2 import InternLM2Tokenizer
from paddlemix.models.internvl2.internvl_chat import MiniMonkeyChatModel
from paddlemix.datasets.internvl_dataset import dynamic_preprocess, dynamic_preprocess2

IMAGENET_MEAN = (0.485, 0.456, 0.406)
IMAGENET_STD = (0.229, 0.224, 0.225)


def build_transform(input_size):
    MEAN, STD = IMAGENET_MEAN, IMAGENET_STD
    transform = T.Compose([
        # T.Lambda(lambda img: img.convert('RGB') if img.mode != 'RGB' else img),
        T.Resize((input_size, input_size), interpolation='bicubic'),
        T.ToTensor(),
        T.Normalize(mean=MEAN, std=STD)
    ])
    return transform


def load_image(image_file, input_size=448, min_num=1, max_num=12):
    image = Image.open(image_file).convert('RGB')
    transform = build_transform(input_size=input_size)
    images, target_aspect_ratio = dynamic_preprocess(image, image_size=input_size, use_thumbnail=True, min_num=min_num, max_num=max_num, return_target_aspect_ratio=True)
    pixel_values = [transform(image) for image in images]
    pixel_values = paddle.stack(pixel_values)
    return pixel_values, target_aspect_ratio


def load_image2(image_file, input_size=448, min_num=1, max_num=12, target_aspect_ratio=None):
    image = Image.open(image_file).convert('RGB')
    transform = build_transform(input_size=input_size)
    images = dynamic_preprocess2(image, image_size=input_size, use_thumbnail=True, min_num=min_num, max_num=max_num, prior_aspect_ratio=target_aspect_ratio)
    pixel_values = [transform(image) for image in images]
    pixel_values = paddle.stack(pixel_values)
    return pixel_values


def main(args):
    assert args.image_path is not None and args.image_path != 'None'
    pixel_values, target_aspect_ratio = load_image(args.image_path, min_num=4, max_num=12)
    pixel_values = pixel_values.to(paddle.bfloat16)
    pixel_values2 = load_image2(args.image_path, min_num=3, max_num=7, target_aspect_ratio=target_aspect_ratio).to(paddle.bfloat16)
    pixel_values = paddle.concat([pixel_values2[:-1], pixel_values[:-1], pixel_values2[-1:]], 0)
    pixel_values = pixel_values.to(paddle.bfloat16)

    args.text = '<image>\n' + args.text

    # init model and tokenizer
    MODEL_PATH = args.model_name_or_path
    tokenizer = InternLM2Tokenizer.from_pretrained(MODEL_PATH)
    # TODO:
    tokenizer.added_tokens_encoder = {'<unk>': 0, '<s>': 1, '</s>': 2, '<|plugin|>': 92538, '<|interpreter|>': 92539, '<|action_end|>': 92540, '<|action_start|>': 92541, '<|im_end|>': 92542, '<|im_start|>': 92543, '<img>': 92544, '</img>': 92545, '<IMG_CONTEXT>': 92546, '<quad>': 92547, '</quad>': 92548, '<ref>': 92549, '</ref>': 92550, '<box>': 92551, '</box>': 92552}
    tokenizer.added_tokens_decoder = {v: k for k, v in tokenizer.added_tokens_encoder.items()}

    print('tokenizer:\n', tokenizer)
    print('len(tokenizer): ', len(tokenizer))

    model = MiniMonkeyChatModel.from_pretrained(MODEL_PATH).eval()

    generation_config = dict(max_new_tokens=512, do_sample=False)

    with paddle.no_grad():
        response, history = model.chat(tokenizer, pixel_values, target_aspect_ratio, args.text, generation_config, use_scm=True, history=None, return_history=True)
        print(f'User: {args.text}\nAssistant: {response}')


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--model_name_or_path",
        type=str,
        default="HUST-VLRLab/Mini-Monkey",
        help="pretrained ckpt and tokenizer",
    )
    parser.add_argument("--image_path", type=str, default=None)
    parser.add_argument("--text", type=str, default='Please describe the image shortly.', required=True)
    args = parser.parse_args()
    main(args)