| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | # |
| 4 | # Adapted from Qwen2.5-VL implementation |
| 5 | # Copyright 2025 The vLLM team. |
| 6 | # Copyright 2025 XLANG Lab, The University of Hong Kong |
| 7 | |
| 8 | """Inference-only OpenCUA-7B model compatible with HuggingFace weights.""" |
| 9 | |
| 10 | import typing |
| 11 | from collections.abc import Mapping, Sequence |
| 12 | from typing import Any |
| 13 | |
| 14 | import torch |
| 15 | import torch.nn as nn |
| 16 | from transformers import BatchFeature |
| 17 | from transformers.models.qwen2_vl import ( |
| 18 | Qwen2VLImageProcessor, |
| 19 | Qwen2VLProcessor, |
| 20 | Qwen2VLVideoProcessor, |
| 21 | ) |
| 22 | |
| 23 | from vllm.config import VllmConfig |
| 24 | from vllm.multimodal import MULTIMODAL_REGISTRY |
| 25 | from vllm.multimodal.inputs import ( |
| 26 | MultiModalFieldConfig, |
| 27 | MultiModalKwargsItems, |
| 28 | ) |
| 29 | from vllm.multimodal.parse import MultiModalDataItems |
| 30 | from vllm.multimodal.processing import ( |
| 31 | BaseDummyInputsBuilder, |
| 32 | BaseMultiModalProcessor, |
| 33 | PromptReplacement, |
| 34 | PromptUpdate, |
| 35 | ) |
| 36 | from vllm.tokenizers import TokenizerLike |
| 37 | |
| 38 | from .qwen2_5_vl import ( |
| 39 | Qwen2_5_VisionTransformer as OpenCUAVisionTransformer, |
| 40 | ) |
| 41 | from .qwen2_5_vl import ( |
| 42 | Qwen2_5_VLForConditionalGeneration, |
| 43 | ) |
| 44 | from .qwen2_vl import ( |
| 45 | Qwen2VLDummyInputsBuilder, |
| 46 | Qwen2VLMultiModalDataParser, |
| 47 | Qwen2VLProcessingInfo, |
| 48 | _create_qwen2vl_field_factory, |
| 49 | ) |
| 50 | from .utils import ( |
| 51 | WeightsMapper, |
| 52 | init_vllm_registered_model, |
| 53 | maybe_prefix, |
| 54 | ) |
| 55 | |
| 56 | |
| 57 | class OpenCUAProcessingInfo(Qwen2VLProcessingInfo): |
| 58 | def get_data_parser(self): |
| 59 | return Qwen2VLMultiModalDataParser( |
| 60 | self.get_hf_config().vision_config.spatial_merge_size, |
| 61 | expected_hidden_size=self._get_expected_hidden_size(), |
| 62 | allow_missing_mm_embeddings=self.allow_missing_mm_embeddings, |
| 63 | ) |
| 64 | |
| 65 | def get_hf_config(self): |
| 66 | return self.ctx.get_hf_config() |
| 67 | |
| 68 | def get_supported_mm_limits(self) -> Mapping[str, int | None]: |
| 69 | return {"image": None} |
| 70 | |
| 71 | def get_hf_processor(self, **kwargs: object): |
| 72 | """Load OpenCUA processor.""" |
| 73 | tokenizer = self.get_tokenizer() |
| 74 | vision_config = self.ctx.get_hf_image_processor_config() |
| 75 | return OpenCUAProcessor( |
| 76 | vision_config=vision_config, |
| 77 | tokenizer=tokenizer, |
| 78 | **kwargs, |
| 79 | ) |
| 80 | |
| 81 | |
| 82 | class OpenCUAProcessor(Qwen2VLProcessor): |
| 83 | def check_argument_for_proper_class(self, attribute_name: str, arg: object) -> None: |
| 84 | if attribute_name == "tokenizer": |
| 85 | return |
| 86 | return super().check_argument_for_proper_class(attribute_name, arg) |
| 87 | |
| 88 | def __init__( |
| 89 | self, |
| 90 | vision_config: dict, |
| 91 | tokenizer: TokenizerLike, |
| 92 | **kwargs, |
| 93 | ): |
| 94 | image_processor = Qwen2VLImageProcessor(**vision_config) |
| 95 | video_processor = Qwen2VLVideoProcessor(**vision_config) |
| 96 | chat_template = kwargs.pop("chat_template", None) |
| 97 | |
| 98 | super().__init__( |
| 99 | image_processor=image_processor, |
| 100 | tokenizer=tokenizer, |
| 101 | video_processor=video_processor, |
| 102 | chat_template=chat_template, |
| 103 | **kwargs, |
| 104 | ) |
| 105 | |
| 106 | self.image_token = "<|media_placeholder|>" |
| 107 | |
| 108 | def __call__( |
| 109 | self, |
| 110 | text=None, |
| 111 | images=None, |
| 112 | return_tensors=None, |
| 113 | **kwargs, |
| 114 | ): |
| 115 | if text is not None: |
| 116 | if not isinstance(text, list): |
| 117 | text = [text] |
| 118 | text_inputs = self.tokenizer(text, **kwargs) |
| 119 | else: |
| 120 | text_inputs = {} |
| 121 | |
| 122 | image_inputs = {} |
| 123 | if images is not None: |
| 124 | if not isinstance(images, list): |
| 125 | images = [images] |
| 126 | if len(images) > 0: |
| 127 | image_inputs = self.image_processor( |
| 128 | images, return_tensors=return_tensors or "pt" |
| 129 | ) |
| 130 | |
| 131 | combined_inputs = {**text_inputs, **image_inputs} |
| 132 | |
| 133 | return BatchFeature(combined_inputs, tensor_type=return_tensors) |
| 134 | |
| 135 | |
| 136 | class OpenCUAMultiModalProcessor(BaseMultiModalProcessor[OpenCUAProcessingInfo]): |
| 137 | def _get_mm_fields_config( |
| 138 | self, |
| 139 | hf_inputs: BatchFeature, |
| 140 | hf_processor_mm_kwargs: Mapping[str, object], |
| 141 | ) -> Mapping[str, MultiModalFieldConfig]: |
| 142 | return _create_qwen2vl_field_factory( |
| 143 | self.info.get_hf_config().vision_config.spatial_merge_size |
| 144 | )(hf_inputs) |
| 145 | |
| 146 | def _get_prompt_updates( |
| 147 | self, |
| 148 | mm_items: MultiModalDataItems, |
| 149 | hf_processor_mm_kwargs: Mapping[str, Any], |
| 150 | out_mm_kwargs: MultiModalKwargsItems, |
| 151 | ) -> Sequence[PromptUpdate]: |
| 152 | hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs) |
| 153 | image_processor = self.info.get_image_processor(**hf_processor_mm_kwargs) |
| 154 | tokenizer = self.info.get_tokenizer() |
| 155 | vocab = tokenizer.get_vocab() |
| 156 | hf_config = self.info.get_hf_config() |
| 157 | |
| 158 | image_token_str = getattr(hf_processor, "image_token", "<|media_placeholder|>") |
| 159 | image_token_id = vocab.get( |
| 160 | image_token_str, |
| 161 | getattr(hf_config, "media_placeholder_token_id", 151664), |
| 162 | ) |
| 163 | |
| 164 | merge_length = image_processor.merge_size**2 |
| 165 | |
| 166 | def get_replacement_opencua(item_idx: int): |
| 167 | out_item = out_mm_kwargs["image"][item_idx] |
| 168 | grid_thw = out_item["image_grid_thw"].data |
| 169 | assert isinstance(grid_thw, torch.Tensor) |
| 170 | |
| 171 | num_tokens = int(grid_thw.prod()) // merge_length |
| 172 | return [image_token_id] * num_tokens |
| 173 | |
| 174 | return [ |
| 175 | PromptReplacement( |
| 176 | modality="image", |
| 177 | target=[image_token_id], |
| 178 | replacement=get_replacement_opencua, |
| 179 | ) |
| 180 | ] |
| 181 | |
| 182 | |
| 183 | if typing.TYPE_CHECKING: |
| 184 | _OpenCUADummyInputsBuilderBase = BaseDummyInputsBuilder[OpenCUAProcessingInfo] |
| 185 | else: |
| 186 | _OpenCUADummyInputsBuilderBase = Qwen2VLDummyInputsBuilder |
| 187 | |
| 188 | |
| 189 | class OpenCUADummyInputsBuilder(_OpenCUADummyInputsBuilderBase): |
| 190 | def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: |
| 191 | num_images = mm_counts.get("image", 0) |
| 192 | |
| 193 | image_token = "<|media_placeholder|>" |
| 194 | |
| 195 | return image_token * num_images |
| 196 | |
| 197 | |
| 198 | @MULTIMODAL_REGISTRY.register_processor( |
| 199 | OpenCUAMultiModalProcessor, |
| 200 | info=OpenCUAProcessingInfo, |
| 201 | dummy_inputs=OpenCUADummyInputsBuilder, |
| 202 | ) |
| 203 | class OpenCUAForConditionalGeneration(Qwen2_5_VLForConditionalGeneration): |
| 204 | packed_modules_mapping = { |
| 205 | "qkv_proj": ["q_proj", "k_proj", "v_proj"], |
| 206 | "gate_up_proj": ["gate_proj", "up_proj"], |
| 207 | } |
| 208 | |
| 209 | hf_to_vllm_mapper = WeightsMapper( |
| 210 | orig_to_new_prefix={ |
| 211 | "model.language_model.": "language_model.model.", |
| 212 | "model.visual.": "visual.", |
| 213 | "vision_tower.": "visual.", |
| 214 | "lm_head.": "language_model.lm_head.", |
| 215 | "model.": "language_model.model.", |
| 216 | } |
| 217 | ) |
| 218 | |
| 219 | supports_encoder_tp_data = True |
| 220 | |
| 221 | @classmethod |
| 222 | def get_placeholder_str(cls, modality: str, i: int) -> str | None: |
| 223 | if modality.startswith("image"): |
| 224 | return "<|media_placeholder|>" |
| 225 | raise ValueError("Only image modality is supported") |
| 226 | |
| 227 | def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): |
| 228 | nn.Module.__init__(self) |
| 229 | config = vllm_config.model_config.hf_config |
| 230 | quant_config = vllm_config.quant_config |
| 231 | multimodal_config = vllm_config.model_config.get_multimodal_config() |
| 232 | |
| 233 | self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data" |
| 234 | self.config = config |
| 235 | self.vllm_config = vllm_config |
| 236 | self.multimodal_config = multimodal_config |
| 237 | self.quant_config = quant_config |
| 238 | self.is_multimodal_pruning_enabled = ( |
| 239 | multimodal_config.is_multimodal_pruning_enabled() |
| 240 | ) |
| 241 | |
| 242 | with self._mark_tower_model(vllm_config, "image"): |
| 243 | self.visual = OpenCUAVisionTransformer( |
| 244 | vision_config=config.vision_config, |
| 245 | norm_eps=getattr(config, "rms_norm_eps", 1e-6), |
| 246 | quant_config=self.quant_config, |
| 247 | prefix=maybe_prefix(prefix, "visual"), |
| 248 | ) |
| 249 | |
| 250 | with self._mark_language_model(vllm_config): |
| 251 | self.language_model = init_vllm_registered_model( |
| 252 | vllm_config=vllm_config, |
| 253 | hf_config=config.text_config, |
| 254 | prefix=maybe_prefix(prefix, "language_model"), |
| 255 | architectures=["Qwen2ForCausalLM"], |
| 256 | ) |
| 257 | |
| 258 | self.make_empty_intermediate_tensors = ( |
| 259 | self.language_model.make_empty_intermediate_tensors |
| 260 | ) |