patch all clip variants

huggingface · Jul 25, 2024 · b610212 · b610212
1 parent c2a5c03
commit b610212
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 4 deletions.
diff --git a/optimum/exporters/onnx/model_configs.py b/optimum/exporters/onnx/model_configs.py
@@ -71,7 +71,7 @@
 )
 from .constants import ONNX_DECODER_MERGED_NAME, ONNX_DECODER_NAME, ONNX_DECODER_WITH_PAST_NAME
 from .model_patcher import (
-    CLIPTextModelPatcher,
+    CLIPModelPatcher,
     FalconModelPatcher,
     MistralModelPatcher,
     MusicgenModelPatcher,
@@ -912,10 +912,16 @@ def outputs(self) -> Dict[str, Dict[int, str]]:
 
         return common_outputs
 
+    def patch_model_for_export(
+        self,
+        model: Union["PreTrainedModel", "TFPreTrainedModel", "ModelMixin"],
+        model_kwargs: Optional[Dict[str, Any]] = None,
+    ) -> "ModelPatcher":
+        return CLIPModelPatcher(self, model, model_kwargs=model_kwargs)
+
 
 class CLIPOnnxConfig(TextAndVisionOnnxConfig):
     NORMALIZED_CONFIG_CLASS = CLIPNormalizedConfig
-    DEFAULT_ONNX_OPSET = 14
 
     @property
     def inputs(self) -> Dict[str, Dict[int, str]]:
@@ -934,6 +940,13 @@ def outputs(self) -> Dict[str, Dict[int, str]]:
             "image_embeds": {0: "image_batch_size"},
         }
 
+    def patch_model_for_export(
+        self,
+        model: Union["PreTrainedModel", "TFPreTrainedModel", "ModelMixin"],
+        model_kwargs: Optional[Dict[str, Any]] = None,
+    ) -> "ModelPatcher":
+        return CLIPModelPatcher(self, model, model_kwargs=model_kwargs)
+
 
 class SentenceTransformersCLIPOnnxConfig(CLIPOnnxConfig):
     @property
@@ -984,7 +997,7 @@ def patch_model_for_export(
         model: Union["PreTrainedModel", "TFPreTrainedModel", "ModelMixin"],
         model_kwargs: Optional[Dict[str, Any]] = None,
     ) -> "ModelPatcher":
-        return CLIPTextModelPatcher(self, model, model_kwargs=model_kwargs)
+        return CLIPModelPatcher(self, model, model_kwargs=model_kwargs)
 
 
 class CLIPTextOnnxConfig(CLIPTextWithProjectionOnnxConfig):
@@ -1000,6 +1013,13 @@ def outputs(self) -> Dict[str, Dict[int, str]]:
 
         return common_outputs
 
+    def patch_model_for_export(
+        self,
+        model: Union["PreTrainedModel", "TFPreTrainedModel", "ModelMixin"],
+        model_kwargs: Optional[Dict[str, Any]] = None,
+    ) -> "ModelPatcher":
+        return CLIPModelPatcher(self, model, model_kwargs=model_kwargs)
+
 
 class UNetOnnxConfig(VisionOnnxConfig):
     ATOL_FOR_VALIDATION = 1e-3

diff --git a/optimum/exporters/onnx/model_patcher.py b/optimum/exporters/onnx/model_patcher.py
@@ -1133,7 +1133,7 @@ def __init__(
                 self._update_causal_mask_original = self._model._update_causal_mask
 
 
-class CLIPTextModelPatcher(ModelPatcher):
+class CLIPModelPatcher(ModelPatcher):
     def __enter__(self):
         super().__enter__()