import math
import os
from typing import List, Optional
import json
import torch
import torchvision
from threading import Thread
from copy import deepcopy
from PIL import Image
from transformers import AutoProcessor, Qwen3PreTrainedModel, Qwen3ForCausalLM, TextIteratorStreamer
from .configuration_minicpm import MiniCPMVConfig
# from .modeling_navit_siglip import SiglipVisionTransformer
from .modeling_navit_siglip_fast import SiglipVisionTransformer
from .resampler import Resampler
# from .mlp_connector import MLPConnector
from .mlp_merger import Merger
from .vit_insert_merger import get_vit_insert_merger
class MiniCPMVPreTrainedModel(Qwen3PreTrainedModel):
config_class = MiniCPMVConfig
class MiniCPMV(MiniCPMVPreTrainedModel):
def __init__(self, config):
super().__init__(config)
self.llm = Qwen3ForCausalLM(config)
self.vpm = self.init_vision_module()
self.vision_dim = self.vpm.embed_dim
self.embed_dim = self.llm.config.hidden_size
# self.resampler = self.init_resampler(self.embed_dim, self.vision_dim)
self.resampler = Merger(
hidden_size=self.vision_dim,
llm_embed_dim=self.embed_dim,
)
self.insert_layer_id = int(os.getenv('INSERT_LAYER_ID', 6))
self.vit_merger = get_vit_insert_merger(
hidden_size=self.vision_dim,
intermediate_size=self.vpm.config.intermediate_size,
vpm=self.vpm,
insert_layer_id=self.insert_layer_id,
)
self.processor = None
self.terminators = ['<|im_end|>', '<|endoftext|>']
def init_vision_module(self):
# same as HuggingFaceM4/siglip-so400m-14-980-flash-attn2-navit add tgt_sizes
if self.config._attn_implementation == 'flash_attention_2':
self.config.vision_config._attn_implementation = 'flash_attention_2'
else:
# not suport sdpa
self.config.vision_config._attn_implementation = 'eager'
self.config.vision_config._attn_implementation = 'flash_attention_2' # Hack 评测也使用flash_attention_2
model = SiglipVisionTransformer(self.config.vision_config)
if self.config.drop_vision_last_layer:
model.encoder.layers = model.encoder.layers[:-1]
setattr(model, 'embed_dim', model.embeddings.embed_dim)
setattr(model, 'patch_size', model.embeddings.patch_size)
return model
# def init_resampler(self, embed_dim, vision_dim):
# return Resampler(
# num_queries=self.config.query_num,
# embed_dim=embed_dim,
# num_heads=embed_dim // 128,
# kv_dim=vision_dim,
# adaptive=True,
# batch_infer=self.config.batch_3d_resampler
# )
def get_input_embeddings(self):
return self.llm.get_input_embeddings()
def set_input_embeddings(self, value):
self.llm.embed_tokens = value
def get_output_embeddings(self):
return self.llm.lm_head
def set_output_embeddings(self, new_embeddings):
self.llm.lm_head = new_embeddings
def set_decoder(self, decoder):
self.llm = decoder
def get_decoder(self):
return self.llm
def get_vllm_embedding(self, data):
if 'vision_hidden_states' not in data:
dtype = self.llm.model.embed_tokens.weight.dtype
device = self.llm.model.embed_tokens.weight.device
tgt_sizes = data['tgt_sizes']
pixel_values_list = data['pixel_values']
temporal_ids = data.get('temporal_ids', None)
vision_hidden_states = []
all_pixel_values = []
img_cnt = []
all_temporal_ids = None
# 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) for i in pixel_values])
for pixel_values in pixel_values_list:
img_cnt.append(len(pixel_values))
all_pixel_values.extend(pixel_values) # 42 * L
if temporal_ids is not None:
all_temporal_ids = []
for t in temporal_ids:
all_temporal_ids.extend(t)
# 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
# # print("all_pixel_values.shape: ", 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=device)
# for i in range(B):
# patch_attn_mask[i, 0, :tgt_sizes[i][0] * tgt_sizes[i][1]] = True
# all_pixel_values = all_pixel_values.type(dtype)
all_pixel_values = torch.concat(all_pixel_values, dim=-1).unsqueeze(0)
# print("shape of all_pixel_values: ", all_pixel_values.shape)
import torch.nn.functional as F
cu_seqlens = F.pad(torch.cumsum(tgt_sizes[:, 0] * tgt_sizes[:, 1], dim=0, dtype=torch.int32).cuda(), (1, 0))
max_seqlen = int(torch.max(cu_seqlens[1:] - cu_seqlens[:-1]).item())
vision_embedding, tgt_sizes = self.vpm(all_pixel_values.type(dtype), tgt_sizes=tgt_sizes, cu_seqlens=cu_seqlens, max_seqlens=max_seqlen, vit_merger=self.vit_merger, insert_layer_id=self.insert_layer_id)
vision_embedding = vision_embedding.last_hidden_state
# vision_embedding = self.resampler(vision_embedding, tgt_sizes, all_temporal_ids)
# print("vision_embedding.shape: ", vision_embedding.shape)
# print("tgt_sizes: ", tgt_sizes)
vision_embedding = self.resampler(vision_embedding, tgt_sizes) # 得到一个列表
# print("vision_embedding: ", vision_embedding)
# print("vision_embedding[0].shape: ", vision_embedding[0].shape)
start = 0
for pixel_values in pixel_values_list:
img_cnt = len(pixel_values)
if img_cnt > 0:
vision_hidden_states.append(vision_embedding[start: start + img_cnt])
start += img_cnt
else:
vision_hidden_states.append([])
else: # no image
if self.training:
dummy_image = torch.zeros(
(1, 3, self.patch_size * 8, self.patch_size * 8),
device=device, dtype=dtype
)
tgt_sizes = torch.Tensor([[8, 8]]).type(torch.int32)
vision_embedding, tgt_sizes = self.vpm(dummy_image, vit_merger=self.vit_merger, insert_layer_id=self.insert_layer_id)
vision_embedding = vision_embedding.last_hidden_state
vision_embedding = self.resampler(vision_embedding, tgt_sizes) # 得到一个列表
else:
dummy_feature = []
for _ in range(len(pixel_values_list)):
vision_hidden_states.append(dummy_feature)
else:
vision_hidden_states = data['vision_hidden_states']
if hasattr(self.llm.config, 'scale_emb'):
vllm_embedding = self.llm.model.embed_tokens(data['input_ids']) * self.llm.config.scale_emb
else:
vllm_embedding = self.llm.model.embed_tokens(data['input_ids'])
vision_hidden_states = [i.type(vllm_embedding.dtype) if isinstance(
i, torch.Tensor) else i for i in vision_hidden_states]
bs = len(data['input_ids'])
device = vllm_embedding.device
embed_dim = vllm_embedding.shape[-1]
updated_vllm_embedding = torch.empty_like(vllm_embedding)
for i in range(bs):
cur_vs_hs = vision_hidden_states[i]
cur_vllm_emb = vllm_embedding[i]
if len(cur_vs_hs) == 0:
updated_vllm_embedding[i] = cur_vllm_emb
continue
cur_image_bound = data['image_bound'][i]
# print("cur_image_bound: ", cur_image_bound)
if len(cur_image_bound) > 0:
for index, (bound_start, bound_end) in enumerate(cur_image_bound):
indeces = torch.arange(bound_start, bound_end, dtype=torch.long, device=cur_vllm_emb.device)
# print("indeces: ", indeces)
# print("shape of cur_vs_hs[index]: ", cur_vs_hs[index].shape)
cur_vllm_emb[indeces] = cur_vs_hs[index]
# image_indices = torch.cat([
# torch.arange(r[0], r[1], dtype=torch.long)
# for r in cur_image_bound
# ]).to(device)
# print("image_indices: ", image_indices)
# indices_expanded = image_indices.view(-1, 1).expand(-1, embed_dim)
# # print("cur_vs_hs: ", cur_vs_hs)
# # print("len(cur_vs_hs): ", len(cur_vs_hs))
# print("indices_expanded: ", indices_expanded)
# vision_features = cur_vs_hs[0].view(-1, embed_dim)
# updated_emb = cur_vllm_emb.clone()
# updated_emb.scatter_(0, indices_expanded, vision_features)
updated_emb = cur_vllm_emb
updated_vllm_embedding[i] = updated_emb
elif self.training:
if isinstance(cur_vs_hs, torch.Tensor) and cur_vs_hs.numel() > 0:
dummy_gradient_term = cur_vs_hs.sum() * 0.0
updated_vllm_embedding[i] = cur_vllm_emb + dummy_gradient_term
else:
updated_vllm_embedding[i] = cur_vllm_emb
else:
updated_vllm_embedding[i] = cur_vllm_emb
vllm_embedding = updated_vllm_embedding
# print("vllm_embedding: ", vllm_embedding)
return vllm_embedding, vision_hidden_states
def forward(self, data, **kwargs):
vllm_embedding, vision_hidden_states = self.get_vllm_embedding(data)
position_ids = data["position_ids"]
if position_ids.dtype != torch.int64:
position_ids = position_ids.long()
# compatible with llama factory
for key in ["input_ids", "inputs_embeds", "position_ids"]:
if key in kwargs:
del kwargs[key]
return self.llm(
input_ids=None,
position_ids=position_ids,
inputs_embeds=vllm_embedding,
**kwargs
)
def _decode(self, inputs_embeds, tokenizer, attention_mask, decode_text=False, **kwargs):
terminators = [tokenizer.convert_tokens_to_ids(i) for i in self.terminators]
output = self.llm.generate(
inputs_embeds=inputs_embeds,
pad_token_id=0,
eos_token_id=terminators,
attention_mask=attention_mask,
**kwargs
)
if decode_text:
return self._decode_text(output, tokenizer)
return output
def _decode_stream(self, inputs_embeds, tokenizer, **kwargs):
terminators = [tokenizer.convert_tokens_to_ids(i) for i in self.terminators]
streamer = TextIteratorStreamer(tokenizer=tokenizer)
generation_kwargs = {
'inputs_embeds': inputs_embeds,
'pad_token_id': 0,
'eos_token_id': terminators,
'streamer': streamer
}
generation_kwargs.update(kwargs)
thread = Thread(target=self.llm.generate, kwargs=generation_kwargs)
thread.start()
return streamer
def _decode_text(self, result_ids, tokenizer):
terminators = [tokenizer.convert_tokens_to_ids(i) for i in self.terminators]
result_text = []
for result in result_ids:
result = result[result != 0]
if result[0] == tokenizer.bos_id:
result = result[1:]
if result[-1] in terminators:
result = result[:-1]
result_text.append(tokenizer.decode(result).strip())
return result_text
def generate(
self,
input_ids=None,
pixel_values=None,
tgt_sizes=None,
image_bound=None,
temporal_ids=None,
attention_mask=None,
tokenizer=None,
vision_hidden_states=None,
return_vision_hidden_states=False,
stream=False,
decode_text=False,
**kwargs
):
assert input_ids is not None
assert len(input_ids) == len(pixel_values)
model_inputs = {
"input_ids": input_ids,
"image_bound": image_bound,
"temporal_ids": temporal_ids,
}
if vision_hidden_states is None:
model_inputs["pixel_values"] = pixel_values
model_inputs['tgt_sizes'] = tgt_sizes
else:
model_inputs["vision_hidden_states"] = vision_hidden_states
with torch.inference_mode():
(
model_inputs["inputs_embeds"],
vision_hidden_states,
) = self.get_vllm_embedding(model_inputs)
if stream:
result = self._decode_stream(model_inputs["inputs_embeds"], tokenizer, **kwargs)
else:
result = self._decode(model_inputs["inputs_embeds"], tokenizer, attention_mask, decode_text=decode_text, **kwargs)
if return_vision_hidden_states:
return result, vision_hidden_states
return result
def chat(
self,
image=None,
msgs=None,
tokenizer=None,
processor=None,
vision_hidden_states=None,
max_new_tokens=2048,
min_new_tokens=0,
sampling=True,
max_inp_length=16384,
system_prompt='',
stream=False,
max_slice_nums=None,
use_image_id=None,
temporal_ids=None,
enable_thinking=False,
**kwargs
):
if isinstance(msgs[0], list):
batched = True
else:
batched = False
msgs_list = msgs
images_list = image
if batched is False:
images_list, msgs_list = [images_list], [msgs_list]
else:
assert images_list is None, "Please integrate image to msgs when using batch inference."
images_list = [None] * len(msgs_list)
assert len(images_list) == len(msgs_list), "The batch dim of images_list and msgs_list should be the same."
if processor is None:
if self.processor is None:
self.processor = AutoProcessor.from_pretrained(self.config._name_or_path, trust_remote_code=True)
processor = self.processor
assert self.config.query_num == processor.image_processor.image_feature_size, "These two values should be the same. Check `config.json` and `preprocessor_config.json`."
assert self.config.patch_size == processor.image_processor.patch_size, "These two values should be the same. Check `config.json` and `preprocessor_config.json`."
assert self.config.use_image_id == processor.image_processor.use_image_id, "These two values should be the same. Check `config.json` and `preprocessor_config.json`."
assert self.config.slice_config.max_slice_nums == processor.image_processor.max_slice_nums, "These two values should be the same. Check `config.json` and `preprocessor_config.json`."
assert self.config.slice_mode == processor.image_processor.slice_mode, "These two values should be the same. Check `config.json` and `preprocessor_config.json`."
prompts_lists = []
input_images_lists = []
for image, msgs in zip(images_list, msgs_list):
if isinstance(msgs, str):
msgs = json.loads(msgs)
copy_msgs = deepcopy(msgs)
assert len(msgs) > 0, "msgs is empty"
assert sampling or not stream, "if use stream mode, make sure sampling=True"
if image is not None and isinstance(copy_msgs[0]["content"], str):
copy_msgs[0]["content"] = [image, copy_msgs[0]["content"]]
images = []
for i, msg in enumerate(copy_msgs):
role = msg["role"]
content = msg["content"]
assert role in ["user", "assistant"]
if i == 0:
assert role == "user", "The role of first msg should be user"
if isinstance(content, str):
content = [content]
cur_msgs = []
for c in content:
if isinstance(c, Image.Image):
images.append(c)
cur_msgs.append("(./)")
elif isinstance(c, str):
cur_msgs.append(c)
msg["content"] = "\n".join(cur_msgs)
if system_prompt:
sys_msg = {'role': 'system', 'content': system_prompt}
copy_msgs = [sys_msg] + copy_msgs
prompts_lists.append(processor.tokenizer.apply_chat_template(copy_msgs, tokenize=False, add_generation_prompt=True, enable_thinking=enable_thinking))
input_images_lists.append(images)
# print("prompts_lists: ", prompts_lists)
# print("input_images_lists: ", input_images_lists)
if enable_thinking:
prefill_answer = '\n'
else:
prefill_answer = ''
inputs = processor(
prompts_lists,
input_images_lists,
max_slice_nums=max_slice_nums,
use_image_id=use_image_id,
temporal_ids=temporal_ids,
return_tensors="pt",
max_length=max_inp_length
).to(self.device)
# print("shape of inputs['pixel_values'][0][0]: ", inputs['pixel_values'][0][0].shape)
# print("inputs['pixel_values']: ", inputs['pixel_values'])
# print("inputs: ", inputs)
if sampling:
generation_config = {
"temperature": 0.7,
"do_sample": True,
}
if not enable_thinking:
generation_config.update(
{
"top_p": 0.8,
"top_k": 100,
"repetition_penalty": 1.03
}
)
else:
generation_config = {
"num_beams": 3,
"repetition_penalty": 1.2,
}
if min_new_tokens > 0:
generation_config['min_new_tokens'] = min_new_tokens
generation_config.update(
(k, kwargs[k]) for k in generation_config.keys() & kwargs.keys()
)
# print("generation_config: ", generation_config)
inputs.pop("image_sizes")
with torch.inference_mode():
res = self.generate(
**inputs,
tokenizer=tokenizer,
max_new_tokens=max_new_tokens,
vision_hidden_states=vision_hidden_states,
stream=stream,
decode_text=True,
**generation_config
)
if stream:
def stream_gen():
for text in prefill_answer:
yield text
for text in res:
for term in self.terminators:
text = text.replace(term, '')
yield text
return stream_gen()
else:
if batched:
answer = [prefill_answer + i if prefill_answer else i for i in res]
else:
answer = prefill_answer + res[0] if prefill_answer else '' + res[0]
return answer