maga_transformer/models/minicpmv/minicpmv.py

import json
import os
from typing import Any, Dict, List

import torch
from PIL import Image
from transformers import AutoTokenizer, AutoProcessor
from maga_transformer.config.gpt_init_model_parameters import \
    GptInitModelParameters
from maga_transformer.distribute.worker_info import g_parallel_info
from maga_transformer.model_factory_register import register_model
from maga_transformer.models.qwen_v2 import QWenV2, QWenV2Weight
from maga_transformer.models.multimodal.multimodal_mixin import MultiModalMixin, BaseVitWeights
from maga_transformer.models.multimodal.multimodal_common import MultiModalEmbeddingInterface, mm_lock 
from maga_transformer.utils.multimodal_util import MMUrlType
from maga_transformer.models.minicpmv.modeling_navit_siglip import SiglipVisionTransformer, SiglipVisionConfig
from maga_transformer.models.minicpmv.resampler import Resampler
from maga_transformer.models.multimodal.multimodal_mixin import BaseVitWeights, BaseMultiModalWeightInfo
from maga_transformer.utils.multimodal_util import MMUrlType, vit_emb_cache_, get_bytes_io_from_url

try:
    from decord import VideoReader, cpu
except ModuleNotFoundError:
    VideoReader = None
    cpu = None

def encode_video(video_path, max_num_frames: int = 32):
    def uniform_sample(l, n):
        gap = len(l) / n
        idxs = [int(i * gap + gap / 2) for i in range(n)]
        return [l[i] for i in idxs]

    vr = VideoReader(video_path, ctx=cpu(0))
    sample_fps = round(vr.get_avg_fps() / 1)  # FPS
    frame_idx = [i for i in range(0, len(vr), sample_fps)]
    if len(frame_idx) > max_num_frames:
        frame_idx = uniform_sample(frame_idx, max_num_frames)
    frames = vr.get_batch(frame_idx).asnumpy()
    frames = [Image.fromarray(v.astype('uint8')) for v in frames]
    return frames

class ImageEmbeddingInterface(MultiModalEmbeddingInterface):

    def __init__(self, config: GptInitModelParameters):
        self.config = config
        config = config.mm_related_params.config
        self.vision_config = SiglipVisionConfig(**config)
        self.processor = AutoProcessor.from_pretrained(config['ckpt_path'],
                                                       trust_remote_code=True)
        self.vpm = SiglipVisionTransformer(self.vision_config)
        self.embed_dim = config['llm_hidden_size']
        self.query_num = config['query_num']
        self.vision_dim = self.vision_config.hidden_size
        self.resampler = Resampler(num_queries=self.query_num,
                                   embed_dim=self.embed_dim,
                                   num_heads=self.embed_dim // 128,
                                   kv_dim=self.vision_dim,
                                   adaptive=True)

    @property
    def _device(self):
        return self.vpm.device

    @torch.inference_mode()
    def mm_embedding(self, url: str, mm_type: MMUrlType, **kwargs):
        dtype = self._data_type
        if g_parallel_info.tp_rank > 0:
            return torch.Tensor([])
        cached_res = vit_emb_cache_.check_cache(url)
        if cached_res is None:
            cached_url_res = get_bytes_io_from_url(url)
            cached_url_res = self._mm_preprocess(cached_url_res, mm_type)
            with mm_lock:
                features = self.mm_process(cached_url_res,
                                        mm_type=mm_type,
                                        **kwargs)
            if isinstance(features, list):
                features = torch.stack(features).to(dtype).contiguous()
            vit_emb_cache_.insert_cache(url, features)
            return (features, None)
        else:
            return (cached_res, None)
        
    def _mm_preprocess(self, data, type, **kwargs):
        if type == MMUrlType.IMAGE:
            return Image.open(data).convert("RGB")
        elif type == MMUrlType.VIDEO:
            return encode_video(data)

    @torch.inference_mode()
    def mm_process(self, mm_input, **kwargs):
        mm_type = kwargs.get("mm_type")
        if mm_type == MMUrlType.DEFAULT:
            if isinstance(mm_input, list):
                return self.image_embedding(mm_input)
            else:
                return self.image_embedding([mm_input])
        elif mm_type == MMUrlType.IMAGE:
            if isinstance(mm_input, list):
                raise Exception("expect single image input, but get a list")
            return self.image_embedding([mm_input])
        elif mm_type == MMUrlType.VIDEO:
            if not isinstance(mm_input, list):
                raise Exception("expect video input, but get a single image")
            return self.image_embedding(mm_input)
        else:
            raise Exception("unknown mm url type")

    @torch.no_grad()
    def image_embedding(self, images: List[Any]) -> List[torch.Tensor]:
        data = self.processor.image_processor(images, return_tensors="pt")
        dtype = self._data_type
        tgt_sizes = data['tgt_sizes']
        pixel_values_list = data['pixel_values']
        vision_hidden_states = []
        all_pixel_values = []
        img_cnt = []
        for pixel_values in pixel_values_list:
            img_cnt.append(len(pixel_values))
            all_pixel_values.extend([
                i.flatten(end_dim=1).permute(1, 0).to(self._device)
                for i in pixel_values
            ])

        assert all_pixel_values
        # exist image
        if all_pixel_values:
            tgt_sizes = [
                tgt_size for tgt_size in tgt_sizes
                if isinstance(tgt_size, torch.Tensor)
            ]
            tgt_sizes = torch.vstack(tgt_sizes).type(torch.int32)

            max_patches = torch.max(tgt_sizes[:, 0] * tgt_sizes[:, 1])

            all_pixel_values = torch.nn.utils.rnn.pad_sequence(
                all_pixel_values, batch_first=True, padding_value=0.0)
            B, L, _ = all_pixel_values.shape
            all_pixel_values = all_pixel_values.permute(0, 2, 1).reshape(
                B, 3, -1, L)

            patch_attn_mask = torch.zeros((B, 1, max_patches),
                                          dtype=torch.bool,
                                          device=self._device)
            for i in range(B):
                patch_attn_mask[i,
                                0, :tgt_sizes[i][0] * tgt_sizes[i][1]] = True

            vision_batch_size = 16
            all_pixel_values = all_pixel_values.type(dtype)
            if B > vision_batch_size:
                hs = []
                for i in range(0, B, vision_batch_size):
                    start_idx = i
                    end_idx = i + vision_batch_size
                    tmp_hs = self.vpm(all_pixel_values[start_idx:end_idx],
                                      patch_attention_mask=patch_attn_mask[
                                          start_idx:end_idx],
                                      tgt_sizes=tgt_sizes[start_idx:end_idx]
                                      ).last_hidden_state
                    hs.append(tmp_hs)
                vision_embedding = torch.cat(hs, dim=0)
            else:
                vision_embedding = self.vpm(
                    all_pixel_values,
                    patch_attention_mask=patch_attn_mask,
                    tgt_sizes=tgt_sizes).last_hidden_state
            vision_embedding = self.resampler(vision_embedding, tgt_sizes)

            start = 0
            for pixel_values in pixel_values_list:
                img_cnt = len(pixel_values)
                if img_cnt > 0:
                    for i in range(img_cnt):
                        vision_hidden_states.append(vision_embedding[start +
                                                                     i])
                    start += img_cnt
                else:
                    vision_hidden_states.append([])
        # print('embedding:', vision_hidden_states)
        # print('embedding shape:', [v.shape for v in vision_hidden_states])
        return vision_hidden_states


