From 9fc0e2ac7dc8f18ef0df760149fb99068b3c2a78 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CElla?= Date: Thu, 21 May 2026 14:54:20 +0200 Subject: [PATCH 01/40] deprecate export_pytorch_via_onnx --- optimum/exporters/openvino/convert.py | 67 --------------------------- 1 file changed, 67 deletions(-) diff --git a/optimum/exporters/openvino/convert.py b/optimum/exporters/openvino/convert.py index a1a928501d..b65a841e17 100644 --- a/optimum/exporters/openvino/convert.py +++ b/optimum/exporters/openvino/convert.py @@ -235,73 +235,6 @@ def export( raise RuntimeError("You either provided a non-PyTorch model or the PyTorch library is not installed.") -def export_pytorch_via_onnx( - model: Union["PreTrainedModel", "ModelMixin"], - config: "OnnxConfig", - opset: int, - output: Path, - device: str = "cpu", - input_shapes: Optional[Dict] = None, - model_kwargs: Optional[Dict[str, Any]] = None, - ov_config: Optional["OVConfig"] = None, - library_name: Optional[str] = None, -): - """ - Exports a PyTorch model to an OpenVINO Intermediate Representation via ONNX export. - - Args: - model ([`PreTrainedModel`]): - The model to export. - config ([`~exporters.onnx.config.OnnxConfig`]): - The configuration associated with the exported model. - opset (`int`): - The version of the ONNX operator set to use. - output (`Path`): - Directory to store the exported model. - device (`str`, defaults to `"cpu"`): - The device on which the model will be exported. Either `cpu` or `cuda`. Only PyTorch is supported for - export on CUDA devices. - input_shapes (`optional[Dict]`, defaults to `None`): - If specified, allows to use specific shapes for the example input provided to the exporter. - model_kwargs (optional[Dict[str, Any]], defaults to `None`): - Additional kwargs for model export. - ov_config (`OVConfig`, *optional*): - The configuration containing the parameters related to quantization. - - Returns: - `Tuple[List[str], List[str], bool]`: A tuple with an ordered list of the model's inputs, and the named inputs from - the ONNX configuration and boolean flag - was legacy ONNX path were applied to model or not. - """ - import torch - - from optimum.exporters.onnx.convert import export_pytorch as export_pytorch_to_onnx - - output = Path(output) - orig_torch_onnx_export = torch.onnx.export - torch.onnx.export = functools.partial(orig_torch_onnx_export, do_constant_folding=False) - model.config.torchscript = False - model.config.return_dict = True - onnx_output = output.with_suffix(".onnx") - input_names, output_names = export_pytorch_to_onnx( - model, config, opset, onnx_output, device, input_shapes, model_kwargs - ) - torch.onnx.export = orig_torch_onnx_export - ov_model = convert_model(str(onnx_output)) - - library_name = _infer_library_from_model_or_model_class(model=model, library_name=library_name) - - _save_model( - ov_model, - output.parent / OV_XML_FILE_NAME if output.suffix != ".xml" else output, - ov_config=ov_config, - library_name=library_name, - config=config, - ) - del ov_model - gc.collect() - return input_names, output_names, True - - def export_pytorch( model: Union["PreTrainedModel", "ModelMixin"], config: "OnnxConfig", From 9a2f46f9fee7c0e7f5ecf2b2bab0bf86519056f2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CElla?= Date: Thu, 21 May 2026 14:54:39 +0200 Subject: [PATCH 02/40] remove optimum-onnx dep --- optimum/exporters/openvino/__init__.py | 2 +- optimum/exporters/openvino/__main__.py | 2 +- .../exporters/openvino/_traceable_cache.py | 107 + .../openvino/_traceable_decorator.py | 175 + optimum/exporters/openvino/base.py | 739 +++ optimum/exporters/openvino/config.py | 347 ++ .../exporters/openvino/input_generators.py | 1816 +++++++ optimum/exporters/openvino/model_configs.py | 4715 ++++++++++------- optimum/exporters/openvino/model_patcher.py | 1583 +++++- optimum/exporters/openvino/utils.py | 2 +- setup.py | 3 +- 11 files changed, 7611 insertions(+), 1880 deletions(-) create mode 100644 optimum/exporters/openvino/_traceable_cache.py create mode 100644 optimum/exporters/openvino/_traceable_decorator.py create mode 100644 optimum/exporters/openvino/base.py create mode 100644 optimum/exporters/openvino/config.py create mode 100644 optimum/exporters/openvino/input_generators.py diff --git a/optimum/exporters/openvino/__init__.py b/optimum/exporters/openvino/__init__.py index 94ea4f103b..320a15e38a 100644 --- a/optimum/exporters/openvino/__init__.py +++ b/optimum/exporters/openvino/__init__.py @@ -15,7 +15,7 @@ import optimum.exporters.openvino.model_configs from .__main__ import main_export -from .convert import export, export_from_model, export_models, export_pytorch_via_onnx +from .convert import export, export_from_model, export_models from .stateful import ensure_stateful_is_available, patch_stateful diff --git a/optimum/exporters/openvino/__main__.py b/optimum/exporters/openvino/__main__.py index 1fb3f5d43c..a962205462 100644 --- a/optimum/exporters/openvino/__main__.py +++ b/optimum/exporters/openvino/__main__.py @@ -28,7 +28,7 @@ from transformers.utils import is_torch_available from openvino import Core, Type, save_model -from optimum.exporters.onnx.base import OnnxConfig +from optimum.exporters.openvino.base import OnnxConfig from optimum.exporters.tasks import TasksManager from optimum.intel.utils.import_utils import ( DIFFUSERS_IMPORT_ERROR, diff --git a/optimum/exporters/openvino/_traceable_cache.py b/optimum/exporters/openvino/_traceable_cache.py new file mode 100644 index 0000000000..059463f9cf --- /dev/null +++ b/optimum/exporters/openvino/_traceable_cache.py @@ -0,0 +1,107 @@ +# Copyright 2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +import logging +from typing import Any + +import torch + + +logger = logging.getLogger(__name__) + + +# Simply removing the nn.Module, same as in https://github.com/huggingface/transformers/pull/35873 +class TraceableCache: + """Base, abstract class for all caches. The actual data structure is specific to each subclass.""" + + def __init__(self): + super().__init__() + + def update( + self, + key_states: torch.Tensor, + value_states: torch.Tensor, + layer_idx: int, + cache_kwargs: dict[str, Any] | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Updates the cache with the new `key_states` and `value_states` for the layer `layer_idx`. + + Parameters: + key_states (`torch.Tensor`): + The new key states to cache. + value_states (`torch.Tensor`): + The new value states to cache. + layer_idx (`int`): + The index of the layer to cache the states for. + cache_kwargs (`Dict[str, Any]`, `optional`): + Additional arguments for the cache subclass. These are specific to each subclass and allow new types of + cache to be created. + + Return: + A tuple containing the updated key and value states. + """ + raise NotImplementedError("Make sure to implement `update` in a subclass.") + + def get_seq_length(self, layer_idx: int | None = 0) -> int: + """Returns the sequence length of the cached states. A layer index can be optionally passed.""" + # TODO: deprecate this function in favor of `cache_position` + raise NotImplementedError("Make sure to implement `get_seq_length` in a subclass.") + + # Deprecate in favor of max-cache-shape because we want to be specific by what we mean with "max_length" + # Prev some cache objects didn't have "max_length" (SlidingWindowCache or SinkCache) because the cache object technically handles + # infinite amount of tokens. In the codebase what we really need to check is the max capacity of certain cache instances, so + # we change naming to be more explicit + def get_max_length(self) -> int | None: + logger.warning_once( + "`get_max_cache()` is deprecated for all Cache classes. Use `get_max_cache_shape()` instead. " + "Calling `get_max_cache()` will raise error from v4.48" + ) + return self.get_max_cache_shape() + + def get_max_cache_shape(self) -> int | None: + """Returns the maximum sequence length (i.e. max capacity) of the cache object.""" + raise NotImplementedError("Make sure to implement `get_max_cache_shape` in a subclass.") + + def get_usable_length(self, new_seq_length: int, layer_idx: int | None = 0) -> int: + """Given the sequence length of the new inputs, returns the usable length of the cache.""" + # Cache without size limit -> all cache is usable + # Cache with size limit -> if the length cache plus the length of the new inputs is larger the maximum cache + # length, we will need to evict part of the cache (and thus not all cache is usable) + max_length = self.get_max_cache_shape() + previous_seq_length = self.get_seq_length(layer_idx) + if max_length is not None and previous_seq_length + new_seq_length > max_length: + return max_length - new_seq_length + return previous_seq_length + + def reorder_cache(self, beam_idx: torch.LongTensor): + """Reorders the cache for beam search, given the selected beam indices.""" + for layer_idx in range(len(self.key_cache)): + if self.key_cache[layer_idx] != []: + device = self.key_cache[layer_idx].device + self.key_cache[layer_idx] = self.key_cache[layer_idx].index_select(0, beam_idx.to(device)) + if self.value_cache[layer_idx] != []: + device = self.value_cache[layer_idx].device + self.value_cache[layer_idx] = self.value_cache[layer_idx].index_select(0, beam_idx.to(device)) + + @property + def seen_tokens(self): + logger.warning_once( + "The `seen_tokens` attribute is deprecated and will be removed in v4.41. Use the `cache_position` " + "model input instead." + ) + if hasattr(self, "_seen_tokens"): + return self._seen_tokens + else: + return None diff --git a/optimum/exporters/openvino/_traceable_decorator.py b/optimum/exporters/openvino/_traceable_decorator.py new file mode 100644 index 0000000000..0a2087dff7 --- /dev/null +++ b/optimum/exporters/openvino/_traceable_decorator.py @@ -0,0 +1,175 @@ +# Copyright 2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import warnings +from collections import defaultdict +from functools import wraps + +from transformers.utils.generic import logger + + +try: + # transformers>=5.2 + from transformers.utils.output_capturing import _CAN_RECORD_REGISTRY, OutputRecorder +except ImportError: + from transformers.utils.generic import _CAN_RECORD_REGISTRY, OutputRecorder + + +# This is a fixed version of transformers.utils.generic.check_model_inputs +# that fixes issues related to onnx export and tracing +# - adds support for positional args (use_cache), without which use_cache end up being passed twice +# - fixes issue with default capture_flags being None for some models +def traceable_check_model_inputs(func): + @wraps(func) + def wrapper(self, *args, **kwargs): + use_cache = ( + kwargs["use_cache"] if kwargs.get("use_cache") is not None else getattr(self.config, "use_cache", None) + ) + if use_cache is not None: + if getattr(self, "gradient_checkpointing", False) and self.training and use_cache: + logger.warning_once( + "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`." + ) + use_cache = False + + # Prevent passing use_cache twice + if "use_cache" in func.__code__.co_varnames: + use_cache_idx = func.__code__.co_varnames.index("use_cache") - 1 # minus 1 for 'self' + if len(args) > use_cache_idx: + args = list(args) + args[use_cache_idx] = use_cache + args = tuple(args) + else: + kwargs["use_cache"] = use_cache + + return_dict = kwargs.pop("return_dict", None) + if return_dict is None: + return_dict = getattr(self.config, "return_dict", True) + + all_args = kwargs.copy() + if "kwargs" in all_args: + for k, v in all_args["kwargs"].items(): + all_args[k] = v + + capture_flags = _CAN_RECORD_REGISTRY.get(str(self.__class__)) or {} # there is a weak ref for executorch + + recordable_keys = { + f"output_{k}": all_args.get( + f"output_{k}", + getattr( + self.config, + f"output_{k}", + all_args.get("output_attentions", getattr(self.config, "output_attentions", False)), + ), + ) + for k in capture_flags + } + + # We let cross attentions to be saved separately because some models add `cross-attn` layer + # when certain conditions are met. Let's output cross attention if attentions are requested (for BC) + if "output_attentions" in recordable_keys: + recordable_keys["output_cross_attentions"] = recordable_keys["output_attentions"] + + collected_outputs = defaultdict(tuple) + monkey_patched_layers = [] + + # Check attention implementation is properly set for capturing attention outputs + if recordable_keys.get("output_attentions", False): + supported_attn = ["eager", "eager_paged", "flex_attention"] + config_attn = getattr(self.config, "_attn_implementation", None) + sub_configs = [getattr(self.config, key, None) for key in self.config.sub_configs] + sub_configs_attn = [ + getattr(config, "_attn_implementation", None) for config in sub_configs if config is not None + ] + if config_attn not in supported_attn or any(attn not in supported_attn for attn in sub_configs_attn): + warnings.warn( + f"`output_attentions=True` is not supported with `attn_implementation` other than {supported_attn}. " + "Please use `model.set_attn_implementation('eager')` to enable capturing attention outputs.", + UserWarning, + stacklevel=2, + ) + + def make_capture_wrapper(module, orig_forward, key, index): + @wraps(orig_forward) + def wrapped_forward(*args, **kwargs): + if key == "hidden_states" and len(collected_outputs[key]) == 0: + collected_outputs[key] += (args[0],) + output = orig_forward(*args, **kwargs) + if not isinstance(output, tuple): + collected_outputs[key] += (output,) + elif output[index] is not None: + if key not in collected_outputs: + collected_outputs[key] = (output[index],) + else: + collected_outputs[key] += (output[index],) + return output + + return wrapped_forward + + if any(recordable_keys.values()): + capture_tasks = [] + for key, layer_specs in capture_flags.items(): + if not recordable_keys.get(f"output_{key}", False): + continue + if not isinstance(layer_specs, list): + layer_specs = [layer_specs] + for specs in layer_specs: + if not isinstance(specs, OutputRecorder): + index = 0 if "hidden_states" in key else 1 + class_name = None if not isinstance(specs, str) else specs + target_class = specs if not isinstance(specs, str) else None + specs = OutputRecorder(target_class=target_class, index=index, class_name=class_name) + capture_tasks.append((key, specs)) + + for name, module in self.named_modules(): + for key, specs in capture_tasks: + # The second check is for multimodals where only backbone layer suffix is available + if (specs.target_class is not None and isinstance(module, specs.target_class)) or ( + specs.class_name is not None and name.endswith(specs.class_name) + ): + if specs.layer_name is not None and specs.layer_name not in name: + continue + # Monkey patch forward + original_forward = module.forward + module.forward = make_capture_wrapper(module, original_forward, key, specs.index) + monkey_patched_layers.append((module, original_forward)) + + outputs = func(self, *args, **kwargs) + # Restore original forward methods + for module, original_forward in monkey_patched_layers: + module.forward = original_forward + + # Inject collected outputs into model output + for key in collected_outputs: + if key == "hidden_states": + if hasattr(outputs, "vision_hidden_states"): + collected_outputs[key] = collected_outputs[key][:-1] + collected_outputs[key] += (outputs.vision_hidden_states,) + elif hasattr(outputs, "last_hidden_state"): + collected_outputs[key] = collected_outputs[key][:-1] + collected_outputs[key] += (outputs.last_hidden_state,) + + outputs[key] = collected_outputs[key] + elif key == "attentions": + if isinstance(capture_flags[key], list) and len(capture_flags[key]) == 2: + outputs[key] = collected_outputs[key][0::2] + outputs["cross_" + key] = collected_outputs[key][1::2] + else: + outputs[key] = collected_outputs[key] + else: + outputs[key] = collected_outputs[key] + if return_dict is False: + outputs = outputs.to_tuple() + return outputs + + return wrapper diff --git a/optimum/exporters/openvino/base.py b/optimum/exporters/openvino/base.py new file mode 100644 index 0000000000..d232785c1e --- /dev/null +++ b/optimum/exporters/openvino/base.py @@ -0,0 +1,739 @@ +# Copyright 2022 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""ONNX configuration base classes.""" + +from __future__ import annotations + +import enum +import gc +import inspect +import itertools +import re +from abc import ABC +from collections import OrderedDict +from collections.abc import Iterable +from pathlib import Path +from typing import Any, ClassVar + +import numpy as np +from transformers import PretrainedConfig, PreTrainedModel + +from optimum.exporters.base import ExporterConfig +from optimum.exporters.openvino.model_patcher import ModelPatcher, PatchingSpec +from optimum.utils import DEFAULT_DUMMY_SHAPES, DummyInputGenerator, DummySeq2SeqPastKeyValuesGenerator, logging +from optimum.utils.doc import add_dynamic_docstring +from optimum.utils.import_utils import ( + is_onnx_available, + is_onnxruntime_available, + is_transformers_version, +) + + +logger = logging.get_logger(__name__) + + +GENERATE_DUMMY_DOCSTRING = r""" + Generates the dummy inputs necessary for tracing the model. If not explicitly specified, default input shapes are used. + + Args: + framework (`str`, defaults to `"pt"`): + The framework for which to create the dummy inputs. + batch_size (`int`, defaults to {batch_size}): + The batch size to use in the dummy inputs. + sequence_length (`int`, defaults to {sequence_length}): + The sequence length to use in the dummy inputs. + num_choices (`int`, defaults to {num_choices}): + The number of candidate answers provided for multiple choice task. + image_width (`int`, defaults to {width}): + The width to use in the dummy inputs for vision tasks. + image_height (`int`, defaults to {height}): + The height to use in the dummy inputs for vision tasks. + num_channels (`int`, defaults to {num_channels}): + The number of channels to use in the dummpy inputs for vision tasks. + feature_size (`int`, defaults to {feature_size}): + The number of features to use in the dummpy inputs for audio tasks in case it is not raw audio. + This is for example the number of STFT bins or MEL bins. + nb_max_frames (`int`, defaults to {nb_max_frames}): + The number of frames to use in the dummpy inputs for audio tasks in case the input is not raw audio. + audio_sequence_length (`int`, defaults to {audio_sequence_length}): + The number of frames to use in the dummpy inputs for audio tasks in case the input is raw audio. + + Returns: + `Dict`: A dictionary mapping the input names to dummy tensors in the proper framework format. +""" + + +class OnnxConfig(ExporterConfig, ABC): + DEFAULT_ONNX_OPSET = 18 + VARIANTS: ClassVar[dict[str, str]] = {"default": "The default ONNX variant."} + DEFAULT_VARIANT = "default" + PATCHING_SPECS: list[PatchingSpec] | None = None + _MODEL_PATCHER = ModelPatcher + + _TASK_TO_COMMON_OUTPUTS = { # noqa: RUF012 + "audio-classification": OrderedDict({"logits": {0: "batch_size"}}), + "audio-frame-classification": OrderedDict({"logits": {0: "batch_size", 1: "sequence_length"}}), + "automatic-speech-recognition": OrderedDict({"logits": {0: "batch_size", 1: "sequence_length"}}), + "audio-xvector": OrderedDict({"logits": {0: "batch_size"}, "embeddings": {0: "batch_size"}}), + "depth-estimation": OrderedDict({"predicted_depth": {0: "batch_size", 1: "height", 2: "width"}}), + "document-question-answering": OrderedDict({"logits": {0: "batch_size", 1: "sequence_length"}}), + "feature-extraction": OrderedDict({"last_hidden_state": {0: "batch_size", 1: "sequence_length"}}), + "fill-mask": OrderedDict({"logits": {0: "batch_size", 1: "sequence_length"}}), + "image-classification": OrderedDict({"logits": {0: "batch_size"}}), + "image-segmentation": OrderedDict({"logits": {0: "batch_size", 2: "height", 3: "width"}}), + "image-to-text": OrderedDict({"logits": {0: "batch_size", 1: "sequence_length"}}), + "image-to-image": OrderedDict( + {"reconstruction": {0: "batch_size", 1: "num_channels", 2: "height", 3: "width"}} + ), + "keypoint-detection": OrderedDict( + {"heatmaps": {0: "batch_size", 1: "num_keypoints", 2: "height", 3: "width"}} + ), + "mask-generation": OrderedDict({"logits": {0: "batch_size"}}), + "masked-im": OrderedDict( + {"reconstruction" if is_transformers_version(">=", "4.29.0") else "logits": {0: "batch_size"}} + ), + "multiple-choice": OrderedDict({"logits": {0: "batch_size", 1: "num_choices"}}), + "object-detection": OrderedDict( + { + "logits": {0: "batch_size", 1: "num_queries"}, + "pred_boxes": {0: "batch_size", 1: "num_queries"}, + } + ), + "question-answering": OrderedDict( + { + "start_logits": {0: "batch_size", 1: "sequence_length"}, + "end_logits": {0: "batch_size", 1: "sequence_length"}, + } + ), + "semantic-segmentation": OrderedDict({"logits": {0: "batch_size", 1: "num_labels", 2: "height", 3: "width"}}), + "text2text-generation": OrderedDict({"logits": {0: "batch_size", 1: "decoder_sequence_length"}}), + "text-classification": OrderedDict({"logits": {0: "batch_size"}}), + "text-generation": OrderedDict({"logits": {0: "batch_size", 1: "sequence_length"}}), + "time-series-forecasting": OrderedDict({"prediction_outputs": {0: "batch_size"}}), + "token-classification": OrderedDict({"logits": {0: "batch_size", 1: "sequence_length"}}), + "visual-question-answering": OrderedDict({"logits": {0: "batch_size", 1: "sequence_length"}}), + "zero-shot-image-classification": OrderedDict( + { + "logits_per_image": {0: "image_batch_size", 1: "text_batch_size"}, + "logits_per_text": {0: "text_batch_size", 1: "image_batch_size"}, + "text_embeds": {0: "text_batch_size"}, + "image_embeds": {0: "image_batch_size"}, + } + ), + "zero-shot-object-detection": OrderedDict( + { + "logits": {0: "batch_size", 1: "num_queries"}, + "pred_boxes": {0: "batch_size", 1: "num_queries"}, + "text_embeds": {0: "text_batch_size"}, + "image_embeds": {0: "image_batch_size"}, + } + ), + } + + def __init__( + self, + config: PretrainedConfig, + task: str = "feature-extraction", + preprocessors: list[Any] | None = None, + int_dtype: str = "int64", + float_dtype: str = "fp32", + ): + super().__init__(config=config, task=task, int_dtype=int_dtype, float_dtype=float_dtype) + + self.variant = "default" + self._preprocessors = preprocessors + + @property + def variant(self) -> str: + """For a given ONNX config, the variant of the model to export. + + This property allows to define variants of a given model, in case + different users would like to export the model differently (with different inputs/outputs, model split in several ONNX or not, etc.). + """ + return self._variant + + @variant.setter + def variant(self, value: str): + if value == "default" and hasattr(self, "DEFAULT_VARIANT"): + value = self.DEFAULT_VARIANT + if value not in self.VARIANTS: + raise ValueError(f"The variant {value} is not supported for the ONNX config {self.__class__.__name__}.") + self._variant = value + + def fix_dynamic_axes( + self, model_path: Path, device: str = "cpu", dtype: str | None = None, input_shapes: dict | None = None + ): + """Fixes potential issues with dynamic axes. + + During the export, ONNX will infer some axes to be dynamic which are actually static. This method is called + right after the export to fix such issues. + + Args: + model_path (`Path`): + The path of the freshly exported ONNX model. + device (`str`, defaults to `"cpu"`): + The device on which the model will be run. This is used to determine the ONNX Runtime provider. + dtype (`Optional[str]`, defaults to `None`): + The data type of the model inputs. If `None`, it will be inferred from the model inputs. + input_shapes (`Optional[Dict[str, Any]]`, defaults to `None`): + The shapes of the model inputs. If `None`, it will be inferred from the model inputs. + """ + if not (is_onnx_available() and is_onnxruntime_available()): + raise RuntimeError( + "The onnx and onnxruntime packages are necessary to fix the dynamic shapes of the exported model. " + "You can install them by doing: pip install onnx onnxruntime" + ) + + import onnx + from onnxruntime import GraphOptimizationLevel, InferenceSession, SessionOptions + + allowed_dynamic_axes = set() + for input_ in self.inputs.values(): + allowed_dynamic_axes |= set(input_.values()) + for output in self.outputs.values(): + allowed_dynamic_axes |= set(output.values()) + + if device.startswith("cuda"): + providers = ["CUDAExecutionProvider"] + else: + providers = ["CPUExecutionProvider"] + + session_options = SessionOptions() + session_options.graph_optimization_level = GraphOptimizationLevel.ORT_DISABLE_ALL # no need to optimize here + session = InferenceSession(model_path.as_posix(), providers=providers, sess_options=session_options) + + to_fix = [] + for output_idx, node in enumerate(session.get_outputs()): + for idx, axis in enumerate(node.shape): + if isinstance(axis, str) and axis not in allowed_dynamic_axes: + to_fix.append((output_idx, idx)) + + # We branch here to avoid doing an unnecessary forward pass. + if to_fix: + if input_shapes is None: + input_shapes = {} + + onnx_input_names = [inp.name for inp in session.get_inputs()] + dummy_inputs = self.generate_dummy_inputs(framework="np", **input_shapes) + dummy_inputs = self.generate_dummy_inputs_for_validation(dummy_inputs, onnx_input_names) + dummy_inputs = self.rename_ambiguous_inputs(dummy_inputs) + + onnx_inputs = {} + for name, value in dummy_inputs.items(): + if isinstance(value, (list, tuple)): + value = self.flatten_output_collection_property(name, value) + onnx_inputs.update(dict(value.items())) + else: + onnx_inputs[name] = value + + for name, value in onnx_inputs.items(): + if value.dtype == np.float32 and dtype == "fp16": + onnx_inputs[name] = onnx_inputs[name].astype(np.float16) + + outputs = session.run(None, onnx_inputs) + del session + + onnx_model = onnx.load(model_path.as_posix(), load_external_data=False) + + for output_idx, dim_idx in to_fix: + dims = onnx_model.graph.output[output_idx].type.tensor_type.shape.dim + dims[dim_idx].dim_value = outputs[output_idx].shape[dim_idx] + + onnx.save( + onnx_model, + model_path.as_posix(), + convert_attribute=True, + ) + del onnx_model + gc.collect() + + @property + def torch_to_onnx_input_map(self) -> dict[str, str]: + """Dictionary mapping input names from the PyTorch model to input names from the exported ONNX model. + + Override the function when the input names and the exported ONNX input names are different. + + Returns: + `Dict[str, str]`: A dictionary mapping the PyTorch model input names to the exported ONNX model input names. + """ + return {} + + @property + def torch_to_onnx_output_map(self) -> dict[str, str]: + """Dictionary mapping output names from the PyTorch model to output names from the exported ONNX model. + + Override the function when the output names and the exported ONNX output names are different. + + Returns: + `Dict[str, str]`: A dictionary mapping the PyTorch model output names to the exported ONNX model output names. + """ + return {} + + def rename_ambiguous_inputs(self, inputs) -> dict[str, dict[int, str]]: + """Updates the input names of the model to export. + + Override the function when the model input names are ambiguous or too generic. + + Returns: + `Dict[str, Dict[int, str]]`: Updated inputs. + """ + return inputs + + def ordered_inputs(self, model: PreTrainedModel) -> dict[str, dict[int, str]]: + """Re-orders the inputs using the model forward pass signature. + + Args: + model ([`transformers.PreTrainedModel`]): + The model for which we will use the OnnxConfig. + + Returns: + `Dict[str, Dict[int, str]]`: The properly ordered inputs. + """ + inputs = self.inputs + inputs = self.rename_ambiguous_inputs(inputs) + + ordered_inputs = {} + if hasattr(model, "forward"): + sig = inspect.signature(model.forward) + else: + sig = inspect.signature(model.call) + + for param in sig.parameters: + param_regex = re.compile(rf"{param}(\..*)?$") + to_insert = [] + for name, dynamic_axes in inputs.items(): + if re.match(param_regex, name): + to_insert.append((name, dynamic_axes)) + # TODO: figure out a smart way of re-ordering potential nested structures. + # to_insert = sorted(to_insert, key=lambda t: t[0]) + for name, dynamic_axes in to_insert: + name = self.torch_to_onnx_input_map.get(name, name) + ordered_inputs[name] = dynamic_axes + return ordered_inputs + + @classmethod + def flatten_output_collection_property(cls, name: str, field: Iterable[Any]) -> dict[str, Any]: + """Flattens any potential nested structure expanding the name of the field with the index of the element within the structure. + + Args: + name (`str`): + The name of the nested structure. + field (`Iterable[Any]`): + The structure to potentially flattened. + + Returns: + `Dict[str, Any]`: Outputs with flattened structure and key mapping this new structure. + + """ + if isinstance(field[0], (list, tuple)): + return {f"{name}.{idx}": item for idx, item in enumerate(itertools.chain.from_iterable(field))} + else: + return {f"{name}.{idx}": item for idx, item in enumerate(field)} + + def generate_dummy_inputs_for_validation( + self, reference_model_inputs: dict[str, Any], onnx_input_names: list[str] + ) -> dict[str, Any]: + """Generates inputs for ONNX Runtime using the reference model inputs. + + Override this to run inference with seq2seq + models which have the encoder and decoder exported as separate ONNX files. + + Args: + reference_model_inputs (`Dict[str, Tensor]`): + Reference inputs for the model. + onnx_input_names (`Optional[List[str]]`, defaults to `None`): + Names of the actual inputs to the ONNX model. This argument may be required as an unused + input to the model is automatically removed by torch.onnx.export (e.g. encoder_outputs in the decoder with past) + + Returns: + `Dict[str, Tensor]`: The mapping holding the kwargs to provide to the model's forward function + """ + return reference_model_inputs + + def patch_model_for_export( + self, model: PreTrainedModel, model_kwargs: dict[str, Any] | None = None + ) -> ModelPatcher: + return self._MODEL_PATCHER(self, model, model_kwargs=model_kwargs) + + +class OnnxConfigWithPast(OnnxConfig, ABC): + """Inherits from [`~exporters.onnx.OnnxConfig`]. A base class to handle the ONNX configuration of decoder-only models.""" + + PAD_ATTENTION_MASK_TO_PAST: bool = False + SUPPORTS_PAST: bool = True + + def __init__( + self, + config: PretrainedConfig, + task: str = "feature-extraction", + int_dtype: str = "int64", + float_dtype: str = "fp32", + use_past: bool = False, + use_past_in_inputs: bool = False, + preprocessors: list[Any] | None = None, + ): + self.use_past = use_past + self.use_past_in_inputs = use_past_in_inputs + + self.is_merged = False + self.use_cache_branch = None + super().__init__( + config=config, + task=task, + int_dtype=int_dtype, + float_dtype=float_dtype, + preprocessors=preprocessors, + ) + + @property + def outputs(self) -> dict[str, dict[int, str]]: + if not self.use_past_in_inputs: + common_outputs = super().outputs + # In the other cases, the sequence_length axis is not dynamic, always of length 1 + elif self.task == "feature-extraction": + common_outputs = OrderedDict({"last_hidden_state": {0: "batch_size"}}) + else: + common_outputs = OrderedDict({"logits": {0: "batch_size", 1: "sequence_length"}}) + if self.use_past: + # When exporting decoder models with use_cache=True, both the decoder without past and with past have the KV cache as an output. + self.add_past_key_values(common_outputs, direction="outputs") + return common_outputs + + @property + def values_override(self) -> dict[str, Any] | None: + if hasattr(self._config, "use_cache"): + return {"use_cache": self.use_past} + + @add_dynamic_docstring(text=GENERATE_DUMMY_DOCSTRING, dynamic_elements=DEFAULT_DUMMY_SHAPES) + def generate_dummy_inputs(self, framework: str = "pt", **kwargs): + dummy_inputs_generators = self._create_dummy_input_generator_classes(**kwargs) + + dummy_inputs = {} + input_names = [key for key in self.inputs if not key.startswith("past_key_values")] + if self.use_past_in_inputs and self.use_cache_branch is not False: + input_names.append("past_key_values") + + for input_name in input_names: + input_was_inserted = False + for dummy_input_gen in dummy_inputs_generators: + if dummy_input_gen.supports_input(input_name): + dummy_inputs[input_name] = self.overwrite_shape_and_generate_input( + dummy_input_gen, + input_name, + framework, + input_shapes=kwargs, + ) + input_was_inserted = True + break + if not input_was_inserted: + raise RuntimeError( + f'Could not generate dummy input for "{input_name}". Try adding a proper dummy input generator to the model ONNX config.' + ) + + # refer to https://github.com/huggingface/optimum/pull/764 + if ( + self.use_past_in_inputs + and self.PAD_ATTENTION_MASK_TO_PAST + and self.use_cache_branch is not False + and "attention_mask" in dummy_inputs + and self.task == "text-generation" + ): + seq_len = dummy_inputs["input_ids"].shape[1] + past_seq_len = dummy_inputs["past_key_values"][0][1].shape[-2] + dummy_inputs["attention_mask"] = DummyInputGenerator.pad_input_on_dim( + dummy_inputs["attention_mask"], desired_length=past_seq_len + seq_len, dim=1 + ) + + return dummy_inputs + + def overwrite_shape_and_generate_input( + self, dummy_input_gen: DummyInputGenerator, input_name: str, framework: str, input_shapes: dict + ): + """The shape passed to the dummy input generator may not always be correct for all of the inputs it manages. + + This method allows + to overwrite some shapes, and generate the dummy input. This should probably be refactored more elegantly. + """ + # models from TextSeq2SeqOnnxConfig use decoder_input_ids as input name + # while models from TextDecoderOnnxConfig use input_ids, hence the check for both + + # NOTE: The check `self.task != "text-generation" is added following the use of a single ONNX for both without/with KV cache, without subgraphs. + # This overwrite may be moved to OnnxSeq2SeqConfigWithPast, but I am afraid it would break encoder-decoder models. + if ( + self.use_past + and self.use_past_in_inputs + and self.use_cache_branch is not False + and input_name in ["decoder_input_ids", "input_ids", "position_ids"] + and self.task != "text-generation" + ): + sequence_length = dummy_input_gen.sequence_length + # Use a sequence length of 1 when the KV cache is already populated. + dummy_input_gen.sequence_length = 1 + dummy_input = dummy_input_gen.generate( + input_name, framework=framework, int_dtype=self.int_dtype, float_dtype=self.float_dtype + ) + dummy_input_gen.sequence_length = sequence_length + else: + dummy_input = dummy_input_gen.generate( + input_name, framework=framework, int_dtype=self.int_dtype, float_dtype=self.float_dtype + ) + + return dummy_input + + def add_past_key_values(self, inputs_or_outputs: dict[str, dict[int, str]], direction: str): + """Fills `input_or_outputs` mapping with past_key_values dynamic axes considering the direction. + + Args: + inputs_or_outputs (`Dict[str, Dict[int, str]]`): + The mapping to fill. + direction (`str`): + either "inputs" or "outputs", it specifies whether `input_or_outputs` is the input mapping or the + output mapping, this is important for axes naming. + """ + if direction not in ["inputs", "outputs"]: + raise ValueError(f'direction must either be "inputs" or "outputs", but {direction} was given') + + if direction == "inputs": + decoder_sequence_name = "past_sequence_length" + name = "past_key_values" + else: + decoder_sequence_name = "past_sequence_length + sequence_length" + name = "present" + + for i in range(self._normalized_config.num_layers): + inputs_or_outputs[f"{name}.{i}.key"] = {0: "batch_size", 2: decoder_sequence_name} + inputs_or_outputs[f"{name}.{i}.value"] = {0: "batch_size", 2: decoder_sequence_name} + + def flatten_past_key_values(self, flattened_output, name, idx, t): + flattened_output[f"{name}.{idx}.key"] = t[0] + flattened_output[f"{name}.{idx}.value"] = t[1] + + def flatten_output_collection_property(self, name: str, field: Iterable[Any]) -> dict[str, Any]: + flattened_output = {} + if name in ["present", "past_key_values"]: + for idx, t in enumerate(field): + self.flatten_past_key_values(flattened_output, name, idx, t) + else: + flattened_output = super().flatten_output_collection_property(name, field) + + return flattened_output + + def generate_dummy_inputs_for_validation( + self, reference_model_inputs: dict[str, Any], onnx_input_names: list[str] + ) -> dict[str, Any]: + if self.is_merged is True and self.use_cache_branch is True: + reference_model_inputs["use_cache_branch"] = DummyInputGenerator.constant_tensor(shape=[1], value=True) + elif self.is_merged is True and self.use_cache_branch is False: + reference_model_inputs["use_cache_branch"] = DummyInputGenerator.constant_tensor(shape=[1], value=False) + + # We don't support optional inputs for now, so even though the non-cache branch is used, + # dummy past key values are necessary + batch_size = reference_model_inputs["input_ids"].shape[0] + pkv_generator = self.DUMMY_PKV_GENERATOR_CLASS( + task=self.task, normalized_config=self._normalized_config, sequence_length=1, batch_size=batch_size + ) + reference_model_inputs["past_key_values"] = pkv_generator.generate( + "past_key_values", framework="pt", int_dtype=self.int_dtype, float_dtype=self.float_dtype + ) + + return super().generate_dummy_inputs_for_validation(reference_model_inputs, onnx_input_names) + + +class ConfigBehavior(str, enum.Enum): + """Specifies the behavior of the [`~exporters.onnx.base.OnnxSeq2SeqConfigWithPast`]. + + - MONOLITH: the config can be used to export the whole seq2seq model as a single file. + - ENCODER: the config can be used to export the encoder part of the seq2seq model. + - DECODER: the config can be used to export the decoder part of the seq2seq model. + """ + + MONOLITH = "monolith" + ENCODER = "encoder" + DECODER = "decoder" + + +class OnnxSeq2SeqConfigWithPast(OnnxConfigWithPast): + """Inherits from [`~exporters.onnx.OnnxConfigWithPast`]. A base class to handle the ONNX configuration of encoder-decoder models.""" + + DUMMY_PKV_GENERATOR_CLASS = DummySeq2SeqPastKeyValuesGenerator + + def __init__( + self, + config: PretrainedConfig, + task: str = "feature-extraction", + int_dtype: str = "int64", + float_dtype: str = "fp32", + use_past: bool = False, + use_past_in_inputs: bool = False, + behavior: ConfigBehavior = ConfigBehavior.MONOLITH, + preprocessors: list[Any] | None = None, + ): + super().__init__( + config=config, + task=task, + int_dtype=int_dtype, + float_dtype=float_dtype, + use_past=use_past, + use_past_in_inputs=use_past_in_inputs, + preprocessors=preprocessors, + ) + self._behavior = behavior + + if self._behavior is ConfigBehavior.ENCODER: + self.task = "feature-extraction" + self.use_past_in_inputs = False + + def with_behavior( + self, + behavior: str | ConfigBehavior, + use_past: bool = False, + use_past_in_inputs: bool = False, + ) -> OnnxSeq2SeqConfigWithPast: + """Creates a copy of the current OnnxConfig but with a different `ConfigBehavior` and `use_past` value. + + Args: + behavior ([`ConfigBehavior`]): + The behavior to use for the new instance. + use_past (`bool`, defaults to `False`): + Whether or not the ONNX config to instantiate is for a model using KV cache. + use_past_in_inputs (`bool`, defaults to `False`): + Whether the KV cache is to be passed as an input to the ONNX. + + Returns: + `OnnxSeq2SeqConfigWithPast` + """ + if isinstance(behavior, str) and not isinstance(behavior, ConfigBehavior): + behavior = ConfigBehavior(behavior) + + onnx_config = self.__class__( + self._config, + task=self.task, + int_dtype=self.int_dtype, + float_dtype=self.float_dtype, + use_past=use_past, + use_past_in_inputs=use_past_in_inputs, + behavior=behavior, + preprocessors=self._preprocessors, + ) + onnx_config.variant = self.variant + return onnx_config + + @property + def torch_to_onnx_input_map(self) -> dict[str, str]: + if self._behavior is ConfigBehavior.DECODER: + return { + "decoder_input_ids": "input_ids", + "decoder_attention_mask": "attention_mask", + "encoder_outputs": "encoder_hidden_states", + "attention_mask": "encoder_attention_mask", + } + return {} + + @property + def outputs(self) -> dict[str, dict[int, str]]: + common_outputs = super(OnnxConfigWithPast, self).outputs + # Renaming the outputs axes properly. + for name, axes_names in common_outputs.items(): + if self._behavior is ConfigBehavior.ENCODER or "encoder" in name: + sequence_name = "encoder_sequence_length" + else: + sequence_name = "decoder_sequence_length" + + new_axes_names = {} + for axis_idx, axis_name in axes_names.items(): + if "sequence" in axis_name: + if self.use_past_in_inputs is False or self.is_merged is True: + new_axes_names[axis_idx] = sequence_name + else: + # Trick to force it since ONNX sometimes infer a dynamic axis where it's not. + new_axes_names[axis_idx] = "1" + else: + new_axes_names[axis_idx] = axis_name + common_outputs[name] = new_axes_names + + if self.use_past: + # When exporting decoder models with use_cache=True, both the decoder without past and with past have the KV cache as an output. + self.add_past_key_values(common_outputs, direction="outputs") + + return common_outputs + + def add_past_key_values(self, inputs_or_outputs: dict[str, dict[int, str]], direction: str): + if direction not in ["inputs", "outputs"]: + raise ValueError(f'direction must either be "inputs" or "outputs", but {direction} was given') + + if direction == "inputs": + decoder_sequence_name = "past_decoder_sequence_length" + name = "past_key_values" + else: + decoder_sequence_name = "past_decoder_sequence_length + decoder_sequence_length" + name = "present" + + for i in range(self._normalized_config.decoder_num_layers): + inputs_or_outputs[f"{name}.{i}.decoder.key"] = {0: "batch_size", 2: decoder_sequence_name} + inputs_or_outputs[f"{name}.{i}.decoder.value"] = {0: "batch_size", 2: decoder_sequence_name} + + if ( + self.is_merged is True + or (self._behavior is ConfigBehavior.DECODER and not self.use_past_in_inputs) + or direction == "inputs" + ): + inputs_or_outputs[f"{name}.{i}.encoder.key"] = {0: "batch_size", 2: "encoder_sequence_length"} + inputs_or_outputs[f"{name}.{i}.encoder.value"] = {0: "batch_size", 2: "encoder_sequence_length"} + + def flatten_past_key_values(self, flattened_output, name, idx, t): + if len(t) not in [2, 4]: + raise ValueError( + "past_key_values to flatten should be of length 2 (self-attention only) or 4 (self and cross attention)." + ) + + flattened_output[f"{name}.{idx}.decoder.key"] = t[0] + flattened_output[f"{name}.{idx}.decoder.value"] = t[1] + if len(t) == 4: + flattened_output[f"{name}.{idx}.encoder.key"] = t[2] + flattened_output[f"{name}.{idx}.encoder.value"] = t[3] + + def generate_dummy_inputs(self, framework: str = "pt", **kwargs): + dummy_inputs = super().generate_dummy_inputs(framework=framework, **kwargs) + + if ( + self.use_past_in_inputs + and self.PAD_ATTENTION_MASK_TO_PAST + and self.use_cache_branch is not False + and "decoder_attention_mask" in dummy_inputs + ): + seq_len = dummy_inputs["decoder_input_ids"].shape[1] + past_seq_len = dummy_inputs["past_key_values"][0][1].shape[-2] + dummy_inputs["decoder_attention_mask"] = DummyInputGenerator.pad_input_on_dim( + dummy_inputs["decoder_attention_mask"], desired_length=past_seq_len + seq_len, dim=1 + ) + + return dummy_inputs + + def generate_dummy_inputs_for_validation( + self, reference_model_inputs: dict[str, Any], onnx_input_names: list[str] + ) -> dict[str, Any]: + if self._behavior is ConfigBehavior.DECODER: + if "decoder_input_ids" in reference_model_inputs: + reference_model_inputs["input_ids"] = reference_model_inputs.pop("decoder_input_ids") + if "attention_mask" in reference_model_inputs: + reference_model_inputs["encoder_attention_mask"] = reference_model_inputs.pop("attention_mask") + if "decoder_attention_mask" in reference_model_inputs: + reference_model_inputs["attention_mask"] = reference_model_inputs.pop("decoder_attention_mask") + if "encoder_outputs" in reference_model_inputs: + if "encoder_hidden_states" in onnx_input_names: + reference_model_inputs["encoder_hidden_states"] = reference_model_inputs.pop("encoder_outputs")[0] + else: + reference_model_inputs.pop("encoder_outputs") + + return super().generate_dummy_inputs_for_validation(reference_model_inputs, onnx_input_names) diff --git a/optimum/exporters/openvino/config.py b/optimum/exporters/openvino/config.py new file mode 100644 index 0000000000..f53230dd57 --- /dev/null +++ b/optimum/exporters/openvino/config.py @@ -0,0 +1,347 @@ +# Copyright 2022 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Common ONNX configuration classes that handle most of the features for building model specific configurations.""" + +from __future__ import annotations + +from collections import OrderedDict +from collections.abc import Iterable +from typing import Any + +from transformers import PretrainedConfig + +from optimum.exporters.openvino.base import ConfigBehavior, OnnxConfig, OnnxConfigWithPast, OnnxSeq2SeqConfigWithPast +from optimum.exporters.tasks import TasksManager +from optimum.utils import ( + DummyAudioInputGenerator, + DummyBboxInputGenerator, + DummyInputGenerator, + DummyPastKeyValuesGenerator, + DummySeq2SeqDecoderTextInputGenerator, + DummySeq2SeqPastKeyValuesGenerator, + DummyTextInputGenerator, + DummyVisionInputGenerator, + logging, +) + + +logger = logging.get_logger(__name__) + + +class TextEncoderOnnxConfig(OnnxConfig): + """Handles encoder-based text architectures.""" + + DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator,) + + +class TextDecoderOnnxConfig(OnnxConfigWithPast): + """Handles decoder-based text architectures.""" + + PAD_ATTENTION_MASK_TO_PAST = True + DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, DummyPastKeyValuesGenerator) + DUMMY_PKV_GENERATOR_CLASS = DummyPastKeyValuesGenerator + + def __init__( + self, + config: PretrainedConfig, + task: str = "feature-extraction", + int_dtype: str = "int64", + float_dtype: str = "fp32", + use_past: bool = False, + use_past_in_inputs: bool = False, + preprocessors: list[Any] | None = None, + ): + super().__init__( + config=config, + task=task, + int_dtype=int_dtype, + float_dtype=float_dtype, + use_past=use_past, + use_past_in_inputs=use_past_in_inputs, + preprocessors=preprocessors, + ) + + @property + def inputs(self) -> dict[str, dict[int, str]]: + if self.use_past_in_inputs: + common_inputs = {"input_ids": {0: "batch_size", 1: "sequence_length"}} + common_inputs["attention_mask"] = {0: "batch_size", 1: "past_sequence_length + sequence_length"} + self.add_past_key_values(common_inputs, direction="inputs") + else: + common_inputs = { + "input_ids": {0: "batch_size", 1: "sequence_length"}, + "attention_mask": {0: "batch_size", 1: "sequence_length"}, + } + return common_inputs + + @property + def outputs(self) -> dict[str, dict[int, str]]: + if self.is_merged is False: + common_outputs = super().outputs + else: + # in the merged case, we need to allow the `sequence_length` to be variable, as it is not 1 + # during the first pass without past key values + common_outputs = OrderedDict({"logits": {0: "batch_size", 1: "sequence_length"}}) + self.add_past_key_values(common_outputs, direction="outputs") + return common_outputs + + +class TextDecoderWithPositionIdsOnnxConfig(TextDecoderOnnxConfig): + @property + def inputs(self) -> dict[str, dict[int, str]]: + common_inputs = super().inputs + + # Decoders based on GPT2 require a position_ids input to avoid generating wrong position_ids in the model itself: + # https://github.com/huggingface/transformers/blob/v4.33.1/src/transformers/models/gpt2/modeling_gpt2.py#L802 + if self.task in {"text-generation", "feature-extraction"}: + common_inputs["position_ids"] = {0: "batch_size", 1: "sequence_length"} + + return common_inputs + + +class TextSeq2SeqOnnxConfig(OnnxSeq2SeqConfigWithPast): + """Handles encoder-decoder-based text architectures.""" + + DUMMY_INPUT_GENERATOR_CLASSES = ( + DummyTextInputGenerator, + DummySeq2SeqDecoderTextInputGenerator, + DummySeq2SeqPastKeyValuesGenerator, + ) + + @property + def inputs(self) -> dict[str, dict[int, str]]: + common_inputs = {} + if self._behavior in {ConfigBehavior.ENCODER, ConfigBehavior.MONOLITH}: + common_inputs["input_ids"] = {0: "batch_size", 1: "encoder_sequence_length"} + else: + common_inputs["encoder_outputs"] = {0: "batch_size", 1: "encoder_sequence_length"} + common_inputs["attention_mask"] = {0: "batch_size", 1: "encoder_sequence_length"} + + if self._behavior in {ConfigBehavior.DECODER, ConfigBehavior.MONOLITH}: + common_inputs["decoder_input_ids"] = {0: "batch_size", 1: "decoder_sequence_length"} + if self.use_past_in_inputs: + self.add_past_key_values(common_inputs, direction="inputs") + + return common_inputs + + def _create_dummy_input_generator_classes(self, **kwargs) -> list[DummyInputGenerator]: + dummy_text_input_generator = self.DUMMY_INPUT_GENERATOR_CLASSES[0]( + self.task, self._normalized_config, **kwargs + ) + dummy_decoder_text_input_generator = self.DUMMY_INPUT_GENERATOR_CLASSES[1]( + self.task, + self._normalized_config, + **kwargs, + ) + dummy_seq2seq_past_key_values_generator = self.DUMMY_INPUT_GENERATOR_CLASSES[2]( + self.task, + self._normalized_config, + encoder_sequence_length=dummy_text_input_generator.sequence_length, + **kwargs, + ) + dummy_inputs_generators = [ + dummy_text_input_generator, + dummy_decoder_text_input_generator, + dummy_seq2seq_past_key_values_generator, + ] + + return dummy_inputs_generators + + +class VisionOnnxConfig(OnnxConfig): + """Handles vision architectures.""" + + DUMMY_INPUT_GENERATOR_CLASSES = (DummyVisionInputGenerator,) + + +class TextAndVisionOnnxConfig(OnnxConfig): + """Handles multi-modal text and vision architectures.""" + + DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, DummyVisionInputGenerator, DummyBboxInputGenerator) + + +class AudioOnnxConfig(OnnxConfig): + """Handles audio architectures.""" + + DUMMY_INPUT_GENERATOR_CLASSES = (DummyAudioInputGenerator,) + + @property + def inputs(self) -> dict[str, dict[int, str]]: + return {"input_values": {0: "batch_size", 1: "sequence_length"}} + + +class AudioToTextOnnxConfig(OnnxSeq2SeqConfigWithPast): + DUMMY_INPUT_GENERATOR_CLASSES = ( + DummyAudioInputGenerator, + DummySeq2SeqDecoderTextInputGenerator, + DummySeq2SeqPastKeyValuesGenerator, + ) + + @property + def inputs(self) -> dict[str, dict[int, str]]: + common_inputs = {} + + if self._behavior in {ConfigBehavior.ENCODER, ConfigBehavior.MONOLITH}: + common_inputs["input_features"] = {0: "batch_size", 1: "feature_size", 2: "encoder_sequence_length"} + else: + common_inputs["encoder_outputs"] = {0: "batch_size", 1: "encoder_sequence_length"} + + if self._behavior in {ConfigBehavior.DECODER, ConfigBehavior.MONOLITH}: + common_inputs["decoder_input_ids"] = {0: "batch_size", 1: "decoder_sequence_length"} + if self.use_past_in_inputs: + self.add_past_key_values(common_inputs, direction="inputs") + + return common_inputs + + +class EncoderDecoderBaseOnnxConfig(OnnxSeq2SeqConfigWithPast): + DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator,) + + def __init__( + self, + config: PretrainedConfig, + task: str = "feature-extraction", + int_dtype: str = "int64", + float_dtype: str = "fp32", + use_past: bool = False, + use_past_in_inputs: bool = False, + behavior: ConfigBehavior = ConfigBehavior.MONOLITH, + preprocessors: list[Any] | None = None, + ): + super().__init__( + config=config, + task=task, + int_dtype=int_dtype, + float_dtype=float_dtype, + use_past=use_past, + use_past_in_inputs=use_past_in_inputs, + behavior=behavior, + preprocessors=preprocessors, + ) + + self.is_decoder_with_past = False + + # Set up the encoder ONNX config. + encoder_onnx_config_constructor = TasksManager.get_exporter_config_constructor( + exporter="onnx", + task="feature-extraction", + model_type=config.encoder.model_type, + library_name="transformers", + ) + self._encoder_onnx_config = encoder_onnx_config_constructor( + config.encoder, int_dtype=int_dtype, float_dtype=float_dtype, preprocessors=preprocessors + ) + self._normalized_config.ENCODER_NORMALIZED_CONFIG_CLASS = self._encoder_onnx_config._normalized_config + + # Set up the decoder ONNX config. + decoder_onnx_config_constructor = TasksManager.get_exporter_config_constructor( + exporter="onnx", + task="feature-extraction", + model_type=config.decoder.model_type, + library_name="transformers", + ) + kwargs = {} + if issubclass(decoder_onnx_config_constructor.func, OnnxConfigWithPast): + self.is_decoder_with_past = True + kwargs["use_past"] = use_past + else: + self.use_past = False + + if use_past and not self.is_decoder_with_past: + raise ValueError( + f"The decoder part of the encoder-decoder model is {config.decoder.model_type} which does not need " + "past key values." + ) + + self._decoder_onnx_config = decoder_onnx_config_constructor( + config.decoder, int_dtype=int_dtype, float_dtype=float_dtype, preprocessors=preprocessors, **kwargs + ) + if issubclass(decoder_onnx_config_constructor.func, OnnxSeq2SeqConfigWithPast): + self._decoder_onnx_config = self._decoder_onnx_config.with_behavior( + self._behavior, use_past=kwargs["use_past"], use_past_in_inputs=use_past_in_inputs + ) + + self._normalized_config.DECODER_NORMALIZED_CONFIG_CLASS = self._decoder_onnx_config._normalized_config + self._normalized_config.DECODER_NORMALIZED_CONFIG_CLASS = self._decoder_onnx_config._normalized_config + self._normalized_config.DECODER_NORMALIZED_CONFIG_CLASS.encoder_num_attention_heads = ( + self._decoder_onnx_config._normalized_config.num_attention_heads + ) + self._normalized_config.DECODER_NORMALIZED_CONFIG_CLASS.decoder_num_attention_heads = ( + self._decoder_onnx_config._normalized_config.num_attention_heads + ) + + if isinstance(self._decoder_onnx_config, OnnxSeq2SeqConfigWithPast): + self._past_key_values_generator = ( + DummySeq2SeqDecoderTextInputGenerator, + DummySeq2SeqPastKeyValuesGenerator, + ) + else: + self._past_key_values_generator = ( + DummySeq2SeqDecoderTextInputGenerator, + DummyPastKeyValuesGenerator, + ) + + self.DUMMY_INPUT_GENERATOR_CLASSES += self._past_key_values_generator + + @property + def inputs(self) -> dict[str, dict[int, str]]: + common_inputs = {} + if self._behavior in {ConfigBehavior.ENCODER, ConfigBehavior.MONOLITH}: + common_inputs["input_ids"] = {0: "batch_size", 1: "encoder_sequence_length"} + else: + common_inputs["encoder_outputs"] = {0: "batch_size", 1: "encoder_sequence_length"} + common_inputs["attention_mask"] = {0: "batch_size", 1: "encoder_sequence_length"} + + if self._behavior in {ConfigBehavior.DECODER, ConfigBehavior.MONOLITH}: + common_inputs["decoder_input_ids"] = {0: "batch_size", 1: "decoder_sequence_length"} + if self.use_past_in_inputs: + self.add_past_key_values(common_inputs, direction="inputs") + + return common_inputs + + def add_past_key_values(self, inputs_or_outputs: dict[str, dict[int, str]], direction: str): + if self.is_decoder_with_past: + return self._decoder_onnx_config.add_past_key_values(inputs_or_outputs, direction) + + def flatten_past_key_values(self, flattened_output, name, idx, t): + if self.is_decoder_with_past: + return self._decoder_onnx_config.flatten_past_key_values(flattened_output, name, idx, t) + + def flatten_output_collection_property(self, name: str, field: Iterable[Any]) -> dict[str, Any]: + return self._decoder_onnx_config.flatten_output_collection_property(name, field) + + def generate_dummy_inputs_for_validation( + self, reference_model_inputs: dict[str, Any], onnx_input_names: list[str] + ) -> dict[str, Any]: + if self._behavior is ConfigBehavior.ENCODER: + return self._encoder_onnx_config.generate_dummy_inputs_for_validation( + reference_model_inputs, onnx_input_names + ) + else: + if self._behavior is ConfigBehavior.DECODER: + if "decoder_input_ids" in reference_model_inputs: + reference_model_inputs["input_ids"] = reference_model_inputs.pop("decoder_input_ids") + if "attention_mask" in reference_model_inputs: + reference_model_inputs["encoder_attention_mask"] = reference_model_inputs.pop("attention_mask") + if "encoder_outputs" in reference_model_inputs: + if "encoder_hidden_states" in onnx_input_names: + reference_model_inputs["encoder_hidden_states"] = reference_model_inputs.pop( + "encoder_outputs" + )[0] + else: + reference_model_inputs.pop("encoder_outputs") + + return self._decoder_onnx_config.generate_dummy_inputs_for_validation( + reference_model_inputs, onnx_input_names + ) diff --git a/optimum/exporters/openvino/input_generators.py b/optimum/exporters/openvino/input_generators.py new file mode 100644 index 0000000000..7323acedad --- /dev/null +++ b/optimum/exporters/openvino/input_generators.py @@ -0,0 +1,1816 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Optional, Tuple + +import torch + +from optimum.utils import ( + DEFAULT_DUMMY_SHAPES, + DummyAudioInputGenerator, + DummyInputGenerator, + DummyPastKeyValuesGenerator, + DummySeq2SeqDecoderTextInputGenerator, + DummySeq2SeqPastKeyValuesGenerator, + DummyTextInputGenerator, + DummyTimestepInputGenerator, + DummyTransformerTextInputGenerator, + DummyVisionInputGenerator, + FalconDummyPastKeyValuesGenerator, + MistralDummyPastKeyValuesGenerator, + NormalizedTextConfig, + is_transformers_version, +) +from optimum.utils.input_generators import DTYPE_MAPPER +from optimum.utils.normalized_config import NormalizedConfig, NormalizedVisionConfig + +from ...intel.utils.import_utils import is_diffusers_version + + +class GPTBigCodeDummyPastKeyValuesGenerator(DummyPastKeyValuesGenerator): + def __init__(self, task: str, normalized_config: NormalizedTextConfig, **kwargs): + super().__init__(task=task, normalized_config=normalized_config, **kwargs) + self.multi_query = normalized_config.multi_query + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + if is_transformers_version("<", "4.54"): + if self.multi_query: + shape = ( + self.batch_size, + self.sequence_length, + self.hidden_size // self.num_attention_heads * 2, + ) + else: + shape = ( + self.batch_size, + self.num_attention_heads, + self.sequence_length, + self.hidden_size // self.num_attention_heads * 2, + ) + pkv = [ + self.random_float_tensor(shape, framework=framework, dtype=float_dtype) for _ in range(self.num_layers) + ] + + else: + if self.multi_query: + shape = ( + self.batch_size, + 1, + self.sequence_length, + self.hidden_size // self.num_attention_heads, + ) + else: + shape = ( + self.batch_size, + self.num_attention_heads, + self.sequence_length, + self.hidden_size // self.num_attention_heads, + ) + pkv = [ + ( + 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) + ] + + return pkv + + +class DummyMoonshineAudioInputGenerator(DummyAudioInputGenerator): + SUPPORTED_INPUT_NAMES = ("input_values", "attention_mask") + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + if input_name == "input_values": # raw waveform + return self.random_float_tensor( + shape=[self.batch_size, self.sequence_length], + min_value=-1, + max_value=1, + framework=framework, + dtype=float_dtype, + ) + elif input_name == "attention_mask": # attention mask + return self.random_mask_tensor( + shape=[self.batch_size, self.sequence_length], + framework=framework, + dtype=int_dtype, + ) + else: + raise ValueError(f"Unsupported input name: {input_name}") + + +class DummySanaTransforemerTextInputGenerator(DummyTransformerTextInputGenerator): + SUPPORTED_INPUT_NAMES = ("encoder_hidden_states", "encoder_attention_mask") + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + if input_name == "encoder_attention_mask": + return self.random_mask_tensor( + shape=[self.batch_size, self.sequence_length], + framework=framework, + dtype=int_dtype, + ) + else: + return super().generate( + input_name=input_name, framework=framework, int_dtype=int_dtype, float_dtype=float_dtype + ) + + +class DummyQwen3VLLMInputGenerator(DummyTextInputGenerator): + SUPPORTED_INPUT_NAMES = ( + "input_ids", + "attention_mask", + "token_type_ids", + "position_ids", + "visual_pos_masks", + "deepstack_visual_embeds", + ) + + def __init__( + self, + task: str, + normalized_config: NormalizedTextConfig, + batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], + sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], + num_choices: int = DEFAULT_DUMMY_SHAPES["num_choices"], + random_batch_size_range: Optional[Tuple[int, int]] = None, + random_sequence_length_range: Optional[Tuple[int, int]] = None, + random_num_choices_range: Optional[Tuple[int, int]] = None, + padding_side: str = "right", + **kwargs, + ): + super().__init__( + task=task, + normalized_config=normalized_config, + batch_size=batch_size, + sequence_length=sequence_length, + num_choices=num_choices, + random_batch_size_range=random_batch_size_range, + random_sequence_length_range=random_sequence_length_range, + random_num_choices_range=random_num_choices_range, + padding_side=padding_side, + **kwargs, + ) + self.embed_dim = normalized_config.hidden_size + self.num_layers = len(self.normalized_config.deepstack_visual_indexes) + + def generate( + self, + input_name: str, + framework: str = "pt", + int_dtype: str = "int64", + float_dtype: str = "fp32", + bool_dtype: str = "bool", + ): + if input_name == "deepstack_visual_embeds": + return self.random_float_tensor( + [self.num_layers, 2 * self.sequence_length, self.embed_dim], framework=framework, dtype=float_dtype + ) + if input_name == "visual_pos_masks": + return self.constant_tensor( + shape=[self.batch_size, self.sequence_length], + framework=framework, + value=1, + dtype=DTYPE_MAPPER.pt(bool_dtype), + ) + return super().generate(input_name, framework, int_dtype, float_dtype) + + +class OVMiniCPM3DummyPastKeyValuesGenerator(MistralDummyPastKeyValuesGenerator): + def __init__( + self, + task: str, + normalized_config: NormalizedTextConfig, + batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], + sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], + random_batch_size_range: Optional[Tuple[int, int]] = None, + random_sequence_length_range: Optional[Tuple[int, int]] = None, + **kwargs, + ): + super().__init__( + task=task, + normalized_config=normalized_config, + batch_size=batch_size, + sequence_length=sequence_length, + random_batch_size_range=random_batch_size_range, + random_sequence_length_range=random_sequence_length_range, + **kwargs, + ) + self.v_head_dim = getattr(normalized_config, "v_head_dim", self.hidden_size // self.num_attention_heads) + self.k_head_dim = normalized_config.qk_nope_head_dim + normalized_config.qk_rope_head_dim + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + v_shape = ( + self.batch_size, + self.num_key_value_heads, + self.sequence_length, + self.v_head_dim, + ) + k_shape = (self.batch_size, self.num_key_value_heads, self.sequence_length, self.k_head_dim) + return [ + ( + self.random_float_tensor(k_shape, framework=framework, dtype=float_dtype), + self.random_float_tensor(v_shape, framework=framework, dtype=float_dtype), + ) + for _ in range(self.num_layers) + ] + + +class ChatGLM2DummyPastKeyValuesGenerator(DummyPastKeyValuesGenerator): + def __init__( + self, + task: str, + normalized_config: NormalizedTextConfig, + batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], + sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], + random_batch_size_range: Optional[Tuple[int, int]] = None, + random_sequence_length_range: Optional[Tuple[int, int]] = None, + **kwargs, + ): + super().__init__( + task=task, + normalized_config=normalized_config, + batch_size=batch_size, + sequence_length=sequence_length, + random_batch_size_range=random_batch_size_range, + random_sequence_length_range=random_sequence_length_range, + ) + self.multi_query_group_num = normalized_config.multi_query_group_num + self.head_dim = normalized_config.kv_channels + self.standart_cache_layout = hasattr(normalized_config, "rope_ratio") + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + if not self.standart_cache_layout: + pkv_shape = ( + self.sequence_length, + self.batch_size, + self.multi_query_group_num, + self.head_dim, + ) + else: + pkv_shape = ( + self.batch_size, + self.multi_query_group_num, + self.sequence_length, + self.head_dim, + ) + return [ + ( + self.random_float_tensor(pkv_shape, framework=framework, dtype=float_dtype), + self.random_float_tensor(pkv_shape, framework=framework, dtype=float_dtype), + ) + for _ in range(self.num_layers) + ] + + +class Eagle3DummyGenerator(DummyInputGenerator): + """ + Dummy input generator for Eagle-3 speculative decoding. + + This generator produces synthetic `hidden_states` tensors that mimic the + intermediate hidden-state outputs of a *main (target) model*, which are + required by the Eagle-3 draft model during speculative decoding. + """ + + SUPPORTED_INPUT_NAMES = ("hidden_states",) + + def __init__( + self, + task: str, + normalized_config: NormalizedTextConfig, + batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], + sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], + **kwargs, + ): + self.batch_size = batch_size + self.sequence_length = sequence_length + self.hidden_size = normalized_config.hidden_size + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + # hidden_states is provided as a concatenation of three hidden-layer outputs from the main model + shape = ( + self.batch_size, + self.sequence_length, + self.hidden_size * 3, + ) + return self.random_float_tensor(shape, framework=framework, dtype=float_dtype) + + +class Eagle3VLMDummyGenerator(DummyInputGenerator): + """ + Dummy input generator for VLM Eagle-3 speculative decoding. + + Produces `inputs_embeds` (float) and 3D `position_ids` (MRoPE) + required by VLM Eagle-3 draft models targeting Qwen3-VL. + """ + + SUPPORTED_INPUT_NAMES = ("inputs_embeds", "position_ids") + + def __init__( + self, + task: str, + normalized_config: NormalizedTextConfig, + batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], + sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], + **kwargs, + ): + self.batch_size = batch_size + self.sequence_length = sequence_length + self.hidden_size = normalized_config.hidden_size + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + if input_name == "inputs_embeds": + shape = (self.batch_size, self.sequence_length, self.hidden_size) + return self.random_float_tensor(shape, framework=framework, dtype=float_dtype) + if input_name == "position_ids": + # The rotary embedding is MRoPE (Multimodal RoPE) + # MRoPE encodes position along three independent axes: temporal, height, and width + # https://github.com/Tencent/AngelSlim/blob/main/angelslim/compressor/speculative/train/models/draft/llama_eagle3.py#L211 + shape = (3, self.batch_size, self.sequence_length) + return self.random_int_tensor(shape, max_value=self.sequence_length, framework=framework, dtype=int_dtype) + + +class QwenDummyPastKeyValuesGenerator(DummyPastKeyValuesGenerator): + def __init__( + self, + task: str, + normalized_config: NormalizedTextConfig, + batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], + sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], + random_batch_size_range: Optional[Tuple[int, int]] = None, + random_sequence_length_range: Optional[Tuple[int, int]] = None, + **kwargs, + ): + super().__init__( + task=task, + normalized_config=normalized_config, + batch_size=batch_size, + sequence_length=sequence_length, + random_batch_size_range=random_batch_size_range, + random_sequence_length_range=random_sequence_length_range, + ) + self.kv_channels = normalized_config.kv_channels + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + past_key_shape = (self.batch_size, self.sequence_length, self.num_attention_heads, self.kv_channels) + past_value_shape = (self.batch_size, self.sequence_length, self.num_attention_heads, self.kv_channels) + return [ + ( + self.random_float_tensor(past_key_shape, framework=framework, dtype=float_dtype), + self.random_float_tensor(past_value_shape, framework=framework, dtype=float_dtype), + ) + for _ in range(self.num_layers) + ] + + +class OVFalconDummyPastKeyValuesGenerator(FalconDummyPastKeyValuesGenerator): + def __init__( + self, + task: str, + normalized_config: NormalizedTextConfig, + batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], + sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], + random_batch_size_range: Optional[Tuple[int, int]] = None, + random_sequence_length_range: Optional[Tuple[int, int]] = None, + **kwargs, + ): + super().__init__( + task=task, + normalized_config=normalized_config, + batch_size=batch_size, + sequence_length=sequence_length, + random_batch_size_range=random_batch_size_range, + random_sequence_length_range=random_sequence_length_range, + **kwargs, + ) + if normalized_config.new_decoder_architecture: + self.num_kv_heads = normalized_config.num_attention_heads + else: + self.num_kv_heads = normalized_config.num_kv_heads if not normalized_config.multi_query else 1 + + self.head_dim = self.hidden_size // self.num_attention_heads + + +class AquilaDummyPastKeyValuesGenerator(DummyPastKeyValuesGenerator): + def __init__( + self, + task: str, + normalized_config: NormalizedTextConfig, + batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], + sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], + random_batch_size_range: Optional[Tuple[int, int]] = None, + random_sequence_length_range: Optional[Tuple[int, int]] = None, + **kwargs, + ): + super().__init__( + task, + normalized_config, + batch_size, + sequence_length, + random_batch_size_range, + random_sequence_length_range, + **kwargs, + ) + self.num_key_value_heads = getattr( + normalized_config, "num_key_value_heads", normalized_config.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.hidden_size // self.num_attention_heads, + ) + 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) + ] + + +class OVMistralDummyPastKeyValuesGenerator(MistralDummyPastKeyValuesGenerator): + def __init__( + self, + task: str, + normalized_config: NormalizedTextConfig, + batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], + sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], + random_batch_size_range: Optional[Tuple[int, int]] = None, + random_sequence_length_range: Optional[Tuple[int, int]] = None, + **kwargs, + ): + super().__init__( + task=task, + normalized_config=normalized_config, + batch_size=batch_size, + sequence_length=sequence_length, + random_batch_size_range=random_batch_size_range, + random_sequence_length_range=random_sequence_length_range, + **kwargs, + ) + self.head_dim = getattr(normalized_config, "head_dim", 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) + ] + + +class Gemma4DummyPastKeyValuesGenerator(DummyPastKeyValuesGenerator): + def __init__( + self, + task: str, + normalized_config: NormalizedTextConfig, + batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], + sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], + random_batch_size_range: Optional[Tuple[int, int]] = None, + random_sequence_length_range: Optional[Tuple[int, int]] = None, + **kwargs, + ): + super().__init__( + task=task, + normalized_config=normalized_config, + batch_size=batch_size, + sequence_length=sequence_length, + random_batch_size_range=random_batch_size_range, + random_sequence_length_range=random_sequence_length_range, + ) + self.num_key_value_heads = normalized_config.num_key_value_heads + self.head_dim = normalized_config.head_dim + self.global_head_dim = getattr(normalized_config.config, "global_head_dim", self.head_dim) + self.layer_types = normalized_config.config.layer_types + self.num_kv_shared_layers = normalized_config.config.num_kv_shared_layers + self.sliding_window = normalized_config.config.sliding_window + # Full-attention layers use fewer KV heads than sliding-attention layers (e.g. 2 vs 8 for 26B-A4B) + self.num_global_key_value_heads = ( + getattr(normalized_config.config, "num_global_key_value_heads", None) or self.num_key_value_heads + ) + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + # some layers do not produce their own KV-cache, they use the shared KV-cache + if self.num_kv_shared_layers > 0: + layer_types = self.layer_types[: -self.num_kv_shared_layers] + else: + layer_types = self.layer_types + past_kv_values = [] + for layer_type in layer_types: + if layer_type == "sliding_attention": + shape = ( + self.batch_size, + self.num_key_value_heads, + self.sliding_window, + self.head_dim, + ) + else: + shape = ( + self.batch_size, + self.num_global_key_value_heads, + self.sequence_length, + self.global_head_dim, + ) + past_kv_value = ( + self.random_float_tensor(shape, framework=framework, dtype=float_dtype), + self.random_float_tensor(shape, framework=framework, dtype=float_dtype), + ) + past_kv_values.append(past_kv_value) + + return past_kv_values + + +class DeciDummyPastKeyValuesGenerator(DummyPastKeyValuesGenerator): + def __init__( + self, + task: str, + normalized_config: NormalizedTextConfig, + batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], + sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], + random_batch_size_range: Optional[Tuple[int, int]] = None, + random_sequence_length_range: Optional[Tuple[int, int]] = None, + **kwargs, + ): + super().__init__( + task=task, + normalized_config=normalized_config, + batch_size=batch_size, + sequence_length=sequence_length, + random_batch_size_range=random_batch_size_range, + random_sequence_length_range=random_sequence_length_range, + ) + self.num_key_value_heads_per_layer = normalized_config.num_key_value_heads_per_layer + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + past_key_values = [] + + for layer_id in range(self.num_layers): + shape = ( + self.batch_size, + self.num_key_value_heads_per_layer[layer_id], + self.sequence_length, + self.hidden_size // self.num_attention_heads, + ) + past_key_values.append( + ( + self.random_float_tensor(shape, framework=framework, dtype=float_dtype), + self.random_float_tensor(shape, framework=framework, dtype=float_dtype), + ) + ) + return past_key_values + + +class DummyLLavaMultiModalProjectorInputGenerator(DummyInputGenerator): + SUPPORTED_INPUT_NAMES = ["image_features"] + + def __init__( + self, + task: str, + normalized_config: NormalizedTextConfig, + batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], + random_batch_size_range: Optional[Tuple[int, int]] = None, + **kwargs, + ): + self.task = task + + self.batch_size = batch_size + self.hidden_size = normalized_config.hidden_size + self.num_patches = (normalized_config.image_size // normalized_config.patch_size) ** 2 + self.normalized_config = normalized_config + + def generate( + self, + input_name: str, + framework: str = "pt", + int_dtype: str = "int64", + float_dtype: str = "fp32", + ): + shape = [self.batch_size, self.num_patches, self.hidden_size] + return self.random_float_tensor(shape, framework=framework, dtype=float_dtype) + + +class PooledProjectionsDummyInputGenerator(DummyInputGenerator): + SUPPORTED_INPUT_NAMES = ["pooled_projections"] + + def __init__( + self, + task: str, + normalized_config: NormalizedConfig, + batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], + random_batch_size_range: Optional[Tuple[int, int]] = None, + **kwargs, + ): + self.task = task + self.batch_size = batch_size + self.pooled_projection_dim = normalized_config.config.pooled_projection_dim + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + shape = [self.batch_size, self.pooled_projection_dim] + return self.random_float_tensor(shape, framework=framework, dtype=float_dtype) + + +class DummyTransformerTimestpsInputGenerator(DummyTimestepInputGenerator): + SUPPORTED_INPUT_NAMES = ("timestep", "text_embeds", "time_ids", "timestep_cond", "guidance") + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + if input_name in ["timestep", "guidance"]: + shape = [self.batch_size] + return self.random_float_tensor(shape, max_value=self.vocab_size, framework=framework, dtype=float_dtype) + return super().generate(input_name, framework, int_dtype, float_dtype) + + +class DummyUnetVisionInputGenerator(DummyVisionInputGenerator): + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + if input_name not in ["sample", "latent_sample"]: + return super().generate(input_name, framework, int_dtype, float_dtype) + # add height and width discount for enable any resolution generation + return self.random_float_tensor( + shape=[self.batch_size, self.num_channels, self.height - 1, self.width - 1], + framework=framework, + dtype=float_dtype, + ) + + +class DummyUnetTimestepInputGenerator(DummyTimestepInputGenerator): + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + if input_name != "timestep": + return super().generate(input_name, framework, int_dtype, float_dtype) + shape = [self.batch_size] + return self.random_int_tensor(shape, max_value=self.vocab_size, framework=framework, dtype=int_dtype) + + +class DummySanaTimestepInputGenerator(DummyTimestepInputGenerator): + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + if input_name != "timestep": + return super().generate(input_name, framework, int_dtype, float_dtype) + shape = [self.batch_size] + return self.random_int_tensor(shape, max_value=self.vocab_size, framework=framework, dtype=float_dtype) + + +class DummyUnetEncoderInputGenerator(DummySeq2SeqDecoderTextInputGenerator): + def __init__( + self, + task: str, + normalized_config: NormalizedTextConfig, + batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], + sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], + num_choices: int = DEFAULT_DUMMY_SHAPES["num_choices"], + random_batch_size_range: Optional[Tuple[int, int]] = None, + random_sequence_length_range: Optional[Tuple[int, int]] = None, + random_num_choices_range: Optional[Tuple[int, int]] = None, + **kwargs, + ): + super().__init__( + task, + normalized_config, + batch_size=batch_size, + sequence_length=sequence_length, + num_choices=num_choices, + random_batch_size_range=random_batch_size_range, + random_sequence_length_range=random_sequence_length_range, + random_num_choices_range=random_num_choices_range, + **kwargs, + ) + if hasattr(normalized_config.config, "model_max_length"): + self.sequence_length = normalized_config.config.model_max_length + + +class DummySanaSeq2SeqDecoderTextWithEncMaskInputGenerator(DummySeq2SeqDecoderTextInputGenerator): + SUPPORTED_INPUT_NAMES = ( + "decoder_input_ids", + "decoder_attention_mask", + "encoder_outputs", + "encoder_hidden_states", + "encoder_attention_mask", + ) + + +# DummySanaTransformerVisionInputGenerator inherits from DummyUnetVisionInputGenerator (defined above) +class DummySanaTransformerVisionInputGenerator(DummyUnetVisionInputGenerator): + SUPPORTED_INPUT_NAMES = ( + "pixel_values", + "pixel_mask", + "sample", + "latent_sample", + "guidance", + ) + + def __init__( + self, + task: str, + normalized_config: NormalizedVisionConfig, + batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], + num_channels: int = DEFAULT_DUMMY_SHAPES["num_channels"], + width: int = DEFAULT_DUMMY_SHAPES["width"] // 8, + height: int = DEFAULT_DUMMY_SHAPES["height"] // 8, + # Reduce img shape by 4 for FLUX to reduce memory usage on conversion + **kwargs, + ): + super().__init__(task, normalized_config, batch_size, num_channels, width=width, height=height, **kwargs) + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + if input_name == "guidance": + return self.random_float_tensor([self.batch_size], framework=framework, dtype=float_dtype) + return super().generate(input_name, framework, int_dtype, float_dtype) + + +class DummyFluxTransformerInputGenerator(DummyVisionInputGenerator): + SUPPORTED_INPUT_NAMES = ( + "pixel_values", + "pixel_mask", + "sample", + "latent_sample", + "hidden_states", + "img_ids", + ) + + def __init__( + self, + task: str, + normalized_config: NormalizedVisionConfig, + batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], + num_channels: int = DEFAULT_DUMMY_SHAPES["num_channels"], + width: int = DEFAULT_DUMMY_SHAPES["width"] // 4, + height: int = DEFAULT_DUMMY_SHAPES["height"] // 4, + # Reduce img shape by 4 for FLUX to reduce memory usage on conversion + **kwargs, + ): + super().__init__(task, normalized_config, batch_size, num_channels, width, height, **kwargs) + if getattr(normalized_config, "in_channels", None): + self.num_channels = normalized_config.in_channels // 4 + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + if input_name in ["hidden_states", "sample"]: + shape = [self.batch_size, (self.height // 2) * (self.width // 2), self.num_channels * 4] + return self.random_float_tensor(shape, framework=framework, dtype=float_dtype) + if input_name == "img_ids": + img_ids_height = self.height // 2 + img_ids_width = self.width // 2 + return self.random_int_tensor( + ( + [self.batch_size, img_ids_height * img_ids_width, 3] + if is_diffusers_version("<", "0.31.0") + else [img_ids_height * img_ids_width, 3] + ), + min_value=0, + max_value=min(img_ids_height, img_ids_width), + framework=framework, + dtype=float_dtype, + ) + + return super().generate(input_name, framework, int_dtype, float_dtype) + + +class DummyFluxTextInputGenerator(DummySeq2SeqDecoderTextInputGenerator): + SUPPORTED_INPUT_NAMES = ( + "decoder_input_ids", + "decoder_attention_mask", + "encoder_outputs", + "encoder_hidden_states", + "txt_ids", + ) + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + if input_name == "txt_ids": + import torch + + shape = ( + [self.batch_size, self.sequence_length, 3] + if is_diffusers_version("<", "0.31.0") + else [self.sequence_length, 3] + ) + dtype = DTYPE_MAPPER.pt(float_dtype) + return torch.full(shape, 0, dtype=dtype) + return super().generate(input_name, framework, int_dtype, float_dtype) + + +class LTXVaeDummyInputGenerator(DummyVisionInputGenerator): + SUPPORTED_INPUT_NAMES = ("pixel_values", "pixel_mask", "sample", "latent_sample", "timestep") + + def __init__( + self, + task: str, + normalized_config: NormalizedVisionConfig, + batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], + num_channels: int = DEFAULT_DUMMY_SHAPES["num_channels"], + width: int = DEFAULT_DUMMY_SHAPES["width"], + height: int = DEFAULT_DUMMY_SHAPES["height"], + num_frames: int = 2, + **kwargs, + ): + super().__init__(task, normalized_config, batch_size, num_channels, width, height, **kwargs) + self.num_frames = num_frames + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + if input_name in ["sample", "latent_sample"]: + return self.random_float_tensor( + [self.batch_size, self.num_channels, self.num_frames, self.height, self.width] + ) + if input_name == "timestep": + return self.random_int_tensor([1], max_value=20, min_value=1, framework=framework, dtype=int_dtype) + + return super().generate(input_name, framework, int_dtype, float_dtype) + + +class LTXTransformerDummyInputGenerator(DummyVisionInputGenerator): + SUPPORTED_INPUT_NAMES = ("hidden_states", "width", "height", "num_frames", "rope_interpolation_scale") + + def __init__( + self, + task: str, + normalized_config: NormalizedVisionConfig, + batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], + num_channels: int = DEFAULT_DUMMY_SHAPES["num_channels"], + width: int = 16, + height: int = 8, + num_frames: int = 2, + frame_rate: int = 10, + **kwargs, + ): + super().__init__(task, normalized_config, batch_size, num_channels, width, height, **kwargs) + self.num_frames = num_frames + self.frame_rate = frame_rate + self.vae_spatial_compression_ratio = normalized_config.config.vae_spatial_compression_ratio + self.vae_temporal_compression_ratio = normalized_config.config.vae_temporal_compression_ratio + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + import torch + + if input_name == "hidden_states": + return self.random_float_tensor( + [self.batch_size, self.num_frames * self.height * self.width, self.num_channels] + ) + if input_name == "width": + return torch.tensor(self.width) + if input_name == "height": + return torch.tensor(self.height) + if input_name == "num_frames": + return torch.tensor(self.num_frames) + if input_name == "rope_interpolation_scale": + import torch + + return torch.tensor( + [ + self.vae_temporal_compression_ratio / self.frame_rate, + self.vae_spatial_compression_ratio, + self.vae_spatial_compression_ratio, + ] + ) + return super().generate(input_name, framework, int_dtype, float_dtype) + + +class DummyMiniCPMVImageInputGenerator(DummyVisionInputGenerator): + SUPPORTED_INPUT_NAMES = ("pixel_values", "patch_attention_mask", "position_ids") + + def __init__( + self, + task: str, + normalized_config: NormalizedVisionConfig, + batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], + num_channels: int = DEFAULT_DUMMY_SHAPES["num_channels"], + width: int = DEFAULT_DUMMY_SHAPES["width"], + height: int = DEFAULT_DUMMY_SHAPES["height"], + **kwargs, + ): + super().__init__(task, normalized_config, batch_size, num_channels, width, height) + self.patch_size = normalized_config.config.patch_size + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + if input_name == "pixel_values": + return self.random_float_tensor( + shape=[ + self.batch_size, + self.num_channels, + self.patch_size, + (self.height * self.width) // self.patch_size, + ], + framework=framework, + dtype=float_dtype, + ) + + if input_name == "patch_attention_mask": + return self.random_int_tensor( + shape=[self.batch_size, 1, (self.height // self.patch_size) * (self.width // self.patch_size)], + framework=framework, + dtype=float_dtype, + min_value=0, + max_value=2, + ) + + if input_name == "position_ids": + return self.random_int_tensor( + shape=[self.batch_size, (self.height // self.patch_size) * (self.width // self.patch_size)], + max_value=self.patch_size, + ) + + +class DummyMiniCPMVResampleInputGenerator(DummyVisionInputGenerator): + SUPPORTED_INPUT_NAMES = ("image_feature", "pos_embed", "key_padding_mask") + + def __init__( + self, + task: str, + normalized_config: NormalizedVisionConfig, + batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], + num_channels: int = DEFAULT_DUMMY_SHAPES["num_channels"], + width: int = DEFAULT_DUMMY_SHAPES["width"], + height: int = DEFAULT_DUMMY_SHAPES["height"], + **kwargs, + ): + super().__init__(task, normalized_config, batch_size, num_channels, width, height) + self.patch_size = normalized_config.config.patch_size + self.hidden_size = normalized_config.config.hidden_size + self.img_hidden_size = normalized_config.config.vision_config.hidden_size + self.feat_size = (normalized_config.config.vision_config.image_size // self.patch_size) * ( + normalized_config.config.vision_config.image_size // self.patch_size + ) + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + if input_name == "image_feature": + return self.random_float_tensor( + shape=[self.batch_size, self.feat_size, self.img_hidden_size], framework=framework, dtype=float_dtype + ) + + if input_name == "key_padding_mask": + return self.constant_tensor( + shape=[self.batch_size, self.feat_size], + framework=framework, + value=1, + dtype=DTYPE_MAPPER.pt(float_dtype), + ) + + if input_name == "pos_embed": + return self.random_float_tensor(shape=[self.feat_size, self.batch_size, self.hidden_size]) + + +class DummyPhi3VisionProjectionInputGenerator(DummyVisionInputGenerator): + SUPPORTED_INPUT_NAMES = ("input",) + + def __init__( + self, + task: str, + normalized_config: NormalizedVisionConfig, + batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], + num_channels: int = DEFAULT_DUMMY_SHAPES["num_channels"], + width: int = 336, + height: int = 336, + crop_size=336, + **kwargs, + ): + self.batch_size = batch_size + self._embed_layer_realization = ( + normalized_config.config.embd_layer["embedding_cls"] + if hasattr(normalized_config.config, "embd_layer") + else "image_audio" + ) + if not hasattr(normalized_config.config, "vision_config"): + self.image_dim_out = ( + normalized_config.config.img_processor.get( + "image_dim_out", normalized_config.config.img_processor.get("hidden_size") + ) + if normalized_config.config.img_processor is not None + else 1152 + ) + if "image_embd_layer" in normalized_config.config.embd_layer: + self.crop_size = normalized_config.config.embd_layer["image_embd_layer"].get("crop_size", crop_size) + else: + self.crop_size = normalized_config.config.embd_layer.get("crop_size", crop_size) + else: + self.image_dim_out = normalized_config.config.vision_config.hidden_size + self.crop_size = normalized_config.config.vision_config.crop_size + self.height = height + self.width = width + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + h = self.height // self.crop_size + w = self.width // self.crop_size + feat_size = (h * w + 1) * 144 + 1 + (h + 1) * 12 + if self._embed_layer_realization in ["linear", "image_audio"]: + shape = [self.batch_size, feat_size, self.image_dim_out] + else: + shape = [self.batch_size, feat_size, self.image_dim_out * 4] + return self.random_float_tensor(shape, framework=framework, dtype=float_dtype) + + +class DummyAudioPhi4MMInputGenerator(DummyInputGenerator): + SUPPORTED_INPUT_NAMES = ("audio_input", "audio_feature", "audio_mask") + + def __init__( + self, + task: str, + normalized_config: NormalizedVisionConfig, + batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], + signal_length=498, + **kwargs, + ): + self.signal_length = signal_length + if hasattr(normalized_config.config, "audio_processor"): + self.audio_chunk_size = ( + signal_length // normalized_config.config.audio_processor["config"]["time_reduction"] + 1 + ) + self.input_size = normalized_config.config.audio_processor["config"]["input_size"] + self.attention_dim = normalized_config.config.audio_processor["config"]["attention_dim"] + else: + self.audio_chunk_size = signal_length // normalized_config.config.audio_config.time_reduction + 1 + self.input_size = normalized_config.config.audio_config.input_size + self.attention_dim = normalized_config.config.audio_config.hidden_size + self.batch_size = batch_size + self.task = task + self.normalized_config = normalized_config + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + if input_name == "audio_input": + return self.random_float_tensor( + [self.batch_size, self.signal_length, self.input_size], framework=framework, dtype=float_dtype + ) + + if input_name == "audio_feature": + return self.random_float_tensor( + [self.batch_size, self.audio_chunk_size, self.attention_dim], framework=framework, dtype=float_dtype + ) + + if input_name == "audio_mask": + return self.random_int_tensor( + [self.batch_size, self.audio_chunk_size, self.audio_chunk_size], + max_value=2, + framework=framework, + dtype="bool", + ) + + +class DummyVisionPositionIdsPhi4InputGenerator(DummyVisionInputGenerator): + SUPPORTED_INPUT_NAMES = ("patch_position_ids", "patch_attention_mask") + + def __init__( + self, + task: str, + normalized_config: NormalizedVisionConfig, + batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], + num_channels: int = DEFAULT_DUMMY_SHAPES["num_channels"], + width: int = DEFAULT_DUMMY_SHAPES["width"], + height: int = DEFAULT_DUMMY_SHAPES["height"], + **kwargs, + ): + super().__init__(task, normalized_config, batch_size, num_channels, width, height, **kwargs) + if hasattr(normalized_config.config, "vision_conifg"): + self.patch_size = getattr(normalized_config.config.vision_config, "patch_size", 14) + else: + self.patch_size = 14 + self.num_patches_per_side = self.height // self.patch_size + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + if input_name == "patch_position_ids": + return self.get_vision_position_ids() + if input_name == "patch_attention_mask": + return self.random_int_tensor( + [self.batch_size, self.height // self.patch_size, self.width // self.patch_size], + framework=framework, + dtype="bool", + max_value=2, + ) + return super().generate(input_name, framework, int_dtype, float_dtype) + + def get_vision_position_ids(self): + # Adopted from https://github.com/huggingface/transformers/blob/v4.51.3/src/transformers/models/phi4_multimodal/modeling_phi4_multimodal.py#L494-L512 + import torch + + batch_size = self.batch_size + max_im_h, max_im_w = self.height, self.width + max_nb_patches_h, max_nb_patches_w = max_im_h // self.patch_size, max_im_w // self.patch_size + boundaries = torch.arange(1 / self.num_patches_per_side, 1.0, 1 / self.num_patches_per_side) + position_ids = torch.full( + size=( + batch_size, + max_nb_patches_h * max_nb_patches_w, + ), + fill_value=0, + ) + patch_attention_mask = torch.ones( + [self.batch_size, self.height // self.patch_size, self.width // self.patch_size], dtype=torch.int64 + ) + patch_attention_mask[0, self.height - 2 :] = 0 + + for batch_idx, p_attn_mask in enumerate(patch_attention_mask): + nb_patches_h = p_attn_mask[:, 0].sum() + nb_patches_w = p_attn_mask[0].sum() + + fractional_coords_h = torch.arange(0, 1 - 1e-6, 1 / nb_patches_h) + fractional_coords_w = torch.arange(0, 1 - 1e-6, 1 / nb_patches_w) + + bucket_coords_h = torch.bucketize(fractional_coords_h, boundaries, right=True) + bucket_coords_w = torch.bucketize(fractional_coords_w, boundaries, right=True) + + pos_ids = (bucket_coords_h[:, None] * self.num_patches_per_side + bucket_coords_w).flatten() + position_ids[batch_idx][p_attn_mask.view(-1)] = pos_ids + return position_ids + + +class DummyQwen2VLLMInputGenerator(DummyTextInputGenerator): + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + generated_input = super().generate(input_name, framework, int_dtype, float_dtype) + if input_name == "position_ids": + return generated_input.unsqueeze(0).expand(3, -1, -1) + return generated_input + + +class DummyQwen3_5LMInputGenerator(DummyTextInputGenerator): + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + generated_input = super().generate(input_name, framework, int_dtype, float_dtype) + if input_name == "position_ids": + return generated_input.unsqueeze(0).expand(4, -1, -1) + return generated_input + + +class DummyQwen2VLVisionEmbedInputGenerator(DummyVisionInputGenerator): + SUPPORTED_INPUT_NAMES = ( + "hidden_states", + "attention_mask", + "window_attention_mask", + "window_index", + "rotary_pos_emb", + ) + + def __init__( + self, + task: str, + normalized_config: NormalizedVisionConfig, + batch_size: int = 1, + num_channels: int = DEFAULT_DUMMY_SHAPES["num_channels"], + width: int = 420, + height: int = 420, + **kwargs, + ): + self.batch_size = batch_size + self.height = height + self.width = width + self.num_channels = num_channels + self.temporal_patch_size = normalized_config.config.temporal_patch_size + self.patch_size = normalized_config.config.patch_size + if normalized_config.use_embed_dim: + self.embed_dim = ( + normalized_config.config.embed_dim + if hasattr(normalized_config.config, "embed_dim") + else normalized_config.hidden_size + ) + else: + self.embed_dim = self.num_channels * self.temporal_patch_size * self.patch_size * self.patch_size + self.num_heads = normalized_config.config.num_heads + self.spatial_merge_size = None + if hasattr(normalized_config.config, "spatial_merge_size"): + self.spatial_merge_size = normalized_config.config.spatial_merge_size + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + grid_h, grid_w = self.height // self.patch_size, self.width // self.patch_size + grid_t = self.batch_size + + if input_name == "hidden_states": + return self.random_float_tensor( + [grid_t * grid_h * grid_w, self.embed_dim], framework=framework, dtype=float_dtype + ) + + if input_name in ["attention_mask", "window_attention_mask"]: + return self.random_mask_tensor( + [1, grid_t * grid_h * grid_w, grid_t * grid_h * grid_w], framework=framework, dtype=float_dtype + ) + + if input_name == "rotary_pos_emb": + dim = self.embed_dim // self.num_heads // 2 + return self.random_float_tensor([grid_h * grid_t * grid_w, dim], framework=framework, dtype=float_dtype) + + if input_name == "window_index": + if self.spatial_merge_size is None: + raise ValueError( + "`spatial_merge_size` parameter is not found in model config. Can not generate dummy input data for `window_index` input" + ) + spatial_merge_unit = self.spatial_merge_size * self.spatial_merge_size + hidden_size = (grid_t * grid_h * grid_w) // spatial_merge_unit + return self.random_int_tensor([hidden_size], max_value=hidden_size) + + +# DummyQwen3VLVisionEmbedInputGenerator inherits from DummyQwen2VLVisionEmbedInputGenerator (defined above) +class DummyQwen3VLVisionEmbedInputGenerator(DummyQwen2VLVisionEmbedInputGenerator): + SUPPORTED_INPUT_NAMES = ( + "hidden_states", + "attention_mask", + "rotary_pos_emb", + "input", + ) + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + grid_h, grid_w = self.height // self.patch_size, self.width // self.patch_size + grid_t = self.batch_size + + if input_name == "hidden_states": + return self.random_float_tensor( + [grid_t * grid_h * grid_w, self.embed_dim], framework=framework, dtype=float_dtype + ) + + if input_name in ["attention_mask"]: + return self.random_mask_tensor( + [1, grid_t * grid_h * grid_w, grid_t * grid_h * grid_w], framework=framework, dtype=float_dtype + ) + + if input_name == "rotary_pos_emb": + dim = self.embed_dim // self.num_heads // 2 + return self.random_float_tensor([grid_h * grid_t * grid_w, dim], framework=framework, dtype=float_dtype) + + if input_name == "input": + return self.constant_tensor( + [4, DEFAULT_DUMMY_SHAPES["sequence_length"]], + framework=framework, + value=0, + dtype=DTYPE_MAPPER.pt(int_dtype), + ) + + +class Qwen3ASRDummySeq2SeqPastKeyValuesGenerator(DummySeq2SeqPastKeyValuesGenerator): + """Custom KV cache generator for Qwen3-ASR with GQA (num_key_value_heads != num_attention_heads). + Qwen3-ASR has no cross-attention, so only self-attention KV cache is generated (2 per layer).""" + + def __init__(self, task, normalized_config, **kwargs): + super().__init__(task, normalized_config, **kwargs) + # Override head count and head_dim for GQA + self.decoder_num_attention_heads = normalized_config.decoder_num_attention_heads + self.decoder_head_dim = getattr(normalized_config, "head_dim", None) + if self.decoder_head_dim is None: + self.decoder_head_dim = self.decoder_hidden_size // normalized_config.num_attention_heads + + def generate(self, input_name, framework="pt", int_dtype="int64", float_dtype="fp32"): + if input_name == "past_key_values": + decoder_shape = ( + self.batch_size, + self.decoder_num_attention_heads, + self.sequence_length, + self.decoder_head_dim, + ) + # Qwen3-ASR has no cross-attention, only self-attention KV cache + return [ + ( + self.random_float_tensor(decoder_shape, framework=framework, dtype=float_dtype), + self.random_float_tensor(decoder_shape, framework=framework, dtype=float_dtype), + ) + for _ in range(self.decoder_num_layers) + ] + return super().generate(input_name, framework=framework, int_dtype=int_dtype, float_dtype=float_dtype) + + +class DummyGemma4VisionInputGenerator(DummyVisionInputGenerator): + SUPPORTED_INPUT_NAMES = ("pixel_values", "image_position_ids") + + def __init__(self, task, normalized_config, batch_size=DEFAULT_DUMMY_SHAPES["batch_size"], **kwargs): + super().__init__(task, normalized_config, batch_size, **kwargs) + self.patch_size = getattr(normalized_config, "patch_size", 16) + self.pooling_kernel_size = getattr(normalized_config, "pooling_kernel_size", 3) + # Gemma4 processor always pads pixel_values to max_soft_tokens * pooling_kernel_size^2 patches. + # The vision model's pooling uses shape-dependent Python operations that get baked in during tracing, + # so the dummy input must match the actual inference shapes. + max_soft_tokens = getattr(normalized_config, "image_seq_length", None) + if max_soft_tokens is None: + max_soft_tokens = getattr(normalized_config, "max_soft_tokens", 280) + self.num_patches = max_soft_tokens * self.pooling_kernel_size**2 + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + if input_name == "pixel_values": + # Gemma4 expects pre-patchified pixel_values: [batch, num_patches, 3 * patch_size^2] + return self.random_float_tensor( + shape=[self.batch_size, self.num_patches, 3 * self.patch_size**2], + framework=framework, + dtype=float_dtype, + ) + if input_name == "image_position_ids": + # Create position ids as a grid. The patch count = h_patches * w_patches + # where both are divisible by pooling_kernel_size for correct pooling. + import math + + k = self.pooling_kernel_size + total_pooled = self.num_patches // (k * k) + # Find roughly square grid for pooled side + pooled_side = int(math.sqrt(total_pooled)) + if pooled_side * pooled_side < total_pooled: + pooled_h = pooled_side + pooled_w = total_pooled // pooled_h + else: + pooled_h = pooled_w = pooled_side + h_patches = pooled_h * k + w_patches = pooled_w * k + pos_ids = torch.stack( + torch.meshgrid(torch.arange(h_patches), torch.arange(w_patches), indexing="ij"), dim=-1 + ).reshape(1, -1, 2) + # Pad to num_patches with -1 (padding position) + if pos_ids.shape[1] < self.num_patches: + pad = torch.full((1, self.num_patches - pos_ids.shape[1], 2), -1, dtype=pos_ids.dtype) + pos_ids = torch.cat([pos_ids, pad], dim=1) + return pos_ids.expand(self.batch_size, -1, -1).clone() + return super().generate(input_name, framework, int_dtype, float_dtype) + + +class DummyVisionPositionIdsInputGenerator(DummyVisionInputGenerator): + SUPPORTED_INPUT_NAMES = ("patch_attention_mask", "patch_position_ids") + + def __init__( + self, + task: str, + normalized_config: NormalizedVisionConfig, + batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], + num_channels: int = DEFAULT_DUMMY_SHAPES["num_channels"], + width: int = DEFAULT_DUMMY_SHAPES["width"], + height: int = DEFAULT_DUMMY_SHAPES["height"], + **kwargs, + ): + super().__init__(task, normalized_config, batch_size, num_channels, width, height, **kwargs) + self.patch_size = normalized_config.config.patch_size + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + if input_name == "patch_attention_mask": + shape = [self.batch_size, self.height // self.patch_size, self.width // self.patch_size] + return self.random_int_tensor(shape, max_value=2, framework=framework, dtype="bool") + if input_name == "patch_position_ids": + max_nb_patches_h, max_nb_patches_w = self.height // self.patch_size, self.width // self.patch_size + shape = [self.batch_size, max_nb_patches_h * max_nb_patches_w] + return self.random_int_tensor( + shape, max_value=min(max_nb_patches_h, max_nb_patches_w), framework=framework, dtype=int_dtype + ) + return super().generate(input_name, framework, int_dtype, float_dtype) + + +class DummySpeechT5OpenVINOInputGenerator(DummyInputGenerator): + SUPPORTED_INPUT_NAMES = ( + "inputs_embeds", + "output_sequence", + "speaker_embeddings", + "spectrogram", + "raw_spectrogram", + "encoder_hidden_states", + ) + + def __init__( + self, + task: str, + normalized_config: NormalizedConfig, + sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], + **kwargs, + ): + self.task = task + self.batch_size = 1 + + self.sequence_length = sequence_length + self.speaker_embedding_dim = normalized_config.speaker_embedding_dim + self.num_mel_bins = normalized_config.num_mel_bins + self.reduction_factor = normalized_config.config.reduction_factor + self.hidden_size = normalized_config.config.hidden_size + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + if input_name in ["output_sequence", "inputs_embeds"]: + shape = [self.batch_size, self.sequence_length, self.num_mel_bins] + elif input_name == "speaker_embeddings": + shape = [self.batch_size, self.speaker_embedding_dim] + elif input_name == "raw_spectrogram": + shape = [self.sequence_length, self.batch_size, self.reduction_factor, self.num_mel_bins] + elif input_name == "encoder_hidden_states": + shape = [self.batch_size, self.sequence_length, self.hidden_size] + elif input_name == "spectrogram": + shape = [self.batch_size, self.sequence_length, self.num_mel_bins] + else: + raise ValueError(f"Unsupported input {input_name} for DummySpeechT5InputGenerator") + + return self.random_float_tensor( + shape=shape, + min_value=0, + max_value=1, + framework=framework, + dtype=float_dtype, + ) + + +class MambaCacheDummyInputGenerator(DummyInputGenerator): + """ + Generates dummy past_ssm_states, past_conv_states and cache_position inputs for Mamba architectures. + """ + + SUPPORTED_INPUT_NAMES = ("cache_params", "cache_position") + + def __init__( + self, + task: str, + normalized_config, + batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], + sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], + **kwargs, + ): + self.normalized_config = normalized_config + self.batch_size = batch_size + self.sequence_length = sequence_length + self.intermediate_size = self.normalized_config.config.intermediate_size + self.ssm_state_size = self.normalized_config.config.state_size + self.conv_kernel_size = self.normalized_config.config.conv_kernel + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + if input_name == "cache_params": + ssm_shape = [self.batch_size, self.intermediate_size, self.ssm_state_size] + conv_shape = [self.batch_size, self.intermediate_size, self.conv_kernel_size] + return [ + ( + self.random_float_tensor(ssm_shape, framework=framework, dtype=float_dtype), + self.random_float_tensor(conv_shape, framework=framework, dtype=float_dtype), + ) + for _ in range(self.normalized_config.num_layers) + ] + elif input_name == "cache_position": + return self.random_int_tensor( + shape=[self.conv_kernel_size], + max_value=self.sequence_length, + framework=framework, + dtype=int_dtype, + ) + + raise ValueError(f"Unsupported input name {input_name}") + + +class Zamba2DummyPastKeyValuesGenerator(DummyPastKeyValuesGenerator): + """ + Generates dummy cache_params inputs for Zamba2 architectures. + """ + + SUPPORTED_INPUT_NAMES = ("cache_params",) + + def __init__( + self, + task: str, + normalized_config, + 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, + ) + + config = normalized_config.config + self.intermediate_size = int(config.mamba_expand * config.hidden_size) + self.conv_kernel_size = config.mamba_d_conv + self.mamba_d_state = config.mamba_d_state + if config.model_type == "zamba2": + self.n_mamba_heads = config.n_mamba_heads + self.mamba_ngroups = config.mamba_ngroups + self.mamba_headdim = config.mamba_headdim + self.head_dim = config.attention_head_dim + # in Zamba2, all layers contain Mamba block + # some of these layers are hybrid so they contain both attention and mamba blocks + self.num_attention_layers = len(config.hybrid_layer_ids) + self.num_mamba_layers = self.num_layers + import logging + + logger = logging.getLogger(__name__) + logger.warning( + "The current support for the 'Zamba2' model type is experimental. " + "Performance is not optimal with high memory consumption. " + "Optimizations and improved support will be available in a future OpenVINO release." + ) + else: + # currently, this else-branch is applied for GraniteMoeHybrid models + self.n_mamba_heads = config.mamba_n_heads + self.mamba_ngroups = config.mamba_n_groups + self.mamba_headdim = config.mamba_d_head + self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + self.num_attention_layers = config.layer_types.count("attention") + self.num_mamba_layers = config.layer_types.count("mamba") + self.num_attention_heads = config.num_key_value_heads + self.sequence_length = 0 + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + past_key_values = [] + for i in range(self.num_mamba_layers): + conv_state_shape = ( + self.batch_size, + self.intermediate_size + 2 * self.mamba_ngroups * self.mamba_d_state, + self.conv_kernel_size, + ) + conv_state = self.random_float_tensor(conv_state_shape, framework=framework, dtype=float_dtype) + past_key_values.append(conv_state) + ssm_state_shape = (self.batch_size, self.n_mamba_heads, self.mamba_headdim, self.mamba_d_state) + ssm_state = self.random_float_tensor(ssm_state_shape, framework=framework, dtype=float_dtype) + past_key_values.append(ssm_state) + + for i in range(self.num_attention_layers): + kv_shape = (self.batch_size, self.num_attention_heads, self.sequence_length, self.head_dim) + k = self.random_float_tensor(kv_shape, framework=framework, dtype=float_dtype) + v = self.random_float_tensor(kv_shape, framework=framework, dtype=float_dtype) + past_key_values.append(k) + past_key_values.append(v) + + return past_key_values + + +class Lfm2DummyPastKeyValuesGenerator(DummyPastKeyValuesGenerator): + """ + Generates dummy past_key_values inputs for Lfm2 architectures. + """ + + SUPPORTED_INPUT_NAMES = ("cache_params",) + + def __init__( + self, + task: str, + normalized_config, + 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, + ) + config = normalized_config.config + self.num_conv_layers = config.layer_types.count("conv") + self.num_atten_layers = config.layer_types.count("full_attention") + self.batch_size = batch_size + self.normalized_config = normalized_config + self.hidden_size = self.normalized_config.hidden_size + self.conv_L_cache = self.normalized_config.conv_L_cache + self.num_key_value_heads = self.normalized_config.num_key_value_heads + self.num_hidden_layers = self.normalized_config.num_hidden_layers + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + past_key_values = [] + + for i in range(self.num_conv_layers): + conv_state_shape = (self.batch_size, self.hidden_size, self.conv_L_cache) + conv_state = self.random_float_tensor(conv_state_shape, framework=framework, dtype=float_dtype) + past_key_values.append(conv_state) + + for i in range(self.num_atten_layers): + shape = ( + self.batch_size, + self.num_key_value_heads, + self.sequence_length, + self.hidden_size // self.num_attention_heads, + ) + + kv_shape = shape # (self.batch_size, self.num_attention_heads, self.sequence_length, self.head_dim) + k = self.random_float_tensor(kv_shape, framework=framework, dtype=float_dtype) + v = self.random_float_tensor(kv_shape, framework=framework, dtype=float_dtype) + past_key_values.append(k) + past_key_values.append(v) + + return past_key_values + + +class DummyVideoChatFlashQwenInputGenerator(DummyVisionInputGenerator): + SUPPORTED_INPUT_NAMES = ("hidden_states", "rotary_pos_emb") + + def __init__( + self, + task: str, + normalized_config: NormalizedVisionConfig, + batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], + num_channels: int = DEFAULT_DUMMY_SHAPES["num_channels"], + width: int = DEFAULT_DUMMY_SHAPES["width"], + height: int = DEFAULT_DUMMY_SHAPES["height"], + visual_seq_length: int = DEFAULT_DUMMY_SHAPES["visual_seq_length"], + **kwargs, + ): + super().__init__(task, normalized_config, batch_size, num_channels, width, height, visual_seq_length, **kwargs) + self.num_frames = normalized_config.config.mm_local_num_frames + self.embed_dim = normalized_config.config.mm_hidden_size + self.height = normalized_config.config.image_size + self.width = normalized_config.config.image_size + self.image_size = (self.height, self.width) + self.patch_size = normalized_config.config.patch_size + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + if input_name == "hidden_states": + return self.random_float_tensor( + shape=[ + self.batch_size, + self.num_channels, + self.num_frames, + self.height, + self.width, + ], + framework=framework, + dtype=float_dtype, + ) + elif input_name == "rotary_pos_emb": + grid_h, grid_w = self.height // self.patch_size, self.width // self.patch_size + grid_t = self.num_frames + # Source: https://huggingface.co/OpenGVLab/VideoChat-Flash-Qwen2_5-7B_InternVideo2-1B/blob/main/vision_tower_builder.py#L523 + # The first dimension of rotary_pos_emb is fixed to 1 in the original model. + # And the second dimension is the total number of tokens for all frames, which is calculated as grid_h * grid_w * grid_t plus 1 for the cls token. + return self.random_float_tensor( + [1, 1 + grid_h * grid_t * grid_w, self.embed_dim], framework=framework, dtype=float_dtype + ) + + +class DummyVideoChatFlashQwenProjectorInputGenerator(DummyInputGenerator): + SUPPORTED_INPUT_NAMES = ["input"] + + def __init__( + self, + task: str, + normalized_config: NormalizedTextConfig, + batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], + random_batch_size_range: Optional[Tuple[int, int]] = None, + **kwargs, + ): + self.task = task + self.batch_size = batch_size + self.hidden_size = normalized_config.config.mm_hidden_size + self.num_patches = normalized_config.config.mm_projector_num_tome_tokens + self.normalized_config = normalized_config + + def generate( + self, + input_name: str, + framework: str = "pt", + int_dtype: str = "int64", + float_dtype: str = "fp32", + ): + shape = [self.batch_size, self.num_patches, self.hidden_size] + return self.random_float_tensor(shape, framework=framework, dtype=float_dtype) + + +class Qwen3NextDummyPastKeyValuesGenerator(DummyPastKeyValuesGenerator): + """ + Generates dummy cache_params inputs for Qwen3-Next architectures. + """ + + SUPPORTED_INPUT_NAMES = ("cache_params",) + + def __init__( + self, + task: str, + normalized_config, + 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, + ) + + config = normalized_config.config + self.num_full_attn_layers = config.layer_types.count("full_attention") + self.num_linear_attn_layers = config.layer_types.count("linear_attention") + self.conv_kernel_size = config.linear_conv_kernel_dim + self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads + self.head_k_dim = config.linear_key_head_dim + self.head_v_dim = config.linear_value_head_dim + self.num_v_heads = config.linear_num_value_heads + self.num_k_heads = config.linear_num_key_heads + self.num_key_value_heads = config.num_key_value_heads + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + cache_params = [] + + for idx in range(self.num_linear_attn_layers): + d_inner = self.num_k_heads * (2 * self.head_k_dim + self.head_v_dim * self.num_v_heads // self.num_k_heads) + conv_state_shape = ( + self.batch_size, + d_inner, + self.conv_kernel_size, + ) + conv_state = self.random_float_tensor(conv_state_shape, framework=framework, dtype=float_dtype) + cache_params.append(conv_state) + num_heads = self.num_v_heads + recurrent_state_shape = (self.batch_size, num_heads, self.head_k_dim, self.head_v_dim) + recurrent_state = self.random_float_tensor(recurrent_state_shape, framework=framework, dtype=float_dtype) + cache_params.append(recurrent_state) + + for idx in range(self.num_full_attn_layers): + kv_shape = (self.batch_size, self.num_key_value_heads, self.sequence_length, self.head_dim) + k = self.random_float_tensor(kv_shape, framework=framework, dtype=float_dtype) + v = self.random_float_tensor(kv_shape, framework=framework, dtype=float_dtype) + cache_params.append(k) + cache_params.append(v) + + return cache_params + + +class Qwen3_5DummyPastKeyValuesGenerator(DummyPastKeyValuesGenerator): + """ + Generates dummy cache_params inputs for Qwen3.5 architectures. + """ + + SUPPORTED_INPUT_NAMES = ("cache_params",) + + def __init__( + self, + task: str, + normalized_config, + 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, + ) + + config = normalized_config.config + self.num_full_attn_layers = config.layer_types.count("full_attention") + self.num_linear_attn_layers = config.layer_types.count("linear_attention") + self.conv_kernel_size = config.linear_conv_kernel_dim + self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + self.head_k_dim = config.linear_key_head_dim + self.head_v_dim = config.linear_value_head_dim + self.num_v_heads = config.linear_num_value_heads + self.num_k_heads = config.linear_num_key_heads + self.num_key_value_heads = config.num_key_value_heads + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + cache_params = [] + + for idx in range(self.num_linear_attn_layers): + d_inner = self.num_k_heads * (2 * self.head_k_dim + self.head_v_dim * self.num_v_heads // self.num_k_heads) + conv_state_shape = ( + self.batch_size, + d_inner, + self.conv_kernel_size, + ) + conv_state = self.random_float_tensor(conv_state_shape, framework=framework, dtype=float_dtype) + cache_params.append(conv_state) + num_heads = self.num_v_heads + recurrent_state_shape = (self.batch_size, num_heads, self.head_k_dim, self.head_v_dim) + recurrent_state = self.random_float_tensor(recurrent_state_shape, framework=framework, dtype=float_dtype) + cache_params.append(recurrent_state) + + for idx in range(self.num_full_attn_layers): + kv_shape = (self.batch_size, self.num_key_value_heads, self.sequence_length, self.head_dim) + k = self.random_float_tensor(kv_shape, framework=framework, dtype=float_dtype) + v = self.random_float_tensor(kv_shape, framework=framework, dtype=float_dtype) + cache_params.append(k) + cache_params.append(v) + + return cache_params + + +class DummyKokoroInputGenerator(DummyInputGenerator): + """Generates dummy inputs for the Kokoro TTS model.""" + + SUPPORTED_INPUT_NAMES = ("input_ids", "ref_s", "speed") + + def __init__( + self, + task: str, + normalized_config: NormalizedConfig, + sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], + **kwargs, + ): + self.task = task + self.batch_size = 1 + self.sequence_length = sequence_length + self.style_dim = getattr(normalized_config, "style_dim", 128) + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + if input_name == "input_ids": + shape = [self.batch_size, self.sequence_length] + input_ids_value = self.random_int_tensor( + shape=shape, min_value=0, max_value=178, framework=framework, dtype=int_dtype + ) + input_ids_value[:, 0] = 0 + input_ids_value[:, -1] = 0 + return input_ids_value + elif input_name == "ref_s": + shape = [self.batch_size, self.style_dim * 2] + return self.random_float_tensor( + shape=shape, min_value=-1, max_value=1, framework=framework, dtype=float_dtype + ) + elif input_name == "speed": + return self.random_int_tensor(shape=[1], min_value=1, max_value=10, framework=framework, dtype=float_dtype) + else: + raise ValueError(f"Unsupported input {input_name} for DummyKokoroInputGenerator") diff --git a/optimum/exporters/openvino/model_configs.py b/optimum/exporters/openvino/model_configs.py index 19d954b62a..aae4b81b58 100644 --- a/optimum/exporters/openvino/model_configs.py +++ b/optimum/exporters/openvino/model_configs.py @@ -15,123 +15,112 @@ import enum import logging import math +import warnings from copy import deepcopy -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union +from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union -import torch +from packaging import version from transformers import AutoConfig, PretrainedConfig, PreTrainedModel -from optimum.exporters.onnx.config import ( +from optimum.exporters.openvino.base import ConfigBehavior, OnnxConfig, OnnxConfigWithPast, OnnxSeq2SeqConfigWithPast +from optimum.exporters.openvino.config import ( + AudioOnnxConfig, AudioToTextOnnxConfig, + EncoderDecoderBaseOnnxConfig, OnnxConfig, + TextAndVisionOnnxConfig, TextDecoderOnnxConfig, TextDecoderWithPositionIdsOnnxConfig, -) -from optimum.exporters.onnx.model_configs import ( - AlbertOnnxConfig, - ASTOnnxConfig, - BartOnnxConfig, - BeitOnnxConfig, - BertOnnxConfig, - BlenderbotOnnxConfig, - BlenderbotSmallOnnxConfig, - BloomOnnxConfig, - CamembertOnnxConfig, - CLIPOnnxConfig, - CLIPTextOnnxConfig, - CLIPTextWithProjectionOnnxConfig, - CLIPVisionModelOnnxConfig, - CodeGenOnnxConfig, - ConvBertOnnxConfig, - ConvNextOnnxConfig, - Data2VecAudioOnnxConfig, - Data2VecTextOnnxConfig, - Data2VecVisionOnnxConfig, - DebertaOnnxConfig, - DebertaV2OnnxConfig, - DeiTOnnxConfig, - DistilBertOnnxConfig, - ElectraOnnxConfig, - EsmOnnxConfig, - FalconOnnxConfig, - FlaubertOnnxConfig, - GemmaOnnxConfig, - GPT2OnnxConfig, - GPTBigCodeOnnxConfig, - GPTJOnnxConfig, - GPTNeoOnnxConfig, - GPTNeoXOnnxConfig, - HubertOnnxConfig, - IBertOnnxConfig, - LevitOnnxConfig, - LlamaOnnxConfig, - MarianOnnxConfig, - MistralOnnxConfig, - MobileBertOnnxConfig, - MobileNetV1OnnxConfig, - MobileNetV2OnnxConfig, - MobileViTOnnxConfig, - MPNetOnnxConfig, - MPTOnnxConfig, - NystromformerOnnxConfig, - Olmo2OnnxConfig, - OPTOnnxConfig, - PegasusOnnxConfig, - PerceiverOnnxConfig, - PhiOnnxConfig, - Pix2StructOnnxConfig, - PoolFormerOnnxConfig, - RemBertOnnxConfig, - ResNetOnnxConfig, - RobertaOnnxConfig, - RoFormerOnnxConfig, - SamOnnxConfig, - SegformerOnnxConfig, - SentenceTransformersTransformerOnnxConfig, - SEWDOnnxConfig, - SEWOnnxConfig, - SiglipOnnxConfig, - SiglipTextOnnxConfig, - SiglipTextWithProjectionOnnxConfig, - SpeechT5OnnxConfig, - SqueezeBertOnnxConfig, - SwinOnnxConfig, - T5OnnxConfig, - UNetOnnxConfig, - UniSpeechOnnxConfig, - UniSpeechSATOnnxConfig, - VaeDecoderOnnxConfig, - VaeEncoderOnnxConfig, - VisionEncoderDecoderOnnxConfig, + TextEncoderOnnxConfig, + TextSeq2SeqOnnxConfig, VisionOnnxConfig, - ViTOnnxConfig, - Wav2Vec2ConformerOnnxConfig, - Wav2Vec2OnnxConfig, - WavLMOnnxConfig, - WhisperOnnxConfig, - XLMOnnxConfig, - XLMRobertaOnnxConfig, ) -from optimum.exporters.onnx.model_patcher import ModelPatcher +from optimum.exporters.openvino.model_patcher import ( + BigBirdPegasusModelPatcher, + CLIPModelPatcher, + CohereModelPatcher, + FluxTransformerModelPatcher, + GptOssModelPatcher, + MetaCLIP2Patcher, + MgpstrModelPatcher, + ModelPatcher, + MoonshineModelPatcher, + MusicgenModelPatcher, + Qwen3MoeModelPatcher, + SAMModelPatcher, + SentenceTransformersCLIPPatcher, + SentenceTransformersTransformerPatcher, + SpeechT5ModelPatcher, + ViTForImageClassificationPatcher, + VitPoseModelPatcher, +) from optimum.exporters.openvino.utils import ONNX_SUPPORTED_ARCHITECTURES from optimum.exporters.tasks import TasksManager -from optimum.utils import DEFAULT_DUMMY_SHAPES -from optimum.utils.input_generators import ( - DTYPE_MAPPER, - DummyAudioInputGenerator, +from optimum.utils import ( + DEFAULT_DUMMY_SHAPES, + ASTDummyAudioInputGenerator, + BartDummyTextInputGenerator, + BloomDummyPastKeyValuesGenerator, + DeepSeekV3DummyPastKeyValuesGenerator, + Dinov2DummyInputGenerator, + DummyCodegenDecoderTextInputGenerator, + DummyDecisionTransformerInputGenerator, + DummyDecoderTextInputGenerator, + DummyEncodecInputGenerator, + DummyFluxTransformerTextInputGenerator, + DummyFluxTransformerVisionInputGenerator, DummyInputGenerator, + DummyIntGenerator, DummyPastKeyValuesGenerator, + DummyPatchTSTInputGenerator, + DummyPix2StructInputGenerator, + DummyPointsGenerator, DummySeq2SeqDecoderTextInputGenerator, DummySeq2SeqPastKeyValuesGenerator, + DummySpeechT5InputGenerator, DummyTextInputGenerator, DummyTimestepInputGenerator, + DummyTransformerTextInputGenerator, + DummyTransformerTimestepInputGenerator, + DummyTransformerVisionInputGenerator, + DummyVisionEmbeddingsGenerator, + DummyVisionEncoderDecoderPastKeyValuesGenerator, DummyVisionInputGenerator, + DummyXPathSeqInputGenerator, FalconDummyPastKeyValuesGenerator, GemmaDummyPastKeyValuesGenerator, + LongformerDummyTextInputGenerator, + MCTCTDummyAudioInputGenerator, + MistralDummyPastKeyValuesGenerator, + NormalizedConfig, + NormalizedEncoderDecoderConfig, + NormalizedSeq2SeqConfig, + NormalizedTextAndVisionConfig, + NormalizedTextConfig, + NormalizedTextConfigWithGQA, + NormalizedTimeSeriesForecastingConfig, + NormalizedVisionConfig, + PerceiverDummyInputGenerator, + Speech2TextDummyAudioInputGenerator, + T5DummySeq2SeqPastKeyValuesGenerator, + VitPoseDummyInputGenerator, + is_diffusers_version, + is_transformers_version, +) +from optimum.utils.input_generators import ( + DummyAudioInputGenerator, + DummyInputGenerator, + DummyPastKeyValuesGenerator, + DummySeq2SeqDecoderTextInputGenerator, + DummySeq2SeqPastKeyValuesGenerator, + DummyTextInputGenerator, + DummyVisionInputGenerator, + GemmaDummyPastKeyValuesGenerator, MistralDummyPastKeyValuesGenerator, ) from optimum.utils.normalized_config import ( NormalizedConfig, + NormalizedConfigManager, NormalizedSeq2SeqConfig, NormalizedTextConfig, NormalizedVisionConfig, @@ -143,6 +132,56 @@ is_openvino_version, is_transformers_version, ) +from .input_generators import ( + AquilaDummyPastKeyValuesGenerator, + ChatGLM2DummyPastKeyValuesGenerator, + DeciDummyPastKeyValuesGenerator, + DummyAudioPhi4MMInputGenerator, + DummyFluxTextInputGenerator, + DummyFluxTransformerInputGenerator, + DummyGemma4VisionInputGenerator, + DummyKokoroInputGenerator, + DummyLLavaMultiModalProjectorInputGenerator, + DummyMiniCPMVImageInputGenerator, + DummyMiniCPMVResampleInputGenerator, + DummyMoonshineAudioInputGenerator, + DummyPhi3VisionProjectionInputGenerator, + DummyQwen2VLLMInputGenerator, + DummyQwen2VLVisionEmbedInputGenerator, + DummyQwen3_5LMInputGenerator, + DummyQwen3VLLMInputGenerator, + DummyQwen3VLVisionEmbedInputGenerator, + DummySanaSeq2SeqDecoderTextWithEncMaskInputGenerator, + DummySanaTimestepInputGenerator, + DummySanaTransforemerTextInputGenerator, + DummySanaTransformerVisionInputGenerator, + DummySpeechT5OpenVINOInputGenerator, + DummyTransformerTimestpsInputGenerator, + DummyUnetEncoderInputGenerator, + DummyUnetTimestepInputGenerator, + DummyUnetVisionInputGenerator, + DummyVideoChatFlashQwenInputGenerator, + DummyVideoChatFlashQwenProjectorInputGenerator, + DummyVisionPositionIdsInputGenerator, + DummyVisionPositionIdsPhi4InputGenerator, + Eagle3DummyGenerator, + Eagle3VLMDummyGenerator, + Gemma4DummyPastKeyValuesGenerator, + GPTBigCodeDummyPastKeyValuesGenerator, + Lfm2DummyPastKeyValuesGenerator, + LTXTransformerDummyInputGenerator, + LTXVaeDummyInputGenerator, + MambaCacheDummyInputGenerator, + OVFalconDummyPastKeyValuesGenerator, + OVMiniCPM3DummyPastKeyValuesGenerator, + OVMistralDummyPastKeyValuesGenerator, + PooledProjectionsDummyInputGenerator, + Qwen3_5DummyPastKeyValuesGenerator, + Qwen3ASRDummySeq2SeqPastKeyValuesGenerator, + Qwen3NextDummyPastKeyValuesGenerator, + QwenDummyPastKeyValuesGenerator, + Zamba2DummyPastKeyValuesGenerator, +) from .model_patcher import ( AfmoeModelPatcher, AquilaModelPatcher, @@ -229,30 +268,6 @@ ) -COMMON_TEXT_TASKS = [ - "feature-extraction", - "fill-mask", - "multiple-choice", - "question-answering", - "text-classification", - "token-classification", -] - - -COMMON_TEXT_GENERATION_TASKS = [ - "feature-extraction", - "feature-extraction-with-past", - "text-generation", - "text-generation-with-past", -] - -COMMON_TEXT2TEXT_GENERATION_TASKS = [ - *COMMON_TEXT_GENERATION_TASKS, - "text2text-generation", - "text2text-generation-with-past", -] - - logger = logging.getLogger(__name__) @@ -357,29 +372,2766 @@ def init_model_configs(): TasksManager._DIFFUSERS_TASKS_TO_MODEL_MAPPINGS["text-to-video"] = {} TasksManager._DIFFUSERS_TASKS_TO_MODEL_MAPPINGS["text-to-video"]["ltx-video"] = "LTXPipeline" - supported_model_types = [ - "_SUPPORTED_MODEL_TYPE", - "_DIFFUSERS_SUPPORTED_MODEL_TYPE", - "_TIMM_SUPPORTED_MODEL_TYPE", - "_SENTENCE_TRANSFORMERS_SUPPORTED_MODEL_TYPE", - ] - # TODO: remove once models from ONNX_SUPPORTED_ARCHITECTURES are deprecated (optimum-intel v1.29) - for supported_models_config in supported_model_types: - supported_models = getattr(TasksManager, supported_models_config) - for model, export_configs in supported_models.items(): - # adding only the architectures that are already supported via optimum-onnx v0.1.0 - if "onnx" not in export_configs or model not in ONNX_SUPPORTED_ARCHITECTURES: - continue - onnx_config = export_configs["onnx"] - supported_models[model]["openvino"] = deepcopy(onnx_config) + supported_model_types = [ + "_SUPPORTED_MODEL_TYPE", + "_DIFFUSERS_SUPPORTED_MODEL_TYPE", + "_TIMM_SUPPORTED_MODEL_TYPE", + "_SENTENCE_TRANSFORMERS_SUPPORTED_MODEL_TYPE", + ] + # TODO: remove once models from ONNX_SUPPORTED_ARCHITECTURES are deprecated (optimum-intel v1.29) + for supported_models_config in supported_model_types: + supported_models = getattr(TasksManager, supported_models_config) + for model, export_configs in supported_models.items(): + # adding only the architectures that are already supported via optimum-onnx v0.1.0 + if "onnx" not in export_configs or model not in ONNX_SUPPORTED_ARCHITECTURES: + continue + onnx_config = export_configs["onnx"] + supported_models[model]["openvino"] = deepcopy(onnx_config) + + setattr(TasksManager, supported_models_config, supported_models) + + +init_model_configs() + + +register_in_tasks_manager = TasksManager.create_register("openvino", overwrite_existing=True) + + +COMMON_TEXT_TASKS = [ + "feature-extraction", + "fill-mask", + "multiple-choice", + "question-answering", + "text-classification", + "token-classification", +] + + +COMMON_TEXT_GENERATION_TASKS = [ + "feature-extraction", + "feature-extraction-with-past", + "text-generation", + "text-generation-with-past", +] + +COMMON_TEXT2TEXT_GENERATION_TASKS = [ + *COMMON_TEXT_GENERATION_TASKS, + "text2text-generation", + "text2text-generation-with-past", +] + + +register_tasks_manager_onnx = TasksManager.create_register("onnx") + + +@register_tasks_manager_onnx("bert", *COMMON_TEXT_TASKS) +class BertOnnxConfig(TextEncoderOnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig + + @property + def inputs(self) -> dict[str, dict[int, str]]: + if self.task == "multiple-choice": + dynamic_axis = {0: "batch_size", 1: "num_choices", 2: "sequence_length"} + else: + dynamic_axis = {0: "batch_size", 1: "sequence_length"} + return { + "input_ids": dynamic_axis, + "attention_mask": dynamic_axis, + "token_type_ids": dynamic_axis, + } + + +@register_tasks_manager_onnx("visual_bert", *["feature-extraction"]) +class VisualBertOnnxConfig(TextAndVisionOnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig + DUMMY_INPUT_GENERATOR_CLASSES = ( + DummyTextInputGenerator, + DummyVisionInputGenerator, + ) + + @property + def inputs(self) -> dict[str, dict[int, str]]: + return { + "input_ids": {0: "batch_size", 1: "sequence_length"}, + "attention_mask": {0: "batch_size", 1: "sequence_length"}, + "visual_embeds": {0: "batch_size", 1: "visual_seq_length", 2: "visual_embedding_dim"}, + "visual_attention_mask": {0: "batch_size", 1: "visual_seq_length"}, + "visual_token_type_ids": {0: "batch_size", 1: "visual_seq_length"}, + } + + @property + def outputs(self) -> dict[str, dict[int, str]]: + return { + "last_hidden_state": {0: "batch_size", 1: "sequence_length + visual_seq_length"}, + } + + +@register_tasks_manager_onnx("albert", *COMMON_TEXT_TASKS) +class AlbertOnnxConfig(BertOnnxConfig): + pass + + +@register_tasks_manager_onnx("convbert", *COMMON_TEXT_TASKS) +class ConvBertOnnxConfig(BertOnnxConfig): + pass + + +@register_tasks_manager_onnx("electra", *COMMON_TEXT_TASKS) +class ElectraOnnxConfig(BertOnnxConfig): + pass + + +@register_tasks_manager_onnx("roformer", *COMMON_TEXT_TASKS) +class RoFormerOnnxConfig(BertOnnxConfig): + pass + + +@register_tasks_manager_onnx("squeezebert", *COMMON_TEXT_TASKS) +class SqueezeBertOnnxConfig(BertOnnxConfig): + pass + + +@register_tasks_manager_onnx("mobilebert", *COMMON_TEXT_TASKS) +class MobileBertOnnxConfig(BertOnnxConfig): + pass + + +@register_tasks_manager_onnx("nystromformer", *COMMON_TEXT_TASKS) +class NystromformerOnnxConfig(BertOnnxConfig): + pass + + +@register_tasks_manager_onnx("xlm", *COMMON_TEXT_TASKS) +class XLMOnnxConfig(BertOnnxConfig): + pass + + +@register_tasks_manager_onnx("splinter", *["feature-extraction", "question-answering"]) +class SplinterOnnxConfig(BertOnnxConfig): + pass + + +@register_tasks_manager_onnx("rembert", *COMMON_TEXT_TASKS) +class RemBertOnnxConfig(BertOnnxConfig): + pass + + +@register_tasks_manager_onnx("longformer", *COMMON_TEXT_TASKS) +class LongformerOnnxConfig(BertOnnxConfig): + DUMMY_INPUT_GENERATOR_CLASSES = (LongformerDummyTextInputGenerator,) + + @property + def inputs(self) -> dict[str, dict[int, str]]: + inputs = super().inputs + + inputs["global_attention_mask"] = inputs["attention_mask"] + + return inputs + + +@register_tasks_manager_onnx("megatron-bert", *COMMON_TEXT_TASKS) +class MegatronBertOnnxConfig(BertOnnxConfig): + pass + + +@register_tasks_manager_onnx("distilbert", *COMMON_TEXT_TASKS) +class DistilBertOnnxConfig(BertOnnxConfig): + @property + def inputs(self) -> dict[str, dict[int, str]]: + if self.task == "multiple-choice": + dynamic_axis = {0: "batch_size", 1: "num_choices", 2: "sequence_length"} + else: + dynamic_axis = {0: "batch_size", 1: "sequence_length"} + return {"input_ids": dynamic_axis, "attention_mask": dynamic_axis} + + +@register_tasks_manager_onnx( + "modernbert", + *[ + "feature-extraction", + "fill-mask", + "text-classification", + "token-classification", + ], +) +class ModernBertOnnxConfig(DistilBertOnnxConfig): + MIN_TRANSFORMERS_VERSION = version.parse("4.48.0") + + +@register_tasks_manager_onnx("mpnet", *COMMON_TEXT_TASKS) +class MPNetOnnxConfig(DistilBertOnnxConfig): + pass + + +@register_tasks_manager_onnx("roberta", *COMMON_TEXT_TASKS) +class RobertaOnnxConfig(DistilBertOnnxConfig): + pass + + +@register_tasks_manager_onnx("camembert", *COMMON_TEXT_TASKS) +class CamembertOnnxConfig(DistilBertOnnxConfig): + pass + + +@register_tasks_manager_onnx("flaubert", *COMMON_TEXT_TASKS) +class FlaubertOnnxConfig(BertOnnxConfig): + pass + + +@register_tasks_manager_onnx("ibert", *COMMON_TEXT_TASKS) +class IBertOnnxConfig(DistilBertOnnxConfig): + pass + + +@register_tasks_manager_onnx("xlm-roberta", *COMMON_TEXT_TASKS) +class XLMRobertaOnnxConfig(DistilBertOnnxConfig): + pass + + +@register_tasks_manager_onnx( + "deberta", + *["feature-extraction", "fill-mask", "text-classification", "token-classification", "question-answering"], +) +class DebertaOnnxConfig(BertOnnxConfig): + @property + def inputs(self) -> dict[str, dict[int, str]]: + common_inputs = super().inputs + if self._config.type_vocab_size == 0: + common_inputs.pop("token_type_ids") + return common_inputs + + +@register_tasks_manager_onnx( + "markuplm", *["feature-extraction", "text-classification", "token-classification", "question-answering"] +) +class MarkupLMOnnxConfig(BertOnnxConfig): + DUMMY_INPUT_GENERATOR_CLASSES = ( + DummyTextInputGenerator, + DummyXPathSeqInputGenerator, + ) + + @property + def inputs(self) -> dict[str, dict[int, str]]: + dynamic_axis = {0: "batch_size", 1: "sequence_length"} + xpath_dynamic_axis = {0: "batch_size", 1: "sequence_length", 2: "max_depth"} + return { + "input_ids": dynamic_axis, + "attention_mask": dynamic_axis, + "token_type_ids": dynamic_axis, + "xpath_subs_seq": xpath_dynamic_axis, + "xpath_tags_seq": xpath_dynamic_axis, + } + + +@register_tasks_manager_onnx("deberta-v2", *COMMON_TEXT_TASKS) +class DebertaV2OnnxConfig(DebertaOnnxConfig): + pass + + +@register_tasks_manager_onnx( + "esm", *["feature-extraction", "fill-mask", "text-classification", "token-classification"] +) +class EsmOnnxConfig(TextEncoderOnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig + + @property + def inputs(self) -> dict[str, dict[int, str]]: + dynamic_axis = {0: "batch_size", 1: "sequence_length"} + return { + "input_ids": dynamic_axis, + "attention_mask": dynamic_axis, + } + + +@register_tasks_manager_onnx("gpt2", *[*COMMON_TEXT_GENERATION_TASKS, "text-classification", "token-classification"]) +class GPT2OnnxConfig(TextDecoderWithPositionIdsOnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args(num_layers="n_layer", num_attention_heads="n_head") + + +@register_tasks_manager_onnx("gptj", *[*COMMON_TEXT_GENERATION_TASKS, "text-classification", "question-answering"]) +class GPTJOnnxConfig(GPT2OnnxConfig): + pass + + +@register_tasks_manager_onnx("codegen", *COMMON_TEXT_GENERATION_TASKS) +class CodeGenOnnxConfig(GPT2OnnxConfig): + pass + + +@register_tasks_manager_onnx("imagegpt", *["feature-extraction", "image-classification"]) +class ImageGPTOnnxConfig(GPT2OnnxConfig): + pass + + +@register_tasks_manager_onnx("decision_transformer", *["feature-extraction", "reinforcement-learning"]) +class DecisionTransformerOnnxConfig(OnnxConfig): + DUMMY_INPUT_GENERATOR_CLASSES = (DummyDecisionTransformerInputGenerator,) + NORMALIZED_CONFIG_CLASS = NormalizedConfig + + @property + def inputs(self) -> dict[str, dict[int, str]]: + return { + "states": {0: "batch_size", 1: "sequence_length"}, + "actions": {0: "batch_size", 1: "sequence_length"}, + "timesteps": {0: "batch_size", 1: "sequence_length"}, + "returns_to_go": {0: "batch_size", 1: "sequence_length"}, + "attention_mask": {0: "batch_size", 1: "sequence_length"}, + } + + @property + def outputs(self) -> dict[str, dict[int, str]]: + return { + "state_preds": {0: "batch_size", 1: "sequence_length"}, + "action_preds": {0: "batch_size", 1: "sequence_length"}, + "return_preds": {0: "batch_size", 1: "sequence_length"}, + "last_hidden_state": {0: "batch_size", 1: "sequence_length"}, + } + + +@register_tasks_manager_onnx("gpt_neo", *[*COMMON_TEXT_GENERATION_TASKS, "text-classification"]) +class GPTNeoOnnxConfig(TextDecoderWithPositionIdsOnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args(num_attention_heads="num_heads") + + +@register_tasks_manager_onnx("gpt_neox", *[*COMMON_TEXT_GENERATION_TASKS, "text-classification"]) +class GPTNeoXOnnxConfig(TextDecoderWithPositionIdsOnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig + + +@register_tasks_manager_onnx("opt", *[*COMMON_TEXT_GENERATION_TASKS, "text-classification", "question-answering"]) +class OPTOnnxConfig( + TextDecoderWithPositionIdsOnnxConfig if is_transformers_version(">=", "4.46.0") else TextDecoderOnnxConfig +): + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig + + +@register_tasks_manager_onnx("llama", *[*COMMON_TEXT_GENERATION_TASKS, "text-classification"]) +class LlamaOnnxConfig(TextDecoderWithPositionIdsOnnxConfig): + DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, MistralDummyPastKeyValuesGenerator) + DUMMY_PKV_GENERATOR_CLASS = MistralDummyPastKeyValuesGenerator + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig + + +@register_tasks_manager_onnx("arcee", *COMMON_TEXT_GENERATION_TASKS) +class ArceeOnnxConfig(LlamaOnnxConfig): + MIN_TRANSFORMERS_VERSION = version.parse("4.53.0") + NORMALIZED_CONFIG_CLASS = NormalizedTextConfigWithGQA + + +@register_tasks_manager_onnx("deepseek_v3", *COMMON_TEXT_GENERATION_TASKS) +class DeepSeekV3OnnxConfig(LlamaOnnxConfig): + MIN_TRANSFORMERS_VERSION = version.parse("4.51.0") + DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, DeepSeekV3DummyPastKeyValuesGenerator) + DUMMY_PKV_GENERATOR_CLASS = DeepSeekV3DummyPastKeyValuesGenerator + + +@register_tasks_manager_onnx("cohere", *COMMON_TEXT_GENERATION_TASKS) +class CohereOnnxConfig(LlamaOnnxConfig): + MIN_TRANSFORMERS_VERSION = version.parse("4.38.0") + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig + _MODEL_PATCHER = CohereModelPatcher + + +@register_tasks_manager_onnx("glm", *COMMON_TEXT_GENERATION_TASKS) +class GLMOnnxConfig(LlamaOnnxConfig): + MIN_TRANSFORMERS_VERSION = version.parse("4.46.0") + + +@register_tasks_manager_onnx("helium", *COMMON_TEXT_GENERATION_TASKS) +class HeliumOnnxConfig(LlamaOnnxConfig): + MIN_TRANSFORMERS_VERSION = version.parse("4.49.0") + + +@register_tasks_manager_onnx("smollm3", *[*COMMON_TEXT_GENERATION_TASKS, "text-classification"]) +class SmolLM3OnnxConfig(LlamaOnnxConfig): + MIN_TRANSFORMERS_VERSION = version.parse("4.53.0") + + +@register_tasks_manager_onnx("stablelm", *COMMON_TEXT_GENERATION_TASKS) +class StableLMOnnxConfig(LlamaOnnxConfig): + MIN_TRANSFORMERS_VERSION = version.parse("4.38.0") + + +@register_tasks_manager_onnx("olmo", *COMMON_TEXT_GENERATION_TASKS) +class OlmoOnnxConfig(LlamaOnnxConfig): + MIN_TRANSFORMERS_VERSION = version.parse("4.40.0") + + +@register_tasks_manager_onnx("olmo2", *COMMON_TEXT_GENERATION_TASKS) +class Olmo2OnnxConfig(OlmoOnnxConfig): + MIN_TRANSFORMERS_VERSION = version.parse("4.47.0") + + +@register_tasks_manager_onnx("qwen2", *[*COMMON_TEXT_GENERATION_TASKS, "text-classification", "token-classification"]) +class Qwen2OnnxConfig(LlamaOnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedTextConfigWithGQA + MIN_TRANSFORMERS_VERSION = version.parse("4.37.0") + + +@register_tasks_manager_onnx("qwen3", *[*COMMON_TEXT_GENERATION_TASKS, "text-classification"]) +class Qwen3OnnxConfig(LlamaOnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedTextConfigWithGQA + MIN_TRANSFORMERS_VERSION = version.parse("4.51.0") + + +@register_tasks_manager_onnx( + "qwen3_moe", *[*COMMON_TEXT_GENERATION_TASKS, "text-classification", "token-classification"] +) +class Qwen3MoeOnnxConfig(LlamaOnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedTextConfigWithGQA + MIN_TRANSFORMERS_VERSION = version.parse("4.51.0") + _MODEL_PATCHER = Qwen3MoeModelPatcher + + +@register_tasks_manager_onnx("gemma", *[*COMMON_TEXT_GENERATION_TASKS, "text-classification"]) +class GemmaOnnxConfig(TextDecoderOnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig + DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, GemmaDummyPastKeyValuesGenerator) + DUMMY_PKV_GENERATOR_CLASS = GemmaDummyPastKeyValuesGenerator + MIN_TRANSFORMERS_VERSION = version.parse("4.38.0") + + +@register_tasks_manager_onnx("gemma2", *[*COMMON_TEXT_GENERATION_TASKS, "text-classification"]) +class Gemma2OnnxConfig(GemmaOnnxConfig): + # Gemma 2 was added in transformers v4.42 using HybridCache + # DynamicCache support was added since v4.53 + MIN_TRANSFORMERS_VERSION = version.parse("4.53.0") + + +@register_tasks_manager_onnx("gemma3_text", *COMMON_TEXT_GENERATION_TASKS) +class Gemma3TextOnnxConfig(GemmaOnnxConfig): + # Gemma 3 was added in transformers v4.50 using HybridCache + # DynamicCache support was added since v4.53 + MIN_TRANSFORMERS_VERSION = version.parse("4.53.0") + + +# we still don't support gemma3 for multimodal feature-extraction(-with-past) and image-text-to-text(-with-past) tasks +@register_tasks_manager_onnx("gemma3", *COMMON_TEXT_GENERATION_TASKS, "text-classification") +class Gemma3OnnxConfig(GemmaOnnxConfig): + # Gemma 3 was added in transformers v4.50 using HybridCache + # DynamicCache support was added since v4.53 + MIN_TRANSFORMERS_VERSION = version.parse("4.53.0") + + def __init__(self, config: PretrainedConfig, **kwargs): + super().__init__(config.text_config, **kwargs) + + +@register_tasks_manager_onnx("gpt_oss", *COMMON_TEXT_GENERATION_TASKS) +class GPTOssOnnxConfig(GemmaOnnxConfig): + MIN_TRANSFORMERS_VERSION = version.parse("4.55.0") + _MODEL_PATCHER = GptOssModelPatcher + + +@register_tasks_manager_onnx("nemotron", *COMMON_TEXT_GENERATION_TASKS) +class NemotronOnnxConfig(GemmaOnnxConfig): + MIN_TRANSFORMERS_VERSION = version.parse("4.48.0") # More stable version than 4.44.0 + NORMALIZED_CONFIG_CLASS = NormalizedTextConfigWithGQA + + +@register_tasks_manager_onnx("granite", *COMMON_TEXT_GENERATION_TASKS) +class GraniteOnnxConfig(LlamaOnnxConfig): + MIN_TRANSFORMERS_VERSION = version.parse("4.45.0") + + +@register_tasks_manager_onnx("phi", *[*COMMON_TEXT_GENERATION_TASKS, "text-classification"]) +class PhiOnnxConfig(TextDecoderWithPositionIdsOnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig + MIN_TRANSFORMERS_VERSION = version.parse("4.36.0") + + +@register_tasks_manager_onnx("phi3", *[*COMMON_TEXT_GENERATION_TASKS, "text-classification"]) +class Phi3OnnxConfig(PhiOnnxConfig): + DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, MistralDummyPastKeyValuesGenerator) + DUMMY_PKV_GENERATOR_CLASS = MistralDummyPastKeyValuesGenerator + NORMALIZED_CONFIG_CLASS = NormalizedTextConfigWithGQA + MIN_TRANSFORMERS_VERSION = version.parse("4.41.0") + + +@register_tasks_manager_onnx("internlm2", *["text-generation", "text-generation-with-past"]) +class InternLM2OnnxConfig(LlamaOnnxConfig): + MIN_TRANSFORMERS_VERSION = version.parse("4.41.0") + + +@register_tasks_manager_onnx("mistral", *[*COMMON_TEXT_GENERATION_TASKS, "text-classification"]) +class MistralOnnxConfig(TextDecoderWithPositionIdsOnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args(num_key_value_heads="num_key_value_heads", allow_new=True) + DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, MistralDummyPastKeyValuesGenerator) + DUMMY_PKV_GENERATOR_CLASS = MistralDummyPastKeyValuesGenerator + + +@register_tasks_manager_onnx("mpt", *[*COMMON_TEXT_GENERATION_TASKS, "text-classification", "token-classification"]) +class MPTOnnxConfig(TextDecoderOnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args( + num_attention_heads="n_heads", hidden_size="d_model", num_layers="n_layers" + ) + + +@register_tasks_manager_onnx("bloom", *[*COMMON_TEXT_GENERATION_TASKS, "text-classification", "token-classification"]) +class BloomOnnxConfig(TextDecoderOnnxConfig): + # Bloom does not require position_ids input. + MIN_TRANSFORMERS_VERSION = version.parse("4.36.0") + DUMMY_PKV_GENERATOR_CLASS = BloomDummyPastKeyValuesGenerator + DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, BloomDummyPastKeyValuesGenerator) + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args(num_layers="n_layer", num_attention_heads="n_head") + + def add_past_key_values(self, inputs_or_outputs: dict[str, dict[int, str]], direction: str): + if is_transformers_version(">=", "4.44"): + super().add_past_key_values(inputs_or_outputs, direction) + else: + if direction not in ["inputs", "outputs"]: + raise ValueError(f'direction must either be "inputs" or "outputs", but {direction} was given') + + if direction == "inputs": + decoder_sequence_name = "past_sequence_length" + name = "past_key_values" + else: + decoder_sequence_name = "past_sequence_length + sequence_length" + name = "present" + + for i in range(self._normalized_config.num_layers): + inputs_or_outputs[f"{name}.{i}.key"] = {0: "batch_size * num_heads", 2: decoder_sequence_name} + inputs_or_outputs[f"{name}.{i}.value"] = {0: "batch_size * num_heads", 1: decoder_sequence_name} + + +@register_tasks_manager_onnx( + "gpt_bigcode", *[*COMMON_TEXT_GENERATION_TASKS, "text-classification", "token-classification"] +) +class GPTBigCodeOnnxConfig(TextDecoderWithPositionIdsOnnxConfig): + DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, GPTBigCodeDummyPastKeyValuesGenerator) + NORMALIZED_CONFIG_CLASS = NormalizedConfigManager.get_normalized_config_class("gpt_bigcode") + DUMMY_PKV_GENERATOR_CLASS = GPTBigCodeDummyPastKeyValuesGenerator + + def add_past_key_values(self, inputs_or_outputs: dict[str, dict[int, str]], direction: str): + if is_transformers_version(">=", "4.54"): + super().add_past_key_values(inputs_or_outputs, direction) + else: + if direction not in ["inputs", "outputs"]: + raise ValueError(f'direction must either be "inputs" or "outputs", but {direction} was given') + + if direction == "inputs": + decoder_sequence_name = "past_sequence_length" + name = "past_key_values" + else: + decoder_sequence_name = "past_sequence_length + sequence_length" + name = "present" + + if self._normalized_config.multi_query: + decoder_sequence_dim = 1 + else: + decoder_sequence_dim = 2 + + for i in range(self._normalized_config.num_layers): + inputs_or_outputs[f"{name}.{i}.key_value"] = { + 0: "batch_size", + decoder_sequence_dim: decoder_sequence_name, + } + + def flatten_past_key_values(self, flattened_output, name, idx, t): + if is_transformers_version(">=", "4.54"): + super().flatten_past_key_values(flattened_output, name, idx, t) + else: + flattened_output[f"{name}.{idx}.key_value"] = t + + +@register_tasks_manager_onnx("falcon", *[*COMMON_TEXT_GENERATION_TASKS, "question-answering", "token-classification"]) +class FalconOnnxConfig(TextDecoderWithPositionIdsOnnxConfig): + DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, FalconDummyPastKeyValuesGenerator) + DUMMY_PKV_GENERATOR_CLASS = FalconDummyPastKeyValuesGenerator + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig + + @property + def inputs(self) -> dict[str, dict[int, str]]: + common_inputs = super().inputs + + if self._config.alibi: + common_inputs.pop("position_ids", None) + + return common_inputs + + +@register_tasks_manager_onnx( + "t5", + *["feature-extraction", "feature-extraction-with-past", "text2text-generation", "text2text-generation-with-past"], +) +class T5OnnxConfig(TextSeq2SeqOnnxConfig): + DUMMY_INPUT_GENERATOR_CLASSES = ( + *TextSeq2SeqOnnxConfig.DUMMY_INPUT_GENERATOR_CLASSES[:-1], + T5DummySeq2SeqPastKeyValuesGenerator, + ) + DUMMY_PKV_GENERATOR_CLASS = T5DummySeq2SeqPastKeyValuesGenerator + NORMALIZED_CONFIG_CLASS = NormalizedSeq2SeqConfig.with_args( + hidden_size="d_model", + num_attention_heads="num_heads", + encoder_num_layers="num_layers", + decoder_num_layers="num_decoder_layers", + key_value_dim="d_kv", + allow_new=True, + ) + + +@register_tasks_manager_onnx( + "mt5", + *["feature-extraction", "feature-extraction-with-past", "text2text-generation", "text2text-generation-with-past"], +) +class MT5OnnxConfig(T5OnnxConfig): + pass + + +@register_tasks_manager_onnx( + "longt5", + *["feature-extraction", "feature-extraction-with-past", "text2text-generation", "text2text-generation-with-past"], +) +class LongT5OnnxConfig(T5OnnxConfig): + pass + + +@register_tasks_manager_onnx( + "m2m_100", + *["feature-extraction", "feature-extraction-with-past", "text2text-generation", "text2text-generation-with-past"], +) +class M2M100OnnxConfig(TextSeq2SeqOnnxConfig): + PAD_ATTENTION_MASK_TO_PAST = True + + NORMALIZED_CONFIG_CLASS = NormalizedSeq2SeqConfig.with_args( + encoder_num_layers="encoder_layers", + decoder_num_layers="decoder_layers", + num_layers="decoder_layers", # Used for the text-generation task past key values input generation. + encoder_num_attention_heads="encoder_attention_heads", + decoder_num_attention_heads="decoder_attention_heads", + eos_token_id="eos_token_id", + ) + DUMMY_INPUT_GENERATOR_CLASSES = ( + BartDummyTextInputGenerator, + { + "feature-extraction": DummySeq2SeqDecoderTextInputGenerator, + "text-generation": DummyDecoderTextInputGenerator, + }, + { + "feature-extraction": DummySeq2SeqPastKeyValuesGenerator, + "text-generation": DummyPastKeyValuesGenerator, + }, + ) + + def _create_dummy_input_generator_classes(self, **kwargs) -> list[DummyInputGenerator]: + dummy_text_input_generator = self.DUMMY_INPUT_GENERATOR_CLASSES[0]( + self.task, self._normalized_config, **kwargs + ) + task = "feature-extraction" if self.task != "text-generation" else "text-generation" + dummy_decoder_text_input_generator = self.DUMMY_INPUT_GENERATOR_CLASSES[1][task]( + self.task, self._normalized_config, **kwargs + ) + if self.task != "text-generation": + kwargs["encoder_sequence_length"] = dummy_text_input_generator.sequence_length + + dummy_seq2seq_past_key_values_generator = self.DUMMY_INPUT_GENERATOR_CLASSES[2][task]( + self.task, self._normalized_config, **kwargs + ) + dummy_inputs_generators = [ + dummy_text_input_generator, + dummy_decoder_text_input_generator, + dummy_seq2seq_past_key_values_generator, + ] + + return dummy_inputs_generators + + @property + def inputs_for_default_and_seq2seq_lm(self): + return super().inputs + + @property + def inputs_for_causal_lm(self): + if self.use_past_in_inputs: + common_inputs = { + "input_ids": {0: "batch_size", 1: "sequence_length"}, + "attention_mask": {0: "batch_size", 1: "past_sequence_length + sequence_length"}, + } + for i in range(self._normalized_config.decoder_num_layers): + common_inputs[f"past_key_values.{i}.key"] = { + 0: "batch_size", + 2: "past_sequence_length", + } + common_inputs[f"past_key_values.{i}.value"] = { + 0: "batch_size", + 2: "past_sequence_length", + } + else: + common_inputs = { + "input_ids": {0: "batch_size", 1: "sequence_length"}, + "attention_mask": {0: "batch_size", 1: "sequence_length"}, + } + + return common_inputs + + @property + def inputs_for_other_tasks(self): + return { + "input_ids": {0: "batch_size", 1: "sequence_length"}, + "attention_mask": {0: "batch_size", 1: "sequence_length"}, + } + + @property + def inputs(self) -> dict[str, dict[int, str]]: + inputs_properties = { + "feature-extraction": self.inputs_for_default_and_seq2seq_lm, + "text2text-generation": self.inputs_for_default_and_seq2seq_lm, + "text-generation": self.inputs_for_causal_lm, + "other": self.inputs_for_other_tasks, + } + return inputs_properties.get(self.task, inputs_properties["other"]) + + @property + def outputs(self) -> dict[str, dict[int, str]]: + if self.task in ["feature-extraction", "text2text-generation"]: + common_outputs = super().outputs + else: + common_outputs = super(OnnxConfigWithPast, self).outputs + if self.use_past: + # When exporting decoder models with use_cache=True, both the decoder without past and with past have the KV cache as an output. + for i in range( + self._normalized_config.encoder_num_layers + if self.task != "text-generation" + else self._normalized_config.decoder_num_layers + ): + common_outputs[f"present.{i}.key"] = {0: "batch_size", 2: "past_sequence_length + sequence_length"} + common_outputs[f"present.{i}.value"] = { + 0: "batch_size", + 2: "past_sequence_length + sequence_length", + } + return common_outputs + + def flatten_past_key_values(self, flattened_output, name, idx, t): + if self.task in ["feature-extraction", "text2text-generation"]: + flattened_output = super().flatten_past_key_values(flattened_output, name, idx, t) + else: + flattened_output = super(OnnxSeq2SeqConfigWithPast, self).flatten_past_key_values( + flattened_output, name, idx, t + ) + + +@register_tasks_manager_onnx( + "bart", *[*COMMON_TEXT2TEXT_GENERATION_TASKS, "text-classification", "question-answering"] +) +class BartOnnxConfig(M2M100OnnxConfig): + pass + + +@register_tasks_manager_onnx( + "mbart", *[*COMMON_TEXT2TEXT_GENERATION_TASKS, "text-classification", "question-answering"] +) +class MBartOnnxConfig(BartOnnxConfig): + pass + + +@register_tasks_manager_onnx("blenderbot", *COMMON_TEXT2TEXT_GENERATION_TASKS) +class BlenderbotOnnxConfig(BartOnnxConfig): + pass + + +@register_tasks_manager_onnx("blenderbot-small", *COMMON_TEXT2TEXT_GENERATION_TASKS) +class BlenderbotSmallOnnxConfig(BartOnnxConfig): + pass + + +@register_tasks_manager_onnx("big_bird", *COMMON_TEXT_TASKS) +class BigBirdOnnxConfig(DistilBertOnnxConfig): + pass + + +@register_tasks_manager_onnx( + "bigbird_pegasus", *[*COMMON_TEXT2TEXT_GENERATION_TASKS, "text-classification", "question-answering"] +) +class BigBirdPegasusOnnxConfig(BartOnnxConfig): + _MODEL_PATCHER = BigBirdPegasusModelPatcher + + +@register_tasks_manager_onnx("pegasus", *COMMON_TEXT2TEXT_GENERATION_TASKS) +class PegasusOnnxConfig(BartOnnxConfig): + pass + + +@register_tasks_manager_onnx("marian", *COMMON_TEXT2TEXT_GENERATION_TASKS) +class MarianOnnxConfig(BartOnnxConfig): + pass + + +@register_tasks_manager_onnx("vit", *["feature-extraction", "image-classification", "masked-im"]) +class ViTOnnxConfig(VisionOnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedVisionConfig + _MODEL_PATCHER = ViTForImageClassificationPatcher + + @property + def inputs(self) -> dict[str, dict[int, str]]: + return {"pixel_values": {0: "batch_size", 2: "height", 3: "width"}} + + @property + def outputs(self) -> dict[str, dict[int, str]]: + common_outputs = super().outputs + + if self.task == "feature-extraction": + common_outputs["last_hidden_state"] = {0: "batch_size"} + + return common_outputs + + +@register_tasks_manager_onnx("vitpose", *["keypoint-detection"]) +class VitPoseOnnxConfig(ViTOnnxConfig): + DUMMY_INPUT_GENERATOR_CLASSES = (VitPoseDummyInputGenerator,) + + _MODEL_PATCHER = VitPoseModelPatcher + + @property + def inputs(self) -> dict[str, dict[int, str]]: + return {"pixel_values": {0: "batch_size"}} + + +@register_tasks_manager_onnx("cvt", *["feature-extraction", "image-classification"]) +class CvTOnnxConfig(ViTOnnxConfig): + pass + + +@register_tasks_manager_onnx("levit", *["feature-extraction", "image-classification"]) +class LevitOnnxConfig(ViTOnnxConfig): + pass + + +@register_tasks_manager_onnx("deit", *["feature-extraction", "image-classification", "masked-im"]) +class DeiTOnnxConfig(ViTOnnxConfig): + pass + + +@register_tasks_manager_onnx("beit", *["feature-extraction", "image-classification"]) +class BeitOnnxConfig(ViTOnnxConfig): + pass + + +@register_tasks_manager_onnx("convnext", *["feature-extraction", "image-classification"]) +class ConvNextOnnxConfig(ViTOnnxConfig): + pass + + +@register_tasks_manager_onnx("convnextv2", *["feature-extraction", "image-classification"]) +class ConvNextV2OnnxConfig(ViTOnnxConfig): + pass + + +@register_tasks_manager_onnx("hiera", *["feature-extraction", "image-classification"]) +class HieraOnnxConfig(ViTOnnxConfig): + pass + + +@register_tasks_manager_onnx("pvt", *["feature-extraction", "image-classification"]) +class PvtOnnxConfig(ViTOnnxConfig): + pass + + +@register_tasks_manager_onnx("vit_mae", *["feature-extraction"]) +class VitMAEOnnxConfig(ViTOnnxConfig): + pass + + +@register_tasks_manager_onnx("vit_msn", *["feature-extraction", "image-classification"]) +class VitMSNOnnxConfig(ViTOnnxConfig): + pass + + +@register_tasks_manager_onnx("dinov2", *["feature-extraction", "image-classification"]) +class Dinov2OnnxConfig(ViTOnnxConfig): + DUMMY_INPUT_GENERATOR_CLASSES = (Dinov2DummyInputGenerator,) + + +@register_tasks_manager_onnx("mobilevit", *["feature-extraction", "image-classification", "image-segmentation"]) +class MobileViTOnnxConfig(ViTOnnxConfig): + pass + + +@register_tasks_manager_onnx("regnet", *["feature-extraction", "image-classification"]) +class RegNetOnnxConfig(ViTOnnxConfig): + pass + + +@register_tasks_manager_onnx("resnet", *["feature-extraction", "image-classification"]) +class ResNetOnnxConfig(ViTOnnxConfig): + pass + + +@register_tasks_manager_onnx("detr", *["feature-extraction", "object-detection", "image-segmentation"]) +class DetrOnnxConfig(ViTOnnxConfig): + @property + def outputs(self) -> dict[str, dict[int, str]]: + if self.task == "image-segmentation": + return { + "logits": {0: "batch_size", 1: "num_queries"}, + "pred_masks": {0: "batch_size", 1: "num_queries"}, + } + else: + return super().outputs + + +@register_tasks_manager_onnx("table-transformer", *["feature-extraction", "object-detection"]) +class TableTransformerOnnxConfig(DetrOnnxConfig): + pass + + +@register_tasks_manager_onnx("yolos", *["feature-extraction", "object-detection"]) +class YolosOnnxConfig(ViTOnnxConfig): + pass + + +@register_tasks_manager_onnx("swin", *["feature-extraction", "image-classification", "masked-im"]) +class SwinOnnxConfig(ViTOnnxConfig): + pass + + +@register_tasks_manager_onnx("swinv2", *["feature-extraction", "image-classification", "masked-im"]) +class SwinV2OnnxConfig(SwinOnnxConfig): + pass + + +@register_tasks_manager_onnx("swin2sr", *["feature-extraction", "image-to-image"]) +class Swin2srOnnxConfig(SwinOnnxConfig): + @property + def outputs(self) -> dict[str, dict[int, str]]: + outputs = super().outputs + + if self.task == "image-to-image": + scale_factor = self._config.upscale + outputs["reconstruction"] = { + 0: "batch_size", + 1: "num_channels", + 2: f"height * {scale_factor}", + 3: f"width * {scale_factor}", + } + + return outputs + + +@register_tasks_manager_onnx( + "dpt", *["feature-extraction", "depth-estimation", "image-segmentation", "semantic-segmentation"] +) +class DptOnnxConfig(ViTOnnxConfig): + pass + + +@register_tasks_manager_onnx("glpn", *["feature-extraction", "depth-estimation"]) +class GlpnOnnxConfig(ViTOnnxConfig): + pass + + +@register_tasks_manager_onnx("poolformer", *["feature-extraction", "image-classification"]) +class PoolFormerOnnxConfig(ViTOnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedVisionConfig + + +@register_tasks_manager_onnx( + "segformer", *["feature-extraction", "image-classification", "image-segmentation", "semantic-segmentation"] +) +class SegformerOnnxConfig(YolosOnnxConfig): + @property + def outputs(self) -> dict[str, dict[int, str]]: + outputs = super().outputs + + if self.task == "image-segmentation": + outputs["logits"] = {0: "batch_size"} + + return outputs + + +@register_tasks_manager_onnx("mobilenet_v1", *["feature-extraction", "image-classification"]) +class MobileNetV1OnnxConfig(ViTOnnxConfig): + @property + def inputs(self) -> dict[str, dict[int, str]]: + return {"pixel_values": {0: "batch_size"}} + + +@register_tasks_manager_onnx("mobilenet_v2", *["feature-extraction", "image-classification"]) +class MobileNetV2OnnxConfig(MobileNetV1OnnxConfig): + pass + + +@register_tasks_manager_onnx("maskformer", *["feature-extraction", "image-segmentation"]) +class MaskFormerOnnxConfig(ViTOnnxConfig): + @property + def outputs(self) -> dict[str, dict[int, str]]: + if self.task == "image-segmentation": + return { + "class_queries_logits": {0: "batch_size", 1: "num_queries"}, + "masks_queries_logits": {0: "batch_size", 1: "num_queries", 2: "height", 3: "width"}, + } + else: + return super().outputs + + @property + def torch_to_onnx_output_map(self) -> dict[str, str]: + return { + "transformer_decoder_last_hidden_state": "last_hidden_state", + } + + +@register_tasks_manager_onnx("donut-swin", *["feature-extraction"]) +class DonutSwinOnnxConfig(ViTOnnxConfig): + pass + + +@register_tasks_manager_onnx("default-timm-config", *["image-classification"], library_name="timm") +class TimmDefaultOnnxConfig(ViTOnnxConfig): + def rename_ambiguous_inputs(self, inputs): + # The input name in the model signature is `x, hence the export input name is updated. + model_inputs = {} + model_inputs["x"] = inputs["pixel_values"] + + return model_inputs + + @property + def torch_to_onnx_input_map(self) -> dict[str, str]: + return {"x": "pixel_values"} + + +@register_tasks_manager_onnx("mgp-str", *["feature-extraction"]) +class MgpstrOnnxConfig(ViTOnnxConfig): + _MODEL_PATCHER = MgpstrModelPatcher + + @property + def outputs(self) -> dict[str, dict[int, str]]: + return { + "char_logits": {0: "batch_size"}, + "bpe_logits": {0: "batch_size"}, + "wp_logits": {0: "batch_size"}, + } + + +@register_tasks_manager_onnx("efficientnet", *["feature-extraction", "image-classification"]) +class EfficientNetOnnxConfig(ViTOnnxConfig): + @property + def outputs(self) -> dict[str, dict[int, str]]: + common_outputs = super().outputs + + if self.task == "image-classification": + common_outputs["logits"] = {0: "batch_size", 1: "num_classes"} + + return common_outputs + + +@register_tasks_manager_onnx( + "transformer", *["feature-extraction", "sentence-similarity"], library_name="sentence_transformers" +) +class SentenceTransformersTransformerOnnxConfig(TextEncoderOnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig + _MODEL_PATCHER = SentenceTransformersTransformerPatcher + + @property + def inputs(self) -> dict[str, dict[int, str]]: + return { + "input_ids": {0: "batch_size", 1: "sequence_length"}, + "attention_mask": {0: "batch_size", 1: "sequence_length"}, + } + + @property + def outputs(self) -> dict[str, dict[int, str]]: + return { + "token_embeddings": {0: "batch_size", 1: "sequence_length"}, + "sentence_embedding": {0: "batch_size"}, + } + + +class CLIPNormalizedConfig(NormalizedTextAndVisionConfig): + TEXT_CONFIG = "text_config" + VISION_CONFIG = "vision_config" + + +@register_tasks_manager_onnx("clip_vision_model", *["feature-extraction"]) +class CLIPVisionModelOnnxConfig(VisionOnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedVisionConfig + _MODEL_PATCHER = CLIPModelPatcher + + @property + def inputs(self) -> dict[str, dict[int, str]]: + return {"pixel_values": {0: "batch_size", 1: "num_channels", 2: "height", 3: "width"}} + + @property + def outputs(self) -> dict[str, dict[int, str]]: + common_outputs = super().outputs + common_outputs["last_hidden_state"] = {0: "batch_size"} + common_outputs["pooler_output"] = {0: "batch_size"} + + return common_outputs + + +@register_tasks_manager_onnx("clip", *["feature-extraction", "zero-shot-image-classification", "image-classification"]) +class CLIPOnnxConfig(TextAndVisionOnnxConfig): + NORMALIZED_CONFIG_CLASS = CLIPNormalizedConfig + _MODEL_PATCHER = CLIPModelPatcher + + @property + def inputs(self) -> dict[str, dict[int, str]]: + inputs = {"pixel_values": {0: "batch_size", 1: "num_channels", 2: "height", 3: "width"}} + + if self.task in ["feature-extraction", "zero-shot-image-classification"]: + inputs.update( + { + "input_ids": {0: "text_batch_size", 1: "sequence_length"}, + "attention_mask": {0: "text_batch_size", 1: "sequence_length"}, + } + ) + + return inputs + + @property + def outputs(self) -> dict[str, dict[int, str]]: + if self.task in ["feature-extraction", "zero-shot-image-classification"]: + return { + "logits_per_image": {0: "image_batch_size", 1: "text_batch_size"}, + "logits_per_text": {0: "text_batch_size", 1: "image_batch_size"}, + "text_embeds": {0: "text_batch_size"}, + "image_embeds": {0: "image_batch_size"}, + } + else: + return super().outputs + + +@register_tasks_manager_onnx( + "clip", *["feature-extraction", "sentence-similarity"], library_name="sentence_transformers" +) +class SentenceTransformersCLIPOnnxConfig(CLIPOnnxConfig): + _MODEL_PATCHER = SentenceTransformersCLIPPatcher + + @property + def outputs(self) -> dict[str, dict[int, str]]: + return { + "text_embeds": {0: "text_batch_size"}, + "image_embeds": {0: "image_batch_size"}, + } + + +@register_tasks_manager_onnx("clip-text-with-projection", *["feature-extraction"], library_name="diffusers") +class CLIPTextWithProjectionOnnxConfig(TextEncoderOnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args( + vocab_size="vocab_size", + sequence_length="max_position_embeddings", + num_layers="num_hidden_layers", + allow_new=True, + ) + _MODEL_PATCHER = CLIPModelPatcher + + @property + def inputs(self) -> dict[str, dict[int, str]]: + return { + "input_ids": {0: "batch_size", 1: "sequence_length"}, + } + + @property + def outputs(self) -> dict[str, dict[int, str]]: + common_outputs = { + "text_embeds": {0: "batch_size", 1: "sequence_length"}, + "last_hidden_state": {0: "batch_size", 1: "sequence_length"}, + } + if self._normalized_config.output_hidden_states: + for i in range(self._normalized_config.num_layers + 1): + common_outputs[f"hidden_states.{i}"] = {0: "batch_size", 1: "sequence_length"} + + return common_outputs + + +@register_tasks_manager_onnx("clip-text", *["feature-extraction"], library_name="diffusers") +class CLIPTextOnnxConfig(CLIPTextWithProjectionOnnxConfig): + _MODEL_PATCHER = CLIPModelPatcher + + @property + def outputs(self) -> dict[str, dict[int, str]]: + common_outputs = { + "last_hidden_state": {0: "batch_size", 1: "sequence_length"}, + "pooler_output": {0: "batch_size"}, + } + + if self._normalized_config.output_hidden_states: + for i in range(self._normalized_config.num_layers + 1): + common_outputs[f"hidden_states.{i}"] = {0: "batch_size", 1: "sequence_length"} + + return common_outputs + + +@register_tasks_manager_onnx( + "metaclip_2", + *["feature-extraction", "zero-shot-image-classification", "image-classification"], + library_name="transformers", +) +class MetaCLIP2OnnxConfig(TextAndVisionOnnxConfig): + NORMALIZED_CONFIG_CLASS = CLIPNormalizedConfig + MIN_TRANSFORMERS_VERSION = version.parse("4.56.2") + VARIANTS = { # noqa: RUF012 + "monolith": "All the MetaClip2 model components are exported as a single model.onnx.", + "split": "The vision model is exported as a separate vision_model.onnx, and the text_model is exported as text_model.onnx", + } + DEFAULT_VARIANT = "monolith" + _MODEL_PATCHER = MetaCLIP2Patcher + + def __init__( + self, + config: PretrainedConfig, + task: str = "feature-extraction", + int_dtype: str = "int64", + float_dtype: str = "fp32", + variant: str = "monolith", + vision_model: bool | None = None, + preprocessors: list[Any] | None = None, + ): + super().__init__( + config=config, + task=task, + int_dtype=int_dtype, + float_dtype=float_dtype, + preprocessors=preprocessors, + ) + self.variant = variant + self.vision_model = vision_model + + @property + def inputs(self) -> dict[str, dict[int, str]]: + if self.variant == "monolith": + inputs = {"pixel_values": {0: "batch_size", 1: "num_channels", 2: "height", 3: "width"}} + if self.task in ["feature-extraction", "zero-shot-image-classification"]: + inputs.update( + { + "input_ids": {0: "text_batch_size", 1: "sequence_length"}, + "attention_mask": {0: "text_batch_size", 1: "sequence_length"}, + } + ) + else: + if self.vision_model: + inputs = {"pixel_values": {0: "batch_size", 1: "num_channels", 2: "height", 3: "width"}} + else: + inputs = { + "input_ids": {0: "text_batch_size", 1: "sequence_length"}, + "attention_mask": {0: "text_batch_size", 1: "sequence_length"}, + } + return inputs + + @property + def outputs(self) -> dict[str, dict[int, str]]: + if self.variant == "split": + if self.vision_model: + return { + "image_embeds": {0: "batch_size"}, + } + else: + return { + "text_embeds": {0: "batch_size"}, + } + else: + if self.task in ["feature-extraction", "zero-shot-image-classification"]: + return { + "logits_per_image": {0: "image_batch_size", 1: "text_batch_size"}, + "logits_per_text": {0: "text_batch_size", 1: "image_batch_size"}, + "text_embeds": {0: "text_batch_size"}, + "image_embeds": {0: "image_batch_size"}, + } + else: + return super().outputs + + +class SiglipNormalizedConfig(CLIPNormalizedConfig): + pass + + +@register_tasks_manager_onnx("chinese_clip", *["feature-extraction", "zero-shot-image-classification"]) +class ChineseCLIPOnnxConfig(CLIPOnnxConfig): + pass + + +@register_tasks_manager_onnx("siglip", *["feature-extraction", "zero-shot-image-classification"]) +class SiglipOnnxConfig(CLIPOnnxConfig): + NORMALIZED_CONFIG_CLASS = SiglipNormalizedConfig + + @property + def inputs(self) -> dict[str, dict[int, str]]: + return { + "input_ids": {0: "text_batch_size", 1: "sequence_length"}, + "pixel_values": {0: "image_batch_size", 1: "num_channels", 2: "height", 3: "width"}, + # NOTE: No attention_mask + } + + +@register_tasks_manager_onnx("siglip-text-with-projection", *["feature-extraction"]) +class SiglipTextWithProjectionOnnxConfig(CLIPTextWithProjectionOnnxConfig): + pass + + +@register_tasks_manager_onnx("siglip-text", *["feature-extraction"]) +class SiglipTextOnnxConfig(CLIPTextOnnxConfig): + pass + + +@register_tasks_manager_onnx("siglip_vision_model", *["feature-extraction"]) +class SiglipVisionModelOnnxConfig(CLIPVisionModelOnnxConfig): + pass + + +@register_tasks_manager_onnx("unet-2d-condition", *["semantic-segmentation"], library_name="diffusers") +class UNetOnnxConfig(VisionOnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args( + image_size="sample_size", + num_channels="in_channels", + hidden_size="cross_attention_dim", + vocab_size="norm_num_groups", + allow_new=True, + ) + + DUMMY_INPUT_GENERATOR_CLASSES = ( + DummyVisionInputGenerator, + DummyTimestepInputGenerator, + DummySeq2SeqDecoderTextInputGenerator, + ) + + @property + def inputs(self) -> dict[str, dict[int, str]]: + common_inputs = { + "sample": {0: "batch_size", 2: "height", 3: "width"}, + "timestep": {}, # a scalar with no dimension + "encoder_hidden_states": {0: "batch_size", 1: "sequence_length"}, + } + + # TODO : add addition_embed_type == text_image, image and image_embeds + # https://github.com/huggingface/diffusers/blob/9366c8f84bfe47099ff047272661786ebb54721d/src/diffusers/models/unets/unet_2d_condition.py#L671 + if getattr(self._normalized_config, "addition_embed_type", None) == "text_time": + common_inputs["text_embeds"] = {0: "batch_size"} + common_inputs["time_ids"] = {0: "batch_size"} + + if getattr(self._normalized_config, "time_cond_proj_dim", None) is not None: + common_inputs["timestep_cond"] = {0: "batch_size"} + + return common_inputs + + @property + def outputs(self) -> dict[str, dict[int, str]]: + return { + "out_sample": {0: "batch_size", 2: "height", 3: "width"}, + } + + @property + def torch_to_onnx_output_map(self) -> dict[str, str]: + return { + "sample": "out_sample", + } + + def generate_dummy_inputs(self, framework: str = "pt", **kwargs): + dummy_inputs = super().generate_dummy_inputs(framework=framework, **kwargs) + dummy_inputs["encoder_hidden_states"] = dummy_inputs["encoder_hidden_states"][0] + + if getattr(self._normalized_config, "addition_embed_type", None) == "text_time": + dummy_inputs["added_cond_kwargs"] = { + "text_embeds": dummy_inputs.pop("text_embeds"), + "time_ids": dummy_inputs.pop("time_ids"), + } + + return dummy_inputs + + def ordered_inputs(self, model) -> dict[str, dict[int, str]]: + inputs = super().ordered_inputs(model=model) + # to fix mismatch between model forward signature and expected inputs + # a dictionary of additional embeddings `added_cond_kwargs` is expected depending on config.addition_embed_type + if getattr(self._normalized_config, "addition_embed_type", None) == "text_time": + inputs["text_embeds"] = self.inputs["text_embeds"] + inputs["time_ids"] = self.inputs["time_ids"] + + return inputs + + +@register_tasks_manager_onnx("vae-encoder", *["semantic-segmentation"], library_name="diffusers") +class VaeEncoderOnnxConfig(VisionOnnxConfig): + ATOL_FOR_VALIDATION = 3e-4 # TODO: this only happens in test_export.py + NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args( + num_channels="in_channels", image_size="sample_size", allow_new=True + ) + + @property + def inputs(self) -> dict[str, dict[int, str]]: + return { + "sample": {0: "batch_size", 2: "sample_height", 3: "sample_width"}, + } + + @property + def outputs(self) -> dict[str, dict[int, str]]: + down_sampling_factor = 2 ** (len(self._normalized_config.down_block_types) - 1) + + return { + "latent_parameters": { + 0: "batch_size", + 2: f"sample_height / {down_sampling_factor}", + 3: f"sample_width / {down_sampling_factor}", + }, + } + + +@register_tasks_manager_onnx("vae-decoder", *["semantic-segmentation"], library_name="diffusers") +class VaeDecoderOnnxConfig(VisionOnnxConfig): + ATOL_FOR_VALIDATION = 3e-4 # TODO: this only happens in test_export.py + NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args(num_channels="latent_channels", allow_new=True) + + @property + def inputs(self) -> dict[str, dict[int, str]]: + return { + "latent_sample": {0: "batch_size", 2: "latent_height", 3: "latent_width"}, + } + + @property + def outputs(self) -> dict[str, dict[int, str]]: + up_sampling_factor = 2 ** (len(self._normalized_config.up_block_types) - 1) + + return { + "sample": { + 0: "batch_size", + 2: f"latent_height * {up_sampling_factor}", + 3: f"latent_width * {up_sampling_factor}", + }, + } + + +@register_tasks_manager_onnx("t5-encoder", *["feature-extraction"], library_name="diffusers") +class T5EncoderOnnxConfig(TextEncoderOnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig + + @property + def inputs(self): + return { + "input_ids": {0: "batch_size", 1: "sequence_length"}, + } + + @property + def outputs(self): + return { + "last_hidden_state": {0: "batch_size", 1: "sequence_length"}, + } + + +@register_tasks_manager_onnx("sd3-transformer-2d", *["semantic-segmentation"], library_name="diffusers") +class SD3TransformerOnnxConfig(VisionOnnxConfig): + DUMMY_INPUT_GENERATOR_CLASSES = ( + DummyTransformerTimestepInputGenerator, + DummyTransformerVisionInputGenerator, + DummyTransformerTextInputGenerator, + ) + + NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args( + image_size="sample_size", + num_channels="in_channels", + vocab_size="attention_head_dim", + hidden_size="joint_attention_dim", + projection_size="pooled_projection_dim", + allow_new=True, + ) + + @property + def inputs(self) -> dict[str, dict[int, str]]: + common_inputs = { + "hidden_states": {0: "batch_size", 2: "height", 3: "width"}, + "encoder_hidden_states": {0: "batch_size", 1: "sequence_length"}, + "pooled_projections": {0: "batch_size"}, + "timestep": {0: "step"}, + } + + return common_inputs + + @property + def outputs(self) -> dict[str, dict[int, str]]: + return { + "out_hidden_states": {0: "batch_size", 2: "height", 3: "width"}, + } + + @property + def torch_to_onnx_output_map(self) -> dict[str, str]: + return { + "sample": "out_hidden_states", + } + + +@register_tasks_manager_onnx("flux-transformer-2d", *["semantic-segmentation"], library_name="diffusers") +class FluxTransformerOnnxConfig(SD3TransformerOnnxConfig): + DUMMY_INPUT_GENERATOR_CLASSES = ( + DummyTransformerTimestepInputGenerator, + DummyFluxTransformerVisionInputGenerator, + DummyFluxTransformerTextInputGenerator, + ) + _MODEL_PATCHER = FluxTransformerModelPatcher + + @property + def inputs(self): + common_inputs = super().inputs + common_inputs["hidden_states"] = {0: "batch_size", 1: "packed_height_width"} + common_inputs["txt_ids"] = ( + {0: "sequence_length"} if is_diffusers_version(">=", "0.31.0") else {0: "batch_size", 1: "sequence_length"} + ) + common_inputs["img_ids"] = ( + {0: "packed_height_width"} + if is_diffusers_version(">=", "0.31.0") + else {0: "batch_size", 1: "packed_height_width"} + ) + + if getattr(self._normalized_config, "guidance_embeds", False): + common_inputs["guidance"] = {0: "batch_size"} + + return common_inputs + + @property + def outputs(self): + return { + "out_hidden_states": {0: "batch_size", 1: "packed_height_width"}, + } + + +@register_tasks_manager_onnx("groupvit", *["feature-extraction"]) +class GroupViTOnnxConfig(CLIPOnnxConfig): + pass + + +@register_tasks_manager_onnx("owlvit", *["feature-extraction", "zero-shot-object-detection"]) +class OwlViTOnnxConfig(CLIPOnnxConfig): + def __init__( + self, + config: PretrainedConfig, + task: str = "feature-extraction", + int_dtype: str = "int64", + float_dtype: str = "fp32", + preprocessors: list[Any] | None = None, + ): + super().__init__( + config=config, + task=task, + int_dtype=int_dtype, + float_dtype=float_dtype, + preprocessors=preprocessors, + ) + if task == "zero-shot-object-detection": + logger.warning( + "The batch size of this model will not be dynamic because non-maximum suppression is performed. " + "Make sure to export the model with the same batch size as the one you will use at inference " + "with `--batch_size N`." + ) + + @property + def inputs(self) -> dict[str, dict[int, str]]: + inputs = {"pixel_values": {0: "batch_size", 1: "num_channels", 2: "height", 3: "width"}} + + if self.task in ["feature-extraction", "zero-shot-object-detection"]: + inputs.update( + { + "input_ids": {0: "text_batch_size", 1: "sequence_length"}, + "attention_mask": {0: "text_batch_size", 1: "sequence_length"}, + } + ) + + return inputs + + @property + def outputs(self) -> dict[str, dict[int, str]]: + outputs = {} + if self.task == "feature-extraction": + outputs["logits_per_image"] = {0: "image_batch_size", 1: "text_batch_size"} + outputs["logits_per_text"] = {0: "text_batch_size", 1: "image_batch_size"} + elif self.task == "zero-shot-object-detection": + outputs["logits"] = {0: "image_batch_size", 2: "num_queries"} + outputs["pred_boxes"] = {0: "image_batch_size", 1: "num_boxes"} + + outputs["text_embeds"] = {0: "text_batch_size", 1: "max_text_queries"} + outputs["image_embeds"] = {0: "image_batch_size"} + return outputs + + +@register_tasks_manager_onnx("owlv2", *["feature-extraction", "zero-shot-object-detection"]) +class OwlV2OnnxConfig(OwlViTOnnxConfig): + MIN_TRANSFORMERS_VERSION = version.parse("4.35.0") + + +@register_tasks_manager_onnx( + "layoutlm", *["feature-extraction", "fill-mask", "text-classification", "token-classification"] +) +class LayoutLMOnnxConfig(TextAndVisionOnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args( + allow_new=True, MAX_2D_POSITION_EMBEDDINGS="max_2d_position_embeddings" + ) + + @property + def inputs(self) -> dict[str, dict[int, str]]: + return { + "input_ids": {0: "batch_size", 1: "sequence_length"}, + "bbox": {0: "batch_size", 1: "sequence_length"}, + "attention_mask": {0: "batch_size", 1: "sequence_length"}, + "token_type_ids": {0: "batch_size", 1: "sequence_length"}, + } + + +@register_tasks_manager_onnx( + "layoutlmv3", *["feature-extraction", "question-answering", "text-classification", "token-classification"] +) +class LayoutLMv3OnnxConfig(TextAndVisionOnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args( + allow_new=True, MAX_2D_POSITION_EMBEDDINGS="max_2d_position_embeddings", image_size="input_size" + ) + + @property + def inputs(self) -> dict[str, dict[int, str]]: + if self.task in ["text-classification", "question-answering"]: + pixel_values_dynamic_axes = {0: "batch_size", 1: "num_channels", 2: "height", 3: "width"} + else: + pixel_values_dynamic_axes = {0: "batch_size", 1: "num_channels"} + return { + "input_ids": {0: "batch_size", 1: "sequence_length"}, + "attention_mask": {0: "batch_size", 1: "sequence_length"}, + "bbox": {0: "batch_size", 1: "sequence_length"}, + "pixel_values": pixel_values_dynamic_axes, + } + + +@register_tasks_manager_onnx( + "lilt", *["feature-extraction", "question-answering", "text-classification", "token-classification"] +) +class LiltOnnxConfig(TextAndVisionOnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args( + allow_new=True, + MAX_2D_POSITION_EMBEDDINGS="max_2d_position_embeddings", + ) + + @property + def inputs(self) -> dict[str, dict[int, str]]: + return { + "input_ids": {0: "batch_size", 1: "sequence_length"}, + "bbox": {0: "batch_size", 1: "sequence_length"}, + "attention_mask": {0: "batch_size", 1: "sequence_length"}, + } + + +@register_tasks_manager_onnx("data2vec-text", *COMMON_TEXT_TASKS) +class Data2VecTextOnnxConfig(DistilBertOnnxConfig): + pass + + +@register_tasks_manager_onnx("data2vec-vision", *["feature-extraction", "image-classification"]) +class Data2VecVisionOnnxConfig(ViTOnnxConfig): + pass + + +@register_tasks_manager_onnx("perceiver", *["fill-mask", "text-classification", "image-classification"]) +class PerceiverOnnxConfig(TextAndVisionOnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig + DUMMY_INPUT_GENERATOR_CLASSES = ( + PerceiverDummyInputGenerator, + *TextAndVisionOnnxConfig.DUMMY_INPUT_GENERATOR_CLASSES, + ) + + def __init__( + self, + config: PretrainedConfig, + task: str = "feature-extraction", + int_dtype: str = "int64", + float_dtype: str = "fp32", + preprocessors: list[Any] | None = None, + ): + super().__init__( + config=config, + task=task, + int_dtype=int_dtype, + float_dtype=float_dtype, + preprocessors=preprocessors, + ) + self.is_generating_dummy_inputs = False + + @property + def inputs_name(self): + if self.is_generating_dummy_inputs: + if self.task in ["fill-mask", "text-classification"]: + return "input_ids" + else: + return "pixel_values" + else: + return "inputs" + + @property + def inputs(self) -> dict[str, dict[int, str]]: + if self.inputs_name in ["input_ids", "inputs"]: + dynamic_axis = {0: "batch_size", 1: "sequence_length"} + return { + "input_ids": dynamic_axis, + "attention_mask": dynamic_axis, + } + else: + dynamic_axis = {0: "batch_size", 1: "sequence_length", 2: "width", 3: "height"} + return { + "pixel_values": dynamic_axis, + } + + @property + def outputs(self) -> dict[str, dict[int, str]]: + outputs = super().outputs + + if "logits" in outputs: + # default is {0: "batch_size", 1: "sequence_length"} where sequence_length is dynamic axis + # but perceiver always return the same max sequence length in the second dimension + outputs["logits"] = {0: "batch_size"} + + return outputs + + def generate_dummy_inputs(self, framework: str = "pt", **kwargs): + self.is_generating_dummy_inputs = True + dummy_inputs = super().generate_dummy_inputs(framework=framework, **kwargs) + dummy_inputs[self.inputs_name] = dummy_inputs.pop(self.inputs_name) + return dummy_inputs + + +@register_tasks_manager_onnx("hubert", *["feature-extraction", "automatic-speech-recognition", "audio-classification"]) +class HubertOnnxConfig(AudioOnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedConfig + + @property + def outputs(self) -> dict[str, dict[int, str]]: + outputs = super().outputs + + # Hubert output formula adapted from: + # https://github.com/huggingface/transformers/blob/v4.55.2/src/transformers/models/hubert/modeling_hubert.py#L721 + if self.task == "automatic-speech-recognition": + sequence_length = "sequence_length" + for kernel_size, stride in zip(self._config.conv_kernel, self._config.conv_stride): + sequence_length = f"( {sequence_length} - {kernel_size} ) // {stride} + 1" + outputs["logits"] = {0: "batch_size", 1: sequence_length} + + return outputs + + +@register_tasks_manager_onnx( + "data2vec-audio", + *[ + "feature-extraction", + "automatic-speech-recognition", + "audio-classification", + "audio-frame-classification", + "audio-xvector", + ], +) +class Data2VecAudioOnnxConfig(HubertOnnxConfig): + pass + + +@register_tasks_manager_onnx( + "wav2vec2", + *[ + "feature-extraction", + "automatic-speech-recognition", + "audio-classification", + "audio-frame-classification", + "audio-xvector", + ], +) +class Wav2Vec2OnnxConfig(HubertOnnxConfig): + pass + + +@register_tasks_manager_onnx( + "wav2vec2-conformer", + *[ + "feature-extraction", + "automatic-speech-recognition", + "audio-classification", + "audio-frame-classification", + "audio-xvector", + ], +) +class Wav2Vec2ConformerOnnxConfig(HubertOnnxConfig): + pass + + +@register_tasks_manager_onnx("sew", *["feature-extraction", "automatic-speech-recognition", "audio-classification"]) +class SEWOnnxConfig(HubertOnnxConfig): + pass + + +@register_tasks_manager_onnx("sew-d", *["feature-extraction", "automatic-speech-recognition", "audio-classification"]) +class SEWDOnnxConfig(HubertOnnxConfig): + pass + + +@register_tasks_manager_onnx( + "unispeech", *["feature-extraction", "automatic-speech-recognition", "audio-classification"] +) +class UniSpeechOnnxConfig(HubertOnnxConfig): + pass + + +@register_tasks_manager_onnx( + "unispeech-sat", + *[ + "feature-extraction", + "automatic-speech-recognition", + "audio-classification", + "audio-frame-classification", + "audio-xvector", + ], +) +class UniSpeechSATOnnxConfig(HubertOnnxConfig): + pass + + +@register_tasks_manager_onnx( + "wavlm", + *[ + "feature-extraction", + "automatic-speech-recognition", + "audio-classification", + "audio-frame-classification", + "audio-xvector", + ], +) +class WavLMOnnxConfig(HubertOnnxConfig): + pass + + +@register_tasks_manager_onnx("mctct", *["feature-extraction", "automatic-speech-recognition"]) +class MCTCTOnnxConfig(AudioOnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args( + input_features_per_channel="input_feat_per_channel", allow_new=True + ) + DUMMY_INPUT_GENERATOR_CLASSES = (MCTCTDummyAudioInputGenerator,) + + @property + def inputs(self) -> dict[str, dict[int, str]]: + return {"input_features": {0: "batch_size", 1: "sequence_length"}} + + @property + def outputs(self) -> dict[str, dict[int, str]]: + outputs = super().outputs + + # mctct output formula adapted from: + # https://github.com/huggingface/transformers/blob/v4.53.3/src/transformers/models/deprecated/mctct/modeling_mctct.py#L455 + if self.task == "automatic-speech-recognition": + sequence_length = "sequence_length" + for kernel_size, stride in zip(self._config.conv_kernel, self._config.conv_stride): + dilation = 1 + padding = kernel_size // 2 + sequence_length = f"( {sequence_length} + 2 * {padding} - {dilation} * ({kernel_size} - 1) - 1 )" + sequence_length = f"( {sequence_length} // {stride} ) + 1" + outputs["logits"] = {0: "batch_size", 1: sequence_length} + + return outputs + + +@register_tasks_manager_onnx("audio-spectrogram-transformer", *["feature-extraction", "audio-classification"]) +class ASTOnnxConfig(OnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args( + num_mel_bins="num_mel_bins", max_length="max_length", allow_new=True + ) + DUMMY_INPUT_GENERATOR_CLASSES = (ASTDummyAudioInputGenerator,) + + @property + def inputs(self) -> dict[str, dict[int, str]]: + return {"input_values": {0: "batch_size"}} + + +@register_tasks_manager_onnx( + "moonshine", + *[ + "feature-extraction", + "feature-extraction-with-past", + "automatic-speech-recognition", + "automatic-speech-recognition-with-past", + ], +) +class MoonshineOnnxConfig(AudioToTextOnnxConfig): + MIN_TRANSFORMERS_VERSION = version.parse("4.48.0") + NORMALIZED_CONFIG_CLASS = NormalizedSeq2SeqConfig + _MODEL_PATCHER = MoonshineModelPatcher + DUMMY_INPUT_GENERATOR_CLASSES = ( + DummyMoonshineAudioInputGenerator, + DummySeq2SeqDecoderTextInputGenerator, + DummySeq2SeqPastKeyValuesGenerator, + ) + + @property + def inputs(self) -> dict[str, dict[int, str]]: + common_inputs = {} + if self._behavior in {ConfigBehavior.ENCODER, ConfigBehavior.MONOLITH}: + common_inputs["input_values"] = {0: "batch_size", 1: "encoder_sequence_length"} + else: + common_inputs["encoder_outputs"] = {0: "batch_size", 1: "encoder_sequence_length"} + common_inputs["attention_mask"] = {0: "batch_size", 1: "encoder_sequence_length"} + + if self._behavior in {ConfigBehavior.DECODER, ConfigBehavior.MONOLITH}: + common_inputs["decoder_input_ids"] = {0: "batch_size", 1: "decoder_sequence_length"} + if self.use_past_in_inputs: + self.add_past_key_values(common_inputs, direction="inputs") + + return common_inputs + + @property + def outputs(self) -> dict[str, dict[int, str]]: + common_outputs = super().outputs + if self._behavior in {ConfigBehavior.ENCODER, ConfigBehavior.MONOLITH}: + if self._behavior is ConfigBehavior.MONOLITH: + output_name = "encoder_last_hidden_state" + else: + output_name = "last_hidden_state" + # Moonshine encoder output formula adapted from: + # transformers.models.moonshine.modeling_moonshine.MoonshinePreTrainedModel._get_feat_extract_output_lengths + # output_conv1_length = int((input_lengths - 127) / 64 + 1) + # output_conv2_length = int((output_conv1_length - 7) / 3 + 1) + # output_conv3_length = int((output_conv2_length - 3) / 2 + 1) + output_sequence_length = "( ( ( encoder_sequence_length - 127 ) // 64 + 1 - 7 ) // 3 + 1 - 3 ) // 2 + 1" + common_outputs[output_name] = {0: "batch_size", 1: output_sequence_length} + return common_outputs + + +@register_tasks_manager_onnx( + "whisper", + *[ + "feature-extraction", + "feature-extraction-with-past", + "audio-classification", + "automatic-speech-recognition", + "automatic-speech-recognition-with-past", + ], +) +class WhisperOnnxConfig(AudioToTextOnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedSeq2SeqConfig.with_args( + encoder_num_layers="encoder_layers", + decoder_num_layers="decoder_layers", + feature_size="num_mel_bins", + allow_new=True, + ) + + @property + def inputs(self) -> dict[str, dict[int, str]]: + if self.task == "audio-classification": + return {"input_features": {0: "batch_size"}} + + common_inputs = super().inputs + if self._behavior in {ConfigBehavior.ENCODER, ConfigBehavior.MONOLITH}: + common_inputs["input_features"] = {0: "batch_size"} # Remove unnecessary dynamic axis. + else: + # the dynamic encoder sequence length is only needed here because the input generator generates + # encoder_outputs with a seq_len=16 but the model expects at inference time seq_len=1500 + # TODO: this can be fixed by generating the correct inputs in the input generator + common_inputs["encoder_outputs"] = {0: "batch_size", 1: "encoder_sequence_length"} + + if self._behavior in {ConfigBehavior.DECODER, ConfigBehavior.MONOLITH}: + if is_transformers_version(">=", "4.43.0") and is_transformers_version("<", "4.46.0"): + # since https://github.com/huggingface/transformers/pull/31166 + if self._behavior is not ConfigBehavior.ENCODER and self.use_past_in_inputs: + common_inputs["cache_position"] = {0: "decoder_sequence_length"} + + return common_inputs + + @property + def outputs(self) -> dict[str, dict[int, str]]: + if self.task == "audio-classification": + return {"logits": {0: "batch_size"}} + + common_outputs = super().outputs + if self._behavior in {ConfigBehavior.ENCODER, ConfigBehavior.MONOLITH}: + if self._behavior is ConfigBehavior.MONOLITH: + output_name = "encoder_last_hidden_state" + else: + output_name = "last_hidden_state" + common_outputs[output_name] = {0: "batch_size"} # Remove unnecessary dynamic axis. + return common_outputs + + +@register_tasks_manager_onnx("musicgen", *["text-to-audio"]) +class MusicgenOnnxConfig(OnnxSeq2SeqConfigWithPast): + # NOTE: Several warnings during the export are not to worry about: + # * for i, indices in enumerate(codes): --> can be unrolled, fixed length (num_quantizers). + # * max_pad = max(padding_left, padding_right) --> does not impact later controlflows. + # if length <= max_pad: --> appears to be always False for Musicgen. + + VARIANTS = { # noqa: RUF012 + "text-conditional-with-past": """Exports Musicgen to ONNX to generate audio samples conditioned on a text prompt (Reference: https://huggingface.co/docs/transformers/model_doc/musicgen#text-conditional-generation). + This uses the decoder KV cache. The following subcomponents are exported: + * text_encoder.onnx: corresponds to the text encoder part in https://github.com/huggingface/transformers/blob/v4.39.1/src/transformers/models/musicgen/modeling_musicgen.py#L1457. + * encodec_decode.onnx: corresponds to the Encodec audio encoder part in https://github.com/huggingface/transformers/blob/v4.39.1/src/transformers/models/musicgen/modeling_musicgen.py#L2472-L2480. + * decoder_model.onnx: The Musicgen decoder, without past key values input, and computing cross attention. Not required at inference (use decoder_model_merged.onnx instead). + * decoder_with_past_model.onnx: The Musicgen decoder, with past_key_values input (KV cache filled), not computing cross attention. Not required at inference (use decoder_model_merged.onnx instead). + * decoder_model_merged.onnx: The two previous models fused in one, to avoid duplicating weights. A boolean input `use_cache_branch` allows to select the branch to use. In the first forward pass where the KV cache is empty, dummy past key values inputs need to be passed and are ignored with use_cache_branch=False. + * build_delay_pattern_mask.onnx: A model taking as input `input_ids`, `pad_token_id`, `max_length`, and building a delayed pattern mask to the input_ids. Implements https://github.com/huggingface/transformers/blob/v4.39.3/src/transformers/models/musicgen/modeling_musicgen.py#L1054.""", + } + # TODO: support audio-prompted generation (audio_encoder_encode.onnx: corresponds to the audio encoder part + # in https://github.com/huggingface/transformers/blob/f01e1609bf4dba146d1347c1368c8c49df8636f6/src/transformers/models/musicgen/modeling_musicgen.py#L2087.) + # With that, we have full Encodec support. + DEFAULT_VARIANT = "text-conditional-with-past" + + NORMALIZED_CONFIG_CLASS = NormalizedEncoderDecoderConfig + + DUMMY_INPUT_GENERATOR_CLASSES = ( + DummyTextInputGenerator, + DummyCodegenDecoderTextInputGenerator, + DummySeq2SeqPastKeyValuesGenerator, + DummyEncodecInputGenerator, + DummyIntGenerator, + ) + DUMMY_PKV_GENERATOR_CLASS = DummySeq2SeqPastKeyValuesGenerator + _MODEL_PATCHER = MusicgenModelPatcher + + def __init__( + self, + config: PretrainedConfig, + task: str = "feature-extraction", + int_dtype: str = "int64", + float_dtype: str = "fp32", + use_past: bool = False, + use_past_in_inputs: bool = False, + behavior: ConfigBehavior = ConfigBehavior.ENCODER, + preprocessors: list[Any] | None = None, + model_part: Literal["text_encoder", "encodec_decode", "decoder", "build_delay_pattern_mask"] | None = None, + variant: str = "text-conditional-with-past", + ): + super().__init__( + config=config, + task=task, + int_dtype=int_dtype, + float_dtype=float_dtype, + use_past=use_past, + use_past_in_inputs=use_past_in_inputs, + behavior=behavior, + preprocessors=preprocessors, + ) + + if ( + model_part in ["text_encoder", "encodec_decode", "build_delay_pattern_mask"] + and behavior != ConfigBehavior.ENCODER + ): + raise ValueError( + f"model_part is {model_part} and behavior is {behavior}. This is not supported, please open an issue at https://github.com/huggingface/optimum/issues." + ) + + if model_part == "decoder" and behavior != ConfigBehavior.DECODER: + raise ValueError( + f"model_part is {model_part} and behavior is {behavior}. This is not supported, please open an issue at https://github.com/huggingface/optimum/issues." + ) + + if behavior == ConfigBehavior.MONOLITH: + raise ValueError( + "Musicgen does not support behavior=ConfigBehavior.MONOLITH. Please open an issue at https://github.com/huggingface/optimum/issues." + ) + + if config.audio_encoder.model_type != "encodec": + raise ValueError( + f"Optimum ONNX export for Musicgen supports only Encodec as the audio encoder, got: {config.audio_encoder.model_type}. Please open an issue at https://github.com/huggingface/optimum/issues." + ) + + # Handling it would require to trace the audio_encoder.decode with torch.jit.script as we than have an unrollable loop. + if config.audio_encoder.chunk_length_s is not None: + raise ValueError( + f"Musicgen ONNX export currently does not support audio_encoder.chunk_length_s not None (got {config.audio_encoder.chunk_length_s}). Please open an issue at https://github.com/huggingface/optimum/issues." + ) + + self.model_part = model_part + if self.model_part == "decoder": + self.use_past = True # without past is not supported, hard-code it here. + + self._normalized_config.ENCODER_NORMALIZED_CONFIG_CLASS = NormalizedTextConfig(self._config.text_encoder) + self._normalized_config.DECODER_NORMALIZED_CONFIG_CLASS = NormalizedConfig(self._config.decoder) + self._normalized_config.decoder_num_layers = self._config.decoder.num_hidden_layers + self._normalized_config.DECODER_NORMALIZED_CONFIG_CLASS.num_layers = self._config.decoder.num_hidden_layers + self._normalized_config.DECODER_NORMALIZED_CONFIG_CLASS.encoder_num_attention_heads = ( + self._config.decoder.num_attention_heads + ) + self._normalized_config.DECODER_NORMALIZED_CONFIG_CLASS.decoder_num_attention_heads = ( + self._config.decoder.num_attention_heads + ) + + @property + def inputs(self) -> dict[str, dict[int, str]]: + # Batched inference is not supported in Transformers. + if self.model_part == "text_encoder": + common_inputs = { + "input_ids": {0: "batch_size", 1: "encoder_sequence_length"}, + "attention_mask": {0: "batch_size", 1: "encoder_sequence_length"}, + } + elif self.model_part == "encodec_decode": + # 0: always 1 for chunk_length_s=None, 2: num_quantizers fixed. + common_inputs = {"audio_codes": {1: "batch_size", 3: "chunk_length"}} + elif self.model_part == "build_delay_pattern_mask": + common_inputs = { + "input_ids": {0: "batch_size_x_num_codebooks"}, + "pad_token_id": {}, + "max_length": {}, + } + elif self._behavior is ConfigBehavior.DECODER: + # Naming it total_batch_size as in case we use guidance_scale, the dimension 0 may be larger than simply the batch_size. + # Reference: https://github.com/huggingface/transformers/blob/31c575bcf13c2b85b65d652dd1b5b401f99be999/src/transformers/models/musicgen/modeling_musicgen.py#L1932-L1935 + common_inputs = { + "decoder_input_ids": {0: "total_batch_size_x_num_codebooks"}, + "encoder_outputs": {0: "total_batch_size", 1: "encoder_sequence_length"}, + # MusicgenForConditionalGeneration maps attention_mask to encoder_attention_mask. + "attention_mask": { + 0: "batch_size", + 1: "encoder_sequence_length", + }, + } + if self.use_past_in_inputs: + # TODO: validate the axis name for attention_mask + # common_inputs["attention_mask"][1] = "past_encoder_sequence_length + sequence_length" + self.add_past_key_values(common_inputs, direction="inputs") + else: + common_inputs["decoder_input_ids"] = { + 0: "total_batch_size_x_num_codebooks", + 1: "decoder_sequence_length", + } + else: + raise ValueError( + "This should not happen. Please open an issue at https://github.com/huggingface/optimum/issues." + ) + + return common_inputs + + @property + def outputs(self) -> dict[str, dict[int, str]]: + common_outputs = {} + + if self.model_part == "text_encoder": + common_outputs = super().outputs + elif self.model_part == "encodec_decode": + common_outputs["audio_values"] = {0: "batch_size", 2: "audio_length"} + elif self.model_part == "build_delay_pattern_mask": + common_outputs["input_ids_edited"] = {0: "total_batch_size_x_num_codebooks"} + common_outputs["delay_pattern_mask"] = {0: "total_batch_size_x_num_codebooks", 1: "max_length"} + elif self._behavior is ConfigBehavior.DECODER: + common_outputs = super().outputs + + # MusicgenForConditionalGeneration output is named logits, not last_hidden_state. + # Rename last_hidden_state -> logits while keeping the order. + common_outputs = { + "logits" if name == "last_hidden_state" else name: value for name, value in common_outputs.items() + } + else: + raise ValueError( + "This should not happen. Please open an issue at https://github.com/huggingface/optimum/issues." + ) + + return common_outputs + + def overwrite_shape_and_generate_input( + self, dummy_input_gen: DummyInputGenerator, input_name: str, framework: str, input_shapes: dict + ): + if self.model_part == "build_delay_pattern_mask" and input_name == "input_ids": + original_batch_size = dummy_input_gen.batch_size + dummy_input_gen.batch_size = ( + original_batch_size * dummy_input_gen.normalized_config.DECODER_NORMALIZED_CONFIG_CLASS.num_codebooks + ) + dummy_input = dummy_input_gen.generate( + input_name, framework=framework, int_dtype=self.int_dtype, float_dtype=self.float_dtype + ) + dummy_input_gen.batch_size = original_batch_size + else: + dummy_input = super().overwrite_shape_and_generate_input( + dummy_input_gen, input_name, framework, input_shapes + ) + + return dummy_input + + +@register_tasks_manager_onnx("speecht5", *["text-to-audio"]) +class SpeechT5OnnxConfig(OnnxSeq2SeqConfigWithPast): + # TODO: Transformers batched generation for Speecht5 is BROKEN (https://github.com/huggingface/transformers/pull/25943), + # so we won't support for now. + NORMALIZED_CONFIG_CLASS = NormalizedSeq2SeqConfig.with_args( + hidden_size="hidden_size", + num_attention_heads="encoder_attention_heads", # TODO: bugged in case encoder and decoder have different number of heads + encoder_num_layers="encoder_layers", + decoder_num_layers="decoder_layers", + allow_new=True, + ) + + DUMMY_INPUT_GENERATOR_CLASSES = ( + DummyTextInputGenerator, + DummySeq2SeqDecoderTextInputGenerator, + DummySeq2SeqPastKeyValuesGenerator, + DummySpeechT5InputGenerator, + ) + DUMMY_PKV_GENERATOR_CLASS = DummySeq2SeqPastKeyValuesGenerator + + VARIANTS = { # noqa: RUF012 + "with-past": "The export follows the Transformers implementation using the KV cache, with the following components exported:\n\t - encoder_model.onnx: corresponds to the encoding part in https://github.com/huggingface/transformers/blob/v4.33.2/src/transformers/models/speecht5/modeling_speecht5.py#L2544-L2556.\n\t - decoder_model.onnx: corresponds to the decoder part in https://github.com/huggingface/transformers/blob/v4.33.2/src/transformers/models/speecht5/modeling_speecht5.py#L2572-L2602.\n\t - decoder_with_past_model.onnx: same as the above, with past_key_values input (KV cache filled).\n\t - decoder_postnet_and_vocoder.onnx: Decoder speech postnet and vocoder (e.g. a SpeechT5HifiGan) to generate speech from the spectrogram, as in https://github.com/huggingface/transformers/blob/v4.33.2/src/transformers/models/speecht5/modeling_speecht5.py#L2605-L2614.", + "without-past": "The same as `with-past`, just without KV cache support. This is not a recommended export as slower than `with-past`.", + } + DEFAULT_VARIANT = "with-past" + _MODEL_PATCHER = SpeechT5ModelPatcher + + def __init__( + self, + config: PretrainedConfig, + task: str = "feature-extraction", + int_dtype: str = "int64", + float_dtype: str = "fp32", + use_past: bool = False, + use_past_in_inputs: bool = False, + behavior: ConfigBehavior = ConfigBehavior.MONOLITH, + preprocessors: list[Any] | None = None, + is_postnet_and_vocoder: bool = False, + ): + super().__init__( + config=config, + task=task, + int_dtype=int_dtype, + float_dtype=float_dtype, + use_past=use_past, + use_past_in_inputs=use_past_in_inputs, + behavior=behavior, + preprocessors=preprocessors, + ) + if float_dtype == "fp16": + raise ValueError( + "The ONNX export of SpeechT5 in float16 is currently not supported due to a bug in PyTorch: https://github.com/pytorch/pytorch/pull/110078. Please open an issue in Optimum if you would like to export SpeechT5 in float16." + ) + self.is_postnet_and_vocoder = is_postnet_and_vocoder + + @property + def inputs(self) -> dict[str, dict[int, str]]: + common_inputs = {} + + # Batched inference is not supported in Transformers. + if self._behavior is ConfigBehavior.ENCODER: + common_inputs["input_ids"] = {1: "encoder_sequence_length"} + elif self._behavior is ConfigBehavior.DECODER: + # NOTE: even when past is used, the decoder takes the full sequence as input as the prenet seem to require it: + # https://github.com/huggingface/transformers/blob/v4.33.2/src/transformers/models/speecht5/modeling_speecht5.py#L2573 + common_inputs["output_sequence"] = {1: "decoder_sequence_length"} + common_inputs["speaker_embeddings"] = {} # No dynamic shape here. + common_inputs["encoder_outputs"] = {1: "encoder_sequence_length"} + common_inputs["encoder_attention_mask"] = {1: "encoder_sequence_length"} + + if self.variant == "with-past" and self.use_past_in_inputs: + self.add_past_key_values(common_inputs, direction="inputs") + elif self.is_postnet_and_vocoder: + common_inputs["spectrogram"] = {0: "n_spectrums x reduction_factor"} + else: + raise ValueError( + "self._behavior is neither encoder or decoder, and is_postnet_and_vocoder=False. This should not happen." + ) + + return common_inputs + + @property + def outputs(self) -> dict[str, dict[int, str]]: + common_outputs = {} + if self._behavior is ConfigBehavior.ENCODER: + common_outputs["encoder_outputs"] = {1: "encoder_sequence_length"} + common_outputs["encoder_attention_mask"] = {1: "encoder_sequence_length"} + elif self._behavior is ConfigBehavior.DECODER: + common_outputs["output_sequence_out"] = {1: "decoder_sequence_length + 1"} + common_outputs["spectrum"] = {} # No dynamic shape here. + common_outputs["prob"] = {} # No dynamic shape here. + + if self.variant == "with-past" and self.use_past: + # When exporting decoder models with use_cache=True, both the decoder without past and with past have the KV cache as an output. + self.add_past_key_values(common_outputs, direction="outputs") + elif self.is_postnet_and_vocoder: + common_outputs["waveform"] = {0: "n_samples"} + else: + raise ValueError( + "self._behavior is neither encoder or decoder, and is_postnet_and_vocoder=False. This should not happen." + ) + + return common_outputs + + def overwrite_shape_and_generate_input( + self, dummy_input_gen: DummyInputGenerator, input_name: str, framework: str, input_shapes: dict + ): + dummy_input_gen.batch_size = 1 + dummy_input = dummy_input_gen.generate( + input_name, framework=framework, int_dtype=self.int_dtype, float_dtype=self.float_dtype + ) + return dummy_input + + +@register_tasks_manager_onnx("vits", *["text-to-audio"]) +class VitsOnnxConfig(TextEncoderOnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig + + @property + def inputs(self) -> dict[str, dict[int, str]]: + return { + "input_ids": {0: "text_batch_size", 1: "sequence_length"}, + "attention_mask": {0: "text_batch_size", 1: "sequence_length"}, + } + + @property + def outputs(self) -> dict[str, dict[int, str]]: + return { + "waveform": {0: "text_batch_size", 1: "n_samples"}, + "spectrogram": {0: "text_batch_size", 2: "num_bins"}, + } + + +@register_tasks_manager_onnx( + "speech_to_text", + *[ + "feature-extraction", + "feature-extraction-with-past", + "automatic-speech-recognition", + "automatic-speech-recognition-with-past", + ], +) +class Speech2TextOnnxConfig(AudioToTextOnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedSeq2SeqConfig.with_args( + decoder_num_layers="decoder_layers", + num_layers="decoder_layers", + input_features_per_channel="input_feat_per_channel", + allow_new=True, + ) + DUMMY_INPUT_GENERATOR_CLASSES = ( + Speech2TextDummyAudioInputGenerator, + *AudioToTextOnnxConfig.DUMMY_INPUT_GENERATOR_CLASSES[1:], + DummyTextInputGenerator, + ) + + @property + def inputs(self) -> dict[str, dict[int, str]]: + common_inputs = {} + + if self._behavior in {ConfigBehavior.ENCODER, ConfigBehavior.MONOLITH}: + common_inputs["input_features"] = {0: "batch_size", 1: "encoder_sequence_length"} + else: + common_inputs["encoder_outputs"] = {0: "batch_size", 1: "encoder_sequence_length"} + common_inputs["attention_mask"] = {0: "batch_size", 1: "encoder_sequence_length"} + + if self._behavior in {ConfigBehavior.DECODER, ConfigBehavior.MONOLITH}: + common_inputs["decoder_input_ids"] = {0: "batch_size", 1: "decoder_sequence_length"} + if self.use_past_in_inputs: + self.add_past_key_values(common_inputs, direction="inputs") + + return common_inputs + + @property + def outputs(self) -> dict[str, dict[int, str]]: + common_outputs = super().outputs + if self._behavior in {ConfigBehavior.ENCODER, ConfigBehavior.MONOLITH}: + if self._behavior is ConfigBehavior.MONOLITH: + output_name = "encoder_last_hidden_state" + else: + output_name = "last_hidden_state" + # Speech2Text encoder output formula adapted from: + # Speech2TextPreTrainedModel._get_feat_extract_output_lengths + # for i in range(self.config.num_conv_layers): + # input_lengths = (input_lengths - 1) // 2 + 1 + downsample_factor = 2 * self._config.num_conv_layers + output_sequence_length = f"( encoder_sequence_length + {downsample_factor} - 1 ) // {downsample_factor}" + common_outputs[output_name] = {0: "batch_size", 1: output_sequence_length} + return common_outputs + + +# TrOCR is a causal model, used as the decoder in some vision encoder-decoder models. +@register_tasks_manager_onnx( + "trocr", + *[ + "feature-extraction", + "feature-extraction-with-past", + ], +) +class TrOCROnnxConfig(TextSeq2SeqOnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedSeq2SeqConfig.with_args( + decoder_num_layers="decoder_layers", + num_layers="decoder_layers", + decoder_num_attention_heads="decoder_attention_heads", + hidden_size="hidden_size", + ) + + +@register_tasks_manager_onnx( + "vision-encoder-decoder", + *[ + "image-to-text", + "image-to-text-with-past", + ], +) +class VisionEncoderDecoderOnnxConfig(EncoderDecoderBaseOnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedEncoderDecoderConfig + DUMMY_INPUT_GENERATOR_CLASSES = (DummyVisionInputGenerator, DummyVisionEncoderDecoderPastKeyValuesGenerator) + + @property + def inputs(self) -> dict[str, dict[int, str]]: + common_inputs = {} + + if self._behavior in {ConfigBehavior.ENCODER, ConfigBehavior.MONOLITH}: + common_inputs["pixel_values"] = {0: "batch_size", 1: "num_channels", 2: "height", 3: "width"} + else: + common_inputs["encoder_outputs"] = {0: "batch_size", 1: "encoder_sequence_length"} + + if self._behavior in {ConfigBehavior.DECODER, ConfigBehavior.MONOLITH}: + common_inputs["decoder_input_ids"] = {0: "batch_size", 1: "decoder_sequence_length"} + if self.use_past_in_inputs: + self.add_past_key_values(common_inputs, direction="inputs") + + return common_inputs + + @property + def outputs(self) -> dict[str, dict[int, str]]: + if self._behavior == ConfigBehavior.ENCODER: + # Some encoders have static sequence length so it is useful to rely on the encoder ONNX config to grab this information. + return self._encoder_onnx_config.outputs + else: + # Ideally, we would want here to have self._decoder_onnx_config.outputs, which is currently not possible + # as we hard-code the task to feature-extraction, that has the wrong output names (e.g. mbart does not support document-question-answering + # so we can not initializer MBartONNXConfig with document-question-answering). + return super().outputs + + +@register_tasks_manager_onnx("sam", *["feature-extraction"]) +class SamOnnxConfig(OnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedEncoderDecoderConfig + DUMMY_INPUT_GENERATOR_CLASSES = (DummyVisionInputGenerator, DummyPointsGenerator, DummyVisionEmbeddingsGenerator) + VARIANTS = { # noqa: RUF012 + "monolith": "All the SAM model components are exported as a single model.onnx.", + "split": "The vision encoder is exported as a separate vision_encoder.onnx, and the prompt encoder and mask decoder are exported as a prompt_encoder_mask_decoder.onnx. This allows to encoder the image only once for multiple point queries.", + } + DEFAULT_VARIANT = "split" + _MODEL_PATCHER = SAMModelPatcher + + def __init__( + self, + config: PretrainedConfig, + task: str = "feature-extraction", + int_dtype: str = "int64", + float_dtype: str = "fp32", + variant: str = "split", + vision_encoder: bool | None = None, + preprocessors: list[Any] | None = None, + ): + super().__init__( + config=config, + task=task, + int_dtype=int_dtype, + float_dtype=float_dtype, + preprocessors=preprocessors, + ) + self.variant = variant + self.vision_encoder = vision_encoder + self._normalized_config.ENCODER_NORMALIZED_CONFIG_CLASS = NormalizedVisionConfig(self._config.vision_config) + + @property + def inputs(self) -> dict[str, dict[int, str]]: + if self.variant == "monolith": + inputs = { + "pixel_values": {0: "batch_size"}, + "input_points": {0: "batch_size", 1: "point_batch_size", 2: "nb_points_per_image"}, + "input_labels": {0: "batch_size", 1: "point_batch_size", 2: "nb_points_per_image"}, + } + else: + if self.vision_encoder: + inputs = {"pixel_values": {0: "batch_size"}} + else: + inputs = { + "image_positional_embeddings": {0: "batch_size"}, + "image_embeddings": {0: "batch_size"}, + "input_points": {0: "batch_size", 1: "point_batch_size", 2: "nb_points_per_image"}, + "input_labels": {0: "batch_size", 1: "point_batch_size", 2: "nb_points_per_image"}, + } + return inputs + + @property + def outputs(self) -> dict[str, dict[int, str]]: + if self.variant == "split" and self.vision_encoder: + return {"image_embeddings": {0: "batch_size"}, "image_positional_embeddings": {0: "batch_size"}} + else: + return { + "iou_scores": {0: "batch_size", 1: "point_batch_size"}, + "pred_masks": {0: "batch_size", 1: "point_batch_size"}, + } + + +class Pix2StructNormalizedConfig(NormalizedSeq2SeqConfig): + ENCODER_NUM_LAYERS = "vision_config.num_hidden_layers" + DECODER_NUM_LAYERS = "text_config.num_layers" + ENCODER_NUM_ATTENTION_HEADS = "vision_config.num_attention_heads" + DECODER_NUM_ATTENTION_HEADS = "text_config.num_heads" + HIDDEN_SIZE = "text_config.hidden_size" + VOCAB_SIZE = "text_config.vocab_size" + + +@register_tasks_manager_onnx( + "pix2struct", + *[ + "image-to-text", + "image-to-text-with-past", + ], +) +class Pix2StructOnnxConfig(OnnxSeq2SeqConfigWithPast): + PAD_ATTENTION_MASK_TO_PAST = True + NORMALIZED_CONFIG_CLASS = Pix2StructNormalizedConfig + DUMMY_INPUT_GENERATOR_CLASSES = ( + DummyTextInputGenerator, + DummySeq2SeqDecoderTextInputGenerator, + DummySeq2SeqPastKeyValuesGenerator, + DummyPix2StructInputGenerator, + ) + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + if is_transformers_version("==", "4.46.0"): + logging.warn_once( + logger, + "Found transformers v4.46.0 while trying to export a Pix2Struct model, " + "this specific version of transformers is broken for this model. Please " + "upgrade to v4.46.1 or higher, or downgrade to v4.45.x.", + ) + + @property + def inputs(self): + common_inputs = {} + + if self._behavior in {ConfigBehavior.ENCODER, ConfigBehavior.MONOLITH}: + common_inputs["flattened_patches"] = {0: "batch_size"} + else: + common_inputs["encoder_outputs"] = {0: "batch_size"} + common_inputs["attention_mask"] = {0: "batch_size"} + + if self._behavior in {ConfigBehavior.DECODER, ConfigBehavior.MONOLITH}: + common_inputs["decoder_input_ids"] = {0: "batch_size", 1: "decoder_sequence_length"} + if self.use_past_in_inputs: + self.add_past_key_values(common_inputs, direction="inputs") + decoder_attention_mask_dim = "past_decoder_sequence_length + decoder_sequence_length" + else: + decoder_attention_mask_dim = "decoder_sequence_length" + common_inputs["decoder_attention_mask"] = {0: "batch_size", 1: decoder_attention_mask_dim} + + return common_inputs + + @property + def outputs(self) -> dict[str, dict[int, str]]: + common_outputs = super().outputs + if self._behavior in {ConfigBehavior.ENCODER, ConfigBehavior.MONOLITH}: + if self._behavior is ConfigBehavior.MONOLITH: + output_name = "encoder_last_hidden_state" + else: + output_name = "last_hidden_state" + common_outputs[output_name] = {0: "batch_size"} # Remove unnecessary dynamic axis. + + return common_outputs + + def _create_dummy_input_generator_classes(self, **kwargs) -> list[DummyInputGenerator]: + if self._preprocessors is None or len(self._preprocessors) < 2: + raise ValueError( + f"Preprocessors for pix2struct need to be available for the ONNX export to infer input static shapes. Got: {self._preprocessors}" + ) + + dummy_inputs_generators = [] + dummy_inputs_generators.append( + self.DUMMY_INPUT_GENERATOR_CLASSES[0](self.task, self._normalized_config, **kwargs) + ) + # A hack for DummyPix2StructInputGenerator to gain access to the preprocessors. + # TODO: we probably pass preprocessors to all dummy input generators. + encoder_sequence_length = self._preprocessors[1].image_processor.max_patches + kwargs["preprocessors"] = self._preprocessors + for cls_ in self.DUMMY_INPUT_GENERATOR_CLASSES[1:]: + dummy_inputs_generators.append( + cls_(self.task, self._normalized_config, encoder_sequence_length=encoder_sequence_length, **kwargs) + ) + + return dummy_inputs_generators + + def overwrite_shape_and_generate_input( + self, dummy_input_gen: DummyInputGenerator, input_name: str, framework: str, input_shapes: dict + ): + if self._preprocessors is None or len(self._preprocessors) < 2: + raise ValueError( + f"Preprocessors for pix2struct need to be available for the ONNX export to infer input static shapes. Got: {self._preprocessors}" + ) + + # it would been simpler if pix2struct dummy input generator took care of generating these as well + if input_name in ["encoder_outputs", "attention_mask"]: + # Pix2struct takes inputs encoder inputs/outputs with a fixed sequence length (max_patches). + original_seq_length = dummy_input_gen.sequence_length + dummy_input_gen.sequence_length = self._preprocessors[1].image_processor.max_patches + dummy_input = dummy_input_gen.generate( + input_name, framework=framework, int_dtype=self.int_dtype, float_dtype=self.float_dtype + ) + dummy_input_gen.sequence_length = original_seq_length + else: + dummy_input = super().overwrite_shape_and_generate_input( + dummy_input_gen, input_name, framework, input_shapes + ) + + return dummy_input + + +@register_tasks_manager_onnx("encoder-decoder", *["text2text-generation", "text2text-generation-with-past"]) +class EncoderDecoderOnnxConfig(EncoderDecoderBaseOnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedEncoderDecoderConfig + + +@register_tasks_manager_onnx("patchtst", *["feature-extraction", "time-series-forecasting"]) +class PatchTSTOnnxConfig(OnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedTimeSeriesForecastingConfig + DUMMY_INPUT_GENERATOR_CLASSES = (DummyPatchTSTInputGenerator,) + + @property + def inputs(self) -> dict[str, dict[int, str]]: + return {"past_values": {0: "batch_size", 1: "sequence_length"}} + + @property + def outputs(self) -> dict[str, dict[int, str]]: + if self.task == "feature-extraction": + return {"last_hidden_state": {0: "batch_size"}} + else: + return super().outputs + + +@register_tasks_manager_onnx("patchtsmixer", *["feature-extraction", "time-series-forecasting"]) +class PatchTSMixerOnnxConfig(PatchTSTOnnxConfig): + pass + + +@register_tasks_manager_onnx("rt_detr", *["object-detection"]) +class RTDetrOnnxConfig(ViTOnnxConfig): + @property + def inputs(self) -> dict[str, dict[int, str]]: + return { + "pixel_values": {0: "batch_size", 2: "height", 3: "width"}, + } + + def _create_dummy_input_generator_classes(self, **kwargs) -> list[DummyInputGenerator]: + min_image_size = int(math.ceil(self._config.num_queries / 32) * 32) + if kwargs["height"] < min_image_size: + warnings.warn( + f"Exporting model with image `height={kwargs['height']}` which is less than " + f"minimal {min_image_size}, setting `height` to {min_image_size}.", + stacklevel=2, + ) + kwargs["height"] = min_image_size + if kwargs["width"] < min_image_size: + warnings.warn( + f"Exporting model with image `width={kwargs['width']}` which is less than " + f"minimal {min_image_size}, setting `width` to {min_image_size}.", + stacklevel=2, + ) + kwargs["width"] = min_image_size + return super()._create_dummy_input_generator_classes(**kwargs) + + +@register_tasks_manager_onnx("rt_detr_v2", *["object-detection"]) +class RTDetrV2OnnxConfig(RTDetrOnnxConfig): + pass + + +@register_tasks_manager_onnx("colpali", *["feature-extraction"]) +class ColPaliOnnxConfig(GemmaOnnxConfig): + DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, DummyVisionInputGenerator) + NORMALIZED_CONFIG_CLASS = NormalizedTextAndVisionConfig.with_args( + allow_new=True, + text_config="text_config", + vision_config="vlm_config.vision_config", + vlm_config="vlm_config", + ) + + VARIANTS = { # noqa: RUF012 + "vision": "Embedding extraction for image.", + "text": "Embedding extraction for text.", + } + DEFAULT_VARIANT = "vision" + + @property + def inputs(self) -> dict[str, dict[int, str]]: + dynamic_axis = {0: "batch_size", 1: "sequence_length"} + if self.variant == "vision": + return { + "input_ids": dynamic_axis, + "attention_mask": dynamic_axis, + "pixel_values": {0: "batch_size"}, + } + else: + return { + "input_ids": dynamic_axis, + "attention_mask": dynamic_axis, + } + + @property + def outputs(self) -> dict[str, dict[int, str]]: + return { + "embeddings": {0: "batch_size", 1: "sequence_length"}, + } + + def generate_dummy_inputs(self, framework: str = "pt", **kwargs): + if self.variant == "vision": + image_token_index = self._normalized_config.vlm_config.image_token_index + num_image_tokens = self._normalized_config.vision_config.num_image_tokens + if "sequence_length" in kwargs: + kwargs["sequence_length"] += num_image_tokens + else: + kwargs["sequence_length"] = DEFAULT_DUMMY_SHAPES["sequence_length"] + num_image_tokens + + dummy_inputs = super().generate_dummy_inputs(framework=framework, **kwargs) + + if self.variant == "vision": + dummy_inputs["input_ids"][:, :num_image_tokens] = image_token_index + return dummy_inputs + + +@register_tasks_manager_onnx("d_fine", *["object-detection"]) +class DFineOnnxConfig(RTDetrOnnxConfig): + MIN_TRANSFORMERS_VERSION = version.parse("4.52.0") + + +@register_tasks_manager_onnx("gemma2-text-encoder", *["feature-extraction"], library_name="diffusers") +class Gemma2TextEncoderOnnxConfig(Gemma2OnnxConfig): + MIN_TRANSFORMERS_VERSION = version.parse("4.42.0") + + +@register_tasks_manager_onnx("sana-transformer", *["semantic-segmentation"], library_name="diffusers") +class SanaTransformerOnnxConfig(SD3TransformerOnnxConfig): + DUMMY_INPUT_GENERATOR_CLASSES = ( + DummyTransformerVisionInputGenerator, + DummySanaTransforemerTextInputGenerator, + DummyTransformerTimestepInputGenerator, + ) + NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args( + hidden_size="caption_channels", + num_channels="in_channels", + allow_new=True, + ) + + @property + def inputs(self) -> dict[str, dict[int, str]]: + return { + "hidden_states": {0: "batch_size", 2: "height", 3: "width"}, + "encoder_hidden_states": {0: "batch_size", 1: "sequence_length"}, + "encoder_attention_mask": {0: "batch_size", 1: "sequence_length"}, + "timestep": {0: "batch_size"}, + } + + +@register_tasks_manager_onnx("dcae-encoder", *["semantic-segmentation"], library_name="diffusers") +class DcaeEncoderOnnxConfig(VaeEncoderOnnxConfig): + @property + def inputs(self) -> dict[str, dict[int, str]]: + return { + "sample": {0: "batch_size", 2: "height", 3: "width"}, + } + + @property + def outputs(self) -> dict[str, dict[int, str]]: + down_sampling_factor = 2 ** (len(self._normalized_config.encoder_block_out_channels) - 1) - setattr(TasksManager, supported_models_config, supported_models) + return { + "latent_sample": { + 0: "batch_size", + 2: f"height / {down_sampling_factor}", + 3: f"width / {down_sampling_factor}", + } + } -init_model_configs() +@register_tasks_manager_onnx("dcae-decoder", *["semantic-segmentation"], library_name="diffusers") +class DcaeDecoderOnnxConfig(VaeDecoderOnnxConfig): + @property + def inputs(self) -> dict[str, dict[int, str]]: + return { + "latent_sample": {0: "batch_size", 2: "latent_height", 3: "latent_width"}, + } + @property + def outputs(self) -> dict[str, dict[int, str]]: + up_sampling_factor = 2 ** (len(self._normalized_config.decoder_block_out_channels) - 1) -register_in_tasks_manager = TasksManager.create_register("openvino", overwrite_existing=True) + return { + "sample": { + 0: "batch_size", + 2: f"latent_height * {up_sampling_factor}", + 3: f"latent_width * {up_sampling_factor}", + } + } @register_in_tasks_manager("baichuan", *["text-generation", "text-generation-with-past"], library_name="transformers") @@ -451,66 +3203,6 @@ def inputs(self) -> Dict[str, Dict[int, str]]: return common_inputs -class DummyQwen3VLLMInputGenerator(DummyTextInputGenerator): - SUPPORTED_INPUT_NAMES = ( - "input_ids", - "attention_mask", - "token_type_ids", - "position_ids", - "visual_pos_masks", - "deepstack_visual_embeds", - ) - - def __init__( - self, - task: str, - normalized_config: NormalizedTextConfig, - batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], - sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], - num_choices: int = DEFAULT_DUMMY_SHAPES["num_choices"], - random_batch_size_range: Optional[Tuple[int, int]] = None, - random_sequence_length_range: Optional[Tuple[int, int]] = None, - random_num_choices_range: Optional[Tuple[int, int]] = None, - padding_side: str = "right", - **kwargs, - ): - super().__init__( - task=task, - normalized_config=normalized_config, - batch_size=batch_size, - sequence_length=sequence_length, - num_choices=num_choices, - random_batch_size_range=random_batch_size_range, - random_sequence_length_range=random_sequence_length_range, - random_num_choices_range=random_num_choices_range, - padding_side=padding_side, - **kwargs, - ) - self.embed_dim = normalized_config.hidden_size - self.num_layers = len(self.normalized_config.deepstack_visual_indexes) - - def generate( - self, - input_name: str, - framework: str = "pt", - int_dtype: str = "int64", - float_dtype: str = "fp32", - bool_dtype: str = "bool", - ): - if input_name == "deepstack_visual_embeds": - return self.random_float_tensor( - [self.num_layers, 2 * self.sequence_length, self.embed_dim], framework=framework, dtype=float_dtype - ) - if input_name == "visual_pos_masks": - return self.constant_tensor( - shape=[self.batch_size, self.sequence_length], - framework=framework, - value=1, - dtype=DTYPE_MAPPER.pt(bool_dtype), - ) - return super().generate(input_name, framework, int_dtype, float_dtype) - - @register_in_tasks_manager( "qwen3_vl_text", *[ @@ -554,46 +3246,6 @@ class MiniCPMOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): _MODEL_PATCHER = MiniCPMModelPatcher -class OVMiniCPM3DummyPastKeyValuesGenerator(MistralDummyPastKeyValuesGenerator): - def __init__( - self, - task: str, - normalized_config: NormalizedTextConfig, - batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], - sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], - random_batch_size_range: Optional[Tuple[int, int]] = None, - random_sequence_length_range: Optional[Tuple[int, int]] = None, - **kwargs, - ): - super().__init__( - task=task, - normalized_config=normalized_config, - batch_size=batch_size, - sequence_length=sequence_length, - random_batch_size_range=random_batch_size_range, - random_sequence_length_range=random_sequence_length_range, - **kwargs, - ) - self.v_head_dim = getattr(normalized_config, "v_head_dim", self.hidden_size // self.num_attention_heads) - self.k_head_dim = normalized_config.qk_nope_head_dim + normalized_config.qk_rope_head_dim - - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - v_shape = ( - self.batch_size, - self.num_key_value_heads, - self.sequence_length, - self.v_head_dim, - ) - k_shape = (self.batch_size, self.num_key_value_heads, self.sequence_length, self.k_head_dim) - return [ - ( - self.random_float_tensor(k_shape, framework=framework, dtype=float_dtype), - self.random_float_tensor(v_shape, framework=framework, dtype=float_dtype), - ) - for _ in range(self.num_layers) - ] - - @register_in_tasks_manager("minicpm3", *["text-generation", "text-generation-with-past"], library_name="transformers") class MiniCPM3OpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): DEFAULT_ONNX_OPSET = 14 @@ -613,53 +3265,6 @@ class StableLMOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): _MODEL_PATCHER = OVDecoderModelPatcher -class ChatGLM2DummyPastKeyValuesGenerator(DummyPastKeyValuesGenerator): - def __init__( - self, - task: str, - normalized_config: NormalizedTextConfig, - batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], - sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], - random_batch_size_range: Optional[Tuple[int, int]] = None, - random_sequence_length_range: Optional[Tuple[int, int]] = None, - **kwargs, - ): - super().__init__( - task=task, - normalized_config=normalized_config, - batch_size=batch_size, - sequence_length=sequence_length, - random_batch_size_range=random_batch_size_range, - random_sequence_length_range=random_sequence_length_range, - ) - self.multi_query_group_num = normalized_config.multi_query_group_num - self.head_dim = normalized_config.kv_channels - self.standart_cache_layout = hasattr(normalized_config, "rope_ratio") - - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - if not self.standart_cache_layout: - pkv_shape = ( - self.sequence_length, - self.batch_size, - self.multi_query_group_num, - self.head_dim, - ) - else: - pkv_shape = ( - self.batch_size, - self.multi_query_group_num, - self.sequence_length, - self.head_dim, - ) - return [ - ( - self.random_float_tensor(pkv_shape, framework=framework, dtype=float_dtype), - self.random_float_tensor(pkv_shape, framework=framework, dtype=float_dtype), - ) - for _ in range(self.num_layers) - ] - - @register_in_tasks_manager("chatglm", *["text-generation", "text-generation-with-past"], library_name="transformers") class ChatGLM2OpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args(vocab_size="padded_vocab_size", num_layers="num_layers") @@ -786,73 +3391,6 @@ def inputs(self) -> Dict[str, Dict[int, str]]: return inputs -class Eagle3DummyGenerator(DummyInputGenerator): - """ - Dummy input generator for Eagle-3 speculative decoding. - - This generator produces synthetic `hidden_states` tensors that mimic the - intermediate hidden-state outputs of a *main (target) model*, which are - required by the Eagle-3 draft model during speculative decoding. - """ - - SUPPORTED_INPUT_NAMES = ("hidden_states",) - - def __init__( - self, - task: str, - normalized_config: NormalizedTextConfig, - batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], - sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], - **kwargs, - ): - self.batch_size = batch_size - self.sequence_length = sequence_length - self.hidden_size = normalized_config.hidden_size - - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - # hidden_states is provided as a concatenation of three hidden-layer outputs from the main model - shape = ( - self.batch_size, - self.sequence_length, - self.hidden_size * 3, - ) - return self.random_float_tensor(shape, framework=framework, dtype=float_dtype) - - -class Eagle3VLMDummyGenerator(DummyInputGenerator): - """ - Dummy input generator for VLM Eagle-3 speculative decoding. - - Produces `inputs_embeds` (float) and 3D `position_ids` (MRoPE) - required by VLM Eagle-3 draft models targeting Qwen3-VL. - """ - - SUPPORTED_INPUT_NAMES = ("inputs_embeds", "position_ids") - - def __init__( - self, - task: str, - normalized_config: NormalizedTextConfig, - batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], - sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], - **kwargs, - ): - self.batch_size = batch_size - self.sequence_length = sequence_length - self.hidden_size = normalized_config.hidden_size - - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - if input_name == "inputs_embeds": - shape = (self.batch_size, self.sequence_length, self.hidden_size) - return self.random_float_tensor(shape, framework=framework, dtype=float_dtype) - if input_name == "position_ids": - # The rotary embedding is MRoPE (Multimodal RoPE) - # MRoPE encodes position along three independent axes: temporal, height, and width - # https://github.com/Tencent/AngelSlim/blob/main/angelslim/compressor/speculative/train/models/draft/llama_eagle3.py#L211 - shape = (3, self.batch_size, self.sequence_length) - return self.random_int_tensor(shape, max_value=self.sequence_length, framework=framework, dtype=int_dtype) - - @register_in_tasks_manager( "llama", *[ @@ -1027,39 +3565,6 @@ class Cohere2OpenVINOConfig(LlamaOpenVINOConfig): MIN_TRANSFORMERS_VERSION = "4.48.0" -class QwenDummyPastKeyValuesGenerator(DummyPastKeyValuesGenerator): - def __init__( - self, - task: str, - normalized_config: NormalizedTextConfig, - batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], - sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], - random_batch_size_range: Optional[Tuple[int, int]] = None, - random_sequence_length_range: Optional[Tuple[int, int]] = None, - **kwargs, - ): - super().__init__( - task=task, - normalized_config=normalized_config, - batch_size=batch_size, - sequence_length=sequence_length, - random_batch_size_range=random_batch_size_range, - random_sequence_length_range=random_sequence_length_range, - ) - self.kv_channels = normalized_config.kv_channels - - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - past_key_shape = (self.batch_size, self.sequence_length, self.num_attention_heads, self.kv_channels) - past_value_shape = (self.batch_size, self.sequence_length, self.num_attention_heads, self.kv_channels) - return [ - ( - self.random_float_tensor(past_key_shape, framework=framework, dtype=float_dtype), - self.random_float_tensor(past_value_shape, framework=framework, dtype=float_dtype), - ) - for _ in range(self.num_layers) - ] - - @register_in_tasks_manager("qwen", *["text-generation", "text-generation-with-past"]) class QwenOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): DEFAULT_ONNX_OPSET = 14 @@ -1235,34 +3740,6 @@ class PhiOpenVINOConfig(PhiOnnxConfig): _MODEL_PATCHER = OVDecoderModelPatcher -class OVFalconDummyPastKeyValuesGenerator(FalconDummyPastKeyValuesGenerator): - def __init__( - self, - task: str, - normalized_config: NormalizedTextConfig, - batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], - sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], - random_batch_size_range: Optional[Tuple[int, int]] = None, - random_sequence_length_range: Optional[Tuple[int, int]] = None, - **kwargs, - ): - super().__init__( - task=task, - normalized_config=normalized_config, - batch_size=batch_size, - sequence_length=sequence_length, - random_batch_size_range=random_batch_size_range, - random_sequence_length_range=random_sequence_length_range, - **kwargs, - ) - if normalized_config.new_decoder_architecture: - self.num_kv_heads = normalized_config.num_attention_heads - else: - self.num_kv_heads = normalized_config.num_kv_heads if not normalized_config.multi_query else 1 - - self.head_dim = self.hidden_size // self.num_attention_heads - - @register_in_tasks_manager( "falcon", *[ @@ -1382,46 +3859,6 @@ def __init__(self, *args, **kwargs): _warn_potential_accuracy_issue_ov_2026_1("xglm") -class AquilaDummyPastKeyValuesGenerator(DummyPastKeyValuesGenerator): - def __init__( - self, - task: str, - normalized_config: NormalizedTextConfig, - batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], - sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], - random_batch_size_range: Optional[Tuple[int, int]] = None, - random_sequence_length_range: Optional[Tuple[int, int]] = None, - **kwargs, - ): - super().__init__( - task, - normalized_config, - batch_size, - sequence_length, - random_batch_size_range, - random_sequence_length_range, - **kwargs, - ) - self.num_key_value_heads = getattr( - normalized_config, "num_key_value_heads", normalized_config.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.hidden_size // self.num_attention_heads, - ) - 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) - ] - - @register_in_tasks_manager("aquila", *["text-generation", "text-generation-with-past"], library_name="transformers") class AquilaMOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): DEFAULT_ONNX_OPSET = 14 @@ -1501,44 +3938,6 @@ class ArcticOpenVINOConfig(MixtralOpenVINOConfig): _MODEL_PATCHER = ArcticModelPatcher -class OVMistralDummyPastKeyValuesGenerator(MistralDummyPastKeyValuesGenerator): - def __init__( - self, - task: str, - normalized_config: NormalizedTextConfig, - batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], - sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], - random_batch_size_range: Optional[Tuple[int, int]] = None, - random_sequence_length_range: Optional[Tuple[int, int]] = None, - **kwargs, - ): - super().__init__( - task=task, - normalized_config=normalized_config, - batch_size=batch_size, - sequence_length=sequence_length, - random_batch_size_range=random_batch_size_range, - random_sequence_length_range=random_sequence_length_range, - **kwargs, - ) - self.head_dim = getattr(normalized_config, "head_dim", 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) - ] - - @register_in_tasks_manager( "mistral", *[ @@ -1614,67 +4013,6 @@ class Gemma3TextOpenVINOConfig(Gemma2OpenVINOConfig): MIN_TRANSFORMERS_VERSION = "4.50.0" -class Gemma4DummyPastKeyValuesGenerator(DummyPastKeyValuesGenerator): - def __init__( - self, - task: str, - normalized_config: NormalizedTextConfig, - batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], - sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], - random_batch_size_range: Optional[Tuple[int, int]] = None, - random_sequence_length_range: Optional[Tuple[int, int]] = None, - **kwargs, - ): - super().__init__( - task=task, - normalized_config=normalized_config, - batch_size=batch_size, - sequence_length=sequence_length, - random_batch_size_range=random_batch_size_range, - random_sequence_length_range=random_sequence_length_range, - ) - self.num_key_value_heads = normalized_config.num_key_value_heads - self.head_dim = normalized_config.head_dim - self.global_head_dim = getattr(normalized_config.config, "global_head_dim", self.head_dim) - self.layer_types = normalized_config.config.layer_types - self.num_kv_shared_layers = normalized_config.config.num_kv_shared_layers - self.sliding_window = normalized_config.config.sliding_window - # Full-attention layers use fewer KV heads than sliding-attention layers (e.g. 2 vs 8 for 26B-A4B) - self.num_global_key_value_heads = ( - getattr(normalized_config.config, "num_global_key_value_heads", None) or self.num_key_value_heads - ) - - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - # some layers do not produce their own KV-cache, they use the shared KV-cache - if self.num_kv_shared_layers > 0: - layer_types = self.layer_types[: -self.num_kv_shared_layers] - else: - layer_types = self.layer_types - past_kv_values = [] - for layer_type in layer_types: - if layer_type == "sliding_attention": - shape = ( - self.batch_size, - self.num_key_value_heads, - self.sliding_window, - self.head_dim, - ) - else: - shape = ( - self.batch_size, - self.num_global_key_value_heads, - self.sequence_length, - self.global_head_dim, - ) - past_kv_value = ( - self.random_float_tensor(shape, framework=framework, dtype=float_dtype), - self.random_float_tensor(shape, framework=framework, dtype=float_dtype), - ) - past_kv_values.append(past_kv_value) - - return past_kv_values - - @register_in_tasks_manager( "gemma4_text", *[ @@ -1714,46 +4052,6 @@ def add_past_key_values(self, inputs_or_outputs: dict[str, dict[int, str]], dire inputs_or_outputs[f"{name}.{i}.value"] = {0: "batch_size", 2: decoder_sequence_name} -class DeciDummyPastKeyValuesGenerator(DummyPastKeyValuesGenerator): - def __init__( - self, - task: str, - normalized_config: NormalizedTextConfig, - batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], - sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], - random_batch_size_range: Optional[Tuple[int, int]] = None, - random_sequence_length_range: Optional[Tuple[int, int]] = None, - **kwargs, - ): - super().__init__( - task=task, - normalized_config=normalized_config, - batch_size=batch_size, - sequence_length=sequence_length, - random_batch_size_range=random_batch_size_range, - random_sequence_length_range=random_sequence_length_range, - ) - self.num_key_value_heads_per_layer = normalized_config.num_key_value_heads_per_layer - - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - past_key_values = [] - - for layer_id in range(self.num_layers): - shape = ( - self.batch_size, - self.num_key_value_heads_per_layer[layer_id], - self.sequence_length, - self.hidden_size // self.num_attention_heads, - ) - past_key_values.append( - ( - self.random_float_tensor(shape, framework=framework, dtype=float_dtype), - self.random_float_tensor(shape, framework=framework, dtype=float_dtype), - ) - ) - return past_key_values - - @register_in_tasks_manager("deci", *["text-generation", "text-generation-with-past"], library_name="transformers") class DeciOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): DEFAULT_ONNX_OPSET = 14 @@ -2189,35 +4487,6 @@ class LlavaNextOpenVINOConfig(LlavaOpenVINOConfig): _OV_2026_1_MODEL_TYPE = "llava_next" -class DummyLLavaMultiModalProjectorInputGenerator(DummyInputGenerator): - SUPPORTED_INPUT_NAMES = ["image_features"] - - def __init__( - self, - task: str, - normalized_config: NormalizedTextConfig, - batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], - random_batch_size_range: Optional[Tuple[int, int]] = None, - **kwargs, - ): - self.task = task - - self.batch_size = batch_size - self.hidden_size = normalized_config.hidden_size - self.num_patches = (normalized_config.image_size // normalized_config.patch_size) ** 2 - self.normalized_config = normalized_config - - def generate( - self, - input_name: str, - framework: str = "pt", - int_dtype: str = "int64", - float_dtype: str = "fp32", - ): - shape = [self.batch_size, self.num_patches, self.hidden_size] - return self.random_float_tensor(shape, framework=framework, dtype=float_dtype) - - class LLavaMultimodalProjectorOpenVINOConfig(OnnxConfig): DUMMY_INPUT_GENERATOR_CLASSES = (DummyLLavaMultiModalProjectorInputGenerator,) NORMALIZED_CONFIG_CLASS = NormalizedVisionConfig @@ -2526,92 +4795,6 @@ def rename_ambiguous_inputs(self, inputs): return super().rename_ambiguous_inputs(inputs) -class PooledProjectionsDummyInputGenerator(DummyInputGenerator): - SUPPORTED_INPUT_NAMES = ["pooled_projections"] - - def __init__( - self, - task: str, - normalized_config: NormalizedConfig, - batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], - random_batch_size_range: Optional[Tuple[int, int]] = None, - **kwargs, - ): - self.task = task - self.batch_size = batch_size - self.pooled_projection_dim = normalized_config.config.pooled_projection_dim - - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - shape = [self.batch_size, self.pooled_projection_dim] - return self.random_float_tensor(shape, framework=framework, dtype=float_dtype) - - -class DummyTransformerTimestpsInputGenerator(DummyTimestepInputGenerator): - SUPPORTED_INPUT_NAMES = ("timestep", "text_embeds", "time_ids", "timestep_cond", "guidance") - - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - if input_name in ["timestep", "guidance"]: - shape = [self.batch_size] - return self.random_float_tensor(shape, max_value=self.vocab_size, framework=framework, dtype=float_dtype) - return super().generate(input_name, framework, int_dtype, float_dtype) - - -class DummyUnetVisionInputGenerator(DummyVisionInputGenerator): - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - if input_name not in ["sample", "latent_sample"]: - return super().generate(input_name, framework, int_dtype, float_dtype) - # add height and width discount for enable any resolution generation - return self.random_float_tensor( - shape=[self.batch_size, self.num_channels, self.height - 1, self.width - 1], - framework=framework, - dtype=float_dtype, - ) - - -class DummyUnetTimestepInputGenerator(DummyTimestepInputGenerator): - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - if input_name != "timestep": - return super().generate(input_name, framework, int_dtype, float_dtype) - shape = [self.batch_size] - return self.random_int_tensor(shape, max_value=self.vocab_size, framework=framework, dtype=int_dtype) - - -class DummySanaTimestepInputGenerator(DummyTimestepInputGenerator): - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - if input_name != "timestep": - return super().generate(input_name, framework, int_dtype, float_dtype) - shape = [self.batch_size] - return self.random_int_tensor(shape, max_value=self.vocab_size, framework=framework, dtype=float_dtype) - - -class DummyUnetEncoderInputGenerator(DummySeq2SeqDecoderTextInputGenerator): - def __init__( - self, - task: str, - normalized_config: NormalizedTextConfig, - batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], - sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], - num_choices: int = DEFAULT_DUMMY_SHAPES["num_choices"], - random_batch_size_range: Optional[Tuple[int, int]] = None, - random_sequence_length_range: Optional[Tuple[int, int]] = None, - random_num_choices_range: Optional[Tuple[int, int]] = None, - **kwargs, - ): - super().__init__( - task, - normalized_config, - batch_size=batch_size, - sequence_length=sequence_length, - num_choices=num_choices, - random_batch_size_range=random_batch_size_range, - random_sequence_length_range=random_sequence_length_range, - random_num_choices_range=random_num_choices_range, - **kwargs, - ) - if hasattr(normalized_config.config, "model_max_length"): - self.sequence_length = normalized_config.config.model_max_length - - @register_in_tasks_manager("unet", *["semantic-segmentation"], library_name="diffusers") @register_in_tasks_manager("unet-2d-condition", *["semantic-segmentation"], library_name="diffusers") class UNetOpenVINOConfig(UNetOnnxConfig): @@ -2678,44 +4861,6 @@ def inputs(self) -> Dict[str, Dict[int, str]]: } -class DummySanaSeq2SeqDecoderTextWithEncMaskInputGenerator(DummySeq2SeqDecoderTextInputGenerator): - SUPPORTED_INPUT_NAMES = ( - "decoder_input_ids", - "decoder_attention_mask", - "encoder_outputs", - "encoder_hidden_states", - "encoder_attention_mask", - ) - - -class DummySanaTransformerVisionInputGenerator(DummyUnetVisionInputGenerator): - SUPPORTED_INPUT_NAMES = ( - "pixel_values", - "pixel_mask", - "sample", - "latent_sample", - "guidance", - ) - - def __init__( - self, - task: str, - normalized_config: NormalizedVisionConfig, - batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], - num_channels: int = DEFAULT_DUMMY_SHAPES["num_channels"], - width: int = DEFAULT_DUMMY_SHAPES["width"] // 8, - height: int = DEFAULT_DUMMY_SHAPES["height"] // 8, - # Reduce img shape by 4 for FLUX to reduce memory usage on conversion - **kwargs, - ): - super().__init__(task, normalized_config, batch_size, num_channels, width=width, height=height, **kwargs) - - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - if input_name == "guidance": - return self.random_float_tensor([self.batch_size], framework=framework, dtype=float_dtype) - return super().generate(input_name, framework, int_dtype, float_dtype) - - @register_in_tasks_manager("sana-transformer", *["semantic-segmentation"], library_name="diffusers") class SanaTransformerOpenVINOConfig(UNetOpenVINOConfig): NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args( @@ -2787,76 +4932,6 @@ def outputs(self) -> Dict[str, Dict[int, str]]: } -class DummyFluxTransformerInputGenerator(DummyVisionInputGenerator): - SUPPORTED_INPUT_NAMES = ( - "pixel_values", - "pixel_mask", - "sample", - "latent_sample", - "hidden_states", - "img_ids", - ) - - def __init__( - self, - task: str, - normalized_config: NormalizedVisionConfig, - batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], - num_channels: int = DEFAULT_DUMMY_SHAPES["num_channels"], - width: int = DEFAULT_DUMMY_SHAPES["width"] // 4, - height: int = DEFAULT_DUMMY_SHAPES["height"] // 4, - # Reduce img shape by 4 for FLUX to reduce memory usage on conversion - **kwargs, - ): - super().__init__(task, normalized_config, batch_size, num_channels, width, height, **kwargs) - if getattr(normalized_config, "in_channels", None): - self.num_channels = normalized_config.in_channels // 4 - - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - if input_name in ["hidden_states", "sample"]: - shape = [self.batch_size, (self.height // 2) * (self.width // 2), self.num_channels * 4] - return self.random_float_tensor(shape, framework=framework, dtype=float_dtype) - if input_name == "img_ids": - img_ids_height = self.height // 2 - img_ids_width = self.width // 2 - return self.random_int_tensor( - ( - [self.batch_size, img_ids_height * img_ids_width, 3] - if is_diffusers_version("<", "0.31.0") - else [img_ids_height * img_ids_width, 3] - ), - min_value=0, - max_value=min(img_ids_height, img_ids_width), - framework=framework, - dtype=float_dtype, - ) - - return super().generate(input_name, framework, int_dtype, float_dtype) - - -class DummyFluxTextInputGenerator(DummySeq2SeqDecoderTextInputGenerator): - SUPPORTED_INPUT_NAMES = ( - "decoder_input_ids", - "decoder_attention_mask", - "encoder_outputs", - "encoder_hidden_states", - "txt_ids", - ) - - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - if input_name == "txt_ids": - import torch - - shape = ( - [self.batch_size, self.sequence_length, 3] - if is_diffusers_version("<", "0.31.0") - else [self.sequence_length, 3] - ) - dtype = DTYPE_MAPPER.pt(float_dtype) - return torch.full(shape, 0, dtype=dtype) - return super().generate(input_name, framework, int_dtype, float_dtype) - - @register_in_tasks_manager("flux-transformer", *["semantic-segmentation"], library_name="diffusers") @register_in_tasks_manager("flux-transformer-2d", *["semantic-segmentation"], library_name="diffusers") class FluxTransformerOpenVINOConfig(SD3TransformerOpenVINOConfig): @@ -2878,40 +4953,12 @@ def inputs(self): ) common_inputs["img_ids"] = ( {0: "batch_size", 1: "packed_height_width"} - if is_diffusers_version("<", "0.31.0") - else {0: "packed_height_width"} - ) - if getattr(self._normalized_config, "guidance_embeds", False): - common_inputs["guidance"] = {0: "batch_size"} - return common_inputs - - -class LTXVaeDummyInputGenerator(DummyVisionInputGenerator): - SUPPORTED_INPUT_NAMES = ("pixel_values", "pixel_mask", "sample", "latent_sample", "timestep") - - def __init__( - self, - task: str, - normalized_config: NormalizedVisionConfig, - batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], - num_channels: int = DEFAULT_DUMMY_SHAPES["num_channels"], - width: int = DEFAULT_DUMMY_SHAPES["width"], - height: int = DEFAULT_DUMMY_SHAPES["height"], - num_frames: int = 2, - **kwargs, - ): - super().__init__(task, normalized_config, batch_size, num_channels, width, height, **kwargs) - self.num_frames = num_frames - - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - if input_name in ["sample", "latent_sample"]: - return self.random_float_tensor( - [self.batch_size, self.num_channels, self.num_frames, self.height, self.width] - ) - if input_name == "timestep": - return self.random_int_tensor([1], max_value=20, min_value=1, framework=framework, dtype=int_dtype) - - return super().generate(input_name, framework, int_dtype, float_dtype) + if is_diffusers_version("<", "0.31.0") + else {0: "packed_height_width"} + ) + if getattr(self._normalized_config, "guidance_embeds", False): + common_inputs["guidance"] = {0: "batch_size"} + return common_inputs @register_in_tasks_manager("ltx-vae-encoder", *["semantic-segmentation"], library_name="diffusers") @@ -2951,53 +4998,6 @@ def outputs(self) -> Dict[str, Dict[int, str]]: } -class LTXTransformerDummyInputGenerator(DummyVisionInputGenerator): - SUPPORTED_INPUT_NAMES = ("hidden_states", "width", "height", "num_frames", "rope_interpolation_scale") - - def __init__( - self, - task: str, - normalized_config: NormalizedVisionConfig, - batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], - num_channels: int = DEFAULT_DUMMY_SHAPES["num_channels"], - width: int = 16, - height: int = 8, - num_frames: int = 2, - frame_rate: int = 10, - **kwargs, - ): - super().__init__(task, normalized_config, batch_size, num_channels, width, height, **kwargs) - self.num_frames = num_frames - self.frame_rate = frame_rate - self.vae_spatial_compression_ratio = normalized_config.config.vae_spatial_compression_ratio - self.vae_temporal_compression_ratio = normalized_config.config.vae_temporal_compression_ratio - - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - import torch - - if input_name == "hidden_states": - return self.random_float_tensor( - [self.batch_size, self.num_frames * self.height * self.width, self.num_channels] - ) - if input_name == "width": - return torch.tensor(self.width) - if input_name == "height": - return torch.tensor(self.height) - if input_name == "num_frames": - return torch.tensor(self.num_frames) - if input_name == "rope_interpolation_scale": - import torch - - return torch.tensor( - [ - self.vae_temporal_compression_ratio / self.frame_rate, - self.vae_spatial_compression_ratio, - self.vae_spatial_compression_ratio, - ] - ) - return super().generate(input_name, framework, int_dtype, float_dtype) - - @register_in_tasks_manager("ltx-video-transformer", *["semantic-segmentation"], library_name="diffusers") class LTXVideoTransformerOpenVINOConfig(SanaTransformerOpenVINOConfig): DUMMY_INPUT_GENERATOR_CLASSES = ( @@ -3026,90 +5026,6 @@ def outputs(self) -> Dict[str, Dict[int, str]]: } -class DummyMiniCPMVImageInputGenerator(DummyVisionInputGenerator): - SUPPORTED_INPUT_NAMES = ("pixel_values", "patch_attention_mask", "position_ids") - - def __init__( - self, - task: str, - normalized_config: NormalizedVisionConfig, - batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], - num_channels: int = DEFAULT_DUMMY_SHAPES["num_channels"], - width: int = DEFAULT_DUMMY_SHAPES["width"], - height: int = DEFAULT_DUMMY_SHAPES["height"], - **kwargs, - ): - super().__init__(task, normalized_config, batch_size, num_channels, width, height) - self.patch_size = normalized_config.config.patch_size - - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - if input_name == "pixel_values": - return self.random_float_tensor( - shape=[ - self.batch_size, - self.num_channels, - self.patch_size, - (self.height * self.width) // self.patch_size, - ], - framework=framework, - dtype=float_dtype, - ) - - if input_name == "patch_attention_mask": - return self.random_int_tensor( - shape=[self.batch_size, 1, (self.height // self.patch_size) * (self.width // self.patch_size)], - framework=framework, - dtype=float_dtype, - min_value=0, - max_value=2, - ) - - if input_name == "position_ids": - return self.random_int_tensor( - shape=[self.batch_size, (self.height // self.patch_size) * (self.width // self.patch_size)], - max_value=self.patch_size, - ) - - -class DummyMiniCPMVResampleInputGenerator(DummyVisionInputGenerator): - SUPPORTED_INPUT_NAMES = ("image_feature", "pos_embed", "key_padding_mask") - - def __init__( - self, - task: str, - normalized_config: NormalizedVisionConfig, - batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], - num_channels: int = DEFAULT_DUMMY_SHAPES["num_channels"], - width: int = DEFAULT_DUMMY_SHAPES["width"], - height: int = DEFAULT_DUMMY_SHAPES["height"], - **kwargs, - ): - super().__init__(task, normalized_config, batch_size, num_channels, width, height) - self.patch_size = normalized_config.config.patch_size - self.hidden_size = normalized_config.config.hidden_size - self.img_hidden_size = normalized_config.config.vision_config.hidden_size - self.feat_size = (normalized_config.config.vision_config.image_size // self.patch_size) * ( - normalized_config.config.vision_config.image_size // self.patch_size - ) - - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - if input_name == "image_feature": - return self.random_float_tensor( - shape=[self.batch_size, self.feat_size, self.img_hidden_size], framework=framework, dtype=float_dtype - ) - - if input_name == "key_padding_mask": - return self.constant_tensor( - shape=[self.batch_size, self.feat_size], - framework=framework, - value=1, - dtype=DTYPE_MAPPER.pt(float_dtype), - ) - - if input_name == "pos_embed": - return self.random_float_tensor(shape=[self.feat_size, self.batch_size, self.hidden_size]) - - class MiniCPMVConfigBehavior(str, enum.Enum): RESAMPLER = "resampler" LANGUAGE = "language" @@ -3259,55 +5175,6 @@ class Phi3VisionConfigBehavior(str, enum.Enum): TEXT_EMBEDDINGS = "text_embeddings" -class DummyPhi3VisionProjectionInputGenerator(DummyVisionInputGenerator): - SUPPORTED_INPUT_NAMES = ("input",) - - def __init__( - self, - task: str, - normalized_config: NormalizedVisionConfig, - batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], - num_channels: int = DEFAULT_DUMMY_SHAPES["num_channels"], - width: int = 336, - height: int = 336, - crop_size=336, - **kwargs, - ): - self.batch_size = batch_size - self._embed_layer_realization = ( - normalized_config.config.embd_layer["embedding_cls"] - if hasattr(normalized_config.config, "embd_layer") - else "image_audio" - ) - if not hasattr(normalized_config.config, "vision_config"): - self.image_dim_out = ( - normalized_config.config.img_processor.get( - "image_dim_out", normalized_config.config.img_processor.get("hidden_size") - ) - if normalized_config.config.img_processor is not None - else 1152 - ) - if "image_embd_layer" in normalized_config.config.embd_layer: - self.crop_size = normalized_config.config.embd_layer["image_embd_layer"].get("crop_size", crop_size) - else: - self.crop_size = normalized_config.config.embd_layer.get("crop_size", crop_size) - else: - self.image_dim_out = normalized_config.config.vision_config.hidden_size - self.crop_size = normalized_config.config.vision_config.crop_size - self.height = height - self.width = width - - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - h = self.height // self.crop_size - w = self.width // self.crop_size - feat_size = (h * w + 1) * 144 + 1 + (h + 1) * 12 - if self._embed_layer_realization in ["linear", "image_audio"]: - shape = [self.batch_size, feat_size, self.image_dim_out] - else: - shape = [self.batch_size, feat_size, self.image_dim_out * 4] - return self.random_float_tensor(shape, framework=framework, dtype=float_dtype) - - @register_in_tasks_manager("phi3_v", *["image-text-to-text"], library_name="transformers") class Phi3VisionOpenVINOConfig(BaseVLMOpenVINOConfig): SUPPORTED_BEHAVIORS = [model_type.value for model_type in Phi3VisionConfigBehavior] @@ -3426,119 +5293,6 @@ def patch_model_for_export(self, model: PreTrainedModel, model_kwargs: Optional[ return super().patch_model_for_export(model, model_kwargs) -class DummyAudioPhi4MMInputGenerator(DummyInputGenerator): - SUPPORTED_INPUT_NAMES = ("audio_input", "audio_feature", "audio_mask") - - def __init__( - self, - task: str, - normalized_config: NormalizedVisionConfig, - batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], - signal_length=498, - **kwargs, - ): - self.signal_length = signal_length - if hasattr(normalized_config.config, "audio_processor"): - self.audio_chunk_size = ( - signal_length // normalized_config.config.audio_processor["config"]["time_reduction"] + 1 - ) - self.input_size = normalized_config.config.audio_processor["config"]["input_size"] - self.attention_dim = normalized_config.config.audio_processor["config"]["attention_dim"] - else: - self.audio_chunk_size = signal_length // normalized_config.config.audio_config.time_reduction + 1 - self.input_size = normalized_config.config.audio_config.input_size - self.attention_dim = normalized_config.config.audio_config.hidden_size - self.batch_size = batch_size - self.task = task - self.normalized_config = normalized_config - - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - if input_name == "audio_input": - return self.random_float_tensor( - [self.batch_size, self.signal_length, self.input_size], framework=framework, dtype=float_dtype - ) - - if input_name == "audio_feature": - return self.random_float_tensor( - [self.batch_size, self.audio_chunk_size, self.attention_dim], framework=framework, dtype=float_dtype - ) - - if input_name == "audio_mask": - return self.random_int_tensor( - [self.batch_size, self.audio_chunk_size, self.audio_chunk_size], - max_value=2, - framework=framework, - dtype="bool", - ) - - -class DummyVisionPositionIdsPhi4InputGenerator(DummyVisionInputGenerator): - SUPPORTED_INPUT_NAMES = ("patch_position_ids", "patch_attention_mask") - - def __init__( - self, - task: str, - normalized_config: NormalizedVisionConfig, - batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], - num_channels: int = DEFAULT_DUMMY_SHAPES["num_channels"], - width: int = DEFAULT_DUMMY_SHAPES["width"], - height: int = DEFAULT_DUMMY_SHAPES["height"], - **kwargs, - ): - super().__init__(task, normalized_config, batch_size, num_channels, width, height, **kwargs) - if hasattr(normalized_config.config, "vision_conifg"): - self.patch_size = getattr(normalized_config.config.vision_config, "patch_size", 14) - else: - self.patch_size = 14 - self.num_patches_per_side = self.height // self.patch_size - - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - if input_name == "patch_position_ids": - return self.get_vision_position_ids() - if input_name == "patch_attention_mask": - return self.random_int_tensor( - [self.batch_size, self.height // self.patch_size, self.width // self.patch_size], - framework=framework, - dtype="bool", - max_value=2, - ) - return super().generate(input_name, framework, int_dtype, float_dtype) - - def get_vision_position_ids(self): - # Adopted from https://github.com/huggingface/transformers/blob/v4.51.3/src/transformers/models/phi4_multimodal/modeling_phi4_multimodal.py#L494-L512 - import torch - - batch_size = self.batch_size - max_im_h, max_im_w = self.height, self.width - max_nb_patches_h, max_nb_patches_w = max_im_h // self.patch_size, max_im_w // self.patch_size - boundaries = torch.arange(1 / self.num_patches_per_side, 1.0, 1 / self.num_patches_per_side) - position_ids = torch.full( - size=( - batch_size, - max_nb_patches_h * max_nb_patches_w, - ), - fill_value=0, - ) - patch_attention_mask = torch.ones( - [self.batch_size, self.height // self.patch_size, self.width // self.patch_size], dtype=torch.int64 - ) - patch_attention_mask[0, self.height - 2 :] = 0 - - for batch_idx, p_attn_mask in enumerate(patch_attention_mask): - nb_patches_h = p_attn_mask[:, 0].sum() - nb_patches_w = p_attn_mask[0].sum() - - fractional_coords_h = torch.arange(0, 1 - 1e-6, 1 / nb_patches_h) - fractional_coords_w = torch.arange(0, 1 - 1e-6, 1 / nb_patches_w) - - bucket_coords_h = torch.bucketize(fractional_coords_h, boundaries, right=True) - bucket_coords_w = torch.bucketize(fractional_coords_w, boundaries, right=True) - - pos_ids = (bucket_coords_h[:, None] * self.num_patches_per_side + bucket_coords_w).flatten() - position_ids[batch_idx][p_attn_mask.view(-1)] = pos_ids - return position_ids - - class Phi4MMConfigBehavior(str, enum.Enum): AUDIO_EMBEDDINGS = "audio_embeddings" AUDIO_ENCODER = "audio_encoder" @@ -3780,123 +5534,6 @@ def rename_ambiguous_inputs(self, inputs): return inputs -class DummyQwen2VLLMInputGenerator(DummyTextInputGenerator): - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - generated_input = super().generate(input_name, framework, int_dtype, float_dtype) - if input_name == "position_ids": - return generated_input.unsqueeze(0).expand(3, -1, -1) - return generated_input - - -class DummyQwen3_5LMInputGenerator(DummyTextInputGenerator): - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - generated_input = super().generate(input_name, framework, int_dtype, float_dtype) - if input_name == "position_ids": - return generated_input.unsqueeze(0).expand(4, -1, -1) - return generated_input - - -class DummyQwen2VLVisionEmbedInputGenerator(DummyVisionInputGenerator): - SUPPORTED_INPUT_NAMES = ( - "hidden_states", - "attention_mask", - "window_attention_mask", - "window_index", - "rotary_pos_emb", - ) - - def __init__( - self, - task: str, - normalized_config: NormalizedVisionConfig, - batch_size: int = 1, - num_channels: int = DEFAULT_DUMMY_SHAPES["num_channels"], - width: int = 420, - height: int = 420, - **kwargs, - ): - self.batch_size = batch_size - self.height = height - self.width = width - self.num_channels = num_channels - self.temporal_patch_size = normalized_config.config.temporal_patch_size - self.patch_size = normalized_config.config.patch_size - if normalized_config.use_embed_dim: - self.embed_dim = ( - normalized_config.config.embed_dim - if hasattr(normalized_config.config, "embed_dim") - else normalized_config.hidden_size - ) - else: - self.embed_dim = self.num_channels * self.temporal_patch_size * self.patch_size * self.patch_size - self.num_heads = normalized_config.config.num_heads - self.spatial_merge_size = None - if hasattr(normalized_config.config, "spatial_merge_size"): - self.spatial_merge_size = normalized_config.config.spatial_merge_size - - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - grid_h, grid_w = self.height // self.patch_size, self.width // self.patch_size - grid_t = self.batch_size - - if input_name == "hidden_states": - return self.random_float_tensor( - [grid_t * grid_h * grid_w, self.embed_dim], framework=framework, dtype=float_dtype - ) - - if input_name in ["attention_mask", "window_attention_mask"]: - return self.random_mask_tensor( - [1, grid_t * grid_h * grid_w, grid_t * grid_h * grid_w], framework=framework, dtype=float_dtype - ) - - if input_name == "rotary_pos_emb": - dim = self.embed_dim // self.num_heads // 2 - return self.random_float_tensor([grid_h * grid_t * grid_w, dim], framework=framework, dtype=float_dtype) - - if input_name == "window_index": - if self.spatial_merge_size is None: - raise ValueError( - "`spatial_merge_size` parameter is not found in model config. Can not generate dummy input data for `window_index` input" - ) - spatial_merge_unit = self.spatial_merge_size * self.spatial_merge_size - hidden_size = (grid_t * grid_h * grid_w) // spatial_merge_unit - return self.random_int_tensor([hidden_size], max_value=hidden_size) - - -class DummyQwen3VLVisionEmbedInputGenerator(DummyQwen2VLVisionEmbedInputGenerator): - SUPPORTED_INPUT_NAMES = ( - "hidden_states", - "attention_mask", - "rotary_pos_emb", - "input", - ) - - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - grid_h, grid_w = self.height // self.patch_size, self.width // self.patch_size - grid_t = self.batch_size - - if input_name == "hidden_states": - return self.random_float_tensor( - [grid_t * grid_h * grid_w, self.embed_dim], framework=framework, dtype=float_dtype - ) - - if input_name in ["attention_mask"]: - return self.random_mask_tensor( - [1, grid_t * grid_h * grid_w, grid_t * grid_h * grid_w], framework=framework, dtype=float_dtype - ) - - if input_name == "rotary_pos_emb": - dim = self.embed_dim // self.num_heads // 2 - return self.random_float_tensor([grid_h * grid_t * grid_w, dim], framework=framework, dtype=float_dtype) - - if input_name == "input": - return self.constant_tensor( - [4, DEFAULT_DUMMY_SHAPES["sequence_length"]], - framework=framework, - value=0, - dtype=DTYPE_MAPPER.pt(int_dtype), - ) - - class QwenVLConfigBehavior(str, enum.Enum): LANGUAGE = "language" VISION_EMBEDDINGS = "vision_embeddings" @@ -4259,37 +5896,6 @@ class WhisperOpenVINOConfig(WhisperOnnxConfig): _MODEL_PATCHER = OVSeq2SeqModelPatcher -class Qwen3ASRDummySeq2SeqPastKeyValuesGenerator(DummySeq2SeqPastKeyValuesGenerator): - """Custom KV cache generator for Qwen3-ASR with GQA (num_key_value_heads != num_attention_heads). - Qwen3-ASR has no cross-attention, so only self-attention KV cache is generated (2 per layer).""" - - def __init__(self, task, normalized_config, **kwargs): - super().__init__(task, normalized_config, **kwargs) - # Override head count and head_dim for GQA - self.decoder_num_attention_heads = normalized_config.decoder_num_attention_heads - self.decoder_head_dim = getattr(normalized_config, "head_dim", None) - if self.decoder_head_dim is None: - self.decoder_head_dim = self.decoder_hidden_size // normalized_config.num_attention_heads - - def generate(self, input_name, framework="pt", int_dtype="int64", float_dtype="fp32"): - if input_name == "past_key_values": - decoder_shape = ( - self.batch_size, - self.decoder_num_attention_heads, - self.sequence_length, - self.decoder_head_dim, - ) - # Qwen3-ASR has no cross-attention, only self-attention KV cache - return [ - ( - self.random_float_tensor(decoder_shape, framework=framework, dtype=float_dtype), - self.random_float_tensor(decoder_shape, framework=framework, dtype=float_dtype), - ) - for _ in range(self.decoder_num_layers) - ] - return super().generate(input_name, framework=framework, int_dtype=int_dtype, float_dtype=float_dtype) - - @register_in_tasks_manager( "qwen3_asr", *[ @@ -4573,54 +6179,6 @@ class Gemma4ConfigBehavior(str, enum.Enum): TEXT_EMBEDDINGS_PER_LAYER = "text_embeddings_per_layer" -class DummyGemma4VisionInputGenerator(DummyVisionInputGenerator): - SUPPORTED_INPUT_NAMES = ("pixel_values", "image_position_ids") - - def __init__(self, task, normalized_config, batch_size=DEFAULT_DUMMY_SHAPES["batch_size"], **kwargs): - super().__init__(task, normalized_config, batch_size, **kwargs) - self.patch_size = getattr(normalized_config, "patch_size", 16) - self.pooling_kernel_size = getattr(normalized_config, "pooling_kernel_size", 3) - # Gemma4 processor always pads pixel_values to max_soft_tokens * pooling_kernel_size^2 patches. - # The vision model's pooling uses shape-dependent Python operations that get baked in during tracing, - # so the dummy input must match the actual inference shapes. - max_soft_tokens = getattr(normalized_config, "image_seq_length", None) - if max_soft_tokens is None: - max_soft_tokens = getattr(normalized_config, "max_soft_tokens", 280) - self.num_patches = max_soft_tokens * self.pooling_kernel_size**2 - - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - if input_name == "pixel_values": - # Gemma4 expects pre-patchified pixel_values: [batch, num_patches, 3 * patch_size^2] - return self.random_float_tensor( - shape=[self.batch_size, self.num_patches, 3 * self.patch_size**2], - framework=framework, - dtype=float_dtype, - ) - if input_name == "image_position_ids": - # Create position ids as a grid. The patch count = h_patches * w_patches - # where both are divisible by pooling_kernel_size for correct pooling. - k = self.pooling_kernel_size - total_pooled = self.num_patches // (k * k) - # Find roughly square grid for pooled side - pooled_side = int(math.sqrt(total_pooled)) - if pooled_side * pooled_side < total_pooled: - pooled_h = pooled_side - pooled_w = total_pooled // pooled_h - else: - pooled_h = pooled_w = pooled_side - h_patches = pooled_h * k - w_patches = pooled_w * k - pos_ids = torch.stack( - torch.meshgrid(torch.arange(h_patches), torch.arange(w_patches), indexing="ij"), dim=-1 - ).reshape(1, -1, 2) - # Pad to num_patches with -1 (padding position) - if pos_ids.shape[1] < self.num_patches: - pad = torch.full((1, self.num_patches - pos_ids.shape[1], 2), -1, dtype=pos_ids.dtype) - pos_ids = torch.cat([pos_ids, pad], dim=1) - return pos_ids.expand(self.batch_size, -1, -1).clone() - return super().generate(input_name, framework, int_dtype, float_dtype) - - @register_in_tasks_manager("gemma4", *["image-text-to-text"], library_name="transformers") class Gemma4OpenVINOConfig(Gemma3OpenVINOConfig): SUPPORTED_BEHAVIORS = [model_type.value for model_type in Gemma4ConfigBehavior] @@ -4808,35 +6366,6 @@ def outputs(self) -> Dict[str, Dict[int, str]]: return super().outputs -class DummyVisionPositionIdsInputGenerator(DummyVisionInputGenerator): - SUPPORTED_INPUT_NAMES = ("patch_attention_mask", "patch_position_ids") - - def __init__( - self, - task: str, - normalized_config: NormalizedVisionConfig, - batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], - num_channels: int = DEFAULT_DUMMY_SHAPES["num_channels"], - width: int = DEFAULT_DUMMY_SHAPES["width"], - height: int = DEFAULT_DUMMY_SHAPES["height"], - **kwargs, - ): - super().__init__(task, normalized_config, batch_size, num_channels, width, height, **kwargs) - self.patch_size = normalized_config.config.patch_size - - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - if input_name == "patch_attention_mask": - shape = [self.batch_size, self.height // self.patch_size, self.width // self.patch_size] - return self.random_int_tensor(shape, max_value=2, framework=framework, dtype="bool") - if input_name == "patch_position_ids": - max_nb_patches_h, max_nb_patches_w = self.height // self.patch_size, self.width // self.patch_size - shape = [self.batch_size, max_nb_patches_h * max_nb_patches_w] - return self.random_int_tensor( - shape, max_value=min(max_nb_patches_h, max_nb_patches_w), framework=framework, dtype=int_dtype - ) - return super().generate(input_name, framework, int_dtype, float_dtype) - - @register_in_tasks_manager("idefics3", *["image-text-to-text"], library_name="transformers") class Idefics3OpenVINOConfig(BaseVLMOpenVINOConfig): DUMMY_INPUT_GENERATOR_CLASSES = (DummyVisionInputGenerator, DummyVisionPositionIdsInputGenerator) @@ -4975,55 +6504,6 @@ class MarianOpenVINOConfig(MarianOnnxConfig): MAX_TRANSFORMERS_VERSION = "4.57.6" -class DummySpeechT5OpenVINOInputGenerator(DummyInputGenerator): - SUPPORTED_INPUT_NAMES = ( - "inputs_embeds", - "output_sequence", - "speaker_embeddings", - "spectrogram", - "raw_spectrogram", - "encoder_hidden_states", - ) - - def __init__( - self, - task: str, - normalized_config: NormalizedConfig, - sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], - **kwargs, - ): - self.task = task - self.batch_size = 1 - - self.sequence_length = sequence_length - self.speaker_embedding_dim = normalized_config.speaker_embedding_dim - self.num_mel_bins = normalized_config.num_mel_bins - self.reduction_factor = normalized_config.config.reduction_factor - self.hidden_size = normalized_config.config.hidden_size - - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - if input_name in ["output_sequence", "inputs_embeds"]: - shape = [self.batch_size, self.sequence_length, self.num_mel_bins] - elif input_name == "speaker_embeddings": - shape = [self.batch_size, self.speaker_embedding_dim] - elif input_name == "raw_spectrogram": - shape = [self.sequence_length, self.batch_size, self.reduction_factor, self.num_mel_bins] - elif input_name == "encoder_hidden_states": - shape = [self.batch_size, self.sequence_length, self.hidden_size] - elif input_name == "spectrogram": - shape = [self.batch_size, self.sequence_length, self.num_mel_bins] - else: - raise ValueError(f"Unsupported input {input_name} for DummySpeechT5InputGenerator") - - return self.random_float_tensor( - shape=shape, - min_value=0, - max_value=1, - framework=framework, - dtype=float_dtype, - ) - - class SpeechT5ConfigBehavior(str, enum.Enum): ENCODER = "encoder" DECODER = "decoder" @@ -5202,50 +6682,6 @@ def patch_model_for_export(self, model: PreTrainedModel, model_kwargs: Optional[ return Llama4ImageEmbeddingsModelPatcher(self, model, model_kwargs) -class MambaCacheDummyInputGenerator(DummyInputGenerator): - """ - Generates dummy past_ssm_states, past_conv_states and cache_position inputs for Mamba architectures. - """ - - SUPPORTED_INPUT_NAMES = ("cache_params", "cache_position") - - def __init__( - self, - task: str, - normalized_config, - batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], - sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], - **kwargs, - ): - self.normalized_config = normalized_config - self.batch_size = batch_size - self.sequence_length = sequence_length - self.intermediate_size = self.normalized_config.config.intermediate_size - self.ssm_state_size = self.normalized_config.config.state_size - self.conv_kernel_size = self.normalized_config.config.conv_kernel - - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - if input_name == "cache_params": - ssm_shape = [self.batch_size, self.intermediate_size, self.ssm_state_size] - conv_shape = [self.batch_size, self.intermediate_size, self.conv_kernel_size] - return [ - ( - self.random_float_tensor(ssm_shape, framework=framework, dtype=float_dtype), - self.random_float_tensor(conv_shape, framework=framework, dtype=float_dtype), - ) - for _ in range(self.normalized_config.num_layers) - ] - elif input_name == "cache_position": - return self.random_int_tensor( - shape=[self.conv_kernel_size], - max_value=self.sequence_length, - framework=framework, - dtype=int_dtype, - ) - - raise ValueError(f"Unsupported input name {input_name}") - - @register_in_tasks_manager( "falcon_mamba", *["text-generation", "text-generation-with-past"], library_name="transformers" ) @@ -5352,82 +6788,6 @@ class VisionEncoderDecoderOpenVINOConfig(VisionEncoderDecoderOnnxConfig): _MODEL_PATCHER = OVSeq2SeqModelPatcher -class Zamba2DummyPastKeyValuesGenerator(DummyPastKeyValuesGenerator): - """ - Generates dummy cache_params inputs for Zamba2 architectures. - """ - - SUPPORTED_INPUT_NAMES = ("cache_params",) - - def __init__( - self, - task: str, - normalized_config, - 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, - ) - - config = normalized_config.config - self.intermediate_size = int(config.mamba_expand * config.hidden_size) - self.conv_kernel_size = config.mamba_d_conv - self.mamba_d_state = config.mamba_d_state - if config.model_type == "zamba2": - self.n_mamba_heads = config.n_mamba_heads - self.mamba_ngroups = config.mamba_ngroups - self.mamba_headdim = config.mamba_headdim - self.head_dim = config.attention_head_dim - # in Zamba2, all layers contain Mamba block - # some of these layers are hybrid so they contain both attention and mamba blocks - self.num_attention_layers = len(config.hybrid_layer_ids) - self.num_mamba_layers = self.num_layers - logger.warning( - "The current support for the 'Zamba2' model type is experimental. " - "Performance is not optimal with high memory consumption. " - "Optimizations and improved support will be available in a future OpenVINO release." - ) - else: - # currently, this else-branch is applied for GraniteMoeHybrid models - self.n_mamba_heads = config.mamba_n_heads - self.mamba_ngroups = config.mamba_n_groups - self.mamba_headdim = config.mamba_d_head - self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) - self.num_attention_layers = config.layer_types.count("attention") - self.num_mamba_layers = config.layer_types.count("mamba") - self.num_attention_heads = config.num_key_value_heads - self.sequence_length = 0 - - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - past_key_values = [] - for i in range(self.num_mamba_layers): - conv_state_shape = ( - self.batch_size, - self.intermediate_size + 2 * self.mamba_ngroups * self.mamba_d_state, - self.conv_kernel_size, - ) - conv_state = self.random_float_tensor(conv_state_shape, framework=framework, dtype=float_dtype) - past_key_values.append(conv_state) - ssm_state_shape = (self.batch_size, self.n_mamba_heads, self.mamba_headdim, self.mamba_d_state) - ssm_state = self.random_float_tensor(ssm_state_shape, framework=framework, dtype=float_dtype) - past_key_values.append(ssm_state) - - for i in range(self.num_attention_layers): - kv_shape = (self.batch_size, self.num_attention_heads, self.sequence_length, self.head_dim) - k = self.random_float_tensor(kv_shape, framework=framework, dtype=float_dtype) - v = self.random_float_tensor(kv_shape, framework=framework, dtype=float_dtype) - past_key_values.append(k) - past_key_values.append(v) - - return past_key_values - - @register_in_tasks_manager("zamba2", *["text-generation", "text-generation-with-past"], library_name="transformers") class Zamba2OpenVINOConfig(MambaOpenVINOConfig): PAD_ATTENTION_MASK_TO_PAST = False @@ -5475,63 +6835,6 @@ def inputs(self) -> Dict[str, Dict[int, str]]: return common_inputs -class Lfm2DummyPastKeyValuesGenerator(DummyPastKeyValuesGenerator): - """ - Generates dummy past_key_values inputs for Lfm2 architectures. - """ - - SUPPORTED_INPUT_NAMES = ("cache_params",) - - def __init__( - self, - task: str, - normalized_config, - 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, - ) - config = normalized_config.config - self.num_conv_layers = config.layer_types.count("conv") - self.num_atten_layers = config.layer_types.count("full_attention") - self.batch_size = batch_size - self.normalized_config = normalized_config - self.hidden_size = self.normalized_config.hidden_size - self.conv_L_cache = self.normalized_config.conv_L_cache - self.num_key_value_heads = self.normalized_config.num_key_value_heads - self.num_hidden_layers = self.normalized_config.num_hidden_layers - - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - past_key_values = [] - - for i in range(self.num_conv_layers): - conv_state_shape = (self.batch_size, self.hidden_size, self.conv_L_cache) - conv_state = self.random_float_tensor(conv_state_shape, framework=framework, dtype=float_dtype) - past_key_values.append(conv_state) - - for i in range(self.num_atten_layers): - shape = ( - self.batch_size, - self.num_key_value_heads, - self.sequence_length, - self.hidden_size // self.num_attention_heads, - ) - - kv_shape = shape # (self.batch_size, self.num_attention_heads, self.sequence_length, self.head_dim) - k = self.random_float_tensor(kv_shape, framework=framework, dtype=float_dtype) - v = self.random_float_tensor(kv_shape, framework=framework, dtype=float_dtype) - past_key_values.append(k) - past_key_values.append(v) - - return past_key_values - - @register_in_tasks_manager( "lfm2", *[ @@ -5978,80 +7281,6 @@ class SiglipTextOpenVINOConfig(SiglipTextOnnxConfig): pass -class DummyVideoChatFlashQwenInputGenerator(DummyVisionInputGenerator): - SUPPORTED_INPUT_NAMES = ("hidden_states", "rotary_pos_emb") - - def __init__( - self, - task: str, - normalized_config: NormalizedVisionConfig, - batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], - num_channels: int = DEFAULT_DUMMY_SHAPES["num_channels"], - width: int = DEFAULT_DUMMY_SHAPES["width"], - height: int = DEFAULT_DUMMY_SHAPES["height"], - visual_seq_length: int = DEFAULT_DUMMY_SHAPES["visual_seq_length"], - **kwargs, - ): - super().__init__(task, normalized_config, batch_size, num_channels, width, height, visual_seq_length, **kwargs) - self.num_frames = normalized_config.config.mm_local_num_frames - self.embed_dim = normalized_config.config.mm_hidden_size - self.height = normalized_config.config.image_size - self.width = normalized_config.config.image_size - self.image_size = (self.height, self.width) - self.patch_size = normalized_config.config.patch_size - - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - if input_name == "hidden_states": - return self.random_float_tensor( - shape=[ - self.batch_size, - self.num_channels, - self.num_frames, - self.height, - self.width, - ], - framework=framework, - dtype=float_dtype, - ) - elif input_name == "rotary_pos_emb": - grid_h, grid_w = self.height // self.patch_size, self.width // self.patch_size - grid_t = self.num_frames - # Source: https://huggingface.co/OpenGVLab/VideoChat-Flash-Qwen2_5-7B_InternVideo2-1B/blob/main/vision_tower_builder.py#L523 - # The first dimension of rotary_pos_emb is fixed to 1 in the original model. - # And the second dimension is the total number of tokens for all frames, which is calculated as grid_h * grid_w * grid_t plus 1 for the cls token. - return self.random_float_tensor( - [1, 1 + grid_h * grid_t * grid_w, self.embed_dim], framework=framework, dtype=float_dtype - ) - - -class DummyVideoChatFlashQwenProjectorInputGenerator(DummyInputGenerator): - SUPPORTED_INPUT_NAMES = ["input"] - - def __init__( - self, - task: str, - normalized_config: NormalizedTextConfig, - batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], - random_batch_size_range: Optional[Tuple[int, int]] = None, - **kwargs, - ): - self.task = task - self.batch_size = batch_size - self.hidden_size = normalized_config.config.mm_hidden_size - self.num_patches = normalized_config.config.mm_projector_num_tome_tokens - self.normalized_config = normalized_config - - def generate( - self, - input_name: str, - framework: str = "pt", - int_dtype: str = "int64", - float_dtype: str = "fp32", - ): - shape = [self.batch_size, self.num_patches, self.hidden_size] - return self.random_float_tensor(shape, framework=framework, dtype=float_dtype) - - class VideoChatFlashQwenProjectorOpenVINOConfig(OnnxConfig): DUMMY_INPUT_GENERATOR_CLASSES = (DummyVideoChatFlashQwenProjectorInputGenerator,) NORMALIZED_CONFIG_CLASS = NormalizedVisionConfig @@ -6204,68 +7433,6 @@ class HunyuanV1DenseOpenVINOConfig(LlamaOpenVINOConfig): DUMMY_PKV_GENERATOR_CLASS = GemmaDummyPastKeyValuesGenerator -class Qwen3NextDummyPastKeyValuesGenerator(DummyPastKeyValuesGenerator): - """ - Generates dummy cache_params inputs for Qwen3-Next architectures. - """ - - SUPPORTED_INPUT_NAMES = ("cache_params",) - - def __init__( - self, - task: str, - normalized_config, - 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, - ) - - config = normalized_config.config - self.num_full_attn_layers = config.layer_types.count("full_attention") - self.num_linear_attn_layers = config.layer_types.count("linear_attention") - self.conv_kernel_size = config.linear_conv_kernel_dim - self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) - self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads - self.head_k_dim = config.linear_key_head_dim - self.head_v_dim = config.linear_value_head_dim - self.num_v_heads = config.linear_num_value_heads - self.num_k_heads = config.linear_num_key_heads - self.num_key_value_heads = config.num_key_value_heads - - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - cache_params = [] - - for idx in range(self.num_linear_attn_layers): - d_inner = self.num_k_heads * (2 * self.head_k_dim + self.head_v_dim * self.num_v_heads // self.num_k_heads) - conv_state_shape = ( - self.batch_size, - d_inner, - self.conv_kernel_size, - ) - conv_state = self.random_float_tensor(conv_state_shape, framework=framework, dtype=float_dtype) - cache_params.append(conv_state) - num_heads = self.num_v_heads - recurrent_state_shape = (self.batch_size, num_heads, self.head_k_dim, self.head_v_dim) - recurrent_state = self.random_float_tensor(recurrent_state_shape, framework=framework, dtype=float_dtype) - cache_params.append(recurrent_state) - - for idx in range(self.num_full_attn_layers): - kv_shape = (self.batch_size, self.num_key_value_heads, self.sequence_length, self.head_dim) - k = self.random_float_tensor(kv_shape, framework=framework, dtype=float_dtype) - v = self.random_float_tensor(kv_shape, framework=framework, dtype=float_dtype) - cache_params.append(k) - cache_params.append(v) - - return cache_params - - @register_in_tasks_manager( "qwen3_next", *["text-generation", "text-generation-with-past"], @@ -6354,67 +7521,6 @@ class LFM2MoeOpenVINOConfig(LFM2OpenVINOConfig): _MODEL_PATCHER = Lfm2MoeModelPatcher -class Qwen3_5DummyPastKeyValuesGenerator(DummyPastKeyValuesGenerator): - """ - Generates dummy cache_params inputs for Qwen3.5 architectures. - """ - - SUPPORTED_INPUT_NAMES = ("cache_params",) - - def __init__( - self, - task: str, - normalized_config, - 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, - ) - - config = normalized_config.config - self.num_full_attn_layers = config.layer_types.count("full_attention") - self.num_linear_attn_layers = config.layer_types.count("linear_attention") - self.conv_kernel_size = config.linear_conv_kernel_dim - self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) - self.head_k_dim = config.linear_key_head_dim - self.head_v_dim = config.linear_value_head_dim - self.num_v_heads = config.linear_num_value_heads - self.num_k_heads = config.linear_num_key_heads - self.num_key_value_heads = config.num_key_value_heads - - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - cache_params = [] - - for idx in range(self.num_linear_attn_layers): - d_inner = self.num_k_heads * (2 * self.head_k_dim + self.head_v_dim * self.num_v_heads // self.num_k_heads) - conv_state_shape = ( - self.batch_size, - d_inner, - self.conv_kernel_size, - ) - conv_state = self.random_float_tensor(conv_state_shape, framework=framework, dtype=float_dtype) - cache_params.append(conv_state) - num_heads = self.num_v_heads - recurrent_state_shape = (self.batch_size, num_heads, self.head_k_dim, self.head_v_dim) - recurrent_state = self.random_float_tensor(recurrent_state_shape, framework=framework, dtype=float_dtype) - cache_params.append(recurrent_state) - - for idx in range(self.num_full_attn_layers): - kv_shape = (self.batch_size, self.num_key_value_heads, self.sequence_length, self.head_dim) - k = self.random_float_tensor(kv_shape, framework=framework, dtype=float_dtype) - v = self.random_float_tensor(kv_shape, framework=framework, dtype=float_dtype) - cache_params.append(k) - cache_params.append(v) - - return cache_params - - @register_in_tasks_manager( "qwen3_5_text", *["text-generation", "text-generation-with-past"], @@ -6649,43 +7755,6 @@ def outputs(self) -> Dict[str, Dict[int, str]]: return super().outputs -class DummyKokoroInputGenerator(DummyInputGenerator): - """Generates dummy inputs for the Kokoro TTS model.""" - - SUPPORTED_INPUT_NAMES = ("input_ids", "ref_s", "speed") - - def __init__( - self, - task: str, - normalized_config: NormalizedConfig, - sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], - **kwargs, - ): - self.task = task - self.batch_size = 1 - self.sequence_length = sequence_length - self.style_dim = getattr(normalized_config, "style_dim", 128) - - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - if input_name == "input_ids": - shape = [self.batch_size, self.sequence_length] - input_ids_value = self.random_int_tensor( - shape=shape, min_value=0, max_value=178, framework=framework, dtype=int_dtype - ) - input_ids_value[:, 0] = 0 - input_ids_value[:, -1] = 0 - return input_ids_value - elif input_name == "ref_s": - shape = [self.batch_size, self.style_dim * 2] - return self.random_float_tensor( - shape=shape, min_value=-1, max_value=1, framework=framework, dtype=float_dtype - ) - elif input_name == "speed": - return self.random_int_tensor(shape=[1], min_value=1, max_value=10, framework=framework, dtype=float_dtype) - else: - raise ValueError(f"Unsupported input {input_name} for DummyKokoroInputGenerator") - - @register_in_tasks_manager( "kokoro", *["text-to-audio"], diff --git a/optimum/exporters/openvino/model_patcher.py b/optimum/exporters/openvino/model_patcher.py index 6fcc9de4de..d7caa35b4d 100644 --- a/optimum/exporters/openvino/model_patcher.py +++ b/optimum/exporters/openvino/model_patcher.py @@ -24,6 +24,7 @@ import torch import torch.nn.functional as F from torch import nn +from transformers import PreTrainedModel from transformers.cache_utils import Cache, DynamicCache, EncoderDecoderCache from transformers.configuration_utils import PretrainedConfig from transformers.generation import GenerationMixin @@ -46,16 +47,7 @@ from transformers.processing_utils import Unpack from transformers.utils import ModelOutput -from optimum.exporters.onnx.base import OnnxConfig -from optimum.exporters.onnx.model_patcher import ( - UNSUPPORTED_OPS_PATCHING_SPEC, - ModelPatcher, - gpt_oss_forward, - override_arguments, -) -from optimum.exporters.onnx.model_patcher import ( - sdpa_mask_without_vmap as _orig_sdpa_mask_without_vmap, -) +# from optimum.exporters.openvino.base import OnnxConfig from optimum.intel.utils.import_utils import ( is_diffusers_version, is_openvino_version, @@ -63,34 +55,1490 @@ is_transformers_version, ) -from ._ov_ops import convert_recurrent_attention_cell +from ._ov_ops import convert_recurrent_attention_cell + + +if is_transformers_version(">=", "4.53"): + from transformers.masking_utils import ALL_MASK_ATTENTION_FUNCTIONS, eager_mask, sdpa_mask + from transformers.models.qwen3_moe.modeling_qwen3_moe import Qwen3MoeSparseMoeBlock +if is_transformers_version(">=", "4.54"): + from transformers.masking_utils import create_causal_mask +if is_transformers_version(">=", "4.56"): + import transformers.masking_utils +if is_transformers_version(">=", "4.57"): + from transformers.models.qwen3_vl.modeling_qwen3_vl import Qwen3VLTextRotaryEmbedding +if is_transformers_version(">=", "5"): + from transformers.modeling_rope_utils import RotaryEmbeddingConfigMixin + +if TYPE_CHECKING: + from transformers.cache_utils import Cache + from transformers.modeling_utils import PreTrainedModel + +if is_transformers_version(">=", "4.54"): + from transformers.utils import TransformersKwargs +else: + TransformersKwargs = object + + +####### + +import dataclasses +import sys +from typing import TYPE_CHECKING + +import transformers +from torch.onnx import symbolic_helper + +from optimum.utils import is_diffusers_version, is_torch_version, is_transformers_version + + +if is_transformers_version(">=", "4.44") and is_transformers_version("<", "4.50"): + from optimum.exporters.openvino._traceable_cache import TraceableCache +if is_transformers_version(">=", "4.54"): + from optimum.exporters.openvino._traceable_decorator import traceable_check_model_inputs +if is_transformers_version(">=", "4.43") and is_transformers_version("<", "4.48"): + from transformers.models.clip.modeling_clip import CLIPAttention, CLIPSdpaAttention +if is_transformers_version(">=", "4.48"): + from transformers.cache_utils import DynamicCache, EncoderDecoderCache + from transformers.models.moonshine.modeling_moonshine import MoonshinePreTrainedModel +if is_transformers_version(">=", "4.53"): + from transformers.masking_utils import ( + ALL_MASK_ATTENTION_FUNCTIONS, + _ignore_causal_mask_sdpa, + and_masks, + causal_mask_function, + eager_mask, + padding_mask_function, + prepare_padding_mask, + sdpa_mask, + ) + from transformers.models.qwen3_moe.modeling_qwen3_moe import Qwen3MoeSparseMoeBlock +if is_transformers_version(">=", "4.53.1"): + from transformers.masking_utils import find_packed_sequence_indices +if is_transformers_version(">=", "4.55"): + from transformers.models.gpt_oss.modeling_gpt_oss import GptOssExperts +if is_transformers_version(">=", "4.56"): + from transformers.cache_utils import DynamicLayer + +if is_diffusers_version(">=", "0.35.0"): + import diffusers.models.transformers.transformer_flux + + +logger = logging.getLogger(__name__) + + +@symbolic_helper.parse_args("v", "v") +def __ior_(g, self: torch._C.Value, other: torch._C.Value) -> torch._C.Value: + return g.op("Or", self, other) + + +torch.onnx.register_custom_op_symbolic("aten::__ior__", __ior_, 14) + +if is_torch_version("<", "2.9"): + # this was fixed in torch in 2.9 https://github.com/pytorch/pytorch/pull/159973 + from torch.onnx import JitScalarType + from torch.onnx.symbolic_opset14 import _attention_scale, _causal_attention_mask + + @symbolic_helper.parse_args("v", "v", "v", "v", "f", "b", "v", "b") + def scaled_dot_product_attention( + g, + query: torch._C.Value, + key: torch._C.Value, + value: torch._C.Value, + attn_mask: torch._C.Value | None = None, + dropout_p: float = 0.0, + is_causal: bool = False, + scale: torch._C.Value | None = None, + enable_gqa: bool = False, + ): + assert (not is_causal) or ( + is_causal and symbolic_helper._is_none(attn_mask) + ), "is_causal and attn_mask cannot be set at the same time" + assert not enable_gqa, "conversion of scaled_dot_product_attention not implemented if enable_gqa is True" + + if symbolic_helper._is_none(scale): + scale = _attention_scale(g, query) + + if is_causal: + attn_mask = _causal_attention_mask(g, query, key) + + # Swap the last two axes of key + # NOTE: onnx-script has different logic here, because the attribute perms in + # transpose needs list of ints + key_shape_builtin = symbolic_helper._get_tensor_rank(key) + key_transposed_axes = list(range(key_shape_builtin)) + key_transposed_axes[-1], key_transposed_axes[-2] = (key_transposed_axes[-2], key_transposed_axes[-1]) + key_transposed = g.op("Transpose", key, perm_i=key_transposed_axes) + + # https://github.com/pytorch/pytorch/blob/12da0c70378b5be9135c6fda62a9863bce4a4818/aten/src/ATen/native/transformers/attention.cpp#L653 + # Scale q, k before matmul for stability see https://tinyurl.com/sudb9s96 for math + query_scaled = g.op("Mul", query, g.op("Sqrt", scale)) + key_transposed_scaled = g.op("Mul", key_transposed, g.op("Sqrt", scale)) + mul_qk = g.op("MatMul", query_scaled, key_transposed_scaled) + + if symbolic_helper._is_none(attn_mask): + mul_qk_add = mul_qk + attn_weight = g.op("Softmax", mul_qk_add, axis_i=-1) + elif JitScalarType.from_value(attn_mask) == JitScalarType.BOOL: + # Turn the Boolean mask to float: attn_mask.masked_fill(not attn_mask, -float('inf')) + const_zero = g.op("Constant", value_t=torch.tensor([0.0])) + const_neg_inf = g.op("Constant", value_t=torch.tensor([-float("inf")])) + attn_mask = g.op("Where", attn_mask, const_zero, const_neg_inf) + mul_qk_add = g.op("Add", mul_qk, attn_mask) + attn_weight = g.op("Softmax", mul_qk_add, axis_i=-1) + # when using scaled dot product attention with a boolean mask, we replace NaN values in attn_weight with 0.0 + attn_weight = g.op( + "Where", g.op("IsNaN", attn_weight), g.op("Constant", value_t=torch.tensor([0.0])), attn_weight + ) + elif JitScalarType.from_value(attn_mask) in ( + JitScalarType.FLOAT, + JitScalarType.HALF, + JitScalarType.BFLOAT16, + ): + mul_qk_add = g.op("Add", mul_qk, attn_mask) + attn_weight = g.op("Softmax", mul_qk_add, axis_i=-1) + else: + raise ValueError(f"Unsupported type for attn_mask: {JitScalarType.from_value(attn_mask)}") + + if dropout_p != 0: + attn_weight = g.op( + "Dropout", + attn_weight, + g.op("Constant", value_t=torch.tensor(dropout_p, dtype=torch.float)), + ) + + return g.op("MatMul", attn_weight, value) + + torch.onnx.register_custom_op_symbolic("aten::scaled_dot_product_attention", scaled_dot_product_attention, 14) + + +def patch_everywhere(attribute_name: str, patch: Any, module_name_prefix: str | None = None): + """Finds all occurrences of `attribute_name` in the loaded modules and patches them with `patch`. + + Args: + attribute_name (`str`): + The name of attribute to patch. + patch (`Any`): + The patch for the attribute. + module_name_prefix (`Optional[str]`, defaults to `None`): + If set, only module names starting with this prefix will be considered for patching. + """ + # sys.modules may be updated while being iterated over, hence the list copy. + for name in list(sys.modules): + module = sys.modules[name] + if module_name_prefix is not None and not name.startswith(module_name_prefix): + continue + if hasattr(module, attribute_name): + setattr(module, attribute_name, patch) + + +def override_arguments(args, kwargs, forward_signature, model_kwargs: dict[str, Any]): + """Override the args and kwargs with the argument values from model_kwargs, following the signature forward_signature corresponding to args and kwargs.""" + args = list(args) + + for argument in model_kwargs: + if argument in forward_signature.parameters: + argument_index = list(forward_signature.parameters.keys()).index(argument) + if argument in kwargs or len(args) <= argument_index: + kwargs[argument] = model_kwargs[argument] + else: + args[argument_index] = model_kwargs[argument] + else: + kwargs[argument] = model_kwargs[argument] + + return args, kwargs + + +def preprocess_encoder_outputs(encoder_outputs): + if is_transformers_version(">=", "4.54") and isinstance(encoder_outputs, (list, tuple)): + encoder_outputs = BaseModelOutput(*encoder_outputs) + + return encoder_outputs + + +def preprocess_past_key_values(past_key_values): + if ( + is_transformers_version(">=", "4.48") + and isinstance(past_key_values, (list, tuple)) + and isinstance(past_key_values[0], (list, tuple)) + ): + if len(past_key_values[0]) == 2: + if hasattr(DynamicCache, "from_legacy_cache"): + past_key_values = DynamicCache.from_legacy_cache(past_key_values) + else: + past_key_values = DynamicCache(past_key_values) + elif len(past_key_values[0]) == 4: + if hasattr(EncoderDecoderCache, "from_legacy_cache"): + past_key_values = EncoderDecoderCache.from_legacy_cache(past_key_values) + else: + past_key_values = EncoderDecoderCache( + DynamicCache([layer[:2] for layer in past_key_values]), + DynamicCache([layer[2:] for layer in past_key_values]), + ) + else: + raise ValueError( + f"past_key_values should have either 2 or 4 elements, but it has {len(past_key_values[0])} elements." + ) + + return past_key_values + + +def postprocess_past_key_values(past_key_values, output_names: list[str]): + if is_transformers_version(">=", "4.48") and isinstance(past_key_values, (EncoderDecoderCache, DynamicCache)): + if hasattr(past_key_values, "to_legacy_cache"): + past_key_values = past_key_values.to_legacy_cache() + elif isinstance(past_key_values, DynamicCache): + past_key_values = [(lay.keys, lay.values) for lay in past_key_values.layers] + elif isinstance(past_key_values, EncoderDecoderCache): + past_key_values = [ + (self_lay.keys, self_lay.values, cross_lay.keys, cross_lay.values) + for self_lay, cross_lay in zip( + past_key_values.self_attention_cache.layers, + past_key_values.cross_attention_cache.layers, + ) + ] + else: + raise NotImplementedError(f"Unable to serialize class {type(past_key_values)}.") + + if ( + isinstance(past_key_values, (list, tuple)) + and isinstance(past_key_values[0], (list, tuple)) + and not any("encoder.key" in output_name for output_name in output_names) + ): + past_key_values = tuple(pkv[:2] for pkv in past_key_values) + + return past_key_values + + +@dataclasses.dataclass +class PatchingSpec: + """Data class that holds patching specifications. + + Args: + o: Module / object where the op to patch is located + name: Name of the op to monkey patch + custom_op: Custom op that patches the original op + orig_op: Original op that is being patched + op_wrapper: Wrapper (optional) that wraps both the original and custom ops. + It is useful for ops that are class or static methods for instance. + """ + + o: Any + name: str + custom_op: Callable + orig_op: Callable | None = None + op_wrapper: Callable | None = None + + +# An ONNX-export-compatible version of `tensor.unfold`. Without this, we get: +# torch.onnx.errors.SymbolicValueError: Unsupported: ONNX export of operator Unfold, input size not accessible. +# See https://github.com/pytorch/pytorch/issues/81871 for more information +def onnx_compatible_unfold(input_tensor, dimension, size, step): + """Custom implementation of torch.unfold without using torch.unfold. + + Args: + input_tensor (torch.Tensor): The input tensor. + dimension (int): The dimension to unfold. + size (int): The size of each slice. + step (int): The step size between slices. + + Returns: + torch.Tensor: The unfolded tensor. + """ + # Check if dimension is within the valid range + if not (-input_tensor.dim() <= dimension < input_tensor.dim()): + raise ValueError( + f"Dimension out of range (expected to be in range of [{-input_tensor.dim()}, {input_tensor.dim() - 1}], but got {dimension})" + ) + + # Normalize negative dimension + dimension = dimension % input_tensor.dim() + + # Compute the shape of the unfolded output + input_size = input_tensor.size(dimension) + num_slices = (input_size - size) // step + 1 + + # Permute dimension to the end for easier indexing + input_tensor = input_tensor.transpose(dimension, -1) + + # Extract slices + slices = [] + for i in range(num_slices): + start = i * step + end = start + size + slices.append(input_tensor[..., start:end]) + + # Stack slices and permute dimensions back + result = torch.stack(slices, dim=-2).transpose(dimension, -2) + return result + + +# An ONNX-export-compatible version of `tensor.repeat_interleave`. +# Without this, we get the following error: https://github.com/pytorch/pytorch/issues/145100 +# NOTE: This implementation is only necessary for export with dynamo=False (dynamo=True works correctly). +# and can be removed once Optimum switches to dynamo-based exports +def onnx_compatible_repeat_interleave(input_tensor, repeats, dim=None, output_size=None): # noqa: D417 + """Custom implementation of torch.repeat_interleave without using torch.repeat_interleave. + + Args: + input_tensor (torch.Tensor): The input tensor. + repeats (int or torch.Tensor): The number of repetitions for each element. + dim (int, optional): The dimension along which to repeat. Defaults to None. + + Returns: + torch.Tensor: The repeated tensor. + """ + if isinstance(repeats, int) or (torch.is_tensor(repeats) and repeats.dim() == 0): + if dim is None: + return input_tensor.flatten().unsqueeze(1).expand(-1, repeats).flatten() + repeats = torch.full((input_tensor.shape[dim],), repeats, dtype=torch.long, device=input_tensor.device) + + if dim is None: + return onnx_compatible_repeat_interleave(input_tensor.flatten(), repeats, 0) + + if dim != 0: + input_tensor = input_tensor.transpose(0, dim) + + # Create expand mask + max_repeats = repeats.max() + expanded = input_tensor.unsqueeze(1).expand(-1, max_repeats, *input_tensor.shape[1:]) + mask = torch.arange(max_repeats, device=input_tensor.device) < repeats.unsqueeze(1) + result = expanded[mask] + + if dim != 0: + result = result.transpose(0, dim) + + return result + + +# Custom implementation of torch.linalg.matrix_norm not using torch.linalg.matrix_norm, torch.norm or torch.linalg.norm. +def onnx_compatible_linalg_norm(x, ord=2, dim=None, keepdim=False, *, dtype=None, out=None) -> torch.Tensor: + if ord != 2: + raise ValueError( + f"Only ord=2 is supported by onnx_compatible_linalg_norm, but got ord={ord}. " + "Please extend this function to support other norms." + ) + + if dim is None: + dim = (-2, -1) + + norm = torch.sqrt(torch.sum(torch.square(x), dim=dim, keepdim=keepdim)) + + if dtype is not None: + norm = norm.to(dtype) + if out is not None: + out.copy_(norm) + + return norm + + +def onnx_compatible_rms_norm(input, normalized_shape, weight=None, eps=None): + if eps is None: + eps = torch.finfo(input.dtype).eps + + axis = -len(normalized_shape) + mean_square = torch.mean(torch.square(input), dim=axis, keepdim=True) + rms = torch.sqrt(mean_square + eps) + output = input / rms + + if weight is not None: + output = output * weight + + return output + + +# A patched version of https://github.com/huggingface/transformers/blob/v4.53.2/src/transformers/masking_utils.py#L602 +# That returns a tensor of zeros with the same shape as position_ids indicating no packed sequence indices. +def find_packed_sequence_indices_patched(position_ids: torch.Tensor) -> torch.Tensor: + return torch.zeros_like(position_ids) + + +if is_transformers_version(">=", "4.53"): + _prepare_padding_mask_slice = "_slice" in inspect.signature(prepare_padding_mask).parameters +else: + _prepare_padding_mask_slice = False + + +# Custom vectorized implementation of sdpa_mask without using vmap +def _orig_sdpa_mask_without_vmap( + batch_size: int, + cache_position: torch.Tensor, + kv_length: int, + kv_offset: int = 0, + mask_function: Callable | None = None, + attention_mask: torch.Tensor | None = None, + local_size: int | None = None, + allow_is_causal_skip: bool = True, + **kwargs, +) -> torch.Tensor | None: + if mask_function is None: + mask_function = causal_mask_function + + q_length = cache_position.shape[0] + # Potentially pad the 2D mask, and slice it correctly + if _prepare_padding_mask_slice: + padding_mask = prepare_padding_mask(attention_mask, kv_length, kv_offset, _slice=False) + else: + padding_mask = prepare_padding_mask(attention_mask, kv_length, kv_offset) + + # Under specific conditions, we can avoid materializing the mask, instead relying on the `is_causal` argument + if allow_is_causal_skip and _ignore_causal_mask_sdpa(padding_mask, q_length, kv_length, kv_offset, local_size): + return None + + # Potentially add the padding 2D mask + if padding_mask is not None: + mask_function = and_masks(mask_function, padding_mask_function(padding_mask)) + + # Create broadcatable indices + device = cache_position.device + q_indices = cache_position[None, None, :, None] + head_indices = torch.arange(1, dtype=torch.long, device=device)[None, :, None, None] + batch_indices = torch.arange(batch_size, dtype=torch.long, device=device)[:, None, None, None] + kv_indices = torch.arange(kv_length, dtype=torch.long, device=device)[None, None, None, :] + kv_offset + # Apply mask function element-wise through broadcasting + causal_mask = mask_function(batch_indices, head_indices, q_indices, kv_indices) + # Expand the mask to match batch size and query length if they weren't used in the mask function + causal_mask = causal_mask.expand(batch_size, -1, q_length, kv_length) + + return causal_mask + + +# Compatibility wrapper for sdpa_mask_without_vmap from optimum. +# The installed optimum version expects (batch_size, cache_position: Tensor, kv_length, ...), +# but transformers >= 5.5 passes (batch_size, q_length: int, kv_length: int, q_offset: int, ...). +def sdpa_mask_without_vmap(batch_size, q_length=None, kv_length=None, q_offset=0, kv_offset=0, **kwargs): + import inspect + + sig = inspect.signature(_orig_sdpa_mask_without_vmap) + if is_transformers_version(">=", "5.5") and "cache_position" in sig.parameters and q_length is not None: + # Old optimum signature: (batch_size, cache_position, kv_length, kv_offset, ...) + cache_position = torch.arange(q_length, dtype=torch.long) + q_offset + kwargs.pop("q_offset", None) + kwargs.pop("allow_is_bidirectional_skip", None) + kwargs.pop("allow_torch_fix", None) + kwargs.pop("use_vmap", None) + kwargs.pop("device", None) + return _orig_sdpa_mask_without_vmap(batch_size, cache_position, kv_length, kv_offset=kv_offset, **kwargs) + else: + return _orig_sdpa_mask_without_vmap( + batch_size, q_length=q_length, kv_length=kv_length, q_offset=q_offset, kv_offset=kv_offset, **kwargs + ) + + +# Adapted from https://github.com/huggingface/transformers/blob/v4.53.0/src/transformers/masking_utils.py#L433 +def eager_mask_without_vmap(*args, **kwargs) -> torch.Tensor: + kwargs.pop("allow_is_causal_skip", None) + dtype = kwargs.get("dtype", torch.float32) + mask = sdpa_mask_without_vmap(*args, allow_is_causal_skip=False, **kwargs) + mask = torch.where(mask, torch.tensor(0.0, device=mask.device, dtype=dtype), torch.finfo(dtype).min) + return mask + + +original_triu = torch.triu +original_tril = torch.tril + + +# Custom implementation of torch.tril that doesn't fail on int32 tensors. +def onnx_compatible_tril(input_tensor: torch.Tensor, *args, **kwargs) -> torch.Tensor: + if input_tensor.dtype == torch.int32: + return original_tril(input_tensor.to(torch.int64), *args, **kwargs).to(torch.int32) + else: + return original_tril(input_tensor, *args, **kwargs) + + +# Custom implementation of torch.triu that doesn't fail on int32 tensors. +def onnx_compatible_triu(input_tensor: torch.Tensor, *args, **kwargs) -> torch.Tensor: + if input_tensor.dtype == torch.int32: + return original_triu(input_tensor.to(torch.int64), *args, **kwargs).to(torch.int32) + else: + return original_triu(input_tensor, *args, **kwargs) + + +original_scaled_dot_product_attention = torch.nn.functional.scaled_dot_product_attention + + +# A patched `torch.nn.functional.scaled_dot_product_attention` that doesn't fail during tracing +# from passing `is_causal` as a tensor (which is usually obtained with tensor shapes comparisons). +def traceable_scaled_dot_product_attention( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_mask: torch.Tensor | None = None, + dropout_p: float = 0.0, + is_causal: bool = False, + **kwargs, +) -> torch.Tensor: + if isinstance(is_causal, torch.Tensor): + is_causal = is_causal.item() + + if "enable_gqa" in kwargs: + kwargs.pop("enable_gqa") + + attn_weights = original_scaled_dot_product_attention( + query=query, key=key, value=value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, **kwargs + ) + + return attn_weights + + +# No-op bfloat16 casting to avoid issues with legacy ONNX export which cast to complex128 +def noop_bfloat16_casting(self): + return self + + +original_movedim = torch.Tensor.movedim + + +def onnx_compatible_movedim(self: torch.Tensor, dim1, dim2) -> torch.Tensor: + dim = self.dim() + if dim1 < 0: + dim1 += dim + if dim2 < 0: + dim2 += dim + return original_movedim(self, dim1, dim2) + + +def patched_dynamic_layer_update( + self, key_states: torch.Tensor, value_states: torch.Tensor, cache_kwargs: dict[str, Any] | None = None +) -> tuple[torch.Tensor, torch.Tensor]: + if self.keys is None: + self.keys = key_states + self.values = value_states + self.device = key_states.device + self.dtype = key_states.dtype + self.is_initialized = True + else: + self.keys = torch.cat([self.keys, key_states], dim=-2) + self.values = torch.cat([self.values, value_states], dim=-2) + return self.keys, self.values + + +UNSUPPORTED_OPS_PATCHING_SPEC = [ + PatchingSpec(torch, "tril", onnx_compatible_tril, torch.tril), + PatchingSpec(torch, "triu", onnx_compatible_triu, torch.triu), + PatchingSpec(torch, "rms_norm", onnx_compatible_rms_norm, torch.rms_norm), + PatchingSpec(torch.Tensor, "unfold", onnx_compatible_unfold, torch.Tensor.unfold), + PatchingSpec(torch.linalg, "norm", onnx_compatible_linalg_norm, torch.linalg.norm), + PatchingSpec(torch.Tensor, "bfloat16", noop_bfloat16_casting, torch.Tensor.bfloat16), + PatchingSpec(torch.Tensor, "movedim", onnx_compatible_movedim, torch.Tensor.movedim), + PatchingSpec(torch.Tensor, "repeat_interleave", onnx_compatible_repeat_interleave, torch.Tensor.repeat_interleave), + # TracerWarning: Using len to get tensor shape might cause the trace to be incorrect. Recommended usage would be tensor.shape[0]. Passing a tensor of different shape might lead to errors or silently give incorrect results. + PatchingSpec(torch.Tensor, "__len__", lambda x: x.shape[0], torch.Tensor.__len__), + PatchingSpec( + torch.nn.functional, + "scaled_dot_product_attention", + traceable_scaled_dot_product_attention, + torch.nn.functional.scaled_dot_product_attention, + ), +] + + +class ModelPatcher: + def __init__( + self, + config: "OnnxConfig", + model: PreTrainedModel, + model_kwargs: dict[str, Any] | None = None, + ): + self._model = model + + patching_specs = config.PATCHING_SPECS or [] + patching_specs.extend(UNSUPPORTED_OPS_PATCHING_SPEC) + + self._patching_specs = [] + for spec in patching_specs: + final_spec = spec + if spec.orig_op is None: + final_spec = dataclasses.replace(spec, orig_op=getattr(spec.o, spec.name)) + self._patching_specs.append(final_spec) + + self.orig_forward_name = "forward" if hasattr(self._model, "forward") else "call" + self.orig_forward = getattr(self._model, self.orig_forward_name) + + if is_transformers_version(">=", "4.54") and hasattr(self.orig_forward, "__wrapped__"): + # the original check_model_inputs has some failing cases that we fix in traceable_check_model_inputs + # we fix those issues in a PR in transformers https://github.com/huggingface/transformers/pull/40811 + # issues are: support for positional args (use_cache for instance) and fix for _CAN_RECORD_REGISTRY + # explicitly mapping to None for some models + self.orig_forward = types.MethodType( + traceable_check_model_inputs(self.orig_forward.__wrapped__), self._model + ) + + self.real_config = config + use_cache = getattr(self.real_config, "use_past", False) + self.model_kwargs = model_kwargs if model_kwargs is not None else {} + + for module in self._model.modules(): + if hasattr(module, "config") and hasattr(module.config, "use_cache"): + module.config.use_cache = use_cache + + @functools.wraps(self.orig_forward) + def patched_forward(*args, **kwargs): + signature = inspect.signature(self.orig_forward) + args, kwargs = override_arguments(args, kwargs, signature, model_kwargs=self.model_kwargs) + + # Transformers doesn't always respect the config.use_cache attribute + # there are cases where setting use_cache to true in every config and + # subconfig of a model still doesn't enable past_key_values in the outputs (gemma3) + # Explicitly setting the use_cache argument of the forward method seems to be the most reliable way + if "use_cache" in signature.parameters: + use_cache_index = list(signature.parameters.keys()).index("use_cache") + if use_cache_index < len(args): + args[use_cache_index] = use_cache + elif "use_cache" in kwargs: + kwargs["use_cache"] = use_cache + + if "past_key_values" in signature.parameters: + # Most models require past_key_values to be a cache instance instead of a tuple now + pkv_index = list(signature.parameters.keys()).index("past_key_values") + if pkv_index < len(args) and args[pkv_index] is not None: + args[pkv_index] = preprocess_past_key_values(args[pkv_index]) + elif kwargs.get("past_key_values") is not None: + kwargs["past_key_values"] = preprocess_past_key_values(kwargs["past_key_values"]) + + if "encoder_outputs" in signature.parameters: + # Some encoder-decoder models started to not accept encoder_outputs as tuple (e.g. moonshine) + encoder_outputs_index = list(signature.parameters.keys()).index("encoder_outputs") + if encoder_outputs_index < len(args) and args[encoder_outputs_index] is not None: + args[encoder_outputs_index] = preprocess_encoder_outputs(args[encoder_outputs_index]) + elif kwargs.get("encoder_outputs") is not None: + kwargs["encoder_outputs"] = preprocess_encoder_outputs(kwargs["encoder_outputs"]) + + outputs = self.orig_forward(*args, **kwargs) + + # This code block handles different cases of the filtered_outputs input to align it with the expected + # format of outputs. It is common for the output type of a model to vary, such as tensor, list, + # tuple, etc. For Transformers models, the output is encapsulated in a ModelOutput object that + # contains the output names of the model. In the case of Timm classification models, the output + # is of type tensor. By default, it is assumed that the output names mentioned in the ONNX config + # match the outputs in order. + filtered_outputs = {} + output_names = list(config.outputs.keys()) + if isinstance(outputs, dict): + for name, value in outputs.items(): + onnx_output_name = config.torch_to_onnx_output_map.get(name, name) + if ( + onnx_output_name in output_names + or (use_cache and name.startswith("past_key_values")) + or any(key.startswith(onnx_output_name) for key in output_names) + ): + filtered_outputs[name] = value + elif isinstance(outputs, (list, tuple)): + filtered_outputs = dict(zip(output_names, outputs)) + else: + if len(output_names) > 1: + num_outputs = len(output_names) + output_names_str = ", ".join(output_names) + raise ValueError( + f"{config.__class__.__name__} expects the model to return {num_outputs} outputs: {output_names_str}, " + f"but the it returned a single output of type {type(outputs)}. Please make sure either that the model " + "returns all the expected outputs, or that the ONNX config is correctly defined with the expected outputs." + ) + output_name = output_names[0] + filtered_outputs[output_name] = outputs + + if filtered_outputs.get("past_key_values") is not None: + filtered_outputs["past_key_values"] = postprocess_past_key_values( + filtered_outputs["past_key_values"], output_names=output_names + ) + + return filtered_outputs + + self.patched_forward = patched_forward + + def patch_ops(self): + for spec in self._patching_specs: + custom_op = spec.custom_op if spec.op_wrapper is None else spec.op_wrapper(spec.custom_op) + setattr(spec.o, spec.name, custom_op) + + def restore_ops(self): + for spec in self._patching_specs: + orig_op = spec.orig_op if spec.op_wrapper is None else spec.op_wrapper(spec.orig_op) + setattr(spec.o, spec.name, orig_op) + + def __enter__(self): + self.patch_ops() + setattr(self._model, self.orig_forward_name, self.patched_forward) + + # This is a workaround for the Cache class in transformers, we replace it + # with traceable cache is because the original one used in transformers + # inherited from nn.Module (for a couple versions), which can't be traced as input. + if is_transformers_version(">=", "4.44") and is_transformers_version("<", "4.50"): + self.original_cache_class = transformers.cache_utils.Cache + transformers.cache_utils.Cache = TraceableCache + + # This is a workaround for mask generation in transformers >= 4.53. + # The masking process uses vmap which is not traceable by TorchScript. + if is_transformers_version(">=", "4.53"): + ALL_MASK_ATTENTION_FUNCTIONS.register("sdpa", sdpa_mask_without_vmap) + ALL_MASK_ATTENTION_FUNCTIONS.register("eager", eager_mask_without_vmap) + + # This is a workaround for the find_packed_sequence_indices function in transformers which + # should only return a tensor of zeros with the same shape as position_ids indicating no packed sequence indices. + # The function uses torch.diff which is not traceable by TorchScript. + if is_transformers_version(">=", "4.53.1"): + self.original_find_packed_sequence_indices = find_packed_sequence_indices + transformers.masking_utils.find_packed_sequence_indices = find_packed_sequence_indices_patched + + # Starting from transformers 4.56.0, DynamicCache uses DynamicLayer which has an update method + # that uses torch.cat to concatenate an empty tensor with the key/value states during the first call. + # This causes issues during TorchScript tracing. + if is_transformers_version(">=", "4.56"): + self.original_dynamic_layer_update = DynamicLayer.update + DynamicLayer.update = patched_dynamic_layer_update + + def __exit__(self, exc_type, exc_value, traceback): + self.restore_ops() + setattr(self._model, self.orig_forward_name, self.orig_forward) + + if is_transformers_version(">=", "4.44") and is_transformers_version("<", "4.50"): + transformers.cache_utils.Cache = self.original_cache_class + + if is_transformers_version(">=", "4.53"): + ALL_MASK_ATTENTION_FUNCTIONS.register("sdpa", sdpa_mask) + ALL_MASK_ATTENTION_FUNCTIONS.register("eager", eager_mask) + + if is_transformers_version(">=", "4.53.1"): + transformers.masking_utils.find_packed_sequence_indices = self.original_find_packed_sequence_indices + + if is_transformers_version(">=", "4.56"): + DynamicLayer.update = self.original_dynamic_layer_update + + def __call__(self, *args, **kwargs): + if getattr(self._model, self.orig_forward_name) is self.orig_forward: + logger.warning("Running the non-patched model") + return self._model(*args, **kwargs) + + +class BigBirdPegasusModelPatcher(ModelPatcher): + def __enter__(self): + super().__enter__() + + if self.real_config._behavior == "encoder" and self._model.config.attention_type == "block_sparse": + logger.warning( + "BigBirdPegasus model is using block sparse attention, which is not supported in ONNX export. " + "The model will be exported with original full attention." + ) + self._model.set_attention_type("original_full") + + def __exit__(self, exc_type, exc_value, traceback): + super().__exit__(exc_type, exc_value, traceback) + + if self.real_config._behavior == "encoder" and self._model.config.attention_type == "block_sparse": + self._model.set_attention_type("block_sparse") + + +class MgpstrModelPatcher(ModelPatcher): + def __init__( + self, + config: "OnnxConfig", + model: PreTrainedModel, + model_kwargs: dict[str, Any] | None = None, + ): + super().__init__(config, model, model_kwargs) + + @functools.wraps(self.orig_forward) + def patched_forward(*args, **kwargs): + signature = inspect.signature(self.orig_forward) + args, kwargs = override_arguments(args, kwargs, signature, model_kwargs=self.model_kwargs) + + # logits is a tuple, so we unpack it and return them as separate outputs + char_logits, bpe_logits, wp_logits = self.orig_forward(*args, **kwargs).logits + + return { + "char_logits": char_logits, + "bpe_logits": bpe_logits, + "wp_logits": wp_logits, + } + + self.patched_forward = patched_forward + + +class SAMModelPatcher(ModelPatcher): + def __init__( + self, + config: "OnnxConfig", + model: PreTrainedModel, + model_kwargs: dict[str, Any] | None = None, + ): + super().__init__(config, model, model_kwargs) + + def patched_forward( + pixel_values=None, + input_points=None, + input_labels=None, + image_embeddings=None, + image_positional_embeddings=None, + return_dict=True, + **kwargs, + ): + if config.variant == "monolith": + return self.orig_forward( + pixel_values=pixel_values, + input_points=input_points, + input_labels=input_labels, + image_embeddings=image_embeddings, + return_dict=return_dict, + **kwargs, + ) + elif config.variant == "split": + # return_dict = get_argument(args, kwargs, signature, "return_dict") + if config.vision_encoder: + # pixel_values = get_argument(args, kwargs, signature, "pixel_values") + image_positional_embeddings = model.get_image_wide_positional_embeddings() + + # repeat with batch size + batch_size = pixel_values.shape[0] + image_positional_embeddings = image_positional_embeddings.repeat(batch_size, 1, 1, 1) + + vision_outputs = model.vision_encoder( + pixel_values, + output_attentions=False, + output_hidden_states=False, + return_dict=return_dict, + ) + image_embeddings = vision_outputs[0] + + if not return_dict: + return (image_embeddings, image_positional_embeddings) + else: + return { + "image_embeddings": image_embeddings, + "image_positional_embeddings": image_positional_embeddings, + } + else: + if input_points is None: + raise ValueError("input_points is required to export the prompt encoder / mask decoder.") + + sparse_embeddings, dense_embeddings = model.prompt_encoder( + input_points=input_points, + input_labels=input_labels, + input_boxes=None, # Not supported in the ONNX export + input_masks=None, # Not supported in the ONNX export + ) + outputs = model.mask_decoder( + image_embeddings=image_embeddings, + image_positional_embeddings=image_positional_embeddings, + sparse_prompt_embeddings=sparse_embeddings, + dense_prompt_embeddings=dense_embeddings, + multimask_output=True, # Not supported in the ONNX export + attention_similarity=None, # Not supported in the ONNX export + target_embedding=None, # Not supported in the ONNX export + ) + low_res_masks, iou_predictions = outputs[:2] + + if not return_dict: + return (iou_predictions, low_res_masks) + else: + return {"iou_scores": iou_predictions, "pred_masks": low_res_masks} + + self.patched_forward = patched_forward + + +def patched_speecht5_prenet_forward( + self, + input_values: torch.Tensor, + speaker_embeddings: torch.Tensor | None = None, +): + # Dropout is always applied, even when evaluating. See ยง2.2 in https://arxiv.org/abs/1712.05884. + + inputs_embeds = input_values + for layer in self.layers: + inputs_embeds = torch.nn.functional.relu(layer(inputs_embeds)) + + # NOTE: we patch the prenet to avoid using torch.nn.functional.dropout, that is exported as a `Dropout` node in the ONNX + # that is ignored during inference by some runtimes as ONNX Runtime. + # Reference: https://github.com/microsoft/onnxruntime/issues/9333 & https://github.com/microsoft/onnxruntime/issues/5549 + mask = torch.rand(inputs_embeds.shape, device=inputs_embeds.device) > self.config.speech_decoder_prenet_dropout + inputs_embeds = inputs_embeds * mask / (1 - self.config.speech_decoder_prenet_dropout) + + # inputs_embeds = nn.functional.dropout( + # inputs_embeds, self.config.speech_decoder_prenet_dropout, training=True + # ) + + inputs_embeds = self.final_layer(inputs_embeds) + inputs_embeds = self.encode_positions(inputs_embeds) + + if speaker_embeddings is not None: + speaker_embeddings = torch.nn.functional.normalize(speaker_embeddings) + speaker_embeddings = speaker_embeddings.unsqueeze(1) + speaker_embeddings = speaker_embeddings.expand(-1, inputs_embeds.size(1), -1) + inputs_embeds = torch.cat([inputs_embeds, speaker_embeddings], dim=-1) + inputs_embeds = torch.nn.functional.relu(self.speaker_embeds_layer(inputs_embeds)) + + return inputs_embeds + + +class SpeechT5ModelPatcher(ModelPatcher): + def __enter__(self): + super().__enter__() + + self.original_speecht5_prenet_forward = self._model.speecht5.decoder.prenet.forward + self._model.speecht5.decoder.prenet.forward = types.MethodType( + patched_speecht5_prenet_forward, self._model.speecht5.decoder.prenet + ) + + def __exit__(self, exc_type, exc_value, traceback): + super().__exit__(exc_type, exc_value, traceback) + + self._model.speecht5.decoder.prenet.forward = types.MethodType( + self.original_speecht5_prenet_forward, self._model.speecht5.decoder.prenet + ) + + def __init__( + self, + config: "OnnxConfig", + model: PreTrainedModel, + model_kwargs: dict[str, Any], + ): + super().__init__(config, model, model_kwargs) + model.vocoder = model_kwargs["vocoder_model"].eval() + + def patched_forward( + input_ids=None, + speaker_embeddings=None, + encoder_outputs=None, + past_key_values=None, + output_sequence=None, + spectrogram=None, + encoder_attention_mask=None, + ): + if past_key_values is not None: + past_key_values = preprocess_past_key_values(past_key_values) + + if self.real_config._behavior == "encoder": + encoder_attention_mask = torch.ones_like(input_ids) + encoder_out = model.speecht5.encoder(input_values=input_ids, attention_mask=encoder_attention_mask) + # downsample encoder attention mask + if isinstance(model.speecht5.encoder, SpeechT5EncoderWithSpeechPrenet): + encoder_attention_mask = model.speecht5.encoder.prenet._get_feature_vector_attention_mask( + encoder_out[0].shape[1], encoder_attention_mask + ) + outputs = { + "encoder_outputs": encoder_out.last_hidden_state, + "encoder_attention_mask": encoder_attention_mask, + } + elif self.real_config._behavior == "decoder": + use_cache = self.real_config.use_past + encoder_hidden_states = encoder_outputs[0] + decoder_hidden_states = model.speecht5.decoder.prenet(output_sequence, speaker_embeddings) + # Run the decoder layers on the last element of the prenet output. + decoder_out = model.speecht5.decoder.wrapped_decoder( + hidden_states=decoder_hidden_states[:, -1:], + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + past_key_values=past_key_values, + output_attentions=False, + use_cache=use_cache, + ) + last_decoder_output = decoder_out.last_hidden_state[0, -1] + past_key_values = decoder_out.past_key_values + # Predict the new mel spectrum for this step in the sequence. + spectrum = model.speech_decoder_postnet.feat_out(last_decoder_output) + spectrum = spectrum.view(model.config.reduction_factor, model.config.num_mel_bins) + # NOTE: extending the spectrogram should is to be handled outside of the ONNX. + # spectrogram.append(spectrum) + # Extend the output sequence with the new mel spectrum. + output_sequence = torch.cat( + (output_sequence, spectrum[-1].view(1, 1, model.config.num_mel_bins)), dim=1 + ) + # Predict the probability that this is the stop token. + prob = torch.sigmoid(model.speech_decoder_postnet.prob_out(last_decoder_output)) + outputs = { + "output_sequence_out": output_sequence, + "spectrum": spectrum, + "prob": prob, + "past_key_values": past_key_values, + } + elif self.real_config.is_postnet_and_vocoder: + # NOTE: the following concatenation is expected to be handled outside of the ONNX: + # spectrogram = torch.cat(spectrogram, dim=0).unsqueeze(0) + spectrogram = spectrogram.unsqueeze(0) + spectrogram = model.speech_decoder_postnet.postnet(spectrogram) + spectrogram = spectrogram.squeeze(0) + waveform = model.vocoder(spectrogram) + outputs = {"waveform": waveform} + else: + raise ValueError("Should not happen") + + if outputs.get("past_key_values") is not None: + outputs["past_key_values"] = postprocess_past_key_values( + outputs["past_key_values"], output_names=list(config.outputs.keys()) + ) + + return outputs + + self.patched_forward = patched_forward + + +class SentenceTransformersTransformerPatcher(ModelPatcher): + def __init__( + self, + config: "OnnxConfig", + model: PreTrainedModel, + model_kwargs: dict[str, Any], + ): + super().__init__(config, model, model_kwargs) + + def patched_forward(input_ids, attention_mask): + result = self.orig_forward({"input_ids": input_ids, "attention_mask": attention_mask}) + + if "input_ids" in result: + del result["input_ids"] + if "attention_mask" in result: + del result["attention_mask"] + if "all_layer_embeddings" in result: + del result["all_layer_embeddings"] + + return result + + self.patched_forward = patched_forward -if is_transformers_version(">=", "4.53"): - from transformers.masking_utils import ALL_MASK_ATTENTION_FUNCTIONS, eager_mask, sdpa_mask - from transformers.models.qwen3_moe.modeling_qwen3_moe import Qwen3MoeSparseMoeBlock -if is_transformers_version(">=", "4.54"): - from transformers.masking_utils import create_causal_mask -if is_transformers_version(">=", "4.56"): - import transformers.masking_utils -if is_transformers_version(">=", "4.57"): - from transformers.models.qwen3_vl.modeling_qwen3_vl import Qwen3VLTextRotaryEmbedding -if is_transformers_version(">=", "5"): - from transformers.modeling_rope_utils import RotaryEmbeddingConfigMixin +class SentenceTransformersCLIPPatcher(ModelPatcher): + def __init__( + self, + config: "OnnxConfig", + model: PreTrainedModel, + model_kwargs: dict[str, Any], + ): + super().__init__(config, model, model_kwargs) -if TYPE_CHECKING: - from transformers.cache_utils import Cache - from transformers.modeling_utils import PreTrainedModel + def patched_forward(input_ids, attention_mask, pixel_values): + vision_outputs = model[0].model.vision_model(pixel_values=pixel_values) + image_embeds = model[0].model.visual_projection(vision_outputs[1]) - from optimum.exporters.onnx.config import OnnxConfig + text_outputs = model[0].model.text_model(input_ids=input_ids, attention_mask=attention_mask) + text_embeds = model[0].model.text_projection(text_outputs[1]) -if is_transformers_version(">=", "4.54"): - from transformers.utils import TransformersKwargs -else: - TransformersKwargs = object + if len(model) > 1: + image_embeds = model[1:](image_embeds) + text_embeds = model[1:](text_embeds) + return {"text_embeds": text_embeds, "image_embeds": image_embeds} -logger = logging.getLogger(__name__) + self.patched_forward = patched_forward + + +# Triu with possible dynamic `diagonal` argument. Not possible with torch.triu unfortunately. +def triu_onnx(x, diagonal=0): + l, w = x.shape + arrange_rows = torch.arange(l, device=x.device) + + arrange_cols = torch.arange(w, device=x.device) + mask = arrange_cols.expand(l, w) + + arrange_rows = arrange_rows[:, None] + diagonal + mask = mask >= arrange_rows + return x.masked_fill(mask == 0, 0) + + +def patched_build_delay_pattern_mask(self, input_ids: torch.Tensor, pad_token_id: int, max_length: int | None = None): + # (bsz * num_codebooks, seq_len) -> (bsz, num_codebooks, seq_len) + input_ids = input_ids.reshape(-1, self.num_codebooks, input_ids.shape[-1]) + bsz, num_codebooks, seq_len = input_ids.shape + + max_length = max_length if max_length is not None else self.generation_config.max_length + input_ids_shifted = torch.ones((bsz, num_codebooks, max_length), dtype=torch.long, device=input_ids.device) * -1 + + channel_codebooks = num_codebooks // 2 if self.config.audio_channels == 2 else num_codebooks + # we only apply the mask if we have a large enough seq len - otherwise we return as is + if max_length < 2 * channel_codebooks - 1: + raise NotImplementedError("Not supported in ONNX export. Please open an issue in Optimum repository.") + + # fill the shifted ids with the prompt entries, offset by the codebook idx + for codebook in range(channel_codebooks): + if self.config.audio_channels == 1: + # mono channel - loop over the codebooks one-by-one + input_ids_shifted[:, codebook, codebook : seq_len + codebook] = input_ids[:, codebook] + else: + # left/right channels are interleaved in the generated codebooks, so handle one then the other + input_ids_shifted[:, 2 * codebook, codebook : seq_len + codebook] = input_ids[:, 2 * codebook] + input_ids_shifted[:, 2 * codebook + 1, codebook : seq_len + codebook] = input_ids[:, 2 * codebook + 1] + + # construct a pattern mask that indicates the positions of padding tokens for each codebook + # first fill the upper triangular part (the EOS padding) + # NOTE: We could use torch.bool here, but PyTorch the complains with `The exported ONNX model failed ONNX shape inference.` + # Using int8 leads to `Could not find an implementation for Where` + delay_pattern = triu_onnx( + torch.ones((channel_codebooks, max_length), dtype=torch.int32), diagonal=max_length - channel_codebooks + 1 + ) + + # NOTE: We could use torch.bool here, but PyTorch the complains with `The exported ONNX model failed ONNX shape inference.` + # Using int32 leads to `Could not find an implementation for Trilu`, hence int64 here + + # then fill the lower triangular part (the BOS padding) + delay_pattern = delay_pattern + torch.tril(torch.ones((channel_codebooks, max_length), dtype=torch.int64)) + delay_pattern = delay_pattern.to(torch.bool) + + if self.config.audio_channels == 2: + # for left/right channel we need to duplicate every row of the pattern mask in an interleaved fashion + delay_pattern = delay_pattern.repeat_interleave(2, dim=0) + + mask = ~delay_pattern.to(input_ids.device) + input_ids = mask * input_ids_shifted + ~mask * pad_token_id + + # find the first position to start generating - this is the first place we have the -1 token + # and will always be in the first codebook (since it has no codebook offset) + first_codebook_ids = input_ids[:, 0, :] + start_ids = (first_codebook_ids == -1).nonzero()[:, 1] + + # TODO: Is this OK? + first_start_id = start_ids.min() + + # (bsz * num_codebooks, seq_len) -> (bsz, num_codebooks, seq_len) + pattern_mask = input_ids.reshape(bsz * num_codebooks, -1) + input_ids_edited = input_ids[..., :first_start_id].reshape(bsz * num_codebooks, -1) + return {"input_ids_edited": input_ids_edited, "delay_pattern_mask": pattern_mask} + + +class MusicgenModelPatcher(ModelPatcher): + def __enter__(self): + self.patch_ops() + if self.real_config.model_part == "build_delay_pattern_mask": + # For build_delay_pattern_mask, we need to override the signature too. + self._model.forward = types.MethodType(patched_build_delay_pattern_mask, self._model) + else: + setattr(self._model, self.orig_forward_name, self.patched_forward) + + def __exit__(self, exc_type, exc_value, traceback): + self.restore_ops() + if self.real_config.model_part == "build_delay_pattern_mask": + self._model.forward = self.original_decoder_forward + else: + setattr(self._model, self.orig_forward_name, self.orig_forward) + + def __init__( + self, + config: "OnnxConfig", + model: PreTrainedModel, + model_kwargs: dict[str, Any] | None = None, + ): + super().__init__(config, model, model_kwargs) + + if config.model_part == "build_delay_pattern_mask": + self.original_decoder_forward = self.orig_forward + elif config.model_part == "encodec_decode": + # EncodecModel.forward -> EncodecModel.decode + @functools.wraps(self.orig_forward) + def patched_forward( + input_values: torch.Tensor | None = None, + padding_mask: torch.Tensor | None = None, + audio_codes: torch.Tensor | None = None, + bandwidth: float | None = None, + audio_scales: torch.Tensor | None = None, + return_dict: bool | None = None, + ): + chunk_length = self.real_config._config.audio_encoder.chunk_length + if chunk_length is None: + if audio_scales is not None: + audio_scales = audio_scales[0] + + if len(audio_codes) != 1: + raise ValueError(f"Expected one frame, got {len(audio_codes)}") + audio_values = self._model._decode_frame(audio_codes[0], audio_scales) + else: + raise ValueError("Not supported, a meaningful error should have been raised ahead.") + decoded_frames = [] + + for frame, scale in zip(audio_codes, audio_scales): + frames = self._model._decode_frame(frame, scale) + decoded_frames.append(frames) + + audio_values = self._model._linear_overlap_add(decoded_frames, self.config.chunk_stride or 1) + + # truncate based on padding mask + if padding_mask is not None and padding_mask.shape[-1] < audio_values.shape[-1]: + audio_values = audio_values[..., : padding_mask.shape[-1]] + + return {"audio_values": audio_values} + + self.patched_forward = patched_forward + + +class MetaCLIP2Patcher(ModelPatcher): + def __init__( + self, + config: "OnnxConfig", + model: PreTrainedModel, + model_kwargs: dict[str, Any] | None = None, + ): + super().__init__(config, model, model_kwargs) + + def patched_forward(input_ids=None, pixel_values=None, attention_mask=None): + if config.variant == "monolith": + return self.orig_forward(input_ids=input_ids, pixel_values=pixel_values, attention_mask=attention_mask) + + if config.variant == "split": + if config.vision_model: + image_embeds = model.get_image_features(pixel_values) + return {"image_embeds": image_embeds} + + text_embeds = model.get_text_features(input_ids, attention_mask) + return { + "text_embeds": text_embeds, + } + + self.patched_forward = patched_forward + + +class CLIPModelPatcher(ModelPatcher): + def __enter__(self): + super().__enter__() + if is_transformers_version(">=", "4.43") and is_transformers_version("<", "4.48"): + self.original_sdpa_forward = CLIPSdpaAttention.forward + CLIPSdpaAttention.forward = CLIPAttention.forward + + def __exit__(self, exc_type, exc_value, traceback): + super().__exit__(exc_type, exc_value, traceback) + if is_transformers_version(">=", "4.43") and is_transformers_version("<", "4.48"): + CLIPSdpaAttention.forward = self.original_sdpa_forward + + +class VitPoseModelPatcher(ModelPatcher): + def __init__( + self, + config: "OnnxConfig", + model: PreTrainedModel, + model_kwargs: dict[str, Any] | None = None, + ): + # Set dataset_index (defaulting to COCO=0), otherwise we will get an error like: + # ValueError: dataset_index must be provided when using multiple experts (num_experts=6). Please provide dataset_index to the forward pass. + if model.config.backbone_config.num_experts > 1: + model_kwargs["dataset_index"] = torch.tensor(0, device=model.device) + + super().__init__(config, model, model_kwargs) + + +# https://github.com/huggingface/transformers/blob/v4.53.0/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py#L228 +def qwen3_moe_forward_patched(self, hidden_states: torch.Tensor) -> torch.Tensor: + batch_size, sequence_length, hidden_dim = hidden_states.shape + hidden_states = hidden_states.view(-1, hidden_dim) + # router_logits: (batch * sequence_length, n_experts) + router_logits = self.gate(hidden_states) + + routing_weights = torch.nn.functional.softmax(router_logits, dim=1, dtype=torch.float) + routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1) + if self.norm_topk_prob: # only diff with mixtral sparse moe block! + routing_weights /= routing_weights.sum(dim=-1, keepdim=True) + # we cast back to the input dtype + routing_weights = routing_weights.to(hidden_states.dtype) + + final_hidden_states = torch.zeros( + (batch_size * sequence_length, hidden_dim), dtype=hidden_states.dtype, device=hidden_states.device + ) + + # One hot encode the selected experts to create an expert mask + # this will be used to easily index which expert is going to be sollicitated + expert_mask = torch.nn.functional.one_hot(selected_experts, num_classes=self.num_experts).permute(2, 1, 0) + + # TODO: we loop over all possible experts instead of hit ones to avoid issues in graph execution. + # expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero() + # Loop over all available experts in the model and perform the computation on each expert + for expert_idx in range(self.num_experts): + expert_layer = self.experts[expert_idx] + idx, top_x = torch.where(expert_mask[expert_idx].squeeze(0)) + + # Index the correct hidden states and compute the expert hidden state for + # the current expert. We need to make sure to multiply the output hidden + # states by `routing_weights` on the corresponding tokens (top-1 and top-2) + current_state = hidden_states[None, top_x].reshape(-1, hidden_dim) + current_hidden_states = expert_layer(current_state) * routing_weights[top_x, idx, None] + + # However `index_add_` only support torch tensors for indexing so we'll use + # the `top_x` tensor here. + final_hidden_states.index_add_(0, top_x, current_hidden_states.to(hidden_states.dtype)) + final_hidden_states = final_hidden_states.reshape(batch_size, sequence_length, hidden_dim) + return final_hidden_states, router_logits + + +class Qwen3MoeModelPatcher(ModelPatcher): + def __enter__(self): + super().__enter__() + + # This is a workaround for the Qwen3 Moe Sparse block that is not compatible with ONNX export. + # The forward method of the Moe Sparse block is patched to avoid looping only on the experts that are selected + # by the router, which fails during execution in ONNX Runtime. + # TODO: investigate more on this issue. + if is_transformers_version(">=", "4.53"): + self.original_moe_forward = Qwen3MoeSparseMoeBlock.forward + Qwen3MoeSparseMoeBlock.forward = qwen3_moe_forward_patched + + def __exit__(self, exc_type, exc_value, traceback): + super().__exit__(exc_type, exc_value, traceback) + + if is_transformers_version(">=", "4.53"): + Qwen3MoeSparseMoeBlock.forward = self.original_moe_forward + + +# This is a traceable version of the original function, +# the original results in a constant integer due to the use of int(expr) +def _get_feat_extract_output_lengths_patched(self, input_lengths: torch.LongTensor): + output_conv1_length = (input_lengths - 127) // 64 + 1 + output_conv2_length = (output_conv1_length - 7) // 3 + 1 + output_conv3_length = (output_conv2_length - 3) // 2 + 1 + return output_conv3_length + + +class MoonshineModelPatcher(ModelPatcher): + def __enter__(self): + super().__enter__() + + if is_transformers_version(">=", "4.48"): + self.original_feat_extract_output_lengths = MoonshinePreTrainedModel._get_feat_extract_output_lengths + MoonshinePreTrainedModel._get_feat_extract_output_lengths = _get_feat_extract_output_lengths_patched + + def __exit__(self, exc_type, exc_value, traceback): + super().__exit__(exc_type, exc_value, traceback) + + if is_transformers_version(">=", "4.48"): + MoonshinePreTrainedModel._get_feat_extract_output_lengths = self.original_feat_extract_output_lengths + del self.original_feat_extract_output_lengths + + +# This is a traceabe of the original function, +# the original results in a constant shape due to the use of *x.shape[:-1] +def patched_apply_rotary_emb( + x: torch.Tensor, + freqs_cis: torch.Tensor | tuple[torch.Tensor], + use_real: bool = True, + use_real_unbind_dim: int = -1, + sequence_dim: int = 2, +): + if use_real: + cos, sin = freqs_cis # [S, D] + if sequence_dim == 2: + cos = cos[None, None, :, :] + sin = sin[None, None, :, :] + elif sequence_dim == 1: + cos = cos[None, :, None, :] + sin = sin[None, :, None, :] + else: + raise ValueError(f"`sequence_dim={sequence_dim}` but should be 1 or 2.") + + cos, sin = cos.to(x.device), sin.to(x.device) + + if use_real_unbind_dim == -1: + # Used for flux, cogvideox, hunyuan-dit + # x_real, x_imag = x.reshape(*x.shape[:-1], -1, 2).unbind(-1) # [B, H, S, D//2] + # We avoid using reshape here because for some reason it gets exported with constant shape. + x_real = x[..., 0::2] + x_imag = x[..., 1::2] + x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(3) + elif use_real_unbind_dim == -2: + # Used for Stable Audio, OmniGen, CogView4 and Cosmos + # x_real, x_imag = x.reshape(*x.shape[:-1], 2, -1).unbind(-2) # [B, H, S, D//2] + # We avoid using reshape here because for some reason it gets exported with constant shape. + x_real = x[..., 0::2, :] + x_imag = x[..., 1::2, :] + x_rotated = torch.cat([-x_imag, x_real], dim=-1) + else: + raise ValueError(f"`use_real_unbind_dim={use_real_unbind_dim}` but should be -1 or -2.") + + out = (x.float() * cos + x_rotated.float() * sin).to(x.dtype) + + return out + else: + # used for lumina + x_rotated = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2)) + freqs_cis = freqs_cis.unsqueeze(2) + x_out = torch.view_as_real(x_rotated * freqs_cis).flatten(3) + + return x_out.type_as(x) + + +class FluxTransformerModelPatcher(ModelPatcher): + def __enter__(self): + super().__enter__() + + if is_diffusers_version(">=", "0.35.0"): + self.original_apply_rotary_emb = diffusers.models.transformers.transformer_flux.apply_rotary_emb + diffusers.models.transformers.transformer_flux.apply_rotary_emb = patched_apply_rotary_emb + + def __exit__(self, exc_type, exc_value, traceback): + super().__exit__(exc_type, exc_value, traceback) + + if is_diffusers_version(">=", "0.35.0"): + diffusers.models.transformers.transformer_flux.apply_rotary_emb = self.original_apply_rotary_emb + del self.original_apply_rotary_emb + + +def patched_cohere_rotary_forward(self, x, position_ids): + # Get batch size and sequence length for manual expansion + batch_size, _ = position_ids.shape[:2] + + # Instead of using expand, manually repeat the tensor. + # Problem with expand: it creates a view with shared memory rather than copying data, + # which causes ONNX export issues with dynamic shapes and view operations. + # Using repeat() ensures actual memory allocation and data copying for ONNX compatibility. + # original: inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1) + inv_freq_base = self.inv_freq[None, :, None].float() # Shape: [1, freq_dim, 1] + inv_freq_expanded = inv_freq_base.repeat(batch_size, 1, 1) # Shape: [batch_size, freq_dim, 1] + + position_ids_expanded = position_ids[:, None, :].float() + device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu" + + with torch.autocast(device_type=device_type, enabled=False): # Force float32 + freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2) + emb = freqs.repeat_interleave(2, dim=-1) # diff from Llama: we interleave() instead of cat() + cos = emb.cos() * self.attention_scaling + sin = emb.sin() * self.attention_scaling + + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + + +class CohereModelPatcher(ModelPatcher): + def __enter__(self): + super().__enter__() + + if is_transformers_version(">=", "4.38.0"): + from transformers.models.cohere.modeling_cohere import CohereRotaryEmbedding + + self.original_forward = CohereRotaryEmbedding.forward + CohereRotaryEmbedding.forward = patched_cohere_rotary_forward + + def __exit__(self, exc_type, exc_value, traceback): + super().__exit__(exc_type, exc_value, traceback) + + if is_transformers_version(">=", "4.38.0"): + from transformers.models.cohere.modeling_cohere import CohereRotaryEmbedding + + CohereRotaryEmbedding.forward = self.original_forward + + +# Copied from https://github.com/huggingface/transformers/blob/v4.56.0/src/transformers/models/gpt_oss/modeling_gpt_oss.py#L81 +def gpt_oss_forward(self, hidden_states: torch.Tensor, router_indices=None, routing_weights=None) -> torch.Tensor: + batch_size = hidden_states.shape[0] + hidden_states = hidden_states.reshape(-1, self.hidden_size) + num_experts = routing_weights.shape[1] + hidden_states = hidden_states.repeat(num_experts, 1) + hidden_states = hidden_states.view(num_experts, -1, self.hidden_size) + gate_up = torch.bmm(hidden_states, self.gate_up_proj) + self.gate_up_proj_bias[..., None, :] + gate, up = gate_up[..., ::2], gate_up[..., 1::2] + gate = gate.clamp(min=None, max=self.limit) + up = up.clamp(min=-self.limit, max=self.limit) + glu = gate * torch.sigmoid(gate * self.alpha) + next_states = torch.bmm(((up + 1) * glu), self.down_proj) + next_states = next_states + self.down_proj_bias[..., None, :] + next_states = next_states.view(num_experts, batch_size, -1, self.hidden_size) + next_states = next_states * routing_weights.transpose(0, 1).view(num_experts, batch_size, -1)[..., None] + next_states = next_states.sum(dim=0) + return next_states + + +class GptOssModelPatcher(ModelPatcher): + def __enter__(self): + super().__enter__() + + if is_transformers_version(">=", "4.55.0"): + self.original_gpt_oss_forward = GptOssExperts.forward + GptOssExperts.forward = gpt_oss_forward + + def __exit__(self, exc_type, exc_value, traceback): + super().__exit__(exc_type, exc_value, traceback) + + if is_transformers_version(">=", "4.55.0"): + GptOssExperts.forward = self.original_gpt_oss_forward def postprocess_past_key_values(past_key_values): @@ -115,28 +1563,6 @@ def _get_model_attribute(model, name): return getattr(target, name) -# Compatibility wrapper for sdpa_mask_without_vmap from optimum. -# The installed optimum version expects (batch_size, cache_position: Tensor, kv_length, ...), -# but transformers >= 5.5 passes (batch_size, q_length: int, kv_length: int, q_offset: int, ...). -def sdpa_mask_without_vmap(batch_size, q_length=None, kv_length=None, q_offset=0, kv_offset=0, **kwargs): - import inspect - - sig = inspect.signature(_orig_sdpa_mask_without_vmap) - if is_transformers_version(">=", "5.5") and "cache_position" in sig.parameters and q_length is not None: - # Old optimum signature: (batch_size, cache_position, kv_length, kv_offset, ...) - cache_position = torch.arange(q_length, dtype=torch.long) + q_offset - kwargs.pop("q_offset", None) - kwargs.pop("allow_is_bidirectional_skip", None) - kwargs.pop("allow_torch_fix", None) - kwargs.pop("use_vmap", None) - kwargs.pop("device", None) - return _orig_sdpa_mask_without_vmap(batch_size, cache_position, kv_length, kv_offset=kv_offset, **kwargs) - else: - return _orig_sdpa_mask_without_vmap( - batch_size, q_length=q_length, kv_length=kv_length, q_offset=q_offset, kv_offset=kv_offset, **kwargs - ) - - for idx, spec in enumerate(UNSUPPORTED_OPS_PATCHING_SPEC): if spec.name in { # onnx-exporter-specific fixes @@ -9937,3 +11363,54 @@ def __enter__(self): def __exit__(self, exc_type, exc_value, traceback): super().__exit__(exc_type, exc_value, traceback) self._model.forward = self._model._orig_forward + + +################################################################################################################################################################ + + +from transformers.models.vit.modeling_vit import ViTPatchEmbeddings + + +def patched_vit_patch_embedding_forward( + self, pixel_values: torch.Tensor, interpolate_pos_encoding: bool = False +) -> torch.Tensor: + # _batch_size, num_channels, height, width = pixel_values.shape + # Unexpted error. + # TypeError: cond must be a bool, but got ? + # torch._check( + # num_channels == self.num_channels, + # lambda: ( + # "Make sure that the channel dimension of the pixel values match with the one set in the configuration." + # f" Expected {self.num_channels} but got {num_channels}." + # ) + # ) + # This check fails if dynamic shapes are not properly set up. + # Let's drop it. + # if not interpolate_pos_encoding: + # torch._check( + # height == self.image_size[0] and width == self.image_size[1], + # lambda:( + # f"Input image size ({height}*{width}) doesn't match model" + # f" ({self.image_size[0]}*{self.image_size[1]})." + # ) + # ) + embeddings = self.projection(pixel_values).flatten(2).transpose(1, 2) + return embeddings + + +class ViTForImageClassificationPatcher(ModelPatcher): + def __enter__(self): + super().__enter__() + + if is_transformers_version(">=", "4.36.0"): + self.original_forward = ViTPatchEmbeddings.forward + ViTPatchEmbeddings.forward = patched_vit_patch_embedding_forward + + def __exit__(self, exc_type, exc_value, traceback): + super().__exit__(exc_type, exc_value, traceback) + + if is_transformers_version(">=", "4.36.0"): + ViTPatchEmbeddings.forward = self.original_forward + + +################################################################################################################################################################ diff --git a/optimum/exporters/openvino/utils.py b/optimum/exporters/openvino/utils.py index a1756e27a6..0186437047 100644 --- a/optimum/exporters/openvino/utils.py +++ b/optimum/exporters/openvino/utils.py @@ -24,7 +24,7 @@ from openvino import Dimension, PartialShape, Symbol from openvino.utils.types import get_element_type -from optimum.exporters.onnx.base import OnnxConfig +from optimum.exporters.openvino.base import OnnxConfig from optimum.exporters.tasks import TasksManager from optimum.intel.utils.import_utils import is_safetensors_available from optimum.utils import is_diffusers_available diff --git a/setup.py b/setup.py index 66947d1943..ebee6d9cdd 100644 --- a/setup.py +++ b/setup.py @@ -28,7 +28,8 @@ INSTALL_REQUIRE = [ "torch>=2.1", - "optimum-onnx@git+https://github.com/huggingface/optimum-onnx.git@transformers-v5", + "optimum@git+https://github.com/huggingface/optimum.git", + "onnx", "transformers>=4.45,<5.1", "setuptools", "huggingface-hub>=0.23.2,<2.0", From 8c5c5fbeb85a7d90b7ae44784dfe97fb91e138c3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CElla?= Date: Thu, 21 May 2026 17:16:32 +0200 Subject: [PATCH 03/40] style --- .../exporters/openvino/input_generators.py | 3 +- optimum/exporters/openvino/model_configs.py | 192 +++++++++--------- 2 files changed, 94 insertions(+), 101 deletions(-) diff --git a/optimum/exporters/openvino/input_generators.py b/optimum/exporters/openvino/input_generators.py index 7323acedad..ad5cbadb05 100644 --- a/optimum/exporters/openvino/input_generators.py +++ b/optimum/exporters/openvino/input_generators.py @@ -16,6 +16,7 @@ import torch +from optimum.intel.utils.import_utils import is_diffusers_version from optimum.utils import ( DEFAULT_DUMMY_SHAPES, DummyAudioInputGenerator, @@ -35,8 +36,6 @@ from optimum.utils.input_generators import DTYPE_MAPPER from optimum.utils.normalized_config import NormalizedConfig, NormalizedVisionConfig -from ...intel.utils.import_utils import is_diffusers_version - class GPTBigCodeDummyPastKeyValuesGenerator(DummyPastKeyValuesGenerator): def __init__(self, task: str, normalized_config: NormalizedTextConfig, **kwargs): diff --git a/optimum/exporters/openvino/model_configs.py b/optimum/exporters/openvino/model_configs.py index aae4b81b58..139864233a 100644 --- a/optimum/exporters/openvino/model_configs.py +++ b/optimum/exporters/openvino/model_configs.py @@ -35,104 +35,7 @@ TextSeq2SeqOnnxConfig, VisionOnnxConfig, ) -from optimum.exporters.openvino.model_patcher import ( - BigBirdPegasusModelPatcher, - CLIPModelPatcher, - CohereModelPatcher, - FluxTransformerModelPatcher, - GptOssModelPatcher, - MetaCLIP2Patcher, - MgpstrModelPatcher, - ModelPatcher, - MoonshineModelPatcher, - MusicgenModelPatcher, - Qwen3MoeModelPatcher, - SAMModelPatcher, - SentenceTransformersCLIPPatcher, - SentenceTransformersTransformerPatcher, - SpeechT5ModelPatcher, - ViTForImageClassificationPatcher, - VitPoseModelPatcher, -) -from optimum.exporters.openvino.utils import ONNX_SUPPORTED_ARCHITECTURES -from optimum.exporters.tasks import TasksManager -from optimum.utils import ( - DEFAULT_DUMMY_SHAPES, - ASTDummyAudioInputGenerator, - BartDummyTextInputGenerator, - BloomDummyPastKeyValuesGenerator, - DeepSeekV3DummyPastKeyValuesGenerator, - Dinov2DummyInputGenerator, - DummyCodegenDecoderTextInputGenerator, - DummyDecisionTransformerInputGenerator, - DummyDecoderTextInputGenerator, - DummyEncodecInputGenerator, - DummyFluxTransformerTextInputGenerator, - DummyFluxTransformerVisionInputGenerator, - DummyInputGenerator, - DummyIntGenerator, - DummyPastKeyValuesGenerator, - DummyPatchTSTInputGenerator, - DummyPix2StructInputGenerator, - DummyPointsGenerator, - DummySeq2SeqDecoderTextInputGenerator, - DummySeq2SeqPastKeyValuesGenerator, - DummySpeechT5InputGenerator, - DummyTextInputGenerator, - DummyTimestepInputGenerator, - DummyTransformerTextInputGenerator, - DummyTransformerTimestepInputGenerator, - DummyTransformerVisionInputGenerator, - DummyVisionEmbeddingsGenerator, - DummyVisionEncoderDecoderPastKeyValuesGenerator, - DummyVisionInputGenerator, - DummyXPathSeqInputGenerator, - FalconDummyPastKeyValuesGenerator, - GemmaDummyPastKeyValuesGenerator, - LongformerDummyTextInputGenerator, - MCTCTDummyAudioInputGenerator, - MistralDummyPastKeyValuesGenerator, - NormalizedConfig, - NormalizedEncoderDecoderConfig, - NormalizedSeq2SeqConfig, - NormalizedTextAndVisionConfig, - NormalizedTextConfig, - NormalizedTextConfigWithGQA, - NormalizedTimeSeriesForecastingConfig, - NormalizedVisionConfig, - PerceiverDummyInputGenerator, - Speech2TextDummyAudioInputGenerator, - T5DummySeq2SeqPastKeyValuesGenerator, - VitPoseDummyInputGenerator, - is_diffusers_version, - is_transformers_version, -) -from optimum.utils.input_generators import ( - DummyAudioInputGenerator, - DummyInputGenerator, - DummyPastKeyValuesGenerator, - DummySeq2SeqDecoderTextInputGenerator, - DummySeq2SeqPastKeyValuesGenerator, - DummyTextInputGenerator, - DummyVisionInputGenerator, - GemmaDummyPastKeyValuesGenerator, - MistralDummyPastKeyValuesGenerator, -) -from optimum.utils.normalized_config import ( - NormalizedConfig, - NormalizedConfigManager, - NormalizedSeq2SeqConfig, - NormalizedTextConfig, - NormalizedVisionConfig, -) - -from ...intel.utils.import_utils import ( - is_diffusers_available, - is_diffusers_version, - is_openvino_version, - is_transformers_version, -) -from .input_generators import ( +from optimum.exporters.openvino.input_generators import ( AquilaDummyPastKeyValuesGenerator, ChatGLM2DummyPastKeyValuesGenerator, DeciDummyPastKeyValuesGenerator, @@ -182,7 +85,7 @@ QwenDummyPastKeyValuesGenerator, Zamba2DummyPastKeyValuesGenerator, ) -from .model_patcher import ( +from optimum.exporters.openvino.model_patcher import ( AfmoeModelPatcher, AquilaModelPatcher, ArcticModelPatcher, @@ -192,12 +95,15 @@ BlenderbotSmallModelPatcher, BloomModelPatcher, ChatGLMModelPatcher, + CLIPModelPatcher, CodeGenModelPatcher, + CohereModelPatcher, CommonImageEmbeddingsModelPatcher, DBRXModelPatcher, DeciLMModelPatcher, DeepseekPatcher, FalconModelPatcher, + FluxTransformerModelPatcher, FluxTransfromerModelPatcher, Gemma2ModelPatcher, Gemma3LMModelPatcher, @@ -228,13 +134,18 @@ MairaImageEmbeddingModelPatcher, MambaPatcher, MarianModelPatcher, + MetaCLIP2Patcher, + MgpstrModelPatcher, MiniCPM3Patcher, MiniCPMModelPatcher, MiniCPMVImageEmbeddingsModelPatcher, MiniCPMVResamplerModelPatcher, MistralModelPatcher, MixtralModelPatcher, + ModelPatcher, + MoonshineModelPatcher, MPTModelPatcher, + MusicgenModelPatcher, OVDecoderModelPatcher, OVSeq2SeqModelPatcher, OVSpeechT5ModelPatcher, @@ -260,12 +171,95 @@ Qwen3VLLanguageModelPatcher, Qwen3VLVisionEmbMergerPatcher, QwenModelPatcher, + SAMModelPatcher, SanaTextEncoderModelPatcher, + SentenceTransformersCLIPPatcher, + SentenceTransformersTransformerPatcher, + SpeechT5ModelPatcher, VideoChatFlashQwenVisionEmbeddingModelPatcher, + ViTForImageClassificationPatcher, + VitPoseModelPatcher, XverseModelPatcher, Zamba2ModelPatcher, _get_model_attribute, ) +from optimum.exporters.openvino.utils import ONNX_SUPPORTED_ARCHITECTURES +from optimum.exporters.tasks import TasksManager +from optimum.intel.utils.import_utils import ( + is_diffusers_available, + is_diffusers_version, + is_openvino_version, + is_transformers_version, +) +from optimum.utils import ( + DEFAULT_DUMMY_SHAPES, + ASTDummyAudioInputGenerator, + BartDummyTextInputGenerator, + BloomDummyPastKeyValuesGenerator, + DeepSeekV3DummyPastKeyValuesGenerator, + Dinov2DummyInputGenerator, + DummyCodegenDecoderTextInputGenerator, + DummyDecisionTransformerInputGenerator, + DummyDecoderTextInputGenerator, + DummyEncodecInputGenerator, + DummyFluxTransformerTextInputGenerator, + DummyFluxTransformerVisionInputGenerator, + DummyInputGenerator, + DummyIntGenerator, + DummyPastKeyValuesGenerator, + DummyPatchTSTInputGenerator, + DummyPix2StructInputGenerator, + DummyPointsGenerator, + DummySeq2SeqDecoderTextInputGenerator, + DummySeq2SeqPastKeyValuesGenerator, + DummySpeechT5InputGenerator, + DummyTextInputGenerator, + DummyTimestepInputGenerator, + DummyTransformerTextInputGenerator, + DummyTransformerTimestepInputGenerator, + DummyTransformerVisionInputGenerator, + DummyVisionEmbeddingsGenerator, + DummyVisionEncoderDecoderPastKeyValuesGenerator, + DummyVisionInputGenerator, + DummyXPathSeqInputGenerator, + FalconDummyPastKeyValuesGenerator, + GemmaDummyPastKeyValuesGenerator, + LongformerDummyTextInputGenerator, + MCTCTDummyAudioInputGenerator, + MistralDummyPastKeyValuesGenerator, + NormalizedConfig, + NormalizedEncoderDecoderConfig, + NormalizedSeq2SeqConfig, + NormalizedTextAndVisionConfig, + NormalizedTextConfig, + NormalizedTextConfigWithGQA, + NormalizedTimeSeriesForecastingConfig, + NormalizedVisionConfig, + PerceiverDummyInputGenerator, + Speech2TextDummyAudioInputGenerator, + T5DummySeq2SeqPastKeyValuesGenerator, + VitPoseDummyInputGenerator, + is_diffusers_version, + is_transformers_version, +) +from optimum.utils.input_generators import ( + DummyAudioInputGenerator, + DummyInputGenerator, + DummyPastKeyValuesGenerator, + DummySeq2SeqDecoderTextInputGenerator, + DummySeq2SeqPastKeyValuesGenerator, + DummyTextInputGenerator, + DummyVisionInputGenerator, + GemmaDummyPastKeyValuesGenerator, + MistralDummyPastKeyValuesGenerator, +) +from optimum.utils.normalized_config import ( + NormalizedConfig, + NormalizedConfigManager, + NormalizedSeq2SeqConfig, + NormalizedTextConfig, + NormalizedVisionConfig, +) logger = logging.getLogger(__name__) From 7439b8bd4f1b3263739f8f69566f606611cb1b39 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CElla?= Date: Thu, 21 May 2026 18:41:58 +0200 Subject: [PATCH 04/40] remove redundant postprocess_past_key_values --- optimum/exporters/openvino/model_patcher.py | 17 ----------------- 1 file changed, 17 deletions(-) diff --git a/optimum/exporters/openvino/model_patcher.py b/optimum/exporters/openvino/model_patcher.py index d7caa35b4d..22e3cc0682 100644 --- a/optimum/exporters/openvino/model_patcher.py +++ b/optimum/exporters/openvino/model_patcher.py @@ -1541,23 +1541,6 @@ def __exit__(self, exc_type, exc_value, traceback): GptOssExperts.forward = self.original_gpt_oss_forward -def postprocess_past_key_values(past_key_values): - if isinstance(past_key_values, (EncoderDecoderCache, DynamicCache)): - if hasattr(past_key_values, "to_legacy_cache"): - past_key_values = past_key_values.to_legacy_cache() - elif isinstance(past_key_values, DynamicCache): - past_key_values = [(lay.keys, lay.values) for lay in past_key_values.layers] - elif isinstance(past_key_values, EncoderDecoderCache): - past_key_values = [ - (self_lay.keys, self_lay.values, cross_lay.keys, cross_lay.values) - for self_lay, cross_lay in zip( - past_key_values.self_attention_cache.layers, - past_key_values.cross_attention_cache.layers, - ) - ] - return past_key_values - - def _get_model_attribute(model, name): target = getattr(model, "model", model) if is_transformers_version(">=", "5") else model return getattr(target, name) From 3f3b218e2bbfc6e052b11247ef28d55adb465703 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CElla?= Date: Tue, 26 May 2026 16:37:46 +0200 Subject: [PATCH 05/40] onnx to ov config --- .../exporters/openvino/input_generators.py | 2 +- optimum/exporters/openvino/model_configs.py | 4322 +++++------------ optimum/exporters/openvino/model_patcher.py | 250 +- 3 files changed, 1258 insertions(+), 3316 deletions(-) diff --git a/optimum/exporters/openvino/input_generators.py b/optimum/exporters/openvino/input_generators.py index ad5cbadb05..82993cb5a9 100644 --- a/optimum/exporters/openvino/input_generators.py +++ b/optimum/exporters/openvino/input_generators.py @@ -1352,7 +1352,7 @@ def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int return super().generate(input_name, framework, int_dtype, float_dtype) -class DummySpeechT5OpenVINOInputGenerator(DummyInputGenerator): +class DummySpeechT5InputGenerator(DummyInputGenerator): SUPPORTED_INPUT_NAMES = ( "inputs_embeds", "output_sequence", diff --git a/optimum/exporters/openvino/model_configs.py b/optimum/exporters/openvino/model_configs.py index 139864233a..379ce6353e 100644 --- a/optimum/exporters/openvino/model_configs.py +++ b/optimum/exporters/openvino/model_configs.py @@ -14,12 +14,9 @@ import enum import logging -import math -import warnings from copy import deepcopy -from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union +from typing import Any, Dict, List, Optional, Union -from packaging import version from transformers import AutoConfig, PretrainedConfig, PreTrainedModel from optimum.exporters.openvino.base import ConfigBehavior, OnnxConfig, OnnxConfigWithPast, OnnxSeq2SeqConfigWithPast @@ -47,7 +44,6 @@ DummyLLavaMultiModalProjectorInputGenerator, DummyMiniCPMVImageInputGenerator, DummyMiniCPMVResampleInputGenerator, - DummyMoonshineAudioInputGenerator, DummyPhi3VisionProjectionInputGenerator, DummyQwen2VLLMInputGenerator, DummyQwen2VLVisionEmbedInputGenerator, @@ -56,9 +52,8 @@ DummyQwen3VLVisionEmbedInputGenerator, DummySanaSeq2SeqDecoderTextWithEncMaskInputGenerator, DummySanaTimestepInputGenerator, - DummySanaTransforemerTextInputGenerator, DummySanaTransformerVisionInputGenerator, - DummySpeechT5OpenVINOInputGenerator, + DummySpeechT5InputGenerator, DummyTransformerTimestpsInputGenerator, DummyUnetEncoderInputGenerator, DummyUnetTimestepInputGenerator, @@ -97,14 +92,12 @@ ChatGLMModelPatcher, CLIPModelPatcher, CodeGenModelPatcher, - CohereModelPatcher, CommonImageEmbeddingsModelPatcher, DBRXModelPatcher, DeciLMModelPatcher, DeepseekPatcher, FalconModelPatcher, FluxTransformerModelPatcher, - FluxTransfromerModelPatcher, Gemma2ModelPatcher, Gemma3LMModelPatcher, Gemma4ImageEmbeddingsModelPatcher, @@ -134,8 +127,6 @@ MairaImageEmbeddingModelPatcher, MambaPatcher, MarianModelPatcher, - MetaCLIP2Patcher, - MgpstrModelPatcher, MiniCPM3Patcher, MiniCPMModelPatcher, MiniCPMVImageEmbeddingsModelPatcher, @@ -143,12 +134,9 @@ MistralModelPatcher, MixtralModelPatcher, ModelPatcher, - MoonshineModelPatcher, MPTModelPatcher, - MusicgenModelPatcher, OVDecoderModelPatcher, OVSeq2SeqModelPatcher, - OVSpeechT5ModelPatcher, PegasusModelPatcher, PersimmonModelPatcher, Phi3ModelPatcher, @@ -173,12 +161,9 @@ QwenModelPatcher, SAMModelPatcher, SanaTextEncoderModelPatcher, - SentenceTransformersCLIPPatcher, SentenceTransformersTransformerPatcher, SpeechT5ModelPatcher, VideoChatFlashQwenVisionEmbeddingModelPatcher, - ViTForImageClassificationPatcher, - VitPoseModelPatcher, XverseModelPatcher, Zamba2ModelPatcher, _get_model_attribute, @@ -191,77 +176,63 @@ is_openvino_version, is_transformers_version, ) -from optimum.utils import ( - DEFAULT_DUMMY_SHAPES, +from optimum.utils.input_generators import ( ASTDummyAudioInputGenerator, BartDummyTextInputGenerator, BloomDummyPastKeyValuesGenerator, - DeepSeekV3DummyPastKeyValuesGenerator, - Dinov2DummyInputGenerator, - DummyCodegenDecoderTextInputGenerator, - DummyDecisionTransformerInputGenerator, + DummyAudioInputGenerator, DummyDecoderTextInputGenerator, - DummyEncodecInputGenerator, - DummyFluxTransformerTextInputGenerator, - DummyFluxTransformerVisionInputGenerator, DummyInputGenerator, - DummyIntGenerator, DummyPastKeyValuesGenerator, - DummyPatchTSTInputGenerator, DummyPix2StructInputGenerator, DummyPointsGenerator, DummySeq2SeqDecoderTextInputGenerator, DummySeq2SeqPastKeyValuesGenerator, - DummySpeechT5InputGenerator, DummyTextInputGenerator, - DummyTimestepInputGenerator, - DummyTransformerTextInputGenerator, - DummyTransformerTimestepInputGenerator, - DummyTransformerVisionInputGenerator, DummyVisionEmbeddingsGenerator, DummyVisionEncoderDecoderPastKeyValuesGenerator, DummyVisionInputGenerator, - DummyXPathSeqInputGenerator, - FalconDummyPastKeyValuesGenerator, GemmaDummyPastKeyValuesGenerator, - LongformerDummyTextInputGenerator, - MCTCTDummyAudioInputGenerator, + GPTBigCodeDummyPastKeyValuesGenerator, MistralDummyPastKeyValuesGenerator, - NormalizedConfig, - NormalizedEncoderDecoderConfig, - NormalizedSeq2SeqConfig, - NormalizedTextAndVisionConfig, - NormalizedTextConfig, - NormalizedTextConfigWithGQA, - NormalizedTimeSeriesForecastingConfig, - NormalizedVisionConfig, PerceiverDummyInputGenerator, - Speech2TextDummyAudioInputGenerator, T5DummySeq2SeqPastKeyValuesGenerator, - VitPoseDummyInputGenerator, - is_diffusers_version, - is_transformers_version, -) -from optimum.utils.input_generators import ( - DummyAudioInputGenerator, - DummyInputGenerator, - DummyPastKeyValuesGenerator, - DummySeq2SeqDecoderTextInputGenerator, - DummySeq2SeqPastKeyValuesGenerator, - DummyTextInputGenerator, - DummyVisionInputGenerator, - GemmaDummyPastKeyValuesGenerator, - MistralDummyPastKeyValuesGenerator, ) from optimum.utils.normalized_config import ( NormalizedConfig, NormalizedConfigManager, + NormalizedEncoderDecoderConfig, NormalizedSeq2SeqConfig, + NormalizedTextAndVisionConfig, NormalizedTextConfig, NormalizedVisionConfig, ) +COMMON_TEXT_TASKS = [ + "feature-extraction", + "fill-mask", + "multiple-choice", + "question-answering", + "text-classification", + "token-classification", +] + + +COMMON_TEXT_GENERATION_TASKS = [ + "feature-extraction", + "feature-extraction-with-past", + "text-generation", + "text-generation-with-past", +] + +COMMON_TEXT2TEXT_GENERATION_TASKS = [ + *COMMON_TEXT_GENERATION_TASKS, + "text2text-generation", + "text2text-generation-with-past", +] + + logger = logging.getLogger(__name__) @@ -274,10 +245,6 @@ def _warn_potential_accuracy_issue_ov_2026_1(model_type: str, min_transformers_v logger.warning(f"Model type '{model_type}' may have potential accuracy issues with OpenVINO >= 2026.1.0.") -if TYPE_CHECKING: - from transformers.modeling_utils import PreTrainedModel # noqa: F811 - - def init_model_configs(): if "open_clip" not in TasksManager._LIBRARY_TO_SUPPORTED_MODEL_TYPES: TasksManager._LIBRARY_TO_SUPPORTED_MODEL_TYPES["open_clip"] = {} @@ -391,2932 +358,195 @@ def init_model_configs(): register_in_tasks_manager = TasksManager.create_register("openvino", overwrite_existing=True) -COMMON_TEXT_TASKS = [ - "feature-extraction", - "fill-mask", - "multiple-choice", - "question-answering", - "text-classification", - "token-classification", -] - +@register_in_tasks_manager("baichuan", *["text-generation", "text-generation-with-past"], library_name="transformers") +class BaichaunOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): + DEFAULT_ONNX_OPSET = 13 + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args( + num_layers="num_hidden_layers", num_attention_heads="num_attention_heads", hidden_size="hidden_size" + ) + _MODEL_PATCHER = BaichuanModelPatcher + MAX_TRANSFORMERS_VERSION = "4.57.6" -COMMON_TEXT_GENERATION_TASKS = [ - "feature-extraction", - "feature-extraction-with-past", - "text-generation", - "text-generation-with-past", -] -COMMON_TEXT2TEXT_GENERATION_TASKS = [ - *COMMON_TEXT_GENERATION_TASKS, - "text2text-generation", - "text2text-generation-with-past", -] +@register_in_tasks_manager( + "qwen2", + *[ + "text-generation", + "text-generation-with-past", + "feature-extraction", + "feature-extraction-with-past", + "text-classification", + "token-classification", + ], + library_name="transformers", +) +class Qwen2OpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): + DEFAULT_ONNX_OPSET = 14 + DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, MistralDummyPastKeyValuesGenerator) + DUMMY_PKV_GENERATOR_CLASS = MistralDummyPastKeyValuesGenerator + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig + _MODEL_PATCHER = OVDecoderModelPatcher -register_tasks_manager_onnx = TasksManager.create_register("onnx") +@register_in_tasks_manager("qwen2_moe", *["text-generation", "text-generation-with-past"], library_name="transformers") +class Qwen2MoEOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): + DEFAULT_ONNX_OPSET = 14 + DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, MistralDummyPastKeyValuesGenerator) + DUMMY_PKV_GENERATOR_CLASS = MistralDummyPastKeyValuesGenerator + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig + _MODEL_PATCHER = Qwen2MoEPatcher -@register_tasks_manager_onnx("bert", *COMMON_TEXT_TASKS) -class BertOnnxConfig(TextEncoderOnnxConfig): +@register_in_tasks_manager( + "qwen3", + *[ + "text-generation", + "text-generation-with-past", + "feature-extraction", + "feature-extraction-with-past", + "text-classification", + ], + library_name="transformers", +) +class Qwen3OpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): + MIN_TRANSFORMERS_VERSION = "4.51.0" + DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, GemmaDummyPastKeyValuesGenerator) + DUMMY_PKV_GENERATOR_CLASS = GemmaDummyPastKeyValuesGenerator NORMALIZED_CONFIG_CLASS = NormalizedTextConfig + _MODEL_PATCHER = OVDecoderModelPatcher @property - def inputs(self) -> dict[str, dict[int, str]]: - if self.task == "multiple-choice": - dynamic_axis = {0: "batch_size", 1: "num_choices", 2: "sequence_length"} + def inputs(self) -> Dict[str, Dict[int, str]]: + if self.task in ["feature-extraction"]: + common_inputs = { + "input_ids": {0: "batch_size", 1: "sequence_length"}, + "attention_mask": {0: "batch_size", 1: "sequence_length"}, + } else: - dynamic_axis = {0: "batch_size", 1: "sequence_length"} - return { - "input_ids": dynamic_axis, - "attention_mask": dynamic_axis, - "token_type_ids": dynamic_axis, - } + common_inputs = super().inputs + return common_inputs -@register_tasks_manager_onnx("visual_bert", *["feature-extraction"]) -class VisualBertOnnxConfig(TextAndVisionOnnxConfig): - NORMALIZED_CONFIG_CLASS = NormalizedTextConfig - DUMMY_INPUT_GENERATOR_CLASSES = ( - DummyTextInputGenerator, - DummyVisionInputGenerator, - ) +@register_in_tasks_manager( + "qwen3_vl_text", + *[ + "text-generation", + "text-generation-with-past", + ], + library_name="transformers", +) +class Qwen3VLTextOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): + MIN_TRANSFORMERS_VERSION = "4.57.0" - @property - def inputs(self) -> dict[str, dict[int, str]]: - return { - "input_ids": {0: "batch_size", 1: "sequence_length"}, - "attention_mask": {0: "batch_size", 1: "sequence_length"}, - "visual_embeds": {0: "batch_size", 1: "visual_seq_length", 2: "visual_embedding_dim"}, - "visual_attention_mask": {0: "batch_size", 1: "visual_seq_length"}, - "visual_token_type_ids": {0: "batch_size", 1: "visual_seq_length"}, - } + DUMMY_INPUT_GENERATOR_CLASSES = (DummyQwen3VLLMInputGenerator, GemmaDummyPastKeyValuesGenerator) + DUMMY_PKV_GENERATOR_CLASS = GemmaDummyPastKeyValuesGenerator + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig + _MODEL_PATCHER = OVDecoderModelPatcher @property - def outputs(self) -> dict[str, dict[int, str]]: - return { - "last_hidden_state": {0: "batch_size", 1: "sequence_length + visual_seq_length"}, - } + def inputs(self) -> Dict[str, Dict[int, str]]: + common_inputs = super().inputs + common_inputs["visual_pos_masks"] = {0: "batch_size", 1: "sequence_length"} + common_inputs["deepstack_visual_embeds"] = {0: "num_layers", 1: "visual_seqlen"} + return common_inputs -@register_tasks_manager_onnx("albert", *COMMON_TEXT_TASKS) -class AlbertOnnxConfig(BertOnnxConfig): - pass +@register_in_tasks_manager( + "qwen3_moe", + *["text-generation", "text-generation-with-past", "feature-extraction", "feature-extraction-with-past"], + library_name="transformers", +) +class Qwen3MoEOpenVINOConfig(Qwen3OpenVINOConfig): + _MODEL_PATCHER = Qwen3MoeModelPatcher -@register_tasks_manager_onnx("convbert", *COMMON_TEXT_TASKS) -class ConvBertOnnxConfig(BertOnnxConfig): - pass +@register_in_tasks_manager("minicpm", *["text-generation", "text-generation-with-past"], library_name="transformers") +class MiniCPMOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): + DEFAULT_ONNX_OPSET = 14 + MAX_TRANSFORMERS_VERSION = "4.53.3" + DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, MistralDummyPastKeyValuesGenerator) + DUMMY_PKV_GENERATOR_CLASS = MistralDummyPastKeyValuesGenerator + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig + _MODEL_PATCHER = MiniCPMModelPatcher -@register_tasks_manager_onnx("electra", *COMMON_TEXT_TASKS) -class ElectraOnnxConfig(BertOnnxConfig): - pass +@register_in_tasks_manager("minicpm3", *["text-generation", "text-generation-with-past"], library_name="transformers") +class MiniCPM3OpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): + DEFAULT_ONNX_OPSET = 14 + MAX_TRANSFORMERS_VERSION = "4.53.3" + DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, OVMiniCPM3DummyPastKeyValuesGenerator) + DUMMY_PKV_GENERATOR_CLASS = OVMiniCPM3DummyPastKeyValuesGenerator + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig + _MODEL_PATCHER = MiniCPM3Patcher -@register_tasks_manager_onnx("roformer", *COMMON_TEXT_TASKS) -class RoFormerOnnxConfig(BertOnnxConfig): - pass +@register_in_tasks_manager("stablelm", *["text-generation", "text-generation-with-past"], library_name="transformers") +class StableLMOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): + DEFAULT_ONNX_OPSET = 14 + DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, MistralDummyPastKeyValuesGenerator) + DUMMY_PKV_GENERATOR_CLASS = MistralDummyPastKeyValuesGenerator + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig + _MODEL_PATCHER = OVDecoderModelPatcher -@register_tasks_manager_onnx("squeezebert", *COMMON_TEXT_TASKS) -class SqueezeBertOnnxConfig(BertOnnxConfig): - pass +@register_in_tasks_manager("chatglm", *["text-generation", "text-generation-with-past"], library_name="transformers") +class ChatGLM2OpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args(vocab_size="padded_vocab_size", num_layers="num_layers") + DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, ChatGLM2DummyPastKeyValuesGenerator) + DUMMY_PKV_GENERATOR_CLASS = ChatGLM2DummyPastKeyValuesGenerator + _MODEL_PATCHER = ChatGLMModelPatcher + MAX_TRANSFORMERS_VERSION = "4.55.4" + def generate_dummy_inputs(self, framework: str = "pt", **kwargs): + dummy_inputs_generators = self._create_dummy_input_generator_classes(**kwargs) -@register_tasks_manager_onnx("mobilebert", *COMMON_TEXT_TASKS) -class MobileBertOnnxConfig(BertOnnxConfig): - pass + dummy_inputs = {} + input_names = [key for key in self.inputs.keys() if not key.startswith("past_key_values")] + if self.use_past_in_inputs and self.use_cache_branch is not False: + input_names.append("past_key_values") + for input_name in input_names: + input_was_inserted = False + for dummy_input_gen in dummy_inputs_generators: + if dummy_input_gen.supports_input(input_name): + dummy_inputs[input_name] = self.overwrite_shape_and_generate_input( + dummy_input_gen, + input_name, + framework, + input_shapes=kwargs, + ) + input_was_inserted = True + break + if not input_was_inserted: + raise RuntimeError( + f'Could not generate dummy input for "{input_name}". Try adding a proper dummy input generator to the model ONNX config.' + ) -@register_tasks_manager_onnx("nystromformer", *COMMON_TEXT_TASKS) -class NystromformerOnnxConfig(BertOnnxConfig): - pass + # refer to https://github.com/huggingface/optimum/pull/764 + if ( + self.use_past_in_inputs + and self.PAD_ATTENTION_MASK_TO_PAST + and self.use_cache_branch is not False + and "attention_mask" in dummy_inputs + ): + # Obtain the past sequence length from the value instead of the key (Bloom). ChatGLM has seq_len in 0 dim instead of -2 + seq_len_dim = 0 if not hasattr(self._normalized_config, "rope_ratio") else -2 + past_present_length = ( + dummy_inputs["input_ids"].shape[1] + dummy_inputs["past_key_values"][0][1].shape[seq_len_dim] + ) + dummy_inputs["attention_mask"] = DummyInputGenerator.pad_input_on_dim( + dummy_inputs["attention_mask"], + desired_length=past_present_length, + dim=1, + dtype=dummy_inputs["attention_mask"].dtype, + ) -@register_tasks_manager_onnx("xlm", *COMMON_TEXT_TASKS) -class XLMOnnxConfig(BertOnnxConfig): - pass + return dummy_inputs - -@register_tasks_manager_onnx("splinter", *["feature-extraction", "question-answering"]) -class SplinterOnnxConfig(BertOnnxConfig): - pass - - -@register_tasks_manager_onnx("rembert", *COMMON_TEXT_TASKS) -class RemBertOnnxConfig(BertOnnxConfig): - pass - - -@register_tasks_manager_onnx("longformer", *COMMON_TEXT_TASKS) -class LongformerOnnxConfig(BertOnnxConfig): - DUMMY_INPUT_GENERATOR_CLASSES = (LongformerDummyTextInputGenerator,) - - @property - def inputs(self) -> dict[str, dict[int, str]]: - inputs = super().inputs - - inputs["global_attention_mask"] = inputs["attention_mask"] - - return inputs - - -@register_tasks_manager_onnx("megatron-bert", *COMMON_TEXT_TASKS) -class MegatronBertOnnxConfig(BertOnnxConfig): - pass - - -@register_tasks_manager_onnx("distilbert", *COMMON_TEXT_TASKS) -class DistilBertOnnxConfig(BertOnnxConfig): - @property - def inputs(self) -> dict[str, dict[int, str]]: - if self.task == "multiple-choice": - dynamic_axis = {0: "batch_size", 1: "num_choices", 2: "sequence_length"} - else: - dynamic_axis = {0: "batch_size", 1: "sequence_length"} - return {"input_ids": dynamic_axis, "attention_mask": dynamic_axis} - - -@register_tasks_manager_onnx( - "modernbert", - *[ - "feature-extraction", - "fill-mask", - "text-classification", - "token-classification", - ], -) -class ModernBertOnnxConfig(DistilBertOnnxConfig): - MIN_TRANSFORMERS_VERSION = version.parse("4.48.0") - - -@register_tasks_manager_onnx("mpnet", *COMMON_TEXT_TASKS) -class MPNetOnnxConfig(DistilBertOnnxConfig): - pass - - -@register_tasks_manager_onnx("roberta", *COMMON_TEXT_TASKS) -class RobertaOnnxConfig(DistilBertOnnxConfig): - pass - - -@register_tasks_manager_onnx("camembert", *COMMON_TEXT_TASKS) -class CamembertOnnxConfig(DistilBertOnnxConfig): - pass - - -@register_tasks_manager_onnx("flaubert", *COMMON_TEXT_TASKS) -class FlaubertOnnxConfig(BertOnnxConfig): - pass - - -@register_tasks_manager_onnx("ibert", *COMMON_TEXT_TASKS) -class IBertOnnxConfig(DistilBertOnnxConfig): - pass - - -@register_tasks_manager_onnx("xlm-roberta", *COMMON_TEXT_TASKS) -class XLMRobertaOnnxConfig(DistilBertOnnxConfig): - pass - - -@register_tasks_manager_onnx( - "deberta", - *["feature-extraction", "fill-mask", "text-classification", "token-classification", "question-answering"], -) -class DebertaOnnxConfig(BertOnnxConfig): - @property - def inputs(self) -> dict[str, dict[int, str]]: - common_inputs = super().inputs - if self._config.type_vocab_size == 0: - common_inputs.pop("token_type_ids") - return common_inputs - - -@register_tasks_manager_onnx( - "markuplm", *["feature-extraction", "text-classification", "token-classification", "question-answering"] -) -class MarkupLMOnnxConfig(BertOnnxConfig): - DUMMY_INPUT_GENERATOR_CLASSES = ( - DummyTextInputGenerator, - DummyXPathSeqInputGenerator, - ) - - @property - def inputs(self) -> dict[str, dict[int, str]]: - dynamic_axis = {0: "batch_size", 1: "sequence_length"} - xpath_dynamic_axis = {0: "batch_size", 1: "sequence_length", 2: "max_depth"} - return { - "input_ids": dynamic_axis, - "attention_mask": dynamic_axis, - "token_type_ids": dynamic_axis, - "xpath_subs_seq": xpath_dynamic_axis, - "xpath_tags_seq": xpath_dynamic_axis, - } - - -@register_tasks_manager_onnx("deberta-v2", *COMMON_TEXT_TASKS) -class DebertaV2OnnxConfig(DebertaOnnxConfig): - pass - - -@register_tasks_manager_onnx( - "esm", *["feature-extraction", "fill-mask", "text-classification", "token-classification"] -) -class EsmOnnxConfig(TextEncoderOnnxConfig): - NORMALIZED_CONFIG_CLASS = NormalizedTextConfig - - @property - def inputs(self) -> dict[str, dict[int, str]]: - dynamic_axis = {0: "batch_size", 1: "sequence_length"} - return { - "input_ids": dynamic_axis, - "attention_mask": dynamic_axis, - } - - -@register_tasks_manager_onnx("gpt2", *[*COMMON_TEXT_GENERATION_TASKS, "text-classification", "token-classification"]) -class GPT2OnnxConfig(TextDecoderWithPositionIdsOnnxConfig): - NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args(num_layers="n_layer", num_attention_heads="n_head") - - -@register_tasks_manager_onnx("gptj", *[*COMMON_TEXT_GENERATION_TASKS, "text-classification", "question-answering"]) -class GPTJOnnxConfig(GPT2OnnxConfig): - pass - - -@register_tasks_manager_onnx("codegen", *COMMON_TEXT_GENERATION_TASKS) -class CodeGenOnnxConfig(GPT2OnnxConfig): - pass - - -@register_tasks_manager_onnx("imagegpt", *["feature-extraction", "image-classification"]) -class ImageGPTOnnxConfig(GPT2OnnxConfig): - pass - - -@register_tasks_manager_onnx("decision_transformer", *["feature-extraction", "reinforcement-learning"]) -class DecisionTransformerOnnxConfig(OnnxConfig): - DUMMY_INPUT_GENERATOR_CLASSES = (DummyDecisionTransformerInputGenerator,) - NORMALIZED_CONFIG_CLASS = NormalizedConfig - - @property - def inputs(self) -> dict[str, dict[int, str]]: - return { - "states": {0: "batch_size", 1: "sequence_length"}, - "actions": {0: "batch_size", 1: "sequence_length"}, - "timesteps": {0: "batch_size", 1: "sequence_length"}, - "returns_to_go": {0: "batch_size", 1: "sequence_length"}, - "attention_mask": {0: "batch_size", 1: "sequence_length"}, - } - - @property - def outputs(self) -> dict[str, dict[int, str]]: - return { - "state_preds": {0: "batch_size", 1: "sequence_length"}, - "action_preds": {0: "batch_size", 1: "sequence_length"}, - "return_preds": {0: "batch_size", 1: "sequence_length"}, - "last_hidden_state": {0: "batch_size", 1: "sequence_length"}, - } - - -@register_tasks_manager_onnx("gpt_neo", *[*COMMON_TEXT_GENERATION_TASKS, "text-classification"]) -class GPTNeoOnnxConfig(TextDecoderWithPositionIdsOnnxConfig): - NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args(num_attention_heads="num_heads") - - -@register_tasks_manager_onnx("gpt_neox", *[*COMMON_TEXT_GENERATION_TASKS, "text-classification"]) -class GPTNeoXOnnxConfig(TextDecoderWithPositionIdsOnnxConfig): - NORMALIZED_CONFIG_CLASS = NormalizedTextConfig - - -@register_tasks_manager_onnx("opt", *[*COMMON_TEXT_GENERATION_TASKS, "text-classification", "question-answering"]) -class OPTOnnxConfig( - TextDecoderWithPositionIdsOnnxConfig if is_transformers_version(">=", "4.46.0") else TextDecoderOnnxConfig -): - NORMALIZED_CONFIG_CLASS = NormalizedTextConfig - - -@register_tasks_manager_onnx("llama", *[*COMMON_TEXT_GENERATION_TASKS, "text-classification"]) -class LlamaOnnxConfig(TextDecoderWithPositionIdsOnnxConfig): - DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, MistralDummyPastKeyValuesGenerator) - DUMMY_PKV_GENERATOR_CLASS = MistralDummyPastKeyValuesGenerator - NORMALIZED_CONFIG_CLASS = NormalizedTextConfig - - -@register_tasks_manager_onnx("arcee", *COMMON_TEXT_GENERATION_TASKS) -class ArceeOnnxConfig(LlamaOnnxConfig): - MIN_TRANSFORMERS_VERSION = version.parse("4.53.0") - NORMALIZED_CONFIG_CLASS = NormalizedTextConfigWithGQA - - -@register_tasks_manager_onnx("deepseek_v3", *COMMON_TEXT_GENERATION_TASKS) -class DeepSeekV3OnnxConfig(LlamaOnnxConfig): - MIN_TRANSFORMERS_VERSION = version.parse("4.51.0") - DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, DeepSeekV3DummyPastKeyValuesGenerator) - DUMMY_PKV_GENERATOR_CLASS = DeepSeekV3DummyPastKeyValuesGenerator - - -@register_tasks_manager_onnx("cohere", *COMMON_TEXT_GENERATION_TASKS) -class CohereOnnxConfig(LlamaOnnxConfig): - MIN_TRANSFORMERS_VERSION = version.parse("4.38.0") - NORMALIZED_CONFIG_CLASS = NormalizedTextConfig - _MODEL_PATCHER = CohereModelPatcher - - -@register_tasks_manager_onnx("glm", *COMMON_TEXT_GENERATION_TASKS) -class GLMOnnxConfig(LlamaOnnxConfig): - MIN_TRANSFORMERS_VERSION = version.parse("4.46.0") - - -@register_tasks_manager_onnx("helium", *COMMON_TEXT_GENERATION_TASKS) -class HeliumOnnxConfig(LlamaOnnxConfig): - MIN_TRANSFORMERS_VERSION = version.parse("4.49.0") - - -@register_tasks_manager_onnx("smollm3", *[*COMMON_TEXT_GENERATION_TASKS, "text-classification"]) -class SmolLM3OnnxConfig(LlamaOnnxConfig): - MIN_TRANSFORMERS_VERSION = version.parse("4.53.0") - - -@register_tasks_manager_onnx("stablelm", *COMMON_TEXT_GENERATION_TASKS) -class StableLMOnnxConfig(LlamaOnnxConfig): - MIN_TRANSFORMERS_VERSION = version.parse("4.38.0") - - -@register_tasks_manager_onnx("olmo", *COMMON_TEXT_GENERATION_TASKS) -class OlmoOnnxConfig(LlamaOnnxConfig): - MIN_TRANSFORMERS_VERSION = version.parse("4.40.0") - - -@register_tasks_manager_onnx("olmo2", *COMMON_TEXT_GENERATION_TASKS) -class Olmo2OnnxConfig(OlmoOnnxConfig): - MIN_TRANSFORMERS_VERSION = version.parse("4.47.0") - - -@register_tasks_manager_onnx("qwen2", *[*COMMON_TEXT_GENERATION_TASKS, "text-classification", "token-classification"]) -class Qwen2OnnxConfig(LlamaOnnxConfig): - NORMALIZED_CONFIG_CLASS = NormalizedTextConfigWithGQA - MIN_TRANSFORMERS_VERSION = version.parse("4.37.0") - - -@register_tasks_manager_onnx("qwen3", *[*COMMON_TEXT_GENERATION_TASKS, "text-classification"]) -class Qwen3OnnxConfig(LlamaOnnxConfig): - NORMALIZED_CONFIG_CLASS = NormalizedTextConfigWithGQA - MIN_TRANSFORMERS_VERSION = version.parse("4.51.0") - - -@register_tasks_manager_onnx( - "qwen3_moe", *[*COMMON_TEXT_GENERATION_TASKS, "text-classification", "token-classification"] -) -class Qwen3MoeOnnxConfig(LlamaOnnxConfig): - NORMALIZED_CONFIG_CLASS = NormalizedTextConfigWithGQA - MIN_TRANSFORMERS_VERSION = version.parse("4.51.0") - _MODEL_PATCHER = Qwen3MoeModelPatcher - - -@register_tasks_manager_onnx("gemma", *[*COMMON_TEXT_GENERATION_TASKS, "text-classification"]) -class GemmaOnnxConfig(TextDecoderOnnxConfig): - NORMALIZED_CONFIG_CLASS = NormalizedTextConfig - DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, GemmaDummyPastKeyValuesGenerator) - DUMMY_PKV_GENERATOR_CLASS = GemmaDummyPastKeyValuesGenerator - MIN_TRANSFORMERS_VERSION = version.parse("4.38.0") - - -@register_tasks_manager_onnx("gemma2", *[*COMMON_TEXT_GENERATION_TASKS, "text-classification"]) -class Gemma2OnnxConfig(GemmaOnnxConfig): - # Gemma 2 was added in transformers v4.42 using HybridCache - # DynamicCache support was added since v4.53 - MIN_TRANSFORMERS_VERSION = version.parse("4.53.0") - - -@register_tasks_manager_onnx("gemma3_text", *COMMON_TEXT_GENERATION_TASKS) -class Gemma3TextOnnxConfig(GemmaOnnxConfig): - # Gemma 3 was added in transformers v4.50 using HybridCache - # DynamicCache support was added since v4.53 - MIN_TRANSFORMERS_VERSION = version.parse("4.53.0") - - -# we still don't support gemma3 for multimodal feature-extraction(-with-past) and image-text-to-text(-with-past) tasks -@register_tasks_manager_onnx("gemma3", *COMMON_TEXT_GENERATION_TASKS, "text-classification") -class Gemma3OnnxConfig(GemmaOnnxConfig): - # Gemma 3 was added in transformers v4.50 using HybridCache - # DynamicCache support was added since v4.53 - MIN_TRANSFORMERS_VERSION = version.parse("4.53.0") - - def __init__(self, config: PretrainedConfig, **kwargs): - super().__init__(config.text_config, **kwargs) - - -@register_tasks_manager_onnx("gpt_oss", *COMMON_TEXT_GENERATION_TASKS) -class GPTOssOnnxConfig(GemmaOnnxConfig): - MIN_TRANSFORMERS_VERSION = version.parse("4.55.0") - _MODEL_PATCHER = GptOssModelPatcher - - -@register_tasks_manager_onnx("nemotron", *COMMON_TEXT_GENERATION_TASKS) -class NemotronOnnxConfig(GemmaOnnxConfig): - MIN_TRANSFORMERS_VERSION = version.parse("4.48.0") # More stable version than 4.44.0 - NORMALIZED_CONFIG_CLASS = NormalizedTextConfigWithGQA - - -@register_tasks_manager_onnx("granite", *COMMON_TEXT_GENERATION_TASKS) -class GraniteOnnxConfig(LlamaOnnxConfig): - MIN_TRANSFORMERS_VERSION = version.parse("4.45.0") - - -@register_tasks_manager_onnx("phi", *[*COMMON_TEXT_GENERATION_TASKS, "text-classification"]) -class PhiOnnxConfig(TextDecoderWithPositionIdsOnnxConfig): - NORMALIZED_CONFIG_CLASS = NormalizedTextConfig - MIN_TRANSFORMERS_VERSION = version.parse("4.36.0") - - -@register_tasks_manager_onnx("phi3", *[*COMMON_TEXT_GENERATION_TASKS, "text-classification"]) -class Phi3OnnxConfig(PhiOnnxConfig): - DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, MistralDummyPastKeyValuesGenerator) - DUMMY_PKV_GENERATOR_CLASS = MistralDummyPastKeyValuesGenerator - NORMALIZED_CONFIG_CLASS = NormalizedTextConfigWithGQA - MIN_TRANSFORMERS_VERSION = version.parse("4.41.0") - - -@register_tasks_manager_onnx("internlm2", *["text-generation", "text-generation-with-past"]) -class InternLM2OnnxConfig(LlamaOnnxConfig): - MIN_TRANSFORMERS_VERSION = version.parse("4.41.0") - - -@register_tasks_manager_onnx("mistral", *[*COMMON_TEXT_GENERATION_TASKS, "text-classification"]) -class MistralOnnxConfig(TextDecoderWithPositionIdsOnnxConfig): - NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args(num_key_value_heads="num_key_value_heads", allow_new=True) - DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, MistralDummyPastKeyValuesGenerator) - DUMMY_PKV_GENERATOR_CLASS = MistralDummyPastKeyValuesGenerator - - -@register_tasks_manager_onnx("mpt", *[*COMMON_TEXT_GENERATION_TASKS, "text-classification", "token-classification"]) -class MPTOnnxConfig(TextDecoderOnnxConfig): - NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args( - num_attention_heads="n_heads", hidden_size="d_model", num_layers="n_layers" - ) - - -@register_tasks_manager_onnx("bloom", *[*COMMON_TEXT_GENERATION_TASKS, "text-classification", "token-classification"]) -class BloomOnnxConfig(TextDecoderOnnxConfig): - # Bloom does not require position_ids input. - MIN_TRANSFORMERS_VERSION = version.parse("4.36.0") - DUMMY_PKV_GENERATOR_CLASS = BloomDummyPastKeyValuesGenerator - DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, BloomDummyPastKeyValuesGenerator) - NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args(num_layers="n_layer", num_attention_heads="n_head") - - def add_past_key_values(self, inputs_or_outputs: dict[str, dict[int, str]], direction: str): - if is_transformers_version(">=", "4.44"): - super().add_past_key_values(inputs_or_outputs, direction) - else: - if direction not in ["inputs", "outputs"]: - raise ValueError(f'direction must either be "inputs" or "outputs", but {direction} was given') - - if direction == "inputs": - decoder_sequence_name = "past_sequence_length" - name = "past_key_values" - else: - decoder_sequence_name = "past_sequence_length + sequence_length" - name = "present" - - for i in range(self._normalized_config.num_layers): - inputs_or_outputs[f"{name}.{i}.key"] = {0: "batch_size * num_heads", 2: decoder_sequence_name} - inputs_or_outputs[f"{name}.{i}.value"] = {0: "batch_size * num_heads", 1: decoder_sequence_name} - - -@register_tasks_manager_onnx( - "gpt_bigcode", *[*COMMON_TEXT_GENERATION_TASKS, "text-classification", "token-classification"] -) -class GPTBigCodeOnnxConfig(TextDecoderWithPositionIdsOnnxConfig): - DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, GPTBigCodeDummyPastKeyValuesGenerator) - NORMALIZED_CONFIG_CLASS = NormalizedConfigManager.get_normalized_config_class("gpt_bigcode") - DUMMY_PKV_GENERATOR_CLASS = GPTBigCodeDummyPastKeyValuesGenerator - - def add_past_key_values(self, inputs_or_outputs: dict[str, dict[int, str]], direction: str): - if is_transformers_version(">=", "4.54"): - super().add_past_key_values(inputs_or_outputs, direction) - else: - if direction not in ["inputs", "outputs"]: - raise ValueError(f'direction must either be "inputs" or "outputs", but {direction} was given') - - if direction == "inputs": - decoder_sequence_name = "past_sequence_length" - name = "past_key_values" - else: - decoder_sequence_name = "past_sequence_length + sequence_length" - name = "present" - - if self._normalized_config.multi_query: - decoder_sequence_dim = 1 - else: - decoder_sequence_dim = 2 - - for i in range(self._normalized_config.num_layers): - inputs_or_outputs[f"{name}.{i}.key_value"] = { - 0: "batch_size", - decoder_sequence_dim: decoder_sequence_name, - } - - def flatten_past_key_values(self, flattened_output, name, idx, t): - if is_transformers_version(">=", "4.54"): - super().flatten_past_key_values(flattened_output, name, idx, t) - else: - flattened_output[f"{name}.{idx}.key_value"] = t - - -@register_tasks_manager_onnx("falcon", *[*COMMON_TEXT_GENERATION_TASKS, "question-answering", "token-classification"]) -class FalconOnnxConfig(TextDecoderWithPositionIdsOnnxConfig): - DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, FalconDummyPastKeyValuesGenerator) - DUMMY_PKV_GENERATOR_CLASS = FalconDummyPastKeyValuesGenerator - NORMALIZED_CONFIG_CLASS = NormalizedTextConfig - - @property - def inputs(self) -> dict[str, dict[int, str]]: - common_inputs = super().inputs - - if self._config.alibi: - common_inputs.pop("position_ids", None) - - return common_inputs - - -@register_tasks_manager_onnx( - "t5", - *["feature-extraction", "feature-extraction-with-past", "text2text-generation", "text2text-generation-with-past"], -) -class T5OnnxConfig(TextSeq2SeqOnnxConfig): - DUMMY_INPUT_GENERATOR_CLASSES = ( - *TextSeq2SeqOnnxConfig.DUMMY_INPUT_GENERATOR_CLASSES[:-1], - T5DummySeq2SeqPastKeyValuesGenerator, - ) - DUMMY_PKV_GENERATOR_CLASS = T5DummySeq2SeqPastKeyValuesGenerator - NORMALIZED_CONFIG_CLASS = NormalizedSeq2SeqConfig.with_args( - hidden_size="d_model", - num_attention_heads="num_heads", - encoder_num_layers="num_layers", - decoder_num_layers="num_decoder_layers", - key_value_dim="d_kv", - allow_new=True, - ) - - -@register_tasks_manager_onnx( - "mt5", - *["feature-extraction", "feature-extraction-with-past", "text2text-generation", "text2text-generation-with-past"], -) -class MT5OnnxConfig(T5OnnxConfig): - pass - - -@register_tasks_manager_onnx( - "longt5", - *["feature-extraction", "feature-extraction-with-past", "text2text-generation", "text2text-generation-with-past"], -) -class LongT5OnnxConfig(T5OnnxConfig): - pass - - -@register_tasks_manager_onnx( - "m2m_100", - *["feature-extraction", "feature-extraction-with-past", "text2text-generation", "text2text-generation-with-past"], -) -class M2M100OnnxConfig(TextSeq2SeqOnnxConfig): - PAD_ATTENTION_MASK_TO_PAST = True - - NORMALIZED_CONFIG_CLASS = NormalizedSeq2SeqConfig.with_args( - encoder_num_layers="encoder_layers", - decoder_num_layers="decoder_layers", - num_layers="decoder_layers", # Used for the text-generation task past key values input generation. - encoder_num_attention_heads="encoder_attention_heads", - decoder_num_attention_heads="decoder_attention_heads", - eos_token_id="eos_token_id", - ) - DUMMY_INPUT_GENERATOR_CLASSES = ( - BartDummyTextInputGenerator, - { - "feature-extraction": DummySeq2SeqDecoderTextInputGenerator, - "text-generation": DummyDecoderTextInputGenerator, - }, - { - "feature-extraction": DummySeq2SeqPastKeyValuesGenerator, - "text-generation": DummyPastKeyValuesGenerator, - }, - ) - - def _create_dummy_input_generator_classes(self, **kwargs) -> list[DummyInputGenerator]: - dummy_text_input_generator = self.DUMMY_INPUT_GENERATOR_CLASSES[0]( - self.task, self._normalized_config, **kwargs - ) - task = "feature-extraction" if self.task != "text-generation" else "text-generation" - dummy_decoder_text_input_generator = self.DUMMY_INPUT_GENERATOR_CLASSES[1][task]( - self.task, self._normalized_config, **kwargs - ) - if self.task != "text-generation": - kwargs["encoder_sequence_length"] = dummy_text_input_generator.sequence_length - - dummy_seq2seq_past_key_values_generator = self.DUMMY_INPUT_GENERATOR_CLASSES[2][task]( - self.task, self._normalized_config, **kwargs - ) - dummy_inputs_generators = [ - dummy_text_input_generator, - dummy_decoder_text_input_generator, - dummy_seq2seq_past_key_values_generator, - ] - - return dummy_inputs_generators - - @property - def inputs_for_default_and_seq2seq_lm(self): - return super().inputs - - @property - def inputs_for_causal_lm(self): - if self.use_past_in_inputs: - common_inputs = { - "input_ids": {0: "batch_size", 1: "sequence_length"}, - "attention_mask": {0: "batch_size", 1: "past_sequence_length + sequence_length"}, - } - for i in range(self._normalized_config.decoder_num_layers): - common_inputs[f"past_key_values.{i}.key"] = { - 0: "batch_size", - 2: "past_sequence_length", - } - common_inputs[f"past_key_values.{i}.value"] = { - 0: "batch_size", - 2: "past_sequence_length", - } - else: - common_inputs = { - "input_ids": {0: "batch_size", 1: "sequence_length"}, - "attention_mask": {0: "batch_size", 1: "sequence_length"}, - } - - return common_inputs - - @property - def inputs_for_other_tasks(self): - return { - "input_ids": {0: "batch_size", 1: "sequence_length"}, - "attention_mask": {0: "batch_size", 1: "sequence_length"}, - } - - @property - def inputs(self) -> dict[str, dict[int, str]]: - inputs_properties = { - "feature-extraction": self.inputs_for_default_and_seq2seq_lm, - "text2text-generation": self.inputs_for_default_and_seq2seq_lm, - "text-generation": self.inputs_for_causal_lm, - "other": self.inputs_for_other_tasks, - } - return inputs_properties.get(self.task, inputs_properties["other"]) - - @property - def outputs(self) -> dict[str, dict[int, str]]: - if self.task in ["feature-extraction", "text2text-generation"]: - common_outputs = super().outputs - else: - common_outputs = super(OnnxConfigWithPast, self).outputs - if self.use_past: - # When exporting decoder models with use_cache=True, both the decoder without past and with past have the KV cache as an output. - for i in range( - self._normalized_config.encoder_num_layers - if self.task != "text-generation" - else self._normalized_config.decoder_num_layers - ): - common_outputs[f"present.{i}.key"] = {0: "batch_size", 2: "past_sequence_length + sequence_length"} - common_outputs[f"present.{i}.value"] = { - 0: "batch_size", - 2: "past_sequence_length + sequence_length", - } - return common_outputs - - def flatten_past_key_values(self, flattened_output, name, idx, t): - if self.task in ["feature-extraction", "text2text-generation"]: - flattened_output = super().flatten_past_key_values(flattened_output, name, idx, t) - else: - flattened_output = super(OnnxSeq2SeqConfigWithPast, self).flatten_past_key_values( - flattened_output, name, idx, t - ) - - -@register_tasks_manager_onnx( - "bart", *[*COMMON_TEXT2TEXT_GENERATION_TASKS, "text-classification", "question-answering"] -) -class BartOnnxConfig(M2M100OnnxConfig): - pass - - -@register_tasks_manager_onnx( - "mbart", *[*COMMON_TEXT2TEXT_GENERATION_TASKS, "text-classification", "question-answering"] -) -class MBartOnnxConfig(BartOnnxConfig): - pass - - -@register_tasks_manager_onnx("blenderbot", *COMMON_TEXT2TEXT_GENERATION_TASKS) -class BlenderbotOnnxConfig(BartOnnxConfig): - pass - - -@register_tasks_manager_onnx("blenderbot-small", *COMMON_TEXT2TEXT_GENERATION_TASKS) -class BlenderbotSmallOnnxConfig(BartOnnxConfig): - pass - - -@register_tasks_manager_onnx("big_bird", *COMMON_TEXT_TASKS) -class BigBirdOnnxConfig(DistilBertOnnxConfig): - pass - - -@register_tasks_manager_onnx( - "bigbird_pegasus", *[*COMMON_TEXT2TEXT_GENERATION_TASKS, "text-classification", "question-answering"] -) -class BigBirdPegasusOnnxConfig(BartOnnxConfig): - _MODEL_PATCHER = BigBirdPegasusModelPatcher - - -@register_tasks_manager_onnx("pegasus", *COMMON_TEXT2TEXT_GENERATION_TASKS) -class PegasusOnnxConfig(BartOnnxConfig): - pass - - -@register_tasks_manager_onnx("marian", *COMMON_TEXT2TEXT_GENERATION_TASKS) -class MarianOnnxConfig(BartOnnxConfig): - pass - - -@register_tasks_manager_onnx("vit", *["feature-extraction", "image-classification", "masked-im"]) -class ViTOnnxConfig(VisionOnnxConfig): - NORMALIZED_CONFIG_CLASS = NormalizedVisionConfig - _MODEL_PATCHER = ViTForImageClassificationPatcher - - @property - def inputs(self) -> dict[str, dict[int, str]]: - return {"pixel_values": {0: "batch_size", 2: "height", 3: "width"}} - - @property - def outputs(self) -> dict[str, dict[int, str]]: - common_outputs = super().outputs - - if self.task == "feature-extraction": - common_outputs["last_hidden_state"] = {0: "batch_size"} - - return common_outputs - - -@register_tasks_manager_onnx("vitpose", *["keypoint-detection"]) -class VitPoseOnnxConfig(ViTOnnxConfig): - DUMMY_INPUT_GENERATOR_CLASSES = (VitPoseDummyInputGenerator,) - - _MODEL_PATCHER = VitPoseModelPatcher - - @property - def inputs(self) -> dict[str, dict[int, str]]: - return {"pixel_values": {0: "batch_size"}} - - -@register_tasks_manager_onnx("cvt", *["feature-extraction", "image-classification"]) -class CvTOnnxConfig(ViTOnnxConfig): - pass - - -@register_tasks_manager_onnx("levit", *["feature-extraction", "image-classification"]) -class LevitOnnxConfig(ViTOnnxConfig): - pass - - -@register_tasks_manager_onnx("deit", *["feature-extraction", "image-classification", "masked-im"]) -class DeiTOnnxConfig(ViTOnnxConfig): - pass - - -@register_tasks_manager_onnx("beit", *["feature-extraction", "image-classification"]) -class BeitOnnxConfig(ViTOnnxConfig): - pass - - -@register_tasks_manager_onnx("convnext", *["feature-extraction", "image-classification"]) -class ConvNextOnnxConfig(ViTOnnxConfig): - pass - - -@register_tasks_manager_onnx("convnextv2", *["feature-extraction", "image-classification"]) -class ConvNextV2OnnxConfig(ViTOnnxConfig): - pass - - -@register_tasks_manager_onnx("hiera", *["feature-extraction", "image-classification"]) -class HieraOnnxConfig(ViTOnnxConfig): - pass - - -@register_tasks_manager_onnx("pvt", *["feature-extraction", "image-classification"]) -class PvtOnnxConfig(ViTOnnxConfig): - pass - - -@register_tasks_manager_onnx("vit_mae", *["feature-extraction"]) -class VitMAEOnnxConfig(ViTOnnxConfig): - pass - - -@register_tasks_manager_onnx("vit_msn", *["feature-extraction", "image-classification"]) -class VitMSNOnnxConfig(ViTOnnxConfig): - pass - - -@register_tasks_manager_onnx("dinov2", *["feature-extraction", "image-classification"]) -class Dinov2OnnxConfig(ViTOnnxConfig): - DUMMY_INPUT_GENERATOR_CLASSES = (Dinov2DummyInputGenerator,) - - -@register_tasks_manager_onnx("mobilevit", *["feature-extraction", "image-classification", "image-segmentation"]) -class MobileViTOnnxConfig(ViTOnnxConfig): - pass - - -@register_tasks_manager_onnx("regnet", *["feature-extraction", "image-classification"]) -class RegNetOnnxConfig(ViTOnnxConfig): - pass - - -@register_tasks_manager_onnx("resnet", *["feature-extraction", "image-classification"]) -class ResNetOnnxConfig(ViTOnnxConfig): - pass - - -@register_tasks_manager_onnx("detr", *["feature-extraction", "object-detection", "image-segmentation"]) -class DetrOnnxConfig(ViTOnnxConfig): - @property - def outputs(self) -> dict[str, dict[int, str]]: - if self.task == "image-segmentation": - return { - "logits": {0: "batch_size", 1: "num_queries"}, - "pred_masks": {0: "batch_size", 1: "num_queries"}, - } - else: - return super().outputs - - -@register_tasks_manager_onnx("table-transformer", *["feature-extraction", "object-detection"]) -class TableTransformerOnnxConfig(DetrOnnxConfig): - pass - - -@register_tasks_manager_onnx("yolos", *["feature-extraction", "object-detection"]) -class YolosOnnxConfig(ViTOnnxConfig): - pass - - -@register_tasks_manager_onnx("swin", *["feature-extraction", "image-classification", "masked-im"]) -class SwinOnnxConfig(ViTOnnxConfig): - pass - - -@register_tasks_manager_onnx("swinv2", *["feature-extraction", "image-classification", "masked-im"]) -class SwinV2OnnxConfig(SwinOnnxConfig): - pass - - -@register_tasks_manager_onnx("swin2sr", *["feature-extraction", "image-to-image"]) -class Swin2srOnnxConfig(SwinOnnxConfig): - @property - def outputs(self) -> dict[str, dict[int, str]]: - outputs = super().outputs - - if self.task == "image-to-image": - scale_factor = self._config.upscale - outputs["reconstruction"] = { - 0: "batch_size", - 1: "num_channels", - 2: f"height * {scale_factor}", - 3: f"width * {scale_factor}", - } - - return outputs - - -@register_tasks_manager_onnx( - "dpt", *["feature-extraction", "depth-estimation", "image-segmentation", "semantic-segmentation"] -) -class DptOnnxConfig(ViTOnnxConfig): - pass - - -@register_tasks_manager_onnx("glpn", *["feature-extraction", "depth-estimation"]) -class GlpnOnnxConfig(ViTOnnxConfig): - pass - - -@register_tasks_manager_onnx("poolformer", *["feature-extraction", "image-classification"]) -class PoolFormerOnnxConfig(ViTOnnxConfig): - NORMALIZED_CONFIG_CLASS = NormalizedVisionConfig - - -@register_tasks_manager_onnx( - "segformer", *["feature-extraction", "image-classification", "image-segmentation", "semantic-segmentation"] -) -class SegformerOnnxConfig(YolosOnnxConfig): - @property - def outputs(self) -> dict[str, dict[int, str]]: - outputs = super().outputs - - if self.task == "image-segmentation": - outputs["logits"] = {0: "batch_size"} - - return outputs - - -@register_tasks_manager_onnx("mobilenet_v1", *["feature-extraction", "image-classification"]) -class MobileNetV1OnnxConfig(ViTOnnxConfig): - @property - def inputs(self) -> dict[str, dict[int, str]]: - return {"pixel_values": {0: "batch_size"}} - - -@register_tasks_manager_onnx("mobilenet_v2", *["feature-extraction", "image-classification"]) -class MobileNetV2OnnxConfig(MobileNetV1OnnxConfig): - pass - - -@register_tasks_manager_onnx("maskformer", *["feature-extraction", "image-segmentation"]) -class MaskFormerOnnxConfig(ViTOnnxConfig): - @property - def outputs(self) -> dict[str, dict[int, str]]: - if self.task == "image-segmentation": - return { - "class_queries_logits": {0: "batch_size", 1: "num_queries"}, - "masks_queries_logits": {0: "batch_size", 1: "num_queries", 2: "height", 3: "width"}, - } - else: - return super().outputs - - @property - def torch_to_onnx_output_map(self) -> dict[str, str]: - return { - "transformer_decoder_last_hidden_state": "last_hidden_state", - } - - -@register_tasks_manager_onnx("donut-swin", *["feature-extraction"]) -class DonutSwinOnnxConfig(ViTOnnxConfig): - pass - - -@register_tasks_manager_onnx("default-timm-config", *["image-classification"], library_name="timm") -class TimmDefaultOnnxConfig(ViTOnnxConfig): - def rename_ambiguous_inputs(self, inputs): - # The input name in the model signature is `x, hence the export input name is updated. - model_inputs = {} - model_inputs["x"] = inputs["pixel_values"] - - return model_inputs - - @property - def torch_to_onnx_input_map(self) -> dict[str, str]: - return {"x": "pixel_values"} - - -@register_tasks_manager_onnx("mgp-str", *["feature-extraction"]) -class MgpstrOnnxConfig(ViTOnnxConfig): - _MODEL_PATCHER = MgpstrModelPatcher - - @property - def outputs(self) -> dict[str, dict[int, str]]: - return { - "char_logits": {0: "batch_size"}, - "bpe_logits": {0: "batch_size"}, - "wp_logits": {0: "batch_size"}, - } - - -@register_tasks_manager_onnx("efficientnet", *["feature-extraction", "image-classification"]) -class EfficientNetOnnxConfig(ViTOnnxConfig): - @property - def outputs(self) -> dict[str, dict[int, str]]: - common_outputs = super().outputs - - if self.task == "image-classification": - common_outputs["logits"] = {0: "batch_size", 1: "num_classes"} - - return common_outputs - - -@register_tasks_manager_onnx( - "transformer", *["feature-extraction", "sentence-similarity"], library_name="sentence_transformers" -) -class SentenceTransformersTransformerOnnxConfig(TextEncoderOnnxConfig): - NORMALIZED_CONFIG_CLASS = NormalizedTextConfig - _MODEL_PATCHER = SentenceTransformersTransformerPatcher - - @property - def inputs(self) -> dict[str, dict[int, str]]: - return { - "input_ids": {0: "batch_size", 1: "sequence_length"}, - "attention_mask": {0: "batch_size", 1: "sequence_length"}, - } - - @property - def outputs(self) -> dict[str, dict[int, str]]: - return { - "token_embeddings": {0: "batch_size", 1: "sequence_length"}, - "sentence_embedding": {0: "batch_size"}, - } - - -class CLIPNormalizedConfig(NormalizedTextAndVisionConfig): - TEXT_CONFIG = "text_config" - VISION_CONFIG = "vision_config" - - -@register_tasks_manager_onnx("clip_vision_model", *["feature-extraction"]) -class CLIPVisionModelOnnxConfig(VisionOnnxConfig): - NORMALIZED_CONFIG_CLASS = NormalizedVisionConfig - _MODEL_PATCHER = CLIPModelPatcher - - @property - def inputs(self) -> dict[str, dict[int, str]]: - return {"pixel_values": {0: "batch_size", 1: "num_channels", 2: "height", 3: "width"}} - - @property - def outputs(self) -> dict[str, dict[int, str]]: - common_outputs = super().outputs - common_outputs["last_hidden_state"] = {0: "batch_size"} - common_outputs["pooler_output"] = {0: "batch_size"} - - return common_outputs - - -@register_tasks_manager_onnx("clip", *["feature-extraction", "zero-shot-image-classification", "image-classification"]) -class CLIPOnnxConfig(TextAndVisionOnnxConfig): - NORMALIZED_CONFIG_CLASS = CLIPNormalizedConfig - _MODEL_PATCHER = CLIPModelPatcher - - @property - def inputs(self) -> dict[str, dict[int, str]]: - inputs = {"pixel_values": {0: "batch_size", 1: "num_channels", 2: "height", 3: "width"}} - - if self.task in ["feature-extraction", "zero-shot-image-classification"]: - inputs.update( - { - "input_ids": {0: "text_batch_size", 1: "sequence_length"}, - "attention_mask": {0: "text_batch_size", 1: "sequence_length"}, - } - ) - - return inputs - - @property - def outputs(self) -> dict[str, dict[int, str]]: - if self.task in ["feature-extraction", "zero-shot-image-classification"]: - return { - "logits_per_image": {0: "image_batch_size", 1: "text_batch_size"}, - "logits_per_text": {0: "text_batch_size", 1: "image_batch_size"}, - "text_embeds": {0: "text_batch_size"}, - "image_embeds": {0: "image_batch_size"}, - } - else: - return super().outputs - - -@register_tasks_manager_onnx( - "clip", *["feature-extraction", "sentence-similarity"], library_name="sentence_transformers" -) -class SentenceTransformersCLIPOnnxConfig(CLIPOnnxConfig): - _MODEL_PATCHER = SentenceTransformersCLIPPatcher - - @property - def outputs(self) -> dict[str, dict[int, str]]: - return { - "text_embeds": {0: "text_batch_size"}, - "image_embeds": {0: "image_batch_size"}, - } - - -@register_tasks_manager_onnx("clip-text-with-projection", *["feature-extraction"], library_name="diffusers") -class CLIPTextWithProjectionOnnxConfig(TextEncoderOnnxConfig): - NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args( - vocab_size="vocab_size", - sequence_length="max_position_embeddings", - num_layers="num_hidden_layers", - allow_new=True, - ) - _MODEL_PATCHER = CLIPModelPatcher - - @property - def inputs(self) -> dict[str, dict[int, str]]: - return { - "input_ids": {0: "batch_size", 1: "sequence_length"}, - } - - @property - def outputs(self) -> dict[str, dict[int, str]]: - common_outputs = { - "text_embeds": {0: "batch_size", 1: "sequence_length"}, - "last_hidden_state": {0: "batch_size", 1: "sequence_length"}, - } - if self._normalized_config.output_hidden_states: - for i in range(self._normalized_config.num_layers + 1): - common_outputs[f"hidden_states.{i}"] = {0: "batch_size", 1: "sequence_length"} - - return common_outputs - - -@register_tasks_manager_onnx("clip-text", *["feature-extraction"], library_name="diffusers") -class CLIPTextOnnxConfig(CLIPTextWithProjectionOnnxConfig): - _MODEL_PATCHER = CLIPModelPatcher - - @property - def outputs(self) -> dict[str, dict[int, str]]: - common_outputs = { - "last_hidden_state": {0: "batch_size", 1: "sequence_length"}, - "pooler_output": {0: "batch_size"}, - } - - if self._normalized_config.output_hidden_states: - for i in range(self._normalized_config.num_layers + 1): - common_outputs[f"hidden_states.{i}"] = {0: "batch_size", 1: "sequence_length"} - - return common_outputs - - -@register_tasks_manager_onnx( - "metaclip_2", - *["feature-extraction", "zero-shot-image-classification", "image-classification"], - library_name="transformers", -) -class MetaCLIP2OnnxConfig(TextAndVisionOnnxConfig): - NORMALIZED_CONFIG_CLASS = CLIPNormalizedConfig - MIN_TRANSFORMERS_VERSION = version.parse("4.56.2") - VARIANTS = { # noqa: RUF012 - "monolith": "All the MetaClip2 model components are exported as a single model.onnx.", - "split": "The vision model is exported as a separate vision_model.onnx, and the text_model is exported as text_model.onnx", - } - DEFAULT_VARIANT = "monolith" - _MODEL_PATCHER = MetaCLIP2Patcher - - def __init__( - self, - config: PretrainedConfig, - task: str = "feature-extraction", - int_dtype: str = "int64", - float_dtype: str = "fp32", - variant: str = "monolith", - vision_model: bool | None = None, - preprocessors: list[Any] | None = None, - ): - super().__init__( - config=config, - task=task, - int_dtype=int_dtype, - float_dtype=float_dtype, - preprocessors=preprocessors, - ) - self.variant = variant - self.vision_model = vision_model - - @property - def inputs(self) -> dict[str, dict[int, str]]: - if self.variant == "monolith": - inputs = {"pixel_values": {0: "batch_size", 1: "num_channels", 2: "height", 3: "width"}} - if self.task in ["feature-extraction", "zero-shot-image-classification"]: - inputs.update( - { - "input_ids": {0: "text_batch_size", 1: "sequence_length"}, - "attention_mask": {0: "text_batch_size", 1: "sequence_length"}, - } - ) - else: - if self.vision_model: - inputs = {"pixel_values": {0: "batch_size", 1: "num_channels", 2: "height", 3: "width"}} - else: - inputs = { - "input_ids": {0: "text_batch_size", 1: "sequence_length"}, - "attention_mask": {0: "text_batch_size", 1: "sequence_length"}, - } - return inputs - - @property - def outputs(self) -> dict[str, dict[int, str]]: - if self.variant == "split": - if self.vision_model: - return { - "image_embeds": {0: "batch_size"}, - } - else: - return { - "text_embeds": {0: "batch_size"}, - } - else: - if self.task in ["feature-extraction", "zero-shot-image-classification"]: - return { - "logits_per_image": {0: "image_batch_size", 1: "text_batch_size"}, - "logits_per_text": {0: "text_batch_size", 1: "image_batch_size"}, - "text_embeds": {0: "text_batch_size"}, - "image_embeds": {0: "image_batch_size"}, - } - else: - return super().outputs - - -class SiglipNormalizedConfig(CLIPNormalizedConfig): - pass - - -@register_tasks_manager_onnx("chinese_clip", *["feature-extraction", "zero-shot-image-classification"]) -class ChineseCLIPOnnxConfig(CLIPOnnxConfig): - pass - - -@register_tasks_manager_onnx("siglip", *["feature-extraction", "zero-shot-image-classification"]) -class SiglipOnnxConfig(CLIPOnnxConfig): - NORMALIZED_CONFIG_CLASS = SiglipNormalizedConfig - - @property - def inputs(self) -> dict[str, dict[int, str]]: - return { - "input_ids": {0: "text_batch_size", 1: "sequence_length"}, - "pixel_values": {0: "image_batch_size", 1: "num_channels", 2: "height", 3: "width"}, - # NOTE: No attention_mask - } - - -@register_tasks_manager_onnx("siglip-text-with-projection", *["feature-extraction"]) -class SiglipTextWithProjectionOnnxConfig(CLIPTextWithProjectionOnnxConfig): - pass - - -@register_tasks_manager_onnx("siglip-text", *["feature-extraction"]) -class SiglipTextOnnxConfig(CLIPTextOnnxConfig): - pass - - -@register_tasks_manager_onnx("siglip_vision_model", *["feature-extraction"]) -class SiglipVisionModelOnnxConfig(CLIPVisionModelOnnxConfig): - pass - - -@register_tasks_manager_onnx("unet-2d-condition", *["semantic-segmentation"], library_name="diffusers") -class UNetOnnxConfig(VisionOnnxConfig): - NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args( - image_size="sample_size", - num_channels="in_channels", - hidden_size="cross_attention_dim", - vocab_size="norm_num_groups", - allow_new=True, - ) - - DUMMY_INPUT_GENERATOR_CLASSES = ( - DummyVisionInputGenerator, - DummyTimestepInputGenerator, - DummySeq2SeqDecoderTextInputGenerator, - ) - - @property - def inputs(self) -> dict[str, dict[int, str]]: - common_inputs = { - "sample": {0: "batch_size", 2: "height", 3: "width"}, - "timestep": {}, # a scalar with no dimension - "encoder_hidden_states": {0: "batch_size", 1: "sequence_length"}, - } - - # TODO : add addition_embed_type == text_image, image and image_embeds - # https://github.com/huggingface/diffusers/blob/9366c8f84bfe47099ff047272661786ebb54721d/src/diffusers/models/unets/unet_2d_condition.py#L671 - if getattr(self._normalized_config, "addition_embed_type", None) == "text_time": - common_inputs["text_embeds"] = {0: "batch_size"} - common_inputs["time_ids"] = {0: "batch_size"} - - if getattr(self._normalized_config, "time_cond_proj_dim", None) is not None: - common_inputs["timestep_cond"] = {0: "batch_size"} - - return common_inputs - - @property - def outputs(self) -> dict[str, dict[int, str]]: - return { - "out_sample": {0: "batch_size", 2: "height", 3: "width"}, - } - - @property - def torch_to_onnx_output_map(self) -> dict[str, str]: - return { - "sample": "out_sample", - } - - def generate_dummy_inputs(self, framework: str = "pt", **kwargs): - dummy_inputs = super().generate_dummy_inputs(framework=framework, **kwargs) - dummy_inputs["encoder_hidden_states"] = dummy_inputs["encoder_hidden_states"][0] - - if getattr(self._normalized_config, "addition_embed_type", None) == "text_time": - dummy_inputs["added_cond_kwargs"] = { - "text_embeds": dummy_inputs.pop("text_embeds"), - "time_ids": dummy_inputs.pop("time_ids"), - } - - return dummy_inputs - - def ordered_inputs(self, model) -> dict[str, dict[int, str]]: - inputs = super().ordered_inputs(model=model) - # to fix mismatch between model forward signature and expected inputs - # a dictionary of additional embeddings `added_cond_kwargs` is expected depending on config.addition_embed_type - if getattr(self._normalized_config, "addition_embed_type", None) == "text_time": - inputs["text_embeds"] = self.inputs["text_embeds"] - inputs["time_ids"] = self.inputs["time_ids"] - - return inputs - - -@register_tasks_manager_onnx("vae-encoder", *["semantic-segmentation"], library_name="diffusers") -class VaeEncoderOnnxConfig(VisionOnnxConfig): - ATOL_FOR_VALIDATION = 3e-4 # TODO: this only happens in test_export.py - NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args( - num_channels="in_channels", image_size="sample_size", allow_new=True - ) - - @property - def inputs(self) -> dict[str, dict[int, str]]: - return { - "sample": {0: "batch_size", 2: "sample_height", 3: "sample_width"}, - } - - @property - def outputs(self) -> dict[str, dict[int, str]]: - down_sampling_factor = 2 ** (len(self._normalized_config.down_block_types) - 1) - - return { - "latent_parameters": { - 0: "batch_size", - 2: f"sample_height / {down_sampling_factor}", - 3: f"sample_width / {down_sampling_factor}", - }, - } - - -@register_tasks_manager_onnx("vae-decoder", *["semantic-segmentation"], library_name="diffusers") -class VaeDecoderOnnxConfig(VisionOnnxConfig): - ATOL_FOR_VALIDATION = 3e-4 # TODO: this only happens in test_export.py - NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args(num_channels="latent_channels", allow_new=True) - - @property - def inputs(self) -> dict[str, dict[int, str]]: - return { - "latent_sample": {0: "batch_size", 2: "latent_height", 3: "latent_width"}, - } - - @property - def outputs(self) -> dict[str, dict[int, str]]: - up_sampling_factor = 2 ** (len(self._normalized_config.up_block_types) - 1) - - return { - "sample": { - 0: "batch_size", - 2: f"latent_height * {up_sampling_factor}", - 3: f"latent_width * {up_sampling_factor}", - }, - } - - -@register_tasks_manager_onnx("t5-encoder", *["feature-extraction"], library_name="diffusers") -class T5EncoderOnnxConfig(TextEncoderOnnxConfig): - NORMALIZED_CONFIG_CLASS = NormalizedTextConfig - - @property - def inputs(self): - return { - "input_ids": {0: "batch_size", 1: "sequence_length"}, - } - - @property - def outputs(self): - return { - "last_hidden_state": {0: "batch_size", 1: "sequence_length"}, - } - - -@register_tasks_manager_onnx("sd3-transformer-2d", *["semantic-segmentation"], library_name="diffusers") -class SD3TransformerOnnxConfig(VisionOnnxConfig): - DUMMY_INPUT_GENERATOR_CLASSES = ( - DummyTransformerTimestepInputGenerator, - DummyTransformerVisionInputGenerator, - DummyTransformerTextInputGenerator, - ) - - NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args( - image_size="sample_size", - num_channels="in_channels", - vocab_size="attention_head_dim", - hidden_size="joint_attention_dim", - projection_size="pooled_projection_dim", - allow_new=True, - ) - - @property - def inputs(self) -> dict[str, dict[int, str]]: - common_inputs = { - "hidden_states": {0: "batch_size", 2: "height", 3: "width"}, - "encoder_hidden_states": {0: "batch_size", 1: "sequence_length"}, - "pooled_projections": {0: "batch_size"}, - "timestep": {0: "step"}, - } - - return common_inputs - - @property - def outputs(self) -> dict[str, dict[int, str]]: - return { - "out_hidden_states": {0: "batch_size", 2: "height", 3: "width"}, - } - - @property - def torch_to_onnx_output_map(self) -> dict[str, str]: - return { - "sample": "out_hidden_states", - } - - -@register_tasks_manager_onnx("flux-transformer-2d", *["semantic-segmentation"], library_name="diffusers") -class FluxTransformerOnnxConfig(SD3TransformerOnnxConfig): - DUMMY_INPUT_GENERATOR_CLASSES = ( - DummyTransformerTimestepInputGenerator, - DummyFluxTransformerVisionInputGenerator, - DummyFluxTransformerTextInputGenerator, - ) - _MODEL_PATCHER = FluxTransformerModelPatcher - - @property - def inputs(self): - common_inputs = super().inputs - common_inputs["hidden_states"] = {0: "batch_size", 1: "packed_height_width"} - common_inputs["txt_ids"] = ( - {0: "sequence_length"} if is_diffusers_version(">=", "0.31.0") else {0: "batch_size", 1: "sequence_length"} - ) - common_inputs["img_ids"] = ( - {0: "packed_height_width"} - if is_diffusers_version(">=", "0.31.0") - else {0: "batch_size", 1: "packed_height_width"} - ) - - if getattr(self._normalized_config, "guidance_embeds", False): - common_inputs["guidance"] = {0: "batch_size"} - - return common_inputs - - @property - def outputs(self): - return { - "out_hidden_states": {0: "batch_size", 1: "packed_height_width"}, - } - - -@register_tasks_manager_onnx("groupvit", *["feature-extraction"]) -class GroupViTOnnxConfig(CLIPOnnxConfig): - pass - - -@register_tasks_manager_onnx("owlvit", *["feature-extraction", "zero-shot-object-detection"]) -class OwlViTOnnxConfig(CLIPOnnxConfig): - def __init__( - self, - config: PretrainedConfig, - task: str = "feature-extraction", - int_dtype: str = "int64", - float_dtype: str = "fp32", - preprocessors: list[Any] | None = None, - ): - super().__init__( - config=config, - task=task, - int_dtype=int_dtype, - float_dtype=float_dtype, - preprocessors=preprocessors, - ) - if task == "zero-shot-object-detection": - logger.warning( - "The batch size of this model will not be dynamic because non-maximum suppression is performed. " - "Make sure to export the model with the same batch size as the one you will use at inference " - "with `--batch_size N`." - ) - - @property - def inputs(self) -> dict[str, dict[int, str]]: - inputs = {"pixel_values": {0: "batch_size", 1: "num_channels", 2: "height", 3: "width"}} - - if self.task in ["feature-extraction", "zero-shot-object-detection"]: - inputs.update( - { - "input_ids": {0: "text_batch_size", 1: "sequence_length"}, - "attention_mask": {0: "text_batch_size", 1: "sequence_length"}, - } - ) - - return inputs - - @property - def outputs(self) -> dict[str, dict[int, str]]: - outputs = {} - if self.task == "feature-extraction": - outputs["logits_per_image"] = {0: "image_batch_size", 1: "text_batch_size"} - outputs["logits_per_text"] = {0: "text_batch_size", 1: "image_batch_size"} - elif self.task == "zero-shot-object-detection": - outputs["logits"] = {0: "image_batch_size", 2: "num_queries"} - outputs["pred_boxes"] = {0: "image_batch_size", 1: "num_boxes"} - - outputs["text_embeds"] = {0: "text_batch_size", 1: "max_text_queries"} - outputs["image_embeds"] = {0: "image_batch_size"} - return outputs - - -@register_tasks_manager_onnx("owlv2", *["feature-extraction", "zero-shot-object-detection"]) -class OwlV2OnnxConfig(OwlViTOnnxConfig): - MIN_TRANSFORMERS_VERSION = version.parse("4.35.0") - - -@register_tasks_manager_onnx( - "layoutlm", *["feature-extraction", "fill-mask", "text-classification", "token-classification"] -) -class LayoutLMOnnxConfig(TextAndVisionOnnxConfig): - NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args( - allow_new=True, MAX_2D_POSITION_EMBEDDINGS="max_2d_position_embeddings" - ) - - @property - def inputs(self) -> dict[str, dict[int, str]]: - return { - "input_ids": {0: "batch_size", 1: "sequence_length"}, - "bbox": {0: "batch_size", 1: "sequence_length"}, - "attention_mask": {0: "batch_size", 1: "sequence_length"}, - "token_type_ids": {0: "batch_size", 1: "sequence_length"}, - } - - -@register_tasks_manager_onnx( - "layoutlmv3", *["feature-extraction", "question-answering", "text-classification", "token-classification"] -) -class LayoutLMv3OnnxConfig(TextAndVisionOnnxConfig): - NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args( - allow_new=True, MAX_2D_POSITION_EMBEDDINGS="max_2d_position_embeddings", image_size="input_size" - ) - - @property - def inputs(self) -> dict[str, dict[int, str]]: - if self.task in ["text-classification", "question-answering"]: - pixel_values_dynamic_axes = {0: "batch_size", 1: "num_channels", 2: "height", 3: "width"} - else: - pixel_values_dynamic_axes = {0: "batch_size", 1: "num_channels"} - return { - "input_ids": {0: "batch_size", 1: "sequence_length"}, - "attention_mask": {0: "batch_size", 1: "sequence_length"}, - "bbox": {0: "batch_size", 1: "sequence_length"}, - "pixel_values": pixel_values_dynamic_axes, - } - - -@register_tasks_manager_onnx( - "lilt", *["feature-extraction", "question-answering", "text-classification", "token-classification"] -) -class LiltOnnxConfig(TextAndVisionOnnxConfig): - NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args( - allow_new=True, - MAX_2D_POSITION_EMBEDDINGS="max_2d_position_embeddings", - ) - - @property - def inputs(self) -> dict[str, dict[int, str]]: - return { - "input_ids": {0: "batch_size", 1: "sequence_length"}, - "bbox": {0: "batch_size", 1: "sequence_length"}, - "attention_mask": {0: "batch_size", 1: "sequence_length"}, - } - - -@register_tasks_manager_onnx("data2vec-text", *COMMON_TEXT_TASKS) -class Data2VecTextOnnxConfig(DistilBertOnnxConfig): - pass - - -@register_tasks_manager_onnx("data2vec-vision", *["feature-extraction", "image-classification"]) -class Data2VecVisionOnnxConfig(ViTOnnxConfig): - pass - - -@register_tasks_manager_onnx("perceiver", *["fill-mask", "text-classification", "image-classification"]) -class PerceiverOnnxConfig(TextAndVisionOnnxConfig): - NORMALIZED_CONFIG_CLASS = NormalizedTextConfig - DUMMY_INPUT_GENERATOR_CLASSES = ( - PerceiverDummyInputGenerator, - *TextAndVisionOnnxConfig.DUMMY_INPUT_GENERATOR_CLASSES, - ) - - def __init__( - self, - config: PretrainedConfig, - task: str = "feature-extraction", - int_dtype: str = "int64", - float_dtype: str = "fp32", - preprocessors: list[Any] | None = None, - ): - super().__init__( - config=config, - task=task, - int_dtype=int_dtype, - float_dtype=float_dtype, - preprocessors=preprocessors, - ) - self.is_generating_dummy_inputs = False - - @property - def inputs_name(self): - if self.is_generating_dummy_inputs: - if self.task in ["fill-mask", "text-classification"]: - return "input_ids" - else: - return "pixel_values" - else: - return "inputs" - - @property - def inputs(self) -> dict[str, dict[int, str]]: - if self.inputs_name in ["input_ids", "inputs"]: - dynamic_axis = {0: "batch_size", 1: "sequence_length"} - return { - "input_ids": dynamic_axis, - "attention_mask": dynamic_axis, - } - else: - dynamic_axis = {0: "batch_size", 1: "sequence_length", 2: "width", 3: "height"} - return { - "pixel_values": dynamic_axis, - } - - @property - def outputs(self) -> dict[str, dict[int, str]]: - outputs = super().outputs - - if "logits" in outputs: - # default is {0: "batch_size", 1: "sequence_length"} where sequence_length is dynamic axis - # but perceiver always return the same max sequence length in the second dimension - outputs["logits"] = {0: "batch_size"} - - return outputs - - def generate_dummy_inputs(self, framework: str = "pt", **kwargs): - self.is_generating_dummy_inputs = True - dummy_inputs = super().generate_dummy_inputs(framework=framework, **kwargs) - dummy_inputs[self.inputs_name] = dummy_inputs.pop(self.inputs_name) - return dummy_inputs - - -@register_tasks_manager_onnx("hubert", *["feature-extraction", "automatic-speech-recognition", "audio-classification"]) -class HubertOnnxConfig(AudioOnnxConfig): - NORMALIZED_CONFIG_CLASS = NormalizedConfig - - @property - def outputs(self) -> dict[str, dict[int, str]]: - outputs = super().outputs - - # Hubert output formula adapted from: - # https://github.com/huggingface/transformers/blob/v4.55.2/src/transformers/models/hubert/modeling_hubert.py#L721 - if self.task == "automatic-speech-recognition": - sequence_length = "sequence_length" - for kernel_size, stride in zip(self._config.conv_kernel, self._config.conv_stride): - sequence_length = f"( {sequence_length} - {kernel_size} ) // {stride} + 1" - outputs["logits"] = {0: "batch_size", 1: sequence_length} - - return outputs - - -@register_tasks_manager_onnx( - "data2vec-audio", - *[ - "feature-extraction", - "automatic-speech-recognition", - "audio-classification", - "audio-frame-classification", - "audio-xvector", - ], -) -class Data2VecAudioOnnxConfig(HubertOnnxConfig): - pass - - -@register_tasks_manager_onnx( - "wav2vec2", - *[ - "feature-extraction", - "automatic-speech-recognition", - "audio-classification", - "audio-frame-classification", - "audio-xvector", - ], -) -class Wav2Vec2OnnxConfig(HubertOnnxConfig): - pass - - -@register_tasks_manager_onnx( - "wav2vec2-conformer", - *[ - "feature-extraction", - "automatic-speech-recognition", - "audio-classification", - "audio-frame-classification", - "audio-xvector", - ], -) -class Wav2Vec2ConformerOnnxConfig(HubertOnnxConfig): - pass - - -@register_tasks_manager_onnx("sew", *["feature-extraction", "automatic-speech-recognition", "audio-classification"]) -class SEWOnnxConfig(HubertOnnxConfig): - pass - - -@register_tasks_manager_onnx("sew-d", *["feature-extraction", "automatic-speech-recognition", "audio-classification"]) -class SEWDOnnxConfig(HubertOnnxConfig): - pass - - -@register_tasks_manager_onnx( - "unispeech", *["feature-extraction", "automatic-speech-recognition", "audio-classification"] -) -class UniSpeechOnnxConfig(HubertOnnxConfig): - pass - - -@register_tasks_manager_onnx( - "unispeech-sat", - *[ - "feature-extraction", - "automatic-speech-recognition", - "audio-classification", - "audio-frame-classification", - "audio-xvector", - ], -) -class UniSpeechSATOnnxConfig(HubertOnnxConfig): - pass - - -@register_tasks_manager_onnx( - "wavlm", - *[ - "feature-extraction", - "automatic-speech-recognition", - "audio-classification", - "audio-frame-classification", - "audio-xvector", - ], -) -class WavLMOnnxConfig(HubertOnnxConfig): - pass - - -@register_tasks_manager_onnx("mctct", *["feature-extraction", "automatic-speech-recognition"]) -class MCTCTOnnxConfig(AudioOnnxConfig): - NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args( - input_features_per_channel="input_feat_per_channel", allow_new=True - ) - DUMMY_INPUT_GENERATOR_CLASSES = (MCTCTDummyAudioInputGenerator,) - - @property - def inputs(self) -> dict[str, dict[int, str]]: - return {"input_features": {0: "batch_size", 1: "sequence_length"}} - - @property - def outputs(self) -> dict[str, dict[int, str]]: - outputs = super().outputs - - # mctct output formula adapted from: - # https://github.com/huggingface/transformers/blob/v4.53.3/src/transformers/models/deprecated/mctct/modeling_mctct.py#L455 - if self.task == "automatic-speech-recognition": - sequence_length = "sequence_length" - for kernel_size, stride in zip(self._config.conv_kernel, self._config.conv_stride): - dilation = 1 - padding = kernel_size // 2 - sequence_length = f"( {sequence_length} + 2 * {padding} - {dilation} * ({kernel_size} - 1) - 1 )" - sequence_length = f"( {sequence_length} // {stride} ) + 1" - outputs["logits"] = {0: "batch_size", 1: sequence_length} - - return outputs - - -@register_tasks_manager_onnx("audio-spectrogram-transformer", *["feature-extraction", "audio-classification"]) -class ASTOnnxConfig(OnnxConfig): - NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args( - num_mel_bins="num_mel_bins", max_length="max_length", allow_new=True - ) - DUMMY_INPUT_GENERATOR_CLASSES = (ASTDummyAudioInputGenerator,) - - @property - def inputs(self) -> dict[str, dict[int, str]]: - return {"input_values": {0: "batch_size"}} - - -@register_tasks_manager_onnx( - "moonshine", - *[ - "feature-extraction", - "feature-extraction-with-past", - "automatic-speech-recognition", - "automatic-speech-recognition-with-past", - ], -) -class MoonshineOnnxConfig(AudioToTextOnnxConfig): - MIN_TRANSFORMERS_VERSION = version.parse("4.48.0") - NORMALIZED_CONFIG_CLASS = NormalizedSeq2SeqConfig - _MODEL_PATCHER = MoonshineModelPatcher - DUMMY_INPUT_GENERATOR_CLASSES = ( - DummyMoonshineAudioInputGenerator, - DummySeq2SeqDecoderTextInputGenerator, - DummySeq2SeqPastKeyValuesGenerator, - ) - - @property - def inputs(self) -> dict[str, dict[int, str]]: - common_inputs = {} - if self._behavior in {ConfigBehavior.ENCODER, ConfigBehavior.MONOLITH}: - common_inputs["input_values"] = {0: "batch_size", 1: "encoder_sequence_length"} - else: - common_inputs["encoder_outputs"] = {0: "batch_size", 1: "encoder_sequence_length"} - common_inputs["attention_mask"] = {0: "batch_size", 1: "encoder_sequence_length"} - - if self._behavior in {ConfigBehavior.DECODER, ConfigBehavior.MONOLITH}: - common_inputs["decoder_input_ids"] = {0: "batch_size", 1: "decoder_sequence_length"} - if self.use_past_in_inputs: - self.add_past_key_values(common_inputs, direction="inputs") - - return common_inputs - - @property - def outputs(self) -> dict[str, dict[int, str]]: - common_outputs = super().outputs - if self._behavior in {ConfigBehavior.ENCODER, ConfigBehavior.MONOLITH}: - if self._behavior is ConfigBehavior.MONOLITH: - output_name = "encoder_last_hidden_state" - else: - output_name = "last_hidden_state" - # Moonshine encoder output formula adapted from: - # transformers.models.moonshine.modeling_moonshine.MoonshinePreTrainedModel._get_feat_extract_output_lengths - # output_conv1_length = int((input_lengths - 127) / 64 + 1) - # output_conv2_length = int((output_conv1_length - 7) / 3 + 1) - # output_conv3_length = int((output_conv2_length - 3) / 2 + 1) - output_sequence_length = "( ( ( encoder_sequence_length - 127 ) // 64 + 1 - 7 ) // 3 + 1 - 3 ) // 2 + 1" - common_outputs[output_name] = {0: "batch_size", 1: output_sequence_length} - return common_outputs - - -@register_tasks_manager_onnx( - "whisper", - *[ - "feature-extraction", - "feature-extraction-with-past", - "audio-classification", - "automatic-speech-recognition", - "automatic-speech-recognition-with-past", - ], -) -class WhisperOnnxConfig(AudioToTextOnnxConfig): - NORMALIZED_CONFIG_CLASS = NormalizedSeq2SeqConfig.with_args( - encoder_num_layers="encoder_layers", - decoder_num_layers="decoder_layers", - feature_size="num_mel_bins", - allow_new=True, - ) - - @property - def inputs(self) -> dict[str, dict[int, str]]: - if self.task == "audio-classification": - return {"input_features": {0: "batch_size"}} - - common_inputs = super().inputs - if self._behavior in {ConfigBehavior.ENCODER, ConfigBehavior.MONOLITH}: - common_inputs["input_features"] = {0: "batch_size"} # Remove unnecessary dynamic axis. - else: - # the dynamic encoder sequence length is only needed here because the input generator generates - # encoder_outputs with a seq_len=16 but the model expects at inference time seq_len=1500 - # TODO: this can be fixed by generating the correct inputs in the input generator - common_inputs["encoder_outputs"] = {0: "batch_size", 1: "encoder_sequence_length"} - - if self._behavior in {ConfigBehavior.DECODER, ConfigBehavior.MONOLITH}: - if is_transformers_version(">=", "4.43.0") and is_transformers_version("<", "4.46.0"): - # since https://github.com/huggingface/transformers/pull/31166 - if self._behavior is not ConfigBehavior.ENCODER and self.use_past_in_inputs: - common_inputs["cache_position"] = {0: "decoder_sequence_length"} - - return common_inputs - - @property - def outputs(self) -> dict[str, dict[int, str]]: - if self.task == "audio-classification": - return {"logits": {0: "batch_size"}} - - common_outputs = super().outputs - if self._behavior in {ConfigBehavior.ENCODER, ConfigBehavior.MONOLITH}: - if self._behavior is ConfigBehavior.MONOLITH: - output_name = "encoder_last_hidden_state" - else: - output_name = "last_hidden_state" - common_outputs[output_name] = {0: "batch_size"} # Remove unnecessary dynamic axis. - return common_outputs - - -@register_tasks_manager_onnx("musicgen", *["text-to-audio"]) -class MusicgenOnnxConfig(OnnxSeq2SeqConfigWithPast): - # NOTE: Several warnings during the export are not to worry about: - # * for i, indices in enumerate(codes): --> can be unrolled, fixed length (num_quantizers). - # * max_pad = max(padding_left, padding_right) --> does not impact later controlflows. - # if length <= max_pad: --> appears to be always False for Musicgen. - - VARIANTS = { # noqa: RUF012 - "text-conditional-with-past": """Exports Musicgen to ONNX to generate audio samples conditioned on a text prompt (Reference: https://huggingface.co/docs/transformers/model_doc/musicgen#text-conditional-generation). - This uses the decoder KV cache. The following subcomponents are exported: - * text_encoder.onnx: corresponds to the text encoder part in https://github.com/huggingface/transformers/blob/v4.39.1/src/transformers/models/musicgen/modeling_musicgen.py#L1457. - * encodec_decode.onnx: corresponds to the Encodec audio encoder part in https://github.com/huggingface/transformers/blob/v4.39.1/src/transformers/models/musicgen/modeling_musicgen.py#L2472-L2480. - * decoder_model.onnx: The Musicgen decoder, without past key values input, and computing cross attention. Not required at inference (use decoder_model_merged.onnx instead). - * decoder_with_past_model.onnx: The Musicgen decoder, with past_key_values input (KV cache filled), not computing cross attention. Not required at inference (use decoder_model_merged.onnx instead). - * decoder_model_merged.onnx: The two previous models fused in one, to avoid duplicating weights. A boolean input `use_cache_branch` allows to select the branch to use. In the first forward pass where the KV cache is empty, dummy past key values inputs need to be passed and are ignored with use_cache_branch=False. - * build_delay_pattern_mask.onnx: A model taking as input `input_ids`, `pad_token_id`, `max_length`, and building a delayed pattern mask to the input_ids. Implements https://github.com/huggingface/transformers/blob/v4.39.3/src/transformers/models/musicgen/modeling_musicgen.py#L1054.""", - } - # TODO: support audio-prompted generation (audio_encoder_encode.onnx: corresponds to the audio encoder part - # in https://github.com/huggingface/transformers/blob/f01e1609bf4dba146d1347c1368c8c49df8636f6/src/transformers/models/musicgen/modeling_musicgen.py#L2087.) - # With that, we have full Encodec support. - DEFAULT_VARIANT = "text-conditional-with-past" - - NORMALIZED_CONFIG_CLASS = NormalizedEncoderDecoderConfig - - DUMMY_INPUT_GENERATOR_CLASSES = ( - DummyTextInputGenerator, - DummyCodegenDecoderTextInputGenerator, - DummySeq2SeqPastKeyValuesGenerator, - DummyEncodecInputGenerator, - DummyIntGenerator, - ) - DUMMY_PKV_GENERATOR_CLASS = DummySeq2SeqPastKeyValuesGenerator - _MODEL_PATCHER = MusicgenModelPatcher - - def __init__( - self, - config: PretrainedConfig, - task: str = "feature-extraction", - int_dtype: str = "int64", - float_dtype: str = "fp32", - use_past: bool = False, - use_past_in_inputs: bool = False, - behavior: ConfigBehavior = ConfigBehavior.ENCODER, - preprocessors: list[Any] | None = None, - model_part: Literal["text_encoder", "encodec_decode", "decoder", "build_delay_pattern_mask"] | None = None, - variant: str = "text-conditional-with-past", - ): - super().__init__( - config=config, - task=task, - int_dtype=int_dtype, - float_dtype=float_dtype, - use_past=use_past, - use_past_in_inputs=use_past_in_inputs, - behavior=behavior, - preprocessors=preprocessors, - ) - - if ( - model_part in ["text_encoder", "encodec_decode", "build_delay_pattern_mask"] - and behavior != ConfigBehavior.ENCODER - ): - raise ValueError( - f"model_part is {model_part} and behavior is {behavior}. This is not supported, please open an issue at https://github.com/huggingface/optimum/issues." - ) - - if model_part == "decoder" and behavior != ConfigBehavior.DECODER: - raise ValueError( - f"model_part is {model_part} and behavior is {behavior}. This is not supported, please open an issue at https://github.com/huggingface/optimum/issues." - ) - - if behavior == ConfigBehavior.MONOLITH: - raise ValueError( - "Musicgen does not support behavior=ConfigBehavior.MONOLITH. Please open an issue at https://github.com/huggingface/optimum/issues." - ) - - if config.audio_encoder.model_type != "encodec": - raise ValueError( - f"Optimum ONNX export for Musicgen supports only Encodec as the audio encoder, got: {config.audio_encoder.model_type}. Please open an issue at https://github.com/huggingface/optimum/issues." - ) - - # Handling it would require to trace the audio_encoder.decode with torch.jit.script as we than have an unrollable loop. - if config.audio_encoder.chunk_length_s is not None: - raise ValueError( - f"Musicgen ONNX export currently does not support audio_encoder.chunk_length_s not None (got {config.audio_encoder.chunk_length_s}). Please open an issue at https://github.com/huggingface/optimum/issues." - ) - - self.model_part = model_part - if self.model_part == "decoder": - self.use_past = True # without past is not supported, hard-code it here. - - self._normalized_config.ENCODER_NORMALIZED_CONFIG_CLASS = NormalizedTextConfig(self._config.text_encoder) - self._normalized_config.DECODER_NORMALIZED_CONFIG_CLASS = NormalizedConfig(self._config.decoder) - self._normalized_config.decoder_num_layers = self._config.decoder.num_hidden_layers - self._normalized_config.DECODER_NORMALIZED_CONFIG_CLASS.num_layers = self._config.decoder.num_hidden_layers - self._normalized_config.DECODER_NORMALIZED_CONFIG_CLASS.encoder_num_attention_heads = ( - self._config.decoder.num_attention_heads - ) - self._normalized_config.DECODER_NORMALIZED_CONFIG_CLASS.decoder_num_attention_heads = ( - self._config.decoder.num_attention_heads - ) - - @property - def inputs(self) -> dict[str, dict[int, str]]: - # Batched inference is not supported in Transformers. - if self.model_part == "text_encoder": - common_inputs = { - "input_ids": {0: "batch_size", 1: "encoder_sequence_length"}, - "attention_mask": {0: "batch_size", 1: "encoder_sequence_length"}, - } - elif self.model_part == "encodec_decode": - # 0: always 1 for chunk_length_s=None, 2: num_quantizers fixed. - common_inputs = {"audio_codes": {1: "batch_size", 3: "chunk_length"}} - elif self.model_part == "build_delay_pattern_mask": - common_inputs = { - "input_ids": {0: "batch_size_x_num_codebooks"}, - "pad_token_id": {}, - "max_length": {}, - } - elif self._behavior is ConfigBehavior.DECODER: - # Naming it total_batch_size as in case we use guidance_scale, the dimension 0 may be larger than simply the batch_size. - # Reference: https://github.com/huggingface/transformers/blob/31c575bcf13c2b85b65d652dd1b5b401f99be999/src/transformers/models/musicgen/modeling_musicgen.py#L1932-L1935 - common_inputs = { - "decoder_input_ids": {0: "total_batch_size_x_num_codebooks"}, - "encoder_outputs": {0: "total_batch_size", 1: "encoder_sequence_length"}, - # MusicgenForConditionalGeneration maps attention_mask to encoder_attention_mask. - "attention_mask": { - 0: "batch_size", - 1: "encoder_sequence_length", - }, - } - if self.use_past_in_inputs: - # TODO: validate the axis name for attention_mask - # common_inputs["attention_mask"][1] = "past_encoder_sequence_length + sequence_length" - self.add_past_key_values(common_inputs, direction="inputs") - else: - common_inputs["decoder_input_ids"] = { - 0: "total_batch_size_x_num_codebooks", - 1: "decoder_sequence_length", - } - else: - raise ValueError( - "This should not happen. Please open an issue at https://github.com/huggingface/optimum/issues." - ) - - return common_inputs - - @property - def outputs(self) -> dict[str, dict[int, str]]: - common_outputs = {} - - if self.model_part == "text_encoder": - common_outputs = super().outputs - elif self.model_part == "encodec_decode": - common_outputs["audio_values"] = {0: "batch_size", 2: "audio_length"} - elif self.model_part == "build_delay_pattern_mask": - common_outputs["input_ids_edited"] = {0: "total_batch_size_x_num_codebooks"} - common_outputs["delay_pattern_mask"] = {0: "total_batch_size_x_num_codebooks", 1: "max_length"} - elif self._behavior is ConfigBehavior.DECODER: - common_outputs = super().outputs - - # MusicgenForConditionalGeneration output is named logits, not last_hidden_state. - # Rename last_hidden_state -> logits while keeping the order. - common_outputs = { - "logits" if name == "last_hidden_state" else name: value for name, value in common_outputs.items() - } - else: - raise ValueError( - "This should not happen. Please open an issue at https://github.com/huggingface/optimum/issues." - ) - - return common_outputs - - def overwrite_shape_and_generate_input( - self, dummy_input_gen: DummyInputGenerator, input_name: str, framework: str, input_shapes: dict - ): - if self.model_part == "build_delay_pattern_mask" and input_name == "input_ids": - original_batch_size = dummy_input_gen.batch_size - dummy_input_gen.batch_size = ( - original_batch_size * dummy_input_gen.normalized_config.DECODER_NORMALIZED_CONFIG_CLASS.num_codebooks - ) - dummy_input = dummy_input_gen.generate( - input_name, framework=framework, int_dtype=self.int_dtype, float_dtype=self.float_dtype - ) - dummy_input_gen.batch_size = original_batch_size - else: - dummy_input = super().overwrite_shape_and_generate_input( - dummy_input_gen, input_name, framework, input_shapes - ) - - return dummy_input - - -@register_tasks_manager_onnx("speecht5", *["text-to-audio"]) -class SpeechT5OnnxConfig(OnnxSeq2SeqConfigWithPast): - # TODO: Transformers batched generation for Speecht5 is BROKEN (https://github.com/huggingface/transformers/pull/25943), - # so we won't support for now. - NORMALIZED_CONFIG_CLASS = NormalizedSeq2SeqConfig.with_args( - hidden_size="hidden_size", - num_attention_heads="encoder_attention_heads", # TODO: bugged in case encoder and decoder have different number of heads - encoder_num_layers="encoder_layers", - decoder_num_layers="decoder_layers", - allow_new=True, - ) - - DUMMY_INPUT_GENERATOR_CLASSES = ( - DummyTextInputGenerator, - DummySeq2SeqDecoderTextInputGenerator, - DummySeq2SeqPastKeyValuesGenerator, - DummySpeechT5InputGenerator, - ) - DUMMY_PKV_GENERATOR_CLASS = DummySeq2SeqPastKeyValuesGenerator - - VARIANTS = { # noqa: RUF012 - "with-past": "The export follows the Transformers implementation using the KV cache, with the following components exported:\n\t - encoder_model.onnx: corresponds to the encoding part in https://github.com/huggingface/transformers/blob/v4.33.2/src/transformers/models/speecht5/modeling_speecht5.py#L2544-L2556.\n\t - decoder_model.onnx: corresponds to the decoder part in https://github.com/huggingface/transformers/blob/v4.33.2/src/transformers/models/speecht5/modeling_speecht5.py#L2572-L2602.\n\t - decoder_with_past_model.onnx: same as the above, with past_key_values input (KV cache filled).\n\t - decoder_postnet_and_vocoder.onnx: Decoder speech postnet and vocoder (e.g. a SpeechT5HifiGan) to generate speech from the spectrogram, as in https://github.com/huggingface/transformers/blob/v4.33.2/src/transformers/models/speecht5/modeling_speecht5.py#L2605-L2614.", - "without-past": "The same as `with-past`, just without KV cache support. This is not a recommended export as slower than `with-past`.", - } - DEFAULT_VARIANT = "with-past" - _MODEL_PATCHER = SpeechT5ModelPatcher - - def __init__( - self, - config: PretrainedConfig, - task: str = "feature-extraction", - int_dtype: str = "int64", - float_dtype: str = "fp32", - use_past: bool = False, - use_past_in_inputs: bool = False, - behavior: ConfigBehavior = ConfigBehavior.MONOLITH, - preprocessors: list[Any] | None = None, - is_postnet_and_vocoder: bool = False, - ): - super().__init__( - config=config, - task=task, - int_dtype=int_dtype, - float_dtype=float_dtype, - use_past=use_past, - use_past_in_inputs=use_past_in_inputs, - behavior=behavior, - preprocessors=preprocessors, - ) - if float_dtype == "fp16": - raise ValueError( - "The ONNX export of SpeechT5 in float16 is currently not supported due to a bug in PyTorch: https://github.com/pytorch/pytorch/pull/110078. Please open an issue in Optimum if you would like to export SpeechT5 in float16." - ) - self.is_postnet_and_vocoder = is_postnet_and_vocoder - - @property - def inputs(self) -> dict[str, dict[int, str]]: - common_inputs = {} - - # Batched inference is not supported in Transformers. - if self._behavior is ConfigBehavior.ENCODER: - common_inputs["input_ids"] = {1: "encoder_sequence_length"} - elif self._behavior is ConfigBehavior.DECODER: - # NOTE: even when past is used, the decoder takes the full sequence as input as the prenet seem to require it: - # https://github.com/huggingface/transformers/blob/v4.33.2/src/transformers/models/speecht5/modeling_speecht5.py#L2573 - common_inputs["output_sequence"] = {1: "decoder_sequence_length"} - common_inputs["speaker_embeddings"] = {} # No dynamic shape here. - common_inputs["encoder_outputs"] = {1: "encoder_sequence_length"} - common_inputs["encoder_attention_mask"] = {1: "encoder_sequence_length"} - - if self.variant == "with-past" and self.use_past_in_inputs: - self.add_past_key_values(common_inputs, direction="inputs") - elif self.is_postnet_and_vocoder: - common_inputs["spectrogram"] = {0: "n_spectrums x reduction_factor"} - else: - raise ValueError( - "self._behavior is neither encoder or decoder, and is_postnet_and_vocoder=False. This should not happen." - ) - - return common_inputs - - @property - def outputs(self) -> dict[str, dict[int, str]]: - common_outputs = {} - if self._behavior is ConfigBehavior.ENCODER: - common_outputs["encoder_outputs"] = {1: "encoder_sequence_length"} - common_outputs["encoder_attention_mask"] = {1: "encoder_sequence_length"} - elif self._behavior is ConfigBehavior.DECODER: - common_outputs["output_sequence_out"] = {1: "decoder_sequence_length + 1"} - common_outputs["spectrum"] = {} # No dynamic shape here. - common_outputs["prob"] = {} # No dynamic shape here. - - if self.variant == "with-past" and self.use_past: - # When exporting decoder models with use_cache=True, both the decoder without past and with past have the KV cache as an output. - self.add_past_key_values(common_outputs, direction="outputs") - elif self.is_postnet_and_vocoder: - common_outputs["waveform"] = {0: "n_samples"} - else: - raise ValueError( - "self._behavior is neither encoder or decoder, and is_postnet_and_vocoder=False. This should not happen." - ) - - return common_outputs - - def overwrite_shape_and_generate_input( - self, dummy_input_gen: DummyInputGenerator, input_name: str, framework: str, input_shapes: dict - ): - dummy_input_gen.batch_size = 1 - dummy_input = dummy_input_gen.generate( - input_name, framework=framework, int_dtype=self.int_dtype, float_dtype=self.float_dtype - ) - return dummy_input - - -@register_tasks_manager_onnx("vits", *["text-to-audio"]) -class VitsOnnxConfig(TextEncoderOnnxConfig): - NORMALIZED_CONFIG_CLASS = NormalizedTextConfig - - @property - def inputs(self) -> dict[str, dict[int, str]]: - return { - "input_ids": {0: "text_batch_size", 1: "sequence_length"}, - "attention_mask": {0: "text_batch_size", 1: "sequence_length"}, - } - - @property - def outputs(self) -> dict[str, dict[int, str]]: - return { - "waveform": {0: "text_batch_size", 1: "n_samples"}, - "spectrogram": {0: "text_batch_size", 2: "num_bins"}, - } - - -@register_tasks_manager_onnx( - "speech_to_text", - *[ - "feature-extraction", - "feature-extraction-with-past", - "automatic-speech-recognition", - "automatic-speech-recognition-with-past", - ], -) -class Speech2TextOnnxConfig(AudioToTextOnnxConfig): - NORMALIZED_CONFIG_CLASS = NormalizedSeq2SeqConfig.with_args( - decoder_num_layers="decoder_layers", - num_layers="decoder_layers", - input_features_per_channel="input_feat_per_channel", - allow_new=True, - ) - DUMMY_INPUT_GENERATOR_CLASSES = ( - Speech2TextDummyAudioInputGenerator, - *AudioToTextOnnxConfig.DUMMY_INPUT_GENERATOR_CLASSES[1:], - DummyTextInputGenerator, - ) - - @property - def inputs(self) -> dict[str, dict[int, str]]: - common_inputs = {} - - if self._behavior in {ConfigBehavior.ENCODER, ConfigBehavior.MONOLITH}: - common_inputs["input_features"] = {0: "batch_size", 1: "encoder_sequence_length"} - else: - common_inputs["encoder_outputs"] = {0: "batch_size", 1: "encoder_sequence_length"} - common_inputs["attention_mask"] = {0: "batch_size", 1: "encoder_sequence_length"} - - if self._behavior in {ConfigBehavior.DECODER, ConfigBehavior.MONOLITH}: - common_inputs["decoder_input_ids"] = {0: "batch_size", 1: "decoder_sequence_length"} - if self.use_past_in_inputs: - self.add_past_key_values(common_inputs, direction="inputs") - - return common_inputs - - @property - def outputs(self) -> dict[str, dict[int, str]]: - common_outputs = super().outputs - if self._behavior in {ConfigBehavior.ENCODER, ConfigBehavior.MONOLITH}: - if self._behavior is ConfigBehavior.MONOLITH: - output_name = "encoder_last_hidden_state" - else: - output_name = "last_hidden_state" - # Speech2Text encoder output formula adapted from: - # Speech2TextPreTrainedModel._get_feat_extract_output_lengths - # for i in range(self.config.num_conv_layers): - # input_lengths = (input_lengths - 1) // 2 + 1 - downsample_factor = 2 * self._config.num_conv_layers - output_sequence_length = f"( encoder_sequence_length + {downsample_factor} - 1 ) // {downsample_factor}" - common_outputs[output_name] = {0: "batch_size", 1: output_sequence_length} - return common_outputs - - -# TrOCR is a causal model, used as the decoder in some vision encoder-decoder models. -@register_tasks_manager_onnx( - "trocr", - *[ - "feature-extraction", - "feature-extraction-with-past", - ], -) -class TrOCROnnxConfig(TextSeq2SeqOnnxConfig): - NORMALIZED_CONFIG_CLASS = NormalizedSeq2SeqConfig.with_args( - decoder_num_layers="decoder_layers", - num_layers="decoder_layers", - decoder_num_attention_heads="decoder_attention_heads", - hidden_size="hidden_size", - ) - - -@register_tasks_manager_onnx( - "vision-encoder-decoder", - *[ - "image-to-text", - "image-to-text-with-past", - ], -) -class VisionEncoderDecoderOnnxConfig(EncoderDecoderBaseOnnxConfig): - NORMALIZED_CONFIG_CLASS = NormalizedEncoderDecoderConfig - DUMMY_INPUT_GENERATOR_CLASSES = (DummyVisionInputGenerator, DummyVisionEncoderDecoderPastKeyValuesGenerator) - - @property - def inputs(self) -> dict[str, dict[int, str]]: - common_inputs = {} - - if self._behavior in {ConfigBehavior.ENCODER, ConfigBehavior.MONOLITH}: - common_inputs["pixel_values"] = {0: "batch_size", 1: "num_channels", 2: "height", 3: "width"} - else: - common_inputs["encoder_outputs"] = {0: "batch_size", 1: "encoder_sequence_length"} - - if self._behavior in {ConfigBehavior.DECODER, ConfigBehavior.MONOLITH}: - common_inputs["decoder_input_ids"] = {0: "batch_size", 1: "decoder_sequence_length"} - if self.use_past_in_inputs: - self.add_past_key_values(common_inputs, direction="inputs") - - return common_inputs - - @property - def outputs(self) -> dict[str, dict[int, str]]: - if self._behavior == ConfigBehavior.ENCODER: - # Some encoders have static sequence length so it is useful to rely on the encoder ONNX config to grab this information. - return self._encoder_onnx_config.outputs - else: - # Ideally, we would want here to have self._decoder_onnx_config.outputs, which is currently not possible - # as we hard-code the task to feature-extraction, that has the wrong output names (e.g. mbart does not support document-question-answering - # so we can not initializer MBartONNXConfig with document-question-answering). - return super().outputs - - -@register_tasks_manager_onnx("sam", *["feature-extraction"]) -class SamOnnxConfig(OnnxConfig): - NORMALIZED_CONFIG_CLASS = NormalizedEncoderDecoderConfig - DUMMY_INPUT_GENERATOR_CLASSES = (DummyVisionInputGenerator, DummyPointsGenerator, DummyVisionEmbeddingsGenerator) - VARIANTS = { # noqa: RUF012 - "monolith": "All the SAM model components are exported as a single model.onnx.", - "split": "The vision encoder is exported as a separate vision_encoder.onnx, and the prompt encoder and mask decoder are exported as a prompt_encoder_mask_decoder.onnx. This allows to encoder the image only once for multiple point queries.", - } - DEFAULT_VARIANT = "split" - _MODEL_PATCHER = SAMModelPatcher - - def __init__( - self, - config: PretrainedConfig, - task: str = "feature-extraction", - int_dtype: str = "int64", - float_dtype: str = "fp32", - variant: str = "split", - vision_encoder: bool | None = None, - preprocessors: list[Any] | None = None, - ): - super().__init__( - config=config, - task=task, - int_dtype=int_dtype, - float_dtype=float_dtype, - preprocessors=preprocessors, - ) - self.variant = variant - self.vision_encoder = vision_encoder - self._normalized_config.ENCODER_NORMALIZED_CONFIG_CLASS = NormalizedVisionConfig(self._config.vision_config) - - @property - def inputs(self) -> dict[str, dict[int, str]]: - if self.variant == "monolith": - inputs = { - "pixel_values": {0: "batch_size"}, - "input_points": {0: "batch_size", 1: "point_batch_size", 2: "nb_points_per_image"}, - "input_labels": {0: "batch_size", 1: "point_batch_size", 2: "nb_points_per_image"}, - } - else: - if self.vision_encoder: - inputs = {"pixel_values": {0: "batch_size"}} - else: - inputs = { - "image_positional_embeddings": {0: "batch_size"}, - "image_embeddings": {0: "batch_size"}, - "input_points": {0: "batch_size", 1: "point_batch_size", 2: "nb_points_per_image"}, - "input_labels": {0: "batch_size", 1: "point_batch_size", 2: "nb_points_per_image"}, - } - return inputs - - @property - def outputs(self) -> dict[str, dict[int, str]]: - if self.variant == "split" and self.vision_encoder: - return {"image_embeddings": {0: "batch_size"}, "image_positional_embeddings": {0: "batch_size"}} - else: - return { - "iou_scores": {0: "batch_size", 1: "point_batch_size"}, - "pred_masks": {0: "batch_size", 1: "point_batch_size"}, - } - - -class Pix2StructNormalizedConfig(NormalizedSeq2SeqConfig): - ENCODER_NUM_LAYERS = "vision_config.num_hidden_layers" - DECODER_NUM_LAYERS = "text_config.num_layers" - ENCODER_NUM_ATTENTION_HEADS = "vision_config.num_attention_heads" - DECODER_NUM_ATTENTION_HEADS = "text_config.num_heads" - HIDDEN_SIZE = "text_config.hidden_size" - VOCAB_SIZE = "text_config.vocab_size" - - -@register_tasks_manager_onnx( - "pix2struct", - *[ - "image-to-text", - "image-to-text-with-past", - ], -) -class Pix2StructOnnxConfig(OnnxSeq2SeqConfigWithPast): - PAD_ATTENTION_MASK_TO_PAST = True - NORMALIZED_CONFIG_CLASS = Pix2StructNormalizedConfig - DUMMY_INPUT_GENERATOR_CLASSES = ( - DummyTextInputGenerator, - DummySeq2SeqDecoderTextInputGenerator, - DummySeq2SeqPastKeyValuesGenerator, - DummyPix2StructInputGenerator, - ) - - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - if is_transformers_version("==", "4.46.0"): - logging.warn_once( - logger, - "Found transformers v4.46.0 while trying to export a Pix2Struct model, " - "this specific version of transformers is broken for this model. Please " - "upgrade to v4.46.1 or higher, or downgrade to v4.45.x.", - ) - - @property - def inputs(self): - common_inputs = {} - - if self._behavior in {ConfigBehavior.ENCODER, ConfigBehavior.MONOLITH}: - common_inputs["flattened_patches"] = {0: "batch_size"} - else: - common_inputs["encoder_outputs"] = {0: "batch_size"} - common_inputs["attention_mask"] = {0: "batch_size"} - - if self._behavior in {ConfigBehavior.DECODER, ConfigBehavior.MONOLITH}: - common_inputs["decoder_input_ids"] = {0: "batch_size", 1: "decoder_sequence_length"} - if self.use_past_in_inputs: - self.add_past_key_values(common_inputs, direction="inputs") - decoder_attention_mask_dim = "past_decoder_sequence_length + decoder_sequence_length" - else: - decoder_attention_mask_dim = "decoder_sequence_length" - common_inputs["decoder_attention_mask"] = {0: "batch_size", 1: decoder_attention_mask_dim} - - return common_inputs - - @property - def outputs(self) -> dict[str, dict[int, str]]: - common_outputs = super().outputs - if self._behavior in {ConfigBehavior.ENCODER, ConfigBehavior.MONOLITH}: - if self._behavior is ConfigBehavior.MONOLITH: - output_name = "encoder_last_hidden_state" - else: - output_name = "last_hidden_state" - common_outputs[output_name] = {0: "batch_size"} # Remove unnecessary dynamic axis. - - return common_outputs - - def _create_dummy_input_generator_classes(self, **kwargs) -> list[DummyInputGenerator]: - if self._preprocessors is None or len(self._preprocessors) < 2: - raise ValueError( - f"Preprocessors for pix2struct need to be available for the ONNX export to infer input static shapes. Got: {self._preprocessors}" - ) - - dummy_inputs_generators = [] - dummy_inputs_generators.append( - self.DUMMY_INPUT_GENERATOR_CLASSES[0](self.task, self._normalized_config, **kwargs) - ) - # A hack for DummyPix2StructInputGenerator to gain access to the preprocessors. - # TODO: we probably pass preprocessors to all dummy input generators. - encoder_sequence_length = self._preprocessors[1].image_processor.max_patches - kwargs["preprocessors"] = self._preprocessors - for cls_ in self.DUMMY_INPUT_GENERATOR_CLASSES[1:]: - dummy_inputs_generators.append( - cls_(self.task, self._normalized_config, encoder_sequence_length=encoder_sequence_length, **kwargs) - ) - - return dummy_inputs_generators - - def overwrite_shape_and_generate_input( - self, dummy_input_gen: DummyInputGenerator, input_name: str, framework: str, input_shapes: dict - ): - if self._preprocessors is None or len(self._preprocessors) < 2: - raise ValueError( - f"Preprocessors for pix2struct need to be available for the ONNX export to infer input static shapes. Got: {self._preprocessors}" - ) - - # it would been simpler if pix2struct dummy input generator took care of generating these as well - if input_name in ["encoder_outputs", "attention_mask"]: - # Pix2struct takes inputs encoder inputs/outputs with a fixed sequence length (max_patches). - original_seq_length = dummy_input_gen.sequence_length - dummy_input_gen.sequence_length = self._preprocessors[1].image_processor.max_patches - dummy_input = dummy_input_gen.generate( - input_name, framework=framework, int_dtype=self.int_dtype, float_dtype=self.float_dtype - ) - dummy_input_gen.sequence_length = original_seq_length - else: - dummy_input = super().overwrite_shape_and_generate_input( - dummy_input_gen, input_name, framework, input_shapes - ) - - return dummy_input - - -@register_tasks_manager_onnx("encoder-decoder", *["text2text-generation", "text2text-generation-with-past"]) -class EncoderDecoderOnnxConfig(EncoderDecoderBaseOnnxConfig): - NORMALIZED_CONFIG_CLASS = NormalizedEncoderDecoderConfig - - -@register_tasks_manager_onnx("patchtst", *["feature-extraction", "time-series-forecasting"]) -class PatchTSTOnnxConfig(OnnxConfig): - NORMALIZED_CONFIG_CLASS = NormalizedTimeSeriesForecastingConfig - DUMMY_INPUT_GENERATOR_CLASSES = (DummyPatchTSTInputGenerator,) - - @property - def inputs(self) -> dict[str, dict[int, str]]: - return {"past_values": {0: "batch_size", 1: "sequence_length"}} - - @property - def outputs(self) -> dict[str, dict[int, str]]: - if self.task == "feature-extraction": - return {"last_hidden_state": {0: "batch_size"}} - else: - return super().outputs - - -@register_tasks_manager_onnx("patchtsmixer", *["feature-extraction", "time-series-forecasting"]) -class PatchTSMixerOnnxConfig(PatchTSTOnnxConfig): - pass - - -@register_tasks_manager_onnx("rt_detr", *["object-detection"]) -class RTDetrOnnxConfig(ViTOnnxConfig): - @property - def inputs(self) -> dict[str, dict[int, str]]: - return { - "pixel_values": {0: "batch_size", 2: "height", 3: "width"}, - } - - def _create_dummy_input_generator_classes(self, **kwargs) -> list[DummyInputGenerator]: - min_image_size = int(math.ceil(self._config.num_queries / 32) * 32) - if kwargs["height"] < min_image_size: - warnings.warn( - f"Exporting model with image `height={kwargs['height']}` which is less than " - f"minimal {min_image_size}, setting `height` to {min_image_size}.", - stacklevel=2, - ) - kwargs["height"] = min_image_size - if kwargs["width"] < min_image_size: - warnings.warn( - f"Exporting model with image `width={kwargs['width']}` which is less than " - f"minimal {min_image_size}, setting `width` to {min_image_size}.", - stacklevel=2, - ) - kwargs["width"] = min_image_size - return super()._create_dummy_input_generator_classes(**kwargs) - - -@register_tasks_manager_onnx("rt_detr_v2", *["object-detection"]) -class RTDetrV2OnnxConfig(RTDetrOnnxConfig): - pass - - -@register_tasks_manager_onnx("colpali", *["feature-extraction"]) -class ColPaliOnnxConfig(GemmaOnnxConfig): - DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, DummyVisionInputGenerator) - NORMALIZED_CONFIG_CLASS = NormalizedTextAndVisionConfig.with_args( - allow_new=True, - text_config="text_config", - vision_config="vlm_config.vision_config", - vlm_config="vlm_config", - ) - - VARIANTS = { # noqa: RUF012 - "vision": "Embedding extraction for image.", - "text": "Embedding extraction for text.", - } - DEFAULT_VARIANT = "vision" - - @property - def inputs(self) -> dict[str, dict[int, str]]: - dynamic_axis = {0: "batch_size", 1: "sequence_length"} - if self.variant == "vision": - return { - "input_ids": dynamic_axis, - "attention_mask": dynamic_axis, - "pixel_values": {0: "batch_size"}, - } - else: - return { - "input_ids": dynamic_axis, - "attention_mask": dynamic_axis, - } - - @property - def outputs(self) -> dict[str, dict[int, str]]: - return { - "embeddings": {0: "batch_size", 1: "sequence_length"}, - } - - def generate_dummy_inputs(self, framework: str = "pt", **kwargs): - if self.variant == "vision": - image_token_index = self._normalized_config.vlm_config.image_token_index - num_image_tokens = self._normalized_config.vision_config.num_image_tokens - if "sequence_length" in kwargs: - kwargs["sequence_length"] += num_image_tokens - else: - kwargs["sequence_length"] = DEFAULT_DUMMY_SHAPES["sequence_length"] + num_image_tokens - - dummy_inputs = super().generate_dummy_inputs(framework=framework, **kwargs) - - if self.variant == "vision": - dummy_inputs["input_ids"][:, :num_image_tokens] = image_token_index - return dummy_inputs - - -@register_tasks_manager_onnx("d_fine", *["object-detection"]) -class DFineOnnxConfig(RTDetrOnnxConfig): - MIN_TRANSFORMERS_VERSION = version.parse("4.52.0") - - -@register_tasks_manager_onnx("gemma2-text-encoder", *["feature-extraction"], library_name="diffusers") -class Gemma2TextEncoderOnnxConfig(Gemma2OnnxConfig): - MIN_TRANSFORMERS_VERSION = version.parse("4.42.0") - - -@register_tasks_manager_onnx("sana-transformer", *["semantic-segmentation"], library_name="diffusers") -class SanaTransformerOnnxConfig(SD3TransformerOnnxConfig): - DUMMY_INPUT_GENERATOR_CLASSES = ( - DummyTransformerVisionInputGenerator, - DummySanaTransforemerTextInputGenerator, - DummyTransformerTimestepInputGenerator, - ) - NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args( - hidden_size="caption_channels", - num_channels="in_channels", - allow_new=True, - ) - - @property - def inputs(self) -> dict[str, dict[int, str]]: - return { - "hidden_states": {0: "batch_size", 2: "height", 3: "width"}, - "encoder_hidden_states": {0: "batch_size", 1: "sequence_length"}, - "encoder_attention_mask": {0: "batch_size", 1: "sequence_length"}, - "timestep": {0: "batch_size"}, - } - - -@register_tasks_manager_onnx("dcae-encoder", *["semantic-segmentation"], library_name="diffusers") -class DcaeEncoderOnnxConfig(VaeEncoderOnnxConfig): - @property - def inputs(self) -> dict[str, dict[int, str]]: - return { - "sample": {0: "batch_size", 2: "height", 3: "width"}, - } - - @property - def outputs(self) -> dict[str, dict[int, str]]: - down_sampling_factor = 2 ** (len(self._normalized_config.encoder_block_out_channels) - 1) - - return { - "latent_sample": { - 0: "batch_size", - 2: f"height / {down_sampling_factor}", - 3: f"width / {down_sampling_factor}", - } - } - - -@register_tasks_manager_onnx("dcae-decoder", *["semantic-segmentation"], library_name="diffusers") -class DcaeDecoderOnnxConfig(VaeDecoderOnnxConfig): - @property - def inputs(self) -> dict[str, dict[int, str]]: - return { - "latent_sample": {0: "batch_size", 2: "latent_height", 3: "latent_width"}, - } - - @property - def outputs(self) -> dict[str, dict[int, str]]: - up_sampling_factor = 2 ** (len(self._normalized_config.decoder_block_out_channels) - 1) - - return { - "sample": { - 0: "batch_size", - 2: f"latent_height * {up_sampling_factor}", - 3: f"latent_width * {up_sampling_factor}", - } - } - - -@register_in_tasks_manager("baichuan", *["text-generation", "text-generation-with-past"], library_name="transformers") -class BaichaunOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): - DEFAULT_ONNX_OPSET = 13 - NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args( - num_layers="num_hidden_layers", num_attention_heads="num_attention_heads", hidden_size="hidden_size" - ) - _MODEL_PATCHER = BaichuanModelPatcher - MAX_TRANSFORMERS_VERSION = "4.57.6" - - -@register_in_tasks_manager( - "qwen2", - *[ - "text-generation", - "text-generation-with-past", - "feature-extraction", - "feature-extraction-with-past", - "text-classification", - "token-classification", - ], - library_name="transformers", -) -class Qwen2OpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): - DEFAULT_ONNX_OPSET = 14 - DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, MistralDummyPastKeyValuesGenerator) - DUMMY_PKV_GENERATOR_CLASS = MistralDummyPastKeyValuesGenerator - NORMALIZED_CONFIG_CLASS = NormalizedTextConfig - _MODEL_PATCHER = OVDecoderModelPatcher - - -@register_in_tasks_manager("qwen2_moe", *["text-generation", "text-generation-with-past"], library_name="transformers") -class Qwen2MoEOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): - DEFAULT_ONNX_OPSET = 14 - DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, MistralDummyPastKeyValuesGenerator) - DUMMY_PKV_GENERATOR_CLASS = MistralDummyPastKeyValuesGenerator - NORMALIZED_CONFIG_CLASS = NormalizedTextConfig - _MODEL_PATCHER = Qwen2MoEPatcher - - -@register_in_tasks_manager( - "qwen3", - *[ - "text-generation", - "text-generation-with-past", - "feature-extraction", - "feature-extraction-with-past", - "text-classification", - ], - library_name="transformers", -) -class Qwen3OpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): - MIN_TRANSFORMERS_VERSION = "4.51.0" - DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, GemmaDummyPastKeyValuesGenerator) - DUMMY_PKV_GENERATOR_CLASS = GemmaDummyPastKeyValuesGenerator - NORMALIZED_CONFIG_CLASS = NormalizedTextConfig - _MODEL_PATCHER = OVDecoderModelPatcher - - @property - def inputs(self) -> Dict[str, Dict[int, str]]: - if self.task in ["feature-extraction"]: - common_inputs = { - "input_ids": {0: "batch_size", 1: "sequence_length"}, - "attention_mask": {0: "batch_size", 1: "sequence_length"}, - } - else: - common_inputs = super().inputs - return common_inputs - - -@register_in_tasks_manager( - "qwen3_vl_text", - *[ - "text-generation", - "text-generation-with-past", - ], - library_name="transformers", -) -class Qwen3VLTextOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): - MIN_TRANSFORMERS_VERSION = "4.57.0" - - DUMMY_INPUT_GENERATOR_CLASSES = (DummyQwen3VLLMInputGenerator, GemmaDummyPastKeyValuesGenerator) - DUMMY_PKV_GENERATOR_CLASS = GemmaDummyPastKeyValuesGenerator - NORMALIZED_CONFIG_CLASS = NormalizedTextConfig - _MODEL_PATCHER = OVDecoderModelPatcher - - @property - def inputs(self) -> Dict[str, Dict[int, str]]: - common_inputs = super().inputs - common_inputs["visual_pos_masks"] = {0: "batch_size", 1: "sequence_length"} - common_inputs["deepstack_visual_embeds"] = {0: "num_layers", 1: "visual_seqlen"} - return common_inputs - - -@register_in_tasks_manager( - "qwen3_moe", - *["text-generation", "text-generation-with-past", "feature-extraction", "feature-extraction-with-past"], - library_name="transformers", -) -class Qwen3MoEOpenVINOConfig(Qwen3OpenVINOConfig): - _MODEL_PATCHER = Qwen3MoeModelPatcher - - -@register_in_tasks_manager("minicpm", *["text-generation", "text-generation-with-past"], library_name="transformers") -class MiniCPMOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): - DEFAULT_ONNX_OPSET = 14 - MAX_TRANSFORMERS_VERSION = "4.53.3" - DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, MistralDummyPastKeyValuesGenerator) - DUMMY_PKV_GENERATOR_CLASS = MistralDummyPastKeyValuesGenerator - NORMALIZED_CONFIG_CLASS = NormalizedTextConfig - _MODEL_PATCHER = MiniCPMModelPatcher - - -@register_in_tasks_manager("minicpm3", *["text-generation", "text-generation-with-past"], library_name="transformers") -class MiniCPM3OpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): - DEFAULT_ONNX_OPSET = 14 - MAX_TRANSFORMERS_VERSION = "4.53.3" - DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, OVMiniCPM3DummyPastKeyValuesGenerator) - DUMMY_PKV_GENERATOR_CLASS = OVMiniCPM3DummyPastKeyValuesGenerator - NORMALIZED_CONFIG_CLASS = NormalizedTextConfig - _MODEL_PATCHER = MiniCPM3Patcher - - -@register_in_tasks_manager("stablelm", *["text-generation", "text-generation-with-past"], library_name="transformers") -class StableLMOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): - DEFAULT_ONNX_OPSET = 14 - DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, MistralDummyPastKeyValuesGenerator) - DUMMY_PKV_GENERATOR_CLASS = MistralDummyPastKeyValuesGenerator - NORMALIZED_CONFIG_CLASS = NormalizedTextConfig - _MODEL_PATCHER = OVDecoderModelPatcher - - -@register_in_tasks_manager("chatglm", *["text-generation", "text-generation-with-past"], library_name="transformers") -class ChatGLM2OpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): - NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args(vocab_size="padded_vocab_size", num_layers="num_layers") - DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, ChatGLM2DummyPastKeyValuesGenerator) - DUMMY_PKV_GENERATOR_CLASS = ChatGLM2DummyPastKeyValuesGenerator - _MODEL_PATCHER = ChatGLMModelPatcher - MAX_TRANSFORMERS_VERSION = "4.55.4" - - def generate_dummy_inputs(self, framework: str = "pt", **kwargs): - dummy_inputs_generators = self._create_dummy_input_generator_classes(**kwargs) - - dummy_inputs = {} - input_names = [key for key in self.inputs.keys() if not key.startswith("past_key_values")] - if self.use_past_in_inputs and self.use_cache_branch is not False: - input_names.append("past_key_values") - - for input_name in input_names: - input_was_inserted = False - for dummy_input_gen in dummy_inputs_generators: - if dummy_input_gen.supports_input(input_name): - dummy_inputs[input_name] = self.overwrite_shape_and_generate_input( - dummy_input_gen, - input_name, - framework, - input_shapes=kwargs, - ) - input_was_inserted = True - break - if not input_was_inserted: - raise RuntimeError( - f'Could not generate dummy input for "{input_name}". Try adding a proper dummy input generator to the model ONNX config.' - ) - - # refer to https://github.com/huggingface/optimum/pull/764 - if ( - self.use_past_in_inputs - and self.PAD_ATTENTION_MASK_TO_PAST - and self.use_cache_branch is not False - and "attention_mask" in dummy_inputs - ): - # Obtain the past sequence length from the value instead of the key (Bloom). ChatGLM has seq_len in 0 dim instead of -2 - seq_len_dim = 0 if not hasattr(self._normalized_config, "rope_ratio") else -2 - past_present_length = ( - dummy_inputs["input_ids"].shape[1] + dummy_inputs["past_key_values"][0][1].shape[seq_len_dim] - ) - - dummy_inputs["attention_mask"] = DummyInputGenerator.pad_input_on_dim( - dummy_inputs["attention_mask"], - desired_length=past_present_length, - dim=1, - dtype=dummy_inputs["attention_mask"].dtype, - ) - - return dummy_inputs - - def add_past_key_values(self, inputs_or_outputs: Dict[str, Dict[int, str]], direction: str): - """ - Fills `input_or_outputs` mapping with past_key_values dynamic axes considering the direction. + def add_past_key_values(self, inputs_or_outputs: Dict[str, Dict[int, str]], direction: str): + """ + Fills `input_or_outputs` mapping with past_key_values dynamic axes considering the direction. Args: inputs_or_outputs (`Dict[str, Dict[int, str]]`): The mapping to fill. @@ -3371,7 +601,11 @@ class MixtralOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): ], library_name="transformers", ) -class GemmaOpenVINOConfig(GemmaOnnxConfig): +class GemmaOpenVINOConfig(TextDecoderOnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig + DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, GemmaDummyPastKeyValuesGenerator) + DUMMY_PKV_GENERATOR_CLASS = GemmaDummyPastKeyValuesGenerator + MIN_TRANSFORMERS_VERSION = "4.38.0" _MODEL_PATCHER = OVDecoderModelPatcher @property @@ -3397,7 +631,10 @@ def inputs(self) -> Dict[str, Dict[int, str]]: ], library_name="transformers", ) -class LlamaOpenVINOConfig(LlamaOnnxConfig): +class LlamaOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): + DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, MistralDummyPastKeyValuesGenerator) + DUMMY_PKV_GENERATOR_CLASS = MistralDummyPastKeyValuesGenerator + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig _MODEL_PATCHER = OVDecoderModelPatcher def __init__( @@ -3494,7 +731,10 @@ class GptOssOpenVINOConfig(LlamaOpenVINOConfig): ], library_name="transformers", ) -class BitnetOpenVINOConfig(LlamaOnnxConfig): +class BitnetOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): + DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, MistralDummyPastKeyValuesGenerator) + DUMMY_PKV_GENERATOR_CLASS = MistralDummyPastKeyValuesGenerator + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig MIN_TRANSFORMERS_VERSION = "4.52.1" MAX_TRANSFORMERS_VERSION = "4.57.6" _MODEL_PATCHER = OVDecoderModelPatcher @@ -3679,7 +919,10 @@ class OlmoOpenVINOConfig(LlamaOpenVINOConfig): @register_in_tasks_manager( "mpt", *["text-generation", "text-generation-with-past", "text-classification"], library_name="transformers" ) -class MPTOpenVINOConfig(MPTOnnxConfig): +class MPTOpenVINOConfig(TextDecoderOnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args( + num_attention_heads="n_heads", hidden_size="d_model", num_layers="n_layers" + ) _MODEL_PATCHER = MPTModelPatcher @@ -3694,12 +937,13 @@ class MPTOpenVINOConfig(MPTOnnxConfig): ], library_name="transformers", ) -class Phi3OpenVINOConfig(PhiOnnxConfig): +class Phi3OpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args(num_key_value_heads="num_key_value_heads", allow_new=True) DUMMY_INPUT_GENERATOR_CLASSES = ( MistralDummyPastKeyValuesGenerator, ) + TextDecoderOnnxConfig.DUMMY_INPUT_GENERATOR_CLASSES DUMMY_PKV_GENERATOR_CLASS = MistralDummyPastKeyValuesGenerator - NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args(num_key_value_heads="num_key_value_heads", allow_new=True) + MIN_TRANSFORMERS_VERSION = "4.36.0" _MODEL_PATCHER = Phi3ModelPatcher @@ -3730,7 +974,9 @@ class PhiMoEOpenVINOConfig(Phi3OpenVINOConfig): ], library_name="transformers", ) -class PhiOpenVINOConfig(PhiOnnxConfig): +class PhiOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig + MIN_TRANSFORMERS_VERSION = "4.36.0" _MODEL_PATCHER = OVDecoderModelPatcher @@ -3746,13 +992,23 @@ class PhiOpenVINOConfig(PhiOnnxConfig): ], library_name="transformers", ) -class FalconOpenVINOConfig(FalconOnnxConfig): +class FalconOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig DUMMY_INPUT_GENERATOR_CLASSES = ( OVFalconDummyPastKeyValuesGenerator, ) + TextDecoderOnnxConfig.DUMMY_INPUT_GENERATOR_CLASSES DUMMY_PKV_GENERATOR_CLASS = OVFalconDummyPastKeyValuesGenerator _MODEL_PATCHER = FalconModelPatcher + @property + def inputs(self) -> Dict[str, Dict[int, str]]: + common_inputs = super().inputs + + if self._config.alibi: + common_inputs.pop("position_ids", None) + + return common_inputs + @register_in_tasks_manager( "persimmon", @@ -3790,7 +1046,8 @@ class BioGPTOpenVINOConfig( ], library_name="transformers", ) -class GPTNeoOpenVINOConfig(GPTNeoOnnxConfig): +class GPTNeoOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args(num_attention_heads="num_heads") _MODEL_PATCHER = GptNeoModelPatcher @@ -3805,7 +1062,8 @@ class GPTNeoOpenVINOConfig(GPTNeoOnnxConfig): ], library_name="transformers", ) -class GPTJOpenVINOConfig(GPTJOnnxConfig): +class GPTJOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args(num_layers="n_layer", num_attention_heads="n_head") _MODEL_PATCHER = GptJModelPatcher @@ -3821,9 +1079,32 @@ class GPTJOpenVINOConfig(GPTJOnnxConfig): ], library_name="transformers", ) -class BloomOpenVINOConfig(BloomOnnxConfig): +class BloomOpenVINOConfig(TextDecoderOnnxConfig): + # Bloom does not require position_ids input. + MIN_TRANSFORMERS_VERSION = "4.36.0" + DUMMY_PKV_GENERATOR_CLASS = BloomDummyPastKeyValuesGenerator + DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, BloomDummyPastKeyValuesGenerator) + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args(num_layers="n_layer", num_attention_heads="n_head") _MODEL_PATCHER = BloomModelPatcher + def add_past_key_values(self, inputs_or_outputs: Dict[str, Dict[int, str]], direction: str): + if is_transformers_version(">=", "4.44"): + super().add_past_key_values(inputs_or_outputs, direction) + else: + if direction not in ["inputs", "outputs"]: + raise ValueError(f'direction must either be "inputs" or "outputs", but {direction} was given') + + if direction == "inputs": + decoder_sequence_name = "past_sequence_length" + name = "past_key_values" + else: + decoder_sequence_name = "past_sequence_length + sequence_length" + name = "present" + + for i in range(self._normalized_config.num_layers): + inputs_or_outputs[f"{name}.{i}.key"] = {0: "batch_size * num_heads", 2: decoder_sequence_name} + inputs_or_outputs[f"{name}.{i}.value"] = {0: "batch_size * num_heads", 1: decoder_sequence_name} + @register_in_tasks_manager( "cohere", @@ -3888,7 +1169,8 @@ class InternLMOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): *["feature-extraction", "feature-extraction-with-past", "text-generation", "text-generation-with-past"], library_name="transformers", ) -class CodeGenOpenVINOConfig(CodeGenOnnxConfig): +class CodeGenOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args(num_layers="n_layer", num_attention_heads="n_head") _MODEL_PATCHER = CodeGenModelPatcher @@ -3943,7 +1225,8 @@ class ArcticOpenVINOConfig(MixtralOpenVINOConfig): ], library_name="transformers", ) -class MistralOpenVINOConfig(MistralOnnxConfig): +class MistralOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args(num_key_value_heads="num_key_value_heads", allow_new=True) DUMMY_INPUT_GENERATOR_CLASSES = ( OVMistralDummyPastKeyValuesGenerator, ) + TextDecoderOnnxConfig.DUMMY_INPUT_GENERATOR_CLASSES @@ -3962,7 +1245,8 @@ class MistralOpenVINOConfig(MistralOnnxConfig): ], library_name="transformers", ) -class GPTNeoxOpenVINOConfig(GPTNeoXOnnxConfig): +class GPTNeoxOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig _MODEL_PATCHER = GptNeoxModelPatcher @@ -4056,9 +1340,20 @@ class DeciOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): _MODEL_PATCHER = DeciLMModelPatcher +class CLIPNormalizedConfig(NormalizedTextAndVisionConfig): + TEXT_CONFIG = "text_config" + VISION_CONFIG = "vision_config" + + +class SiglipNormalizedConfig(CLIPNormalizedConfig): + pass + + @register_in_tasks_manager("clip", *["zero-shot-image-classification"], library_name="open_clip") -class OpenCLIPOpenVINOConfig(CLIPOnnxConfig): +class OpenCLIPOpenVINOConfig(TextAndVisionOnnxConfig): DEFAULT_ONNX_OPSET = 14 + NORMALIZED_CONFIG_CLASS = CLIPNormalizedConfig + _MODEL_PATCHER = CLIPModelPatcher @property def inputs(self) -> Dict[str, Dict[int, str]]: @@ -4099,8 +1394,15 @@ def generate_dummy_inputs_for_validation( @register_in_tasks_manager("clip_text_model", *["feature-extraction"], library_name="open_clip") -class OpenCLIPTextOpenVINOConfig(CLIPTextOnnxConfig): +class OpenCLIPTextOpenVINOConfig(TextEncoderOnnxConfig): DEFAULT_ONNX_OPSET = 14 + NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args( + vocab_size="vocab_size", + sequence_length="max_position_embeddings", + num_layers="num_hidden_layers", + allow_new=True, + ) + _MODEL_PATCHER = CLIPModelPatcher @property def inputs(self) -> Dict[str, Dict[int, str]]: @@ -4155,24 +1457,113 @@ def rename_ambiguous_inputs(self, inputs): @register_in_tasks_manager( "clip", *["feature-extraction", "zero-shot-image-classification"], library_name="transformers" ) -class CLIPOpenVINOConfig(CLIPOnnxConfig): - pass +class CLIPOpenVINOConfig(TextAndVisionOnnxConfig): + NORMALIZED_CONFIG_CLASS = CLIPNormalizedConfig + _MODEL_PATCHER = CLIPModelPatcher + + @property + def inputs(self) -> Dict[str, Dict[int, str]]: + inputs = {"pixel_values": {0: "batch_size", 1: "num_channels", 2: "height", 3: "width"}} + + if self.task in ["feature-extraction", "zero-shot-image-classification"]: + inputs.update( + { + "input_ids": {0: "text_batch_size", 1: "sequence_length"}, + "attention_mask": {0: "text_batch_size", 1: "sequence_length"}, + } + ) + + return inputs + + @property + def outputs(self) -> Dict[str, Dict[int, str]]: + if self.task in ["feature-extraction", "zero-shot-image-classification"]: + return { + "logits_per_image": {0: "image_batch_size", 1: "text_batch_size"}, + "logits_per_text": {0: "text_batch_size", 1: "image_batch_size"}, + "text_embeds": {0: "text_batch_size"}, + "image_embeds": {0: "image_batch_size"}, + } + else: + return super().outputs @register_in_tasks_manager("clip_text_model", *["feature-extraction"], library_name="transformers") @register_in_tasks_manager("clip-text", *["feature-extraction"], library_name="diffusers") -class CLIPTextOpenVINOConfig(CLIPTextOnnxConfig): - pass +class CLIPTextOpenVINOConfig(TextEncoderOnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args( + vocab_size="vocab_size", + sequence_length="max_position_embeddings", + num_layers="num_hidden_layers", + allow_new=True, + ) + _MODEL_PATCHER = CLIPModelPatcher + + @property + def inputs(self) -> Dict[str, Dict[int, str]]: + return { + "input_ids": {0: "batch_size", 1: "sequence_length"}, + } + + @property + def outputs(self) -> Dict[str, Dict[int, str]]: + common_outputs = { + "last_hidden_state": {0: "batch_size", 1: "sequence_length"}, + "pooler_output": {0: "batch_size"}, + } + + if self._normalized_config.output_hidden_states: + for i in range(self._normalized_config.num_layers + 1): + common_outputs[f"hidden_states.{i}"] = {0: "batch_size", 1: "sequence_length"} + + return common_outputs @register_in_tasks_manager("clip-text-with-projection", *["feature-extraction"], library_name="diffusers") -class CLIPTextWithProjectionOpenVINOConfig(CLIPTextWithProjectionOnnxConfig): - pass +class CLIPTextWithProjectionOpenVINOConfig(TextEncoderOnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args( + vocab_size="vocab_size", + sequence_length="max_position_embeddings", + num_layers="num_hidden_layers", + allow_new=True, + ) + _MODEL_PATCHER = CLIPModelPatcher + + @property + def inputs(self) -> Dict[str, Dict[int, str]]: + return { + "input_ids": {0: "batch_size", 1: "sequence_length"}, + } + + @property + def outputs(self) -> Dict[str, Dict[int, str]]: + common_outputs = { + "text_embeds": {0: "batch_size", 1: "sequence_length"}, + "last_hidden_state": {0: "batch_size", 1: "sequence_length"}, + } + if self._normalized_config.output_hidden_states: + for i in range(self._normalized_config.num_layers + 1): + common_outputs[f"hidden_states.{i}"] = {0: "batch_size", 1: "sequence_length"} + + return common_outputs @register_in_tasks_manager("clip_vision_model", *["feature-extraction"], library_name="transformers") -class CLIPVisionModelOpenVINOConfig(CLIPVisionModelOnnxConfig): - pass +class CLIPVisionModelOpenVINOConfig(VisionOnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedVisionConfig + _MODEL_PATCHER = CLIPModelPatcher + + @property + def inputs(self) -> Dict[str, Dict[int, str]]: + return {"pixel_values": {0: "batch_size", 1: "num_channels", 2: "height", 3: "width"}} + + @property + def outputs(self) -> Dict[str, Dict[int, str]]: + common_outputs = super().outputs + common_outputs["last_hidden_state"] = {0: "batch_size"} + common_outputs["pooler_output"] = {0: "batch_size"} + + return common_outputs @register_in_tasks_manager( @@ -4187,9 +1578,18 @@ class CLIPVisionModelOpenVINOConfig(CLIPVisionModelOnnxConfig): ], library_name="transformers", ) -class IBertOpenVINOConfig(IBertOnnxConfig): +class IBertOpenVINOConfig(TextEncoderOnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig _MODEL_PATCHER = IBertModelPatcher + @property + def inputs(self) -> Dict[str, Dict[int, str]]: + if self.task == "multiple-choice": + dynamic_axis = {0: "batch_size", 1: "num_choices", 2: "sequence_length"} + else: + dynamic_axis = {0: "batch_size", 1: "sequence_length"} + return {"input_ids": dynamic_axis, "attention_mask": dynamic_axis} + # TODO: this is a very confusing class TBH, why not simply decompose the VLM into components, like diffusion models ? class LMInputEmbedsConfigHelper(TextDecoderWithPositionIdsOnnxConfig): @@ -4791,7 +2191,14 @@ def rename_ambiguous_inputs(self, inputs): @register_in_tasks_manager("unet", *["semantic-segmentation"], library_name="diffusers") @register_in_tasks_manager("unet-2d-condition", *["semantic-segmentation"], library_name="diffusers") -class UNetOpenVINOConfig(UNetOnnxConfig): +class UNetOpenVINOConfig(VisionOnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args( + image_size="sample_size", + num_channels="in_channels", + hidden_size="cross_attention_dim", + vocab_size="norm_num_groups", + allow_new=True, + ) DUMMY_INPUT_GENERATOR_CLASSES = ( DummyUnetVisionInputGenerator, DummyUnetTimestepInputGenerator, @@ -4800,12 +2207,59 @@ class UNetOpenVINOConfig(UNetOnnxConfig): @property def inputs(self) -> Dict[str, Dict[int, str]]: - common_inputs = super().inputs - common_inputs["timestep"] = {0: "batch_size"} + common_inputs = { + "sample": {0: "batch_size", 2: "height", 3: "width"}, + "timestep": {0: "batch_size"}, + "encoder_hidden_states": {0: "batch_size", 1: "sequence_length"}, + } + + # TODO : add addition_embed_type == text_image, image and image_embeds + if getattr(self._normalized_config, "addition_embed_type", None) == "text_time": + common_inputs["text_embeds"] = {0: "batch_size"} + common_inputs["time_ids"] = {0: "batch_size"} + + if getattr(self._normalized_config, "time_cond_proj_dim", None) is not None: + common_inputs["timestep_cond"] = {0: "batch_size"} + if hasattr(self._normalized_config.config, "model_max_length"): common_inputs["encoder_hidden_states"] = {0: "batch_size"} + return common_inputs + @property + def outputs(self) -> Dict[str, Dict[int, str]]: + return { + "out_sample": {0: "batch_size", 2: "height", 3: "width"}, + } + + @property + def torch_to_onnx_output_map(self) -> Dict[str, str]: + return { + "sample": "out_sample", + } + + def generate_dummy_inputs(self, framework: str = "pt", **kwargs): + dummy_inputs = super().generate_dummy_inputs(framework=framework, **kwargs) + dummy_inputs["encoder_hidden_states"] = dummy_inputs["encoder_hidden_states"][0] + + if getattr(self._normalized_config, "addition_embed_type", None) == "text_time": + dummy_inputs["added_cond_kwargs"] = { + "text_embeds": dummy_inputs.pop("text_embeds"), + "time_ids": dummy_inputs.pop("time_ids"), + } + + return dummy_inputs + + def ordered_inputs(self, model) -> Dict[str, Dict[int, str]]: + inputs = super().ordered_inputs(model=model) + # to fix mismatch between model forward signature and expected inputs + # a dictionary of additional embeddings `added_cond_kwargs` is expected depending on config.addition_embed_type + if getattr(self._normalized_config, "addition_embed_type", None) == "text_time": + inputs["text_embeds"] = self.inputs["text_embeds"] + inputs["time_ids"] = self.inputs["time_ids"] + + return inputs + @register_in_tasks_manager("sd3-transformer", *["semantic-segmentation"], library_name="diffusers") @register_in_tasks_manager("sd3-transformer-2d", *["semantic-segmentation"], library_name="diffusers") @@ -4887,17 +2341,62 @@ def rename_ambiguous_inputs(self, inputs): @register_in_tasks_manager("vae-encoder", *["semantic-segmentation"], library_name="diffusers") -class VaeEncoderOpenVINOConfig(VaeEncoderOnnxConfig): - pass +class VaeEncoderOpenVINOConfig(VisionOnnxConfig): + ATOL_FOR_VALIDATION = 3e-4 # TODO: this only happens in test_export.py + NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args( + num_channels="in_channels", image_size="sample_size", allow_new=True + ) + + @property + def inputs(self) -> Dict[str, Dict[int, str]]: + return { + "sample": {0: "batch_size", 2: "sample_height", 3: "sample_width"}, + } + + @property + def outputs(self) -> Dict[str, Dict[int, str]]: + down_sampling_factor = 2 ** (len(self._normalized_config.down_block_types) - 1) + + return { + "latent_parameters": { + 0: "batch_size", + 2: f"sample_height / {down_sampling_factor}", + 3: f"sample_width / {down_sampling_factor}", + }, + } @register_in_tasks_manager("vae-decoder", *["semantic-segmentation"], library_name="diffusers") -class VaeDecoderOpenVINOConfig(VaeDecoderOnnxConfig): - pass +class VaeDecoderOpenVINOConfig(VisionOnnxConfig): + ATOL_FOR_VALIDATION = 3e-4 # TODO: this only happens in test_export.py + NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args(num_channels="latent_channels", allow_new=True) + + @property + def inputs(self) -> Dict[str, Dict[int, str]]: + return { + "latent_sample": {0: "batch_size", 2: "latent_height", 3: "latent_width"}, + } + + @property + def outputs(self) -> Dict[str, Dict[int, str]]: + up_sampling_factor = 2 ** (len(self._normalized_config.up_block_types) - 1) + + return { + "sample": { + 0: "batch_size", + 2: f"latent_height * {up_sampling_factor}", + 3: f"latent_width * {up_sampling_factor}", + }, + } @register_in_tasks_manager("dcae-encoder", *["semantic-segmentation"], library_name="diffusers") -class DcaeEncoderOpenVINOConfig(VaeEncoderOnnxConfig): +class DcaeEncoderOpenVINOConfig(VisionOnnxConfig): + ATOL_FOR_VALIDATION = 3e-4 # TODO: this only happens in test_export.py + NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args( + num_channels="in_channels", image_size="sample_size", allow_new=True + ) + @property def inputs(self) -> Dict[str, Dict[int, str]]: return { @@ -4912,7 +2411,10 @@ def outputs(self) -> Dict[str, Dict[int, str]]: @register_in_tasks_manager("dcae-decoder", *["semantic-segmentation"], library_name="diffusers") -class DcaeDecoderOpenVINOConfig(VaeDecoderOnnxConfig): +class DcaeDecoderOpenVINOConfig(VisionOnnxConfig): + ATOL_FOR_VALIDATION = 3e-4 # TODO: this only happens in test_export.py + NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args(num_channels="latent_channels", allow_new=True) + @property def inputs(self) -> Dict[str, Dict[int, str]]: return { @@ -4935,7 +2437,7 @@ class FluxTransformerOpenVINOConfig(SD3TransformerOpenVINOConfig): DummyFluxTextInputGenerator, PooledProjectionsDummyInputGenerator, ) - _MODEL_PATCHER = FluxTransfromerModelPatcher + _MODEL_PATCHER = FluxTransformerModelPatcher @property def inputs(self): @@ -4956,7 +2458,11 @@ def inputs(self): @register_in_tasks_manager("ltx-vae-encoder", *["semantic-segmentation"], library_name="diffusers") -class LTXVaeEncoderOpenVINOConfig(VaeEncoderOnnxConfig): +class LTXVaeEncoderOpenVINOConfig(VisionOnnxConfig): + ATOL_FOR_VALIDATION = 3e-4 # TODO: this only happens in test_export.py + NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args( + num_channels="in_channels", image_size="sample_size", allow_new=True + ) DUMMY_INPUT_GENERATOR_CLASSES = (LTXVaeDummyInputGenerator,) @property @@ -4973,7 +2479,9 @@ def outputs(self) -> Dict[str, Dict[int, str]]: @register_in_tasks_manager("ltx-vae-decoder", *["semantic-segmentation"], library_name="diffusers") -class LTXVaeDecoderOpenVINOConfig(VaeDecoderOnnxConfig): +class LTXVaeDecoderOpenVINOConfig(VisionOnnxConfig): + ATOL_FOR_VALIDATION = 3e-4 # TODO: this only happens in test_export.py + NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args(num_channels="latent_channels", allow_new=True) DUMMY_INPUT_GENERATOR_CLASSES = (LTXVaeDummyInputGenerator,) @property @@ -5886,9 +3394,51 @@ class GraniteMoEOpenVINOConfig(LlamaOpenVINOConfig): ], library_name="transformers", ) -class WhisperOpenVINOConfig(WhisperOnnxConfig): +class WhisperOpenVINOConfig(AudioToTextOnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedSeq2SeqConfig.with_args( + encoder_num_layers="encoder_layers", + decoder_num_layers="decoder_layers", + feature_size="num_mel_bins", + allow_new=True, + ) _MODEL_PATCHER = OVSeq2SeqModelPatcher + @property + def inputs(self) -> Dict[str, Dict[int, str]]: + if self.task == "audio-classification": + return {"input_features": {0: "batch_size"}} + + common_inputs = super().inputs + if self._behavior in {ConfigBehavior.ENCODER, ConfigBehavior.MONOLITH}: + common_inputs["input_features"] = {0: "batch_size"} # Remove unnecessary dynamic axis. + else: + # the dynamic encoder sequence length is only needed here because the input generator generates + # encoder_outputs with a seq_len=16 but the model expects at inference time seq_len=1500 + # TODO: this can be fixed by generating the correct inputs in the input generator + common_inputs["encoder_outputs"] = {0: "batch_size", 1: "encoder_sequence_length"} + + if self._behavior in {ConfigBehavior.DECODER, ConfigBehavior.MONOLITH}: + if is_transformers_version(">=", "4.43.0") and is_transformers_version("<", "4.46.0"): + # since https://github.com/huggingface/transformers/pull/31166 + if self._behavior is not ConfigBehavior.ENCODER and self.use_past_in_inputs: + common_inputs["cache_position"] = {0: "decoder_sequence_length"} + + return common_inputs + + @property + def outputs(self) -> Dict[str, Dict[int, str]]: + if self.task == "audio-classification": + return {"logits": {0: "batch_size"}} + + common_outputs = super().outputs + if self._behavior in {ConfigBehavior.ENCODER, ConfigBehavior.MONOLITH}: + if self._behavior is ConfigBehavior.MONOLITH: + output_name = "encoder_last_hidden_state" + else: + output_name = "last_hidden_state" + common_outputs[output_name] = {0: "batch_size"} # Remove unnecessary dynamic axis. + return common_outputs + @register_in_tasks_manager( "qwen3_asr", @@ -5982,7 +3532,20 @@ def add_past_key_values(self, inputs_or_outputs: Dict[str, Dict[int, str]], dire *["feature-extraction", "feature-extraction-with-past", "text2text-generation", "text2text-generation-with-past"], library_name="transformers", ) -class T5OpenVINOConfig(T5OnnxConfig): +class T5OpenVINOConfig(TextSeq2SeqOnnxConfig): + DUMMY_INPUT_GENERATOR_CLASSES = ( + *TextSeq2SeqOnnxConfig.DUMMY_INPUT_GENERATOR_CLASSES[:-1], + T5DummySeq2SeqPastKeyValuesGenerator, + ) + DUMMY_PKV_GENERATOR_CLASS = T5DummySeq2SeqPastKeyValuesGenerator + NORMALIZED_CONFIG_CLASS = NormalizedSeq2SeqConfig.with_args( + hidden_size="d_model", + num_attention_heads="num_heads", + encoder_num_layers="num_layers", + decoder_num_layers="num_decoder_layers", + key_value_dim="d_kv", + allow_new=True, + ) _MODEL_PATCHER = OVSeq2SeqModelPatcher @@ -5996,31 +3559,146 @@ class MT5OpenVINOConfig(T5OpenVINOConfig): MAX_TRANSFORMERS_VERSION = "4.57.6" -@register_in_tasks_manager( - "longt5", - *["feature-extraction", "feature-extraction-with-past", "text2text-generation", "text2text-generation-with-past"], - library_name="transformers", -) -class LongT5OpenVINOConfig(T5OpenVINOConfig): - pass +@register_in_tasks_manager( + "longt5", + *["feature-extraction", "feature-extraction-with-past", "text2text-generation", "text2text-generation-with-past"], + library_name="transformers", +) +class LongT5OpenVINOConfig(T5OpenVINOConfig): + pass + + +@register_in_tasks_manager( + "bart", + *[ + "feature-extraction", + "feature-extraction-with-past", + "text-generation", + "text-generation-with-past", + "text2text-generation", + "text2text-generation-with-past", + "text-classification", + "question-answering", + ], + library_name="transformers", +) +class BartOpenVINOConfig(TextSeq2SeqOnnxConfig): + PAD_ATTENTION_MASK_TO_PAST = True + + NORMALIZED_CONFIG_CLASS = NormalizedSeq2SeqConfig.with_args( + encoder_num_layers="encoder_layers", + decoder_num_layers="decoder_layers", + num_layers="decoder_layers", # Used for the text-generation task past key values input generation. + encoder_num_attention_heads="encoder_attention_heads", + decoder_num_attention_heads="decoder_attention_heads", + eos_token_id="eos_token_id", + ) + DUMMY_INPUT_GENERATOR_CLASSES = ( + BartDummyTextInputGenerator, + { + "feature-extraction": DummySeq2SeqDecoderTextInputGenerator, + "text-generation": DummyDecoderTextInputGenerator, + }, + { + "feature-extraction": DummySeq2SeqPastKeyValuesGenerator, + "text-generation": DummyPastKeyValuesGenerator, + }, + ) + _MODEL_PATCHER = OVSeq2SeqModelPatcher + + def _create_dummy_input_generator_classes(self, **kwargs): + dummy_text_input_generator = self.DUMMY_INPUT_GENERATOR_CLASSES[0]( + self.task, self._normalized_config, **kwargs + ) + task = "feature-extraction" if self.task != "text-generation" else "text-generation" + dummy_decoder_text_input_generator = self.DUMMY_INPUT_GENERATOR_CLASSES[1][task]( + self.task, self._normalized_config, **kwargs + ) + if self.task != "text-generation": + kwargs["encoder_sequence_length"] = dummy_text_input_generator.sequence_length + + dummy_seq2seq_past_key_values_generator = self.DUMMY_INPUT_GENERATOR_CLASSES[2][task]( + self.task, self._normalized_config, **kwargs + ) + dummy_inputs_generators = [ + dummy_text_input_generator, + dummy_decoder_text_input_generator, + dummy_seq2seq_past_key_values_generator, + ] + + return dummy_inputs_generators + + @property + def inputs_for_default_and_seq2seq_lm(self): + return super().inputs + + @property + def inputs_for_causal_lm(self): + if self.use_past_in_inputs: + common_inputs = { + "input_ids": {0: "batch_size", 1: "sequence_length"}, + "attention_mask": {0: "batch_size", 1: "past_sequence_length + sequence_length"}, + } + for i in range(self._normalized_config.decoder_num_layers): + common_inputs[f"past_key_values.{i}.key"] = { + 0: "batch_size", + 2: "past_sequence_length", + } + common_inputs[f"past_key_values.{i}.value"] = { + 0: "batch_size", + 2: "past_sequence_length", + } + else: + common_inputs = { + "input_ids": {0: "batch_size", 1: "sequence_length"}, + "attention_mask": {0: "batch_size", 1: "sequence_length"}, + } + + return common_inputs + + @property + def inputs_for_other_tasks(self): + return { + "input_ids": {0: "batch_size", 1: "sequence_length"}, + "attention_mask": {0: "batch_size", 1: "sequence_length"}, + } + @property + def inputs(self) -> dict: + inputs_properties = { + "feature-extraction": self.inputs_for_default_and_seq2seq_lm, + "text2text-generation": self.inputs_for_default_and_seq2seq_lm, + "text-generation": self.inputs_for_causal_lm, + "other": self.inputs_for_other_tasks, + } + return inputs_properties.get(self.task, inputs_properties["other"]) -@register_in_tasks_manager( - "bart", - *[ - "feature-extraction", - "feature-extraction-with-past", - "text-generation", - "text-generation-with-past", - "text2text-generation", - "text2text-generation-with-past", - "text-classification", - "question-answering", - ], - library_name="transformers", -) -class BartOpenVINOConfig(BartOnnxConfig): - _MODEL_PATCHER = OVSeq2SeqModelPatcher + @property + def outputs(self) -> dict: + if self.task in ["feature-extraction", "text2text-generation"]: + common_outputs = super().outputs + else: + common_outputs = super(OnnxConfigWithPast, self).outputs + if self.use_past: + for i in range( + self._normalized_config.encoder_num_layers + if self.task != "text-generation" + else self._normalized_config.decoder_num_layers + ): + common_outputs[f"present.{i}.key"] = {0: "batch_size", 2: "past_sequence_length + sequence_length"} + common_outputs[f"present.{i}.value"] = { + 0: "batch_size", + 2: "past_sequence_length + sequence_length", + } + return common_outputs + + def flatten_past_key_values(self, flattened_output, name, idx, t): + if self.task in ["feature-extraction", "text2text-generation"]: + flattened_output = super().flatten_past_key_values(flattened_output, name, idx, t) + else: + flattened_output = super(OnnxSeq2SeqConfigWithPast, self).flatten_past_key_values( + flattened_output, name, idx, t + ) @register_in_tasks_manager( @@ -6440,7 +4118,7 @@ class SmolVLMOpenVINOConfig(Idefics3OpenVINOConfig): ], library_name="transformers", ) -class BlenderbotOpenVINOConfig(BlenderbotOnnxConfig): +class BlenderbotOpenVINOConfig(BartOpenVINOConfig): _MODEL_PATCHER = BlenderbotModelPatcher @@ -6456,7 +4134,7 @@ class BlenderbotOpenVINOConfig(BlenderbotOnnxConfig): ], library_name="transformers", ) -class BlenderbotSmallOpenVINOConfig(BlenderbotSmallOnnxConfig): +class BlenderbotSmallOpenVINOConfig(BartOpenVINOConfig): _MODEL_PATCHER = BlenderbotSmallModelPatcher @@ -6472,7 +4150,7 @@ class BlenderbotSmallOpenVINOConfig(BlenderbotSmallOnnxConfig): ], library_name="transformers", ) -class PegasusOpenVINOConfig(PegasusOnnxConfig): +class PegasusOpenVINOConfig(BartOpenVINOConfig): _MODEL_PATCHER = PegasusModelPatcher def __init__(self, *args, **kwargs): @@ -6492,7 +4170,7 @@ def __init__(self, *args, **kwargs): ], library_name="transformers", ) -class MarianOpenVINOConfig(MarianOnnxConfig): +class MarianOpenVINOConfig(BartOpenVINOConfig): _MODEL_PATCHER = MarianModelPatcher # TODO (@echarlaix): add v5 support MAX_TRANSFORMERS_VERSION = "4.57.6" @@ -6510,13 +4188,26 @@ class SpeechT5ConfigBehavior(str, enum.Enum): *["text-to-audio", "text-to-audio-with-past"], library_name="transformers", ) -class SpeechT5OpenVINOConfig(SpeechT5OnnxConfig): +class SpeechT5OpenVINOConfig(OnnxSeq2SeqConfigWithPast): + NORMALIZED_CONFIG_CLASS = NormalizedSeq2SeqConfig.with_args( + hidden_size="hidden_size", + num_attention_heads="encoder_attention_heads", + encoder_num_layers="encoder_layers", + decoder_num_layers="decoder_layers", + allow_new=True, + ) DUMMY_INPUT_GENERATOR_CLASSES = ( DummyTextInputGenerator, DummySeq2SeqPastKeyValuesGenerator, - DummySpeechT5OpenVINOInputGenerator, + DummySpeechT5InputGenerator, ) - _MODEL_PATCHER = OVSpeechT5ModelPatcher + DUMMY_PKV_GENERATOR_CLASS = DummySeq2SeqPastKeyValuesGenerator + VARIANTS = { # noqa: RUF012 + "with-past": "The export follows the Transformers implementation using the KV cache.", + "without-past": "The same as `with-past`, just without KV cache support.", + } + DEFAULT_VARIANT = "with-past" + _MODEL_PATCHER = SpeechT5ModelPatcher def __init__( self, @@ -6765,7 +4456,8 @@ def generate_dummy_inputs(self, framework: str = "pt", **kwargs): ], library_name="transformers", ) -class GPT2OpenVINOConfig(GPT2OnnxConfig): +class GPT2OpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args(num_layers="n_layer", num_attention_heads="n_head") _MODEL_PATCHER = OVDecoderModelPatcher @@ -6778,9 +4470,34 @@ class GPT2OpenVINOConfig(GPT2OnnxConfig): "document-question-answering-with-past", ], ) -class VisionEncoderDecoderOpenVINOConfig(VisionEncoderDecoderOnnxConfig): +class VisionEncoderDecoderOpenVINOConfig(EncoderDecoderBaseOnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedEncoderDecoderConfig + DUMMY_INPUT_GENERATOR_CLASSES = (DummyVisionInputGenerator, DummyVisionEncoderDecoderPastKeyValuesGenerator) _MODEL_PATCHER = OVSeq2SeqModelPatcher + @property + def inputs(self) -> dict: + common_inputs = {} + + if self._behavior in {ConfigBehavior.ENCODER, ConfigBehavior.MONOLITH}: + common_inputs["pixel_values"] = {0: "batch_size", 1: "num_channels", 2: "height", 3: "width"} + else: + common_inputs["encoder_outputs"] = {0: "batch_size", 1: "encoder_sequence_length"} + + if self._behavior in {ConfigBehavior.DECODER, ConfigBehavior.MONOLITH}: + common_inputs["decoder_input_ids"] = {0: "batch_size", 1: "decoder_sequence_length"} + if self.use_past_in_inputs: + self.add_past_key_values(common_inputs, direction="inputs") + + return common_inputs + + @property + def outputs(self) -> dict: + if self._behavior == ConfigBehavior.ENCODER: + return self._encoder_onnx_config.outputs + else: + return super().outputs + @register_in_tasks_manager("zamba2", *["text-generation", "text-generation-with-past"], library_name="transformers") class Zamba2OpenVINOConfig(MambaOpenVINOConfig): @@ -6927,8 +4644,15 @@ def inputs(self) -> Dict[str, Dict[int, str]]: @register_in_tasks_manager("audio-spectrogram-transformer", *["feature-extraction", "audio-classification"]) -class ASTOpenVINOConfig(ASTOnnxConfig): - pass +class ASTOpenVINOConfig(OnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args( + num_mel_bins="num_mel_bins", max_length="max_length", allow_new=True + ) + DUMMY_INPUT_GENERATOR_CLASSES = (ASTDummyAudioInputGenerator,) + + @property + def inputs(self) -> dict: + return {"input_values": {0: "batch_size"}} @register_in_tasks_manager( @@ -6949,12 +4673,19 @@ def __init__(self, *args, **kwargs): @register_in_tasks_manager("olmo2", *COMMON_TEXT_GENERATION_TASKS, library_name="transformers") -class Olmo2OOpenVINOConfig(Olmo2OnnxConfig): - pass +class Olmo2OOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): + DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, MistralDummyPastKeyValuesGenerator) + DUMMY_PKV_GENERATOR_CLASS = MistralDummyPastKeyValuesGenerator + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig + MIN_TRANSFORMERS_VERSION = "4.47.0" @register_in_tasks_manager("opt", *[*COMMON_TEXT_GENERATION_TASKS, "text-classification", "question-answering"]) -class OPTOpenVINOConfig(OPTOnnxConfig): +class OPTOpenVINOConfig( + TextDecoderWithPositionIdsOnnxConfig if is_transformers_version(">=", "4.46.0") else TextDecoderOnnxConfig +): + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig + def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) _warn_potential_accuracy_issue_ov_2026_1("opt") @@ -6963,8 +4694,41 @@ def __init__(self, *args, **kwargs): @register_in_tasks_manager( "gpt_bigcode", *[*COMMON_TEXT_GENERATION_TASKS, "text-classification", "token-classification"] ) -class GPTBigCodeOpenVINOConfig(GPTBigCodeOnnxConfig): - pass +class GPTBigCodeOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): + DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, GPTBigCodeDummyPastKeyValuesGenerator) + NORMALIZED_CONFIG_CLASS = NormalizedConfigManager.get_normalized_config_class("gpt_bigcode") + DUMMY_PKV_GENERATOR_CLASS = GPTBigCodeDummyPastKeyValuesGenerator + + def add_past_key_values(self, inputs_or_outputs: dict, direction: str): + if is_transformers_version(">=", "4.54"): + super().add_past_key_values(inputs_or_outputs, direction) + else: + if direction not in ["inputs", "outputs"]: + raise ValueError(f'direction must either be "inputs" or "outputs", but {direction} was given') + + if direction == "inputs": + decoder_sequence_name = "past_sequence_length" + name = "past_key_values" + else: + decoder_sequence_name = "past_sequence_length + sequence_length" + name = "present" + + if self._normalized_config.multi_query: + decoder_sequence_dim = 1 + else: + decoder_sequence_dim = 2 + + for i in range(self._normalized_config.num_layers): + inputs_or_outputs[f"{name}.{i}.key_value"] = { + 0: "batch_size", + decoder_sequence_dim: decoder_sequence_name, + } + + def flatten_past_key_values(self, flattened_output, name, idx, t): + if is_transformers_version(">=", "4.54"): + super().flatten_past_key_values(flattened_output, name, idx, t) + else: + flattened_output[f"{name}.{idx}.key_value"] = t @register_in_tasks_manager( @@ -6974,78 +4738,196 @@ class GPTBigCodeOpenVINOConfig(GPTBigCodeOnnxConfig): "image-to-text-with-past", ], ) -class Pix2StructOpenVINOConfig(Pix2StructOnnxConfig): +class _Pix2StructNormalizedConfig(NormalizedSeq2SeqConfig): + ENCODER_NUM_LAYERS = "vision_config.num_hidden_layers" + DECODER_NUM_LAYERS = "text_config.num_layers" + ENCODER_NUM_ATTENTION_HEADS = "vision_config.num_attention_heads" + DECODER_NUM_ATTENTION_HEADS = "text_config.num_heads" + HIDDEN_SIZE = "text_config.hidden_size" + VOCAB_SIZE = "text_config.vocab_size" + + +class Pix2StructOpenVINOConfig(OnnxSeq2SeqConfigWithPast): + PAD_ATTENTION_MASK_TO_PAST = True + NORMALIZED_CONFIG_CLASS = _Pix2StructNormalizedConfig + DUMMY_INPUT_GENERATOR_CLASSES = ( + DummyTextInputGenerator, + DummySeq2SeqDecoderTextInputGenerator, + DummySeq2SeqPastKeyValuesGenerator, + DummyPix2StructInputGenerator, + ) _MODEL_PATCHER = OVSeq2SeqModelPatcher + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + if is_transformers_version("==", "4.46.0"): + logger.warning( + "Found transformers v4.46.0 while trying to export a Pix2Struct model, " + "this specific version of transformers is broken for this model. Please " + "upgrade to v4.46.1 or higher, or downgrade to v4.45.x.", + ) + + @property + def inputs(self): + common_inputs = {} + + if self._behavior in {ConfigBehavior.ENCODER, ConfigBehavior.MONOLITH}: + common_inputs["flattened_patches"] = {0: "batch_size"} + else: + common_inputs["encoder_outputs"] = {0: "batch_size"} + common_inputs["attention_mask"] = {0: "batch_size"} + + if self._behavior in {ConfigBehavior.DECODER, ConfigBehavior.MONOLITH}: + common_inputs["decoder_input_ids"] = {0: "batch_size", 1: "decoder_sequence_length"} + if self.use_past_in_inputs: + self.add_past_key_values(common_inputs, direction="inputs") + decoder_attention_mask_dim = "past_decoder_sequence_length + decoder_sequence_length" + else: + decoder_attention_mask_dim = "decoder_sequence_length" + common_inputs["decoder_attention_mask"] = {0: "batch_size", 1: decoder_attention_mask_dim} + + return common_inputs + + @property + def outputs(self) -> dict: + common_outputs = super().outputs + if self._behavior in {ConfigBehavior.ENCODER, ConfigBehavior.MONOLITH}: + if self._behavior is ConfigBehavior.MONOLITH: + output_name = "encoder_last_hidden_state" + else: + output_name = "last_hidden_state" + common_outputs[output_name] = {0: "batch_size"} # Remove unnecessary dynamic axis. + + return common_outputs + + def _create_dummy_input_generator_classes(self, **kwargs): + if self._preprocessors is None or len(self._preprocessors) < 2: + raise ValueError( + f"Preprocessors for pix2struct need to be available for the ONNX export to infer input static shapes. Got: {self._preprocessors}" + ) + + dummy_inputs_generators = [] + dummy_inputs_generators.append( + self.DUMMY_INPUT_GENERATOR_CLASSES[0](self.task, self._normalized_config, **kwargs) + ) + encoder_sequence_length = self._preprocessors[1].image_processor.max_patches + kwargs["preprocessors"] = self._preprocessors + for cls_ in self.DUMMY_INPUT_GENERATOR_CLASSES[1:]: + dummy_inputs_generators.append( + cls_(self.task, self._normalized_config, encoder_sequence_length=encoder_sequence_length, **kwargs) + ) + + return dummy_inputs_generators + + def overwrite_shape_and_generate_input(self, dummy_input_gen, input_name: str, framework: str, input_shapes: dict): + if self._preprocessors is None or len(self._preprocessors) < 2: + raise ValueError( + f"Preprocessors for pix2struct need to be available for the ONNX export to infer input static shapes. Got: {self._preprocessors}" + ) + + if input_name in ["encoder_outputs", "attention_mask"]: + original_seq_length = dummy_input_gen.sequence_length + dummy_input_gen.sequence_length = self._preprocessors[1].image_processor.max_patches + dummy_input = dummy_input_gen.generate( + input_name, framework=framework, int_dtype=self.int_dtype, float_dtype=self.float_dtype + ) + dummy_input_gen.sequence_length = original_seq_length + else: + dummy_input = super().overwrite_shape_and_generate_input( + dummy_input_gen, input_name, framework, input_shapes + ) + + return dummy_input + @register_in_tasks_manager("bert", *COMMON_TEXT_TASKS) -class BertOpenVINOConfig(BertOnnxConfig): - pass +class BertOpenVINOConfig(TextEncoderOnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig + + @property + def inputs(self) -> dict: + if self.task == "multiple-choice": + dynamic_axis = {0: "batch_size", 1: "num_choices", 2: "sequence_length"} + else: + dynamic_axis = {0: "batch_size", 1: "sequence_length"} + return { + "input_ids": dynamic_axis, + "attention_mask": dynamic_axis, + "token_type_ids": dynamic_axis, + } @register_in_tasks_manager("albert", *COMMON_TEXT_TASKS) -class AlbertOpenVINOConfig(AlbertOnnxConfig): +class AlbertOpenVINOConfig(BertOpenVINOConfig): pass @register_in_tasks_manager("nystromformer", *COMMON_TEXT_TASKS) -class NystromformerOpenVINOConfig(NystromformerOnnxConfig): +class NystromformerOpenVINOConfig(BertOpenVINOConfig): pass @register_in_tasks_manager("convbert", *COMMON_TEXT_TASKS) -class ConvBertOpenVINOConfig(ConvBertOnnxConfig): +class ConvBertOpenVINOConfig(BertOpenVINOConfig): pass @register_in_tasks_manager("electra", *COMMON_TEXT_TASKS) -class ElectraOpenVINOConfig(ElectraOnnxConfig): +class ElectraOpenVINOConfig(BertOpenVINOConfig): pass @register_in_tasks_manager("roformer", *COMMON_TEXT_TASKS) -class RoFormerOpenVINOConfig(RoFormerOnnxConfig): +class RoFormerOpenVINOConfig(BertOpenVINOConfig): pass @register_in_tasks_manager("squeezebert", *COMMON_TEXT_TASKS) -class SqueezeBertOpenVINOConfig(SqueezeBertOnnxConfig): +class SqueezeBertOpenVINOConfig(BertOpenVINOConfig): pass @register_in_tasks_manager("mobilebert", *COMMON_TEXT_TASKS) -class MobileBertOpenVINOConfig(MobileBertOnnxConfig): +class MobileBertOpenVINOConfig(BertOpenVINOConfig): pass @register_in_tasks_manager("xlm", *COMMON_TEXT_TASKS) -class XLMOpenVINOConfig(XLMOnnxConfig): +class XLMOpenVINOConfig(BertOpenVINOConfig): # TODO (@echarlaix): add v5 support MAX_TRANSFORMERS_VERSION = "4.57.6" @register_in_tasks_manager("xlm-roberta", *COMMON_TEXT_TASKS) -class XLMRobertaOpenVINOConfig(XLMRobertaOnnxConfig): +class XLMRobertaOpenVINOConfig(DistilBertOpenVINOConfig): pass @register_in_tasks_manager("distilbert", *COMMON_TEXT_TASKS) -class DistilBertOpenVINOConfig(DistilBertOnnxConfig): - pass +class DistilBertOpenVINOConfig(TextEncoderOnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig + + @property + def inputs(self) -> dict: + if self.task == "multiple-choice": + dynamic_axis = {0: "batch_size", 1: "num_choices", 2: "sequence_length"} + else: + dynamic_axis = {0: "batch_size", 1: "sequence_length"} + return {"input_ids": dynamic_axis, "attention_mask": dynamic_axis} @register_in_tasks_manager("roberta", *COMMON_TEXT_TASKS) -class RobertaOpenVINOConfig(RobertaOnnxConfig): +class RobertaOpenVINOConfig(DistilBertOpenVINOConfig): pass @register_in_tasks_manager("camembert", *COMMON_TEXT_TASKS) -class CamembertOpenVINOConfig(CamembertOnnxConfig): +class CamembertOpenVINOConfig(DistilBertOpenVINOConfig): pass @register_in_tasks_manager("flaubert", *COMMON_TEXT_TASKS) -class FlaubertOpenVINOConfig(FlaubertOnnxConfig): +class FlaubertOpenVINOConfig(BertOpenVINOConfig): # TODO (@echarlaix): add v5 support MAX_TRANSFORMERS_VERSION = "4.57.6" @@ -7054,12 +4936,17 @@ class FlaubertOpenVINOConfig(FlaubertOnnxConfig): "deberta", *["feature-extraction", "fill-mask", "text-classification", "token-classification", "question-answering"], ) -class DebertaOpenVINOConfig(DebertaOnnxConfig): - pass +class DebertaOpenVINOConfig(BertOpenVINOConfig): + @property + def inputs(self) -> dict: + common_inputs = super().inputs + if self._config.type_vocab_size == 0: + common_inputs.pop("token_type_ids") + return common_inputs @register_in_tasks_manager("deberta-v2", *COMMON_TEXT_TASKS) -class DebertaV2OpenVINOConfig(DebertaV2OnnxConfig): +class DebertaV2OpenVINOConfig(DebertaOpenVINOConfig): pass @@ -7073,95 +4960,185 @@ class DebertaV2OpenVINOConfig(DebertaV2OnnxConfig): "audio-xvector", ], ) -class Data2VecAudioOpenVINOConfig(Data2VecAudioOnnxConfig): +class Data2VecAudioOpenVINOConfig(HubertOpenVINOConfig): pass @register_in_tasks_manager("data2vec-text", *COMMON_TEXT_TASKS) -class Data2VecTextOpenVINOConfig(Data2VecTextOnnxConfig): +class Data2VecTextOpenVINOConfig(DistilBertOpenVINOConfig): # TODO (@echarlaix): add v5 support MAX_TRANSFORMERS_VERSION = "4.57.6" @register_in_tasks_manager("data2vec-vision", *["feature-extraction", "image-classification"]) -class Data2VecVisionOpenVINOConfig(Data2VecVisionOnnxConfig): +class Data2VecVisionOpenVINOConfig(ViTOpenVINOConfig): pass @register_in_tasks_manager("perceiver", *["fill-mask", "text-classification", "image-classification"]) -class PerceiverOpenVINOConfig(PerceiverOnnxConfig): - pass +class PerceiverOpenVINOConfig(TextAndVisionOnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig + DUMMY_INPUT_GENERATOR_CLASSES = ( + PerceiverDummyInputGenerator, + *TextAndVisionOnnxConfig.DUMMY_INPUT_GENERATOR_CLASSES, + ) + + def __init__( + self, + config, + task: str = "feature-extraction", + int_dtype: str = "int64", + float_dtype: str = "fp32", + preprocessors=None, + ): + super().__init__( + config=config, + task=task, + int_dtype=int_dtype, + float_dtype=float_dtype, + preprocessors=preprocessors, + ) + self.is_generating_dummy_inputs = False + + @property + def inputs_name(self): + if self.is_generating_dummy_inputs: + if self.task in ["fill-mask", "text-classification"]: + return "input_ids" + else: + return "pixel_values" + else: + return "inputs" + + @property + def inputs(self) -> dict: + if self.inputs_name in ["input_ids", "inputs"]: + dynamic_axis = {0: "batch_size", 1: "sequence_length"} + return { + "input_ids": dynamic_axis, + "attention_mask": dynamic_axis, + } + else: + dynamic_axis = {0: "batch_size", 1: "sequence_length", 2: "width", 3: "height"} + return { + "pixel_values": dynamic_axis, + } + + @property + def outputs(self) -> dict: + outputs = super().outputs + + if "logits" in outputs: + outputs["logits"] = {0: "batch_size"} + + return outputs + + def generate_dummy_inputs(self, framework: str = "pt", **kwargs): + self.is_generating_dummy_inputs = True + dummy_inputs = super().generate_dummy_inputs(framework=framework, **kwargs) + dummy_inputs[self.inputs_name] = dummy_inputs.pop(self.inputs_name) + return dummy_inputs @register_in_tasks_manager("esm", *["feature-extraction", "fill-mask", "text-classification", "token-classification"]) -class EsmOpenVINOConfig(EsmOnnxConfig): - pass +class EsmOpenVINOConfig(TextEncoderOnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig + + @property + def inputs(self) -> dict: + dynamic_axis = {0: "batch_size", 1: "sequence_length"} + return { + "input_ids": dynamic_axis, + "attention_mask": dynamic_axis, + } @register_in_tasks_manager("mpnet", *COMMON_TEXT_TASKS) -class MPNetOpenVINOConfig(MPNetOnnxConfig): +class MPNetOpenVINOConfig(DistilBertOpenVINOConfig): pass @register_in_tasks_manager("beit", *["feature-extraction", "image-classification"]) -class BeitOpenVINOConfig(BeitOnnxConfig): +class BeitOpenVINOConfig(ViTOpenVINOConfig): pass @register_in_tasks_manager("deit", *["feature-extraction", "image-classification", "masked-im"]) -class DeiTOpenVINOConfig(DeiTOnnxConfig): +class DeiTOpenVINOConfig(ViTOpenVINOConfig): pass @register_in_tasks_manager("levit", *["feature-extraction", "image-classification"]) -class LevitOpenVINOConfig(LevitOnnxConfig): +class LevitOpenVINOConfig(ViTOpenVINOConfig): pass @register_in_tasks_manager("mobilevit", *["feature-extraction", "image-classification", "image-segmentation"]) -class MobileViTOpenVINOConfig(MobileViTOnnxConfig): +class MobileViTOpenVINOConfig(ViTOpenVINOConfig): pass @register_in_tasks_manager("mobilenet_v1", *["feature-extraction", "image-classification"]) -class MobileNetV1OpenVINOConfig(MobileNetV1OnnxConfig): - pass +class MobileNetV1OpenVINOConfig(ViTOpenVINOConfig): + @property + def inputs(self) -> dict: + return {"pixel_values": {0: "batch_size"}} @register_in_tasks_manager("mobilenet_v2", *["feature-extraction", "image-classification"]) -class MobileNetV2OpenVINOConfig(MobileNetV2OnnxConfig): +class MobileNetV2OpenVINOConfig(MobileNetV1OpenVINOConfig): pass @register_in_tasks_manager("poolformer", *["feature-extraction", "image-classification"]) -class PoolFormerOpenVINOConfig(PoolFormerOnnxConfig): - pass +class PoolFormerOpenVINOConfig(ViTOpenVINOConfig): + NORMALIZED_CONFIG_CLASS = NormalizedVisionConfig @register_in_tasks_manager( "segformer", *["feature-extraction", "image-classification", "image-segmentation", "semantic-segmentation"] ) -class SegformerOpenVINOConfig(SegformerOnnxConfig): - pass +class SegformerOpenVINOConfig(ViTOpenVINOConfig): + @property + def outputs(self) -> dict: + outputs = super().outputs + + if self.task == "image-segmentation": + outputs["logits"] = {0: "batch_size"} + + return outputs @register_in_tasks_manager("swin", *["feature-extraction", "image-classification", "masked-im"]) -class SwinOpenVINOConfig(SwinOnnxConfig): +class SwinOpenVINOConfig(ViTOpenVINOConfig): pass @register_in_tasks_manager("vit", *["feature-extraction", "image-classification", "masked-im"]) -class ViTOpenVINOConfig(ViTOnnxConfig): - pass +class ViTOpenVINOConfig(VisionOnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedVisionConfig + + @property + def inputs(self) -> dict: + return {"pixel_values": {0: "batch_size", 2: "height", 3: "width"}} + + @property + def outputs(self) -> dict: + common_outputs = super().outputs + + if self.task == "feature-extraction": + common_outputs["last_hidden_state"] = {0: "batch_size"} + + return common_outputs @register_in_tasks_manager("convnext", *["feature-extraction", "image-classification"]) -class ConvNextOpenVINOConfig(ConvNextOnnxConfig): +class ConvNextOpenVINOConfig(ViTOpenVINOConfig): pass @register_in_tasks_manager("resnet", *["feature-extraction", "image-classification"]) -class ResNetOpenVINOConfig(ResNetOnnxConfig): +class ResNetOpenVINOConfig(ViTOpenVINOConfig): pass @@ -7175,7 +5152,7 @@ class ResNetOpenVINOConfig(ResNetOnnxConfig): "audio-xvector", ], ) -class Wav2Vec2OpenVINOConfig(Wav2Vec2OnnxConfig): +class Wav2Vec2OpenVINOConfig(HubertOpenVINOConfig): pass @@ -7189,29 +5166,43 @@ class Wav2Vec2OpenVINOConfig(Wav2Vec2OnnxConfig): "audio-xvector", ], ) -class Wav2Vec2ConformerOpenVINOConfig(Wav2Vec2ConformerOnnxConfig): +class Wav2Vec2ConformerOpenVINOConfig(HubertOpenVINOConfig): pass @register_in_tasks_manager("hubert", *["feature-extraction", "automatic-speech-recognition", "audio-classification"]) -class HubertOpenVINOConfig(HubertOnnxConfig): - pass +class HubertOpenVINOConfig(AudioOnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedConfig + + @property + def outputs(self) -> dict: + outputs = super().outputs + + # Hubert output formula adapted from: + # https://github.com/huggingface/transformers/blob/v4.55.2/src/transformers/models/hubert/modeling_hubert.py#L721 + if self.task == "automatic-speech-recognition": + sequence_length = "sequence_length" + for kernel_size, stride in zip(self._config.conv_kernel, self._config.conv_stride): + sequence_length = f"( {sequence_length} - {kernel_size} ) // {stride} + 1" + outputs["logits"] = {0: "batch_size", 1: sequence_length} + + return outputs @register_in_tasks_manager("sew", *["feature-extraction", "automatic-speech-recognition", "audio-classification"]) -class SEWOpenVINOConfig(SEWOnnxConfig): +class SEWOpenVINOConfig(HubertOpenVINOConfig): pass @register_in_tasks_manager("sew-d", *["feature-extraction", "automatic-speech-recognition", "audio-classification"]) -class SEWDOpenVINOConfig(SEWDOnnxConfig): +class SEWDOpenVINOConfig(HubertOpenVINOConfig): pass @register_in_tasks_manager( "unispeech", *["feature-extraction", "automatic-speech-recognition", "audio-classification"] ) -class UniSpeechOpenVINOConfig(UniSpeechOnnxConfig): +class UniSpeechOpenVINOConfig(HubertOpenVINOConfig): pass @@ -7225,7 +5216,7 @@ class UniSpeechOpenVINOConfig(UniSpeechOnnxConfig): "audio-xvector", ], ) -class UniSpeechSATOpenVINOConfig(UniSpeechSATOnnxConfig): +class UniSpeechSATOpenVINOConfig(HubertOpenVINOConfig): pass @@ -7239,40 +5230,171 @@ class UniSpeechSATOpenVINOConfig(UniSpeechSATOnnxConfig): "audio-xvector", ], ) -class WavLMOpenVINOConfig(WavLMOnnxConfig): +class WavLMOpenVINOConfig(HubertOpenVINOConfig): pass @register_in_tasks_manager("sam", *["feature-extraction"]) -class SamOpenVINOConfig(SamOnnxConfig): - pass +class SamOpenVINOConfig(OnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedEncoderDecoderConfig + DUMMY_INPUT_GENERATOR_CLASSES = (DummyVisionInputGenerator, DummyPointsGenerator, DummyVisionEmbeddingsGenerator) + VARIANTS = { # noqa: RUF012 + "monolith": "All the SAM model components are exported as a single model.onnx.", + "split": "The vision encoder is exported as a separate vision_encoder.onnx, and the prompt encoder and mask decoder are exported as a prompt_encoder_mask_decoder.onnx. This allows to encoder the image only once for multiple point queries.", + } + DEFAULT_VARIANT = "split" + _MODEL_PATCHER = SAMModelPatcher + + def __init__( + self, + config, + task: str = "feature-extraction", + int_dtype: str = "int64", + float_dtype: str = "fp32", + variant: str = "split", + vision_encoder=None, + preprocessors=None, + ): + super().__init__( + config=config, + task=task, + int_dtype=int_dtype, + float_dtype=float_dtype, + preprocessors=preprocessors, + ) + self.variant = variant + self.vision_encoder = vision_encoder + self._normalized_config.ENCODER_NORMALIZED_CONFIG_CLASS = NormalizedVisionConfig(self._config.vision_config) + + @property + def inputs(self) -> dict: + if self.variant == "monolith": + inputs = { + "pixel_values": {0: "batch_size"}, + "input_points": {0: "batch_size", 1: "point_batch_size", 2: "nb_points_per_image"}, + "input_labels": {0: "batch_size", 1: "point_batch_size", 2: "nb_points_per_image"}, + } + else: + if self.vision_encoder: + inputs = {"pixel_values": {0: "batch_size"}} + else: + inputs = { + "image_positional_embeddings": {0: "batch_size"}, + "image_embeddings": {0: "batch_size"}, + "input_points": {0: "batch_size", 1: "point_batch_size", 2: "nb_points_per_image"}, + "input_labels": {0: "batch_size", 1: "point_batch_size", 2: "nb_points_per_image"}, + } + return inputs + + @property + def outputs(self) -> dict: + if self.variant == "split" and self.vision_encoder: + return {"image_embeddings": {0: "batch_size"}, "image_positional_embeddings": {0: "batch_size"}} + else: + return { + "iou_scores": {0: "batch_size", 1: "point_batch_size"}, + "pred_masks": {0: "batch_size", 1: "point_batch_size"}, + } @register_in_tasks_manager("siglip", *["feature-extraction", "zero-shot-image-classification"]) -class SiglipOpenVINOConfig(SiglipOnnxConfig): - pass +class SiglipOpenVINOConfig(TextAndVisionOnnxConfig): + NORMALIZED_CONFIG_CLASS = SiglipNormalizedConfig + _MODEL_PATCHER = CLIPModelPatcher + + @property + def inputs(self) -> dict: + return { + "input_ids": {0: "text_batch_size", 1: "sequence_length"}, + "pixel_values": {0: "image_batch_size", 1: "num_channels", 2: "height", 3: "width"}, + # NOTE: No attention_mask + } + + @property + def outputs(self) -> dict: + if self.task in ["feature-extraction", "zero-shot-image-classification"]: + return { + "logits_per_image": {0: "image_batch_size", 1: "text_batch_size"}, + "logits_per_text": {0: "text_batch_size", 1: "image_batch_size"}, + "text_embeds": {0: "text_batch_size"}, + "image_embeds": {0: "image_batch_size"}, + } + else: + return super().outputs @register_in_tasks_manager( "transformer", *["feature-extraction", "sentence-similarity"], library_name="sentence_transformers" ) -class SentenceTransformersTransformerOpenVINOConfig(SentenceTransformersTransformerOnnxConfig): - pass +class SentenceTransformersTransformerOpenVINOConfig(TextEncoderOnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig + _MODEL_PATCHER = SentenceTransformersTransformerPatcher + + @property + def inputs(self) -> dict: + return { + "input_ids": {0: "batch_size", 1: "sequence_length"}, + "attention_mask": {0: "batch_size", 1: "sequence_length"}, + } + + @property + def outputs(self) -> dict: + return { + "token_embeddings": {0: "batch_size", 1: "sequence_length"}, + "sentence_embedding": {0: "batch_size"}, + } @register_in_tasks_manager("rembert", *COMMON_TEXT_TASKS) -class RemBertOpenVINOConfig(RemBertOnnxConfig): +class RemBertOpenVINOConfig(BertOpenVINOConfig): pass @register_in_tasks_manager("siglip-text-with-projection", *["feature-extraction"]) -class SiglipTextWithProjectionOpenVINOConfig(SiglipTextWithProjectionOnnxConfig): - pass +class SiglipTextWithProjectionOpenVINOConfig(TextEncoderOnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args( + vocab_size="vocab_size", + sequence_length="max_position_embeddings", + num_layers="num_hidden_layers", + allow_new=True, + ) + _MODEL_PATCHER = CLIPModelPatcher + + @property + def inputs(self) -> dict: + return { + "input_ids": {0: "batch_size", 1: "sequence_length"}, + } + + @property + def outputs(self) -> dict: + common_outputs = { + "text_embeds": {0: "batch_size", 1: "sequence_length"}, + "last_hidden_state": {0: "batch_size", 1: "sequence_length"}, + } + if self._normalized_config.output_hidden_states: + for i in range(self._normalized_config.num_layers + 1): + common_outputs[f"hidden_states.{i}"] = {0: "batch_size", 1: "sequence_length"} + + return common_outputs @register_in_tasks_manager("siglip-text", *["feature-extraction"]) -class SiglipTextOpenVINOConfig(SiglipTextOnnxConfig): - pass +class SiglipTextOpenVINOConfig(SiglipTextWithProjectionOpenVINOConfig): + _MODEL_PATCHER = CLIPModelPatcher + + @property + def outputs(self) -> dict: + common_outputs = { + "last_hidden_state": {0: "batch_size", 1: "sequence_length"}, + "pooler_output": {0: "batch_size"}, + } + + if self._normalized_config.output_hidden_states: + for i in range(self._normalized_config.num_layers + 1): + common_outputs[f"hidden_states.{i}"] = {0: "batch_size", 1: "sequence_length"} + + return common_outputs class VideoChatFlashQwenProjectorOpenVINOConfig(OnnxConfig): diff --git a/optimum/exporters/openvino/model_patcher.py b/optimum/exporters/openvino/model_patcher.py index 22e3cc0682..1ae4608a1d 100644 --- a/optimum/exporters/openvino/model_patcher.py +++ b/optimum/exporters/openvino/model_patcher.py @@ -12,18 +12,22 @@ # See the License for the specific language governing permissions and # limitations under the License. +import dataclasses import functools import inspect import logging import logging as log import math +import sys import types from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable, Dict, List, Optional, Tuple, Union import torch import torch.nn.functional as F +import transformers from torch import nn +from torch.onnx import symbolic_helper from transformers import PreTrainedModel from transformers.cache_utils import Cache, DynamicCache, EncoderDecoderCache from transformers.configuration_utils import PretrainedConfig @@ -33,7 +37,6 @@ BaseModelOutputWithPast, BaseModelOutputWithPooling, ) -from transformers.modeling_utils import PreTrainedModel from transformers.models.llama.configuration_llama import LlamaConfig from transformers.models.llama.modeling_llama import ( LlamaAttention, @@ -47,7 +50,8 @@ from transformers.processing_utils import Unpack from transformers.utils import ModelOutput -# from optimum.exporters.openvino.base import OnnxConfig +from optimum.exporters.openvino._ov_ops import convert_recurrent_attention_cell +from optimum.exporters.openvino.base import OnnxConfig from optimum.intel.utils.import_utils import ( is_diffusers_version, is_openvino_version, @@ -55,49 +59,14 @@ is_transformers_version, ) -from ._ov_ops import convert_recurrent_attention_cell - - -if is_transformers_version(">=", "4.53"): - from transformers.masking_utils import ALL_MASK_ATTENTION_FUNCTIONS, eager_mask, sdpa_mask - from transformers.models.qwen3_moe.modeling_qwen3_moe import Qwen3MoeSparseMoeBlock -if is_transformers_version(">=", "4.54"): - from transformers.masking_utils import create_causal_mask -if is_transformers_version(">=", "4.56"): - import transformers.masking_utils -if is_transformers_version(">=", "4.57"): - from transformers.models.qwen3_vl.modeling_qwen3_vl import Qwen3VLTextRotaryEmbedding -if is_transformers_version(">=", "5"): - from transformers.modeling_rope_utils import RotaryEmbeddingConfigMixin - -if TYPE_CHECKING: - from transformers.cache_utils import Cache - from transformers.modeling_utils import PreTrainedModel - -if is_transformers_version(">=", "4.54"): - from transformers.utils import TransformersKwargs -else: - TransformersKwargs = object - - -####### - -import dataclasses -import sys -from typing import TYPE_CHECKING - -import transformers -from torch.onnx import symbolic_helper - -from optimum.utils import is_diffusers_version, is_torch_version, is_transformers_version +if is_transformers_version(">=", "4.43") and is_transformers_version("<", "4.48"): + from transformers.models.clip.modeling_clip import CLIPAttention, CLIPSdpaAttention if is_transformers_version(">=", "4.44") and is_transformers_version("<", "4.50"): from optimum.exporters.openvino._traceable_cache import TraceableCache -if is_transformers_version(">=", "4.54"): - from optimum.exporters.openvino._traceable_decorator import traceable_check_model_inputs -if is_transformers_version(">=", "4.43") and is_transformers_version("<", "4.48"): - from transformers.models.clip.modeling_clip import CLIPAttention, CLIPSdpaAttention + + if is_transformers_version(">=", "4.48"): from transformers.cache_utils import DynamicCache, EncoderDecoderCache from transformers.models.moonshine.modeling_moonshine import MoonshinePreTrainedModel @@ -113,15 +82,35 @@ sdpa_mask, ) from transformers.models.qwen3_moe.modeling_qwen3_moe import Qwen3MoeSparseMoeBlock + + + if is_transformers_version(">=", "4.53.1"): from transformers.masking_utils import find_packed_sequence_indices + + +if is_transformers_version(">=", "4.54"): + from transformers.masking_utils import create_causal_mask + from transformers.utils import TransformersKwargs + + from optimum.exporters.openvino._traceable_decorator import traceable_check_model_inputs + +else: + TransformersKwargs = object + if is_transformers_version(">=", "4.55"): from transformers.models.gpt_oss.modeling_gpt_oss import GptOssExperts if is_transformers_version(">=", "4.56"): + import transformers.masking_utils from transformers.cache_utils import DynamicLayer -if is_diffusers_version(">=", "0.35.0"): - import diffusers.models.transformers.transformer_flux + +if is_transformers_version(">=", "4.57"): + from transformers.models.qwen3_vl.modeling_qwen3_vl import Qwen3VLTextRotaryEmbedding + + +if is_transformers_version(">=", "5"): + from transformers.modeling_rope_utils import RotaryEmbeddingConfigMixin logger = logging.getLogger(__name__) @@ -969,108 +958,6 @@ def patched_speecht5_prenet_forward( return inputs_embeds -class SpeechT5ModelPatcher(ModelPatcher): - def __enter__(self): - super().__enter__() - - self.original_speecht5_prenet_forward = self._model.speecht5.decoder.prenet.forward - self._model.speecht5.decoder.prenet.forward = types.MethodType( - patched_speecht5_prenet_forward, self._model.speecht5.decoder.prenet - ) - - def __exit__(self, exc_type, exc_value, traceback): - super().__exit__(exc_type, exc_value, traceback) - - self._model.speecht5.decoder.prenet.forward = types.MethodType( - self.original_speecht5_prenet_forward, self._model.speecht5.decoder.prenet - ) - - def __init__( - self, - config: "OnnxConfig", - model: PreTrainedModel, - model_kwargs: dict[str, Any], - ): - super().__init__(config, model, model_kwargs) - model.vocoder = model_kwargs["vocoder_model"].eval() - - def patched_forward( - input_ids=None, - speaker_embeddings=None, - encoder_outputs=None, - past_key_values=None, - output_sequence=None, - spectrogram=None, - encoder_attention_mask=None, - ): - if past_key_values is not None: - past_key_values = preprocess_past_key_values(past_key_values) - - if self.real_config._behavior == "encoder": - encoder_attention_mask = torch.ones_like(input_ids) - encoder_out = model.speecht5.encoder(input_values=input_ids, attention_mask=encoder_attention_mask) - # downsample encoder attention mask - if isinstance(model.speecht5.encoder, SpeechT5EncoderWithSpeechPrenet): - encoder_attention_mask = model.speecht5.encoder.prenet._get_feature_vector_attention_mask( - encoder_out[0].shape[1], encoder_attention_mask - ) - outputs = { - "encoder_outputs": encoder_out.last_hidden_state, - "encoder_attention_mask": encoder_attention_mask, - } - elif self.real_config._behavior == "decoder": - use_cache = self.real_config.use_past - encoder_hidden_states = encoder_outputs[0] - decoder_hidden_states = model.speecht5.decoder.prenet(output_sequence, speaker_embeddings) - # Run the decoder layers on the last element of the prenet output. - decoder_out = model.speecht5.decoder.wrapped_decoder( - hidden_states=decoder_hidden_states[:, -1:], - encoder_hidden_states=encoder_hidden_states, - encoder_attention_mask=encoder_attention_mask, - past_key_values=past_key_values, - output_attentions=False, - use_cache=use_cache, - ) - last_decoder_output = decoder_out.last_hidden_state[0, -1] - past_key_values = decoder_out.past_key_values - # Predict the new mel spectrum for this step in the sequence. - spectrum = model.speech_decoder_postnet.feat_out(last_decoder_output) - spectrum = spectrum.view(model.config.reduction_factor, model.config.num_mel_bins) - # NOTE: extending the spectrogram should is to be handled outside of the ONNX. - # spectrogram.append(spectrum) - # Extend the output sequence with the new mel spectrum. - output_sequence = torch.cat( - (output_sequence, spectrum[-1].view(1, 1, model.config.num_mel_bins)), dim=1 - ) - # Predict the probability that this is the stop token. - prob = torch.sigmoid(model.speech_decoder_postnet.prob_out(last_decoder_output)) - outputs = { - "output_sequence_out": output_sequence, - "spectrum": spectrum, - "prob": prob, - "past_key_values": past_key_values, - } - elif self.real_config.is_postnet_and_vocoder: - # NOTE: the following concatenation is expected to be handled outside of the ONNX: - # spectrogram = torch.cat(spectrogram, dim=0).unsqueeze(0) - spectrogram = spectrogram.unsqueeze(0) - spectrogram = model.speech_decoder_postnet.postnet(spectrogram) - spectrogram = spectrogram.squeeze(0) - waveform = model.vocoder(spectrogram) - outputs = {"waveform": waveform} - else: - raise ValueError("Should not happen") - - if outputs.get("past_key_values") is not None: - outputs["past_key_values"] = postprocess_past_key_values( - outputs["past_key_values"], output_names=list(config.outputs.keys()) - ) - - return outputs - - self.patched_forward = patched_forward - - class SentenceTransformersTransformerPatcher(ModelPatcher): def __init__( self, @@ -1447,22 +1334,6 @@ def patched_apply_rotary_emb( return x_out.type_as(x) -class FluxTransformerModelPatcher(ModelPatcher): - def __enter__(self): - super().__enter__() - - if is_diffusers_version(">=", "0.35.0"): - self.original_apply_rotary_emb = diffusers.models.transformers.transformer_flux.apply_rotary_emb - diffusers.models.transformers.transformer_flux.apply_rotary_emb = patched_apply_rotary_emb - - def __exit__(self, exc_type, exc_value, traceback): - super().__exit__(exc_type, exc_value, traceback) - - if is_diffusers_version(">=", "0.35.0"): - diffusers.models.transformers.transformer_flux.apply_rotary_emb = self.original_apply_rotary_emb - del self.original_apply_rotary_emb - - def patched_cohere_rotary_forward(self, x, position_ids): # Get batch size and sequence length for manual expansion batch_size, _ = position_ids.shape[:2] @@ -5006,7 +4877,7 @@ def rope(pos: torch.Tensor, dim: int, theta: int) -> torch.Tensor: return emb.unsqueeze(1) -class FluxTransfromerModelPatcher(ModelPatcher): +class FluxTransformerModelPatcher(ModelPatcher): def __enter__(self): super().__enter__() if is_diffusers_version("<", "0.31.0"): @@ -7591,7 +7462,7 @@ def speecht5_decoder_layer_forward( return outputs -class OVSpeechT5ModelPatcher(ModelPatcher): +class SpeechT5ModelPatcher(ModelPatcher): def __enter__(self): if self.real_config._behavior != "vocoder": super().__enter__() @@ -11346,54 +11217,3 @@ def __enter__(self): def __exit__(self, exc_type, exc_value, traceback): super().__exit__(exc_type, exc_value, traceback) self._model.forward = self._model._orig_forward - - -################################################################################################################################################################ - - -from transformers.models.vit.modeling_vit import ViTPatchEmbeddings - - -def patched_vit_patch_embedding_forward( - self, pixel_values: torch.Tensor, interpolate_pos_encoding: bool = False -) -> torch.Tensor: - # _batch_size, num_channels, height, width = pixel_values.shape - # Unexpted error. - # TypeError: cond must be a bool, but got ? - # torch._check( - # num_channels == self.num_channels, - # lambda: ( - # "Make sure that the channel dimension of the pixel values match with the one set in the configuration." - # f" Expected {self.num_channels} but got {num_channels}." - # ) - # ) - # This check fails if dynamic shapes are not properly set up. - # Let's drop it. - # if not interpolate_pos_encoding: - # torch._check( - # height == self.image_size[0] and width == self.image_size[1], - # lambda:( - # f"Input image size ({height}*{width}) doesn't match model" - # f" ({self.image_size[0]}*{self.image_size[1]})." - # ) - # ) - embeddings = self.projection(pixel_values).flatten(2).transpose(1, 2) - return embeddings - - -class ViTForImageClassificationPatcher(ModelPatcher): - def __enter__(self): - super().__enter__() - - if is_transformers_version(">=", "4.36.0"): - self.original_forward = ViTPatchEmbeddings.forward - ViTPatchEmbeddings.forward = patched_vit_patch_embedding_forward - - def __exit__(self, exc_type, exc_value, traceback): - super().__exit__(exc_type, exc_value, traceback) - - if is_transformers_version(">=", "4.36.0"): - ViTPatchEmbeddings.forward = self.original_forward - - -################################################################################################################################################################ From 6f235d581c3cb2f70d5a26a1edea066251daee8d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CElla?= Date: Wed, 27 May 2026 15:14:57 +0200 Subject: [PATCH 06/40] remove unused --- .../exporters/openvino/_traceable_cache.py | 2 +- .../openvino/_traceable_decorator.py | 2 +- optimum/exporters/openvino/base.py | 2 +- .../exporters/openvino/input_generators.py | 37 -- optimum/exporters/openvino/model_configs.py | 2 + optimum/exporters/openvino/model_patcher.py | 426 ------------------ 6 files changed, 5 insertions(+), 466 deletions(-) diff --git a/optimum/exporters/openvino/_traceable_cache.py b/optimum/exporters/openvino/_traceable_cache.py index 059463f9cf..cc2486713e 100644 --- a/optimum/exporters/openvino/_traceable_cache.py +++ b/optimum/exporters/openvino/_traceable_cache.py @@ -1,4 +1,4 @@ -# Copyright 2025 The HuggingFace Team. All rights reserved. +# Copyright 2026 The HuggingFace Team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/optimum/exporters/openvino/_traceable_decorator.py b/optimum/exporters/openvino/_traceable_decorator.py index 0a2087dff7..a60025c1eb 100644 --- a/optimum/exporters/openvino/_traceable_decorator.py +++ b/optimum/exporters/openvino/_traceable_decorator.py @@ -1,4 +1,4 @@ -# Copyright 2025 The HuggingFace Team. All rights reserved. +# Copyright 2026 The HuggingFace Team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/optimum/exporters/openvino/base.py b/optimum/exporters/openvino/base.py index d232785c1e..1120989b3c 100644 --- a/optimum/exporters/openvino/base.py +++ b/optimum/exporters/openvino/base.py @@ -1,4 +1,4 @@ -# Copyright 2022 The HuggingFace Team. All rights reserved. +# Copyright 2026 The HuggingFace Team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/optimum/exporters/openvino/input_generators.py b/optimum/exporters/openvino/input_generators.py index 82993cb5a9..2a9b75b29a 100644 --- a/optimum/exporters/openvino/input_generators.py +++ b/optimum/exporters/openvino/input_generators.py @@ -87,43 +87,6 @@ def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int return pkv -class DummyMoonshineAudioInputGenerator(DummyAudioInputGenerator): - SUPPORTED_INPUT_NAMES = ("input_values", "attention_mask") - - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - if input_name == "input_values": # raw waveform - return self.random_float_tensor( - shape=[self.batch_size, self.sequence_length], - min_value=-1, - max_value=1, - framework=framework, - dtype=float_dtype, - ) - elif input_name == "attention_mask": # attention mask - return self.random_mask_tensor( - shape=[self.batch_size, self.sequence_length], - framework=framework, - dtype=int_dtype, - ) - else: - raise ValueError(f"Unsupported input name: {input_name}") - - -class DummySanaTransforemerTextInputGenerator(DummyTransformerTextInputGenerator): - SUPPORTED_INPUT_NAMES = ("encoder_hidden_states", "encoder_attention_mask") - - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - if input_name == "encoder_attention_mask": - return self.random_mask_tensor( - shape=[self.batch_size, self.sequence_length], - framework=framework, - dtype=int_dtype, - ) - else: - return super().generate( - input_name=input_name, framework=framework, int_dtype=int_dtype, float_dtype=float_dtype - ) - class DummyQwen3VLLMInputGenerator(DummyTextInputGenerator): SUPPORTED_INPUT_NAMES = ( diff --git a/optimum/exporters/openvino/model_configs.py b/optimum/exporters/openvino/model_configs.py index 379ce6353e..416d8af472 100644 --- a/optimum/exporters/openvino/model_configs.py +++ b/optimum/exporters/openvino/model_configs.py @@ -333,6 +333,8 @@ def init_model_configs(): TasksManager._DIFFUSERS_TASKS_TO_MODEL_MAPPINGS["text-to-video"] = {} TasksManager._DIFFUSERS_TASKS_TO_MODEL_MAPPINGS["text-to-video"]["ltx-video"] = "LTXPipeline" + # TODO: add warning to state that architectures from ONNX_SUPPORTED_ARCHITECTURES are deprecated + supported_model_types = [ "_SUPPORTED_MODEL_TYPE", "_DIFFUSERS_SUPPORTED_MODEL_TYPE", diff --git a/optimum/exporters/openvino/model_patcher.py b/optimum/exporters/openvino/model_patcher.py index 1ae4608a1d..60244b9ac7 100644 --- a/optimum/exporters/openvino/model_patcher.py +++ b/optimum/exporters/openvino/model_patcher.py @@ -817,31 +817,6 @@ def __exit__(self, exc_type, exc_value, traceback): self._model.set_attention_type("block_sparse") -class MgpstrModelPatcher(ModelPatcher): - def __init__( - self, - config: "OnnxConfig", - model: PreTrainedModel, - model_kwargs: dict[str, Any] | None = None, - ): - super().__init__(config, model, model_kwargs) - - @functools.wraps(self.orig_forward) - def patched_forward(*args, **kwargs): - signature = inspect.signature(self.orig_forward) - args, kwargs = override_arguments(args, kwargs, signature, model_kwargs=self.model_kwargs) - - # logits is a tuple, so we unpack it and return them as separate outputs - char_logits, bpe_logits, wp_logits = self.orig_forward(*args, **kwargs).logits - - return { - "char_logits": char_logits, - "bpe_logits": bpe_logits, - "wp_logits": wp_logits, - } - - self.patched_forward = patched_forward - class SAMModelPatcher(ModelPatcher): def __init__( @@ -982,191 +957,7 @@ def patched_forward(input_ids, attention_mask): self.patched_forward = patched_forward -class SentenceTransformersCLIPPatcher(ModelPatcher): - def __init__( - self, - config: "OnnxConfig", - model: PreTrainedModel, - model_kwargs: dict[str, Any], - ): - super().__init__(config, model, model_kwargs) - - def patched_forward(input_ids, attention_mask, pixel_values): - vision_outputs = model[0].model.vision_model(pixel_values=pixel_values) - image_embeds = model[0].model.visual_projection(vision_outputs[1]) - - text_outputs = model[0].model.text_model(input_ids=input_ids, attention_mask=attention_mask) - text_embeds = model[0].model.text_projection(text_outputs[1]) - - if len(model) > 1: - image_embeds = model[1:](image_embeds) - text_embeds = model[1:](text_embeds) - - return {"text_embeds": text_embeds, "image_embeds": image_embeds} - - self.patched_forward = patched_forward - - -# Triu with possible dynamic `diagonal` argument. Not possible with torch.triu unfortunately. -def triu_onnx(x, diagonal=0): - l, w = x.shape - arrange_rows = torch.arange(l, device=x.device) - - arrange_cols = torch.arange(w, device=x.device) - mask = arrange_cols.expand(l, w) - - arrange_rows = arrange_rows[:, None] + diagonal - mask = mask >= arrange_rows - return x.masked_fill(mask == 0, 0) - - -def patched_build_delay_pattern_mask(self, input_ids: torch.Tensor, pad_token_id: int, max_length: int | None = None): - # (bsz * num_codebooks, seq_len) -> (bsz, num_codebooks, seq_len) - input_ids = input_ids.reshape(-1, self.num_codebooks, input_ids.shape[-1]) - bsz, num_codebooks, seq_len = input_ids.shape - - max_length = max_length if max_length is not None else self.generation_config.max_length - input_ids_shifted = torch.ones((bsz, num_codebooks, max_length), dtype=torch.long, device=input_ids.device) * -1 - - channel_codebooks = num_codebooks // 2 if self.config.audio_channels == 2 else num_codebooks - # we only apply the mask if we have a large enough seq len - otherwise we return as is - if max_length < 2 * channel_codebooks - 1: - raise NotImplementedError("Not supported in ONNX export. Please open an issue in Optimum repository.") - - # fill the shifted ids with the prompt entries, offset by the codebook idx - for codebook in range(channel_codebooks): - if self.config.audio_channels == 1: - # mono channel - loop over the codebooks one-by-one - input_ids_shifted[:, codebook, codebook : seq_len + codebook] = input_ids[:, codebook] - else: - # left/right channels are interleaved in the generated codebooks, so handle one then the other - input_ids_shifted[:, 2 * codebook, codebook : seq_len + codebook] = input_ids[:, 2 * codebook] - input_ids_shifted[:, 2 * codebook + 1, codebook : seq_len + codebook] = input_ids[:, 2 * codebook + 1] - - # construct a pattern mask that indicates the positions of padding tokens for each codebook - # first fill the upper triangular part (the EOS padding) - # NOTE: We could use torch.bool here, but PyTorch the complains with `The exported ONNX model failed ONNX shape inference.` - # Using int8 leads to `Could not find an implementation for Where` - delay_pattern = triu_onnx( - torch.ones((channel_codebooks, max_length), dtype=torch.int32), diagonal=max_length - channel_codebooks + 1 - ) - - # NOTE: We could use torch.bool here, but PyTorch the complains with `The exported ONNX model failed ONNX shape inference.` - # Using int32 leads to `Could not find an implementation for Trilu`, hence int64 here - - # then fill the lower triangular part (the BOS padding) - delay_pattern = delay_pattern + torch.tril(torch.ones((channel_codebooks, max_length), dtype=torch.int64)) - delay_pattern = delay_pattern.to(torch.bool) - - if self.config.audio_channels == 2: - # for left/right channel we need to duplicate every row of the pattern mask in an interleaved fashion - delay_pattern = delay_pattern.repeat_interleave(2, dim=0) - - mask = ~delay_pattern.to(input_ids.device) - input_ids = mask * input_ids_shifted + ~mask * pad_token_id - - # find the first position to start generating - this is the first place we have the -1 token - # and will always be in the first codebook (since it has no codebook offset) - first_codebook_ids = input_ids[:, 0, :] - start_ids = (first_codebook_ids == -1).nonzero()[:, 1] - - # TODO: Is this OK? - first_start_id = start_ids.min() - - # (bsz * num_codebooks, seq_len) -> (bsz, num_codebooks, seq_len) - pattern_mask = input_ids.reshape(bsz * num_codebooks, -1) - input_ids_edited = input_ids[..., :first_start_id].reshape(bsz * num_codebooks, -1) - return {"input_ids_edited": input_ids_edited, "delay_pattern_mask": pattern_mask} - - -class MusicgenModelPatcher(ModelPatcher): - def __enter__(self): - self.patch_ops() - if self.real_config.model_part == "build_delay_pattern_mask": - # For build_delay_pattern_mask, we need to override the signature too. - self._model.forward = types.MethodType(patched_build_delay_pattern_mask, self._model) - else: - setattr(self._model, self.orig_forward_name, self.patched_forward) - - def __exit__(self, exc_type, exc_value, traceback): - self.restore_ops() - if self.real_config.model_part == "build_delay_pattern_mask": - self._model.forward = self.original_decoder_forward - else: - setattr(self._model, self.orig_forward_name, self.orig_forward) - - def __init__( - self, - config: "OnnxConfig", - model: PreTrainedModel, - model_kwargs: dict[str, Any] | None = None, - ): - super().__init__(config, model, model_kwargs) - - if config.model_part == "build_delay_pattern_mask": - self.original_decoder_forward = self.orig_forward - elif config.model_part == "encodec_decode": - # EncodecModel.forward -> EncodecModel.decode - @functools.wraps(self.orig_forward) - def patched_forward( - input_values: torch.Tensor | None = None, - padding_mask: torch.Tensor | None = None, - audio_codes: torch.Tensor | None = None, - bandwidth: float | None = None, - audio_scales: torch.Tensor | None = None, - return_dict: bool | None = None, - ): - chunk_length = self.real_config._config.audio_encoder.chunk_length - if chunk_length is None: - if audio_scales is not None: - audio_scales = audio_scales[0] - - if len(audio_codes) != 1: - raise ValueError(f"Expected one frame, got {len(audio_codes)}") - audio_values = self._model._decode_frame(audio_codes[0], audio_scales) - else: - raise ValueError("Not supported, a meaningful error should have been raised ahead.") - decoded_frames = [] - - for frame, scale in zip(audio_codes, audio_scales): - frames = self._model._decode_frame(frame, scale) - decoded_frames.append(frames) - - audio_values = self._model._linear_overlap_add(decoded_frames, self.config.chunk_stride or 1) - # truncate based on padding mask - if padding_mask is not None and padding_mask.shape[-1] < audio_values.shape[-1]: - audio_values = audio_values[..., : padding_mask.shape[-1]] - - return {"audio_values": audio_values} - - self.patched_forward = patched_forward - - -class MetaCLIP2Patcher(ModelPatcher): - def __init__( - self, - config: "OnnxConfig", - model: PreTrainedModel, - model_kwargs: dict[str, Any] | None = None, - ): - super().__init__(config, model, model_kwargs) - - def patched_forward(input_ids=None, pixel_values=None, attention_mask=None): - if config.variant == "monolith": - return self.orig_forward(input_ids=input_ids, pixel_values=pixel_values, attention_mask=attention_mask) - - if config.variant == "split": - if config.vision_model: - image_embeds = model.get_image_features(pixel_values) - return {"image_embeds": image_embeds} - - text_embeds = model.get_text_features(input_ids, attention_mask) - return { - "text_embeds": text_embeds, - } - - self.patched_forward = patched_forward class CLIPModelPatcher(ModelPatcher): @@ -1182,62 +973,6 @@ def __exit__(self, exc_type, exc_value, traceback): CLIPSdpaAttention.forward = self.original_sdpa_forward -class VitPoseModelPatcher(ModelPatcher): - def __init__( - self, - config: "OnnxConfig", - model: PreTrainedModel, - model_kwargs: dict[str, Any] | None = None, - ): - # Set dataset_index (defaulting to COCO=0), otherwise we will get an error like: - # ValueError: dataset_index must be provided when using multiple experts (num_experts=6). Please provide dataset_index to the forward pass. - if model.config.backbone_config.num_experts > 1: - model_kwargs["dataset_index"] = torch.tensor(0, device=model.device) - - super().__init__(config, model, model_kwargs) - - -# https://github.com/huggingface/transformers/blob/v4.53.0/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py#L228 -def qwen3_moe_forward_patched(self, hidden_states: torch.Tensor) -> torch.Tensor: - batch_size, sequence_length, hidden_dim = hidden_states.shape - hidden_states = hidden_states.view(-1, hidden_dim) - # router_logits: (batch * sequence_length, n_experts) - router_logits = self.gate(hidden_states) - - routing_weights = torch.nn.functional.softmax(router_logits, dim=1, dtype=torch.float) - routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1) - if self.norm_topk_prob: # only diff with mixtral sparse moe block! - routing_weights /= routing_weights.sum(dim=-1, keepdim=True) - # we cast back to the input dtype - routing_weights = routing_weights.to(hidden_states.dtype) - - final_hidden_states = torch.zeros( - (batch_size * sequence_length, hidden_dim), dtype=hidden_states.dtype, device=hidden_states.device - ) - - # One hot encode the selected experts to create an expert mask - # this will be used to easily index which expert is going to be sollicitated - expert_mask = torch.nn.functional.one_hot(selected_experts, num_classes=self.num_experts).permute(2, 1, 0) - - # TODO: we loop over all possible experts instead of hit ones to avoid issues in graph execution. - # expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero() - # Loop over all available experts in the model and perform the computation on each expert - for expert_idx in range(self.num_experts): - expert_layer = self.experts[expert_idx] - idx, top_x = torch.where(expert_mask[expert_idx].squeeze(0)) - - # Index the correct hidden states and compute the expert hidden state for - # the current expert. We need to make sure to multiply the output hidden - # states by `routing_weights` on the corresponding tokens (top-1 and top-2) - current_state = hidden_states[None, top_x].reshape(-1, hidden_dim) - current_hidden_states = expert_layer(current_state) * routing_weights[top_x, idx, None] - - # However `index_add_` only support torch tensors for indexing so we'll use - # the `top_x` tensor here. - final_hidden_states.index_add_(0, top_x, current_hidden_states.to(hidden_states.dtype)) - final_hidden_states = final_hidden_states.reshape(batch_size, sequence_length, hidden_dim) - return final_hidden_states, router_logits - class Qwen3MoeModelPatcher(ModelPatcher): def __enter__(self): @@ -1267,134 +1002,6 @@ def _get_feat_extract_output_lengths_patched(self, input_lengths: torch.LongTens return output_conv3_length -class MoonshineModelPatcher(ModelPatcher): - def __enter__(self): - super().__enter__() - - if is_transformers_version(">=", "4.48"): - self.original_feat_extract_output_lengths = MoonshinePreTrainedModel._get_feat_extract_output_lengths - MoonshinePreTrainedModel._get_feat_extract_output_lengths = _get_feat_extract_output_lengths_patched - - def __exit__(self, exc_type, exc_value, traceback): - super().__exit__(exc_type, exc_value, traceback) - - if is_transformers_version(">=", "4.48"): - MoonshinePreTrainedModel._get_feat_extract_output_lengths = self.original_feat_extract_output_lengths - del self.original_feat_extract_output_lengths - - -# This is a traceabe of the original function, -# the original results in a constant shape due to the use of *x.shape[:-1] -def patched_apply_rotary_emb( - x: torch.Tensor, - freqs_cis: torch.Tensor | tuple[torch.Tensor], - use_real: bool = True, - use_real_unbind_dim: int = -1, - sequence_dim: int = 2, -): - if use_real: - cos, sin = freqs_cis # [S, D] - if sequence_dim == 2: - cos = cos[None, None, :, :] - sin = sin[None, None, :, :] - elif sequence_dim == 1: - cos = cos[None, :, None, :] - sin = sin[None, :, None, :] - else: - raise ValueError(f"`sequence_dim={sequence_dim}` but should be 1 or 2.") - - cos, sin = cos.to(x.device), sin.to(x.device) - - if use_real_unbind_dim == -1: - # Used for flux, cogvideox, hunyuan-dit - # x_real, x_imag = x.reshape(*x.shape[:-1], -1, 2).unbind(-1) # [B, H, S, D//2] - # We avoid using reshape here because for some reason it gets exported with constant shape. - x_real = x[..., 0::2] - x_imag = x[..., 1::2] - x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(3) - elif use_real_unbind_dim == -2: - # Used for Stable Audio, OmniGen, CogView4 and Cosmos - # x_real, x_imag = x.reshape(*x.shape[:-1], 2, -1).unbind(-2) # [B, H, S, D//2] - # We avoid using reshape here because for some reason it gets exported with constant shape. - x_real = x[..., 0::2, :] - x_imag = x[..., 1::2, :] - x_rotated = torch.cat([-x_imag, x_real], dim=-1) - else: - raise ValueError(f"`use_real_unbind_dim={use_real_unbind_dim}` but should be -1 or -2.") - - out = (x.float() * cos + x_rotated.float() * sin).to(x.dtype) - - return out - else: - # used for lumina - x_rotated = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2)) - freqs_cis = freqs_cis.unsqueeze(2) - x_out = torch.view_as_real(x_rotated * freqs_cis).flatten(3) - - return x_out.type_as(x) - - -def patched_cohere_rotary_forward(self, x, position_ids): - # Get batch size and sequence length for manual expansion - batch_size, _ = position_ids.shape[:2] - - # Instead of using expand, manually repeat the tensor. - # Problem with expand: it creates a view with shared memory rather than copying data, - # which causes ONNX export issues with dynamic shapes and view operations. - # Using repeat() ensures actual memory allocation and data copying for ONNX compatibility. - # original: inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1) - inv_freq_base = self.inv_freq[None, :, None].float() # Shape: [1, freq_dim, 1] - inv_freq_expanded = inv_freq_base.repeat(batch_size, 1, 1) # Shape: [batch_size, freq_dim, 1] - - position_ids_expanded = position_ids[:, None, :].float() - device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu" - - with torch.autocast(device_type=device_type, enabled=False): # Force float32 - freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2) - emb = freqs.repeat_interleave(2, dim=-1) # diff from Llama: we interleave() instead of cat() - cos = emb.cos() * self.attention_scaling - sin = emb.sin() * self.attention_scaling - - return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) - - -class CohereModelPatcher(ModelPatcher): - def __enter__(self): - super().__enter__() - - if is_transformers_version(">=", "4.38.0"): - from transformers.models.cohere.modeling_cohere import CohereRotaryEmbedding - - self.original_forward = CohereRotaryEmbedding.forward - CohereRotaryEmbedding.forward = patched_cohere_rotary_forward - - def __exit__(self, exc_type, exc_value, traceback): - super().__exit__(exc_type, exc_value, traceback) - - if is_transformers_version(">=", "4.38.0"): - from transformers.models.cohere.modeling_cohere import CohereRotaryEmbedding - - CohereRotaryEmbedding.forward = self.original_forward - - -# Copied from https://github.com/huggingface/transformers/blob/v4.56.0/src/transformers/models/gpt_oss/modeling_gpt_oss.py#L81 -def gpt_oss_forward(self, hidden_states: torch.Tensor, router_indices=None, routing_weights=None) -> torch.Tensor: - batch_size = hidden_states.shape[0] - hidden_states = hidden_states.reshape(-1, self.hidden_size) - num_experts = routing_weights.shape[1] - hidden_states = hidden_states.repeat(num_experts, 1) - hidden_states = hidden_states.view(num_experts, -1, self.hidden_size) - gate_up = torch.bmm(hidden_states, self.gate_up_proj) + self.gate_up_proj_bias[..., None, :] - gate, up = gate_up[..., ::2], gate_up[..., 1::2] - gate = gate.clamp(min=None, max=self.limit) - up = up.clamp(min=-self.limit, max=self.limit) - glu = gate * torch.sigmoid(gate * self.alpha) - next_states = torch.bmm(((up + 1) * glu), self.down_proj) - next_states = next_states + self.down_proj_bias[..., None, :] - next_states = next_states.view(num_experts, batch_size, -1, self.hidden_size) - next_states = next_states * routing_weights.transpose(0, 1).view(num_experts, batch_size, -1)[..., None] - next_states = next_states.sum(dim=0) - return next_states class GptOssModelPatcher(ModelPatcher): @@ -10342,39 +9949,6 @@ def __exit__(self, exc_type, exc_value, traceback): del sparse_moe_block.down_projs, sparse_moe_block.gate_projs, sparse_moe_block.up_projs -class Gemma4PerLayerInputsGetterModelPatcher(ModelPatcher): - def __init__( - self, - config: "OnnxConfig", - model: Union["PreTrainedModel"], - model_kwargs: Dict[str, Any] = None, - ): - model.__orig_forward = model.forward - - def per_layer_inputs_forward(self, input_ids: torch.Tensor) -> torch.Tensor: - per_layer_inputs_mask = torch.logical_and(input_ids >= 0, input_ids < self.vocab_size_per_layer_input) - per_layer_inputs_tokens = torch.where(per_layer_inputs_mask, input_ids, torch.zeros_like(input_ids)) - per_layer_inputs = self.language_model.get_per_layer_inputs(per_layer_inputs_tokens, None) - return per_layer_inputs - - model.forward = types.MethodType(per_layer_inputs_forward, model) - super().__init__(config, model, model_kwargs) - - def __enter__(self): - super().__enter__() - - def __exit__(self, exc_type, exc_value, traceback): - super().__exit__(exc_type, exc_value, traceback) - self._model.forward = self._model.__orig_forward - - -# OpenVINO has a bug due to which Clamp(-inf, inf) doesn't work correctly: CVS-185473. -# When min == -inf and max == inf, Clamp is equivalent to an identity operation and -# can be removed from the model, which serves as a workaround for the issue. -def patched_gemma4_clippable_linear_forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - hidden_states = self.linear(hidden_states) - return hidden_states - class Gemma4ImageEmbeddingsModelPatcher(CommonImageEmbeddingsModelPatcher): def __init__(self, config, model, model_kwargs): From a9595744e81ef25ff5e07db51fbf29227492d3b4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CElla?= Date: Wed, 27 May 2026 15:29:14 +0200 Subject: [PATCH 07/40] style --- optimum/exporters/openvino/input_generators.py | 3 --- optimum/exporters/openvino/model_patcher.py | 10 ---------- 2 files changed, 13 deletions(-) diff --git a/optimum/exporters/openvino/input_generators.py b/optimum/exporters/openvino/input_generators.py index 2a9b75b29a..036ea52042 100644 --- a/optimum/exporters/openvino/input_generators.py +++ b/optimum/exporters/openvino/input_generators.py @@ -19,14 +19,12 @@ from optimum.intel.utils.import_utils import is_diffusers_version from optimum.utils import ( DEFAULT_DUMMY_SHAPES, - DummyAudioInputGenerator, DummyInputGenerator, DummyPastKeyValuesGenerator, DummySeq2SeqDecoderTextInputGenerator, DummySeq2SeqPastKeyValuesGenerator, DummyTextInputGenerator, DummyTimestepInputGenerator, - DummyTransformerTextInputGenerator, DummyVisionInputGenerator, FalconDummyPastKeyValuesGenerator, MistralDummyPastKeyValuesGenerator, @@ -87,7 +85,6 @@ def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int return pkv - class DummyQwen3VLLMInputGenerator(DummyTextInputGenerator): SUPPORTED_INPUT_NAMES = ( "input_ids", diff --git a/optimum/exporters/openvino/model_patcher.py b/optimum/exporters/openvino/model_patcher.py index 60244b9ac7..e445d1e2c6 100644 --- a/optimum/exporters/openvino/model_patcher.py +++ b/optimum/exporters/openvino/model_patcher.py @@ -69,7 +69,6 @@ if is_transformers_version(">=", "4.48"): from transformers.cache_utils import DynamicCache, EncoderDecoderCache - from transformers.models.moonshine.modeling_moonshine import MoonshinePreTrainedModel if is_transformers_version(">=", "4.53"): from transformers.masking_utils import ( ALL_MASK_ATTENTION_FUNCTIONS, @@ -84,7 +83,6 @@ from transformers.models.qwen3_moe.modeling_qwen3_moe import Qwen3MoeSparseMoeBlock - if is_transformers_version(">=", "4.53.1"): from transformers.masking_utils import find_packed_sequence_indices @@ -817,7 +815,6 @@ def __exit__(self, exc_type, exc_value, traceback): self._model.set_attention_type("block_sparse") - class SAMModelPatcher(ModelPatcher): def __init__( self, @@ -957,9 +954,6 @@ def patched_forward(input_ids, attention_mask): self.patched_forward = patched_forward - - - class CLIPModelPatcher(ModelPatcher): def __enter__(self): super().__enter__() @@ -973,7 +967,6 @@ def __exit__(self, exc_type, exc_value, traceback): CLIPSdpaAttention.forward = self.original_sdpa_forward - class Qwen3MoeModelPatcher(ModelPatcher): def __enter__(self): super().__enter__() @@ -1002,8 +995,6 @@ def _get_feat_extract_output_lengths_patched(self, input_lengths: torch.LongTens return output_conv3_length - - class GptOssModelPatcher(ModelPatcher): def __enter__(self): super().__enter__() @@ -9949,7 +9940,6 @@ def __exit__(self, exc_type, exc_value, traceback): del sparse_moe_block.down_projs, sparse_moe_block.gate_projs, sparse_moe_block.up_projs - class Gemma4ImageEmbeddingsModelPatcher(CommonImageEmbeddingsModelPatcher): def __init__(self, config, model, model_kwargs): super().__init__(config, model, model_kwargs) From e0a04cef65404a0e277ff94b9527721625b790df Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CElla?= Date: Wed, 27 May 2026 15:51:37 +0200 Subject: [PATCH 08/40] moving utils --- optimum/exporters/openvino/base.py | 2 +- optimum/exporters/openvino/base_patcher.py | 751 ++++++++++++++++++++ optimum/exporters/openvino/model_configs.py | 24 +- optimum/exporters/openvino/model_patcher.py | 704 +----------------- 4 files changed, 765 insertions(+), 716 deletions(-) create mode 100644 optimum/exporters/openvino/base_patcher.py diff --git a/optimum/exporters/openvino/base.py b/optimum/exporters/openvino/base.py index 1120989b3c..29c483870b 100644 --- a/optimum/exporters/openvino/base.py +++ b/optimum/exporters/openvino/base.py @@ -30,7 +30,7 @@ from transformers import PretrainedConfig, PreTrainedModel from optimum.exporters.base import ExporterConfig -from optimum.exporters.openvino.model_patcher import ModelPatcher, PatchingSpec +from optimum.exporters.openvino.base_patcher import ModelPatcher, PatchingSpec from optimum.utils import DEFAULT_DUMMY_SHAPES, DummyInputGenerator, DummySeq2SeqPastKeyValuesGenerator, logging from optimum.utils.doc import add_dynamic_docstring from optimum.utils.import_utils import ( diff --git a/optimum/exporters/openvino/base_patcher.py b/optimum/exporters/openvino/base_patcher.py new file mode 100644 index 0000000000..5e1d45923d --- /dev/null +++ b/optimum/exporters/openvino/base_patcher.py @@ -0,0 +1,751 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import dataclasses +import functools +import inspect +import logging +import sys +import types +from typing import Any, Callable + +import torch +import transformers +from torch.onnx import symbolic_helper +from transformers import PreTrainedModel +from transformers.cache_utils import DynamicCache, EncoderDecoderCache + +from optimum.intel.utils.import_utils import is_torch_version, is_transformers_version + + +if is_transformers_version(">=", "4.44") and is_transformers_version("<", "4.50"): + from optimum.exporters.openvino._traceable_cache import TraceableCache + + +if is_transformers_version(">=", "4.48"): + from transformers.cache_utils import DynamicCache, EncoderDecoderCache +if is_transformers_version(">=", "4.53"): + from transformers.masking_utils import ( + ALL_MASK_ATTENTION_FUNCTIONS, + _ignore_causal_mask_sdpa, + and_masks, + causal_mask_function, + eager_mask, + padding_mask_function, + prepare_padding_mask, + sdpa_mask, + ) + + +if is_transformers_version(">=", "4.53.1"): + from transformers.masking_utils import find_packed_sequence_indices + + +if is_transformers_version(">=", "4.54"): + from transformers.utils import TransformersKwargs + + from optimum.exporters.openvino._traceable_decorator import traceable_check_model_inputs +else: + TransformersKwargs = object + + +if is_transformers_version(">=", "4.56"): + import transformers.masking_utils + from transformers.cache_utils import DynamicLayer + + +logger = logging.getLogger(__name__) + + +@symbolic_helper.parse_args("v", "v") +def __ior_(g, self: torch._C.Value, other: torch._C.Value) -> torch._C.Value: + return g.op("Or", self, other) + + +torch.onnx.register_custom_op_symbolic("aten::__ior__", __ior_, 14) + +if is_torch_version("<", "2.9"): + # this was fixed in torch in 2.9 https://github.com/pytorch/pytorch/pull/159973 + from torch.onnx import JitScalarType + from torch.onnx.symbolic_opset14 import _attention_scale, _causal_attention_mask + + @symbolic_helper.parse_args("v", "v", "v", "v", "f", "b", "v", "b") + def scaled_dot_product_attention( + g, + query: torch._C.Value, + key: torch._C.Value, + value: torch._C.Value, + attn_mask: torch._C.Value | None = None, + dropout_p: float = 0.0, + is_causal: bool = False, + scale: torch._C.Value | None = None, + enable_gqa: bool = False, + ): + assert (not is_causal) or ( + is_causal and symbolic_helper._is_none(attn_mask) + ), "is_causal and attn_mask cannot be set at the same time" + assert not enable_gqa, "conversion of scaled_dot_product_attention not implemented if enable_gqa is True" + + if symbolic_helper._is_none(scale): + scale = _attention_scale(g, query) + + if is_causal: + attn_mask = _causal_attention_mask(g, query, key) + + # Swap the last two axes of key + # NOTE: onnx-script has different logic here, because the attribute perms in + # transpose needs list of ints + key_shape_builtin = symbolic_helper._get_tensor_rank(key) + key_transposed_axes = list(range(key_shape_builtin)) + key_transposed_axes[-1], key_transposed_axes[-2] = (key_transposed_axes[-2], key_transposed_axes[-1]) + key_transposed = g.op("Transpose", key, perm_i=key_transposed_axes) + + # https://github.com/pytorch/pytorch/blob/12da0c70378b5be9135c6fda62a9863bce4a4818/aten/src/ATen/native/transformers/attention.cpp#L653 + # Scale q, k before matmul for stability see https://tinyurl.com/sudb9s96 for math + query_scaled = g.op("Mul", query, g.op("Sqrt", scale)) + key_transposed_scaled = g.op("Mul", key_transposed, g.op("Sqrt", scale)) + mul_qk = g.op("MatMul", query_scaled, key_transposed_scaled) + + if symbolic_helper._is_none(attn_mask): + mul_qk_add = mul_qk + attn_weight = g.op("Softmax", mul_qk_add, axis_i=-1) + elif JitScalarType.from_value(attn_mask) == JitScalarType.BOOL: + # Turn the Boolean mask to float: attn_mask.masked_fill(not attn_mask, -float('inf')) + const_zero = g.op("Constant", value_t=torch.tensor([0.0])) + const_neg_inf = g.op("Constant", value_t=torch.tensor([-float("inf")])) + attn_mask = g.op("Where", attn_mask, const_zero, const_neg_inf) + mul_qk_add = g.op("Add", mul_qk, attn_mask) + attn_weight = g.op("Softmax", mul_qk_add, axis_i=-1) + # when using scaled dot product attention with a boolean mask, we replace NaN values in attn_weight with 0.0 + attn_weight = g.op( + "Where", g.op("IsNaN", attn_weight), g.op("Constant", value_t=torch.tensor([0.0])), attn_weight + ) + elif JitScalarType.from_value(attn_mask) in ( + JitScalarType.FLOAT, + JitScalarType.HALF, + JitScalarType.BFLOAT16, + ): + mul_qk_add = g.op("Add", mul_qk, attn_mask) + attn_weight = g.op("Softmax", mul_qk_add, axis_i=-1) + else: + raise ValueError(f"Unsupported type for attn_mask: {JitScalarType.from_value(attn_mask)}") + + if dropout_p != 0: + attn_weight = g.op( + "Dropout", + attn_weight, + g.op("Constant", value_t=torch.tensor(dropout_p, dtype=torch.float)), + ) + + return g.op("MatMul", attn_weight, value) + + torch.onnx.register_custom_op_symbolic("aten::scaled_dot_product_attention", scaled_dot_product_attention, 14) + + +def patch_everywhere(attribute_name: str, patch: Any, module_name_prefix: str | None = None): + """Finds all occurrences of `attribute_name` in the loaded modules and patches them with `patch`. + + Args: + attribute_name (`str`): + The name of attribute to patch. + patch (`Any`): + The patch for the attribute. + module_name_prefix (`Optional[str]`, defaults to `None`): + If set, only module names starting with this prefix will be considered for patching. + """ + # sys.modules may be updated while being iterated over, hence the list copy. + for name in list(sys.modules): + module = sys.modules[name] + if module_name_prefix is not None and not name.startswith(module_name_prefix): + continue + if hasattr(module, attribute_name): + setattr(module, attribute_name, patch) + + +def override_arguments(args, kwargs, forward_signature, model_kwargs: dict[str, Any]): + """Override the args and kwargs with the argument values from model_kwargs, following the signature forward_signature corresponding to args and kwargs.""" + args = list(args) + + for argument in model_kwargs: + if argument in forward_signature.parameters: + argument_index = list(forward_signature.parameters.keys()).index(argument) + if argument in kwargs or len(args) <= argument_index: + kwargs[argument] = model_kwargs[argument] + else: + args[argument_index] = model_kwargs[argument] + else: + kwargs[argument] = model_kwargs[argument] + + return args, kwargs + + +def preprocess_encoder_outputs(encoder_outputs): + if is_transformers_version(">=", "4.54") and isinstance(encoder_outputs, (list, tuple)): + encoder_outputs = BaseModelOutput(*encoder_outputs) + + return encoder_outputs + + +def preprocess_past_key_values(past_key_values): + if ( + is_transformers_version(">=", "4.48") + and isinstance(past_key_values, (list, tuple)) + and isinstance(past_key_values[0], (list, tuple)) + ): + if len(past_key_values[0]) == 2: + if hasattr(DynamicCache, "from_legacy_cache"): + past_key_values = DynamicCache.from_legacy_cache(past_key_values) + else: + past_key_values = DynamicCache(past_key_values) + elif len(past_key_values[0]) == 4: + if hasattr(EncoderDecoderCache, "from_legacy_cache"): + past_key_values = EncoderDecoderCache.from_legacy_cache(past_key_values) + else: + past_key_values = EncoderDecoderCache( + DynamicCache([layer[:2] for layer in past_key_values]), + DynamicCache([layer[2:] for layer in past_key_values]), + ) + else: + raise ValueError( + f"past_key_values should have either 2 or 4 elements, but it has {len(past_key_values[0])} elements." + ) + + return past_key_values + + +def postprocess_past_key_values(past_key_values, output_names: list[str]): + if is_transformers_version(">=", "4.48") and isinstance(past_key_values, (EncoderDecoderCache, DynamicCache)): + if hasattr(past_key_values, "to_legacy_cache"): + past_key_values = past_key_values.to_legacy_cache() + elif isinstance(past_key_values, DynamicCache): + past_key_values = [(lay.keys, lay.values) for lay in past_key_values.layers] + elif isinstance(past_key_values, EncoderDecoderCache): + past_key_values = [ + (self_lay.keys, self_lay.values, cross_lay.keys, cross_lay.values) + for self_lay, cross_lay in zip( + past_key_values.self_attention_cache.layers, + past_key_values.cross_attention_cache.layers, + ) + ] + else: + raise NotImplementedError(f"Unable to serialize class {type(past_key_values)}.") + + if ( + isinstance(past_key_values, (list, tuple)) + and isinstance(past_key_values[0], (list, tuple)) + and not any("encoder.key" in output_name for output_name in output_names) + ): + past_key_values = tuple(pkv[:2] for pkv in past_key_values) + + return past_key_values + + +@dataclasses.dataclass +class PatchingSpec: + """Data class that holds patching specifications. + + Args: + o: Module / object where the op to patch is located + name: Name of the op to monkey patch + custom_op: Custom op that patches the original op + orig_op: Original op that is being patched + op_wrapper: Wrapper (optional) that wraps both the original and custom ops. + It is useful for ops that are class or static methods for instance. + """ + + o: Any + name: str + custom_op: Callable + orig_op: Callable | None = None + op_wrapper: Callable | None = None + + +# An ONNX-export-compatible version of `tensor.unfold`. Without this, we get: +# torch.onnx.errors.SymbolicValueError: Unsupported: ONNX export of operator Unfold, input size not accessible. +# See https://github.com/pytorch/pytorch/issues/81871 for more information +def onnx_compatible_unfold(input_tensor, dimension, size, step): + """Custom implementation of torch.unfold without using torch.unfold. + + Args: + input_tensor (torch.Tensor): The input tensor. + dimension (int): The dimension to unfold. + size (int): The size of each slice. + step (int): The step size between slices. + + Returns: + torch.Tensor: The unfolded tensor. + """ + # Check if dimension is within the valid range + if not (-input_tensor.dim() <= dimension < input_tensor.dim()): + raise ValueError( + f"Dimension out of range (expected to be in range of [{-input_tensor.dim()}, {input_tensor.dim() - 1}], but got {dimension})" + ) + + # Normalize negative dimension + dimension = dimension % input_tensor.dim() + + # Compute the shape of the unfolded output + input_size = input_tensor.size(dimension) + num_slices = (input_size - size) // step + 1 + + # Permute dimension to the end for easier indexing + input_tensor = input_tensor.transpose(dimension, -1) + + # Extract slices + slices = [] + for i in range(num_slices): + start = i * step + end = start + size + slices.append(input_tensor[..., start:end]) + + # Stack slices and permute dimensions back + result = torch.stack(slices, dim=-2).transpose(dimension, -2) + return result + + +# An ONNX-export-compatible version of `tensor.repeat_interleave`. +# Without this, we get the following error: https://github.com/pytorch/pytorch/issues/145100 +# NOTE: This implementation is only necessary for export with dynamo=False (dynamo=True works correctly). +# and can be removed once Optimum switches to dynamo-based exports +def onnx_compatible_repeat_interleave(input_tensor, repeats, dim=None, output_size=None): # noqa: D417 + """Custom implementation of torch.repeat_interleave without using torch.repeat_interleave. + + Args: + input_tensor (torch.Tensor): The input tensor. + repeats (int or torch.Tensor): The number of repetitions for each element. + dim (int, optional): The dimension along which to repeat. Defaults to None. + + Returns: + torch.Tensor: The repeated tensor. + """ + if isinstance(repeats, int) or (torch.is_tensor(repeats) and repeats.dim() == 0): + if dim is None: + return input_tensor.flatten().unsqueeze(1).expand(-1, repeats).flatten() + repeats = torch.full((input_tensor.shape[dim],), repeats, dtype=torch.long, device=input_tensor.device) + + if dim is None: + return onnx_compatible_repeat_interleave(input_tensor.flatten(), repeats, 0) + + if dim != 0: + input_tensor = input_tensor.transpose(0, dim) + + # Create expand mask + max_repeats = repeats.max() + expanded = input_tensor.unsqueeze(1).expand(-1, max_repeats, *input_tensor.shape[1:]) + mask = torch.arange(max_repeats, device=input_tensor.device) < repeats.unsqueeze(1) + result = expanded[mask] + + if dim != 0: + result = result.transpose(0, dim) + + return result + + +# Custom implementation of torch.linalg.matrix_norm not using torch.linalg.matrix_norm, torch.norm or torch.linalg.norm. +def onnx_compatible_linalg_norm(x, ord=2, dim=None, keepdim=False, *, dtype=None, out=None) -> torch.Tensor: + if ord != 2: + raise ValueError( + f"Only ord=2 is supported by onnx_compatible_linalg_norm, but got ord={ord}. " + "Please extend this function to support other norms." + ) + + if dim is None: + dim = (-2, -1) + + norm = torch.sqrt(torch.sum(torch.square(x), dim=dim, keepdim=keepdim)) + + if dtype is not None: + norm = norm.to(dtype) + if out is not None: + out.copy_(norm) + + return norm + + +def onnx_compatible_rms_norm(input, normalized_shape, weight=None, eps=None): + if eps is None: + eps = torch.finfo(input.dtype).eps + + axis = -len(normalized_shape) + mean_square = torch.mean(torch.square(input), dim=axis, keepdim=True) + rms = torch.sqrt(mean_square + eps) + output = input / rms + + if weight is not None: + output = output * weight + + return output + + +# A patched version of https://github.com/huggingface/transformers/blob/v4.53.2/src/transformers/masking_utils.py#L602 +# That returns a tensor of zeros with the same shape as position_ids indicating no packed sequence indices. +def find_packed_sequence_indices_patched(position_ids: torch.Tensor) -> torch.Tensor: + return torch.zeros_like(position_ids) + + +if is_transformers_version(">=", "4.53"): + _prepare_padding_mask_slice = "_slice" in inspect.signature(prepare_padding_mask).parameters +else: + _prepare_padding_mask_slice = False + + +# Custom vectorized implementation of sdpa_mask without using vmap +def _orig_sdpa_mask_without_vmap( + batch_size: int, + cache_position: torch.Tensor, + kv_length: int, + kv_offset: int = 0, + mask_function: Callable | None = None, + attention_mask: torch.Tensor | None = None, + local_size: int | None = None, + allow_is_causal_skip: bool = True, + **kwargs, +) -> torch.Tensor | None: + if mask_function is None: + mask_function = causal_mask_function + + q_length = cache_position.shape[0] + # Potentially pad the 2D mask, and slice it correctly + if _prepare_padding_mask_slice: + padding_mask = prepare_padding_mask(attention_mask, kv_length, kv_offset, _slice=False) + else: + padding_mask = prepare_padding_mask(attention_mask, kv_length, kv_offset) + + # Under specific conditions, we can avoid materializing the mask, instead relying on the `is_causal` argument + if allow_is_causal_skip and _ignore_causal_mask_sdpa(padding_mask, q_length, kv_length, kv_offset, local_size): + return None + + # Potentially add the padding 2D mask + if padding_mask is not None: + mask_function = and_masks(mask_function, padding_mask_function(padding_mask)) + + # Create broadcatable indices + device = cache_position.device + q_indices = cache_position[None, None, :, None] + head_indices = torch.arange(1, dtype=torch.long, device=device)[None, :, None, None] + batch_indices = torch.arange(batch_size, dtype=torch.long, device=device)[:, None, None, None] + kv_indices = torch.arange(kv_length, dtype=torch.long, device=device)[None, None, None, :] + kv_offset + # Apply mask function element-wise through broadcasting + causal_mask = mask_function(batch_indices, head_indices, q_indices, kv_indices) + # Expand the mask to match batch size and query length if they weren't used in the mask function + causal_mask = causal_mask.expand(batch_size, -1, q_length, kv_length) + + return causal_mask + + +# Compatibility wrapper for sdpa_mask_without_vmap from optimum. +# The installed optimum version expects (batch_size, cache_position: Tensor, kv_length, ...), +# but transformers >= 5.5 passes (batch_size, q_length: int, kv_length: int, q_offset: int, ...). +def sdpa_mask_without_vmap(batch_size, q_length=None, kv_length=None, q_offset=0, kv_offset=0, **kwargs): + import inspect + + sig = inspect.signature(_orig_sdpa_mask_without_vmap) + if is_transformers_version(">=", "5.5") and "cache_position" in sig.parameters and q_length is not None: + # Old optimum signature: (batch_size, cache_position, kv_length, kv_offset, ...) + cache_position = torch.arange(q_length, dtype=torch.long) + q_offset + kwargs.pop("q_offset", None) + kwargs.pop("allow_is_bidirectional_skip", None) + kwargs.pop("allow_torch_fix", None) + kwargs.pop("use_vmap", None) + kwargs.pop("device", None) + return _orig_sdpa_mask_without_vmap(batch_size, cache_position, kv_length, kv_offset=kv_offset, **kwargs) + else: + return _orig_sdpa_mask_without_vmap( + batch_size, q_length=q_length, kv_length=kv_length, q_offset=q_offset, kv_offset=kv_offset, **kwargs + ) + + +# Adapted from https://github.com/huggingface/transformers/blob/v4.53.0/src/transformers/masking_utils.py#L433 +def eager_mask_without_vmap(*args, **kwargs) -> torch.Tensor: + kwargs.pop("allow_is_causal_skip", None) + dtype = kwargs.get("dtype", torch.float32) + mask = sdpa_mask_without_vmap(*args, allow_is_causal_skip=False, **kwargs) + mask = torch.where(mask, torch.tensor(0.0, device=mask.device, dtype=dtype), torch.finfo(dtype).min) + return mask + + +original_triu = torch.triu +original_tril = torch.tril + + +# Custom implementation of torch.tril that doesn't fail on int32 tensors. +def onnx_compatible_tril(input_tensor: torch.Tensor, *args, **kwargs) -> torch.Tensor: + if input_tensor.dtype == torch.int32: + return original_tril(input_tensor.to(torch.int64), *args, **kwargs).to(torch.int32) + else: + return original_tril(input_tensor, *args, **kwargs) + + +# Custom implementation of torch.triu that doesn't fail on int32 tensors. +def onnx_compatible_triu(input_tensor: torch.Tensor, *args, **kwargs) -> torch.Tensor: + if input_tensor.dtype == torch.int32: + return original_triu(input_tensor.to(torch.int64), *args, **kwargs).to(torch.int32) + else: + return original_triu(input_tensor, *args, **kwargs) + + +original_scaled_dot_product_attention = torch.nn.functional.scaled_dot_product_attention + + +# A patched `torch.nn.functional.scaled_dot_product_attention` that doesn't fail during tracing +# from passing `is_causal` as a tensor (which is usually obtained with tensor shapes comparisons). +def traceable_scaled_dot_product_attention( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_mask: torch.Tensor | None = None, + dropout_p: float = 0.0, + is_causal: bool = False, + **kwargs, +) -> torch.Tensor: + if isinstance(is_causal, torch.Tensor): + is_causal = is_causal.item() + + if "enable_gqa" in kwargs: + kwargs.pop("enable_gqa") + + attn_weights = original_scaled_dot_product_attention( + query=query, key=key, value=value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, **kwargs + ) + + return attn_weights + + +# No-op bfloat16 casting to avoid issues with legacy ONNX export which cast to complex128 +def noop_bfloat16_casting(self): + return self + + +original_movedim = torch.Tensor.movedim + + +def onnx_compatible_movedim(self: torch.Tensor, dim1, dim2) -> torch.Tensor: + dim = self.dim() + if dim1 < 0: + dim1 += dim + if dim2 < 0: + dim2 += dim + return original_movedim(self, dim1, dim2) + + +def patched_dynamic_layer_update( + self, key_states: torch.Tensor, value_states: torch.Tensor, cache_kwargs: dict[str, Any] | None = None +) -> tuple[torch.Tensor, torch.Tensor]: + if self.keys is None: + self.keys = key_states + self.values = value_states + self.device = key_states.device + self.dtype = key_states.dtype + self.is_initialized = True + else: + self.keys = torch.cat([self.keys, key_states], dim=-2) + self.values = torch.cat([self.values, value_states], dim=-2) + return self.keys, self.values + + +UNSUPPORTED_OPS_PATCHING_SPEC = [ + PatchingSpec(torch, "tril", onnx_compatible_tril, torch.tril), + PatchingSpec(torch, "triu", onnx_compatible_triu, torch.triu), + PatchingSpec(torch, "rms_norm", onnx_compatible_rms_norm, torch.rms_norm), + PatchingSpec(torch.Tensor, "unfold", onnx_compatible_unfold, torch.Tensor.unfold), + PatchingSpec(torch.linalg, "norm", onnx_compatible_linalg_norm, torch.linalg.norm), + PatchingSpec(torch.Tensor, "bfloat16", noop_bfloat16_casting, torch.Tensor.bfloat16), + PatchingSpec(torch.Tensor, "movedim", onnx_compatible_movedim, torch.Tensor.movedim), + PatchingSpec(torch.Tensor, "repeat_interleave", onnx_compatible_repeat_interleave, torch.Tensor.repeat_interleave), + # TracerWarning: Using len to get tensor shape might cause the trace to be incorrect. Recommended usage would be tensor.shape[0]. Passing a tensor of different shape might lead to errors or silently give incorrect results. + PatchingSpec(torch.Tensor, "__len__", lambda x: x.shape[0], torch.Tensor.__len__), + PatchingSpec( + torch.nn.functional, + "scaled_dot_product_attention", + traceable_scaled_dot_product_attention, + torch.nn.functional.scaled_dot_product_attention, + ), +] + + +class ModelPatcher: + def __init__( + self, + config: "OnnxConfig", + model: PreTrainedModel, + model_kwargs: dict[str, Any] | None = None, + ): + self._model = model + + patching_specs = config.PATCHING_SPECS or [] + patching_specs.extend(UNSUPPORTED_OPS_PATCHING_SPEC) + + self._patching_specs = [] + for spec in patching_specs: + final_spec = spec + if spec.orig_op is None: + final_spec = dataclasses.replace(spec, orig_op=getattr(spec.o, spec.name)) + self._patching_specs.append(final_spec) + + self.orig_forward_name = "forward" if hasattr(self._model, "forward") else "call" + self.orig_forward = getattr(self._model, self.orig_forward_name) + + if is_transformers_version(">=", "4.54") and hasattr(self.orig_forward, "__wrapped__"): + # the original check_model_inputs has some failing cases that we fix in traceable_check_model_inputs + # we fix those issues in a PR in transformers https://github.com/huggingface/transformers/pull/40811 + # issues are: support for positional args (use_cache for instance) and fix for _CAN_RECORD_REGISTRY + # explicitly mapping to None for some models + self.orig_forward = types.MethodType( + traceable_check_model_inputs(self.orig_forward.__wrapped__), self._model + ) + + self.real_config = config + use_cache = getattr(self.real_config, "use_past", False) + self.model_kwargs = model_kwargs if model_kwargs is not None else {} + + for module in self._model.modules(): + if hasattr(module, "config") and hasattr(module.config, "use_cache"): + module.config.use_cache = use_cache + + @functools.wraps(self.orig_forward) + def patched_forward(*args, **kwargs): + signature = inspect.signature(self.orig_forward) + args, kwargs = override_arguments(args, kwargs, signature, model_kwargs=self.model_kwargs) + + # Transformers doesn't always respect the config.use_cache attribute + # there are cases where setting use_cache to true in every config and + # subconfig of a model still doesn't enable past_key_values in the outputs (gemma3) + # Explicitly setting the use_cache argument of the forward method seems to be the most reliable way + if "use_cache" in signature.parameters: + use_cache_index = list(signature.parameters.keys()).index("use_cache") + if use_cache_index < len(args): + args[use_cache_index] = use_cache + elif "use_cache" in kwargs: + kwargs["use_cache"] = use_cache + + if "past_key_values" in signature.parameters: + # Most models require past_key_values to be a cache instance instead of a tuple now + pkv_index = list(signature.parameters.keys()).index("past_key_values") + if pkv_index < len(args) and args[pkv_index] is not None: + args[pkv_index] = preprocess_past_key_values(args[pkv_index]) + elif kwargs.get("past_key_values") is not None: + kwargs["past_key_values"] = preprocess_past_key_values(kwargs["past_key_values"]) + + if "encoder_outputs" in signature.parameters: + # Some encoder-decoder models started to not accept encoder_outputs as tuple (e.g. moonshine) + encoder_outputs_index = list(signature.parameters.keys()).index("encoder_outputs") + if encoder_outputs_index < len(args) and args[encoder_outputs_index] is not None: + args[encoder_outputs_index] = preprocess_encoder_outputs(args[encoder_outputs_index]) + elif kwargs.get("encoder_outputs") is not None: + kwargs["encoder_outputs"] = preprocess_encoder_outputs(kwargs["encoder_outputs"]) + + outputs = self.orig_forward(*args, **kwargs) + + # This code block handles different cases of the filtered_outputs input to align it with the expected + # format of outputs. It is common for the output type of a model to vary, such as tensor, list, + # tuple, etc. For Transformers models, the output is encapsulated in a ModelOutput object that + # contains the output names of the model. In the case of Timm classification models, the output + # is of type tensor. By default, it is assumed that the output names mentioned in the ONNX config + # match the outputs in order. + filtered_outputs = {} + output_names = list(config.outputs.keys()) + if isinstance(outputs, dict): + for name, value in outputs.items(): + onnx_output_name = config.torch_to_onnx_output_map.get(name, name) + if ( + onnx_output_name in output_names + or (use_cache and name.startswith("past_key_values")) + or any(key.startswith(onnx_output_name) for key in output_names) + ): + filtered_outputs[name] = value + elif isinstance(outputs, (list, tuple)): + filtered_outputs = dict(zip(output_names, outputs)) + else: + if len(output_names) > 1: + num_outputs = len(output_names) + output_names_str = ", ".join(output_names) + raise ValueError( + f"{config.__class__.__name__} expects the model to return {num_outputs} outputs: {output_names_str}, " + f"but the it returned a single output of type {type(outputs)}. Please make sure either that the model " + "returns all the expected outputs, or that the ONNX config is correctly defined with the expected outputs." + ) + output_name = output_names[0] + filtered_outputs[output_name] = outputs + + if filtered_outputs.get("past_key_values") is not None: + filtered_outputs["past_key_values"] = postprocess_past_key_values( + filtered_outputs["past_key_values"], output_names=output_names + ) + + return filtered_outputs + + self.patched_forward = patched_forward + + def patch_ops(self): + for spec in self._patching_specs: + custom_op = spec.custom_op if spec.op_wrapper is None else spec.op_wrapper(spec.custom_op) + setattr(spec.o, spec.name, custom_op) + + def restore_ops(self): + for spec in self._patching_specs: + orig_op = spec.orig_op if spec.op_wrapper is None else spec.op_wrapper(spec.orig_op) + setattr(spec.o, spec.name, orig_op) + + def __enter__(self): + self.patch_ops() + setattr(self._model, self.orig_forward_name, self.patched_forward) + + # This is a workaround for the Cache class in transformers, we replace it + # with traceable cache is because the original one used in transformers + # inherited from nn.Module (for a couple versions), which can't be traced as input. + if is_transformers_version(">=", "4.44") and is_transformers_version("<", "4.50"): + self.original_cache_class = transformers.cache_utils.Cache + transformers.cache_utils.Cache = TraceableCache + + # This is a workaround for mask generation in transformers >= 4.53. + # The masking process uses vmap which is not traceable by TorchScript. + if is_transformers_version(">=", "4.53"): + ALL_MASK_ATTENTION_FUNCTIONS.register("sdpa", sdpa_mask_without_vmap) + ALL_MASK_ATTENTION_FUNCTIONS.register("eager", eager_mask_without_vmap) + + # This is a workaround for the find_packed_sequence_indices function in transformers which + # should only return a tensor of zeros with the same shape as position_ids indicating no packed sequence indices. + # The function uses torch.diff which is not traceable by TorchScript. + if is_transformers_version(">=", "4.53.1"): + self.original_find_packed_sequence_indices = find_packed_sequence_indices + transformers.masking_utils.find_packed_sequence_indices = find_packed_sequence_indices_patched + + # Starting from transformers 4.56.0, DynamicCache uses DynamicLayer which has an update method + # that uses torch.cat to concatenate an empty tensor with the key/value states during the first call. + # This causes issues during TorchScript tracing. + if is_transformers_version(">=", "4.56"): + self.original_dynamic_layer_update = DynamicLayer.update + DynamicLayer.update = patched_dynamic_layer_update + + def __exit__(self, exc_type, exc_value, traceback): + self.restore_ops() + setattr(self._model, self.orig_forward_name, self.orig_forward) + + if is_transformers_version(">=", "4.44") and is_transformers_version("<", "4.50"): + transformers.cache_utils.Cache = self.original_cache_class + + if is_transformers_version(">=", "4.53"): + ALL_MASK_ATTENTION_FUNCTIONS.register("sdpa", sdpa_mask) + ALL_MASK_ATTENTION_FUNCTIONS.register("eager", eager_mask) + + if is_transformers_version(">=", "4.53.1"): + transformers.masking_utils.find_packed_sequence_indices = self.original_find_packed_sequence_indices + + if is_transformers_version(">=", "4.56"): + DynamicLayer.update = self.original_dynamic_layer_update + + def __call__(self, *args, **kwargs): + if getattr(self._model, self.orig_forward_name) is self.orig_forward: + logger.warning("Running the non-patched model") + return self._model(*args, **kwargs) diff --git a/optimum/exporters/openvino/model_configs.py b/optimum/exporters/openvino/model_configs.py index 416d8af472..9ebab904df 100644 --- a/optimum/exporters/openvino/model_configs.py +++ b/optimum/exporters/openvino/model_configs.py @@ -4733,13 +4733,6 @@ def flatten_past_key_values(self, flattened_output, name, idx, t): flattened_output[f"{name}.{idx}.key_value"] = t -@register_in_tasks_manager( - "pix2struct", - *[ - "image-to-text", - "image-to-text-with-past", - ], -) class _Pix2StructNormalizedConfig(NormalizedSeq2SeqConfig): ENCODER_NUM_LAYERS = "vision_config.num_hidden_layers" DECODER_NUM_LAYERS = "text_config.num_layers" @@ -4749,6 +4742,13 @@ class _Pix2StructNormalizedConfig(NormalizedSeq2SeqConfig): VOCAB_SIZE = "text_config.vocab_size" +@register_in_tasks_manager( + "pix2struct", + *[ + "image-to-text", + "image-to-text-with-past", + ], +) class Pix2StructOpenVINOConfig(OnnxSeq2SeqConfigWithPast): PAD_ATTENTION_MASK_TO_PAST = True NORMALIZED_CONFIG_CLASS = _Pix2StructNormalizedConfig @@ -4900,11 +4900,6 @@ class XLMOpenVINOConfig(BertOpenVINOConfig): MAX_TRANSFORMERS_VERSION = "4.57.6" -@register_in_tasks_manager("xlm-roberta", *COMMON_TEXT_TASKS) -class XLMRobertaOpenVINOConfig(DistilBertOpenVINOConfig): - pass - - @register_in_tasks_manager("distilbert", *COMMON_TEXT_TASKS) class DistilBertOpenVINOConfig(TextEncoderOnnxConfig): NORMALIZED_CONFIG_CLASS = NormalizedTextConfig @@ -4918,6 +4913,11 @@ def inputs(self) -> dict: return {"input_ids": dynamic_axis, "attention_mask": dynamic_axis} +@register_in_tasks_manager("xlm-roberta", *COMMON_TEXT_TASKS) +class XLMRobertaOpenVINOConfig(DistilBertOpenVINOConfig): + pass + + @register_in_tasks_manager("roberta", *COMMON_TEXT_TASKS) class RobertaOpenVINOConfig(DistilBertOpenVINOConfig): pass diff --git a/optimum/exporters/openvino/model_patcher.py b/optimum/exporters/openvino/model_patcher.py index e445d1e2c6..0087b10358 100644 --- a/optimum/exporters/openvino/model_patcher.py +++ b/optimum/exporters/openvino/model_patcher.py @@ -12,13 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. -import dataclasses import functools import inspect import logging import logging as log import math -import sys import types from dataclasses import dataclass from typing import Any, Callable, Dict, List, Optional, Tuple, Union @@ -27,7 +25,6 @@ import torch.nn.functional as F import transformers from torch import nn -from torch.onnx import symbolic_helper from transformers import PreTrainedModel from transformers.cache_utils import Cache, DynamicCache, EncoderDecoderCache from transformers.configuration_utils import PretrainedConfig @@ -51,7 +48,7 @@ from transformers.utils import ModelOutput from optimum.exporters.openvino._ov_ops import convert_recurrent_attention_cell -from optimum.exporters.openvino.base import OnnxConfig +from optimum.exporters.openvino.base_patcher import UNSUPPORTED_OPS_PATCHING_SPEC, ModelPatcher from optimum.intel.utils.import_utils import ( is_diffusers_version, is_openvino_version, @@ -63,36 +60,21 @@ if is_transformers_version(">=", "4.43") and is_transformers_version("<", "4.48"): from transformers.models.clip.modeling_clip import CLIPAttention, CLIPSdpaAttention -if is_transformers_version(">=", "4.44") and is_transformers_version("<", "4.50"): - from optimum.exporters.openvino._traceable_cache import TraceableCache - if is_transformers_version(">=", "4.48"): from transformers.cache_utils import DynamicCache, EncoderDecoderCache if is_transformers_version(">=", "4.53"): from transformers.masking_utils import ( ALL_MASK_ATTENTION_FUNCTIONS, - _ignore_causal_mask_sdpa, - and_masks, - causal_mask_function, eager_mask, - padding_mask_function, - prepare_padding_mask, sdpa_mask, ) from transformers.models.qwen3_moe.modeling_qwen3_moe import Qwen3MoeSparseMoeBlock -if is_transformers_version(">=", "4.53.1"): - from transformers.masking_utils import find_packed_sequence_indices - - if is_transformers_version(">=", "4.54"): from transformers.masking_utils import create_causal_mask from transformers.utils import TransformersKwargs - - from optimum.exporters.openvino._traceable_decorator import traceable_check_model_inputs - else: TransformersKwargs = object @@ -100,7 +82,6 @@ from transformers.models.gpt_oss.modeling_gpt_oss import GptOssExperts if is_transformers_version(">=", "4.56"): import transformers.masking_utils - from transformers.cache_utils import DynamicLayer if is_transformers_version(">=", "4.57"): @@ -114,689 +95,6 @@ logger = logging.getLogger(__name__) -@symbolic_helper.parse_args("v", "v") -def __ior_(g, self: torch._C.Value, other: torch._C.Value) -> torch._C.Value: - return g.op("Or", self, other) - - -torch.onnx.register_custom_op_symbolic("aten::__ior__", __ior_, 14) - -if is_torch_version("<", "2.9"): - # this was fixed in torch in 2.9 https://github.com/pytorch/pytorch/pull/159973 - from torch.onnx import JitScalarType - from torch.onnx.symbolic_opset14 import _attention_scale, _causal_attention_mask - - @symbolic_helper.parse_args("v", "v", "v", "v", "f", "b", "v", "b") - def scaled_dot_product_attention( - g, - query: torch._C.Value, - key: torch._C.Value, - value: torch._C.Value, - attn_mask: torch._C.Value | None = None, - dropout_p: float = 0.0, - is_causal: bool = False, - scale: torch._C.Value | None = None, - enable_gqa: bool = False, - ): - assert (not is_causal) or ( - is_causal and symbolic_helper._is_none(attn_mask) - ), "is_causal and attn_mask cannot be set at the same time" - assert not enable_gqa, "conversion of scaled_dot_product_attention not implemented if enable_gqa is True" - - if symbolic_helper._is_none(scale): - scale = _attention_scale(g, query) - - if is_causal: - attn_mask = _causal_attention_mask(g, query, key) - - # Swap the last two axes of key - # NOTE: onnx-script has different logic here, because the attribute perms in - # transpose needs list of ints - key_shape_builtin = symbolic_helper._get_tensor_rank(key) - key_transposed_axes = list(range(key_shape_builtin)) - key_transposed_axes[-1], key_transposed_axes[-2] = (key_transposed_axes[-2], key_transposed_axes[-1]) - key_transposed = g.op("Transpose", key, perm_i=key_transposed_axes) - - # https://github.com/pytorch/pytorch/blob/12da0c70378b5be9135c6fda62a9863bce4a4818/aten/src/ATen/native/transformers/attention.cpp#L653 - # Scale q, k before matmul for stability see https://tinyurl.com/sudb9s96 for math - query_scaled = g.op("Mul", query, g.op("Sqrt", scale)) - key_transposed_scaled = g.op("Mul", key_transposed, g.op("Sqrt", scale)) - mul_qk = g.op("MatMul", query_scaled, key_transposed_scaled) - - if symbolic_helper._is_none(attn_mask): - mul_qk_add = mul_qk - attn_weight = g.op("Softmax", mul_qk_add, axis_i=-1) - elif JitScalarType.from_value(attn_mask) == JitScalarType.BOOL: - # Turn the Boolean mask to float: attn_mask.masked_fill(not attn_mask, -float('inf')) - const_zero = g.op("Constant", value_t=torch.tensor([0.0])) - const_neg_inf = g.op("Constant", value_t=torch.tensor([-float("inf")])) - attn_mask = g.op("Where", attn_mask, const_zero, const_neg_inf) - mul_qk_add = g.op("Add", mul_qk, attn_mask) - attn_weight = g.op("Softmax", mul_qk_add, axis_i=-1) - # when using scaled dot product attention with a boolean mask, we replace NaN values in attn_weight with 0.0 - attn_weight = g.op( - "Where", g.op("IsNaN", attn_weight), g.op("Constant", value_t=torch.tensor([0.0])), attn_weight - ) - elif JitScalarType.from_value(attn_mask) in ( - JitScalarType.FLOAT, - JitScalarType.HALF, - JitScalarType.BFLOAT16, - ): - mul_qk_add = g.op("Add", mul_qk, attn_mask) - attn_weight = g.op("Softmax", mul_qk_add, axis_i=-1) - else: - raise ValueError(f"Unsupported type for attn_mask: {JitScalarType.from_value(attn_mask)}") - - if dropout_p != 0: - attn_weight = g.op( - "Dropout", - attn_weight, - g.op("Constant", value_t=torch.tensor(dropout_p, dtype=torch.float)), - ) - - return g.op("MatMul", attn_weight, value) - - torch.onnx.register_custom_op_symbolic("aten::scaled_dot_product_attention", scaled_dot_product_attention, 14) - - -def patch_everywhere(attribute_name: str, patch: Any, module_name_prefix: str | None = None): - """Finds all occurrences of `attribute_name` in the loaded modules and patches them with `patch`. - - Args: - attribute_name (`str`): - The name of attribute to patch. - patch (`Any`): - The patch for the attribute. - module_name_prefix (`Optional[str]`, defaults to `None`): - If set, only module names starting with this prefix will be considered for patching. - """ - # sys.modules may be updated while being iterated over, hence the list copy. - for name in list(sys.modules): - module = sys.modules[name] - if module_name_prefix is not None and not name.startswith(module_name_prefix): - continue - if hasattr(module, attribute_name): - setattr(module, attribute_name, patch) - - -def override_arguments(args, kwargs, forward_signature, model_kwargs: dict[str, Any]): - """Override the args and kwargs with the argument values from model_kwargs, following the signature forward_signature corresponding to args and kwargs.""" - args = list(args) - - for argument in model_kwargs: - if argument in forward_signature.parameters: - argument_index = list(forward_signature.parameters.keys()).index(argument) - if argument in kwargs or len(args) <= argument_index: - kwargs[argument] = model_kwargs[argument] - else: - args[argument_index] = model_kwargs[argument] - else: - kwargs[argument] = model_kwargs[argument] - - return args, kwargs - - -def preprocess_encoder_outputs(encoder_outputs): - if is_transformers_version(">=", "4.54") and isinstance(encoder_outputs, (list, tuple)): - encoder_outputs = BaseModelOutput(*encoder_outputs) - - return encoder_outputs - - -def preprocess_past_key_values(past_key_values): - if ( - is_transformers_version(">=", "4.48") - and isinstance(past_key_values, (list, tuple)) - and isinstance(past_key_values[0], (list, tuple)) - ): - if len(past_key_values[0]) == 2: - if hasattr(DynamicCache, "from_legacy_cache"): - past_key_values = DynamicCache.from_legacy_cache(past_key_values) - else: - past_key_values = DynamicCache(past_key_values) - elif len(past_key_values[0]) == 4: - if hasattr(EncoderDecoderCache, "from_legacy_cache"): - past_key_values = EncoderDecoderCache.from_legacy_cache(past_key_values) - else: - past_key_values = EncoderDecoderCache( - DynamicCache([layer[:2] for layer in past_key_values]), - DynamicCache([layer[2:] for layer in past_key_values]), - ) - else: - raise ValueError( - f"past_key_values should have either 2 or 4 elements, but it has {len(past_key_values[0])} elements." - ) - - return past_key_values - - -def postprocess_past_key_values(past_key_values, output_names: list[str]): - if is_transformers_version(">=", "4.48") and isinstance(past_key_values, (EncoderDecoderCache, DynamicCache)): - if hasattr(past_key_values, "to_legacy_cache"): - past_key_values = past_key_values.to_legacy_cache() - elif isinstance(past_key_values, DynamicCache): - past_key_values = [(lay.keys, lay.values) for lay in past_key_values.layers] - elif isinstance(past_key_values, EncoderDecoderCache): - past_key_values = [ - (self_lay.keys, self_lay.values, cross_lay.keys, cross_lay.values) - for self_lay, cross_lay in zip( - past_key_values.self_attention_cache.layers, - past_key_values.cross_attention_cache.layers, - ) - ] - else: - raise NotImplementedError(f"Unable to serialize class {type(past_key_values)}.") - - if ( - isinstance(past_key_values, (list, tuple)) - and isinstance(past_key_values[0], (list, tuple)) - and not any("encoder.key" in output_name for output_name in output_names) - ): - past_key_values = tuple(pkv[:2] for pkv in past_key_values) - - return past_key_values - - -@dataclasses.dataclass -class PatchingSpec: - """Data class that holds patching specifications. - - Args: - o: Module / object where the op to patch is located - name: Name of the op to monkey patch - custom_op: Custom op that patches the original op - orig_op: Original op that is being patched - op_wrapper: Wrapper (optional) that wraps both the original and custom ops. - It is useful for ops that are class or static methods for instance. - """ - - o: Any - name: str - custom_op: Callable - orig_op: Callable | None = None - op_wrapper: Callable | None = None - - -# An ONNX-export-compatible version of `tensor.unfold`. Without this, we get: -# torch.onnx.errors.SymbolicValueError: Unsupported: ONNX export of operator Unfold, input size not accessible. -# See https://github.com/pytorch/pytorch/issues/81871 for more information -def onnx_compatible_unfold(input_tensor, dimension, size, step): - """Custom implementation of torch.unfold without using torch.unfold. - - Args: - input_tensor (torch.Tensor): The input tensor. - dimension (int): The dimension to unfold. - size (int): The size of each slice. - step (int): The step size between slices. - - Returns: - torch.Tensor: The unfolded tensor. - """ - # Check if dimension is within the valid range - if not (-input_tensor.dim() <= dimension < input_tensor.dim()): - raise ValueError( - f"Dimension out of range (expected to be in range of [{-input_tensor.dim()}, {input_tensor.dim() - 1}], but got {dimension})" - ) - - # Normalize negative dimension - dimension = dimension % input_tensor.dim() - - # Compute the shape of the unfolded output - input_size = input_tensor.size(dimension) - num_slices = (input_size - size) // step + 1 - - # Permute dimension to the end for easier indexing - input_tensor = input_tensor.transpose(dimension, -1) - - # Extract slices - slices = [] - for i in range(num_slices): - start = i * step - end = start + size - slices.append(input_tensor[..., start:end]) - - # Stack slices and permute dimensions back - result = torch.stack(slices, dim=-2).transpose(dimension, -2) - return result - - -# An ONNX-export-compatible version of `tensor.repeat_interleave`. -# Without this, we get the following error: https://github.com/pytorch/pytorch/issues/145100 -# NOTE: This implementation is only necessary for export with dynamo=False (dynamo=True works correctly). -# and can be removed once Optimum switches to dynamo-based exports -def onnx_compatible_repeat_interleave(input_tensor, repeats, dim=None, output_size=None): # noqa: D417 - """Custom implementation of torch.repeat_interleave without using torch.repeat_interleave. - - Args: - input_tensor (torch.Tensor): The input tensor. - repeats (int or torch.Tensor): The number of repetitions for each element. - dim (int, optional): The dimension along which to repeat. Defaults to None. - - Returns: - torch.Tensor: The repeated tensor. - """ - if isinstance(repeats, int) or (torch.is_tensor(repeats) and repeats.dim() == 0): - if dim is None: - return input_tensor.flatten().unsqueeze(1).expand(-1, repeats).flatten() - repeats = torch.full((input_tensor.shape[dim],), repeats, dtype=torch.long, device=input_tensor.device) - - if dim is None: - return onnx_compatible_repeat_interleave(input_tensor.flatten(), repeats, 0) - - if dim != 0: - input_tensor = input_tensor.transpose(0, dim) - - # Create expand mask - max_repeats = repeats.max() - expanded = input_tensor.unsqueeze(1).expand(-1, max_repeats, *input_tensor.shape[1:]) - mask = torch.arange(max_repeats, device=input_tensor.device) < repeats.unsqueeze(1) - result = expanded[mask] - - if dim != 0: - result = result.transpose(0, dim) - - return result - - -# Custom implementation of torch.linalg.matrix_norm not using torch.linalg.matrix_norm, torch.norm or torch.linalg.norm. -def onnx_compatible_linalg_norm(x, ord=2, dim=None, keepdim=False, *, dtype=None, out=None) -> torch.Tensor: - if ord != 2: - raise ValueError( - f"Only ord=2 is supported by onnx_compatible_linalg_norm, but got ord={ord}. " - "Please extend this function to support other norms." - ) - - if dim is None: - dim = (-2, -1) - - norm = torch.sqrt(torch.sum(torch.square(x), dim=dim, keepdim=keepdim)) - - if dtype is not None: - norm = norm.to(dtype) - if out is not None: - out.copy_(norm) - - return norm - - -def onnx_compatible_rms_norm(input, normalized_shape, weight=None, eps=None): - if eps is None: - eps = torch.finfo(input.dtype).eps - - axis = -len(normalized_shape) - mean_square = torch.mean(torch.square(input), dim=axis, keepdim=True) - rms = torch.sqrt(mean_square + eps) - output = input / rms - - if weight is not None: - output = output * weight - - return output - - -# A patched version of https://github.com/huggingface/transformers/blob/v4.53.2/src/transformers/masking_utils.py#L602 -# That returns a tensor of zeros with the same shape as position_ids indicating no packed sequence indices. -def find_packed_sequence_indices_patched(position_ids: torch.Tensor) -> torch.Tensor: - return torch.zeros_like(position_ids) - - -if is_transformers_version(">=", "4.53"): - _prepare_padding_mask_slice = "_slice" in inspect.signature(prepare_padding_mask).parameters -else: - _prepare_padding_mask_slice = False - - -# Custom vectorized implementation of sdpa_mask without using vmap -def _orig_sdpa_mask_without_vmap( - batch_size: int, - cache_position: torch.Tensor, - kv_length: int, - kv_offset: int = 0, - mask_function: Callable | None = None, - attention_mask: torch.Tensor | None = None, - local_size: int | None = None, - allow_is_causal_skip: bool = True, - **kwargs, -) -> torch.Tensor | None: - if mask_function is None: - mask_function = causal_mask_function - - q_length = cache_position.shape[0] - # Potentially pad the 2D mask, and slice it correctly - if _prepare_padding_mask_slice: - padding_mask = prepare_padding_mask(attention_mask, kv_length, kv_offset, _slice=False) - else: - padding_mask = prepare_padding_mask(attention_mask, kv_length, kv_offset) - - # Under specific conditions, we can avoid materializing the mask, instead relying on the `is_causal` argument - if allow_is_causal_skip and _ignore_causal_mask_sdpa(padding_mask, q_length, kv_length, kv_offset, local_size): - return None - - # Potentially add the padding 2D mask - if padding_mask is not None: - mask_function = and_masks(mask_function, padding_mask_function(padding_mask)) - - # Create broadcatable indices - device = cache_position.device - q_indices = cache_position[None, None, :, None] - head_indices = torch.arange(1, dtype=torch.long, device=device)[None, :, None, None] - batch_indices = torch.arange(batch_size, dtype=torch.long, device=device)[:, None, None, None] - kv_indices = torch.arange(kv_length, dtype=torch.long, device=device)[None, None, None, :] + kv_offset - # Apply mask function element-wise through broadcasting - causal_mask = mask_function(batch_indices, head_indices, q_indices, kv_indices) - # Expand the mask to match batch size and query length if they weren't used in the mask function - causal_mask = causal_mask.expand(batch_size, -1, q_length, kv_length) - - return causal_mask - - -# Compatibility wrapper for sdpa_mask_without_vmap from optimum. -# The installed optimum version expects (batch_size, cache_position: Tensor, kv_length, ...), -# but transformers >= 5.5 passes (batch_size, q_length: int, kv_length: int, q_offset: int, ...). -def sdpa_mask_without_vmap(batch_size, q_length=None, kv_length=None, q_offset=0, kv_offset=0, **kwargs): - import inspect - - sig = inspect.signature(_orig_sdpa_mask_without_vmap) - if is_transformers_version(">=", "5.5") and "cache_position" in sig.parameters and q_length is not None: - # Old optimum signature: (batch_size, cache_position, kv_length, kv_offset, ...) - cache_position = torch.arange(q_length, dtype=torch.long) + q_offset - kwargs.pop("q_offset", None) - kwargs.pop("allow_is_bidirectional_skip", None) - kwargs.pop("allow_torch_fix", None) - kwargs.pop("use_vmap", None) - kwargs.pop("device", None) - return _orig_sdpa_mask_without_vmap(batch_size, cache_position, kv_length, kv_offset=kv_offset, **kwargs) - else: - return _orig_sdpa_mask_without_vmap( - batch_size, q_length=q_length, kv_length=kv_length, q_offset=q_offset, kv_offset=kv_offset, **kwargs - ) - - -# Adapted from https://github.com/huggingface/transformers/blob/v4.53.0/src/transformers/masking_utils.py#L433 -def eager_mask_without_vmap(*args, **kwargs) -> torch.Tensor: - kwargs.pop("allow_is_causal_skip", None) - dtype = kwargs.get("dtype", torch.float32) - mask = sdpa_mask_without_vmap(*args, allow_is_causal_skip=False, **kwargs) - mask = torch.where(mask, torch.tensor(0.0, device=mask.device, dtype=dtype), torch.finfo(dtype).min) - return mask - - -original_triu = torch.triu -original_tril = torch.tril - - -# Custom implementation of torch.tril that doesn't fail on int32 tensors. -def onnx_compatible_tril(input_tensor: torch.Tensor, *args, **kwargs) -> torch.Tensor: - if input_tensor.dtype == torch.int32: - return original_tril(input_tensor.to(torch.int64), *args, **kwargs).to(torch.int32) - else: - return original_tril(input_tensor, *args, **kwargs) - - -# Custom implementation of torch.triu that doesn't fail on int32 tensors. -def onnx_compatible_triu(input_tensor: torch.Tensor, *args, **kwargs) -> torch.Tensor: - if input_tensor.dtype == torch.int32: - return original_triu(input_tensor.to(torch.int64), *args, **kwargs).to(torch.int32) - else: - return original_triu(input_tensor, *args, **kwargs) - - -original_scaled_dot_product_attention = torch.nn.functional.scaled_dot_product_attention - - -# A patched `torch.nn.functional.scaled_dot_product_attention` that doesn't fail during tracing -# from passing `is_causal` as a tensor (which is usually obtained with tensor shapes comparisons). -def traceable_scaled_dot_product_attention( - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - attn_mask: torch.Tensor | None = None, - dropout_p: float = 0.0, - is_causal: bool = False, - **kwargs, -) -> torch.Tensor: - if isinstance(is_causal, torch.Tensor): - is_causal = is_causal.item() - - if "enable_gqa" in kwargs: - kwargs.pop("enable_gqa") - - attn_weights = original_scaled_dot_product_attention( - query=query, key=key, value=value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, **kwargs - ) - - return attn_weights - - -# No-op bfloat16 casting to avoid issues with legacy ONNX export which cast to complex128 -def noop_bfloat16_casting(self): - return self - - -original_movedim = torch.Tensor.movedim - - -def onnx_compatible_movedim(self: torch.Tensor, dim1, dim2) -> torch.Tensor: - dim = self.dim() - if dim1 < 0: - dim1 += dim - if dim2 < 0: - dim2 += dim - return original_movedim(self, dim1, dim2) - - -def patched_dynamic_layer_update( - self, key_states: torch.Tensor, value_states: torch.Tensor, cache_kwargs: dict[str, Any] | None = None -) -> tuple[torch.Tensor, torch.Tensor]: - if self.keys is None: - self.keys = key_states - self.values = value_states - self.device = key_states.device - self.dtype = key_states.dtype - self.is_initialized = True - else: - self.keys = torch.cat([self.keys, key_states], dim=-2) - self.values = torch.cat([self.values, value_states], dim=-2) - return self.keys, self.values - - -UNSUPPORTED_OPS_PATCHING_SPEC = [ - PatchingSpec(torch, "tril", onnx_compatible_tril, torch.tril), - PatchingSpec(torch, "triu", onnx_compatible_triu, torch.triu), - PatchingSpec(torch, "rms_norm", onnx_compatible_rms_norm, torch.rms_norm), - PatchingSpec(torch.Tensor, "unfold", onnx_compatible_unfold, torch.Tensor.unfold), - PatchingSpec(torch.linalg, "norm", onnx_compatible_linalg_norm, torch.linalg.norm), - PatchingSpec(torch.Tensor, "bfloat16", noop_bfloat16_casting, torch.Tensor.bfloat16), - PatchingSpec(torch.Tensor, "movedim", onnx_compatible_movedim, torch.Tensor.movedim), - PatchingSpec(torch.Tensor, "repeat_interleave", onnx_compatible_repeat_interleave, torch.Tensor.repeat_interleave), - # TracerWarning: Using len to get tensor shape might cause the trace to be incorrect. Recommended usage would be tensor.shape[0]. Passing a tensor of different shape might lead to errors or silently give incorrect results. - PatchingSpec(torch.Tensor, "__len__", lambda x: x.shape[0], torch.Tensor.__len__), - PatchingSpec( - torch.nn.functional, - "scaled_dot_product_attention", - traceable_scaled_dot_product_attention, - torch.nn.functional.scaled_dot_product_attention, - ), -] - - -class ModelPatcher: - def __init__( - self, - config: "OnnxConfig", - model: PreTrainedModel, - model_kwargs: dict[str, Any] | None = None, - ): - self._model = model - - patching_specs = config.PATCHING_SPECS or [] - patching_specs.extend(UNSUPPORTED_OPS_PATCHING_SPEC) - - self._patching_specs = [] - for spec in patching_specs: - final_spec = spec - if spec.orig_op is None: - final_spec = dataclasses.replace(spec, orig_op=getattr(spec.o, spec.name)) - self._patching_specs.append(final_spec) - - self.orig_forward_name = "forward" if hasattr(self._model, "forward") else "call" - self.orig_forward = getattr(self._model, self.orig_forward_name) - - if is_transformers_version(">=", "4.54") and hasattr(self.orig_forward, "__wrapped__"): - # the original check_model_inputs has some failing cases that we fix in traceable_check_model_inputs - # we fix those issues in a PR in transformers https://github.com/huggingface/transformers/pull/40811 - # issues are: support for positional args (use_cache for instance) and fix for _CAN_RECORD_REGISTRY - # explicitly mapping to None for some models - self.orig_forward = types.MethodType( - traceable_check_model_inputs(self.orig_forward.__wrapped__), self._model - ) - - self.real_config = config - use_cache = getattr(self.real_config, "use_past", False) - self.model_kwargs = model_kwargs if model_kwargs is not None else {} - - for module in self._model.modules(): - if hasattr(module, "config") and hasattr(module.config, "use_cache"): - module.config.use_cache = use_cache - - @functools.wraps(self.orig_forward) - def patched_forward(*args, **kwargs): - signature = inspect.signature(self.orig_forward) - args, kwargs = override_arguments(args, kwargs, signature, model_kwargs=self.model_kwargs) - - # Transformers doesn't always respect the config.use_cache attribute - # there are cases where setting use_cache to true in every config and - # subconfig of a model still doesn't enable past_key_values in the outputs (gemma3) - # Explicitly setting the use_cache argument of the forward method seems to be the most reliable way - if "use_cache" in signature.parameters: - use_cache_index = list(signature.parameters.keys()).index("use_cache") - if use_cache_index < len(args): - args[use_cache_index] = use_cache - elif "use_cache" in kwargs: - kwargs["use_cache"] = use_cache - - if "past_key_values" in signature.parameters: - # Most models require past_key_values to be a cache instance instead of a tuple now - pkv_index = list(signature.parameters.keys()).index("past_key_values") - if pkv_index < len(args) and args[pkv_index] is not None: - args[pkv_index] = preprocess_past_key_values(args[pkv_index]) - elif kwargs.get("past_key_values") is not None: - kwargs["past_key_values"] = preprocess_past_key_values(kwargs["past_key_values"]) - - if "encoder_outputs" in signature.parameters: - # Some encoder-decoder models started to not accept encoder_outputs as tuple (e.g. moonshine) - encoder_outputs_index = list(signature.parameters.keys()).index("encoder_outputs") - if encoder_outputs_index < len(args) and args[encoder_outputs_index] is not None: - args[encoder_outputs_index] = preprocess_encoder_outputs(args[encoder_outputs_index]) - elif kwargs.get("encoder_outputs") is not None: - kwargs["encoder_outputs"] = preprocess_encoder_outputs(kwargs["encoder_outputs"]) - - outputs = self.orig_forward(*args, **kwargs) - - # This code block handles different cases of the filtered_outputs input to align it with the expected - # format of outputs. It is common for the output type of a model to vary, such as tensor, list, - # tuple, etc. For Transformers models, the output is encapsulated in a ModelOutput object that - # contains the output names of the model. In the case of Timm classification models, the output - # is of type tensor. By default, it is assumed that the output names mentioned in the ONNX config - # match the outputs in order. - filtered_outputs = {} - output_names = list(config.outputs.keys()) - if isinstance(outputs, dict): - for name, value in outputs.items(): - onnx_output_name = config.torch_to_onnx_output_map.get(name, name) - if ( - onnx_output_name in output_names - or (use_cache and name.startswith("past_key_values")) - or any(key.startswith(onnx_output_name) for key in output_names) - ): - filtered_outputs[name] = value - elif isinstance(outputs, (list, tuple)): - filtered_outputs = dict(zip(output_names, outputs)) - else: - if len(output_names) > 1: - num_outputs = len(output_names) - output_names_str = ", ".join(output_names) - raise ValueError( - f"{config.__class__.__name__} expects the model to return {num_outputs} outputs: {output_names_str}, " - f"but the it returned a single output of type {type(outputs)}. Please make sure either that the model " - "returns all the expected outputs, or that the ONNX config is correctly defined with the expected outputs." - ) - output_name = output_names[0] - filtered_outputs[output_name] = outputs - - if filtered_outputs.get("past_key_values") is not None: - filtered_outputs["past_key_values"] = postprocess_past_key_values( - filtered_outputs["past_key_values"], output_names=output_names - ) - - return filtered_outputs - - self.patched_forward = patched_forward - - def patch_ops(self): - for spec in self._patching_specs: - custom_op = spec.custom_op if spec.op_wrapper is None else spec.op_wrapper(spec.custom_op) - setattr(spec.o, spec.name, custom_op) - - def restore_ops(self): - for spec in self._patching_specs: - orig_op = spec.orig_op if spec.op_wrapper is None else spec.op_wrapper(spec.orig_op) - setattr(spec.o, spec.name, orig_op) - - def __enter__(self): - self.patch_ops() - setattr(self._model, self.orig_forward_name, self.patched_forward) - - # This is a workaround for the Cache class in transformers, we replace it - # with traceable cache is because the original one used in transformers - # inherited from nn.Module (for a couple versions), which can't be traced as input. - if is_transformers_version(">=", "4.44") and is_transformers_version("<", "4.50"): - self.original_cache_class = transformers.cache_utils.Cache - transformers.cache_utils.Cache = TraceableCache - - # This is a workaround for mask generation in transformers >= 4.53. - # The masking process uses vmap which is not traceable by TorchScript. - if is_transformers_version(">=", "4.53"): - ALL_MASK_ATTENTION_FUNCTIONS.register("sdpa", sdpa_mask_without_vmap) - ALL_MASK_ATTENTION_FUNCTIONS.register("eager", eager_mask_without_vmap) - - # This is a workaround for the find_packed_sequence_indices function in transformers which - # should only return a tensor of zeros with the same shape as position_ids indicating no packed sequence indices. - # The function uses torch.diff which is not traceable by TorchScript. - if is_transformers_version(">=", "4.53.1"): - self.original_find_packed_sequence_indices = find_packed_sequence_indices - transformers.masking_utils.find_packed_sequence_indices = find_packed_sequence_indices_patched - - # Starting from transformers 4.56.0, DynamicCache uses DynamicLayer which has an update method - # that uses torch.cat to concatenate an empty tensor with the key/value states during the first call. - # This causes issues during TorchScript tracing. - if is_transformers_version(">=", "4.56"): - self.original_dynamic_layer_update = DynamicLayer.update - DynamicLayer.update = patched_dynamic_layer_update - - def __exit__(self, exc_type, exc_value, traceback): - self.restore_ops() - setattr(self._model, self.orig_forward_name, self.orig_forward) - - if is_transformers_version(">=", "4.44") and is_transformers_version("<", "4.50"): - transformers.cache_utils.Cache = self.original_cache_class - - if is_transformers_version(">=", "4.53"): - ALL_MASK_ATTENTION_FUNCTIONS.register("sdpa", sdpa_mask) - ALL_MASK_ATTENTION_FUNCTIONS.register("eager", eager_mask) - - if is_transformers_version(">=", "4.53.1"): - transformers.masking_utils.find_packed_sequence_indices = self.original_find_packed_sequence_indices - - if is_transformers_version(">=", "4.56"): - DynamicLayer.update = self.original_dynamic_layer_update - - def __call__(self, *args, **kwargs): - if getattr(self._model, self.orig_forward_name) is self.orig_forward: - logger.warning("Running the non-patched model") - return self._model(*args, **kwargs) - - class BigBirdPegasusModelPatcher(ModelPatcher): def __enter__(self): super().__enter__() From 289617fb1f7f9ea1e99c8a5f4f678d4236b3cebd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CElla?= Date: Wed, 27 May 2026 16:10:26 +0200 Subject: [PATCH 09/40] move configs --- optimum/exporters/openvino/model_configs.py | 74 ++++++++++----------- optimum/exporters/openvino/model_patcher.py | 2 +- 2 files changed, 38 insertions(+), 38 deletions(-) diff --git a/optimum/exporters/openvino/model_configs.py b/optimum/exporters/openvino/model_configs.py index 9ebab904df..a0550e7300 100644 --- a/optimum/exporters/openvino/model_configs.py +++ b/optimum/exporters/openvino/model_configs.py @@ -4952,6 +4952,25 @@ class DebertaV2OpenVINOConfig(DebertaOpenVINOConfig): pass +@register_in_tasks_manager("hubert", *["feature-extraction", "automatic-speech-recognition", "audio-classification"]) +class HubertOpenVINOConfig(AudioOnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedConfig + + @property + def outputs(self) -> dict: + outputs = super().outputs + + # Hubert output formula adapted from: + # https://github.com/huggingface/transformers/blob/v4.55.2/src/transformers/models/hubert/modeling_hubert.py#L721 + if self.task == "automatic-speech-recognition": + sequence_length = "sequence_length" + for kernel_size, stride in zip(self._config.conv_kernel, self._config.conv_stride): + sequence_length = f"( {sequence_length} - {kernel_size} ) // {stride} + 1" + outputs["logits"] = {0: "batch_size", 1: sequence_length} + + return outputs + + @register_in_tasks_manager( "data2vec-audio", *[ @@ -4972,6 +4991,24 @@ class Data2VecTextOpenVINOConfig(DistilBertOpenVINOConfig): MAX_TRANSFORMERS_VERSION = "4.57.6" +@register_in_tasks_manager("vit", *["feature-extraction", "image-classification", "masked-im"]) +class ViTOpenVINOConfig(VisionOnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedVisionConfig + + @property + def inputs(self) -> dict: + return {"pixel_values": {0: "batch_size", 2: "height", 3: "width"}} + + @property + def outputs(self) -> dict: + common_outputs = super().outputs + + if self.task == "feature-extraction": + common_outputs["last_hidden_state"] = {0: "batch_size"} + + return common_outputs + + @register_in_tasks_manager("data2vec-vision", *["feature-extraction", "image-classification"]) class Data2VecVisionOpenVINOConfig(ViTOpenVINOConfig): pass @@ -5116,24 +5153,6 @@ class SwinOpenVINOConfig(ViTOpenVINOConfig): pass -@register_in_tasks_manager("vit", *["feature-extraction", "image-classification", "masked-im"]) -class ViTOpenVINOConfig(VisionOnnxConfig): - NORMALIZED_CONFIG_CLASS = NormalizedVisionConfig - - @property - def inputs(self) -> dict: - return {"pixel_values": {0: "batch_size", 2: "height", 3: "width"}} - - @property - def outputs(self) -> dict: - common_outputs = super().outputs - - if self.task == "feature-extraction": - common_outputs["last_hidden_state"] = {0: "batch_size"} - - return common_outputs - - @register_in_tasks_manager("convnext", *["feature-extraction", "image-classification"]) class ConvNextOpenVINOConfig(ViTOpenVINOConfig): pass @@ -5172,25 +5191,6 @@ class Wav2Vec2ConformerOpenVINOConfig(HubertOpenVINOConfig): pass -@register_in_tasks_manager("hubert", *["feature-extraction", "automatic-speech-recognition", "audio-classification"]) -class HubertOpenVINOConfig(AudioOnnxConfig): - NORMALIZED_CONFIG_CLASS = NormalizedConfig - - @property - def outputs(self) -> dict: - outputs = super().outputs - - # Hubert output formula adapted from: - # https://github.com/huggingface/transformers/blob/v4.55.2/src/transformers/models/hubert/modeling_hubert.py#L721 - if self.task == "automatic-speech-recognition": - sequence_length = "sequence_length" - for kernel_size, stride in zip(self._config.conv_kernel, self._config.conv_stride): - sequence_length = f"( {sequence_length} - {kernel_size} ) // {stride} + 1" - outputs["logits"] = {0: "batch_size", 1: sequence_length} - - return outputs - - @register_in_tasks_manager("sew", *["feature-extraction", "automatic-speech-recognition", "audio-classification"]) class SEWOpenVINOConfig(HubertOpenVINOConfig): pass diff --git a/optimum/exporters/openvino/model_patcher.py b/optimum/exporters/openvino/model_patcher.py index 0087b10358..ee4f9ff54b 100644 --- a/optimum/exporters/openvino/model_patcher.py +++ b/optimum/exporters/openvino/model_patcher.py @@ -48,7 +48,7 @@ from transformers.utils import ModelOutput from optimum.exporters.openvino._ov_ops import convert_recurrent_attention_cell -from optimum.exporters.openvino.base_patcher import UNSUPPORTED_OPS_PATCHING_SPEC, ModelPatcher +from optimum.exporters.openvino.base_patcher import UNSUPPORTED_OPS_PATCHING_SPEC, ModelPatcher, sdpa_mask_without_vmap from optimum.intel.utils.import_utils import ( is_diffusers_version, is_openvino_version, From 3de4f6ac9804f91f52067ba7dbcb17ffc831f3ef Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CElla?= Date: Wed, 27 May 2026 16:14:06 +0200 Subject: [PATCH 10/40] rename for clarity --- optimum/exporters/openvino/base.py | 2 +- optimum/exporters/openvino/model_patcher.py | 2 +- .../exporters/openvino/{base_patcher.py => patching_utils.py} | 0 3 files changed, 2 insertions(+), 2 deletions(-) rename optimum/exporters/openvino/{base_patcher.py => patching_utils.py} (100%) diff --git a/optimum/exporters/openvino/base.py b/optimum/exporters/openvino/base.py index 29c483870b..5eee6301fe 100644 --- a/optimum/exporters/openvino/base.py +++ b/optimum/exporters/openvino/base.py @@ -30,7 +30,7 @@ from transformers import PretrainedConfig, PreTrainedModel from optimum.exporters.base import ExporterConfig -from optimum.exporters.openvino.base_patcher import ModelPatcher, PatchingSpec +from optimum.exporters.openvino.patching_utils import ModelPatcher, PatchingSpec from optimum.utils import DEFAULT_DUMMY_SHAPES, DummyInputGenerator, DummySeq2SeqPastKeyValuesGenerator, logging from optimum.utils.doc import add_dynamic_docstring from optimum.utils.import_utils import ( diff --git a/optimum/exporters/openvino/model_patcher.py b/optimum/exporters/openvino/model_patcher.py index ee4f9ff54b..ca8ce90702 100644 --- a/optimum/exporters/openvino/model_patcher.py +++ b/optimum/exporters/openvino/model_patcher.py @@ -48,7 +48,7 @@ from transformers.utils import ModelOutput from optimum.exporters.openvino._ov_ops import convert_recurrent_attention_cell -from optimum.exporters.openvino.base_patcher import UNSUPPORTED_OPS_PATCHING_SPEC, ModelPatcher, sdpa_mask_without_vmap +from optimum.exporters.openvino.patching_utils import UNSUPPORTED_OPS_PATCHING_SPEC, ModelPatcher, sdpa_mask_without_vmap from optimum.intel.utils.import_utils import ( is_diffusers_version, is_openvino_version, diff --git a/optimum/exporters/openvino/base_patcher.py b/optimum/exporters/openvino/patching_utils.py similarity index 100% rename from optimum/exporters/openvino/base_patcher.py rename to optimum/exporters/openvino/patching_utils.py From b820ff548b50aadada93b9e494396e073b86626c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CElla?= Date: Wed, 27 May 2026 18:10:07 +0200 Subject: [PATCH 11/40] fix imports --- optimum/exporters/openvino/convert.py | 2 +- optimum/exporters/openvino/model_patcher.py | 2 +- optimum/intel/openvino/modeling_timm.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/optimum/exporters/openvino/convert.py b/optimum/exporters/openvino/convert.py index b65a841e17..152b8d9ca7 100644 --- a/optimum/exporters/openvino/convert.py +++ b/optimum/exporters/openvino/convert.py @@ -92,7 +92,7 @@ if TYPE_CHECKING: - from optimum.exporters.onnx.base import OnnxConfig + from optimum.exporters.openvino.base import OnnxConfig from optimum.intel.openvino.configuration import OVConfig diff --git a/optimum/exporters/openvino/model_patcher.py b/optimum/exporters/openvino/model_patcher.py index ca8ce90702..8af1d07cfc 100644 --- a/optimum/exporters/openvino/model_patcher.py +++ b/optimum/exporters/openvino/model_patcher.py @@ -5376,7 +5376,7 @@ def gemma4_lm_forward( logits_to_keep: Union[int, torch.Tensor] = 0, **lm_kwargs, ): - from optimum.exporters.onnx.model_patcher import preprocess_past_key_values + from optimum.exporters.openvino.patching_utils import preprocess_past_key_values output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions output_hidden_states = ( diff --git a/optimum/intel/openvino/modeling_timm.py b/optimum/intel/openvino/modeling_timm.py index f2566c8c4a..ce60323aac 100644 --- a/optimum/intel/openvino/modeling_timm.py +++ b/optimum/intel/openvino/modeling_timm.py @@ -40,7 +40,7 @@ from transformers.modeling_outputs import ImageClassifierOutput from transformers.utils import TensorType -from optimum.exporters.onnx.config import VisionOnnxConfig +from optimum.exporters.openvino.config import VisionOnnxConfig from optimum.utils import NormalizedVisionConfig from .utils import _is_timm_ov_dir From 802397ca130c8176ee209ca8a4b363003bbb7121 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CElla?= Date: Wed, 27 May 2026 18:14:56 +0200 Subject: [PATCH 12/40] fix --- optimum/exporters/openvino/model_patcher.py | 6 +++++- tests/openvino/test_export.py | 4 ++-- tests/openvino/utils_tests.py | 7 +++++++ 3 files changed, 14 insertions(+), 3 deletions(-) diff --git a/optimum/exporters/openvino/model_patcher.py b/optimum/exporters/openvino/model_patcher.py index 8af1d07cfc..b83fc6f597 100644 --- a/optimum/exporters/openvino/model_patcher.py +++ b/optimum/exporters/openvino/model_patcher.py @@ -48,7 +48,11 @@ from transformers.utils import ModelOutput from optimum.exporters.openvino._ov_ops import convert_recurrent_attention_cell -from optimum.exporters.openvino.patching_utils import UNSUPPORTED_OPS_PATCHING_SPEC, ModelPatcher, sdpa_mask_without_vmap +from optimum.exporters.openvino.patching_utils import ( + UNSUPPORTED_OPS_PATCHING_SPEC, + ModelPatcher, + sdpa_mask_without_vmap, +) from optimum.intel.utils.import_utils import ( is_diffusers_version, is_openvino_version, diff --git a/tests/openvino/test_export.py b/tests/openvino/test_export.py index fbca6ace56..93c60b8781 100644 --- a/tests/openvino/test_export.py +++ b/tests/openvino/test_export.py @@ -24,11 +24,11 @@ MODEL_NAMES, OPENVINO_DEVICE, REMOTE_CODE_MODELS, + SDPA_ARCHS_ONNX_EXPORT_NOT_SUPPORTED, ) -from optimum.exporters.onnx.constants import SDPA_ARCHS_ONNX_EXPORT_NOT_SUPPORTED -from optimum.exporters.onnx.model_configs import BertOnnxConfig from optimum.exporters.openvino import export_from_model, main_export +from optimum.exporters.openvino.model_configs import BertOnnxConfig from optimum.exporters.tasks import TasksManager from optimum.intel import ( OVFluxPipeline, diff --git a/tests/openvino/utils_tests.py b/tests/openvino/utils_tests.py index 8a842c3159..5edbd3a8d5 100644 --- a/tests/openvino/utils_tests.py +++ b/tests/openvino/utils_tests.py @@ -583,6 +583,13 @@ def _create_tiny_kokoro_model(): REMOTE_CODE_MODELS += ("afmoe",) +SDPA_ARCHS_ONNX_EXPORT_NOT_SUPPORTED = [ + "bart", + "musicgen", + "whisper", +] + + def get_num_quantized_nodes(model): num_fake_nodes = 0 types_map = { From 67fc3b935ac1cb015e614856954e46c91e6b7a07 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CElla?= Date: Wed, 27 May 2026 18:18:02 +0200 Subject: [PATCH 13/40] fix --- tests/openvino/test_export.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/openvino/test_export.py b/tests/openvino/test_export.py index 93c60b8781..45b606c7e1 100644 --- a/tests/openvino/test_export.py +++ b/tests/openvino/test_export.py @@ -28,7 +28,7 @@ ) from optimum.exporters.openvino import export_from_model, main_export -from optimum.exporters.openvino.model_configs import BertOnnxConfig +from optimum.exporters.openvino.model_configs import BertOpenVINOConfig from optimum.exporters.tasks import TasksManager from optimum.intel import ( OVFluxPipeline, @@ -369,7 +369,7 @@ def test_compare_openvino_onnx_supported_architectures(self): class CustomExportModelTest(unittest.TestCase): def test_custom_export_config_model(self): - class BertOnnxConfigWithPooler(BertOnnxConfig): + class BertOpenVINOConfigWithPooler(BertOpenVINOConfig): @property def outputs(self): if self.task == "feature-extraction-with-pooler": @@ -386,7 +386,7 @@ def outputs(self): model_id = "sentence-transformers/all-MiniLM-L6-v2" config = AutoConfig.from_pretrained(model_id) - custom_export_configs = {"model": BertOnnxConfigWithPooler(config, task=custom_task)} + custom_export_configs = {"model": BertOpenVINOConfigWithPooler(config, task=custom_task)} with TemporaryDirectory() as tmpdirname: main_export( From 2c574e888d4becf013a89f0a4bcef7ca01458083 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CElla?= Date: Wed, 27 May 2026 18:25:28 +0200 Subject: [PATCH 14/40] deprecate ONNX_SUPPORTED_ARCHITECTURES --- optimum/exporters/openvino/convert.py | 43 +++++++++++---------- optimum/exporters/openvino/model_configs.py | 22 ----------- 2 files changed, 22 insertions(+), 43 deletions(-) diff --git a/optimum/exporters/openvino/convert.py b/optimum/exporters/openvino/convert.py index 152b8d9ca7..b2b0c640e6 100644 --- a/optimum/exporters/openvino/convert.py +++ b/optimum/exporters/openvino/convert.py @@ -29,6 +29,23 @@ from openvino import Model, save_model from openvino.exceptions import OVTypeError from openvino.tools.ovc import convert_model +from optimum.exporters.openvino.utils import ( + MULTI_MODAL_TEXT_GENERATION_MODELS, + ONNX_SUPPORTED_ARCHITECTURES, + OV_XML_FILE_NAME, + _get_dynamic_shapes_info, + _get_input_info, + _get_kokoro_submodels_fn_and_export_configs, + _get_model_dtype, + _get_open_clip_submodels_fn_and_export_configs, + _normalize_dummy_inputs, + allow_skip_tracing_check, + clear_class_registry, + remove_none_from_dummy_inputs, + save_config, + save_preprocessors, + set_simplified_chat_template, +) from optimum.exporters.tasks import TasksManager from optimum.exporters.utils import ( DECODER_NAME, @@ -50,11 +67,12 @@ _torch_version, _transformers_version, compare_versions, + is_diffusers_available, + is_nncf_available, is_transformers_version, ) -from optimum.utils import DEFAULT_DUMMY_SHAPES, is_diffusers_available +from optimum.utils import DEFAULT_DUMMY_SHAPES -from ...intel.utils.import_utils import is_nncf_available from ...intel.utils.modeling_utils import _infer_library_from_model_or_model_class from .stateful import ( ensure_export_task_support_stateful, @@ -62,23 +80,6 @@ ensure_stateful_is_available, patch_stateful, ) -from .utils import ( - MULTI_MODAL_TEXT_GENERATION_MODELS, - ONNX_SUPPORTED_ARCHITECTURES, - OV_XML_FILE_NAME, - _get_dynamic_shapes_info, - _get_input_info, - _get_kokoro_submodels_fn_and_export_configs, - _get_model_dtype, - _get_open_clip_submodels_fn_and_export_configs, - _normalize_dummy_inputs, - allow_skip_tracing_check, - clear_class_registry, - remove_none_from_dummy_inputs, - save_config, - save_preprocessors, - set_simplified_chat_template, -) logger = logging.getLogger(__name__) @@ -606,8 +607,8 @@ def export_from_model( model_type = getattr(model.config, "model_type", None) or "" if model_type in ONNX_SUPPORTED_ARCHITECTURES: - logger.warning( - f"The OpenVINO export of {model_type} models is not officially supported by optimum-intel, export at your own risks." + raise ValueError( + f"The OpenVINO export of {model_type} models is not officially supported by optimum-intel and has been deprecated." ) custom_architecture = library_name == "transformers" and model_type not in TasksManager._SUPPORTED_MODEL_TYPE diff --git a/optimum/exporters/openvino/model_configs.py b/optimum/exporters/openvino/model_configs.py index a0550e7300..20794f4fc1 100644 --- a/optimum/exporters/openvino/model_configs.py +++ b/optimum/exporters/openvino/model_configs.py @@ -14,7 +14,6 @@ import enum import logging -from copy import deepcopy from typing import Any, Dict, List, Optional, Union from transformers import AutoConfig, PretrainedConfig, PreTrainedModel @@ -168,7 +167,6 @@ Zamba2ModelPatcher, _get_model_attribute, ) -from optimum.exporters.openvino.utils import ONNX_SUPPORTED_ARCHITECTURES from optimum.exporters.tasks import TasksManager from optimum.intel.utils.import_utils import ( is_diffusers_available, @@ -333,26 +331,6 @@ def init_model_configs(): TasksManager._DIFFUSERS_TASKS_TO_MODEL_MAPPINGS["text-to-video"] = {} TasksManager._DIFFUSERS_TASKS_TO_MODEL_MAPPINGS["text-to-video"]["ltx-video"] = "LTXPipeline" - # TODO: add warning to state that architectures from ONNX_SUPPORTED_ARCHITECTURES are deprecated - - supported_model_types = [ - "_SUPPORTED_MODEL_TYPE", - "_DIFFUSERS_SUPPORTED_MODEL_TYPE", - "_TIMM_SUPPORTED_MODEL_TYPE", - "_SENTENCE_TRANSFORMERS_SUPPORTED_MODEL_TYPE", - ] - # TODO: remove once models from ONNX_SUPPORTED_ARCHITECTURES are deprecated (optimum-intel v1.29) - for supported_models_config in supported_model_types: - supported_models = getattr(TasksManager, supported_models_config) - for model, export_configs in supported_models.items(): - # adding only the architectures that are already supported via optimum-onnx v0.1.0 - if "onnx" not in export_configs or model not in ONNX_SUPPORTED_ARCHITECTURES: - continue - onnx_config = export_configs["onnx"] - supported_models[model]["openvino"] = deepcopy(onnx_config) - - setattr(TasksManager, supported_models_config, supported_models) - init_model_configs() From 397ef7859904b13d5c7ee223e6b7cc82b69037b7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CElla?= Date: Wed, 27 May 2026 18:41:51 +0200 Subject: [PATCH 15/40] remove unnecessary OVMistralDummyPastKeyValuesGenerator --- .../exporters/openvino/input_generators.py | 39 +------------------ optimum/exporters/openvino/model_configs.py | 6 +-- optimum/exporters/openvino/model_patcher.py | 26 ++----------- optimum/exporters/openvino/patching_utils.py | 1 + 4 files changed, 7 insertions(+), 65 deletions(-) diff --git a/optimum/exporters/openvino/input_generators.py b/optimum/exporters/openvino/input_generators.py index 036ea52042..238e83a4ff 100644 --- a/optimum/exporters/openvino/input_generators.py +++ b/optimum/exporters/openvino/input_generators.py @@ -344,6 +344,7 @@ def __init__( **kwargs, ): super().__init__( + task=task, normalized_config=normalized_config, batch_size=batch_size, @@ -400,44 +401,6 @@ def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int ] -class OVMistralDummyPastKeyValuesGenerator(MistralDummyPastKeyValuesGenerator): - def __init__( - self, - task: str, - normalized_config: NormalizedTextConfig, - batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], - sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], - random_batch_size_range: Optional[Tuple[int, int]] = None, - random_sequence_length_range: Optional[Tuple[int, int]] = None, - **kwargs, - ): - super().__init__( - task=task, - normalized_config=normalized_config, - batch_size=batch_size, - sequence_length=sequence_length, - random_batch_size_range=random_batch_size_range, - random_sequence_length_range=random_sequence_length_range, - **kwargs, - ) - self.head_dim = getattr(normalized_config, "head_dim", 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) - ] - - class Gemma4DummyPastKeyValuesGenerator(DummyPastKeyValuesGenerator): def __init__( self, diff --git a/optimum/exporters/openvino/model_configs.py b/optimum/exporters/openvino/model_configs.py index 20794f4fc1..78f71c0eae 100644 --- a/optimum/exporters/openvino/model_configs.py +++ b/optimum/exporters/openvino/model_configs.py @@ -71,7 +71,6 @@ MambaCacheDummyInputGenerator, OVFalconDummyPastKeyValuesGenerator, OVMiniCPM3DummyPastKeyValuesGenerator, - OVMistralDummyPastKeyValuesGenerator, PooledProjectionsDummyInputGenerator, Qwen3_5DummyPastKeyValuesGenerator, Qwen3ASRDummySeq2SeqPastKeyValuesGenerator, @@ -191,7 +190,6 @@ DummyVisionEncoderDecoderPastKeyValuesGenerator, DummyVisionInputGenerator, GemmaDummyPastKeyValuesGenerator, - GPTBigCodeDummyPastKeyValuesGenerator, MistralDummyPastKeyValuesGenerator, PerceiverDummyInputGenerator, T5DummySeq2SeqPastKeyValuesGenerator, @@ -1208,9 +1206,9 @@ class ArcticOpenVINOConfig(MixtralOpenVINOConfig): class MistralOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args(num_key_value_heads="num_key_value_heads", allow_new=True) DUMMY_INPUT_GENERATOR_CLASSES = ( - OVMistralDummyPastKeyValuesGenerator, + MistralDummyPastKeyValuesGenerator, ) + TextDecoderOnnxConfig.DUMMY_INPUT_GENERATOR_CLASSES - DUMMY_PKV_GENERATOR_CLASS = OVMistralDummyPastKeyValuesGenerator + DUMMY_PKV_GENERATOR_CLASS = MistralDummyPastKeyValuesGenerator _MODEL_PATCHER = MistralModelPatcher diff --git a/optimum/exporters/openvino/model_patcher.py b/optimum/exporters/openvino/model_patcher.py index b83fc6f597..62428ae553 100644 --- a/optimum/exporters/openvino/model_patcher.py +++ b/optimum/exporters/openvino/model_patcher.py @@ -48,9 +48,12 @@ from transformers.utils import ModelOutput from optimum.exporters.openvino._ov_ops import convert_recurrent_attention_cell +from optimum.exporters.openvino.base import OnnxConfig from optimum.exporters.openvino.patching_utils import ( UNSUPPORTED_OPS_PATCHING_SPEC, ModelPatcher, + postprocess_past_key_values, + preprocess_past_key_values, sdpa_mask_without_vmap, ) from optimum.intel.utils.import_utils import ( @@ -269,25 +272,6 @@ def __exit__(self, exc_type, exc_value, traceback): CLIPSdpaAttention.forward = self.original_sdpa_forward -class Qwen3MoeModelPatcher(ModelPatcher): - def __enter__(self): - super().__enter__() - - # This is a workaround for the Qwen3 Moe Sparse block that is not compatible with ONNX export. - # The forward method of the Moe Sparse block is patched to avoid looping only on the experts that are selected - # by the router, which fails during execution in ONNX Runtime. - # TODO: investigate more on this issue. - if is_transformers_version(">=", "4.53"): - self.original_moe_forward = Qwen3MoeSparseMoeBlock.forward - Qwen3MoeSparseMoeBlock.forward = qwen3_moe_forward_patched - - def __exit__(self, exc_type, exc_value, traceback): - super().__exit__(exc_type, exc_value, traceback) - - if is_transformers_version(">=", "4.53"): - Qwen3MoeSparseMoeBlock.forward = self.original_moe_forward - - # This is a traceable version of the original function, # the original results in a constant integer due to the use of int(expr) def _get_feat_extract_output_lengths_patched(self, input_lengths: torch.LongTensor): @@ -5380,8 +5364,6 @@ def gemma4_lm_forward( logits_to_keep: Union[int, torch.Tensor] = 0, **lm_kwargs, ): - from optimum.exporters.openvino.patching_utils import preprocess_past_key_values - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions output_hidden_states = ( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states @@ -9282,8 +9264,6 @@ def patched_encoder_forward(inputs_embeds, attention_mask=None, pixel_position_i **kwargs, ) - from transformers.modeling_outputs import BaseModelOutputWithPast - return BaseModelOutputWithPast(last_hidden_state=hidden_states) self._orig_encoder_forward = orig_encoder_forward diff --git a/optimum/exporters/openvino/patching_utils.py b/optimum/exporters/openvino/patching_utils.py index 5e1d45923d..cd4bb2d41c 100644 --- a/optimum/exporters/openvino/patching_utils.py +++ b/optimum/exporters/openvino/patching_utils.py @@ -25,6 +25,7 @@ from torch.onnx import symbolic_helper from transformers import PreTrainedModel from transformers.cache_utils import DynamicCache, EncoderDecoderCache +from transformers.modeling_outputs import BaseModelOutput from optimum.intel.utils.import_utils import is_torch_version, is_transformers_version From a272f93df7d6e15d867245feaeea00851be31eb4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CElla?= Date: Wed, 27 May 2026 18:48:26 +0200 Subject: [PATCH 16/40] add missing --- .../exporters/openvino/input_generators.py | 1 - optimum/exporters/openvino/model_patcher.py | 57 +++++++------------ 2 files changed, 22 insertions(+), 36 deletions(-) diff --git a/optimum/exporters/openvino/input_generators.py b/optimum/exporters/openvino/input_generators.py index 238e83a4ff..da2fc7a7f7 100644 --- a/optimum/exporters/openvino/input_generators.py +++ b/optimum/exporters/openvino/input_generators.py @@ -344,7 +344,6 @@ def __init__( **kwargs, ): super().__init__( - task=task, normalized_config=normalized_config, batch_size=batch_size, diff --git a/optimum/exporters/openvino/model_patcher.py b/optimum/exporters/openvino/model_patcher.py index 62428ae553..f4a346dfed 100644 --- a/optimum/exporters/openvino/model_patcher.py +++ b/optimum/exporters/openvino/model_patcher.py @@ -52,6 +52,7 @@ from optimum.exporters.openvino.patching_utils import ( UNSUPPORTED_OPS_PATCHING_SPEC, ModelPatcher, + override_arguments, postprocess_past_key_values, preprocess_past_key_values, sdpa_mask_without_vmap, @@ -85,8 +86,7 @@ else: TransformersKwargs = object -if is_transformers_version(">=", "4.55"): - from transformers.models.gpt_oss.modeling_gpt_oss import GptOssExperts + if is_transformers_version(">=", "4.56"): import transformers.masking_utils @@ -102,24 +102,6 @@ logger = logging.getLogger(__name__) -class BigBirdPegasusModelPatcher(ModelPatcher): - def __enter__(self): - super().__enter__() - - if self.real_config._behavior == "encoder" and self._model.config.attention_type == "block_sparse": - logger.warning( - "BigBirdPegasus model is using block sparse attention, which is not supported in ONNX export. " - "The model will be exported with original full attention." - ) - self._model.set_attention_type("original_full") - - def __exit__(self, exc_type, exc_value, traceback): - super().__exit__(exc_type, exc_value, traceback) - - if self.real_config._behavior == "encoder" and self._model.config.attention_type == "block_sparse": - self._model.set_attention_type("block_sparse") - - class SAMModelPatcher(ModelPatcher): def __init__( self, @@ -281,21 +263,6 @@ def _get_feat_extract_output_lengths_patched(self, input_lengths: torch.LongTens return output_conv3_length -class GptOssModelPatcher(ModelPatcher): - def __enter__(self): - super().__enter__() - - if is_transformers_version(">=", "4.55.0"): - self.original_gpt_oss_forward = GptOssExperts.forward - GptOssExperts.forward = gpt_oss_forward - - def __exit__(self, exc_type, exc_value, traceback): - super().__exit__(exc_type, exc_value, traceback) - - if is_transformers_version(">=", "4.55.0"): - GptOssExperts.forward = self.original_gpt_oss_forward - - def _get_model_attribute(model, name): target = getattr(model, "model", model) if is_transformers_version(">=", "5") else model return getattr(target, name) @@ -8148,6 +8115,26 @@ def __exit__(self, exc_type, exc_value, traceback): conv_layer.slow_forward = conv_layer._orig_forward +# Copied from https://github.com/huggingface/transformers/blob/v4.56.0/src/transformers/models/gpt_oss/modeling_gpt_oss.py#L81 +def gpt_oss_forward(self, hidden_states: torch.Tensor, router_indices=None, routing_weights=None) -> torch.Tensor: + batch_size = hidden_states.shape[0] + hidden_states = hidden_states.reshape(-1, self.hidden_size) + num_experts = routing_weights.shape[1] + hidden_states = hidden_states.repeat(num_experts, 1) + hidden_states = hidden_states.view(num_experts, -1, self.hidden_size) + gate_up = torch.bmm(hidden_states, self.gate_up_proj) + self.gate_up_proj_bias[..., None, :] + gate, up = gate_up[..., ::2], gate_up[..., 1::2] + gate = gate.clamp(min=None, max=self.limit) + up = up.clamp(min=-self.limit, max=self.limit) + glu = gate * torch.sigmoid(gate * self.alpha) + next_states = torch.bmm(((up + 1) * glu), self.down_proj) + next_states = next_states + self.down_proj_bias[..., None, :] + next_states = next_states.view(num_experts, batch_size, -1, self.hidden_size) + next_states = next_states * routing_weights.transpose(0, 1).view(num_experts, batch_size, -1)[..., None] + next_states = next_states.sum(dim=0) + return next_states + + class GptOssModelPatcher(OVDecoderModelPatcher): def __enter__(self): super().__enter__() From 16926be0d926870d8df4927e321ada3a0df0daeb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CElla?= Date: Wed, 27 May 2026 18:55:36 +0200 Subject: [PATCH 17/40] add missing patched_gemma4_clippable_linear_forward --- optimum/exporters/openvino/model_configs.py | 1 - optimum/exporters/openvino/model_patcher.py | 8 ++++++++ 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/optimum/exporters/openvino/model_configs.py b/optimum/exporters/openvino/model_configs.py index 78f71c0eae..a5bf6d79a8 100644 --- a/optimum/exporters/openvino/model_configs.py +++ b/optimum/exporters/openvino/model_configs.py @@ -23,7 +23,6 @@ AudioOnnxConfig, AudioToTextOnnxConfig, EncoderDecoderBaseOnnxConfig, - OnnxConfig, TextAndVisionOnnxConfig, TextDecoderOnnxConfig, TextDecoderWithPositionIdsOnnxConfig, diff --git a/optimum/exporters/openvino/model_patcher.py b/optimum/exporters/openvino/model_patcher.py index f4a346dfed..72a49f9af2 100644 --- a/optimum/exporters/openvino/model_patcher.py +++ b/optimum/exporters/openvino/model_patcher.py @@ -9211,6 +9211,14 @@ def __exit__(self, exc_type, exc_value, traceback): del sparse_moe_block.down_projs, sparse_moe_block.gate_projs, sparse_moe_block.up_projs +# OpenVINO has a bug due to which Clamp(-inf, inf) doesn't work correctly: CVS-185473. +# When min == -inf and max == inf, Clamp is equivalent to an identity operation and +# can be removed from the model, which serves as a workaround for the issue. +def patched_gemma4_clippable_linear_forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.linear(hidden_states) + return hidden_states + + class Gemma4ImageEmbeddingsModelPatcher(CommonImageEmbeddingsModelPatcher): def __init__(self, config, model, model_kwargs): super().__init__(config, model, model_kwargs) From 63b3269292f6a193cf3f8b16bab0e99ff24dc14f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CElla?= Date: Wed, 27 May 2026 18:57:29 +0200 Subject: [PATCH 18/40] fix --- optimum/exporters/openvino/patching_utils.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/optimum/exporters/openvino/patching_utils.py b/optimum/exporters/openvino/patching_utils.py index cd4bb2d41c..03873a4373 100644 --- a/optimum/exporters/openvino/patching_utils.py +++ b/optimum/exporters/openvino/patching_utils.py @@ -27,6 +27,7 @@ from transformers.cache_utils import DynamicCache, EncoderDecoderCache from transformers.modeling_outputs import BaseModelOutput +from optimum.exporters.base import ExporterConfig from optimum.intel.utils.import_utils import is_torch_version, is_transformers_version @@ -578,7 +579,7 @@ def patched_dynamic_layer_update( class ModelPatcher: def __init__( self, - config: "OnnxConfig", + config: ExporterConfig, model: PreTrainedModel, model_kwargs: dict[str, Any] | None = None, ): From 7366916cde411c7e67ff9b842f0588006dae75ee Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CElla?= Date: Thu, 28 May 2026 10:25:39 +0200 Subject: [PATCH 19/40] postprocess_past_key_values --- optimum/exporters/openvino/patching_utils.py | 18 +++--------------- 1 file changed, 3 insertions(+), 15 deletions(-) diff --git a/optimum/exporters/openvino/patching_utils.py b/optimum/exporters/openvino/patching_utils.py index 03873a4373..a7e43fc6b7 100644 --- a/optimum/exporters/openvino/patching_utils.py +++ b/optimum/exporters/openvino/patching_utils.py @@ -226,8 +226,8 @@ def preprocess_past_key_values(past_key_values): return past_key_values -def postprocess_past_key_values(past_key_values, output_names: list[str]): - if is_transformers_version(">=", "4.48") and isinstance(past_key_values, (EncoderDecoderCache, DynamicCache)): +def postprocess_past_key_values(past_key_values): + if isinstance(past_key_values, (EncoderDecoderCache, DynamicCache)): if hasattr(past_key_values, "to_legacy_cache"): past_key_values = past_key_values.to_legacy_cache() elif isinstance(past_key_values, DynamicCache): @@ -240,16 +240,6 @@ def postprocess_past_key_values(past_key_values, output_names: list[str]): past_key_values.cross_attention_cache.layers, ) ] - else: - raise NotImplementedError(f"Unable to serialize class {type(past_key_values)}.") - - if ( - isinstance(past_key_values, (list, tuple)) - and isinstance(past_key_values[0], (list, tuple)) - and not any("encoder.key" in output_name for output_name in output_names) - ): - past_key_values = tuple(pkv[:2] for pkv in past_key_values) - return past_key_values @@ -681,9 +671,7 @@ def patched_forward(*args, **kwargs): filtered_outputs[output_name] = outputs if filtered_outputs.get("past_key_values") is not None: - filtered_outputs["past_key_values"] = postprocess_past_key_values( - filtered_outputs["past_key_values"], output_names=output_names - ) + filtered_outputs["past_key_values"] = postprocess_past_key_values(filtered_outputs["past_key_values"]) return filtered_outputs From 7b4510d69ba8b2cc155123c1af3cd650c64ab317 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CElla?= Date: Thu, 28 May 2026 15:30:39 +0200 Subject: [PATCH 20/40] add donut swin config --- optimum/exporters/openvino/config.py | 10 +++++----- optimum/exporters/openvino/model_configs.py | 5 +++++ optimum/exporters/openvino/utils.py | 1 - 3 files changed, 10 insertions(+), 6 deletions(-) diff --git a/optimum/exporters/openvino/config.py b/optimum/exporters/openvino/config.py index f53230dd57..8ea3873b5b 100644 --- a/optimum/exporters/openvino/config.py +++ b/optimum/exporters/openvino/config.py @@ -1,4 +1,4 @@ -# Copyright 2022 The HuggingFace Team. All rights reserved. +# Copyright 2026 The HuggingFace Team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -232,9 +232,9 @@ def __init__( self.is_decoder_with_past = False - # Set up the encoder ONNX config. + # Set up the encoder openvino config. encoder_onnx_config_constructor = TasksManager.get_exporter_config_constructor( - exporter="onnx", + exporter="openvino", task="feature-extraction", model_type=config.encoder.model_type, library_name="transformers", @@ -244,9 +244,9 @@ def __init__( ) self._normalized_config.ENCODER_NORMALIZED_CONFIG_CLASS = self._encoder_onnx_config._normalized_config - # Set up the decoder ONNX config. + # Set up the decoder openvino config. decoder_onnx_config_constructor = TasksManager.get_exporter_config_constructor( - exporter="onnx", + exporter="openvino", task="feature-extraction", model_type=config.decoder.model_type, library_name="transformers", diff --git a/optimum/exporters/openvino/model_configs.py b/optimum/exporters/openvino/model_configs.py index a5bf6d79a8..595516d651 100644 --- a/optimum/exporters/openvino/model_configs.py +++ b/optimum/exporters/openvino/model_configs.py @@ -5128,6 +5128,11 @@ class SwinOpenVINOConfig(ViTOpenVINOConfig): pass +@register_in_tasks_manager("donut-swin", *["feature-extraction"]) +class DonutSwinOpenVINOConfig(ViTOpenVINOConfig): + pass + + @register_in_tasks_manager("convnext", *["feature-extraction", "image-classification"]) class ConvNextOpenVINOConfig(ViTOpenVINOConfig): pass diff --git a/optimum/exporters/openvino/utils.py b/optimum/exporters/openvino/utils.py index 0186437047..2832ebbb9f 100644 --- a/optimum/exporters/openvino/utils.py +++ b/optimum/exporters/openvino/utils.py @@ -363,7 +363,6 @@ def _get_kokoro_submodels(model): "default-timm-config", "detr", "dinov2", - "donut-swin", "dpt", "efficientnet", "encoder-decoder", From 116b0d0b33882a1c335f3d5e418f112f7fb48dda Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CElla?= Date: Thu, 28 May 2026 16:14:48 +0200 Subject: [PATCH 21/40] fix speecht5 --- optimum/exporters/openvino/model_configs.py | 25 +++++++++++++++------ 1 file changed, 18 insertions(+), 7 deletions(-) diff --git a/optimum/exporters/openvino/model_configs.py b/optimum/exporters/openvino/model_configs.py index 595516d651..5745e389fd 100644 --- a/optimum/exporters/openvino/model_configs.py +++ b/optimum/exporters/openvino/model_configs.py @@ -4166,6 +4166,13 @@ class SpeechT5ConfigBehavior(str, enum.Enum): library_name="transformers", ) class SpeechT5OpenVINOConfig(OnnxSeq2SeqConfigWithPast): + DUMMY_INPUT_GENERATOR_CLASSES = ( + DummyTextInputGenerator, + DummySeq2SeqPastKeyValuesGenerator, + DummySpeechT5InputGenerator, + ) + _MODEL_PATCHER = SpeechT5ModelPatcher + NORMALIZED_CONFIG_CLASS = NormalizedSeq2SeqConfig.with_args( hidden_size="hidden_size", num_attention_heads="encoder_attention_heads", @@ -4173,18 +4180,12 @@ class SpeechT5OpenVINOConfig(OnnxSeq2SeqConfigWithPast): decoder_num_layers="decoder_layers", allow_new=True, ) - DUMMY_INPUT_GENERATOR_CLASSES = ( - DummyTextInputGenerator, - DummySeq2SeqPastKeyValuesGenerator, - DummySpeechT5InputGenerator, - ) DUMMY_PKV_GENERATOR_CLASS = DummySeq2SeqPastKeyValuesGenerator VARIANTS = { # noqa: RUF012 "with-past": "The export follows the Transformers implementation using the KV cache.", "without-past": "The same as `with-past`, just without KV cache support.", } DEFAULT_VARIANT = "with-past" - _MODEL_PATCHER = SpeechT5ModelPatcher def __init__( self, @@ -4206,9 +4207,10 @@ def __init__( use_past_in_inputs=use_past_in_inputs, behavior=behavior, preprocessors=preprocessors, - is_postnet_and_vocoder=False, ) + self.is_postnet_and_vocoder = False + def add_past_key_values(self, inputs_or_outputs: Dict[str, Dict[int, str]], direction: str): if direction not in ["inputs", "outputs"]: raise ValueError(f'direction must either be "inputs" or "outputs", but {direction} was given') @@ -4312,6 +4314,15 @@ def with_behavior( "self._behavior is neither encoder, decoder, postnet, or vocoder. This should not happen." ) + def overwrite_shape_and_generate_input( + self, dummy_input_gen: DummyInputGenerator, input_name: str, framework: str, input_shapes: dict + ): + dummy_input_gen.batch_size = 1 + dummy_input = dummy_input_gen.generate( + input_name, framework=framework, int_dtype=self.int_dtype, float_dtype=self.float_dtype + ) + return dummy_input + @register_in_tasks_manager( "llama4_text", *["text-generation", "text-generation-with-past"], library_name="transformers" From 4af9cff9e139198c253966970ee66f22bacd7590 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CElla?= Date: Thu, 28 May 2026 16:20:57 +0200 Subject: [PATCH 22/40] add TrOCROpenVINOConfig --- optimum/exporters/openvino/model_configs.py | 16 ++++++++++++++++ optimum/exporters/openvino/utils.py | 1 - 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/optimum/exporters/openvino/model_configs.py b/optimum/exporters/openvino/model_configs.py index 5745e389fd..19c5ac9866 100644 --- a/optimum/exporters/openvino/model_configs.py +++ b/optimum/exporters/openvino/model_configs.py @@ -5905,3 +5905,19 @@ def outputs(self) -> Dict[str, Dict[int, str]]: "waveform": {0: "batch_size", 1: "audio_length"}, "phonemes": {0: "batch_size", 1: "phoneme_length"}, } + + +@register_in_tasks_manager( + "trocr", + *[ + "feature-extraction", + "feature-extraction-with-past", + ], +) +class TrOCROpenVINOConfig(TextSeq2SeqOnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedSeq2SeqConfig.with_args( + decoder_num_layers="decoder_layers", + num_layers="decoder_layers", + decoder_num_attention_heads="decoder_attention_heads", + hidden_size="hidden_size", + ) diff --git a/optimum/exporters/openvino/utils.py b/optimum/exporters/openvino/utils.py index 2832ebbb9f..164c11f2f1 100644 --- a/optimum/exporters/openvino/utils.py +++ b/optimum/exporters/openvino/utils.py @@ -400,7 +400,6 @@ def _get_kokoro_submodels(model): "swin2sr", "swinv2", "table-transformer", - "trocr", "visual_bert", "vit_mae", "vit_msn", From fd705ddd34138f9f192a43bfb340eb62aa58d90e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CElla?= Date: Thu, 28 May 2026 16:51:19 +0200 Subject: [PATCH 23/40] remove onnx --- .../openvino/_traceable_decorator.py | 2 +- optimum/exporters/openvino/base.py | 96 +------------------ setup.py | 1 - 3 files changed, 2 insertions(+), 97 deletions(-) diff --git a/optimum/exporters/openvino/_traceable_decorator.py b/optimum/exporters/openvino/_traceable_decorator.py index a60025c1eb..e12f6f2fad 100644 --- a/optimum/exporters/openvino/_traceable_decorator.py +++ b/optimum/exporters/openvino/_traceable_decorator.py @@ -26,7 +26,7 @@ # This is a fixed version of transformers.utils.generic.check_model_inputs -# that fixes issues related to onnx export and tracing +# that fixes issues related to export and tracing # - adds support for positional args (use_cache), without which use_cache end up being passed twice # - fixes issue with default capture_flags being None for some models def traceable_check_model_inputs(func): diff --git a/optimum/exporters/openvino/base.py b/optimum/exporters/openvino/base.py index 5eee6301fe..196cfc463e 100644 --- a/optimum/exporters/openvino/base.py +++ b/optimum/exporters/openvino/base.py @@ -33,11 +33,6 @@ from optimum.exporters.openvino.patching_utils import ModelPatcher, PatchingSpec from optimum.utils import DEFAULT_DUMMY_SHAPES, DummyInputGenerator, DummySeq2SeqPastKeyValuesGenerator, logging from optimum.utils.doc import add_dynamic_docstring -from optimum.utils.import_utils import ( - is_onnx_available, - is_onnxruntime_available, - is_transformers_version, -) logger = logging.get_logger(__name__) @@ -100,9 +95,7 @@ class OnnxConfig(ExporterConfig, ABC): {"heatmaps": {0: "batch_size", 1: "num_keypoints", 2: "height", 3: "width"}} ), "mask-generation": OrderedDict({"logits": {0: "batch_size"}}), - "masked-im": OrderedDict( - {"reconstruction" if is_transformers_version(">=", "4.29.0") else "logits": {0: "batch_size"}} - ), + "masked-im": OrderedDict({"reconstruction": {0: "batch_size"}}), "multiple-choice": OrderedDict({"logits": {0: "batch_size", 1: "num_choices"}}), "object-detection": OrderedDict( { @@ -171,93 +164,6 @@ def variant(self, value: str): raise ValueError(f"The variant {value} is not supported for the ONNX config {self.__class__.__name__}.") self._variant = value - def fix_dynamic_axes( - self, model_path: Path, device: str = "cpu", dtype: str | None = None, input_shapes: dict | None = None - ): - """Fixes potential issues with dynamic axes. - - During the export, ONNX will infer some axes to be dynamic which are actually static. This method is called - right after the export to fix such issues. - - Args: - model_path (`Path`): - The path of the freshly exported ONNX model. - device (`str`, defaults to `"cpu"`): - The device on which the model will be run. This is used to determine the ONNX Runtime provider. - dtype (`Optional[str]`, defaults to `None`): - The data type of the model inputs. If `None`, it will be inferred from the model inputs. - input_shapes (`Optional[Dict[str, Any]]`, defaults to `None`): - The shapes of the model inputs. If `None`, it will be inferred from the model inputs. - """ - if not (is_onnx_available() and is_onnxruntime_available()): - raise RuntimeError( - "The onnx and onnxruntime packages are necessary to fix the dynamic shapes of the exported model. " - "You can install them by doing: pip install onnx onnxruntime" - ) - - import onnx - from onnxruntime import GraphOptimizationLevel, InferenceSession, SessionOptions - - allowed_dynamic_axes = set() - for input_ in self.inputs.values(): - allowed_dynamic_axes |= set(input_.values()) - for output in self.outputs.values(): - allowed_dynamic_axes |= set(output.values()) - - if device.startswith("cuda"): - providers = ["CUDAExecutionProvider"] - else: - providers = ["CPUExecutionProvider"] - - session_options = SessionOptions() - session_options.graph_optimization_level = GraphOptimizationLevel.ORT_DISABLE_ALL # no need to optimize here - session = InferenceSession(model_path.as_posix(), providers=providers, sess_options=session_options) - - to_fix = [] - for output_idx, node in enumerate(session.get_outputs()): - for idx, axis in enumerate(node.shape): - if isinstance(axis, str) and axis not in allowed_dynamic_axes: - to_fix.append((output_idx, idx)) - - # We branch here to avoid doing an unnecessary forward pass. - if to_fix: - if input_shapes is None: - input_shapes = {} - - onnx_input_names = [inp.name for inp in session.get_inputs()] - dummy_inputs = self.generate_dummy_inputs(framework="np", **input_shapes) - dummy_inputs = self.generate_dummy_inputs_for_validation(dummy_inputs, onnx_input_names) - dummy_inputs = self.rename_ambiguous_inputs(dummy_inputs) - - onnx_inputs = {} - for name, value in dummy_inputs.items(): - if isinstance(value, (list, tuple)): - value = self.flatten_output_collection_property(name, value) - onnx_inputs.update(dict(value.items())) - else: - onnx_inputs[name] = value - - for name, value in onnx_inputs.items(): - if value.dtype == np.float32 and dtype == "fp16": - onnx_inputs[name] = onnx_inputs[name].astype(np.float16) - - outputs = session.run(None, onnx_inputs) - del session - - onnx_model = onnx.load(model_path.as_posix(), load_external_data=False) - - for output_idx, dim_idx in to_fix: - dims = onnx_model.graph.output[output_idx].type.tensor_type.shape.dim - dims[dim_idx].dim_value = outputs[output_idx].shape[dim_idx] - - onnx.save( - onnx_model, - model_path.as_posix(), - convert_attribute=True, - ) - del onnx_model - gc.collect() - @property def torch_to_onnx_input_map(self) -> dict[str, str]: """Dictionary mapping input names from the PyTorch model to input names from the exported ONNX model. diff --git a/setup.py b/setup.py index ebee6d9cdd..f144fc0a10 100644 --- a/setup.py +++ b/setup.py @@ -29,7 +29,6 @@ INSTALL_REQUIRE = [ "torch>=2.1", "optimum@git+https://github.com/huggingface/optimum.git", - "onnx", "transformers>=4.45,<5.1", "setuptools", "huggingface-hub>=0.23.2,<2.0", From 008cf69530b2d1980a377e0cae45ba48cbe2f4f4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CElla?= Date: Thu, 28 May 2026 16:56:03 +0200 Subject: [PATCH 24/40] style --- optimum/exporters/openvino/base.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/optimum/exporters/openvino/base.py b/optimum/exporters/openvino/base.py index 196cfc463e..43e29b14c5 100644 --- a/optimum/exporters/openvino/base.py +++ b/optimum/exporters/openvino/base.py @@ -16,17 +16,14 @@ from __future__ import annotations import enum -import gc import inspect import itertools import re from abc import ABC from collections import OrderedDict from collections.abc import Iterable -from pathlib import Path from typing import Any, ClassVar -import numpy as np from transformers import PretrainedConfig, PreTrainedModel from optimum.exporters.base import ExporterConfig From b46c5fc305f4ed42cb103984075ca1bf29ded3e6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CElla?= Date: Thu, 28 May 2026 17:16:28 +0200 Subject: [PATCH 25/40] protobuf --- setup.py | 1 + 1 file changed, 1 insertion(+) diff --git a/setup.py b/setup.py index f144fc0a10..b6f2110b0b 100644 --- a/setup.py +++ b/setup.py @@ -69,6 +69,7 @@ "decord", "imageio", "kokoro", + "protobuf", ] QUALITY_REQUIRE = ["black~=23.1", "ruff==0.4.4"] From fe6792a44d5959e018c5b1f41173f7fd366af8d6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CElla?= Date: Thu, 28 May 2026 18:11:28 +0200 Subject: [PATCH 26/40] remove opset --- optimum/exporters/openvino/__main__.py | 6 +- optimum/exporters/openvino/base.py | 95 ++++--- optimum/exporters/openvino/config.py | 66 ++--- optimum/exporters/openvino/convert.py | 42 ++- optimum/exporters/openvino/model_configs.py | 254 +++++++++---------- optimum/exporters/openvino/model_patcher.py | 108 ++++---- optimum/exporters/openvino/patching_utils.py | 6 +- optimum/exporters/openvino/utils.py | 10 +- 8 files changed, 273 insertions(+), 314 deletions(-) diff --git a/optimum/exporters/openvino/__main__.py b/optimum/exporters/openvino/__main__.py index a962205462..f3d8fed9de 100644 --- a/optimum/exporters/openvino/__main__.py +++ b/optimum/exporters/openvino/__main__.py @@ -28,7 +28,7 @@ from transformers.utils import is_torch_available from openvino import Core, Type, save_model -from optimum.exporters.openvino.base import OnnxConfig +from optimum.exporters.openvino.base import OpenVINOConfig from optimum.exporters.tasks import TasksManager from optimum.intel.utils.import_utils import ( DIFFUSERS_IMPORT_ERROR, @@ -227,7 +227,7 @@ def main_export( local_files_only: bool = False, token: Optional[Union[bool, str]] = None, model_kwargs: Optional[Dict[str, Any]] = None, - custom_export_configs: Optional[Dict[str, "OnnxConfig"]] = None, + custom_export_configs: Optional[Dict[str, "OpenVINOConfig"]] = None, fn_get_submodels: Optional[Callable] = None, ov_config: "OVConfig" = None, stateful: bool = True, @@ -283,7 +283,7 @@ def main_export( the export. This argument should be used along the `custom_export_configs` argument in case, for example, the model inputs/outputs are changed (for example, if `model_kwargs={"output_attentions": True}` is passed). - custom_export_configs (`Optional[Dict[str, OnnxConfig]]`, defaults to `None`): + custom_export_configs (`Optional[Dict[str, OpenVINOConfig]]`, defaults to `None`): Experimental usage: override the default export config used for the given model. This argument may be useful for advanced users that desire a finer-grained control on the export. An example is available [here](https://huggingface.co/docs/optimum/main/en/exporters/onnx/usage_guides/export_a_model). fn_get_submodels (`Optional[Callable]`, defaults to `None`): Experimental usage: Override the default submodels that are used at the export. This is diff --git a/optimum/exporters/openvino/base.py b/optimum/exporters/openvino/base.py index 43e29b14c5..fb4eff5453 100644 --- a/optimum/exporters/openvino/base.py +++ b/optimum/exporters/openvino/base.py @@ -11,7 +11,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""ONNX configuration base classes.""" +"""OpenVINO configuration base classes.""" from __future__ import annotations @@ -66,9 +66,8 @@ """ -class OnnxConfig(ExporterConfig, ABC): - DEFAULT_ONNX_OPSET = 18 - VARIANTS: ClassVar[dict[str, str]] = {"default": "The default ONNX variant."} +class OpenVINOConfig(ExporterConfig, ABC): + VARIANTS: ClassVar[dict[str, str]] = {"default": "The default OpenVINO variant."} DEFAULT_VARIANT = "default" PATCHING_SPECS: list[PatchingSpec] | None = None _MODEL_PATCHER = ModelPatcher @@ -146,10 +145,10 @@ def __init__( @property def variant(self) -> str: - """For a given ONNX config, the variant of the model to export. + """For a given OpenVINO config, the variant of the model to export. This property allows to define variants of a given model, in case - different users would like to export the model differently (with different inputs/outputs, model split in several ONNX or not, etc.). + different users would like to export the model differently (with different inputs/outputs, model split in several OpenVINO or not, etc.). """ return self._variant @@ -158,28 +157,28 @@ def variant(self, value: str): if value == "default" and hasattr(self, "DEFAULT_VARIANT"): value = self.DEFAULT_VARIANT if value not in self.VARIANTS: - raise ValueError(f"The variant {value} is not supported for the ONNX config {self.__class__.__name__}.") + raise ValueError(f"The variant {value} is not supported for the OpenVINO config {self.__class__.__name__}.") self._variant = value @property - def torch_to_onnx_input_map(self) -> dict[str, str]: - """Dictionary mapping input names from the PyTorch model to input names from the exported ONNX model. + def torch_to_ov_input_map(self) -> dict[str, str]: + """Dictionary mapping input names from the PyTorch model to input names from the exported OpenVINO model. - Override the function when the input names and the exported ONNX input names are different. + Override the function when the input names and the exported OpenVINO input names are different. Returns: - `Dict[str, str]`: A dictionary mapping the PyTorch model input names to the exported ONNX model input names. + `Dict[str, str]`: A dictionary mapping the PyTorch model input names to the exported OpenVINO model input names. """ return {} @property - def torch_to_onnx_output_map(self) -> dict[str, str]: - """Dictionary mapping output names from the PyTorch model to output names from the exported ONNX model. + def torch_to_ov_output_map(self) -> dict[str, str]: + """Dictionary mapping output names from the PyTorch model to output names from the exported OpenVINO model. - Override the function when the output names and the exported ONNX output names are different. + Override the function when the output names and the exported OpenVINO output names are different. Returns: - `Dict[str, str]`: A dictionary mapping the PyTorch model output names to the exported ONNX model output names. + `Dict[str, str]`: A dictionary mapping the PyTorch model output names to the exported OpenVINO model output names. """ return {} @@ -198,7 +197,7 @@ def ordered_inputs(self, model: PreTrainedModel) -> dict[str, dict[int, str]]: Args: model ([`transformers.PreTrainedModel`]): - The model for which we will use the OnnxConfig. + The model for which we will use the OpenVINOConfig. Returns: `Dict[str, Dict[int, str]]`: The properly ordered inputs. @@ -221,7 +220,7 @@ def ordered_inputs(self, model: PreTrainedModel) -> dict[str, dict[int, str]]: # TODO: figure out a smart way of re-ordering potential nested structures. # to_insert = sorted(to_insert, key=lambda t: t[0]) for name, dynamic_axes in to_insert: - name = self.torch_to_onnx_input_map.get(name, name) + name = self.torch_to_ov_input_map.get(name, name) ordered_inputs[name] = dynamic_axes return ordered_inputs @@ -245,18 +244,18 @@ def flatten_output_collection_property(cls, name: str, field: Iterable[Any]) -> return {f"{name}.{idx}": item for idx, item in enumerate(field)} def generate_dummy_inputs_for_validation( - self, reference_model_inputs: dict[str, Any], onnx_input_names: list[str] + self, reference_model_inputs: dict[str, Any], ov_input_names: list[str] ) -> dict[str, Any]: - """Generates inputs for ONNX Runtime using the reference model inputs. + """Generates inputs for OpenVINO Runtime using the reference model inputs. Override this to run inference with seq2seq - models which have the encoder and decoder exported as separate ONNX files. + models which have the encoder and decoder exported as separate OpenVINO files. Args: reference_model_inputs (`Dict[str, Tensor]`): Reference inputs for the model. - onnx_input_names (`Optional[List[str]]`, defaults to `None`): - Names of the actual inputs to the ONNX model. This argument may be required as an unused + ov_input_names (`Optional[List[str]]`, defaults to `None`): + Names of the actual inputs to the OpenVINO model. This argument may be required as an unused input to the model is automatically removed by torch.onnx.export (e.g. encoder_outputs in the decoder with past) Returns: @@ -270,8 +269,8 @@ def patch_model_for_export( return self._MODEL_PATCHER(self, model, model_kwargs=model_kwargs) -class OnnxConfigWithPast(OnnxConfig, ABC): - """Inherits from [`~exporters.onnx.OnnxConfig`]. A base class to handle the ONNX configuration of decoder-only models.""" +class OpenVINOConfigWithPast(OpenVINOConfig, ABC): + """Inherits from [`~exporters.onnx.OpenVINOConfig`]. A base class to handle the OpenVINO configuration of decoder-only models.""" PAD_ATTENTION_MASK_TO_PAST: bool = False SUPPORTS_PAST: bool = True @@ -341,7 +340,7 @@ def generate_dummy_inputs(self, framework: str = "pt", **kwargs): break if not input_was_inserted: raise RuntimeError( - f'Could not generate dummy input for "{input_name}". Try adding a proper dummy input generator to the model ONNX config.' + f'Could not generate dummy input for "{input_name}". Try adding a proper dummy input generator to the model OpenVINO config.' ) # refer to https://github.com/huggingface/optimum/pull/764 @@ -368,11 +367,11 @@ def overwrite_shape_and_generate_input( This method allows to overwrite some shapes, and generate the dummy input. This should probably be refactored more elegantly. """ - # models from TextSeq2SeqOnnxConfig use decoder_input_ids as input name - # while models from TextDecoderOnnxConfig use input_ids, hence the check for both + # models from TextSeq2SeqOpenVINOConfig use decoder_input_ids as input name + # while models from TextDecoderOpenVINOConfig use input_ids, hence the check for both - # NOTE: The check `self.task != "text-generation" is added following the use of a single ONNX for both without/with KV cache, without subgraphs. - # This overwrite may be moved to OnnxSeq2SeqConfigWithPast, but I am afraid it would break encoder-decoder models. + # NOTE: The check `self.task != "text-generation" is added following the use of a single OpenVINO for both without/with KV cache, without subgraphs. + # This overwrite may be moved to OpenVINOSeq2SeqConfigWithPast, but I am afraid it would break encoder-decoder models. if ( self.use_past and self.use_past_in_inputs @@ -433,7 +432,7 @@ def flatten_output_collection_property(self, name: str, field: Iterable[Any]) -> return flattened_output def generate_dummy_inputs_for_validation( - self, reference_model_inputs: dict[str, Any], onnx_input_names: list[str] + self, reference_model_inputs: dict[str, Any], ov_input_names: list[str] ) -> dict[str, Any]: if self.is_merged is True and self.use_cache_branch is True: reference_model_inputs["use_cache_branch"] = DummyInputGenerator.constant_tensor(shape=[1], value=True) @@ -450,11 +449,11 @@ def generate_dummy_inputs_for_validation( "past_key_values", framework="pt", int_dtype=self.int_dtype, float_dtype=self.float_dtype ) - return super().generate_dummy_inputs_for_validation(reference_model_inputs, onnx_input_names) + return super().generate_dummy_inputs_for_validation(reference_model_inputs, ov_input_names) class ConfigBehavior(str, enum.Enum): - """Specifies the behavior of the [`~exporters.onnx.base.OnnxSeq2SeqConfigWithPast`]. + """Specifies the behavior of the [`~exporters.onnx.base.OpenVINOSeq2SeqConfigWithPast`]. - MONOLITH: the config can be used to export the whole seq2seq model as a single file. - ENCODER: the config can be used to export the encoder part of the seq2seq model. @@ -466,8 +465,8 @@ class ConfigBehavior(str, enum.Enum): DECODER = "decoder" -class OnnxSeq2SeqConfigWithPast(OnnxConfigWithPast): - """Inherits from [`~exporters.onnx.OnnxConfigWithPast`]. A base class to handle the ONNX configuration of encoder-decoder models.""" +class OpenVINOSeq2SeqConfigWithPast(OpenVINOConfigWithPast): + """Inherits from [`~exporters.onnx.OpenVINOConfigWithPast`]. A base class to handle the OpenVINO configuration of encoder-decoder models.""" DUMMY_PKV_GENERATOR_CLASS = DummySeq2SeqPastKeyValuesGenerator @@ -502,24 +501,24 @@ def with_behavior( behavior: str | ConfigBehavior, use_past: bool = False, use_past_in_inputs: bool = False, - ) -> OnnxSeq2SeqConfigWithPast: - """Creates a copy of the current OnnxConfig but with a different `ConfigBehavior` and `use_past` value. + ) -> OpenVINOSeq2SeqConfigWithPast: + """Creates a copy of the current OpenVINOConfig but with a different `ConfigBehavior` and `use_past` value. Args: behavior ([`ConfigBehavior`]): The behavior to use for the new instance. use_past (`bool`, defaults to `False`): - Whether or not the ONNX config to instantiate is for a model using KV cache. + Whether or not the OpenVINO config to instantiate is for a model using KV cache. use_past_in_inputs (`bool`, defaults to `False`): - Whether the KV cache is to be passed as an input to the ONNX. + Whether the KV cache is to be passed as an input to the OpenVINO. Returns: - `OnnxSeq2SeqConfigWithPast` + `OpenVINOSeq2SeqConfigWithPast` """ if isinstance(behavior, str) and not isinstance(behavior, ConfigBehavior): behavior = ConfigBehavior(behavior) - onnx_config = self.__class__( + ov_config = self.__class__( self._config, task=self.task, int_dtype=self.int_dtype, @@ -529,11 +528,11 @@ def with_behavior( behavior=behavior, preprocessors=self._preprocessors, ) - onnx_config.variant = self.variant - return onnx_config + ov_config.variant = self.variant + return ov_config @property - def torch_to_onnx_input_map(self) -> dict[str, str]: + def torch_to_ov_input_map(self) -> dict[str, str]: if self._behavior is ConfigBehavior.DECODER: return { "decoder_input_ids": "input_ids", @@ -545,7 +544,7 @@ def torch_to_onnx_input_map(self) -> dict[str, str]: @property def outputs(self) -> dict[str, dict[int, str]]: - common_outputs = super(OnnxConfigWithPast, self).outputs + common_outputs = super(OpenVINOConfigWithPast, self).outputs # Renaming the outputs axes properly. for name, axes_names in common_outputs.items(): if self._behavior is ConfigBehavior.ENCODER or "encoder" in name: @@ -559,7 +558,7 @@ def outputs(self) -> dict[str, dict[int, str]]: if self.use_past_in_inputs is False or self.is_merged is True: new_axes_names[axis_idx] = sequence_name else: - # Trick to force it since ONNX sometimes infer a dynamic axis where it's not. + # Trick to force it since OpenVINO sometimes infer a dynamic axis where it's not. new_axes_names[axis_idx] = "1" else: new_axes_names[axis_idx] = axis_name @@ -624,7 +623,7 @@ def generate_dummy_inputs(self, framework: str = "pt", **kwargs): return dummy_inputs def generate_dummy_inputs_for_validation( - self, reference_model_inputs: dict[str, Any], onnx_input_names: list[str] + self, reference_model_inputs: dict[str, Any], ov_input_names: list[str] ) -> dict[str, Any]: if self._behavior is ConfigBehavior.DECODER: if "decoder_input_ids" in reference_model_inputs: @@ -634,9 +633,9 @@ def generate_dummy_inputs_for_validation( if "decoder_attention_mask" in reference_model_inputs: reference_model_inputs["attention_mask"] = reference_model_inputs.pop("decoder_attention_mask") if "encoder_outputs" in reference_model_inputs: - if "encoder_hidden_states" in onnx_input_names: + if "encoder_hidden_states" in ov_input_names: reference_model_inputs["encoder_hidden_states"] = reference_model_inputs.pop("encoder_outputs")[0] else: reference_model_inputs.pop("encoder_outputs") - return super().generate_dummy_inputs_for_validation(reference_model_inputs, onnx_input_names) + return super().generate_dummy_inputs_for_validation(reference_model_inputs, ov_input_names) diff --git a/optimum/exporters/openvino/config.py b/optimum/exporters/openvino/config.py index 8ea3873b5b..1910659e27 100644 --- a/optimum/exporters/openvino/config.py +++ b/optimum/exporters/openvino/config.py @@ -11,7 +11,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""Common ONNX configuration classes that handle most of the features for building model specific configurations.""" +"""Common OpenVINO configuration classes that handle most of the features for building model specific configurations.""" from __future__ import annotations @@ -21,7 +21,7 @@ from transformers import PretrainedConfig -from optimum.exporters.openvino.base import ConfigBehavior, OnnxConfig, OnnxConfigWithPast, OnnxSeq2SeqConfigWithPast +from optimum.exporters.openvino.base import ConfigBehavior, OpenVINOConfig, OpenVINOConfigWithPast, OpenVINOSeq2SeqConfigWithPast from optimum.exporters.tasks import TasksManager from optimum.utils import ( DummyAudioInputGenerator, @@ -39,13 +39,13 @@ logger = logging.get_logger(__name__) -class TextEncoderOnnxConfig(OnnxConfig): +class TextEncoderOpenVINOConfig(OpenVINOConfig): """Handles encoder-based text architectures.""" DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator,) -class TextDecoderOnnxConfig(OnnxConfigWithPast): +class TextDecoderOpenVINOConfig(OpenVINOConfigWithPast): """Handles decoder-based text architectures.""" PAD_ATTENTION_MASK_TO_PAST = True @@ -97,7 +97,7 @@ def outputs(self) -> dict[str, dict[int, str]]: return common_outputs -class TextDecoderWithPositionIdsOnnxConfig(TextDecoderOnnxConfig): +class TextDecoderWithPositionIdsOpenVINOConfig(TextDecoderOpenVINOConfig): @property def inputs(self) -> dict[str, dict[int, str]]: common_inputs = super().inputs @@ -110,7 +110,7 @@ def inputs(self) -> dict[str, dict[int, str]]: return common_inputs -class TextSeq2SeqOnnxConfig(OnnxSeq2SeqConfigWithPast): +class TextSeq2SeqOpenVINOConfig(OpenVINOSeq2SeqConfigWithPast): """Handles encoder-decoder-based text architectures.""" DUMMY_INPUT_GENERATOR_CLASSES = ( @@ -159,19 +159,19 @@ def _create_dummy_input_generator_classes(self, **kwargs) -> list[DummyInputGene return dummy_inputs_generators -class VisionOnnxConfig(OnnxConfig): +class VisionOpenVINOConfig(OpenVINOConfig): """Handles vision architectures.""" DUMMY_INPUT_GENERATOR_CLASSES = (DummyVisionInputGenerator,) -class TextAndVisionOnnxConfig(OnnxConfig): +class TextAndVisionOpenVINOConfig(OpenVINOConfig): """Handles multi-modal text and vision architectures.""" DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, DummyVisionInputGenerator, DummyBboxInputGenerator) -class AudioOnnxConfig(OnnxConfig): +class AudioOpenVINOConfig(OpenVINOConfig): """Handles audio architectures.""" DUMMY_INPUT_GENERATOR_CLASSES = (DummyAudioInputGenerator,) @@ -181,7 +181,7 @@ def inputs(self) -> dict[str, dict[int, str]]: return {"input_values": {0: "batch_size", 1: "sequence_length"}} -class AudioToTextOnnxConfig(OnnxSeq2SeqConfigWithPast): +class AudioToTextOpenVINOConfig(OpenVINOSeq2SeqConfigWithPast): DUMMY_INPUT_GENERATOR_CLASSES = ( DummyAudioInputGenerator, DummySeq2SeqDecoderTextInputGenerator, @@ -205,7 +205,7 @@ def inputs(self) -> dict[str, dict[int, str]]: return common_inputs -class EncoderDecoderBaseOnnxConfig(OnnxSeq2SeqConfigWithPast): +class EncoderDecoderBaseOpenVINOConfig(OpenVINOSeq2SeqConfigWithPast): DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator,) def __init__( @@ -233,26 +233,26 @@ def __init__( self.is_decoder_with_past = False # Set up the encoder openvino config. - encoder_onnx_config_constructor = TasksManager.get_exporter_config_constructor( + encoder_ov_config_constructor = TasksManager.get_exporter_config_constructor( exporter="openvino", task="feature-extraction", model_type=config.encoder.model_type, library_name="transformers", ) - self._encoder_onnx_config = encoder_onnx_config_constructor( + self._encoder_ov_config = encoder_ov_config_constructor( config.encoder, int_dtype=int_dtype, float_dtype=float_dtype, preprocessors=preprocessors ) - self._normalized_config.ENCODER_NORMALIZED_CONFIG_CLASS = self._encoder_onnx_config._normalized_config + self._normalized_config.ENCODER_NORMALIZED_CONFIG_CLASS = self._encoder_ov_config._normalized_config # Set up the decoder openvino config. - decoder_onnx_config_constructor = TasksManager.get_exporter_config_constructor( + decoder_ov_config_constructor = TasksManager.get_exporter_config_constructor( exporter="openvino", task="feature-extraction", model_type=config.decoder.model_type, library_name="transformers", ) kwargs = {} - if issubclass(decoder_onnx_config_constructor.func, OnnxConfigWithPast): + if issubclass(decoder_ov_config_constructor.func, OpenVINOConfigWithPast): self.is_decoder_with_past = True kwargs["use_past"] = use_past else: @@ -264,24 +264,24 @@ def __init__( "past key values." ) - self._decoder_onnx_config = decoder_onnx_config_constructor( + self._decoder_ov_config = decoder_ov_config_constructor( config.decoder, int_dtype=int_dtype, float_dtype=float_dtype, preprocessors=preprocessors, **kwargs ) - if issubclass(decoder_onnx_config_constructor.func, OnnxSeq2SeqConfigWithPast): - self._decoder_onnx_config = self._decoder_onnx_config.with_behavior( + if issubclass(decoder_ov_config_constructor.func, OpenVINOSeq2SeqConfigWithPast): + self._decoder_ov_config = self._decoder_ov_config.with_behavior( self._behavior, use_past=kwargs["use_past"], use_past_in_inputs=use_past_in_inputs ) - self._normalized_config.DECODER_NORMALIZED_CONFIG_CLASS = self._decoder_onnx_config._normalized_config - self._normalized_config.DECODER_NORMALIZED_CONFIG_CLASS = self._decoder_onnx_config._normalized_config + self._normalized_config.DECODER_NORMALIZED_CONFIG_CLASS = self._decoder_ov_config._normalized_config + self._normalized_config.DECODER_NORMALIZED_CONFIG_CLASS = self._decoder_ov_config._normalized_config self._normalized_config.DECODER_NORMALIZED_CONFIG_CLASS.encoder_num_attention_heads = ( - self._decoder_onnx_config._normalized_config.num_attention_heads + self._decoder_ov_config._normalized_config.num_attention_heads ) self._normalized_config.DECODER_NORMALIZED_CONFIG_CLASS.decoder_num_attention_heads = ( - self._decoder_onnx_config._normalized_config.num_attention_heads + self._decoder_ov_config._normalized_config.num_attention_heads ) - if isinstance(self._decoder_onnx_config, OnnxSeq2SeqConfigWithPast): + if isinstance(self._decoder_ov_config, OpenVINOSeq2SeqConfigWithPast): self._past_key_values_generator = ( DummySeq2SeqDecoderTextInputGenerator, DummySeq2SeqPastKeyValuesGenerator, @@ -312,21 +312,21 @@ def inputs(self) -> dict[str, dict[int, str]]: def add_past_key_values(self, inputs_or_outputs: dict[str, dict[int, str]], direction: str): if self.is_decoder_with_past: - return self._decoder_onnx_config.add_past_key_values(inputs_or_outputs, direction) + return self._decoder_ov_config.add_past_key_values(inputs_or_outputs, direction) def flatten_past_key_values(self, flattened_output, name, idx, t): if self.is_decoder_with_past: - return self._decoder_onnx_config.flatten_past_key_values(flattened_output, name, idx, t) + return self._decoder_ov_config.flatten_past_key_values(flattened_output, name, idx, t) def flatten_output_collection_property(self, name: str, field: Iterable[Any]) -> dict[str, Any]: - return self._decoder_onnx_config.flatten_output_collection_property(name, field) + return self._decoder_ov_config.flatten_output_collection_property(name, field) def generate_dummy_inputs_for_validation( - self, reference_model_inputs: dict[str, Any], onnx_input_names: list[str] + self, reference_model_inputs: dict[str, Any], ov_input_names: list[str] ) -> dict[str, Any]: if self._behavior is ConfigBehavior.ENCODER: - return self._encoder_onnx_config.generate_dummy_inputs_for_validation( - reference_model_inputs, onnx_input_names + return self._encoder_ov_config.generate_dummy_inputs_for_validation( + reference_model_inputs, ov_input_names ) else: if self._behavior is ConfigBehavior.DECODER: @@ -335,13 +335,13 @@ def generate_dummy_inputs_for_validation( if "attention_mask" in reference_model_inputs: reference_model_inputs["encoder_attention_mask"] = reference_model_inputs.pop("attention_mask") if "encoder_outputs" in reference_model_inputs: - if "encoder_hidden_states" in onnx_input_names: + if "encoder_hidden_states" in ov_input_names: reference_model_inputs["encoder_hidden_states"] = reference_model_inputs.pop( "encoder_outputs" )[0] else: reference_model_inputs.pop("encoder_outputs") - return self._decoder_onnx_config.generate_dummy_inputs_for_validation( - reference_model_inputs, onnx_input_names + return self._decoder_ov_config.generate_dummy_inputs_for_validation( + reference_model_inputs, ov_input_names ) diff --git a/optimum/exporters/openvino/convert.py b/optimum/exporters/openvino/convert.py index b2b0c640e6..946824401b 100644 --- a/optimum/exporters/openvino/convert.py +++ b/optimum/exporters/openvino/convert.py @@ -93,14 +93,14 @@ if TYPE_CHECKING: - from optimum.exporters.openvino.base import OnnxConfig + from optimum.exporters.openvino.base import OpenVINOConfig from optimum.intel.openvino.configuration import OVConfig def _set_runtime_options( models_and_export_configs: Dict[ str, - Tuple[Union["PreTrainedModel", "ModelMixin", "DiffusionPipeline"], "OnnxConfig"], + Tuple[Union["PreTrainedModel", "ModelMixin", "DiffusionPipeline"], "OpenVINOConfig"], ], task: str, library_name: str, @@ -129,7 +129,7 @@ def _save_model( path: str, ov_config: Optional["OVConfig"] = None, library_name: Optional[str] = None, - config: "OnnxConfig" = None, + config: "OpenVINOConfig" = None, ): compress_to_fp16 = ov_config is not None and ov_config.dtype == "fp16" model = _add_version_info_to_model(model, library_name) @@ -147,9 +147,8 @@ def _save_model( def export( model: Union["PreTrainedModel", "ModelMixin", "DiffusionPipeline"], - config: "OnnxConfig", + config: "OpenVINOConfig", output: Path, - opset: Optional[int] = None, device: str = "cpu", input_shapes: Optional[Dict] = None, model_kwargs: Optional[Dict[str, Any]] = None, @@ -164,12 +163,10 @@ def export( Args: model ([`PreTrainedModel`]): The model to export. - config ([`~exporters.onnx.config.OnnxConfig`]): - The ONNX configuration associated with the exported model. + config ([`~exporters.onnx.config.OpenVINOConfig`]): + The OpenVINO configuration associated with the exported model. output (`Path`): Directory to store the exported model. - opset (`Optional[int]`, defaults to `None`): - The version of the ONNX operator set to use. device (`str`, *optional*, defaults to `cpu`): The device on which the model will be exported. Either `cpu` or `cuda`. Only PyTorch is supported for export on CUDA devices. @@ -182,7 +179,7 @@ def export( Returns: `Tuple[List[str], List[str]]`: A tuple with an ordered list of the model's inputs, and the named inputs from - the ONNX configuration. + the OpenVINO configuration. """ if not is_torch_available(): raise ImportError( @@ -222,7 +219,6 @@ def export( return export_pytorch( model, config, - opset, output, device=device, input_shapes=input_shapes, @@ -238,8 +234,7 @@ def export( def export_pytorch( model: Union["PreTrainedModel", "ModelMixin"], - config: "OnnxConfig", - opset: int, + config: "OpenVINOConfig", output: Path, device: str = "cpu", input_shapes: Optional[Dict] = None, @@ -255,10 +250,8 @@ def export_pytorch( Args: model ([`PreTrainedModel`]): The model to export. - config ([`~exporters.onnx.config.OnnxConfig`]): + config ([`~exporters.onnx.config.OpenVINOConfig`]): The configuration associated with the exported model. - opset (`int`): - The version of the ONNX operator set to use. output (`Path`): Directory to store the exported model. device (`str`, defaults to `"cpu"`): @@ -275,7 +268,7 @@ def export_pytorch( Returns: `Tuple[List[str], List[str], bool]`: A tuple with an ordered list of the model's inputs, and the named inputs from - the ONNX configuration and boolean flag - was legacy ONNX path were applied to model or not. + the OpenVINO configuration and boolean flag - was legacy OpenVINO path were applied to model or not. """ import torch from torch.utils._pytree import tree_map @@ -416,10 +409,9 @@ def ts_patched_forward(*args, **kwargs): def export_models( models_and_export_configs: Dict[ - str, Tuple[Union["PreTrainedModel", "ModelMixin", "DiffusionPipeline"], "OnnxConfig"] + str, Tuple[Union["PreTrainedModel", "ModelMixin", "DiffusionPipeline"], "OpenVINOConfig"] ], output_dir: Path, - opset: Optional[int] = None, output_names: Optional[List[str]] = None, device: str = "cpu", input_shapes: Optional[Dict] = None, @@ -433,9 +425,8 @@ def export_models( Export the models to OpenVINO IR format Args: - models_and_export_configs (Dict[ str, Tuple[Union["PreTrainedModel", "ModelMixin"], "OnnxConfig"]): + models_and_export_configs (Dict[ str, Tuple[Union["PreTrainedModel", "ModelMixin"], "OpenVINOConfig"]): output_dir (Path): output directory for saving models - opset (Optional[int], optional, Default to None): ONNX export opset output_names (Optional[List[str]], optional, Defaults to None): model output names device (str, optional, Defaults to "cpu"): The device on which the model will be exported. Either `cpu` or `cuda`. Only PyTorch is supported for @@ -453,7 +444,7 @@ def export_models( ValueError: if custom names set not equal of number of models Returns: - list of input_names and output_names from ONNX configuration + list of input_names and output_names from OpenVINO configuration """ outputs = [] @@ -473,7 +464,6 @@ def export_models( model=submodel, config=sub_export_config, output=output_path, - opset=opset, device=device, input_shapes=input_shapes, model_kwargs=model_kwargs, @@ -580,9 +570,8 @@ def export_from_model( task: Optional[str] = None, ov_config: Optional["OVConfig"] = None, stateful: bool = True, - opset: Optional[int] = None, model_kwargs: Optional[Dict[str, Any]] = None, - custom_export_configs: Optional[Dict[str, "OnnxConfig"]] = None, + custom_export_configs: Optional[Dict[str, "OpenVINOConfig"]] = None, fn_get_submodels: Optional[Callable] = None, preprocessors: List = None, device: str = "cpu", @@ -658,7 +647,7 @@ def export_from_model( and not getattr(model, "_supports_cache_class", is_transformers_version(">=", "4.54")) ): stateful = False - # TODO: support onnx_config.py in the model repo + # TODO: support ov_config.py in the model repo if custom_architecture and custom_export_configs is None: raise ValueError( f"Trying to export a {model_type} model, that is a custom or unsupported architecture, but no custom export configuration was passed as `custom_export_configs`. Please refer to https://huggingface.co/docs/optimum/main/en/exporters/onnx/usage_guides/export_a_model#custom-export-of-transformers-models for an example on how to export custom models. Please open an issue at https://github.com/huggingface/optimum-intel/issues if you would like the model type {model_type} to be supported natively in the OpenVINO export." @@ -822,7 +811,6 @@ def export_from_model( device=device, ov_config=ov_config, stateful=stateful_submodels, - opset=opset, model_kwargs=model_kwargs, patch_16bit_model=patch_16bit_model, library_name=library_name, diff --git a/optimum/exporters/openvino/model_configs.py b/optimum/exporters/openvino/model_configs.py index 19c5ac9866..7ff180d4da 100644 --- a/optimum/exporters/openvino/model_configs.py +++ b/optimum/exporters/openvino/model_configs.py @@ -18,17 +18,17 @@ from transformers import AutoConfig, PretrainedConfig, PreTrainedModel -from optimum.exporters.openvino.base import ConfigBehavior, OnnxConfig, OnnxConfigWithPast, OnnxSeq2SeqConfigWithPast +from optimum.exporters.openvino.base import ConfigBehavior, OpenVINOConfig, OpenVINOConfigWithPast, OpenVINOSeq2SeqConfigWithPast from optimum.exporters.openvino.config import ( - AudioOnnxConfig, - AudioToTextOnnxConfig, - EncoderDecoderBaseOnnxConfig, - TextAndVisionOnnxConfig, - TextDecoderOnnxConfig, - TextDecoderWithPositionIdsOnnxConfig, - TextEncoderOnnxConfig, - TextSeq2SeqOnnxConfig, - VisionOnnxConfig, + AudioOpenVINOConfig, + AudioToTextOpenVINOConfig, + EncoderDecoderBaseOpenVINOConfig, + TextAndVisionOpenVINOConfig, + TextDecoderOpenVINOConfig, + TextDecoderWithPositionIdsOpenVINOConfig, + TextEncoderOpenVINOConfig, + TextSeq2SeqOpenVINOConfig, + VisionOpenVINOConfig, ) from optimum.exporters.openvino.input_generators import ( AquilaDummyPastKeyValuesGenerator, @@ -336,8 +336,7 @@ def init_model_configs(): @register_in_tasks_manager("baichuan", *["text-generation", "text-generation-with-past"], library_name="transformers") -class BaichaunOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): - DEFAULT_ONNX_OPSET = 13 +class BaichaunOpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args( num_layers="num_hidden_layers", num_attention_heads="num_attention_heads", hidden_size="hidden_size" ) @@ -357,8 +356,7 @@ class BaichaunOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): ], library_name="transformers", ) -class Qwen2OpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): - DEFAULT_ONNX_OPSET = 14 +class Qwen2OpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, MistralDummyPastKeyValuesGenerator) DUMMY_PKV_GENERATOR_CLASS = MistralDummyPastKeyValuesGenerator NORMALIZED_CONFIG_CLASS = NormalizedTextConfig @@ -366,8 +364,7 @@ class Qwen2OpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): @register_in_tasks_manager("qwen2_moe", *["text-generation", "text-generation-with-past"], library_name="transformers") -class Qwen2MoEOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): - DEFAULT_ONNX_OPSET = 14 +class Qwen2MoEOpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, MistralDummyPastKeyValuesGenerator) DUMMY_PKV_GENERATOR_CLASS = MistralDummyPastKeyValuesGenerator NORMALIZED_CONFIG_CLASS = NormalizedTextConfig @@ -385,7 +382,7 @@ class Qwen2MoEOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): ], library_name="transformers", ) -class Qwen3OpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): +class Qwen3OpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): MIN_TRANSFORMERS_VERSION = "4.51.0" DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, GemmaDummyPastKeyValuesGenerator) DUMMY_PKV_GENERATOR_CLASS = GemmaDummyPastKeyValuesGenerator @@ -412,7 +409,7 @@ def inputs(self) -> Dict[str, Dict[int, str]]: ], library_name="transformers", ) -class Qwen3VLTextOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): +class Qwen3VLTextOpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): MIN_TRANSFORMERS_VERSION = "4.57.0" DUMMY_INPUT_GENERATOR_CLASSES = (DummyQwen3VLLMInputGenerator, GemmaDummyPastKeyValuesGenerator) @@ -438,8 +435,7 @@ class Qwen3MoEOpenVINOConfig(Qwen3OpenVINOConfig): @register_in_tasks_manager("minicpm", *["text-generation", "text-generation-with-past"], library_name="transformers") -class MiniCPMOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): - DEFAULT_ONNX_OPSET = 14 +class MiniCPMOpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): MAX_TRANSFORMERS_VERSION = "4.53.3" DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, MistralDummyPastKeyValuesGenerator) DUMMY_PKV_GENERATOR_CLASS = MistralDummyPastKeyValuesGenerator @@ -448,8 +444,7 @@ class MiniCPMOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): @register_in_tasks_manager("minicpm3", *["text-generation", "text-generation-with-past"], library_name="transformers") -class MiniCPM3OpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): - DEFAULT_ONNX_OPSET = 14 +class MiniCPM3OpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): MAX_TRANSFORMERS_VERSION = "4.53.3" DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, OVMiniCPM3DummyPastKeyValuesGenerator) DUMMY_PKV_GENERATOR_CLASS = OVMiniCPM3DummyPastKeyValuesGenerator @@ -458,8 +453,7 @@ class MiniCPM3OpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): @register_in_tasks_manager("stablelm", *["text-generation", "text-generation-with-past"], library_name="transformers") -class StableLMOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): - DEFAULT_ONNX_OPSET = 14 +class StableLMOpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, MistralDummyPastKeyValuesGenerator) DUMMY_PKV_GENERATOR_CLASS = MistralDummyPastKeyValuesGenerator NORMALIZED_CONFIG_CLASS = NormalizedTextConfig @@ -467,7 +461,7 @@ class StableLMOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): @register_in_tasks_manager("chatglm", *["text-generation", "text-generation-with-past"], library_name="transformers") -class ChatGLM2OpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): +class ChatGLM2OpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args(vocab_size="padded_vocab_size", num_layers="num_layers") DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, ChatGLM2DummyPastKeyValuesGenerator) DUMMY_PKV_GENERATOR_CLASS = ChatGLM2DummyPastKeyValuesGenerator @@ -496,7 +490,7 @@ def generate_dummy_inputs(self, framework: str = "pt", **kwargs): break if not input_was_inserted: raise RuntimeError( - f'Could not generate dummy input for "{input_name}". Try adding a proper dummy input generator to the model ONNX config.' + f'Could not generate dummy input for "{input_name}". Try adding a proper dummy input generator to the model OpenVINO config.' ) # refer to https://github.com/huggingface/optimum/pull/764 @@ -556,12 +550,10 @@ def add_past_key_values(self, inputs_or_outputs: Dict[str, Dict[int, str]], dire @register_in_tasks_manager("mixtral", *["text-generation", "text-generation-with-past"], library_name="transformers") -class MixtralOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): - # The ONNX export of this architecture needs the Trilu operator support, available since opset 14 - DEFAULT_ONNX_OPSET = 14 +class MixtralOpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): DUMMY_INPUT_GENERATOR_CLASSES = ( MistralDummyPastKeyValuesGenerator, - ) + TextDecoderOnnxConfig.DUMMY_INPUT_GENERATOR_CLASSES + ) + TextDecoderOpenVINOConfig.DUMMY_INPUT_GENERATOR_CLASSES DUMMY_PKV_GENERATOR_CLASS = MistralDummyPastKeyValuesGenerator NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args(num_key_value_heads="num_key_value_heads", allow_new=True) _MODEL_PATCHER = MixtralModelPatcher @@ -578,7 +570,7 @@ class MixtralOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): ], library_name="transformers", ) -class GemmaOpenVINOConfig(TextDecoderOnnxConfig): +class GemmaOpenVINOConfig(TextDecoderOpenVINOConfig): NORMALIZED_CONFIG_CLASS = NormalizedTextConfig DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, GemmaDummyPastKeyValuesGenerator) DUMMY_PKV_GENERATOR_CLASS = GemmaDummyPastKeyValuesGenerator @@ -587,7 +579,7 @@ class GemmaOpenVINOConfig(TextDecoderOnnxConfig): @property def inputs(self) -> Dict[str, Dict[int, str]]: - # position_ids was removed from optimum-onnx's gemma config because + # position_ids was removed from optimum-openvino's gemma config because # it's not necessary (it's correctly generated inside the model) # but openvino genai requires it to be present to work properly inputs = super().inputs @@ -608,7 +600,7 @@ def inputs(self) -> Dict[str, Dict[int, str]]: ], library_name="transformers", ) -class LlamaOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): +class LlamaOpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, MistralDummyPastKeyValuesGenerator) DUMMY_PKV_GENERATOR_CLASS = MistralDummyPastKeyValuesGenerator NORMALIZED_CONFIG_CLASS = NormalizedTextConfig @@ -708,7 +700,7 @@ class GptOssOpenVINOConfig(LlamaOpenVINOConfig): ], library_name="transformers", ) -class BitnetOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): +class BitnetOpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, MistralDummyPastKeyValuesGenerator) DUMMY_PKV_GENERATOR_CLASS = MistralDummyPastKeyValuesGenerator NORMALIZED_CONFIG_CLASS = NormalizedTextConfig @@ -777,8 +769,7 @@ class Cohere2OpenVINOConfig(LlamaOpenVINOConfig): @register_in_tasks_manager("qwen", *["text-generation", "text-generation-with-past"]) -class QwenOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): - DEFAULT_ONNX_OPSET = 14 +class QwenOpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): MAX_TRANSFORMERS_VERSION = "4.55.4" NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args( num_layers="num_hidden_layers", num_attention_heads="num_attention_heads", hidden_size="hidden_size" @@ -809,7 +800,7 @@ def generate_dummy_inputs(self, framework: str = "pt", **kwargs): break if not input_was_inserted: raise RuntimeError( - f'Could not generate dummy input for "{input_name}". Try adding a proper dummy input generator to the model ONNX config.' + f'Could not generate dummy input for "{input_name}". Try adding a proper dummy input generator to the model OpenVINO config.' ) # refer to https://github.com/huggingface/optimum/pull/764 @@ -859,8 +850,7 @@ def add_past_key_values(self, inputs_or_outputs: Dict[str, Dict[int, str]], dire @register_in_tasks_manager( "starcoder2", *["text-generation", "text-generation-with-past"], library_name="transformers" ) -class Starcoder2OpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): - DEFAULT_ONNX_OPSET = 14 +class Starcoder2OpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, MistralDummyPastKeyValuesGenerator) DUMMY_PKV_GENERATOR_CLASS = MistralDummyPastKeyValuesGenerator NORMALIZED_CONFIG_CLASS = NormalizedTextConfig @@ -868,8 +858,7 @@ class Starcoder2OpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): @register_in_tasks_manager("internlm2", *["text-generation", "text-generation-with-past"], library_name="transformers") -class InternLM2OpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): - DEFAULT_ONNX_OPSET = 14 +class InternLM2OpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): MAX_TRANSFORMERS_VERSION = "4.57.6" DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, MistralDummyPastKeyValuesGenerator) DUMMY_PKV_GENERATOR_CLASS = MistralDummyPastKeyValuesGenerator @@ -878,8 +867,7 @@ class InternLM2OpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): @register_in_tasks_manager("orion", *["text-generation", "text-generation-with-past"], library_name="transformers") -class OrionOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): - DEFAULT_ONNX_OPSET = 14 +class OrionOpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): MAX_TRANSFORMERS_VERSION = "4.57.6" DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, MistralDummyPastKeyValuesGenerator) DUMMY_PKV_GENERATOR_CLASS = MistralDummyPastKeyValuesGenerator @@ -889,14 +877,13 @@ class OrionOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): @register_in_tasks_manager("olmo", *["text-generation", "text-generation-with-past"], library_name="transformers") class OlmoOpenVINOConfig(LlamaOpenVINOConfig): - DEFAULT_ONNX_OPSET = 14 NORMALIZED_CONFIG_CLASS = NormalizedTextConfig @register_in_tasks_manager( "mpt", *["text-generation", "text-generation-with-past", "text-classification"], library_name="transformers" ) -class MPTOpenVINOConfig(TextDecoderOnnxConfig): +class MPTOpenVINOConfig(TextDecoderOpenVINOConfig): NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args( num_attention_heads="n_heads", hidden_size="d_model", num_layers="n_layers" ) @@ -914,11 +901,11 @@ class MPTOpenVINOConfig(TextDecoderOnnxConfig): ], library_name="transformers", ) -class Phi3OpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): +class Phi3OpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args(num_key_value_heads="num_key_value_heads", allow_new=True) DUMMY_INPUT_GENERATOR_CLASSES = ( MistralDummyPastKeyValuesGenerator, - ) + TextDecoderOnnxConfig.DUMMY_INPUT_GENERATOR_CLASSES + ) + TextDecoderOpenVINOConfig.DUMMY_INPUT_GENERATOR_CLASSES DUMMY_PKV_GENERATOR_CLASS = MistralDummyPastKeyValuesGenerator MIN_TRANSFORMERS_VERSION = "4.36.0" _MODEL_PATCHER = Phi3ModelPatcher @@ -951,7 +938,7 @@ class PhiMoEOpenVINOConfig(Phi3OpenVINOConfig): ], library_name="transformers", ) -class PhiOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): +class PhiOpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): NORMALIZED_CONFIG_CLASS = NormalizedTextConfig MIN_TRANSFORMERS_VERSION = "4.36.0" _MODEL_PATCHER = OVDecoderModelPatcher @@ -969,11 +956,11 @@ class PhiOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): ], library_name="transformers", ) -class FalconOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): +class FalconOpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): NORMALIZED_CONFIG_CLASS = NormalizedTextConfig DUMMY_INPUT_GENERATOR_CLASSES = ( OVFalconDummyPastKeyValuesGenerator, - ) + TextDecoderOnnxConfig.DUMMY_INPUT_GENERATOR_CLASSES + ) + TextDecoderOpenVINOConfig.DUMMY_INPUT_GENERATOR_CLASSES DUMMY_PKV_GENERATOR_CLASS = OVFalconDummyPastKeyValuesGenerator _MODEL_PATCHER = FalconModelPatcher @@ -998,15 +985,14 @@ def inputs(self) -> Dict[str, Dict[int, str]]: ], library_name="transformers", ) -class PersimmonOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): - DEFAULT_ONNX_OPSET = 14 +class PersimmonOpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): NORMALIZED_CONFIG_CLASS = NormalizedTextConfig _MODEL_PATCHER = PersimmonModelPatcher @register_in_tasks_manager("biogpt", *["text-generation", "text-generation-with-past"], library_name="transformers") class BioGPTOpenVINOConfig( - TextDecoderWithPositionIdsOnnxConfig if is_transformers_version(">=", "4.52.0") else TextDecoderOnnxConfig + TextDecoderWithPositionIdsOpenVINOConfig if is_transformers_version(">=", "4.52.0") else TextDecoderOpenVINOConfig ): NORMALIZED_CONFIG_CLASS = NormalizedTextConfig _MODEL_PATCHER = OVDecoderModelPatcher @@ -1023,7 +1009,7 @@ class BioGPTOpenVINOConfig( ], library_name="transformers", ) -class GPTNeoOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): +class GPTNeoOpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args(num_attention_heads="num_heads") _MODEL_PATCHER = GptNeoModelPatcher @@ -1039,7 +1025,7 @@ class GPTNeoOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): ], library_name="transformers", ) -class GPTJOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): +class GPTJOpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args(num_layers="n_layer", num_attention_heads="n_head") _MODEL_PATCHER = GptJModelPatcher @@ -1056,7 +1042,7 @@ class GPTJOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): ], library_name="transformers", ) -class BloomOpenVINOConfig(TextDecoderOnnxConfig): +class BloomOpenVINOConfig(TextDecoderOpenVINOConfig): # Bloom does not require position_ids input. MIN_TRANSFORMERS_VERSION = "4.36.0" DUMMY_PKV_GENERATOR_CLASS = BloomDummyPastKeyValuesGenerator @@ -1099,8 +1085,7 @@ class CohereOpenVINOConfig(LlamaOpenVINOConfig): @register_in_tasks_manager("xglm", *["text-generation", "text-generation-with-past"], library_name="transformers") -class XGLMConfig(TextDecoderWithPositionIdsOnnxConfig): - DEFAULT_ONNX_OPSET = 13 +class XGLMConfig(TextDecoderWithPositionIdsOpenVINOConfig): NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args( num_attention_heads="attention_heads", hidden_size="d_model" ) @@ -1112,8 +1097,7 @@ def __init__(self, *args, **kwargs): @register_in_tasks_manager("aquila", *["text-generation", "text-generation-with-past"], library_name="transformers") -class AquilaMOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): - DEFAULT_ONNX_OPSET = 14 +class AquilaMOpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): MAX_TRANSFORMERS_VERSION = "4.57.6" DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, AquilaDummyPastKeyValuesGenerator) DUMMY_PKV_GENERATOR_CLASS = AquilaDummyPastKeyValuesGenerator @@ -1122,8 +1106,7 @@ class AquilaMOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): @register_in_tasks_manager("xverse", *["text-generation", "text-generation-with-past"], library_name="transformers") -class XverseMOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): - DEFAULT_ONNX_OPSET = 14 +class XverseMOpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): MAX_TRANSFORMERS_VERSION = "4.57.6" DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, DummyPastKeyValuesGenerator) DUMMY_PKV_GENERATOR_CLASS = DummyPastKeyValuesGenerator @@ -1132,8 +1115,7 @@ class XverseMOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): @register_in_tasks_manager("internlm", *["text-generation", "text-generation-with-past"], library_name="transformers") -class InternLMOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): - DEFAULT_ONNX_OPSET = 14 +class InternLMOpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): MAX_TRANSFORMERS_VERSION = "4.57.6" DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, DummyPastKeyValuesGenerator) DUMMY_PKV_GENERATOR_CLASS = DummyPastKeyValuesGenerator @@ -1146,7 +1128,7 @@ class InternLMOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): *["feature-extraction", "feature-extraction-with-past", "text-generation", "text-generation-with-past"], library_name="transformers", ) -class CodeGenOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): +class CodeGenOpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args(num_layers="n_layer", num_attention_heads="n_head") _MODEL_PATCHER = CodeGenModelPatcher @@ -1156,8 +1138,7 @@ class CodeGenOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): *["text-generation", "text-generation-with-past"], library_name="transformers", ) -class DBRXOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): - DEFAULT_ONNX_OPSET = 14 +class DBRXOpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): MAX_TRANSFORMERS_VERSION = "4.57.6" NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args( num_attention_heads="n_heads", @@ -1176,8 +1157,7 @@ class DBRXOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): *["text-generation", "text-generation-with-past"], library_name="transformers", ) -class JaisOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): - DEFAULT_ONNX_OPSET = 14 +class JaisOpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): MAX_TRANSFORMERS_VERSION = "4.57.6" NORMALIZED_CONFIG_CLASS = NormalizedTextConfig DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, DummyPastKeyValuesGenerator) @@ -1202,11 +1182,11 @@ class ArcticOpenVINOConfig(MixtralOpenVINOConfig): ], library_name="transformers", ) -class MistralOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): +class MistralOpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args(num_key_value_heads="num_key_value_heads", allow_new=True) DUMMY_INPUT_GENERATOR_CLASSES = ( MistralDummyPastKeyValuesGenerator, - ) + TextDecoderOnnxConfig.DUMMY_INPUT_GENERATOR_CLASSES + ) + TextDecoderOpenVINOConfig.DUMMY_INPUT_GENERATOR_CLASSES DUMMY_PKV_GENERATOR_CLASS = MistralDummyPastKeyValuesGenerator _MODEL_PATCHER = MistralModelPatcher @@ -1222,7 +1202,7 @@ class MistralOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): ], library_name="transformers", ) -class GPTNeoxOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): +class GPTNeoxOpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): NORMALIZED_CONFIG_CLASS = NormalizedTextConfig _MODEL_PATCHER = GptNeoxModelPatcher @@ -1230,9 +1210,8 @@ class GPTNeoxOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): @register_in_tasks_manager( "gpt_neox_japanese", *["text-generation", "text-generation-with-past"], library_name="transformers" ) -class GPTNeoxJapaneseOpenVINOConfig(TextDecoderOnnxConfig): +class GPTNeoxJapaneseOpenVINOConfig(TextDecoderOpenVINOConfig): # GPTNeoxJapanese does not require position_ids input. - DEFAULT_ONNX_OPSET = 13 NORMALIZED_CONFIG_CLASS = NormalizedTextConfig _MODEL_PATCHER = GptNeoxModelPatcher @@ -1308,8 +1287,7 @@ def add_past_key_values(self, inputs_or_outputs: dict[str, dict[int, str]], dire @register_in_tasks_manager("deci", *["text-generation", "text-generation-with-past"], library_name="transformers") -class DeciOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): - DEFAULT_ONNX_OPSET = 14 +class DeciOpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): MAX_TRANSFORMERS_VERSION = "4.57.6" DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, DeciDummyPastKeyValuesGenerator) DUMMY_PKV_GENERATOR_CLASS = DeciDummyPastKeyValuesGenerator @@ -1327,8 +1305,7 @@ class SiglipNormalizedConfig(CLIPNormalizedConfig): @register_in_tasks_manager("clip", *["zero-shot-image-classification"], library_name="open_clip") -class OpenCLIPOpenVINOConfig(TextAndVisionOnnxConfig): - DEFAULT_ONNX_OPSET = 14 +class OpenCLIPOpenVINOConfig(TextAndVisionOpenVINOConfig): NORMALIZED_CONFIG_CLASS = CLIPNormalizedConfig _MODEL_PATCHER = CLIPModelPatcher @@ -1359,20 +1336,19 @@ def generate_dummy_inputs(self, framework: str = "pt", **kwargs): return super().generate_dummy_inputs(framework, **kwargs) def generate_dummy_inputs_for_validation( - self, reference_model_inputs: Dict[str, Any], onnx_input_names: Optional[List[str]] = None + self, reference_model_inputs: Dict[str, Any], ov_input_names: Optional[List[str]] = None ) -> Dict[str, Any]: if "attention_mask" in reference_model_inputs: reference_model_inputs.pop("attention_mask") - if "image" in onnx_input_names and "pixel_values" in reference_model_inputs: + if "image" in ov_input_names and "pixel_values" in reference_model_inputs: reference_model_inputs["image"] = reference_model_inputs.pop("pixel_values") - if "text" in onnx_input_names and "input_ids" in reference_model_inputs: + if "text" in ov_input_names and "input_ids" in reference_model_inputs: reference_model_inputs["text"] = reference_model_inputs.pop("input_ids") return super().generate_dummy_inputs_for_validation(reference_model_inputs) @register_in_tasks_manager("clip_text_model", *["feature-extraction"], library_name="open_clip") -class OpenCLIPTextOpenVINOConfig(TextEncoderOnnxConfig): - DEFAULT_ONNX_OPSET = 14 +class OpenCLIPTextOpenVINOConfig(TextEncoderOpenVINOConfig): NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args( vocab_size="vocab_size", sequence_length="max_position_embeddings", @@ -1408,9 +1384,7 @@ def generate_dummy_inputs(self, framework: str = "pt", **kwargs): @register_in_tasks_manager("clip_vision_model", *["feature-extraction"], library_name="open_clip") -class OpenCLIPVisualOpenVINOConfig(VisionOnnxConfig): - DEFAULT_ONNX_OPSET = 14 - +class OpenCLIPVisualOpenVINOConfig(VisionOpenVINOConfig): NORMALIZED_CONFIG_CLASS = NormalizedVisionConfig @property @@ -1434,7 +1408,7 @@ def rename_ambiguous_inputs(self, inputs): @register_in_tasks_manager( "clip", *["feature-extraction", "zero-shot-image-classification"], library_name="transformers" ) -class CLIPOpenVINOConfig(TextAndVisionOnnxConfig): +class CLIPOpenVINOConfig(TextAndVisionOpenVINOConfig): NORMALIZED_CONFIG_CLASS = CLIPNormalizedConfig _MODEL_PATCHER = CLIPModelPatcher @@ -1467,7 +1441,7 @@ def outputs(self) -> Dict[str, Dict[int, str]]: @register_in_tasks_manager("clip_text_model", *["feature-extraction"], library_name="transformers") @register_in_tasks_manager("clip-text", *["feature-extraction"], library_name="diffusers") -class CLIPTextOpenVINOConfig(TextEncoderOnnxConfig): +class CLIPTextOpenVINOConfig(TextEncoderOpenVINOConfig): NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args( vocab_size="vocab_size", sequence_length="max_position_embeddings", @@ -1497,7 +1471,7 @@ def outputs(self) -> Dict[str, Dict[int, str]]: @register_in_tasks_manager("clip-text-with-projection", *["feature-extraction"], library_name="diffusers") -class CLIPTextWithProjectionOpenVINOConfig(TextEncoderOnnxConfig): +class CLIPTextWithProjectionOpenVINOConfig(TextEncoderOpenVINOConfig): NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args( vocab_size="vocab_size", sequence_length="max_position_embeddings", @@ -1526,7 +1500,7 @@ def outputs(self) -> Dict[str, Dict[int, str]]: @register_in_tasks_manager("clip_vision_model", *["feature-extraction"], library_name="transformers") -class CLIPVisionModelOpenVINOConfig(VisionOnnxConfig): +class CLIPVisionModelOpenVINOConfig(VisionOpenVINOConfig): NORMALIZED_CONFIG_CLASS = NormalizedVisionConfig _MODEL_PATCHER = CLIPModelPatcher @@ -1555,7 +1529,7 @@ def outputs(self) -> Dict[str, Dict[int, str]]: ], library_name="transformers", ) -class IBertOpenVINOConfig(TextEncoderOnnxConfig): +class IBertOpenVINOConfig(TextEncoderOpenVINOConfig): NORMALIZED_CONFIG_CLASS = NormalizedTextConfig _MODEL_PATCHER = IBertModelPatcher @@ -1569,7 +1543,7 @@ def inputs(self) -> Dict[str, Dict[int, str]]: # TODO: this is a very confusing class TBH, why not simply decompose the VLM into components, like diffusion models ? -class LMInputEmbedsConfigHelper(TextDecoderWithPositionIdsOnnxConfig): +class LMInputEmbedsConfigHelper(TextDecoderWithPositionIdsOpenVINOConfig): def __init__(self, export_config, patcher_cls=None, dummy_input_generator=None, inputs_update=None): self.orig_export_config = export_config if dummy_input_generator is not None: @@ -1577,7 +1551,6 @@ def __init__(self, export_config, patcher_cls=None, dummy_input_generator=None, dummy_input_generator, ) + export_config.DUMMY_INPUT_GENERATOR_CLASSES self.DUMMY_INPUT_GENERATOR_CLASSES = export_config.DUMMY_INPUT_GENERATOR_CLASSES - self.DEFAULT_ONNX_OPSET = export_config.DEFAULT_ONNX_OPSET self.DUMMY_PKV_GENERATOR_CLASS = export_config.DUMMY_PKV_GENERATOR_CLASS self._config = export_config._config self._normalized_config = export_config._normalized_config @@ -1638,7 +1611,7 @@ def generate_dummy_inputs(self, framework: str = "pt", **kwargs): return dummy_inputs -class InputEmbedOpenVINOConfig(TextDecoderOnnxConfig): +class InputEmbedOpenVINOConfig(TextDecoderOpenVINOConfig): NORMALIZED_CONFIG_CLASS = NormalizedTextConfig _MODEL_PATCHER = InputEmbeddingPatcher @@ -1717,7 +1690,7 @@ class VLMConfigBehavior(str, enum.Enum): LANGUAGE = "language" -class BaseVLMOpenVINOConfig(OnnxConfig): +class BaseVLMOpenVINOConfig(OpenVINOConfig): SUPPORTED_BEHAVIORS = [model_type.value for model_type in VLMConfigBehavior] NORMALIZED_CONFIG_CLASS = NormalizedVisionConfig DUMMY_INPUT_GENERATOR_CLASSES = (DummyVisionInputGenerator,) @@ -1858,7 +1831,7 @@ class LlavaNextOpenVINOConfig(LlavaOpenVINOConfig): _OV_2026_1_MODEL_TYPE = "llava_next" -class LLavaMultimodalProjectorOpenVINOConfig(OnnxConfig): +class LLavaMultimodalProjectorOpenVINOConfig(OpenVINOConfig): DUMMY_INPUT_GENERATOR_CLASSES = (DummyLLavaMultiModalProjectorInputGenerator,) NORMALIZED_CONFIG_CLASS = NormalizedVisionConfig @@ -2168,7 +2141,7 @@ def rename_ambiguous_inputs(self, inputs): @register_in_tasks_manager("unet", *["semantic-segmentation"], library_name="diffusers") @register_in_tasks_manager("unet-2d-condition", *["semantic-segmentation"], library_name="diffusers") -class UNetOpenVINOConfig(VisionOnnxConfig): +class UNetOpenVINOConfig(VisionOpenVINOConfig): NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args( image_size="sample_size", num_channels="in_channels", @@ -2210,7 +2183,7 @@ def outputs(self) -> Dict[str, Dict[int, str]]: } @property - def torch_to_onnx_output_map(self) -> Dict[str, str]: + def torch_to_ov_output_map(self) -> Dict[str, str]: return { "sample": "out_sample", } @@ -2318,7 +2291,7 @@ def rename_ambiguous_inputs(self, inputs): @register_in_tasks_manager("vae-encoder", *["semantic-segmentation"], library_name="diffusers") -class VaeEncoderOpenVINOConfig(VisionOnnxConfig): +class VaeEncoderOpenVINOConfig(VisionOpenVINOConfig): ATOL_FOR_VALIDATION = 3e-4 # TODO: this only happens in test_export.py NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args( num_channels="in_channels", image_size="sample_size", allow_new=True @@ -2344,7 +2317,7 @@ def outputs(self) -> Dict[str, Dict[int, str]]: @register_in_tasks_manager("vae-decoder", *["semantic-segmentation"], library_name="diffusers") -class VaeDecoderOpenVINOConfig(VisionOnnxConfig): +class VaeDecoderOpenVINOConfig(VisionOpenVINOConfig): ATOL_FOR_VALIDATION = 3e-4 # TODO: this only happens in test_export.py NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args(num_channels="latent_channels", allow_new=True) @@ -2368,7 +2341,7 @@ def outputs(self) -> Dict[str, Dict[int, str]]: @register_in_tasks_manager("dcae-encoder", *["semantic-segmentation"], library_name="diffusers") -class DcaeEncoderOpenVINOConfig(VisionOnnxConfig): +class DcaeEncoderOpenVINOConfig(VisionOpenVINOConfig): ATOL_FOR_VALIDATION = 3e-4 # TODO: this only happens in test_export.py NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args( num_channels="in_channels", image_size="sample_size", allow_new=True @@ -2388,7 +2361,7 @@ def outputs(self) -> Dict[str, Dict[int, str]]: @register_in_tasks_manager("dcae-decoder", *["semantic-segmentation"], library_name="diffusers") -class DcaeDecoderOpenVINOConfig(VisionOnnxConfig): +class DcaeDecoderOpenVINOConfig(VisionOpenVINOConfig): ATOL_FOR_VALIDATION = 3e-4 # TODO: this only happens in test_export.py NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args(num_channels="latent_channels", allow_new=True) @@ -2435,7 +2408,7 @@ def inputs(self): @register_in_tasks_manager("ltx-vae-encoder", *["semantic-segmentation"], library_name="diffusers") -class LTXVaeEncoderOpenVINOConfig(VisionOnnxConfig): +class LTXVaeEncoderOpenVINOConfig(VisionOpenVINOConfig): ATOL_FOR_VALIDATION = 3e-4 # TODO: this only happens in test_export.py NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args( num_channels="in_channels", image_size="sample_size", allow_new=True @@ -2456,7 +2429,7 @@ def outputs(self) -> Dict[str, Dict[int, str]]: @register_in_tasks_manager("ltx-vae-decoder", *["semantic-segmentation"], library_name="diffusers") -class LTXVaeDecoderOpenVINOConfig(VisionOnnxConfig): +class LTXVaeDecoderOpenVINOConfig(VisionOpenVINOConfig): ATOL_FOR_VALIDATION = 3e-4 # TODO: this only happens in test_export.py NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args(num_channels="latent_channels", allow_new=True) DUMMY_INPUT_GENERATOR_CLASSES = (LTXVaeDummyInputGenerator,) @@ -3371,7 +3344,7 @@ class GraniteMoEOpenVINOConfig(LlamaOpenVINOConfig): ], library_name="transformers", ) -class WhisperOpenVINOConfig(AudioToTextOnnxConfig): +class WhisperOpenVINOConfig(AudioToTextOpenVINOConfig): NORMALIZED_CONFIG_CLASS = NormalizedSeq2SeqConfig.with_args( encoder_num_layers="encoder_layers", decoder_num_layers="decoder_layers", @@ -3425,7 +3398,7 @@ def outputs(self) -> Dict[str, Dict[int, str]]: ], library_name="transformers", ) -class Qwen3ASROpenVINOConfig(AudioToTextOnnxConfig): +class Qwen3ASROpenVINOConfig(AudioToTextOpenVINOConfig): """OpenVINO export config for Qwen3-ASR model.""" DUMMY_INPUT_GENERATOR_CLASSES = ( @@ -3509,9 +3482,9 @@ def add_past_key_values(self, inputs_or_outputs: Dict[str, Dict[int, str]], dire *["feature-extraction", "feature-extraction-with-past", "text2text-generation", "text2text-generation-with-past"], library_name="transformers", ) -class T5OpenVINOConfig(TextSeq2SeqOnnxConfig): +class T5OpenVINOConfig(TextSeq2SeqOpenVINOConfig): DUMMY_INPUT_GENERATOR_CLASSES = ( - *TextSeq2SeqOnnxConfig.DUMMY_INPUT_GENERATOR_CLASSES[:-1], + *TextSeq2SeqOpenVINOConfig.DUMMY_INPUT_GENERATOR_CLASSES[:-1], T5DummySeq2SeqPastKeyValuesGenerator, ) DUMMY_PKV_GENERATOR_CLASS = T5DummySeq2SeqPastKeyValuesGenerator @@ -3559,7 +3532,7 @@ class LongT5OpenVINOConfig(T5OpenVINOConfig): ], library_name="transformers", ) -class BartOpenVINOConfig(TextSeq2SeqOnnxConfig): +class BartOpenVINOConfig(TextSeq2SeqOpenVINOConfig): PAD_ATTENTION_MASK_TO_PAST = True NORMALIZED_CONFIG_CLASS = NormalizedSeq2SeqConfig.with_args( @@ -3655,7 +3628,7 @@ def outputs(self) -> dict: if self.task in ["feature-extraction", "text2text-generation"]: common_outputs = super().outputs else: - common_outputs = super(OnnxConfigWithPast, self).outputs + common_outputs = super(OpenVINOConfigWithPast, self).outputs if self.use_past: for i in range( self._normalized_config.encoder_num_layers @@ -3673,7 +3646,7 @@ def flatten_past_key_values(self, flattened_output, name, idx, t): if self.task in ["feature-extraction", "text2text-generation"]: flattened_output = super().flatten_past_key_values(flattened_output, name, idx, t) else: - flattened_output = super(OnnxSeq2SeqConfigWithPast, self).flatten_past_key_values( + flattened_output = super(OpenVINOSeq2SeqConfigWithPast, self).flatten_past_key_values( flattened_output, name, idx, t ) @@ -4165,7 +4138,7 @@ class SpeechT5ConfigBehavior(str, enum.Enum): *["text-to-audio", "text-to-audio-with-past"], library_name="transformers", ) -class SpeechT5OpenVINOConfig(OnnxSeq2SeqConfigWithPast): +class SpeechT5OpenVINOConfig(OpenVINOSeq2SeqConfigWithPast): DUMMY_INPUT_GENERATOR_CLASSES = ( DummyTextInputGenerator, DummySeq2SeqPastKeyValuesGenerator, @@ -4359,7 +4332,7 @@ def patch_model_for_export(self, model: PreTrainedModel, model_kwargs: Optional[ "falcon_mamba", *["text-generation", "text-generation-with-past"], library_name="transformers" ) @register_in_tasks_manager("mamba", *["text-generation", "text-generation-with-past"], library_name="transformers") -class MambaOpenVINOConfig(TextDecoderOnnxConfig): +class MambaOpenVINOConfig(TextDecoderOpenVINOConfig): DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, MambaCacheDummyInputGenerator) DUMMY_PKV_GENERATOR_CLASS = MambaCacheDummyInputGenerator NORMALIZED_CONFIG_CLASS = NormalizedTextConfig @@ -4426,7 +4399,7 @@ def generate_dummy_inputs(self, framework: str = "pt", **kwargs): break if not input_was_inserted: raise RuntimeError( - f'Could not generate dummy input for "{input_name}". Try adding a proper dummy input generator to the model ONNX config.' + f'Could not generate dummy input for "{input_name}". Try adding a proper dummy input generator to the model OpenVINO config.' ) return dummy_inputs @@ -4444,7 +4417,7 @@ def generate_dummy_inputs(self, framework: str = "pt", **kwargs): ], library_name="transformers", ) -class GPT2OpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): +class GPT2OpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args(num_layers="n_layer", num_attention_heads="n_head") _MODEL_PATCHER = OVDecoderModelPatcher @@ -4458,7 +4431,7 @@ class GPT2OpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): "document-question-answering-with-past", ], ) -class VisionEncoderDecoderOpenVINOConfig(EncoderDecoderBaseOnnxConfig): +class VisionEncoderDecoderOpenVINOConfig(EncoderDecoderBaseOpenVINOConfig): NORMALIZED_CONFIG_CLASS = NormalizedEncoderDecoderConfig DUMMY_INPUT_GENERATOR_CLASSES = (DummyVisionInputGenerator, DummyVisionEncoderDecoderPastKeyValuesGenerator) _MODEL_PATCHER = OVSeq2SeqModelPatcher @@ -4482,7 +4455,7 @@ def inputs(self) -> dict: @property def outputs(self) -> dict: if self._behavior == ConfigBehavior.ENCODER: - return self._encoder_onnx_config.outputs + return self._encoder_ov_config.outputs else: return super().outputs @@ -4632,7 +4605,7 @@ def inputs(self) -> Dict[str, Dict[int, str]]: @register_in_tasks_manager("audio-spectrogram-transformer", *["feature-extraction", "audio-classification"]) -class ASTOpenVINOConfig(OnnxConfig): +class ASTOpenVINOConfig(OpenVINOConfig): NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args( num_mel_bins="num_mel_bins", max_length="max_length", allow_new=True ) @@ -4661,7 +4634,7 @@ def __init__(self, *args, **kwargs): @register_in_tasks_manager("olmo2", *COMMON_TEXT_GENERATION_TASKS, library_name="transformers") -class Olmo2OOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): +class Olmo2OOpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, MistralDummyPastKeyValuesGenerator) DUMMY_PKV_GENERATOR_CLASS = MistralDummyPastKeyValuesGenerator NORMALIZED_CONFIG_CLASS = NormalizedTextConfig @@ -4670,7 +4643,7 @@ class Olmo2OOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): @register_in_tasks_manager("opt", *[*COMMON_TEXT_GENERATION_TASKS, "text-classification", "question-answering"]) class OPTOpenVINOConfig( - TextDecoderWithPositionIdsOnnxConfig if is_transformers_version(">=", "4.46.0") else TextDecoderOnnxConfig + TextDecoderWithPositionIdsOpenVINOConfig if is_transformers_version(">=", "4.46.0") else TextDecoderOpenVINOConfig ): NORMALIZED_CONFIG_CLASS = NormalizedTextConfig @@ -4682,7 +4655,7 @@ def __init__(self, *args, **kwargs): @register_in_tasks_manager( "gpt_bigcode", *[*COMMON_TEXT_GENERATION_TASKS, "text-classification", "token-classification"] ) -class GPTBigCodeOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): +class GPTBigCodeOpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, GPTBigCodeDummyPastKeyValuesGenerator) NORMALIZED_CONFIG_CLASS = NormalizedConfigManager.get_normalized_config_class("gpt_bigcode") DUMMY_PKV_GENERATOR_CLASS = GPTBigCodeDummyPastKeyValuesGenerator @@ -4735,7 +4708,7 @@ class _Pix2StructNormalizedConfig(NormalizedSeq2SeqConfig): "image-to-text-with-past", ], ) -class Pix2StructOpenVINOConfig(OnnxSeq2SeqConfigWithPast): +class Pix2StructOpenVINOConfig(OpenVINOSeq2SeqConfigWithPast): PAD_ATTENTION_MASK_TO_PAST = True NORMALIZED_CONFIG_CLASS = _Pix2StructNormalizedConfig DUMMY_INPUT_GENERATOR_CLASSES = ( @@ -4791,7 +4764,7 @@ def outputs(self) -> dict: def _create_dummy_input_generator_classes(self, **kwargs): if self._preprocessors is None or len(self._preprocessors) < 2: raise ValueError( - f"Preprocessors for pix2struct need to be available for the ONNX export to infer input static shapes. Got: {self._preprocessors}" + f"Preprocessors for pix2struct need to be available for the OpenVINO export to infer input static shapes. Got: {self._preprocessors}" ) dummy_inputs_generators = [] @@ -4810,7 +4783,7 @@ def _create_dummy_input_generator_classes(self, **kwargs): def overwrite_shape_and_generate_input(self, dummy_input_gen, input_name: str, framework: str, input_shapes: dict): if self._preprocessors is None or len(self._preprocessors) < 2: raise ValueError( - f"Preprocessors for pix2struct need to be available for the ONNX export to infer input static shapes. Got: {self._preprocessors}" + f"Preprocessors for pix2struct need to be available for the OpenVINO export to infer input static shapes. Got: {self._preprocessors}" ) if input_name in ["encoder_outputs", "attention_mask"]: @@ -4829,7 +4802,7 @@ def overwrite_shape_and_generate_input(self, dummy_input_gen, input_name: str, f @register_in_tasks_manager("bert", *COMMON_TEXT_TASKS) -class BertOpenVINOConfig(TextEncoderOnnxConfig): +class BertOpenVINOConfig(TextEncoderOpenVINOConfig): NORMALIZED_CONFIG_CLASS = NormalizedTextConfig @property @@ -4887,7 +4860,7 @@ class XLMOpenVINOConfig(BertOpenVINOConfig): @register_in_tasks_manager("distilbert", *COMMON_TEXT_TASKS) -class DistilBertOpenVINOConfig(TextEncoderOnnxConfig): +class DistilBertOpenVINOConfig(TextEncoderOpenVINOConfig): NORMALIZED_CONFIG_CLASS = NormalizedTextConfig @property @@ -4939,7 +4912,7 @@ class DebertaV2OpenVINOConfig(DebertaOpenVINOConfig): @register_in_tasks_manager("hubert", *["feature-extraction", "automatic-speech-recognition", "audio-classification"]) -class HubertOpenVINOConfig(AudioOnnxConfig): +class HubertOpenVINOConfig(AudioOpenVINOConfig): NORMALIZED_CONFIG_CLASS = NormalizedConfig @property @@ -4978,7 +4951,7 @@ class Data2VecTextOpenVINOConfig(DistilBertOpenVINOConfig): @register_in_tasks_manager("vit", *["feature-extraction", "image-classification", "masked-im"]) -class ViTOpenVINOConfig(VisionOnnxConfig): +class ViTOpenVINOConfig(VisionOpenVINOConfig): NORMALIZED_CONFIG_CLASS = NormalizedVisionConfig @property @@ -5001,11 +4974,11 @@ class Data2VecVisionOpenVINOConfig(ViTOpenVINOConfig): @register_in_tasks_manager("perceiver", *["fill-mask", "text-classification", "image-classification"]) -class PerceiverOpenVINOConfig(TextAndVisionOnnxConfig): +class PerceiverOpenVINOConfig(TextAndVisionOpenVINOConfig): NORMALIZED_CONFIG_CLASS = NormalizedTextConfig DUMMY_INPUT_GENERATOR_CLASSES = ( PerceiverDummyInputGenerator, - *TextAndVisionOnnxConfig.DUMMY_INPUT_GENERATOR_CLASSES, + *TextAndVisionOpenVINOConfig.DUMMY_INPUT_GENERATOR_CLASSES, ) def __init__( @@ -5066,7 +5039,7 @@ def generate_dummy_inputs(self, framework: str = "pt", **kwargs): @register_in_tasks_manager("esm", *["feature-extraction", "fill-mask", "text-classification", "token-classification"]) -class EsmOpenVINOConfig(TextEncoderOnnxConfig): +class EsmOpenVINOConfig(TextEncoderOpenVINOConfig): NORMALIZED_CONFIG_CLASS = NormalizedTextConfig @property @@ -5228,7 +5201,7 @@ class WavLMOpenVINOConfig(HubertOpenVINOConfig): @register_in_tasks_manager("sam", *["feature-extraction"]) -class SamOpenVINOConfig(OnnxConfig): +class SamOpenVINOConfig(OpenVINOConfig): NORMALIZED_CONFIG_CLASS = NormalizedEncoderDecoderConfig DUMMY_INPUT_GENERATOR_CLASSES = (DummyVisionInputGenerator, DummyPointsGenerator, DummyVisionEmbeddingsGenerator) VARIANTS = { # noqa: RUF012 @@ -5291,7 +5264,7 @@ def outputs(self) -> dict: @register_in_tasks_manager("siglip", *["feature-extraction", "zero-shot-image-classification"]) -class SiglipOpenVINOConfig(TextAndVisionOnnxConfig): +class SiglipOpenVINOConfig(TextAndVisionOpenVINOConfig): NORMALIZED_CONFIG_CLASS = SiglipNormalizedConfig _MODEL_PATCHER = CLIPModelPatcher @@ -5319,7 +5292,7 @@ def outputs(self) -> dict: @register_in_tasks_manager( "transformer", *["feature-extraction", "sentence-similarity"], library_name="sentence_transformers" ) -class SentenceTransformersTransformerOpenVINOConfig(TextEncoderOnnxConfig): +class SentenceTransformersTransformerOpenVINOConfig(TextEncoderOpenVINOConfig): NORMALIZED_CONFIG_CLASS = NormalizedTextConfig _MODEL_PATCHER = SentenceTransformersTransformerPatcher @@ -5344,7 +5317,7 @@ class RemBertOpenVINOConfig(BertOpenVINOConfig): @register_in_tasks_manager("siglip-text-with-projection", *["feature-extraction"]) -class SiglipTextWithProjectionOpenVINOConfig(TextEncoderOnnxConfig): +class SiglipTextWithProjectionOpenVINOConfig(TextEncoderOpenVINOConfig): NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args( vocab_size="vocab_size", sequence_length="max_position_embeddings", @@ -5390,7 +5363,7 @@ def outputs(self) -> dict: return common_outputs -class VideoChatFlashQwenProjectorOpenVINOConfig(OnnxConfig): +class VideoChatFlashQwenProjectorOpenVINOConfig(OpenVINOConfig): DUMMY_INPUT_GENERATOR_CLASSES = (DummyVideoChatFlashQwenProjectorInputGenerator,) NORMALIZED_CONFIG_CLASS = NormalizedVisionConfig @@ -5611,7 +5584,7 @@ def generate_dummy_inputs(self, framework: str = "pt", **kwargs): break if not input_was_inserted: raise RuntimeError( - f'Could not generate dummy input for "{input_name}". Try adding a proper dummy input generator to the model ONNX config.' + f'Could not generate dummy input for "{input_name}". Try adding a proper dummy input generator to the model OpenVINO config.' ) return dummy_inputs @@ -5698,7 +5671,7 @@ def generate_dummy_inputs(self, framework: str = "pt", **kwargs): break if not input_was_inserted: raise RuntimeError( - f'Could not generate dummy input for "{input_name}". Try adding a proper dummy input generator to the model ONNX config.' + f'Could not generate dummy input for "{input_name}". Try adding a proper dummy input generator to the model OpenVINO config.' ) return dummy_inputs @@ -5869,8 +5842,7 @@ def outputs(self) -> Dict[str, Dict[int, str]]: *["text-to-audio"], library_name="kokoro", ) -class KokoroOpenVINOConfig(OnnxConfig): - DEFAULT_ONNX_OPSET = 14 +class KokoroOpenVINOConfig(OpenVINOConfig): DUMMY_INPUT_GENERATOR_CLASSES = (DummyKokoroInputGenerator,) NORMALIZED_CONFIG_CLASS = NormalizedConfig _MODEL_PATCHER = KokoroModelPatcher @@ -5914,7 +5886,7 @@ def outputs(self) -> Dict[str, Dict[int, str]]: "feature-extraction-with-past", ], ) -class TrOCROpenVINOConfig(TextSeq2SeqOnnxConfig): +class TrOCROpenVINOConfig(TextSeq2SeqOpenVINOConfig): NORMALIZED_CONFIG_CLASS = NormalizedSeq2SeqConfig.with_args( decoder_num_layers="decoder_layers", num_layers="decoder_layers", diff --git a/optimum/exporters/openvino/model_patcher.py b/optimum/exporters/openvino/model_patcher.py index 72a49f9af2..92ac740992 100644 --- a/optimum/exporters/openvino/model_patcher.py +++ b/optimum/exporters/openvino/model_patcher.py @@ -48,7 +48,7 @@ from transformers.utils import ModelOutput from optimum.exporters.openvino._ov_ops import convert_recurrent_attention_cell -from optimum.exporters.openvino.base import OnnxConfig +from optimum.exporters.openvino.base import OpenVINOConfig from optimum.exporters.openvino.patching_utils import ( UNSUPPORTED_OPS_PATCHING_SPEC, ModelPatcher, @@ -105,7 +105,7 @@ class SAMModelPatcher(ModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: PreTrainedModel, model_kwargs: dict[str, Any] | None = None, ): @@ -161,17 +161,17 @@ def patched_forward( sparse_embeddings, dense_embeddings = model.prompt_encoder( input_points=input_points, input_labels=input_labels, - input_boxes=None, # Not supported in the ONNX export - input_masks=None, # Not supported in the ONNX export + input_boxes=None, # Not supported in the OpenVINO export + input_masks=None, # Not supported in the OpenVINO export ) outputs = model.mask_decoder( image_embeddings=image_embeddings, image_positional_embeddings=image_positional_embeddings, sparse_prompt_embeddings=sparse_embeddings, dense_prompt_embeddings=dense_embeddings, - multimask_output=True, # Not supported in the ONNX export - attention_similarity=None, # Not supported in the ONNX export - target_embedding=None, # Not supported in the ONNX export + multimask_output=True, # Not supported in the OpenVINO export + attention_similarity=None, # Not supported in the OpenVINO export + target_embedding=None, # Not supported in the OpenVINO export ) low_res_masks, iou_predictions = outputs[:2] @@ -194,8 +194,8 @@ def patched_speecht5_prenet_forward( for layer in self.layers: inputs_embeds = torch.nn.functional.relu(layer(inputs_embeds)) - # NOTE: we patch the prenet to avoid using torch.nn.functional.dropout, that is exported as a `Dropout` node in the ONNX - # that is ignored during inference by some runtimes as ONNX Runtime. + # NOTE: we patch the prenet to avoid using torch.nn.functional.dropout, that is exported as a `Dropout` node in the OpenVINO + # that is ignored during inference by some runtimes as OpenVINO Runtime. # Reference: https://github.com/microsoft/onnxruntime/issues/9333 & https://github.com/microsoft/onnxruntime/issues/5549 mask = torch.rand(inputs_embeds.shape, device=inputs_embeds.device) > self.config.speech_decoder_prenet_dropout inputs_embeds = inputs_embeds * mask / (1 - self.config.speech_decoder_prenet_dropout) @@ -220,7 +220,7 @@ def patched_speecht5_prenet_forward( class SentenceTransformersTransformerPatcher(ModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: PreTrainedModel, model_kwargs: dict[str, Any], ): @@ -270,7 +270,7 @@ def _get_model_attribute(model, name): for idx, spec in enumerate(UNSUPPORTED_OPS_PATCHING_SPEC): if spec.name in { - # onnx-exporter-specific fixes + # openvino-exporter-specific fixes "triu", "tril", "norm", @@ -828,7 +828,7 @@ def _glm4_core_attention_forward(self, query_layer, key_layer, value_layer, atte class ChatGLMModelPatcher(OVDecoderModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Dict[str, Any], ): @@ -1212,7 +1212,7 @@ def _qwen_attention_forward( class QwenModelPatcher(OVDecoderModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Dict[str, Any], ): @@ -1373,7 +1373,7 @@ def apply_rotary_pos_emb(q, k, cos, sin, position_ids): class BaichuanModelPatcher(OVDecoderModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Dict[str, Any], ): @@ -3306,7 +3306,7 @@ def __exit__(self, exc_type, exc_value, traceback): class Gemma2ModelPatcher(OVDecoderModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Optional[Dict[str, Any]] = None, ): @@ -3495,7 +3495,7 @@ def __exit__(self, exc_type, exc_value, traceback): class IBertModelPatcher(ModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Dict[str, Any], ): @@ -3513,7 +3513,7 @@ def __init__( class InternVLChatImageEmbeddingModelPatcher(ModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Dict[str, Any], ): @@ -3536,7 +3536,7 @@ def __exit__(self, exc_type, exc_value, traceback): class InternVL2ChatLangModelPatcher(OVDecoderModelPatcher): - def __init__(self, config: "OnnxConfig", model: "PreTrainedModel", model_kwargs: Dict[str, Any]): + def __init__(self, config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Dict[str, Any]): model_type = model.config.model_type patcher_for_model_type = { "llama": OVDecoderModelPatcher, @@ -3656,7 +3656,7 @@ def maira_vision_embed_forward(self, pixel_values): class LlavaImageEmbeddingModelPatcher(ModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Dict[str, Any], ): @@ -3672,7 +3672,7 @@ def __exit__(self, exc_type, exc_value, traceback): class MairaImageEmbeddingModelPatcher(ModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Dict[str, Any], ): @@ -3689,7 +3689,7 @@ def __exit__(self, exc_type, exc_value, traceback): class LlavaNextVideoImageEmbeddingModelPatcher(ModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Dict[str, Any], ): @@ -3906,7 +3906,7 @@ def _minicpmv_siglip_transformer_forward( class MiniCPMVResamplerModelPatcher(ModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Dict[str, Any], ): @@ -3923,7 +3923,7 @@ def __exit__(self, exc_type, exc_value, traceback): class MiniCPMVImageEmbeddingsModelPatcher(ModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Dict[str, Any], ): @@ -3954,7 +3954,7 @@ def __exit__(self, exc_type, exc_value, traceback): class LlavaQwen2ImageEmbeddingsModelPatcher(ModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Dict[str, Any], ): @@ -3972,7 +3972,7 @@ def __exit__(self, exc_type, exc_value, traceback): class InputEmbeddingPatcher(ModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Dict[str, Any], ): @@ -3997,7 +3997,7 @@ def phi3_vision_embeddings_forward(self, pixel_values: torch.FloatTensor): class Phi3VisionImageEmbeddingsPatcher(ModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Dict[str, Any], ): @@ -4449,7 +4449,7 @@ def deepseek_moe_infer(self, x, topk_ids, topk_weight): class Qwen2VLLanguageModelPatcher(OVDecoderModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Dict[str, Any] = None, ): @@ -4492,7 +4492,7 @@ def __exit__(self, exc_type, exc_value, traceback): class Qwen3VLLanguageModelPatcher(OVDecoderModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: Union["PreTrainedModel"], model_kwargs: Optional[Dict[str, Any]] = None, ): @@ -4658,7 +4658,7 @@ def block_forward( class Qwen2VLVisionEmbMergerPatcher(ModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Dict[str, Any] = None, ): @@ -4692,7 +4692,7 @@ def __exit__(self, exc_type, exc_value, traceback): class Qwen2_5_VLVisionEmbMergerPatcher(ModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Dict[str, Any] = None, ): @@ -4754,7 +4754,7 @@ def __exit__(self, exc_type, exc_value, traceback): class Qwen3VLVisionEmbMergerPatcher(ModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: Union["PreTrainedModel"], model_kwargs: Dict[str, Any] = None, ): @@ -4867,7 +4867,7 @@ def __exit__(self, exc_type, exc_value, traceback): class OVSeq2SeqModelPatcher(ModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Optional[Dict[str, Any]] = None, ): @@ -4913,7 +4913,7 @@ def patched_forward(*args, **kwargs): outputs = self.super_patched_forward(*args, **kwargs) - # the optimum-onnx seq2seq model patcher only converts to tuple starting from 4.48 + # the seq2seq model patcher only converts to tuple starting from 4.48 if isinstance(outputs.get("past_key_values"), (DynamicCache, EncoderDecoderCache)): outputs["past_key_values"] = postprocess_past_key_values(outputs["past_key_values"]) @@ -4992,7 +4992,7 @@ def __exit__(self, exc_type, exc_value, traceback): class MiniCPMModelPatcher(OVDecoderModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Optional[Dict[str, Any]] = None, ): @@ -5007,7 +5007,7 @@ def __init__( class CommonImageEmbeddingsModelPatcher(ModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Dict[str, Any], ): @@ -5091,7 +5091,7 @@ def _gemma3_mm_update_causal_mask( class Gemma3LMModelPatcher(OVDecoderModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Optional[Dict[str, Any]] = None, ): @@ -5525,7 +5525,7 @@ def __exit__(self, exc_type, exc_value, traceback): class Idefics3ImageEmbeddingsModelPatcher(ModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Optional[Dict[str, Any]] = None, ): @@ -6341,7 +6341,7 @@ def __exit__(self, exc_type, exc_value, traceback): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Dict[str, Any], ): @@ -6456,7 +6456,7 @@ def patched_vocoder_forward(spectrogram: torch.FloatTensor): class Phi4MMLanguageModelPatcher(OVDecoderModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Optional[Dict[str, Any]] = None, ): @@ -6497,7 +6497,7 @@ def __exit__(self, exc_type, exc_value, traceback): class Phi4MMAudioForwardEmbeddingsPatcher(ModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Optional[Dict[str, Any]] = None, ): @@ -6521,7 +6521,7 @@ def __exit__(self, exc_type, exc_value, traceback): class Phi4MMAudioEncoderPatcher(ModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Optional[Dict[str, Any]] = None, ): @@ -6562,7 +6562,7 @@ def __exit__(self, exc_type, exc_value, traceback): class Phi4MMVisionEmbeddingsPatcher(ModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Optional[Dict[str, Any]] = None, ): @@ -6871,7 +6871,7 @@ def __exit__(self, exc_type, exc_value, traceback): class Llama4ImageEmbeddingsModelPatcher(ModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Dict[str, Any], ): @@ -7236,7 +7236,7 @@ def mamba_mixer_forward( class MambaPatcher(ModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Optional[Dict[str, Any]] = None, ): @@ -7743,7 +7743,7 @@ def segment_sum(input_tensor): class Zamba2ModelPatcher(ModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Optional[Dict[str, Any]] = None, ): @@ -7956,7 +7956,7 @@ def lfm2_short_conv_forward_patched( class Lfm2ModelPatcher(OVDecoderModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Optional[Dict[str, Any]] = None, ): @@ -8201,7 +8201,7 @@ def granite_moe_hybrid_update_causal_mask( class GraniteMoeHybridModelPatcher(OVDecoderModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Optional[Dict[str, Any]] = None, ): @@ -8383,7 +8383,7 @@ def __enter__(self): if self.real_config._behavior == "encoder" and self._model.config.attention_type == "block_sparse": logger.warning( - "BigBirdPegasus model is using block sparse attention, which is not supported in ONNX export. " + "BigBirdPegasus model is using block sparse attention, which is not supported in OpenVINO export. " "The model will be exported with original full attention." ) self._model.set_attention_type("original_full") @@ -8483,7 +8483,7 @@ def __exit__(self, exc_type, exc_value, traceback): class VideoChatFlashQwenVisionEmbeddingModelPatcher(ModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Dict[str, Any] = None, ): @@ -9033,7 +9033,7 @@ def forward( class Qwen3NextModelPatcher(OVDecoderModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Optional[Dict[str, Any]] = None, ): @@ -9487,7 +9487,7 @@ def apply_mask_to_padding_states(hidden_states, attention_mask): class Qwen3_5ModelPatcher(OVDecoderModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Optional[Dict[str, Any]] = None, ): @@ -9672,7 +9672,7 @@ def __exit__(self, exc_type, exc_value, traceback): class Qwen3_5VisionEmbMergerPatcher(ModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Dict[str, Any] = None, ): @@ -9746,7 +9746,7 @@ def patched_qwen3_5_moe_sparse_moe_block(self, hidden_states: torch.Tensor) -> t class Qwen3_5MoeModelPatcher(Qwen3_5ModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Optional[Dict[str, Any]] = None, ): @@ -9782,7 +9782,7 @@ class Qwen3ASRModelPatcher(OVSeq2SeqModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Optional[Dict[str, Any]] = None, ): diff --git a/optimum/exporters/openvino/patching_utils.py b/optimum/exporters/openvino/patching_utils.py index a7e43fc6b7..c9b7d094f8 100644 --- a/optimum/exporters/openvino/patching_utils.py +++ b/optimum/exporters/openvino/patching_utils.py @@ -649,11 +649,11 @@ def patched_forward(*args, **kwargs): output_names = list(config.outputs.keys()) if isinstance(outputs, dict): for name, value in outputs.items(): - onnx_output_name = config.torch_to_onnx_output_map.get(name, name) + ov_output_name = config.torch_to_ov_output_map.get(name, name) if ( - onnx_output_name in output_names + ov_output_name in output_names or (use_cache and name.startswith("past_key_values")) - or any(key.startswith(onnx_output_name) for key in output_names) + or any(key.startswith(ov_output_name) for key in output_names) ): filtered_outputs[name] = value elif isinstance(outputs, (list, tuple)): diff --git a/optimum/exporters/openvino/utils.py b/optimum/exporters/openvino/utils.py index 164c11f2f1..b6c04b7df4 100644 --- a/optimum/exporters/openvino/utils.py +++ b/optimum/exporters/openvino/utils.py @@ -24,7 +24,7 @@ from openvino import Dimension, PartialShape, Symbol from openvino.utils.types import get_element_type -from optimum.exporters.openvino.base import OnnxConfig +from optimum.exporters.openvino.base import OpenVINOConfig from optimum.exporters.tasks import TasksManager from optimum.intel.utils.import_utils import is_safetensors_available from optimum.utils import is_diffusers_available @@ -87,7 +87,7 @@ def flattenize_inputs(inputs: List[Any]): def _get_input_info( - model: Union["PreTrainedModel", "ModelMixin"], config: OnnxConfig, dummy_inputs: Dict[str, Any] + model: Union["PreTrainedModel", "ModelMixin"], config: OpenVINOConfig, dummy_inputs: Dict[str, Any] ) -> List[InputInfo]: sig = inspect.signature(model.forward) if hasattr(model, "forward") else inspect.signature(model.call) inputs = config.ordered_inputs(model) @@ -129,7 +129,7 @@ def _get_input_info( def _get_dynamic_shapes_info( - model: Union["PreTrainedModel", "ModelMixin"], config: OnnxConfig, dummy_inputs: Dict[str, Any] + model: Union["PreTrainedModel", "ModelMixin"], config: OpenVINOConfig, dummy_inputs: Dict[str, Any] ) -> List[InputInfo]: import torch @@ -248,7 +248,7 @@ def _get_open_clip_submodels_fn_and_export_configs( library_name: str = "open_clip", task: Optional[str] = None, preprocessors: List = None, - custom_export_configs: Dict[str, "OnnxConfig"] = None, + custom_export_configs: Dict[str, "OpenVINOConfig"] = None, fn_get_submodels: Callable = None, ): custom_export = {} @@ -295,7 +295,7 @@ def _get_kokoro_submodels_fn_and_export_configs( library_name: str = "kokoro", task: Optional[str] = None, preprocessors: List = None, - custom_export_configs: Dict[str, "OnnxConfig"] = None, + custom_export_configs: Dict[str, "OpenVINOConfig"] = None, fn_get_submodels: Callable = None, ): export_config_constructor = TasksManager.get_exporter_config_constructor( From 1f98e9e89d29b930cc7a6937f7ea9b7011e95979 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CElla?= Date: Thu, 28 May 2026 18:23:26 +0200 Subject: [PATCH 27/40] fix --- optimum/intel/openvino/modeling_timm.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/optimum/intel/openvino/modeling_timm.py b/optimum/intel/openvino/modeling_timm.py index ce60323aac..ad0d96231b 100644 --- a/optimum/intel/openvino/modeling_timm.py +++ b/optimum/intel/openvino/modeling_timm.py @@ -40,7 +40,7 @@ from transformers.modeling_outputs import ImageClassifierOutput from transformers.utils import TensorType -from optimum.exporters.openvino.config import VisionOnnxConfig +from optimum.exporters.openvino.config import VisionOpenVINOConfig from optimum.utils import NormalizedVisionConfig from .utils import _is_timm_ov_dir @@ -79,7 +79,7 @@ def from_pretrained( return cls.from_dict(config_dict, **kwargs) -class TimmOnnxConfig(VisionOnnxConfig): +class TimmOnnxConfig(VisionOpenVINOConfig): DEFAULT_TIMM_ONNX_OPSET = 13 outputs = OrderedDict([("logits", {0: "batch_size"})]) NORMALIZED_CONFIG_CLASS = NormalizedVisionConfig From b18eb3f944c5ad1d3d0fd0558c676dd0661fb0c7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CElla?= Date: Thu, 28 May 2026 18:25:26 +0200 Subject: [PATCH 28/40] style --- optimum/exporters/openvino/base.py | 4 +++- optimum/exporters/openvino/config.py | 15 ++++++++------- optimum/exporters/openvino/model_configs.py | 7 ++++++- 3 files changed, 17 insertions(+), 9 deletions(-) diff --git a/optimum/exporters/openvino/base.py b/optimum/exporters/openvino/base.py index fb4eff5453..b5f31ca4ac 100644 --- a/optimum/exporters/openvino/base.py +++ b/optimum/exporters/openvino/base.py @@ -157,7 +157,9 @@ def variant(self, value: str): if value == "default" and hasattr(self, "DEFAULT_VARIANT"): value = self.DEFAULT_VARIANT if value not in self.VARIANTS: - raise ValueError(f"The variant {value} is not supported for the OpenVINO config {self.__class__.__name__}.") + raise ValueError( + f"The variant {value} is not supported for the OpenVINO config {self.__class__.__name__}." + ) self._variant = value @property diff --git a/optimum/exporters/openvino/config.py b/optimum/exporters/openvino/config.py index 1910659e27..f248795b72 100644 --- a/optimum/exporters/openvino/config.py +++ b/optimum/exporters/openvino/config.py @@ -21,7 +21,12 @@ from transformers import PretrainedConfig -from optimum.exporters.openvino.base import ConfigBehavior, OpenVINOConfig, OpenVINOConfigWithPast, OpenVINOSeq2SeqConfigWithPast +from optimum.exporters.openvino.base import ( + ConfigBehavior, + OpenVINOConfig, + OpenVINOConfigWithPast, + OpenVINOSeq2SeqConfigWithPast, +) from optimum.exporters.tasks import TasksManager from optimum.utils import ( DummyAudioInputGenerator, @@ -325,9 +330,7 @@ def generate_dummy_inputs_for_validation( self, reference_model_inputs: dict[str, Any], ov_input_names: list[str] ) -> dict[str, Any]: if self._behavior is ConfigBehavior.ENCODER: - return self._encoder_ov_config.generate_dummy_inputs_for_validation( - reference_model_inputs, ov_input_names - ) + return self._encoder_ov_config.generate_dummy_inputs_for_validation(reference_model_inputs, ov_input_names) else: if self._behavior is ConfigBehavior.DECODER: if "decoder_input_ids" in reference_model_inputs: @@ -342,6 +345,4 @@ def generate_dummy_inputs_for_validation( else: reference_model_inputs.pop("encoder_outputs") - return self._decoder_ov_config.generate_dummy_inputs_for_validation( - reference_model_inputs, ov_input_names - ) + return self._decoder_ov_config.generate_dummy_inputs_for_validation(reference_model_inputs, ov_input_names) diff --git a/optimum/exporters/openvino/model_configs.py b/optimum/exporters/openvino/model_configs.py index 7ff180d4da..e7aec66630 100644 --- a/optimum/exporters/openvino/model_configs.py +++ b/optimum/exporters/openvino/model_configs.py @@ -18,7 +18,12 @@ from transformers import AutoConfig, PretrainedConfig, PreTrainedModel -from optimum.exporters.openvino.base import ConfigBehavior, OpenVINOConfig, OpenVINOConfigWithPast, OpenVINOSeq2SeqConfigWithPast +from optimum.exporters.openvino.base import ( + ConfigBehavior, + OpenVINOConfig, + OpenVINOConfigWithPast, + OpenVINOSeq2SeqConfigWithPast, +) from optimum.exporters.openvino.config import ( AudioOpenVINOConfig, AudioToTextOpenVINOConfig, From 6404ebea86d986f17f37a3230a5d2812929ec94d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CElla?= Date: Fri, 29 May 2026 10:14:39 +0200 Subject: [PATCH 29/40] fix timm export --- optimum/intel/openvino/modeling.py | 4 ++-- optimum/intel/openvino/modeling_base.py | 1 - optimum/intel/openvino/modeling_timm.py | 2 +- 3 files changed, 3 insertions(+), 4 deletions(-) diff --git a/optimum/intel/openvino/modeling.py b/optimum/intel/openvino/modeling.py index 9fc568e042..7abaf0bccf 100644 --- a/optimum/intel/openvino/modeling.py +++ b/optimum/intel/openvino/modeling.py @@ -571,14 +571,14 @@ def from_pretrained( "To load a timm model, please make sure to upgrade your `timm` version to at least 0.9.0, you can upgrade it by running `pip install --upgrade timm`" ) - from .modeling_timm import TimmConfig, TimmForImageClassification, TimmOnnxConfig + from .modeling_timm import TimmConfig, TimmForImageClassification, TimmOpenVINOConfig config = TimmConfig.from_pretrained(model_id, **kwargs) # If locally saved timm model, directly load if local_timm_model: return super()._from_pretrained(model_id=model_id, config=config) model = TimmForImageClassification.from_pretrained(model_id, **kwargs) - onnx_config = TimmOnnxConfig(model.config) + onnx_config = TimmOpenVINOConfig(model.config) return cls._to_load(model=model, config=config, onnx_config=onnx_config, stateful=False, **kwargs) else: diff --git a/optimum/intel/openvino/modeling_base.py b/optimum/intel/openvino/modeling_base.py index fd4231418e..a735e1c1ad 100644 --- a/optimum/intel/openvino/modeling_base.py +++ b/optimum/intel/openvino/modeling_base.py @@ -928,7 +928,6 @@ def _to_load( export( model=model, config=onnx_config, - opset=onnx_config.DEFAULT_ONNX_OPSET, output=save_dir_path / cls._all_ov_model_paths["model"], stateful=stateful, ) diff --git a/optimum/intel/openvino/modeling_timm.py b/optimum/intel/openvino/modeling_timm.py index ad0d96231b..d80e9b5c84 100644 --- a/optimum/intel/openvino/modeling_timm.py +++ b/optimum/intel/openvino/modeling_timm.py @@ -79,7 +79,7 @@ def from_pretrained( return cls.from_dict(config_dict, **kwargs) -class TimmOnnxConfig(VisionOpenVINOConfig): +class TimmOpenVINOConfig(VisionOpenVINOConfig): DEFAULT_TIMM_ONNX_OPSET = 13 outputs = OrderedDict([("logits", {0: "batch_size"})]) NORMALIZED_CONFIG_CLASS = NormalizedVisionConfig From 89ccfbf1a9d591beb26c96294c0ee39848ff38a4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CElla?= Date: Fri, 29 May 2026 15:45:17 +0200 Subject: [PATCH 30/40] fix --- tests/openvino/test_quantization.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/openvino/test_quantization.py b/tests/openvino/test_quantization.py index 979e1c246d..d286718c26 100644 --- a/tests/openvino/test_quantization.py +++ b/tests/openvino/test_quantization.py @@ -2076,7 +2076,7 @@ def preprocess_function(examples, tokenizer): tokenizer = AutoTokenizer.from_pretrained(model_name) quantizer = OVQuantizer.from_pretrained(transformers_model, device=OPENVINO_DEVICE) calibration_dataset = quantizer.get_calibration_dataset( - "squadshifts", + "ludwigschmidt/squadshifts", dataset_config_name="new_wiki", preprocess_function=partial(preprocess_function, tokenizer=tokenizer), num_samples=10, From c423c91ef3440df7439914558eee14397049bf1e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CElla?= Date: Fri, 29 May 2026 16:31:13 +0200 Subject: [PATCH 31/40] remove unecessary eager_mask_without_vmap --- optimum/exporters/openvino/model_patcher.py | 18 +----------------- optimum/exporters/openvino/patching_utils.py | 11 +++++++++-- 2 files changed, 10 insertions(+), 19 deletions(-) diff --git a/optimum/exporters/openvino/model_patcher.py b/optimum/exporters/openvino/model_patcher.py index 92ac740992..66e6f4186d 100644 --- a/optimum/exporters/openvino/model_patcher.py +++ b/optimum/exporters/openvino/model_patcher.py @@ -52,10 +52,10 @@ from optimum.exporters.openvino.patching_utils import ( UNSUPPORTED_OPS_PATCHING_SPEC, ModelPatcher, + eager_mask_without_vmap, override_arguments, postprocess_past_key_values, preprocess_past_key_values, - sdpa_mask_without_vmap, ) from optimum.intel.utils.import_utils import ( is_diffusers_version, @@ -507,22 +507,6 @@ def patch_cos_sin_cached_fp32(model): ) -# Adapted from https://github.com/huggingface/transformers/blob/v4.53.0/src/transformers/masking_utils.py#L433 -# Specifically for OpenVINO, we use torch.finfo(torch.float16).min instead of torch.finfo(dtype).min -def eager_mask_without_vmap(*args, **kwargs) -> Optional[torch.Tensor]: - kwargs.pop("allow_is_causal_skip", None) - dtype = kwargs.get("dtype", torch.float32) - mask = sdpa_mask_without_vmap(*args, allow_is_causal_skip=False, **kwargs) - # we use torch.finfo(torch.float16).min instead torch.finfo(dtype).min to avoid an overflow but not - # sure this is the right way to handle this, we are basically pretending that -65,504 is -inf - mask = torch.where( - mask, - torch.tensor(0.0, device=mask.device, dtype=dtype), - torch.tensor(torch.finfo(torch.float16).min, device=mask.device, dtype=dtype), - ) - return mask - - # Adapted from https://github.com/huggingface/transformers/blob/3c307e380ad07ca16903a39e09a47d532cb782d9/src/transformers/models/phimoe/modular_phimoe.py#L57 def _longrope_forward(self, x, position_ids=None, layer_type=None): # _compute_longrope_parameters https://github.com/huggingface/transformers/blob/v5.0.0/src/transformers/modeling_rope_utils.py#L391 diff --git a/optimum/exporters/openvino/patching_utils.py b/optimum/exporters/openvino/patching_utils.py index c9b7d094f8..e1477e0d9c 100644 --- a/optimum/exporters/openvino/patching_utils.py +++ b/optimum/exporters/openvino/patching_utils.py @@ -459,11 +459,18 @@ def sdpa_mask_without_vmap(batch_size, q_length=None, kv_length=None, q_offset=0 # Adapted from https://github.com/huggingface/transformers/blob/v4.53.0/src/transformers/masking_utils.py#L433 -def eager_mask_without_vmap(*args, **kwargs) -> torch.Tensor: +# Specifically for OpenVINO, we use torch.finfo(torch.float16).min instead of torch.finfo(dtype).min +def eager_mask_without_vmap(*args, **kwargs) -> Optional[torch.Tensor]: kwargs.pop("allow_is_causal_skip", None) dtype = kwargs.get("dtype", torch.float32) mask = sdpa_mask_without_vmap(*args, allow_is_causal_skip=False, **kwargs) - mask = torch.where(mask, torch.tensor(0.0, device=mask.device, dtype=dtype), torch.finfo(dtype).min) + # we use torch.finfo(torch.float16).min instead torch.finfo(dtype).min to avoid an overflow but not + # sure this is the right way to handle this, we are basically pretending that -65,504 is -inf + mask = torch.where( + mask, + torch.tensor(0.0, device=mask.device, dtype=dtype), + torch.tensor(torch.finfo(torch.float16).min, device=mask.device, dtype=dtype), + ) return mask From 4907bb4a921830a60412ce5de6a30ab6f3615dd1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CElla?= Date: Fri, 29 May 2026 17:19:56 +0200 Subject: [PATCH 32/40] remove patching --- optimum/exporters/openvino/patching_utils.py | 111 +------------------ 1 file changed, 2 insertions(+), 109 deletions(-) diff --git a/optimum/exporters/openvino/patching_utils.py b/optimum/exporters/openvino/patching_utils.py index e1477e0d9c..dec47776e7 100644 --- a/optimum/exporters/openvino/patching_utils.py +++ b/optimum/exporters/openvino/patching_utils.py @@ -16,19 +16,17 @@ import functools import inspect import logging -import sys import types from typing import Any, Callable import torch import transformers -from torch.onnx import symbolic_helper from transformers import PreTrainedModel from transformers.cache_utils import DynamicCache, EncoderDecoderCache from transformers.modeling_outputs import BaseModelOutput from optimum.exporters.base import ExporterConfig -from optimum.intel.utils.import_utils import is_torch_version, is_transformers_version +from optimum.intel.utils.import_utils import is_transformers_version if is_transformers_version(">=", "4.44") and is_transformers_version("<", "4.50"): @@ -70,111 +68,6 @@ logger = logging.getLogger(__name__) -@symbolic_helper.parse_args("v", "v") -def __ior_(g, self: torch._C.Value, other: torch._C.Value) -> torch._C.Value: - return g.op("Or", self, other) - - -torch.onnx.register_custom_op_symbolic("aten::__ior__", __ior_, 14) - -if is_torch_version("<", "2.9"): - # this was fixed in torch in 2.9 https://github.com/pytorch/pytorch/pull/159973 - from torch.onnx import JitScalarType - from torch.onnx.symbolic_opset14 import _attention_scale, _causal_attention_mask - - @symbolic_helper.parse_args("v", "v", "v", "v", "f", "b", "v", "b") - def scaled_dot_product_attention( - g, - query: torch._C.Value, - key: torch._C.Value, - value: torch._C.Value, - attn_mask: torch._C.Value | None = None, - dropout_p: float = 0.0, - is_causal: bool = False, - scale: torch._C.Value | None = None, - enable_gqa: bool = False, - ): - assert (not is_causal) or ( - is_causal and symbolic_helper._is_none(attn_mask) - ), "is_causal and attn_mask cannot be set at the same time" - assert not enable_gqa, "conversion of scaled_dot_product_attention not implemented if enable_gqa is True" - - if symbolic_helper._is_none(scale): - scale = _attention_scale(g, query) - - if is_causal: - attn_mask = _causal_attention_mask(g, query, key) - - # Swap the last two axes of key - # NOTE: onnx-script has different logic here, because the attribute perms in - # transpose needs list of ints - key_shape_builtin = symbolic_helper._get_tensor_rank(key) - key_transposed_axes = list(range(key_shape_builtin)) - key_transposed_axes[-1], key_transposed_axes[-2] = (key_transposed_axes[-2], key_transposed_axes[-1]) - key_transposed = g.op("Transpose", key, perm_i=key_transposed_axes) - - # https://github.com/pytorch/pytorch/blob/12da0c70378b5be9135c6fda62a9863bce4a4818/aten/src/ATen/native/transformers/attention.cpp#L653 - # Scale q, k before matmul for stability see https://tinyurl.com/sudb9s96 for math - query_scaled = g.op("Mul", query, g.op("Sqrt", scale)) - key_transposed_scaled = g.op("Mul", key_transposed, g.op("Sqrt", scale)) - mul_qk = g.op("MatMul", query_scaled, key_transposed_scaled) - - if symbolic_helper._is_none(attn_mask): - mul_qk_add = mul_qk - attn_weight = g.op("Softmax", mul_qk_add, axis_i=-1) - elif JitScalarType.from_value(attn_mask) == JitScalarType.BOOL: - # Turn the Boolean mask to float: attn_mask.masked_fill(not attn_mask, -float('inf')) - const_zero = g.op("Constant", value_t=torch.tensor([0.0])) - const_neg_inf = g.op("Constant", value_t=torch.tensor([-float("inf")])) - attn_mask = g.op("Where", attn_mask, const_zero, const_neg_inf) - mul_qk_add = g.op("Add", mul_qk, attn_mask) - attn_weight = g.op("Softmax", mul_qk_add, axis_i=-1) - # when using scaled dot product attention with a boolean mask, we replace NaN values in attn_weight with 0.0 - attn_weight = g.op( - "Where", g.op("IsNaN", attn_weight), g.op("Constant", value_t=torch.tensor([0.0])), attn_weight - ) - elif JitScalarType.from_value(attn_mask) in ( - JitScalarType.FLOAT, - JitScalarType.HALF, - JitScalarType.BFLOAT16, - ): - mul_qk_add = g.op("Add", mul_qk, attn_mask) - attn_weight = g.op("Softmax", mul_qk_add, axis_i=-1) - else: - raise ValueError(f"Unsupported type for attn_mask: {JitScalarType.from_value(attn_mask)}") - - if dropout_p != 0: - attn_weight = g.op( - "Dropout", - attn_weight, - g.op("Constant", value_t=torch.tensor(dropout_p, dtype=torch.float)), - ) - - return g.op("MatMul", attn_weight, value) - - torch.onnx.register_custom_op_symbolic("aten::scaled_dot_product_attention", scaled_dot_product_attention, 14) - - -def patch_everywhere(attribute_name: str, patch: Any, module_name_prefix: str | None = None): - """Finds all occurrences of `attribute_name` in the loaded modules and patches them with `patch`. - - Args: - attribute_name (`str`): - The name of attribute to patch. - patch (`Any`): - The patch for the attribute. - module_name_prefix (`Optional[str]`, defaults to `None`): - If set, only module names starting with this prefix will be considered for patching. - """ - # sys.modules may be updated while being iterated over, hence the list copy. - for name in list(sys.modules): - module = sys.modules[name] - if module_name_prefix is not None and not name.startswith(module_name_prefix): - continue - if hasattr(module, attribute_name): - setattr(module, attribute_name, patch) - - def override_arguments(args, kwargs, forward_signature, model_kwargs: dict[str, Any]): """Override the args and kwargs with the argument values from model_kwargs, following the signature forward_signature corresponding to args and kwargs.""" args = list(args) @@ -460,7 +353,7 @@ def sdpa_mask_without_vmap(batch_size, q_length=None, kv_length=None, q_offset=0 # Adapted from https://github.com/huggingface/transformers/blob/v4.53.0/src/transformers/masking_utils.py#L433 # Specifically for OpenVINO, we use torch.finfo(torch.float16).min instead of torch.finfo(dtype).min -def eager_mask_without_vmap(*args, **kwargs) -> Optional[torch.Tensor]: +def eager_mask_without_vmap(*args, **kwargs): kwargs.pop("allow_is_causal_skip", None) dtype = kwargs.get("dtype", torch.float32) mask = sdpa_mask_without_vmap(*args, allow_is_causal_skip=False, **kwargs) From 7e5185b55819329f14c3342a73cebf9cbf58022f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CElla?= Date: Fri, 29 May 2026 17:28:26 +0200 Subject: [PATCH 33/40] remove noop_bfloat16_casting --- optimum/exporters/openvino/patching_utils.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/optimum/exporters/openvino/patching_utils.py b/optimum/exporters/openvino/patching_utils.py index dec47776e7..e09c28d108 100644 --- a/optimum/exporters/openvino/patching_utils.py +++ b/optimum/exporters/openvino/patching_utils.py @@ -414,11 +414,6 @@ def traceable_scaled_dot_product_attention( return attn_weights -# No-op bfloat16 casting to avoid issues with legacy ONNX export which cast to complex128 -def noop_bfloat16_casting(self): - return self - - original_movedim = torch.Tensor.movedim @@ -452,7 +447,6 @@ def patched_dynamic_layer_update( PatchingSpec(torch, "rms_norm", onnx_compatible_rms_norm, torch.rms_norm), PatchingSpec(torch.Tensor, "unfold", onnx_compatible_unfold, torch.Tensor.unfold), PatchingSpec(torch.linalg, "norm", onnx_compatible_linalg_norm, torch.linalg.norm), - PatchingSpec(torch.Tensor, "bfloat16", noop_bfloat16_casting, torch.Tensor.bfloat16), PatchingSpec(torch.Tensor, "movedim", onnx_compatible_movedim, torch.Tensor.movedim), PatchingSpec(torch.Tensor, "repeat_interleave", onnx_compatible_repeat_interleave, torch.Tensor.repeat_interleave), # TracerWarning: Using len to get tensor shape might cause the trace to be incorrect. Recommended usage would be tensor.shape[0]. Passing a tensor of different shape might lead to errors or silently give incorrect results. From 1770c90bcd465ab6b0574e78f670f1b73f367aa1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CElla?= Date: Tue, 2 Jun 2026 11:20:03 +0200 Subject: [PATCH 34/40] clean PatchingSpec --- optimum/exporters/openvino/model_patcher.py | 16 -- optimum/exporters/openvino/patching_utils.py | 189 ------------------- 2 files changed, 205 deletions(-) diff --git a/optimum/exporters/openvino/model_patcher.py b/optimum/exporters/openvino/model_patcher.py index 66e6f4186d..b84e8c89ed 100644 --- a/optimum/exporters/openvino/model_patcher.py +++ b/optimum/exporters/openvino/model_patcher.py @@ -50,7 +50,6 @@ from optimum.exporters.openvino._ov_ops import convert_recurrent_attention_cell from optimum.exporters.openvino.base import OpenVINOConfig from optimum.exporters.openvino.patching_utils import ( - UNSUPPORTED_OPS_PATCHING_SPEC, ModelPatcher, eager_mask_without_vmap, override_arguments, @@ -268,21 +267,6 @@ def _get_model_attribute(model, name): return getattr(target, name) -for idx, spec in enumerate(UNSUPPORTED_OPS_PATCHING_SPEC): - if spec.name in { - # openvino-exporter-specific fixes - "triu", - "tril", - "norm", - "unfold", - "movedim", - "rms_norm", - "repeat_interleave", - "scaled_dot_product_attention", - }: - UNSUPPORTED_OPS_PATCHING_SPEC.pop(idx) - - # Original code: https://github.com/huggingface/transformers/blob/v5.0.0/src/transformers/integrations/moe.py#L98 # Method needs to be patched to match the pattern that OpenVINO MoE optimization transformation expect. # The difference is that this method packs hidden states and routing weight to dense representation, diff --git a/optimum/exporters/openvino/patching_utils.py b/optimum/exporters/openvino/patching_utils.py index e09c28d108..d8ce0a2c05 100644 --- a/optimum/exporters/openvino/patching_utils.py +++ b/optimum/exporters/openvino/patching_utils.py @@ -156,123 +156,6 @@ class PatchingSpec: op_wrapper: Callable | None = None -# An ONNX-export-compatible version of `tensor.unfold`. Without this, we get: -# torch.onnx.errors.SymbolicValueError: Unsupported: ONNX export of operator Unfold, input size not accessible. -# See https://github.com/pytorch/pytorch/issues/81871 for more information -def onnx_compatible_unfold(input_tensor, dimension, size, step): - """Custom implementation of torch.unfold without using torch.unfold. - - Args: - input_tensor (torch.Tensor): The input tensor. - dimension (int): The dimension to unfold. - size (int): The size of each slice. - step (int): The step size between slices. - - Returns: - torch.Tensor: The unfolded tensor. - """ - # Check if dimension is within the valid range - if not (-input_tensor.dim() <= dimension < input_tensor.dim()): - raise ValueError( - f"Dimension out of range (expected to be in range of [{-input_tensor.dim()}, {input_tensor.dim() - 1}], but got {dimension})" - ) - - # Normalize negative dimension - dimension = dimension % input_tensor.dim() - - # Compute the shape of the unfolded output - input_size = input_tensor.size(dimension) - num_slices = (input_size - size) // step + 1 - - # Permute dimension to the end for easier indexing - input_tensor = input_tensor.transpose(dimension, -1) - - # Extract slices - slices = [] - for i in range(num_slices): - start = i * step - end = start + size - slices.append(input_tensor[..., start:end]) - - # Stack slices and permute dimensions back - result = torch.stack(slices, dim=-2).transpose(dimension, -2) - return result - - -# An ONNX-export-compatible version of `tensor.repeat_interleave`. -# Without this, we get the following error: https://github.com/pytorch/pytorch/issues/145100 -# NOTE: This implementation is only necessary for export with dynamo=False (dynamo=True works correctly). -# and can be removed once Optimum switches to dynamo-based exports -def onnx_compatible_repeat_interleave(input_tensor, repeats, dim=None, output_size=None): # noqa: D417 - """Custom implementation of torch.repeat_interleave without using torch.repeat_interleave. - - Args: - input_tensor (torch.Tensor): The input tensor. - repeats (int or torch.Tensor): The number of repetitions for each element. - dim (int, optional): The dimension along which to repeat. Defaults to None. - - Returns: - torch.Tensor: The repeated tensor. - """ - if isinstance(repeats, int) or (torch.is_tensor(repeats) and repeats.dim() == 0): - if dim is None: - return input_tensor.flatten().unsqueeze(1).expand(-1, repeats).flatten() - repeats = torch.full((input_tensor.shape[dim],), repeats, dtype=torch.long, device=input_tensor.device) - - if dim is None: - return onnx_compatible_repeat_interleave(input_tensor.flatten(), repeats, 0) - - if dim != 0: - input_tensor = input_tensor.transpose(0, dim) - - # Create expand mask - max_repeats = repeats.max() - expanded = input_tensor.unsqueeze(1).expand(-1, max_repeats, *input_tensor.shape[1:]) - mask = torch.arange(max_repeats, device=input_tensor.device) < repeats.unsqueeze(1) - result = expanded[mask] - - if dim != 0: - result = result.transpose(0, dim) - - return result - - -# Custom implementation of torch.linalg.matrix_norm not using torch.linalg.matrix_norm, torch.norm or torch.linalg.norm. -def onnx_compatible_linalg_norm(x, ord=2, dim=None, keepdim=False, *, dtype=None, out=None) -> torch.Tensor: - if ord != 2: - raise ValueError( - f"Only ord=2 is supported by onnx_compatible_linalg_norm, but got ord={ord}. " - "Please extend this function to support other norms." - ) - - if dim is None: - dim = (-2, -1) - - norm = torch.sqrt(torch.sum(torch.square(x), dim=dim, keepdim=keepdim)) - - if dtype is not None: - norm = norm.to(dtype) - if out is not None: - out.copy_(norm) - - return norm - - -def onnx_compatible_rms_norm(input, normalized_shape, weight=None, eps=None): - if eps is None: - eps = torch.finfo(input.dtype).eps - - axis = -len(normalized_shape) - mean_square = torch.mean(torch.square(input), dim=axis, keepdim=True) - rms = torch.sqrt(mean_square + eps) - output = input / rms - - if weight is not None: - output = output * weight - - return output - - # A patched version of https://github.com/huggingface/transformers/blob/v4.53.2/src/transformers/masking_utils.py#L602 # That returns a tensor of zeros with the same shape as position_ids indicating no packed sequence indices. def find_packed_sequence_indices_patched(position_ids: torch.Tensor) -> torch.Tensor: @@ -367,65 +250,6 @@ def eager_mask_without_vmap(*args, **kwargs): return mask -original_triu = torch.triu -original_tril = torch.tril - - -# Custom implementation of torch.tril that doesn't fail on int32 tensors. -def onnx_compatible_tril(input_tensor: torch.Tensor, *args, **kwargs) -> torch.Tensor: - if input_tensor.dtype == torch.int32: - return original_tril(input_tensor.to(torch.int64), *args, **kwargs).to(torch.int32) - else: - return original_tril(input_tensor, *args, **kwargs) - - -# Custom implementation of torch.triu that doesn't fail on int32 tensors. -def onnx_compatible_triu(input_tensor: torch.Tensor, *args, **kwargs) -> torch.Tensor: - if input_tensor.dtype == torch.int32: - return original_triu(input_tensor.to(torch.int64), *args, **kwargs).to(torch.int32) - else: - return original_triu(input_tensor, *args, **kwargs) - - -original_scaled_dot_product_attention = torch.nn.functional.scaled_dot_product_attention - - -# A patched `torch.nn.functional.scaled_dot_product_attention` that doesn't fail during tracing -# from passing `is_causal` as a tensor (which is usually obtained with tensor shapes comparisons). -def traceable_scaled_dot_product_attention( - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - attn_mask: torch.Tensor | None = None, - dropout_p: float = 0.0, - is_causal: bool = False, - **kwargs, -) -> torch.Tensor: - if isinstance(is_causal, torch.Tensor): - is_causal = is_causal.item() - - if "enable_gqa" in kwargs: - kwargs.pop("enable_gqa") - - attn_weights = original_scaled_dot_product_attention( - query=query, key=key, value=value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, **kwargs - ) - - return attn_weights - - -original_movedim = torch.Tensor.movedim - - -def onnx_compatible_movedim(self: torch.Tensor, dim1, dim2) -> torch.Tensor: - dim = self.dim() - if dim1 < 0: - dim1 += dim - if dim2 < 0: - dim2 += dim - return original_movedim(self, dim1, dim2) - - def patched_dynamic_layer_update( self, key_states: torch.Tensor, value_states: torch.Tensor, cache_kwargs: dict[str, Any] | None = None ) -> tuple[torch.Tensor, torch.Tensor]: @@ -442,21 +266,8 @@ def patched_dynamic_layer_update( UNSUPPORTED_OPS_PATCHING_SPEC = [ - PatchingSpec(torch, "tril", onnx_compatible_tril, torch.tril), - PatchingSpec(torch, "triu", onnx_compatible_triu, torch.triu), - PatchingSpec(torch, "rms_norm", onnx_compatible_rms_norm, torch.rms_norm), - PatchingSpec(torch.Tensor, "unfold", onnx_compatible_unfold, torch.Tensor.unfold), - PatchingSpec(torch.linalg, "norm", onnx_compatible_linalg_norm, torch.linalg.norm), - PatchingSpec(torch.Tensor, "movedim", onnx_compatible_movedim, torch.Tensor.movedim), - PatchingSpec(torch.Tensor, "repeat_interleave", onnx_compatible_repeat_interleave, torch.Tensor.repeat_interleave), # TracerWarning: Using len to get tensor shape might cause the trace to be incorrect. Recommended usage would be tensor.shape[0]. Passing a tensor of different shape might lead to errors or silently give incorrect results. PatchingSpec(torch.Tensor, "__len__", lambda x: x.shape[0], torch.Tensor.__len__), - PatchingSpec( - torch.nn.functional, - "scaled_dot_product_attention", - traceable_scaled_dot_product_attention, - torch.nn.functional.scaled_dot_product_attention, - ), ] From db9d6d2a08ffc8b11550114060f7f718c82172b8 Mon Sep 17 00:00:00 2001 From: Ella Charlaix <80481427+echarlaix@users.noreply.github.com> Date: Tue, 2 Jun 2026 12:18:46 +0200 Subject: [PATCH 35/40] Update optimum/exporters/openvino/base.py Co-authored-by: Roman Kazantsev --- optimum/exporters/openvino/base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/optimum/exporters/openvino/base.py b/optimum/exporters/openvino/base.py index b5f31ca4ac..601a26fcca 100644 --- a/optimum/exporters/openvino/base.py +++ b/optimum/exporters/openvino/base.py @@ -455,7 +455,7 @@ def generate_dummy_inputs_for_validation( class ConfigBehavior(str, enum.Enum): - """Specifies the behavior of the [`~exporters.onnx.base.OpenVINOSeq2SeqConfigWithPast`]. + """Specifies the behavior of the [`~exporters.openvino.base.OpenVINOSeq2SeqConfigWithPast`]. - MONOLITH: the config can be used to export the whole seq2seq model as a single file. - ENCODER: the config can be used to export the encoder part of the seq2seq model. From 003695ca3302a53c1cd0ab8afa51d8ea8ed86eeb Mon Sep 17 00:00:00 2001 From: Ella Charlaix <80481427+echarlaix@users.noreply.github.com> Date: Tue, 2 Jun 2026 12:19:03 +0200 Subject: [PATCH 36/40] Update optimum/exporters/openvino/base.py Co-authored-by: Roman Kazantsev --- optimum/exporters/openvino/base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/optimum/exporters/openvino/base.py b/optimum/exporters/openvino/base.py index 601a26fcca..8c09341f08 100644 --- a/optimum/exporters/openvino/base.py +++ b/optimum/exporters/openvino/base.py @@ -468,7 +468,7 @@ class ConfigBehavior(str, enum.Enum): class OpenVINOSeq2SeqConfigWithPast(OpenVINOConfigWithPast): - """Inherits from [`~exporters.onnx.OpenVINOConfigWithPast`]. A base class to handle the OpenVINO configuration of encoder-decoder models.""" + """Inherits from [`~exporters.openvino.OpenVINOConfigWithPast`]. A base class to handle the OpenVINO configuration of encoder-decoder models.""" DUMMY_PKV_GENERATOR_CLASS = DummySeq2SeqPastKeyValuesGenerator From 0568cb1e17b1119d3ddaf11945042c74a414d7ca Mon Sep 17 00:00:00 2001 From: Ella Charlaix <80481427+echarlaix@users.noreply.github.com> Date: Tue, 2 Jun 2026 12:19:14 +0200 Subject: [PATCH 37/40] Update optimum/exporters/openvino/convert.py Co-authored-by: Roman Kazantsev --- optimum/exporters/openvino/convert.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/optimum/exporters/openvino/convert.py b/optimum/exporters/openvino/convert.py index 946824401b..12662598f3 100644 --- a/optimum/exporters/openvino/convert.py +++ b/optimum/exporters/openvino/convert.py @@ -597,7 +597,7 @@ def export_from_model( if model_type in ONNX_SUPPORTED_ARCHITECTURES: raise ValueError( - f"The OpenVINO export of {model_type} models is not officially supported by optimum-intel and has been deprecated." + f"The OpenVINO export of {model_type} models is not officially supported by optimum-intel and has been removed." ) custom_architecture = library_name == "transformers" and model_type not in TasksManager._SUPPORTED_MODEL_TYPE From 3f47ab12f0ee0787332364f5dcfed5f994011fe5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CElla?= Date: Tue, 2 Jun 2026 12:46:14 +0200 Subject: [PATCH 38/40] remove onnx reference --- optimum/exporters/openvino/convert.py | 4 ++-- optimum/exporters/openvino/model_configs.py | 4 ++-- optimum/exporters/openvino/utils.py | 3 +-- 3 files changed, 5 insertions(+), 6 deletions(-) diff --git a/optimum/exporters/openvino/convert.py b/optimum/exporters/openvino/convert.py index 12662598f3..dbe8e403cd 100644 --- a/optimum/exporters/openvino/convert.py +++ b/optimum/exporters/openvino/convert.py @@ -163,7 +163,7 @@ def export( Args: model ([`PreTrainedModel`]): The model to export. - config ([`~exporters.onnx.config.OpenVINOConfig`]): + config ([`~exporters.openvino.config.OpenVINOConfig`]): The OpenVINO configuration associated with the exported model. output (`Path`): Directory to store the exported model. @@ -250,7 +250,7 @@ def export_pytorch( Args: model ([`PreTrainedModel`]): The model to export. - config ([`~exporters.onnx.config.OpenVINOConfig`]): + config ([`~exporters.openvino.config.OpenVINOConfig`]): The configuration associated with the exported model. output (`Path`): Directory to store the exported model. diff --git a/optimum/exporters/openvino/model_configs.py b/optimum/exporters/openvino/model_configs.py index e7aec66630..32ee9ed288 100644 --- a/optimum/exporters/openvino/model_configs.py +++ b/optimum/exporters/openvino/model_configs.py @@ -5210,8 +5210,8 @@ class SamOpenVINOConfig(OpenVINOConfig): NORMALIZED_CONFIG_CLASS = NormalizedEncoderDecoderConfig DUMMY_INPUT_GENERATOR_CLASSES = (DummyVisionInputGenerator, DummyPointsGenerator, DummyVisionEmbeddingsGenerator) VARIANTS = { # noqa: RUF012 - "monolith": "All the SAM model components are exported as a single model.onnx.", - "split": "The vision encoder is exported as a separate vision_encoder.onnx, and the prompt encoder and mask decoder are exported as a prompt_encoder_mask_decoder.onnx. This allows to encoder the image only once for multiple point queries.", + "monolith": "All the SAM model components are exported as a single model.", + "split": "The vision encoder is exported as a separate vision_encoder, and the prompt encoder and mask decoder are exported as a prompt_encoder_mask_decoder. This allows to encoder the image only once for multiple point queries.", } DEFAULT_VARIANT = "split" _MODEL_PATCHER = SAMModelPatcher diff --git a/optimum/exporters/openvino/utils.py b/optimum/exporters/openvino/utils.py index b6c04b7df4..77ec6c7cd7 100644 --- a/optimum/exporters/openvino/utils.py +++ b/optimum/exporters/openvino/utils.py @@ -350,8 +350,7 @@ def _get_kokoro_submodels(model): "qwen3_5_moe_text", ] -# All transformers, diffusers, timm and sentence transformers models that are supported via optimum-onnx OnnxConfigs but that have currently no test -# TODO: add tests for all models that are compatible and remove support for all others +# All transformers, diffusers, timm and sentence transformers models that were supported via optimum-onnx OnnxConfigs for which support is now removed ONNX_SUPPORTED_ARCHITECTURES = { "big_bird", "chinese_clip", From 2a5ce6a5644c64d66535d16421def23ad32c0f2b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CElla?= Date: Tue, 2 Jun 2026 12:51:43 +0200 Subject: [PATCH 39/40] remove patched_speecht5_prenet_forward --- optimum/exporters/openvino/model_patcher.py | 34 --------------------- 1 file changed, 34 deletions(-) diff --git a/optimum/exporters/openvino/model_patcher.py b/optimum/exporters/openvino/model_patcher.py index b84e8c89ed..199bbce903 100644 --- a/optimum/exporters/openvino/model_patcher.py +++ b/optimum/exporters/openvino/model_patcher.py @@ -182,40 +182,6 @@ def patched_forward( self.patched_forward = patched_forward -def patched_speecht5_prenet_forward( - self, - input_values: torch.Tensor, - speaker_embeddings: torch.Tensor | None = None, -): - # Dropout is always applied, even when evaluating. See ยง2.2 in https://arxiv.org/abs/1712.05884. - - inputs_embeds = input_values - for layer in self.layers: - inputs_embeds = torch.nn.functional.relu(layer(inputs_embeds)) - - # NOTE: we patch the prenet to avoid using torch.nn.functional.dropout, that is exported as a `Dropout` node in the OpenVINO - # that is ignored during inference by some runtimes as OpenVINO Runtime. - # Reference: https://github.com/microsoft/onnxruntime/issues/9333 & https://github.com/microsoft/onnxruntime/issues/5549 - mask = torch.rand(inputs_embeds.shape, device=inputs_embeds.device) > self.config.speech_decoder_prenet_dropout - inputs_embeds = inputs_embeds * mask / (1 - self.config.speech_decoder_prenet_dropout) - - # inputs_embeds = nn.functional.dropout( - # inputs_embeds, self.config.speech_decoder_prenet_dropout, training=True - # ) - - inputs_embeds = self.final_layer(inputs_embeds) - inputs_embeds = self.encode_positions(inputs_embeds) - - if speaker_embeddings is not None: - speaker_embeddings = torch.nn.functional.normalize(speaker_embeddings) - speaker_embeddings = speaker_embeddings.unsqueeze(1) - speaker_embeddings = speaker_embeddings.expand(-1, inputs_embeds.size(1), -1) - inputs_embeds = torch.cat([inputs_embeds, speaker_embeddings], dim=-1) - inputs_embeds = torch.nn.functional.relu(self.speaker_embeds_layer(inputs_embeds)) - - return inputs_embeds - - class SentenceTransformersTransformerPatcher(ModelPatcher): def __init__( self, From 6f997758971b4c4373e9c60c2ca68404733a2f7b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CElla?= Date: Tue, 2 Jun 2026 15:14:55 +0200 Subject: [PATCH 40/40] rename onnx to openvino --- optimum/intel/openvino/modeling.py | 8 ++++---- optimum/intel/openvino/modeling_base.py | 6 +++--- optimum/intel/openvino/quantization.py | 3 --- 3 files changed, 7 insertions(+), 10 deletions(-) diff --git a/optimum/intel/openvino/modeling.py b/optimum/intel/openvino/modeling.py index 7abaf0bccf..91cf0149ec 100644 --- a/optimum/intel/openvino/modeling.py +++ b/optimum/intel/openvino/modeling.py @@ -578,9 +578,9 @@ def from_pretrained( if local_timm_model: return super()._from_pretrained(model_id=model_id, config=config) model = TimmForImageClassification.from_pretrained(model_id, **kwargs) - onnx_config = TimmOpenVINOConfig(model.config) + openvino_config = TimmOpenVINOConfig(model.config) - return cls._to_load(model=model, config=config, onnx_config=onnx_config, stateful=False, **kwargs) + return cls._to_load(model=model, config=config, exporter_config=openvino_config, stateful=False, **kwargs) else: return super().from_pretrained( model_id=model_id, @@ -716,7 +716,7 @@ def forward( @add_start_docstrings( """ - Onnx Model with a language modeling head on top for Connectionist Temporal Classification (CTC). + OpenVINO Model with a language modeling head on top for Connectionist Temporal Classification (CTC). """, MODEL_START_DOCSTRING, ) @@ -796,7 +796,7 @@ def forward( @add_start_docstrings( """ - Onnx Model with an XVector feature extraction head on top for tasks like Speaker Verification. + OpenVINO Model with an XVector feature extraction head on top for tasks like Speaker Verification. """, MODEL_START_DOCSTRING, ) diff --git a/optimum/intel/openvino/modeling_base.py b/optimum/intel/openvino/modeling_base.py index a735e1c1ad..40c9314d76 100644 --- a/optimum/intel/openvino/modeling_base.py +++ b/optimum/intel/openvino/modeling_base.py @@ -905,7 +905,7 @@ def _to_load( cls, model, config: PretrainedConfig, - onnx_config: ExportConfig, + exporter_config: ExportConfig, token: Optional[Union[bool, str]] = None, revision: Optional[str] = None, force_download: bool = False, @@ -924,10 +924,10 @@ def _to_load( ) compile_only = False - # Export the model to the ONNX format + # Export the model to the OpenVINO format export( model=model, - config=onnx_config, + config=exporter_config, output=save_dir_path / cls._all_ov_model_paths["model"], stateful=stateful, ) diff --git a/optimum/intel/openvino/quantization.py b/optimum/intel/openvino/quantization.py index 30005894f5..eb421646dc 100644 --- a/optimum/intel/openvino/quantization.py +++ b/optimum/intel/openvino/quantization.py @@ -1307,7 +1307,6 @@ def quantize( ] = None, save_directory: Optional[Union[str, Path]] = None, ov_config: OVConfig = None, - file_name: Optional[str] = None, batch_size: int = 1, data_collator: Optional[DataCollator] = None, remove_unused_columns: bool = False, @@ -1326,8 +1325,6 @@ def quantize( ov_config (`OVConfig`, *optional*): The configuration containing the parameters related to quantization. If not provided, 8-bit symmetric weight-only quantization will be applied. - file_name (`str`, *optional*): - The model file name to use when saving the model. Overwrites the default file name `"model.onnx"`. batch_size (`int`, defaults to 1): The number of calibration samples to load per batch. data_collator (`DataCollator`, *optional*):