From 58e8dacf36e3a5c5ae62c9518e05e851d5269a28 Mon Sep 17 00:00:00 2001 From: "Kazantsev, Roman" Date: Wed, 15 Apr 2026 21:36:58 +0000 Subject: [PATCH] Remove optimum-onnx deps --- optimum/exporters/ipex/model_config.py | 2 +- optimum/exporters/openvino/__init__.py | 2 +- optimum/exporters/openvino/__main__.py | 2 +- .../openvino/_onnx_compat/__init__.py | 20 + .../openvino/_onnx_compat/_traceable_cache.py | 107 + .../_onnx_compat/_traceable_decorator.py | 175 + .../exporters/openvino/_onnx_compat/base.py | 775 +++++ .../exporters/openvino/_onnx_compat/config.py | 383 +++ .../openvino/_onnx_compat/constants.py | 40 + .../openvino/_onnx_compat/input_generators.py | 110 + .../openvino/_onnx_compat/model_configs.py | 2856 +++++++++++++++++ .../openvino/_onnx_compat/model_patcher.py | 1461 +++++++++ .../_onnx_compat/model_patcher_dynamo.py | 59 + optimum/exporters/openvino/convert.py | 69 +- optimum/exporters/openvino/model_configs.py | 6 +- optimum/exporters/openvino/model_patcher.py | 6 +- optimum/exporters/openvino/utils.py | 2 +- optimum/intel/openvino/modeling_timm.py | 2 +- setup.py | 2 +- 19 files changed, 5999 insertions(+), 80 deletions(-) create mode 100644 optimum/exporters/openvino/_onnx_compat/__init__.py create mode 100644 optimum/exporters/openvino/_onnx_compat/_traceable_cache.py create mode 100644 optimum/exporters/openvino/_onnx_compat/_traceable_decorator.py create mode 100644 optimum/exporters/openvino/_onnx_compat/base.py create mode 100644 optimum/exporters/openvino/_onnx_compat/config.py create mode 100644 optimum/exporters/openvino/_onnx_compat/constants.py create mode 100644 optimum/exporters/openvino/_onnx_compat/input_generators.py create mode 100644 optimum/exporters/openvino/_onnx_compat/model_configs.py create mode 100644 optimum/exporters/openvino/_onnx_compat/model_patcher.py create mode 100644 optimum/exporters/openvino/_onnx_compat/model_patcher_dynamo.py diff --git a/optimum/exporters/ipex/model_config.py b/optimum/exporters/ipex/model_config.py index 5cd7f1e79d..e338ab1223 100755 --- a/optimum/exporters/ipex/model_config.py +++ b/optimum/exporters/ipex/model_config.py @@ -14,7 +14,7 @@ from typing import Optional, Tuple -from optimum.exporters.onnx.model_configs import ( +from optimum.exporters.openvino._onnx_compat.model_configs import ( FalconOnnxConfig, GPT2OnnxConfig, LlamaOnnxConfig, 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 f8637d2d1e..631dd65b88 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._onnx_compat.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/_onnx_compat/__init__.py b/optimum/exporters/openvino/_onnx_compat/__init__.py new file mode 100644 index 0000000000..e45d29c26e --- /dev/null +++ b/optimum/exporters/openvino/_onnx_compat/__init__.py @@ -0,0 +1,20 @@ +# Copyright 2024 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. + +""" +Compatibility module providing alternatives to optimum-onnx exports. + +This module contains classes and functions originally from optimum.exporters.onnx +that are needed by optimum-intel, extracted here to remove the optimum-onnx dependency. +""" diff --git a/optimum/exporters/openvino/_onnx_compat/_traceable_cache.py b/optimum/exporters/openvino/_onnx_compat/_traceable_cache.py new file mode 100644 index 0000000000..059463f9cf --- /dev/null +++ b/optimum/exporters/openvino/_onnx_compat/_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/_onnx_compat/_traceable_decorator.py b/optimum/exporters/openvino/_onnx_compat/_traceable_decorator.py new file mode 100644 index 0000000000..0a2087dff7 --- /dev/null +++ b/optimum/exporters/openvino/_onnx_compat/_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/_onnx_compat/base.py b/optimum/exporters/openvino/_onnx_compat/base.py new file mode 100644 index 0000000000..f99ff6315b --- /dev/null +++ b/optimum/exporters/openvino/_onnx_compat/base.py @@ -0,0 +1,775 @@ +# 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 os +import re +from abc import ABC +from collections import OrderedDict +from collections.abc import Iterable +from pathlib import Path +from typing import TYPE_CHECKING, Any, ClassVar + +import numpy as np +from transformers.utils import is_accelerate_available, is_torch_available + + +if is_torch_available(): + import torch.nn as nn + +from optimum.exporters.base import ExporterConfig +from optimum.exporters.openvino._onnx_compat.constants import ONNX_DECODER_MERGED_NAME, ONNX_DECODER_NAME, ONNX_DECODER_WITH_PAST_NAME +from optimum.exporters.openvino._onnx_compat.model_patcher import ModelPatcher +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, +) + + +if is_accelerate_available(): + from accelerate.utils import find_tied_parameters + +if TYPE_CHECKING: + from transformers import PretrainedConfig, PreTrainedModel + + from optimum.exporters.openvino._onnx_compat.model_patcher import PatchingSpec + + +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" + ) + + from onnxruntime import GraphOptimizationLevel, InferenceSession, SessionOptions + + import onnx + + 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 post_process_exported_models( + self, + path: Path, + models_and_onnx_configs: dict[str, tuple[PreTrainedModel, OnnxConfig]], + onnx_files_subpaths: list[str], + ): + """Performs any model-specific post-processing on the exported models.""" + return models_and_onnx_configs, onnx_files_subpaths + + 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 post_process_exported_models( + self, + path: Path, + models_and_onnx_configs: dict[str, tuple[PreTrainedModel, OnnxConfig]], + onnx_files_subpaths: list[str], + ): + models_and_onnx_configs, onnx_files_subpaths = super().post_process_exported_models( + path, models_and_onnx_configs, onnx_files_subpaths + ) + return models_and_onnx_configs, onnx_files_subpaths + + 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/_onnx_compat/config.py b/optimum/exporters/openvino/_onnx_compat/config.py new file mode 100644 index 0000000000..60d29797ce --- /dev/null +++ b/optimum/exporters/openvino/_onnx_compat/config.py @@ -0,0 +1,383 @@ +# 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 pathlib import Path +from typing import TYPE_CHECKING, Any + +from optimum.exporters.openvino._onnx_compat.base import ConfigBehavior, OnnxConfig, OnnxConfigWithPast, OnnxSeq2SeqConfigWithPast +from optimum.exporters.openvino._onnx_compat.constants import ONNX_DECODER_MERGED_NAME, ONNX_DECODER_NAME, ONNX_DECODER_WITH_PAST_NAME +from optimum.exporters.tasks import TasksManager +from optimum.utils import ( + DummyAudioInputGenerator, + DummyBboxInputGenerator, + DummyInputGenerator, + DummyPastKeyValuesGenerator, + DummySeq2SeqDecoderTextInputGenerator, + DummySeq2SeqPastKeyValuesGenerator, + DummyTextInputGenerator, + DummyVisionInputGenerator, + logging, +) + + +if TYPE_CHECKING: + from transformers import PretrainedConfig, PreTrainedModel + + +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 + + def post_process_exported_models( + self, + path: Path, + models_and_onnx_configs: dict[str, tuple[PreTrainedModel, OnnxConfig]], + onnx_files_subpaths: list[str], + ): + models_and_onnx_configs, onnx_files_subpaths = super().post_process_exported_models( + path, models_and_onnx_configs, onnx_files_subpaths + ) + return models_and_onnx_configs, onnx_files_subpaths + + +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 + ) + + def post_process_exported_models( + self, + path: Path, + models_and_onnx_configs: dict[str, tuple[PreTrainedModel, OnnxConfig]], + onnx_files_subpaths: list[str], + ): + models_and_onnx_configs, onnx_files_subpaths = super().post_process_exported_models( + path, models_and_onnx_configs, onnx_files_subpaths + ) + if self.use_past is True and len(models_and_onnx_configs) == 3: + models_and_onnx_configs[ONNX_DECODER_NAME][1]._decoder_onnx_config.is_merged = True + models_and_onnx_configs[ONNX_DECODER_NAME][1]._decoder_onnx_config.use_cache_branch = False + + # Past key values won't be generated by default, but added in the input + models_and_onnx_configs[ONNX_DECODER_NAME][1]._decoder_onnx_config.use_past_in_inputs = True + + models_and_onnx_configs[ONNX_DECODER_WITH_PAST_NAME][1]._decoder_onnx_config.use_cache_branch = True + models_and_onnx_configs[ONNX_DECODER_WITH_PAST_NAME][1]._decoder_onnx_config.is_merged = True + + return models_and_onnx_configs, onnx_files_subpaths diff --git a/optimum/exporters/openvino/_onnx_compat/constants.py b/optimum/exporters/openvino/_onnx_compat/constants.py new file mode 100644 index 0000000000..a996bac53e --- /dev/null +++ b/optimum/exporters/openvino/_onnx_compat/constants.py @@ -0,0 +1,40 @@ +# Copyright 2023 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. + +# 2 GB +EXTERNAL_DATA_FORMAT_SIZE_LIMIT = 2 * 1024 * 1024 * 1024 + +ONNX_ENCODER_NAME = "encoder_model" +ONNX_DECODER_NAME = "decoder_model" +ONNX_DECODER_WITH_PAST_NAME = "decoder_with_past_model" +ONNX_DECODER_MERGED_NAME = "decoder_model_merged" + +UNPICKABLE_ARCHS = [ + "encodec", + "hubert", + "sew", + "sew-d", + "speecht5", + "unispeech", + "unispeech-sat", + "wav2vec2", + "wav2vec2-conformer", + "wavlm", +] + +SDPA_ARCHS_ONNX_EXPORT_NOT_SUPPORTED = [ + "bart", + "musicgen", + "whisper", +] diff --git a/optimum/exporters/openvino/_onnx_compat/input_generators.py b/optimum/exporters/openvino/_onnx_compat/input_generators.py new file mode 100644 index 0000000000..e11ebd16e6 --- /dev/null +++ b/optimum/exporters/openvino/_onnx_compat/input_generators.py @@ -0,0 +1,110 @@ +# 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 + +from optimum.utils import ( + DummyAudioInputGenerator, + DummyPastKeyValuesGenerator, + DummyTransformerTextInputGenerator, + NormalizedTextConfig, + is_transformers_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 + ) diff --git a/optimum/exporters/openvino/_onnx_compat/model_configs.py b/optimum/exporters/openvino/_onnx_compat/model_configs.py new file mode 100644 index 0000000000..8141310077 --- /dev/null +++ b/optimum/exporters/openvino/_onnx_compat/model_configs.py @@ -0,0 +1,2856 @@ +# 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. +"""Model specific ONNX configurations.""" + +from __future__ import annotations + +import math +import warnings +from typing import TYPE_CHECKING, Any, Literal + +from packaging import version + +from optimum.exporters.openvino._onnx_compat.base import ConfigBehavior, OnnxConfig, OnnxConfigWithPast, OnnxSeq2SeqConfigWithPast +from optimum.exporters.openvino._onnx_compat.config import ( + AudioOnnxConfig, + AudioToTextOnnxConfig, + EncoderDecoderBaseOnnxConfig, + TextAndVisionOnnxConfig, + TextDecoderOnnxConfig, + TextDecoderWithPositionIdsOnnxConfig, + TextEncoderOnnxConfig, + TextSeq2SeqOnnxConfig, + VisionOnnxConfig, +) +from optimum.exporters.openvino._onnx_compat.input_generators import ( + DummyMoonshineAudioInputGenerator, + DummySanaTransforemerTextInputGenerator, + GPTBigCodeDummyPastKeyValuesGenerator, +) +from optimum.exporters.openvino._onnx_compat.model_patcher import ( + BigBirdPegasusModelPatcher, + CLIPModelPatcher, + CohereModelPatcher, + FluxTransformerModelPatcher, + GptOssModelPatcher, + MetaCLIP2Patcher, + MgpstrModelPatcher, + MoonshineModelPatcher, + MusicgenModelPatcher, + Qwen3MoeModelPatcher, + SAMModelPatcher, + SentenceTransformersCLIPPatcher, + SentenceTransformersTransformerPatcher, + SpeechT5ModelPatcher, + VitPoseModelPatcher, +) +from optimum.exporters.openvino._onnx_compat.model_patcher_dynamo import ViTForImageClassificationPatcher +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, + logging, +) +from optimum.utils.normalized_config import NormalizedConfigManager + + +if TYPE_CHECKING: + from transformers import PretrainedConfig + + +logger = logging.get_logger(__name__) + + +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) + + 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}", + } + } diff --git a/optimum/exporters/openvino/_onnx_compat/model_patcher.py b/optimum/exporters/openvino/_onnx_compat/model_patcher.py new file mode 100644 index 0000000000..11200d8458 --- /dev/null +++ b/optimum/exporters/openvino/_onnx_compat/model_patcher.py @@ -0,0 +1,1461 @@ +# 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. +from __future__ import annotations + +import dataclasses +import functools +import inspect +import sys +import types +from typing import TYPE_CHECKING, Any, Callable + +import torch +import transformers +from torch.onnx import symbolic_helper +from transformers.modeling_outputs import BaseModelOutput +from transformers.models.speecht5.modeling_speecht5 import SpeechT5EncoderWithSpeechPrenet + +from optimum.utils import is_diffusers_version, is_torch_version, is_transformers_version, logging + + +if is_transformers_version(">=", "4.44") and is_transformers_version("<", "4.50"): + from optimum.exporters.openvino._onnx_compat._traceable_cache import TraceableCache +if is_transformers_version(">=", "4.54"): + from optimum.exporters.openvino._onnx_compat._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 + +if TYPE_CHECKING: + from transformers import PreTrainedModel + + from optimum.exporters.openvino._onnx_compat.base import OnnxConfig + + +logger = logging.get_logger(__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 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 + + +# 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 + + +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): + 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 diff --git a/optimum/exporters/openvino/_onnx_compat/model_patcher_dynamo.py b/optimum/exporters/openvino/_onnx_compat/model_patcher_dynamo.py new file mode 100644 index 0000000000..5427e8dec3 --- /dev/null +++ b/optimum/exporters/openvino/_onnx_compat/model_patcher_dynamo.py @@ -0,0 +1,59 @@ +# 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. +import torch +from transformers.models.vit.modeling_vit import ViTPatchEmbeddings + +from .model_patcher import ModelPatcher, is_transformers_version + + +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/convert.py b/optimum/exporters/openvino/convert.py index fddd840b7d..ea85a5c27d 100644 --- a/optimum/exporters/openvino/convert.py +++ b/optimum/exporters/openvino/convert.py @@ -90,7 +90,7 @@ if TYPE_CHECKING: - from optimum.exporters.onnx.base import OnnxConfig + from optimum.exporters.openvino._onnx_compat.base import OnnxConfig from optimum.intel.openvino.configuration import OVConfig @@ -233,73 +233,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", diff --git a/optimum/exporters/openvino/model_configs.py b/optimum/exporters/openvino/model_configs.py index 2e24a3fdc6..8a63eeefaf 100644 --- a/optimum/exporters/openvino/model_configs.py +++ b/optimum/exporters/openvino/model_configs.py @@ -19,8 +19,8 @@ from transformers import AutoConfig, PretrainedConfig, PreTrainedModel -from optimum.exporters.onnx.config import OnnxConfig, TextDecoderOnnxConfig, TextDecoderWithPositionIdsOnnxConfig -from optimum.exporters.onnx.model_configs import ( +from optimum.exporters.openvino._onnx_compat.config import OnnxConfig, TextDecoderOnnxConfig, TextDecoderWithPositionIdsOnnxConfig +from optimum.exporters.openvino._onnx_compat.model_configs import ( AlbertOnnxConfig, ASTOnnxConfig, BartOnnxConfig, @@ -105,7 +105,7 @@ XLMOnnxConfig, XLMRobertaOnnxConfig, ) -from optimum.exporters.onnx.model_patcher import ModelPatcher +from optimum.exporters.openvino._onnx_compat.model_patcher import ModelPatcher from optimum.exporters.openvino.utils import ONNX_SUPPORTED_ARCHITECTURES from optimum.exporters.tasks import TasksManager from optimum.utils import DEFAULT_DUMMY_SHAPES diff --git a/optimum/exporters/openvino/model_patcher.py b/optimum/exporters/openvino/model_patcher.py index 4ae2203561..706ce716c7 100644 --- a/optimum/exporters/openvino/model_patcher.py +++ b/optimum/exporters/openvino/model_patcher.py @@ -46,8 +46,8 @@ 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 ( +from optimum.exporters.openvino._onnx_compat.base import OnnxConfig +from optimum.exporters.openvino._onnx_compat.model_patcher import ( UNSUPPORTED_OPS_PATCHING_SPEC, ModelPatcher, gpt_oss_forward, @@ -78,7 +78,7 @@ from transformers.cache_utils import Cache from transformers.modeling_utils import PreTrainedModel - from optimum.exporters.onnx.config import OnnxConfig + from optimum.exporters.openvino._onnx_compat.config import OnnxConfig if is_transformers_version(">=", "4.54"): from transformers.utils import TransformersKwargs diff --git a/optimum/exporters/openvino/utils.py b/optimum/exporters/openvino/utils.py index af2f1edaba..a32fb6832a 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._onnx_compat.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/optimum/intel/openvino/modeling_timm.py b/optimum/intel/openvino/modeling_timm.py index f2566c8c4a..774e3dbe8e 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._onnx_compat.config import VisionOnnxConfig from optimum.utils import NormalizedVisionConfig from .utils import _is_timm_ov_dir diff --git a/setup.py b/setup.py index 5f9efdf679..7b6872f36d 100644 --- a/setup.py +++ b/setup.py @@ -28,7 +28,7 @@ INSTALL_REQUIRE = [ "torch>=2.1", - "optimum-onnx@git+https://github.com/huggingface/optimum-onnx.git@transformers-v5", + "optimum>=2.1.0", "transformers>=4.45,<5.1", "setuptools", "huggingface-hub>=0.23.2,<2.0",