694 lines
26 KiB
Python
694 lines
26 KiB
Python
import json
|
|
import unittest
|
|
from io import BytesIO
|
|
from typing import Optional
|
|
|
|
import requests
|
|
import torch
|
|
|
|
# Compatibility shim: Kimi-VL dynamic module expects PytorchGELUTanh which may
|
|
# be missing in transformers==4.57.1. Inject a lightweight implementation so
|
|
# the model can import successfully without downgrading transformers.
|
|
import transformers.activations as _hf_activations
|
|
from PIL import Image
|
|
from transformers import (
|
|
AutoModel,
|
|
AutoProcessor,
|
|
Gemma3ForConditionalGeneration,
|
|
Qwen2_5_VLForConditionalGeneration,
|
|
)
|
|
|
|
from sglang.test.ci.ci_register import register_cuda_ci
|
|
|
|
if not hasattr(_hf_activations, "PytorchGELUTanh"):
|
|
|
|
class PytorchGELUTanh(torch.nn.Module):
|
|
def forward(self, x):
|
|
return torch.nn.functional.gelu(x, approximate="tanh")
|
|
|
|
_hf_activations.PytorchGELUTanh = PytorchGELUTanh
|
|
_hf_activations.ACT2FN.setdefault(
|
|
"pytorch_gelu_tanh",
|
|
lambda x: torch.nn.functional.gelu(x, approximate="tanh"),
|
|
)
|
|
|
|
from sglang import Engine
|
|
from sglang.srt.entrypoints.openai.protocol import ChatCompletionRequest
|
|
from sglang.srt.parser.conversation import generate_chat_conv
|
|
from sglang.srt.utils.hf_transformers_utils import _fix_added_tokens_encoding
|
|
|
|
register_cuda_ci(est_time=447, suite="stage-b-test-1-gpu-large")
|
|
|
|
IMAGE_MAN_IRONING_URL = "https://raw.githubusercontent.com/sgl-project/sgl-test-files/refs/heads/main/images/man_ironing_on_back_of_suv.png"
|
|
IMAGE_SGL_LOGO_URL = "https://raw.githubusercontent.com/sgl-project/sgl-test-files/refs/heads/main/images/sgl_logo.png"
|
|
|
|
|
|
class VLMInputTestBase:
|
|
model_path = None
|
|
chat_template = None
|
|
processor = None
|
|
visual = None # Should be a callable for precomputed embeddings
|
|
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
assert cls.model_path is not None, "Set model_path in subclass"
|
|
assert cls.chat_template is not None, "Set chat_template in subclass"
|
|
cls.image_urls = [IMAGE_MAN_IRONING_URL, IMAGE_SGL_LOGO_URL]
|
|
cls.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
cls.main_image = []
|
|
for image_url in cls.image_urls:
|
|
response = requests.get(image_url)
|
|
cls.main_image.append(Image.open(BytesIO(response.content)))
|
|
cls.processor = AutoProcessor.from_pretrained(
|
|
cls.model_path, trust_remote_code=True, use_fast=True
|
|
)
|
|
_fix_added_tokens_encoding(cls.processor.tokenizer)
|
|
cls._init_visual()
|
|
|
|
@classmethod
|
|
def _init_visual(cls):
|
|
"""Override in subclass to set up cls.visual as a callable for precomputed embeddings."""
|
|
raise NotImplementedError
|
|
|
|
def setUp(self):
|
|
self.engine = Engine(
|
|
model_path=self.model_path,
|
|
chat_template=self.chat_template,
|
|
device=self.device.type,
|
|
mem_fraction_static=0.8,
|
|
enable_multimodal=True,
|
|
disable_cuda_graph=True,
|
|
trust_remote_code=True,
|
|
)
|
|
|
|
def tearDown(self):
|
|
self.engine.shutdown()
|
|
|
|
def verify_response(self, output):
|
|
# The goal is to check that the model roughly understands:
|
|
# - image 1: taxi / car scene
|
|
# - image 2: SGL logo / company
|
|
# We intentionally keep the check keyword-based and loose to avoid
|
|
# overfitting to a specific phrasing.
|
|
out_text = output["text"].lower()
|
|
|
|
assert any(w in out_text for w in ("taxi", "cab", "car")), out_text
|
|
|
|
has_sg_or_logo_side = any(
|
|
kw in out_text
|
|
for kw in (
|
|
"sg ",
|
|
"sgl",
|
|
" sgl",
|
|
"logo",
|
|
"software guidance",
|
|
"labs",
|
|
"laborator",
|
|
"company",
|
|
" text",
|
|
)
|
|
)
|
|
assert has_sg_or_logo_side, out_text
|
|
|
|
def get_completion_request(self) -> ChatCompletionRequest:
|
|
json_structure = {
|
|
"model": self.model_path,
|
|
"messages": [
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{"type": "image_url", "image_url": {"url": self.image_urls[0]}},
|
|
{"type": "image_url", "image_url": {"url": self.image_urls[1]}},
|
|
{
|
|
"type": "text",
|
|
"text": "Describe both the first image and the second image in detail separately.", # update prompt, ensure kimi-vl understands the images separately.
|
|
},
|
|
],
|
|
}
|
|
],
|
|
}
|
|
json_str = json.dumps(json_structure)
|
|
return ChatCompletionRequest.model_validate_json(json_str)
|
|
|
|
def get_processor_output(self, req: Optional[ChatCompletionRequest] = None):
|
|
if req is None:
|
|
req = self.get_completion_request()
|
|
conv = generate_chat_conv(req, template_name=self.chat_template)
|
|
text = conv.get_prompt()
|
|
|
|
# Process inputs using processor
|
|
inputs = self.processor(
|
|
text=[text],
|
|
images=self.main_image,
|
|
return_tensors="pt",
|
|
).to(self.device)
|
|
|
|
return inputs, text
|
|
|
|
async def test_accepts_image(self):
|
|
req = self.get_completion_request()
|
|
conv = generate_chat_conv(req, template_name=self.chat_template)
|
|
text = conv.get_prompt()
|
|
output = await self.engine.async_generate(
|
|
prompt=text,
|
|
image_data=self.main_image,
|
|
sampling_params=dict(temperature=0.0, max_new_tokens=512),
|
|
)
|
|
self.verify_response(output)
|
|
|
|
async def test_accepts_precomputed_embeddings(self):
|
|
req = self.get_completion_request()
|
|
processor_output, _ = self.get_processor_output(req=req)
|
|
|
|
with torch.inference_mode():
|
|
precomputed_embeddings = self.__class__.visual(processor_output)
|
|
|
|
output = await self.engine.async_generate(
|
|
input_ids=processor_output["input_ids"][0].detach().cpu().tolist(),
|
|
image_data=[
|
|
self._precomputed_image_data(processor_output, precomputed_embeddings)
|
|
],
|
|
sampling_params=dict(temperature=0.0, max_new_tokens=512),
|
|
)
|
|
self.verify_response(output)
|
|
|
|
async def test_accepts_processor_output(self):
|
|
req = self.get_completion_request()
|
|
processor_output, prompt = self.get_processor_output(req=req)
|
|
output = await self.engine.async_generate(
|
|
input_ids=processor_output["input_ids"][0].detach().cpu().tolist(),
|
|
image_data=[self._processor_output_image_data(processor_output)],
|
|
sampling_params=dict(temperature=0.0, max_new_tokens=512),
|
|
)
|
|
self.verify_response(output)
|
|
|
|
def _precomputed_image_data(self, processor_output, precomputed_embeddings):
|
|
"""This should not be overridden."""
|
|
return dict(
|
|
processor_output,
|
|
format="precomputed_embedding",
|
|
feature=precomputed_embeddings,
|
|
)
|
|
|
|
def _processor_output_image_data(self, processor_output):
|
|
"""Override in subclass to pass the correct set of arguments."""
|
|
raise NotImplementedError
|
|
|
|
|
|
class TestQwenVLUnderstandsImage(VLMInputTestBase, unittest.IsolatedAsyncioTestCase):
|
|
model_path = "Qwen/Qwen2.5-VL-3B-Instruct"
|
|
chat_template = "qwen2-vl"
|
|
|
|
@classmethod
|
|
def _init_visual(cls):
|
|
model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
|
|
cls.model_path, torch_dtype=torch.bfloat16
|
|
).eval()
|
|
# In transformers v5, .visual moved under .model
|
|
visual = model.model.visual
|
|
cls.visual_model = visual.to(cls.device)
|
|
|
|
# In transformers v5, the visual encoder returns BaseModelOutputWithPooling;
|
|
# pooler_output has the spatially-merged embeddings we need.
|
|
def visual(processor_output):
|
|
out = cls.visual_model(
|
|
processor_output["pixel_values"], processor_output["image_grid_thw"]
|
|
)
|
|
return out.pooler_output if hasattr(out, "pooler_output") else out
|
|
|
|
cls.visual = visual
|
|
|
|
def _processor_output_image_data(self, processor_output):
|
|
return dict(processor_output, format="processor_output")
|
|
|
|
|
|
class TestGemmaUnderstandsImage(VLMInputTestBase, unittest.IsolatedAsyncioTestCase):
|
|
model_path = "google/gemma-3-4b-it"
|
|
chat_template = "gemma-it"
|
|
|
|
@classmethod
|
|
def _init_visual(cls):
|
|
model = Gemma3ForConditionalGeneration.from_pretrained(
|
|
cls.model_path, torch_dtype=torch.bfloat16
|
|
)
|
|
base_model = model.model
|
|
|
|
cls.vision_tower = base_model.vision_tower.eval().to(cls.device)
|
|
|
|
if hasattr(base_model, "multi_modal_projector"):
|
|
cls.mm_projector = base_model.multi_modal_projector.eval().to(cls.device)
|
|
else:
|
|
cls.mm_projector = model.multi_modal_projector.eval().to(cls.device)
|
|
|
|
cls.visual = lambda processor_output: cls.mm_projector(
|
|
cls.vision_tower(
|
|
pixel_values=processor_output["pixel_values"]
|
|
).last_hidden_state
|
|
)
|
|
|
|
def _processor_output_image_data(self, processor_output):
|
|
return dict(processor_output, format="processor_output")
|
|
|
|
|
|
# Updated Kimi-VL test to use the new input format.
|
|
class TestKimiVLImageUnderstandsImage(
|
|
VLMInputTestBase, unittest.IsolatedAsyncioTestCase
|
|
):
|
|
model_path = "moonshotai/Kimi-VL-A3B-Instruct"
|
|
chat_template = "kimi-vl"
|
|
|
|
@classmethod
|
|
def _init_visual(cls):
|
|
import inspect
|
|
|
|
from transformers import AutoConfig
|
|
from transformers.dynamic_module_utils import get_class_from_dynamic_module
|
|
|
|
config = AutoConfig.from_pretrained(cls.model_path, trust_remote_code=True)
|
|
|
|
# Transformers v5 auto-populates rope_scaling with
|
|
# {"rope_theta": ..., "rope_type": "default"} even when the original
|
|
# config had rope_scaling: null. The remote KimiVL code branches on
|
|
# `if self.config.rope_scaling is None` so we must reset it.
|
|
tc = getattr(config, "text_config", None)
|
|
if tc is not None:
|
|
rs = getattr(tc, "rope_scaling", None)
|
|
if isinstance(rs, dict) and rs.get("rope_type") == "default":
|
|
tc.rope_scaling = None
|
|
|
|
# Transformers v5 calls tie_weights(recompute_mapping=False) in
|
|
# post_init, but KimiVL's tie_weights doesn't accept that kwarg.
|
|
auto_map = getattr(config, "auto_map", {})
|
|
model_ref = auto_map.get("AutoModel")
|
|
if model_ref:
|
|
model_cls = get_class_from_dynamic_module(model_ref, cls.model_path)
|
|
orig_tie = model_cls.tie_weights
|
|
if "recompute_mapping" not in inspect.signature(orig_tie).parameters:
|
|
|
|
def _patched_tie(self, **kwargs):
|
|
return orig_tie(self)
|
|
|
|
model_cls.tie_weights = _patched_tie
|
|
|
|
model = AutoModel.from_pretrained(
|
|
cls.model_path, config=config, trust_remote_code=True
|
|
)
|
|
cls.vision_tower = model.vision_tower.eval().to(cls.device)
|
|
cls.mm_projector = model.multi_modal_projector.eval().to(cls.device)
|
|
_vt_dtype = next(cls.vision_tower.parameters()).dtype
|
|
|
|
cls.visual = lambda tokenizer_output: cls.mm_projector(
|
|
cls.vision_tower(
|
|
pixel_values=tokenizer_output["pixel_values"].to(_vt_dtype),
|
|
grid_hws=tokenizer_output["image_grid_hws"],
|
|
)
|
|
)
|
|
|
|
def _processor_output_image_data(self, processor_output):
|
|
return dict(processor_output, format="processor_output")
|
|
|
|
|
|
# not for CI: too large
|
|
# class TestLlama4ImageUnderstandsImage(
|
|
# VLMInputTestBase, unittest.IsolatedAsyncioTestCase
|
|
# ):
|
|
# # Allow overriding via env for local/offline runs.
|
|
# model_path = "meta-llama/Llama-4-Scout-17B-16E-Instruct"
|
|
# chat_template = "llama-4"
|
|
|
|
# def setUp(self):
|
|
# if torch.cuda.device_count() < 4:
|
|
# self.skipTest("Skipping Llama-4 test: requires 4 GPUs for TP=4")
|
|
# self.engine = Engine(
|
|
# model_path=self.model_path,
|
|
# trust_remote_code=True,
|
|
# chat_template=self.chat_template,
|
|
# enable_multimodal=True,
|
|
# mem_fraction_static=0.8,
|
|
# tp_size=4,
|
|
# attention_backend="fa3",
|
|
# context_length=65536,
|
|
# )
|
|
|
|
# @classmethod
|
|
# def _init_visual(cls):
|
|
# model = AutoModel.from_pretrained(
|
|
# cls.model_path,
|
|
# trust_remote_code=True,
|
|
# torch_dtype="auto",
|
|
# force_download=True,
|
|
# )
|
|
# cls.vision_tower = model.vision_model.eval().to(cls.device)
|
|
# cls.mm_projector = model.multi_modal_projector.eval().to(cls.device)
|
|
|
|
# cls.visual = lambda tokenizer_output: cls.mm_projector(
|
|
# cls.vision_tower(
|
|
# pixel_values=tokenizer_output["pixel_values"],
|
|
# ).last_hidden_state.flatten(0, -2)
|
|
# )
|
|
|
|
# def _processor_output_image_data(self, processor_output):
|
|
# # Llama-4 vision expects processor_output format with pixel_values
|
|
# return dict(processor_output, format="processor_output")
|
|
|
|
|
|
# class TestLlavaUnderstandsImage(VLMInputTestBase, unittest.IsolatedAsyncioTestCase):
|
|
# model_path = "llava-hf/llava-1.5-7b-hf"
|
|
# chat_template = "vicuna_v1.1"
|
|
|
|
# @classmethod
|
|
# def _init_visual(cls):
|
|
# from transformers import LlavaForConditionalGeneration
|
|
|
|
# model = LlavaForConditionalGeneration.from_pretrained(
|
|
# cls.model_path,
|
|
# torch_dtype=torch.float16,
|
|
# low_cpu_mem_usage=True,
|
|
# )
|
|
# cls.vision_tower = model.vision_tower.eval().to(cls.device)
|
|
# cls.multi_modal_projector = model.multi_modal_projector.eval().to(cls.device)
|
|
# cls.config = model.config
|
|
|
|
# def visual_func(processor_output):
|
|
# pixel_values = processor_output["pixel_values"].to(
|
|
# cls.device, dtype=torch.float16
|
|
# )
|
|
|
|
# vision_outputs = cls.vision_tower(pixel_values, output_hidden_states=True)
|
|
# image_features = vision_outputs.hidden_states[-2]
|
|
|
|
# if cls.config.vision_feature_select_strategy == "default":
|
|
# image_features = image_features[:, 1:]
|
|
# elif cls.config.vision_feature_select_strategy == "full":
|
|
# image_features = image_features
|
|
|
|
# image_features = cls.multi_modal_projector(image_features)
|
|
# return image_features
|
|
|
|
# cls.visual = visual_func
|
|
|
|
# def _processor_output_image_data(self, processor_output):
|
|
# return dict(processor_output, format="processor_output")
|
|
|
|
|
|
class TestInternVLUnderstandsImage(VLMInputTestBase, unittest.IsolatedAsyncioTestCase):
|
|
model_path = "OpenGVLab/InternVL2-2B"
|
|
chat_template = "internvl-2-5"
|
|
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
assert cls.model_path is not None, "Set model_path in subclass"
|
|
assert cls.chat_template is not None, "Set chat_template in subclass"
|
|
cls.image_urls = [IMAGE_MAN_IRONING_URL, IMAGE_SGL_LOGO_URL]
|
|
cls.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
cls.main_image = []
|
|
for image_url in cls.image_urls:
|
|
response = requests.get(image_url)
|
|
cls.main_image.append(Image.open(BytesIO(response.content)))
|
|
|
|
# InternVL models (2, 3, 3.5) do not ship a standard HuggingFace
|
|
# Processor; AutoProcessor.from_pretrained returns a bare tokenizer.
|
|
# Use AutoTokenizer explicitly so the intent is clear.
|
|
from transformers import AutoTokenizer
|
|
|
|
cls.processor = AutoTokenizer.from_pretrained(
|
|
cls.model_path, trust_remote_code=True
|
|
)
|
|
cls._init_visual()
|
|
|
|
@classmethod
|
|
def _init_visual(cls):
|
|
try:
|
|
model = AutoModel.from_pretrained(
|
|
cls.model_path,
|
|
trust_remote_code=True,
|
|
torch_dtype=torch.bfloat16,
|
|
low_cpu_mem_usage=False,
|
|
)
|
|
except RuntimeError as e:
|
|
if "meta" not in str(e):
|
|
raise
|
|
# Transformers v5 always uses meta tensors for init, which breaks
|
|
# models calling .item() in __init__ (e.g. InternVL's drop_path_rate).
|
|
# Fall back to from_config + manual weight loading.
|
|
import gc
|
|
import glob
|
|
import os
|
|
|
|
from huggingface_hub import snapshot_download
|
|
from safetensors.torch import load_file
|
|
from transformers import AutoConfig
|
|
|
|
config = AutoConfig.from_pretrained(cls.model_path, trust_remote_code=True)
|
|
with torch.device("cpu"):
|
|
model = AutoModel.from_config(
|
|
config,
|
|
trust_remote_code=True,
|
|
torch_dtype=torch.bfloat16,
|
|
)
|
|
model_dir = snapshot_download(cls.model_path)
|
|
for f in sorted(glob.glob(os.path.join(model_dir, "*.safetensors"))):
|
|
shard = load_file(f)
|
|
model.load_state_dict(shard, strict=False)
|
|
del shard
|
|
gc.collect()
|
|
|
|
cls.vision_model = model.vision_model.eval().to(cls.device)
|
|
cls.mlp1 = model.mlp1.eval().to(cls.device)
|
|
|
|
config = model.config
|
|
cls.internvl_config = config
|
|
image_size = getattr(config, "force_image_size", None) or (
|
|
config.vision_config.image_size
|
|
)
|
|
patch_size = config.vision_config.patch_size
|
|
cls.num_image_token = int(
|
|
(image_size // patch_size) ** 2 * (config.downsample_ratio**2)
|
|
)
|
|
cls.internvl_image_size = image_size
|
|
cls.internvl_downsample_ratio = config.downsample_ratio
|
|
cls.internvl_ps_version = config.ps_version
|
|
cls.internvl_select_layer = config.select_layer
|
|
|
|
del model
|
|
|
|
def pixel_shuffle(x, scale_factor):
|
|
n, w, h, c = x.size()
|
|
x = x.view(n, w, int(h * scale_factor), int(c / scale_factor))
|
|
x = x.permute(0, 2, 1, 3).contiguous()
|
|
x = x.view(
|
|
n,
|
|
int(h * scale_factor),
|
|
int(w * scale_factor),
|
|
int(c / (scale_factor * scale_factor)),
|
|
)
|
|
if cls.internvl_ps_version != "v1":
|
|
x = x.permute(0, 2, 1, 3).contiguous()
|
|
return x
|
|
|
|
def visual_func(processor_output):
|
|
pixel_values = processor_output["pixel_values"].to(
|
|
cls.device, dtype=torch.bfloat16
|
|
)
|
|
if cls.internvl_select_layer == -1:
|
|
vit_embeds = cls.vision_model(
|
|
pixel_values=pixel_values,
|
|
output_hidden_states=False,
|
|
return_dict=True,
|
|
).last_hidden_state
|
|
else:
|
|
vit_embeds = cls.vision_model(
|
|
pixel_values=pixel_values,
|
|
output_hidden_states=True,
|
|
return_dict=True,
|
|
).hidden_states[cls.internvl_select_layer]
|
|
vit_embeds = vit_embeds[:, 1:, :]
|
|
|
|
h = w = int(vit_embeds.shape[1] ** 0.5)
|
|
vit_embeds = vit_embeds.reshape(vit_embeds.shape[0], h, w, -1)
|
|
vit_embeds = pixel_shuffle(
|
|
vit_embeds, scale_factor=cls.internvl_downsample_ratio
|
|
)
|
|
vit_embeds = vit_embeds.reshape(
|
|
vit_embeds.shape[0], -1, vit_embeds.shape[-1]
|
|
)
|
|
vit_embeds = cls.mlp1(vit_embeds)
|
|
return vit_embeds
|
|
|
|
cls.visual = visual_func
|
|
|
|
def get_processor_output(self, req=None):
|
|
"""Override to handle InternVL's custom preprocessing.
|
|
|
|
Uses shared ``image_to_pixel_values`` from ``internvl_utils`` for
|
|
image preprocessing (dynamic tiling + normalize) and expands
|
|
``<IMG_CONTEXT>`` placeholders into ``<img>`` + context tokens +
|
|
``</img>`` — mirroring the logic in
|
|
``InternVLProcessor.process_internlm2_mm_data_async``.
|
|
"""
|
|
from sglang.srt.multimodal.internvl_utils import image_to_pixel_values
|
|
from sglang.srt.multimodal.processors.internvl import InternVLProcessor
|
|
|
|
if req is None:
|
|
req = self.get_completion_request()
|
|
conv = generate_chat_conv(req, template_name=self.chat_template)
|
|
text = conv.get_prompt()
|
|
|
|
# Preprocess images using the shared utility (dynamic tiling +
|
|
# bicubic resize + ImageNet normalize), same pipeline as the engine.
|
|
all_pixel_values = []
|
|
num_patches_list = []
|
|
for img in self.main_image:
|
|
pv = image_to_pixel_values(
|
|
img,
|
|
input_size=self.internvl_image_size,
|
|
max_num_tiles=InternVLProcessor.IMAGE_MAX_NUM,
|
|
use_thumbnail=True,
|
|
)
|
|
all_pixel_values.append(pv)
|
|
num_patches_list.append(pv.shape[0])
|
|
|
|
pixel_values = torch.cat(all_pixel_values, dim=0).to(self.device)
|
|
|
|
# Expand each <IMG_CONTEXT> placeholder into <img> + <IMG_CONTEXT>*N + </img>.
|
|
# This mirrors InternVLProcessor.process_internlm2_mm_data_async.
|
|
ph = "<<<__IMG_PH__>>>"
|
|
expanded_text = text.replace(InternVLProcessor.IMG_CONTEXT, ph)
|
|
for num_patches in num_patches_list:
|
|
image_tokens = (
|
|
InternVLProcessor.IMG_START
|
|
+ InternVLProcessor.IMG_CONTEXT * (self.num_image_token * num_patches)
|
|
+ InternVLProcessor.IMG_END
|
|
)
|
|
expanded_text = expanded_text.replace(ph, image_tokens, 1)
|
|
# Remove any remaining placeholders (more placeholders than images)
|
|
expanded_text = expanded_text.replace(ph, "")
|
|
|
|
# Tokenize the expanded text
|
|
input_ids = self.processor(expanded_text, return_tensors="pt")["input_ids"]
|
|
|
|
return {
|
|
"input_ids": input_ids,
|
|
"pixel_values": pixel_values,
|
|
}, text
|
|
|
|
def _processor_output_image_data(self, processor_output):
|
|
return dict(processor_output, format="processor_output")
|
|
|
|
|
|
class TestMiniCPMVUnderstandsImage(VLMInputTestBase, unittest.IsolatedAsyncioTestCase):
|
|
model_path = "openbmb/MiniCPM-V-4"
|
|
chat_template = "minicpmv"
|
|
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
assert cls.model_path is not None, "Set model_path in subclass"
|
|
assert cls.chat_template is not None, "Set chat_template in subclass"
|
|
cls.image_urls = [IMAGE_MAN_IRONING_URL, IMAGE_SGL_LOGO_URL]
|
|
cls.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
cls.main_image = []
|
|
for image_url in cls.image_urls:
|
|
response = requests.get(image_url)
|
|
cls.main_image.append(Image.open(BytesIO(response.content)))
|
|
|
|
cls.processor = AutoProcessor.from_pretrained(
|
|
cls.model_path, trust_remote_code=True
|
|
)
|
|
_fix_added_tokens_encoding(cls.processor.tokenizer)
|
|
cls._init_visual()
|
|
|
|
@classmethod
|
|
def _init_visual(cls):
|
|
try:
|
|
model = AutoModel.from_pretrained(
|
|
cls.model_path, trust_remote_code=True, torch_dtype=torch.bfloat16
|
|
)
|
|
except (AttributeError, RuntimeError) as e:
|
|
err = str(e)
|
|
if "all_tied_weights_keys" not in err and "meta" not in err:
|
|
raise
|
|
# Transformers v5: remote model code may lack all_tied_weights_keys
|
|
# or meta-tensor init may break .item() calls. Fall back to
|
|
# from_config + manual weight loading.
|
|
import gc
|
|
import glob
|
|
import os
|
|
|
|
from huggingface_hub import snapshot_download
|
|
from safetensors.torch import load_file
|
|
from transformers import AutoConfig
|
|
|
|
config = AutoConfig.from_pretrained(cls.model_path, trust_remote_code=True)
|
|
with torch.device("cpu"):
|
|
model = AutoModel.from_config(
|
|
config,
|
|
trust_remote_code=True,
|
|
torch_dtype=torch.bfloat16,
|
|
)
|
|
model_dir = snapshot_download(cls.model_path)
|
|
for f in sorted(glob.glob(os.path.join(model_dir, "*.safetensors"))):
|
|
shard = load_file(f)
|
|
model.load_state_dict(shard, strict=False)
|
|
del shard
|
|
gc.collect()
|
|
|
|
cls.vpm_model = model.vpm.eval().to(cls.device)
|
|
cls.resampler_model = model.resampler.eval().to(cls.device)
|
|
del model
|
|
|
|
def visual_func(processor_output):
|
|
pixel_values = processor_output["pixel_values"]
|
|
tgt_sizes = processor_output["tgt_sizes"]
|
|
|
|
pixel_values_flat = []
|
|
tgt_sizes_flat = []
|
|
for pixel_b, tgt_b in zip(pixel_values, tgt_sizes):
|
|
if isinstance(pixel_b, (list, tuple)):
|
|
for pixel_n, tgt_n in zip(pixel_b, tgt_b):
|
|
pixel_values_flat.append(pixel_n)
|
|
tgt_sizes_flat.append(tgt_n)
|
|
else:
|
|
pixel_values_flat.append(pixel_b)
|
|
tgt_sizes_flat.append(tgt_b)
|
|
|
|
tgt_sizes_tensor = torch.stack(tgt_sizes_flat, dim=0)
|
|
device = cls.vpm_model.embeddings.position_embedding.weight.device
|
|
dtype = cls.vpm_model.embeddings.position_embedding.weight.dtype
|
|
|
|
all_pixel_values_lst = [
|
|
i.flatten(end_dim=1).permute(1, 0) for i in pixel_values_flat
|
|
]
|
|
max_patches = int(
|
|
(tgt_sizes_tensor[:, 0] * tgt_sizes_tensor[:, 1]).max().item()
|
|
)
|
|
all_pixel_values = torch.nn.utils.rnn.pad_sequence(
|
|
all_pixel_values_lst, 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=device
|
|
)
|
|
tgt_sizes_dev = tgt_sizes_tensor.to(device)
|
|
mask_shapes = tgt_sizes_dev[:, 0] * tgt_sizes_dev[:, 1]
|
|
patch_attn_mask[:, 0, :] = torch.arange(
|
|
max_patches, device=device
|
|
).unsqueeze(0) < mask_shapes.unsqueeze(1)
|
|
|
|
vision_output = cls.vpm_model(
|
|
all_pixel_values.type(dtype),
|
|
patch_attention_mask=patch_attn_mask,
|
|
tgt_sizes=tgt_sizes_tensor,
|
|
)
|
|
vision_embedding = vision_output.last_hidden_state
|
|
return cls.resampler_model(vision_embedding, tgt_sizes_tensor)
|
|
|
|
cls.visual = visual_func
|
|
|
|
def _processor_output_image_data(self, processor_output):
|
|
return dict(processor_output, format="processor_output")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|