class MiniCPMVVitWeight(BaseVitWeights):

    def _set_weight_prefix(self):
        self._ckpt_prefix = ""
        self._ft_prefix = "self.mm_part."


class MiniCPMVWeightInfo(QWenV2Weight, BaseMultiModalWeightInfo):

    def __init__(self, config, tp_size, tp_rank):
        QWenV2Weight.__init__(self, config, tp_size, tp_rank, prefix="llm.")
        BaseMultiModalWeightInfo.__init__(self, config)

    def _get_weight_info(self):
        weights = super()._get_weight_info()
        weights = self._get_vit_info(weights)
        return weights


class MiniCPMV(QWenV2, MultiModalMixin):

    def __init__(self, config: GptInitModelParameters):
        QWenV2.__init__(self, config)
        self.config.mm_sep_tokens = [
            [self.tokenizer.im_start_id, self.tokenizer.im_end_id],
            [self.tokenizer.slice_start_id, self.tokenizer.slice_end_id]
        ]

    def _init_multimodal(self, config: GptInitModelParameters):
        self.mm_part = ImageEmbeddingInterface(config)
        config.mm_related_params.vit_weights = MiniCPMVVitWeight({
            "vpm":
            self.mm_part.vpm,
            "resampler":
            self.mm_part.resampler
        })

    @staticmethod
    def get_weight_cls():
        return MiniCPMVWeightInfo

    @classmethod
    def get_tokenizer(cls, config: GptInitModelParameters):
        tokenizer = AutoTokenizer.from_pretrained(config.tokenizer_path,
                                                  verbose=False,
                                                  trust_remote_code=True,
                                                  use_fast=True)
        return tokenizer

    @classmethod
    def _create_config(cls, ckpt_path: str):
        config = GptInitModelParameters(head_num=0,
                                        head_num_kv=0,
                                        size_per_head=0,
                                        layer_num=0,
                                        inter_size=0,
                                        vocab_size=0,
                                        max_seq_len=8192,
                                        ckpt_path=ckpt_path,
                                        rotary_embedding_dim=128,
                                        rotary_embedding_style=1,
                                        activation_type='SiGLU',
                                        has_pre_decoder_layernorm=False,
                                        has_post_decoder_layernorm=True,
                                        norm_type='rmsnorm')
        config_path = os.path.join(ckpt_path, 'config.json')
        if os.path.exists(config_path):
            with open(config_path) as reader:
                content = reader.read()
                config_json = json.loads(content)
                QWenV2._from_config_json(config, config_json)
                MiniCPMV._init_vit_params(config, config_json)
        else:
            raise Exception("no config.json found")
        return config

    @staticmethod
    def _init_vit_params(config: GptInitModelParameters,
                         config_json: Dict[str, Any]):
        config.mm_related_params.config = config_json["vision_config"]
        config.mm_related_params.config["llm_hidden_size"] = config_json[
            "hidden_size"]
        config.mm_related_params.config["query_num"] = config_json["query_num"]
        config.mm_related_params.config["ckpt_path"] = config.ckpt_path


register_model('minicpmv', MiniCPMV, ["MiniCPMV"])