diff --git a/docs/source/openvino/models.mdx b/docs/source/openvino/models.mdx index 902dd2ad75..b21ff609fd 100644 --- a/docs/source/openvino/models.mdx +++ b/docs/source/openvino/models.mdx @@ -88,6 +88,7 @@ Here is the list of the supported architectures : - InternLM2 - InternVL2 - Jais +- Jina VLM - LeViT - LFM2 - LFM2-MoE diff --git a/optimum/exporters/openvino/input_generators.py b/optimum/exporters/openvino/input_generators.py index 34f75923ba..45c0b18b67 100644 --- a/optimum/exporters/openvino/input_generators.py +++ b/optimum/exporters/openvino/input_generators.py @@ -2124,3 +2124,97 @@ def generate(self, input_name, framework="pt", int_dtype="int64", float_dtype="f return self.random_float_tensor([seq_len, self.embed_dim], framework=framework, dtype=float_dtype) return super().generate(input_name, framework, int_dtype, float_dtype) + + +class DummyJinaVLMVisionInputGenerator(DummyInputGenerator): + """Dummy inputs for the JinaVLM (model_type='jvlm') vision embeddings sub-model. + + The JinaVLM processor produces image tensors as pre-extracted patches rather than raw + ``pixel_values``: + + * ``image_patches`` of shape ``(batch_size, n_crops, n_patches, n_pixels)`` where + ``n_pixels = n_channels * patch_size ** 2`` and ``n_patches = (input_h // patch_size) * + (input_w // patch_size)``; + * ``image_masks`` of shape ``(batch_size, n_crops, n_patches)`` describing padded patches. + """ + + SUPPORTED_INPUT_NAMES = ("image_patches", "image_masks") + + def __init__( + self, + task: str, + normalized_config: NormalizedVisionConfig, + batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], + **kwargs, + ): + self.task = task + self.normalized_config = normalized_config + self.batch_size = batch_size + config = normalized_config.config + self.n_channels = config.n_channels + self.patch_size = config.patch_size + input_size = config.input_size + input_h, input_w = (input_size[0], input_size[1]) if input_size is not None else (self.patch_size, self.patch_size) + self.n_patches = (input_h // self.patch_size) * (input_w // self.patch_size) + self.n_pixels = self.n_channels * self.patch_size * self.patch_size + # A single crop is enough to trace the vision graph; the runtime supports dynamic n_crops. + self.n_crops = 1 + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + if input_name == "image_patches": + return self.random_float_tensor( + shape=[self.batch_size, self.n_crops, self.n_patches, self.n_pixels], + framework=framework, + dtype=float_dtype, + ) + if input_name == "image_masks": + return self.random_int_tensor( + shape=[self.batch_size, self.n_crops, self.n_patches], + min_value=0, + max_value=2, + framework=framework, + dtype=int_dtype, + ) + raise ValueError(f"Unsupported input name {input_name} for DummyJinaVLMVisionInputGenerator") + + +class JinaVLMDummyPastKeyValuesGenerator(DummyPastKeyValuesGenerator): + """Past key/values generator for the JinaVLM (model_type='jvlm') language model. + + The JinaVLM text decoder uses grouped-query attention with the number of key/value heads and + the per-head dimension defined in ``text_config.block_config.attn_config``. + """ + + def __init__( + self, + task: str, + normalized_config: NormalizedConfig, + batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], + sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], + **kwargs, + ): + super().__init__( + task=task, + normalized_config=normalized_config, + batch_size=batch_size, + sequence_length=sequence_length, + **kwargs, + ) + attn_config = normalized_config.config.block_config.attn_config + self.num_key_value_heads = attn_config.n_kv_heads or attn_config.n_heads + self.head_dim = attn_config.head_dim or (self.hidden_size // self.num_attention_heads) + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + shape = ( + self.batch_size, + self.num_key_value_heads, + self.sequence_length, + self.head_dim, + ) + return [ + ( + self.random_float_tensor(shape, framework=framework, dtype=float_dtype), + self.random_float_tensor(shape, framework=framework, dtype=float_dtype), + ) + for _ in range(self.num_layers) + ] diff --git a/optimum/exporters/openvino/model_configs.py b/optimum/exporters/openvino/model_configs.py index 1bb72ff0e4..e36088ac13 100644 --- a/optimum/exporters/openvino/model_configs.py +++ b/optimum/exporters/openvino/model_configs.py @@ -45,6 +45,7 @@ DummyFluxTransformerInputGenerator, DummyGemma4UnifiedVisionInputGenerator, DummyGemma4VisionInputGenerator, + DummyJinaVLMVisionInputGenerator, DummyKokoroInputGenerator, DummyLLavaMultiModalProjectorInputGenerator, DummyMiniCPMVImageInputGenerator, @@ -77,6 +78,7 @@ FunASRDummyAudioInputGenerator, Gemma4DummyPastKeyValuesGenerator, GPTBigCodeDummyPastKeyValuesGenerator, + JinaVLMDummyPastKeyValuesGenerator, Lfm2DummyPastKeyValuesGenerator, LTX2AudioVaeDecoderDummyInputGenerator, LTX2ConnectorsDummyInputGenerator, @@ -132,6 +134,8 @@ InternVL2ChatLangModelPatcher, InternVLChatImageEmbeddingModelPatcher, JaisModelPatcher, + JinaVLMLanguageModelPatcher, + JinaVLMVisionEmbeddingsModelPatcher, KokoroModelPatcher, Lfm2ModelPatcher, Lfm2MoeModelPatcher, @@ -1974,6 +1978,22 @@ def patch_model_for_export(self, model: PreTrainedModel, model_kwargs: Optional[ return CommonImageEmbeddingsModelPatcher(self, model, model_kwargs) +class JinaVLMNormalizedTextConfig(NormalizedTextConfig): + NUM_LAYERS = "num_hidden_layers" + HIDDEN_SIZE = "hidden_size" + VOCAB_SIZE = "vocab_size" + + @property + def num_attention_heads(self): + return self.config.block_config.attn_config.n_heads + + +class JinaVLMLanguageOpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): + DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, JinaVLMDummyPastKeyValuesGenerator) + DUMMY_PKV_GENERATOR_CLASS = JinaVLMDummyPastKeyValuesGenerator + NORMALIZED_CONFIG_CLASS = JinaVLMNormalizedTextConfig + + @register_in_tasks_manager("llava", *["image-text-to-text"], library_name="transformers") class LlavaOpenVINOConfig(BaseVLMOpenVINOConfig): _OV_2026_1_MODEL_TYPE = "llava" @@ -4619,6 +4639,111 @@ def __init__( self._normalized_config = self.NORMALIZED_CONFIG_CLASS(self._config) +@register_in_tasks_manager("jvlm", *["image-text-to-text"], library_name="transformers") +class JinaVLMOpenVINOConfig(BaseVLMOpenVINOConfig): + MIN_TRANSFORMERS_VERSION = "4.57.0" + DUMMY_INPUT_GENERATOR_CLASSES = (DummyJinaVLMVisionInputGenerator,) + + def __init__( + self, + config: "PretrainedConfig", + task: str = "feature-extraction", + int_dtype: str = "int64", + float_dtype: str = "fp32", + behavior: VLMConfigBehavior = VLMConfigBehavior.VISION_EMBEDDINGS, + preprocessors: Optional[List[Any]] = None, + **kwargs, + ): + super().__init__( + config=config, + task=task, + int_dtype=int_dtype, + float_dtype=float_dtype, + preprocessors=preprocessors, + ) + self._orig_config = config + if self._behavior == VLMConfigBehavior.VISION_EMBEDDINGS and hasattr(config, "vision_config"): + self._config = config.vision_config + self._normalized_config = NormalizedVisionConfig(self._config) + + @property + def inputs(self) -> Dict[str, Dict[int, str]]: + if self._behavior != VLMConfigBehavior.VISION_EMBEDDINGS: + return {} + return { + "image_patches": {0: "batch_size", 1: "num_crops", 2: "num_patches"}, + "image_masks": {0: "batch_size", 1: "num_crops", 2: "num_patches"}, + } + + @property + def outputs(self) -> Dict[str, Dict[int, str]]: + if self._behavior != VLMConfigBehavior.VISION_EMBEDDINGS: + return {} + return {"last_hidden_state": {0: "batch_size", 1: "num_image_tokens"}} + + def with_behavior(self, behavior: Union[str, VLMConfigBehavior]): + if isinstance(behavior, str) and not isinstance(behavior, VLMConfigBehavior): + behavior = VLMConfigBehavior(behavior) + + if behavior == VLMConfigBehavior.TEXT_EMBEDDINGS: + InputEmbedOpenVINOConfig.NORMALIZED_CONFIG_CLASS = JinaVLMNormalizedTextConfig + return InputEmbedOpenVINOConfig( + self._orig_config.text_config, + task="feature-extraction", + int_dtype=self.int_dtype, + float_dtype=self.float_dtype, + ) + + if behavior == VLMConfigBehavior.LANGUAGE: + internal_config = JinaVLMLanguageOpenVINOConfig( + self._orig_config.text_config, + task="text-generation", + use_past=True, + use_past_in_inputs=True, + int_dtype=self.int_dtype, + float_dtype=self.float_dtype, + ) + export_config = LMInputEmbedsConfigHelper( + internal_config, + patcher_cls=JinaVLMLanguageModelPatcher, + ) + export_config._normalized_config = internal_config._normalized_config + return export_config + + if behavior == VLMConfigBehavior.VISION_EMBEDDINGS: + return self.__class__( + self._orig_config, + task=self.task, + int_dtype=self.int_dtype, + float_dtype=self.float_dtype, + behavior=behavior, + preprocessors=self._preprocessors, + ) + + def get_model_for_behavior(self, model, behavior: Union[str, VLMConfigBehavior]): + if isinstance(behavior, str) and not isinstance(behavior, VLMConfigBehavior): + behavior = VLMConfigBehavior(behavior) + + if behavior == VLMConfigBehavior.LANGUAGE: + # JinaVLMForConditionalGeneration owns both the text decoder and the lm_head; the + # language patcher wraps its forward to consume `inputs_embeds` and emit logits. + return model + + if behavior == VLMConfigBehavior.VISION_EMBEDDINGS: + return model + + if behavior == VLMConfigBehavior.TEXT_EMBEDDINGS: + text_embedding = model.get_input_embeddings() + text_embedding.config = model.config.text_config + return text_embedding + + def patch_model_for_export(self, model: PreTrainedModel, model_kwargs: Optional[Dict[str, Any]] = None): + model_kwargs = model_kwargs or {} + if self._behavior != VLMConfigBehavior.VISION_EMBEDDINGS: + return super().patch_model_for_export(model, model_kwargs) + return JinaVLMVisionEmbeddingsModelPatcher(self, model, model_kwargs) + + @register_in_tasks_manager("gemma3", *["image-text-to-text"], library_name="transformers") class Gemma3OpenVINOConfig(BaseVLMOpenVINOConfig): def __init__( diff --git a/optimum/exporters/openvino/model_patcher.py b/optimum/exporters/openvino/model_patcher.py index a93cf65a7d..77ecd980a6 100644 --- a/optimum/exporters/openvino/model_patcher.py +++ b/optimum/exporters/openvino/model_patcher.py @@ -3063,6 +3063,149 @@ def __exit__(self, exc_type, exc_value, traceback): self._model.forward = self._model.__orig_forward +def jina_vlm_vision_embed_forward(self, image_patches, image_masks=None): + # Runs the JinaVLM vision tower + vision-language connector and returns the flattened image + # embeddings. This mirrors the tensor operations in `JinaVLMForConditionalGeneration._encode_images` + # and `JinaVLMVisionModel.forward` from the original PyTorch remote code: + # https://huggingface.co/jinaai/jina-vlm/blob/ddfa80b180f87f59873fd1cea352dec51183ab88/modeling_jvlm.py#L525 + # (`_encode_images`) and + # https://huggingface.co/jinaai/jina-vlm/blob/ddfa80b180f87f59873fd1cea352dec51183ab88/modeling_jvlm.py#L223 + # (`JinaVLMVisionModel.forward`). + # It replaces the multi-dim `torch.all(..., keepdim=True)` padding-mask reduction with an + # OpenVINO-traceable equivalent, and drops the `image_input_idx` bookkeeping (only used for + # reshaping/sorting). At inference time the runtime scatters these features into the text + # embeddings using the processor-provided `image_input_idx`. + inner = self.model + vision_model = inner.vision_model + batch_size, n_crops, n_patches, n_pixels = image_patches.shape + + flat_patches = image_patches.reshape(batch_size * n_crops, n_patches, n_pixels) + # A crop is fully-padded when every pixel equals -1. Reduce the two trailing dims explicitly + # so the resulting mask keeps shape (batch*crops, 1, 1) after tracing. + is_pad = (flat_patches == -1).all(dim=-1).all(dim=-1) # (batch*crops,) + mask = (~is_pad).reshape(batch_size * n_crops, 1, 1).to(flat_patches.dtype) + + out = vision_model.get_visual_features(flat_patches) + image_features = out.hidden_states + features = [] + for layer in vision_model.vit_layers: + feats = image_features[layer] + if vision_model.n_prefix_tokens > 0: + feats = feats[:, 1:] + features.append(feats) + image_features = torch.cat(features, dim=-1) + image_features = image_features * mask + image_features = image_features.reshape(batch_size, n_crops, n_patches, -1).contiguous() + + if image_masks is not None: + image_masks = image_masks.reshape(batch_size, n_crops, n_patches) + image_features = vision_model.vl_connector(image_features, image_masks) + + # image_features: (batch_size, n_crops, n_pooled_patches, hidden_size) + _, __, n_pooled_patches, hidden_size = image_features.shape + # Add the image-patch token embedding, exactly as in the reference implementation. + patch_token = torch.tensor([[151938]], dtype=torch.long, device=image_features.device) + patch_embed = inner.language_model.embedding(patch_token) + image_features = image_features.reshape(batch_size, n_crops * n_pooled_patches, hidden_size) + image_features = image_features + patch_embed.reshape(1, 1, hidden_size) + return image_features + + +class JinaVLMVisionEmbeddingsModelPatcher(ModelPatcher): + def __init__( + self, + config: "OpenVINOConfig", + model: "PreTrainedModel", + model_kwargs: Dict[str, Any], + ): + model.__orig_forward = model.forward + model.forward = types.MethodType(jina_vlm_vision_embed_forward, model) + # Force eager attention so tracing does not rely on flash/sdpa backends inside the + # vision tower and the vision-language connector. + self._orig_attn_implementation = model.config._attn_implementation + model.config._attn_implementation = "sdpa" + if getattr(model.config, "vision_config", None) is not None: + model.config.vision_config._attn_implementation = "sdpa" + super().__init__(config, model, model_kwargs) + + def __exit__(self, exc_type, exc_value, traceback): + super().__exit__(exc_type, exc_value, traceback) + self._model.forward = self._model.__orig_forward + self._model.config._attn_implementation = self._orig_attn_implementation + if getattr(self._model.config, "vision_config", None) is not None: + self._model.config.vision_config._attn_implementation = self._orig_attn_implementation + + +def jina_vlm_language_model_forward( + self, + inputs_embeds, + attention_mask=None, + position_ids=None, + past_key_values=None, + **kwargs, +): + # `self` is the top-level JinaVLMForConditionalGeneration. Route embeddings through the text + # decoder + lm_head, returning logits and the updated cache. The text decoder already handles + # `torch.jit.is_tracing()` for cache creation, so tracing produces a valid graph. + # This mirrors the tensor operations in `JinaVLMForConditionalGeneration.forward` from the + # original PyTorch remote code: + # https://huggingface.co/jinaai/jina-vlm/blob/ddfa80b180f87f59873fd1cea352dec51183ab88/modeling_jvlm.py#L749 + seq_len = inputs_embeds.shape[1] + past_len = past_key_values.get_seq_length() if past_key_values is not None else 0 + # JinaVLM builds its causal mask from the full attention_mask length (past + current). At + # runtime OVModelWithEmbedForCausalLM already supplies the full-length mask; during tracing the + # exporter provides a mask padded only to the past length, so extend it to (past + seq) here to + # keep the attention narrow operations consistent with the key/value length. + if attention_mask is not None: + expected_len = past_len + seq_len + if attention_mask.shape[-1] < expected_len: + pad = torch.ones( + (attention_mask.shape[0], expected_len - attention_mask.shape[-1]), + dtype=attention_mask.dtype, + device=attention_mask.device, + ) + attention_mask = torch.cat([attention_mask, pad], dim=-1) + outputs = self.model.language_model( + input_ids=None, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=True, + ) + logits = self.lm_head(outputs.last_hidden_state) + return {"logits": logits, "past_key_values": outputs.past_key_values} + + +class JinaVLMLanguageModelPatcher(OVDecoderModelPatcher): + def __init__( + self, + config: "OpenVINOConfig", + model: "PreTrainedModel", + model_kwargs: Dict[str, Any], + ): + model.__orig_forward = model.forward + model.forward = types.MethodType(jina_vlm_language_model_forward, model) + self._orig_attn_implementation = model.config._attn_implementation + model.config._attn_implementation = "sdpa" + if getattr(model.config, "text_config", None) is not None: + model.config.text_config._attn_implementation = "sdpa" + super().__init__(config, model, model_kwargs) + + def __enter__(self): + # Skip OVDecoderModelPatcher rope/causal-mask patches: JinaVLM uses a custom text decoder + # that does not expose `_update_causal_mask` or the standard rope cache attributes. + ModelPatcher.__enter__(self) + return self + + def __exit__(self, exc_type, exc_value, traceback): + ModelPatcher.__exit__(self, exc_type, exc_value, traceback) + self._model.forward = self._model.__orig_forward + self._model.config._attn_implementation = self._orig_attn_implementation + if getattr(self._model.config, "text_config", None) is not None: + self._model.config.text_config._attn_implementation = self._orig_attn_implementation + + class LlavaNextVideoImageEmbeddingModelPatcher(ModelPatcher): def __init__( self, diff --git a/optimum/exporters/openvino/utils.py b/optimum/exporters/openvino/utils.py index 55c6d852e2..dd84e03951 100644 --- a/optimum/exporters/openvino/utils.py +++ b/optimum/exporters/openvino/utils.py @@ -339,6 +339,7 @@ def _get_kokoro_submodels(model): "minicpmo", "videochat_flash_qwen", "qwen3_omni_moe", + "jvlm", ] SSM_MODELS = [ diff --git a/optimum/intel/openvino/modeling_visual_language.py b/optimum/intel/openvino/modeling_visual_language.py index 1fe784455c..a9f7c3fafb 100644 --- a/optimum/intel/openvino/modeling_visual_language.py +++ b/optimum/intel/openvino/modeling_visual_language.py @@ -29,6 +29,7 @@ AutoConfig, AutoImageProcessor, AutoModel, + AutoModelForCausalLM, AutoModelForImageTextToText, GenerationConfig, GenerationMixin, @@ -333,6 +334,8 @@ def __init__(self, model: ov.Model, parent_model: OVBaseModel) -> None: self._main_input = "images" elif model_has_input_output_name(self.model, "hidden_states"): self._main_input = "hidden_states" + elif model_has_input_output_name(self.model, "image_patches"): + self._main_input = "image_patches" else: self._main_input = "pixel_values" @@ -5640,6 +5643,132 @@ def _update_model_kwargs_for_generation( return model_kwargs +class _OVJinaVLMForCausalLM(OVModelForVisualCausalLM): + # JinaVLM's remote code registers JinaVLMForConditionalGeneration under AutoModelForCausalLM. + auto_model_class = AutoModelForCausalLM + + def get_vision_embeddings(self, pixel_values, input_ids=None, **kwargs): + # During decoding (single new token) no image is processed. + if input_ids is not None and input_ids.shape[1] == 1: + return None + image_masks = kwargs.get("image_masks") + image_features = self.vision_embeddings(pixel_values, image_masks=image_masks).last_hidden_state + return image_features + + def merge_vision_text_embeddings( + self, vision_embeds, inputs_embeds, input_ids=None, attention_mask=None, position_ids=None, **kwargs + ): + # JinaVLM interleaves image features into the text embeddings at positions given by + # `image_input_idx` (index-based scatter), matching the reference + # `_interleave_image_and_text_embeddings`. + image_embeds = torch.from_numpy(vision_embeds) if isinstance(vision_embeds, np.ndarray) else vision_embeds + inputs_embeds = torch.from_numpy(inputs_embeds) if isinstance(inputs_embeds, np.ndarray) else inputs_embeds + image_input_idx = kwargs.get("image_input_idx") + if image_input_idx is None: + raise ValueError("`image_input_idx` is required to merge JinaVLM image and text embeddings") + if isinstance(image_input_idx, np.ndarray): + image_input_idx = torch.from_numpy(image_input_idx) + + batch_size = inputs_embeds.shape[0] + hidden_size = inputs_embeds.shape[-1] + image_embeds = image_embeds.reshape(batch_size, -1, hidden_size).to(inputs_embeds.dtype) + image_input_idx = image_input_idx.reshape(batch_size, -1) + + valid = image_input_idx >= 0 + batch_idx = torch.arange(batch_size, device=inputs_embeds.device) + batch_idx = torch.tile(batch_idx[:, None], [1, image_input_idx.shape[1]]) + inputs_embeds = inputs_embeds.clone() + inputs_embeds[batch_idx[valid], image_input_idx[valid]] = image_embeds[valid] + + return inputs_embeds, attention_mask, position_ids + + def prepare_inputs_for_generation( + self, + input_ids, + past_key_values=None, + inputs_embeds=None, + pixel_values=None, + image_sizes=None, + attention_mask=None, + **kwargs, + ): + model_inputs = super().prepare_inputs_for_generation( + input_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + pixel_values=pixel_values, + image_sizes=image_sizes, + attention_mask=attention_mask, + **kwargs, + ) + # Carry the JinaVLM-specific image tensors through the first generation step only; after + # the prefill step the image inputs are no longer required. + if past_key_values is None: + model_inputs["image_patches"] = kwargs.get("image_patches") + model_inputs["image_masks"] = kwargs.get("image_masks") + model_inputs["image_input_idx"] = kwargs.get("image_input_idx") + return model_inputs + + def forward( + self, + input_ids, + pixel_values=None, + past_key_values=None, + inputs_embeds=None, + image_sizes=None, + attention_mask=None, + position_ids=None, + image_patches=None, + image_masks=None, + image_input_idx=None, + **kwargs, + ): + if pixel_values is None: + pixel_values = image_patches + return super().forward( + input_ids=input_ids, + pixel_values=pixel_values, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + image_sizes=image_sizes, + attention_mask=attention_mask, + position_ids=position_ids, + image_masks=image_masks, + image_input_idx=image_input_idx, + **kwargs, + ) + + @staticmethod + def preprocess_inputs( + text: Optional[str] = None, + image: Optional["Image"] = None, + processor: Optional[AutoImageProcessor] = None, + tokenizer: Optional[PreTrainedTokenizer] = None, + config: Optional[PretrainedConfig] = None, + video: Optional["VideoInput"] = None, + audio: Optional[np.ndarray] = None, + ): + if processor is None: + raise ValueError("processor is required") + if video is not None: + raise ValueError("Video input is not supported") + if audio is not None: + raise ValueError("Audio input is not supported") + if getattr(processor, "chat_template", None) is None: + raise ValueError("JinaVLM requires a chat template to build image-text inputs.") + + images = None + content = [] + if image is not None: + images = image if isinstance(image, list) else [image] + content.extend([{"type": "image"}] * len(images)) + content.append({"type": "text", "text": text}) + conversation = [{"role": "user", "content": content}] + prompt = processor.apply_chat_template(conversation, add_generation_prompt=True, tokenize=False) + inputs = processor(images=images, text=prompt, return_tensors="pt") + return inputs + + class _OVGotOCR2ForCausalLM(OVModelForVisualCausalLM): def get_vision_embeddings(self, pixel_values, input_ids, **kwargs): if input_ids is not None and input_ids.shape[1] == 1 and kwargs.get("past_key_values") is not None: @@ -7377,6 +7506,7 @@ def generate(self, *args, **kwargs): "qwen2_5_vl": _OVQwen2_5_VLForCausalLM, "qwen2_5_vl_text": _OVQwen2_5_VLForCausalLM, "got_ocr2": _OVGotOCR2ForCausalLM, + "jvlm": _OVJinaVLMForCausalLM, "gemma3": _OVGemma3ForCausalLM, "gemma3n": _OVGemma4ForCausalLM, "gemma4": _OVGemma4ForCausalLM, diff --git a/tests/openvino/test_export.py b/tests/openvino/test_export.py index d9928449a3..51f7adc988 100644 --- a/tests/openvino/test_export.py +++ b/tests/openvino/test_export.py @@ -90,6 +90,7 @@ class ExportModelTest(unittest.TestCase): "stable-diffusion-xl-refiner": OVStableDiffusionXLImg2ImgPipeline, "latent-consistency": OVLatentConsistencyModelPipeline, "llava": OVModelForVisualCausalLM, + "jvlm": OVModelForVisualCausalLM, "sam": OVSamModel, "speecht5": OVModelForTextToSpeechSeq2Seq, "clip": OVModelForZeroShotImageClassification, @@ -141,7 +142,7 @@ class ExportModelTest(unittest.TestCase): "ltx2": {"text_encoder": "8.0", "vae_encoder": "8.0", "vae_decoder": "8.0"}, } - GENERATIVE_MODELS = ("pix2struct", "t5", "bart", "gpt2", "whisper", "llava", "speecht5") + GENERATIVE_MODELS = ("pix2struct", "t5", "bart", "gpt2", "whisper", "llava", "speecht5", "jvlm") def _openvino_export( self, @@ -166,7 +167,7 @@ def _openvino_export( model_class = TasksManager.get_model_class_for_task(task, library=library_name) model = model_class(f"hf_hub:{model_name}", pretrained=True, exportable=True) TasksManager.standardize_model_attributes(model_name, model, library_name=library_name) - elif model_type in ["llava", "videochat_flash_qwen"]: + elif model_type in ["llava", "videochat_flash_qwen", "jvlm"]: model = MODEL_TYPE_TO_CLS_MAPPING[model_type].auto_model_class.from_pretrained( model_name, **loading_kwargs ) @@ -281,7 +282,9 @@ def test_export_with_custom_gen_config(self, model_type): task = auto_model.export_feature model_name = MODEL_NAMES[model_type] loading_kwargs = {"attn_implementation": "eager"} if model_type in SDPA_ARCHS_ONNX_EXPORT_NOT_SUPPORTED else {} - if model_type == "llava": + if model_type in REMOTE_CODE_MODELS: + loading_kwargs["trust_remote_code"] = True + if model_type in ["llava", "jvlm"]: model = MODEL_TYPE_TO_CLS_MAPPING[model_type].auto_model_class.from_pretrained( model_name, **loading_kwargs ) @@ -311,7 +314,10 @@ def test_export_with_custom_gen_config(self, model_type): ) use_cache = supported_task.endswith("-with-past") - ov_model = auto_model.from_pretrained(tmpdirname, use_cache=use_cache) + trust_remote_code = model_type in REMOTE_CODE_MODELS + ov_model = auto_model.from_pretrained( + tmpdirname, use_cache=use_cache, trust_remote_code=trust_remote_code + ) self.assertIsInstance(ov_model, OVBaseModel) self.assertTrue(ov_model.can_generate()) self.assertTrue(ov_model.generation_config is not None) @@ -321,7 +327,9 @@ def test_export_with_custom_gen_config(self, model_type): # check that generate config remains after repeated saving with TemporaryDirectory() as tmpdirname2: ov_model.save_pretrained(tmpdirname2) - ov_model = auto_model.from_pretrained(tmpdirname2, use_cache=use_cache) + ov_model = auto_model.from_pretrained( + tmpdirname2, use_cache=use_cache, trust_remote_code=trust_remote_code + ) self.assertIsInstance(ov_model, OVBaseModel) self.assertTrue(ov_model.can_generate()) self.assertTrue(ov_model.generation_config is not None) diff --git a/tests/openvino/test_exporters_cli.py b/tests/openvino/test_exporters_cli.py index c57789a4c5..88533fbca5 100644 --- a/tests/openvino/test_exporters_cli.py +++ b/tests/openvino/test_exporters_cli.py @@ -136,6 +136,7 @@ class OVCLIExportTestCase(unittest.TestCase): ("text-generation-with-past", "mamba"), ("text-generation-with-past", "falcon_mamba"), ("text-to-image", "flux.2-klein"), + ("image-text-to-text", "jvlm"), ] # filter architectures depending on min/max transformers supported versions SUPPORTED_ARCHITECTURES = [ @@ -166,6 +167,7 @@ class OVCLIExportTestCase(unittest.TestCase): ), # Tokenizers fail to convert on 2025.4, ticket: CVS-176880 "lfm2_moe": 2, "llava": 2, + "jvlm": 2, "sana": 2, "ltx-video": 2, "ltx2": 2, diff --git a/tests/openvino/test_seq2seq.py b/tests/openvino/test_seq2seq.py index 371b729b5e..31745128de 100644 --- a/tests/openvino/test_seq2seq.py +++ b/tests/openvino/test_seq2seq.py @@ -600,6 +600,7 @@ class OVModelForVisualCausalLMIntegrationTest(OVSeq2SeqTestMixin): "qwen3_5", "qwen3_5_moe", "qwen3_omni_moe", + "jvlm", ] SUPPORT_VIDEO = ["llava_next_video", "qwen2_vl", "qwen2_5_vl", "qwen3_vl", "videochat_flash_qwen"] SUPPORT_AUDIO = ["qwen3_omni_moe"] @@ -633,6 +634,7 @@ class OVModelForVisualCausalLMIntegrationTest(OVSeq2SeqTestMixin): "phi4mm", "videochat_flash_qwen", "gemma3n", + "jvlm", ] IMAGE = Image.open( requests.get( diff --git a/tests/openvino/utils_tests.py b/tests/openvino/utils_tests.py index 74418694cc..ee5aa93c19 100644 --- a/tests/openvino/utils_tests.py +++ b/tests/openvino/utils_tests.py @@ -144,6 +144,102 @@ def _create_tiny_kokoro_model(): return str(output_dir) +def _create_tiny_jvlm_model(): + """Generate a tiny random JinaVLM (model_type='jvlm') model for testing and return its path. + + JinaVLM is a trust-remote-code image-text-to-text architecture that is not published as a + tiny fixture on the Hub, so the fixture is synthesized locally from the original config and + remote-code assets (no original weights are downloaded). The reduced model preserves the real + architecture: model_type ('jvlm'), architectures (JinaVLMForConditionalGeneration), the nested + text/vision/vl_connector configs and flags, the special-token vocabulary, the fixed vision + patch geometry, GQA/head-dim coupling and the RoPE contract. The result is cached under the + system temp dir so subsequent calls are cheap. + """ + from transformers import AutoConfig, AutoModelForCausalLM, AutoProcessor + + original_model_id = "jinaai/jina-vlm" + cache_marker_value = "tiny-jvlm-v1" + output_dir = Path(tempfile.gettempdir()) / "optimum_intel_tiny_random_jvlm" + marker = output_dir / ".tiny_cache_marker" + + def _cache_is_valid(): + if not ((output_dir / "config.json").exists() and marker.exists()): + return False + try: + if marker.read_text().strip() != cache_marker_value: + return False + cfg = AutoConfig.from_pretrained(output_dir, trust_remote_code=True) + except Exception: + return False + return ( + cfg.model_type == "jvlm" + and cfg.architectures == ["JinaVLMForConditionalGeneration"] + and cfg.text_config.hidden_size == cfg.vision_config.output_size + ) + + if _cache_is_valid(): + return str(output_dir) + + if output_dir.exists(): + import shutil + + shutil.rmtree(output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + + torch.manual_seed(SEED) + + config = AutoConfig.from_pretrained(original_model_id, trust_remote_code=True) + cfg = config.to_dict() + text = cfg["text_config"] + vision = cfg["vision_config"] + + # Tiny dimensions preserving divisibility / coupling invariants. + text["hidden_size"] = 64 + text["n_layers"] = 2 + text["num_hidden_layers"] = 2 + text["block_config"]["attn_config"].update({"n_heads": 4, "n_kv_heads": 2, "head_dim": 16}) + text["block_config"]["ffn_config"]["size"] = 128 + # keep vocab_size so image special-token ids remain valid embedding indices + + vision["hidden_size"] = 64 + # vit_layers references [-4, -10] -> need at least 10 vision layers + vision["n_layers"] = 10 + vision["output_size"] = text["hidden_size"] # invariant: vision output == text hidden + vision["block_config"]["attn_config"].update({"n_heads": 4, "head_dim": 16}) + vision["block_config"]["ffn_config"]["size"] = 128 + conn = vision["vl_connector_config"] + conn["attn_pooling_config"].update({"n_heads": 4, "head_dim": 16}) + conn["mlp_projector_config"]["size"] = 128 + + cfg["dtype"] = "float32" + cfg["torch_dtype"] = "float32" + for sub in (text, vision): + sub["dtype"] = "float32" + sub["torch_dtype"] = "float32" + cfg["text_config"] = text + cfg["vision_config"] = vision + + tiny_config = config.__class__.from_dict(cfg) + model = AutoModelForCausalLM.from_config(tiny_config, trust_remote_code=True).to(torch.float32) + with torch.no_grad(): + model.lm_head.weight.normal_(mean=0.0, std=0.2) + model.get_input_embeddings().weight.normal_(mean=0.0, std=0.2) + model.save_pretrained(output_dir, safe_serialization=True) + + processor = AutoProcessor.from_pretrained(original_model_id, trust_remote_code=True, use_fast=False) + processor.save_pretrained(output_dir) + + try: + from transformers import GenerationConfig + + GenerationConfig.from_pretrained(original_model_id).save_pretrained(output_dir) + except Exception: + pass + + marker.write_text(cache_marker_value) + return str(output_dir) + + SEED = 42 F32_CONFIG = {"INFERENCE_PRECISION_HINT": "f32"} @@ -238,6 +334,7 @@ def _create_tiny_kokoro_model(): "internlm2": "optimum-intel-internal-testing/tiny-random-internlm2", "internvl_chat": "optimum-intel-internal-testing/tiny-random-internvl2", "jais": "optimum-intel-internal-testing/tiny-random-jais", + "jvlm": _create_tiny_jvlm_model(), "kokoro": _create_tiny_kokoro_model(), "levit": "optimum-intel-internal-testing/tiny-random-LevitModel", "lfm2": "optimum-intel-internal-testing/tiny-random-lfm2", @@ -479,6 +576,11 @@ def _resolve_cached_model_paths(model_names: dict) -> dict: "text_embeddings_model": 1, "vision_embeddings_model": 9, }, + "jvlm": { + "lm_model": 18, + "text_embeddings_model": 1, + "vision_embeddings_model": 38, + }, "llava_next": { "lm_model": 30, "text_embeddings_model": 1, @@ -667,6 +769,7 @@ def _resolve_cached_model_paths(model_names: dict) -> dict: "qwen3_asr", "fun_asr", "videochat_flash_qwen", + "jvlm", ) if is_transformers_version("<", "5"): @@ -685,6 +788,7 @@ def _resolve_cached_model_paths(model_names: dict) -> dict: "qwen3_moe": "OVModelForCausalLM", "llama4": "OVModelForCausalLM", "llava": "OVModelForVisualCausalLM", + "jvlm": "OVModelForVisualCausalLM", "qwen3_5_moe": "OVModelForVisualCausalLM", "gemma4_moe": "OVModelForVisualCausalLM", "gemma4_unified": "OVModelForVisualCausalLM",