Skip to content
Draft
Show file tree
Hide file tree
Changes from 2 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions docs/source/openvino/models.mdx
Original file line number Diff line number Diff line change
Expand Up @@ -88,6 +88,7 @@ Here is the list of the supported architectures :
- InternLM2
- InternVL2
- Jais
- Jina VLM
- LeViT
- LFM2
- LFM2-MoE
Expand Down
94 changes: 94 additions & 0 deletions optimum/exporters/openvino/input_generators.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
]
125 changes: 125 additions & 0 deletions optimum/exporters/openvino/model_configs.py
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,7 @@
DummyFluxTransformerInputGenerator,
DummyGemma4UnifiedVisionInputGenerator,
DummyGemma4VisionInputGenerator,
DummyJinaVLMVisionInputGenerator,
DummyKokoroInputGenerator,
DummyLLavaMultiModalProjectorInputGenerator,
DummyMiniCPMVImageInputGenerator,
Expand Down Expand Up @@ -77,6 +78,7 @@
FunASRDummyAudioInputGenerator,
Gemma4DummyPastKeyValuesGenerator,
GPTBigCodeDummyPastKeyValuesGenerator,
JinaVLMDummyPastKeyValuesGenerator,
Lfm2DummyPastKeyValuesGenerator,
LTX2AudioVaeDecoderDummyInputGenerator,
LTX2ConnectorsDummyInputGenerator,
Expand Down Expand Up @@ -132,6 +134,8 @@
InternVL2ChatLangModelPatcher,
InternVLChatImageEmbeddingModelPatcher,
JaisModelPatcher,
JinaVLMLanguageModelPatcher,
JinaVLMVisionEmbeddingsModelPatcher,
KokoroModelPatcher,
Lfm2ModelPatcher,
Lfm2MoeModelPatcher,
Expand Down Expand Up @@ -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"
Expand Down Expand Up @@ -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__(
Expand Down
Loading
Loading