Add openvino support for OpenCLIP (#857)

* Add support for OpenCLIP

* fix comments

* fix comments

* rename class

* remove patching

---------

Co-authored-by: Ella Charlaix <[email protected]>

Author: sbalandi

optimum/commands/export/openvino.py

@@ -22,6 +22,7 @@
 
 from ...exporters import TasksManager
 from ...intel.utils.import_utils import DIFFUSERS_IMPORT_ERROR, is_diffusers_available
+from ...intel.utils.modeling_utils import _infer_library_from_model_name_or_path
 from ...utils.save_utils import maybe_load_preprocessors, maybe_save_preprocessors
 from ..base import BaseOptimumCLICommand, CommandInfo
 
@@ -77,7 +78,7 @@ def parse_args_openvino(parser: "ArgumentParser"):
     optional_group.add_argument(
         "--library",
         type=str,
-        choices=["transformers", "diffusers", "timm", "sentence_transformers"],
+        choices=["transformers", "diffusers", "timm", "sentence_transformers", "open_clip"],
         default=None,
         help="The library used to load the model before export. If not provided, will attempt to infer the local checkpoint's library",
     )
@@ -234,7 +235,7 @@ def run(self):
 
         if self.args.library is None:
             # TODO: add revision, subfolder and token to args
-            library_name = TasksManager._infer_library_from_model_name_or_path(
+            library_name = _infer_library_from_model_name_or_path(
                 model_name_or_path=self.args.model, cache_dir=self.args.cache_dir
             )
             if library_name == "sentence_transformers":
@@ -280,7 +281,7 @@ def run(self):
 
         quantization_config = ov_config.quantization_config if ov_config else None
         quantize_with_dataset = quantization_config and getattr(quantization_config, "dataset", None) is not None
-        task = infer_task(self.args.task, self.args.model)
+        task = infer_task(self.args.task, self.args.model, library_name=library_name)
 
         if library_name == "diffusers" and quantize_with_dataset:
             if not is_diffusers_available():

optimum/exporters/openvino/__main__.py

@@ -36,6 +36,10 @@
     is_openvino_version,
     is_transformers_version,
 )
+from optimum.intel.utils.modeling_utils import (
+    _infer_library_from_model_name_or_path,
+    _OpenClipForZeroShotImageClassification,
+)
 from optimum.utils.save_utils import maybe_load_preprocessors
 
 from .utils import _MAX_UNCOMPRESSED_SIZE, clear_class_registry
@@ -59,25 +63,29 @@ def infer_task(
     revision: Optional[str] = None,
     cache_dir: str = HUGGINGFACE_HUB_CACHE,
     token: Optional[Union[bool, str]] = None,
+    library_name: Optional[str] = None,
 ):
     task = TasksManager.map_from_synonym(task)
     if task == "auto":
-        try:
-            task = TasksManager._infer_task_from_model_name_or_path(
-                model_name_or_path=model_name_or_path,
-                subfolder=subfolder,
-                revision=revision,
-                cache_dir=cache_dir,
-                token=token,
-            )
-        except KeyError as e:
-            raise KeyError(
-                f"The task could not be automatically inferred. Please provide the argument --task with the relevant task from {', '.join(TasksManager.get_all_tasks())}. Detailed error: {e}"
-            )
-        except RequestsConnectionError as e:
-            raise RequestsConnectionError(
-                f"The task could not be automatically inferred as this is available only for models hosted on the Hugging Face Hub. Please provide the argument --task with the relevant task from {', '.join(TasksManager.get_all_tasks())}. Detailed error: {e}"
-            )
+        if library_name == "open_clip":
+            task = "zero-shot-image-classification"
+        else:
+            try:
+                task = TasksManager._infer_task_from_model_name_or_path(
+                    model_name_or_path=model_name_or_path,
+                    subfolder=subfolder,
+                    revision=revision,
+                    cache_dir=cache_dir,
+                    token=token,
+                )
+            except KeyError as e:
+                raise KeyError(
+                    f"The task could not be automatically inferred. Please provide the argument --task with the relevant task from {', '.join(TasksManager.get_all_tasks())}. Detailed error: {e}"
+                )
+            except RequestsConnectionError as e:
+                raise RequestsConnectionError(
+                    f"The task could not be automatically inferred as this is available only for models hosted on the Hugging Face Hub. Please provide the argument --task with the relevant task from {', '.join(TasksManager.get_all_tasks())}. Detailed error: {e}"
+                )
     return task
 
 
@@ -182,16 +190,13 @@ def main_export(
             raise ValueError("You cannot use both `use_auth_token` and `token` arguments at the same time.")
         token = use_auth_token
 
-    original_task = task
-    task = infer_task(
-        task, model_name_or_path, subfolder=subfolder, revision=revision, cache_dir=cache_dir, token=token
-    )
-    framework = TasksManager.determine_framework(
-        model_name_or_path, subfolder=subfolder, revision=revision, cache_dir=cache_dir, token=token
-    )
+    if framework is None:
+        framework = TasksManager.determine_framework(
+            model_name_or_path, subfolder=subfolder, revision=revision, cache_dir=cache_dir, token=token
+        )
 
     if library_name is None:
-        library_name = TasksManager._infer_library_from_model_name_or_path(
+        library_name = _infer_library_from_model_name_or_path(
             model_name_or_path=model_name_or_path,
             subfolder=subfolder,
             revision=revision,
@@ -205,6 +210,17 @@ def main_export(
             )
             library_name = "transformers"
 
+    original_task = task
+    task = infer_task(
+        task,
+        model_name_or_path,
+        subfolder=subfolder,
+        revision=revision,
+        cache_dir=cache_dir,
+        token=token,
+        library_name=library_name,
+    )
+
     do_gptq_patching = False
     custom_architecture = False
     patch_16bit = False
@@ -305,21 +321,24 @@ class StoreAttr(object):
 
         GPTQQuantizer.post_init_model = post_init_model
 
-    model = TasksManager.get_model_from_task(
-        task,
-        model_name_or_path,
-        subfolder=subfolder,
-        revision=revision,
-        cache_dir=cache_dir,
-        token=token,
-        local_files_only=local_files_only,
-        force_download=force_download,
-        trust_remote_code=trust_remote_code,
-        framework=framework,
-        device=device,
-        library_name=library_name,
-        **loading_kwargs,
-    )
+    if library_name == "open_clip":
+        model = _OpenClipForZeroShotImageClassification.from_pretrained(model_name_or_path, cache_dir=cache_dir)
+    else:
+        model = TasksManager.get_model_from_task(
+            task,
+            model_name_or_path,
+            subfolder=subfolder,
+            revision=revision,
+            cache_dir=cache_dir,
+            token=token,
+            local_files_only=local_files_only,
+            force_download=force_download,
+            trust_remote_code=trust_remote_code,
+            framework=framework,
+            device=device,
+            library_name=library_name,
+            **loading_kwargs,
+        )
 
     needs_pad_token_id = task == "text-classification" and getattr(model.config, "pad_token_id", None) is None
 

optimum/exporters/openvino/convert.py

@@ -33,6 +33,7 @@
 from optimum.exporters.utils import _get_submodels_and_export_configs
 from optimum.intel.utils.import_utils import (
     _nncf_version,
+    _open_clip_version,
     _optimum_intel_version,
     _optimum_version,
     _timm_version,
@@ -44,11 +45,13 @@
 from optimum.utils.save_utils import maybe_save_preprocessors
 
 from ...intel.utils.import_utils import is_nncf_available
+from ...intel.utils.modeling_utils import _infer_library_from_model_or_model_class
 from .model_patcher import patch_model_with_bettertransformer
 from .stateful import ensure_export_task_support_stateful, ensure_stateful_is_available, patch_stateful
 from .utils import (
     OV_XML_FILE_NAME,
     _get_input_info,
+    _get_open_clip_submodels_fn_and_export_configs,
     clear_class_registry,
     remove_none_from_dummy_inputs,
 )
@@ -183,11 +186,7 @@ def export_tensorflow(
     input_names, output_names = export_tensorflow_onnx(model, config, opset, onnx_path)
     ov_model = convert_model(str(onnx_path))
 
-    if model.__class__.__module__.startswith("optimum"):
-        # for wrapped models
-        library_name = TasksManager._infer_library_from_model_or_model_class(model=model.model)
-    else:
-        library_name = TasksManager._infer_library_from_model_or_model_class(model=model)
+    library_name = _infer_library_from_model_or_model_class(model=model)
 
     _save_model(
         ov_model,
@@ -248,11 +247,7 @@ def export_pytorch_via_onnx(
     torch.onnx.export = orig_torch_onnx_export
     ov_model = convert_model(str(onnx_output))
 
-    if model.__class__.__module__.startswith("optimum"):
-        # for wrapped models
-        library_name = TasksManager._infer_library_from_model_or_model_class(model=model.model)
-    else:
-        library_name = TasksManager._infer_library_from_model_or_model_class(model=model)
+    library_name = _infer_library_from_model_or_model_class(model=model)
 
     _save_model(
         ov_model,
@@ -422,11 +417,7 @@ def ts_patched_forward(*args, **kwargs):
         if stateful:
             patch_stateful(model.config, ov_model)
 
-        if model.__module__.startswith("optimum"):
-            # for wrapped models like timm in optimum.intel.openvino.modeling_timm
-            library_name = TasksManager._infer_library_from_model_or_model_class(model=model.model)
-        else:
-            library_name = TasksManager._infer_library_from_model_or_model_class(model=model)
+        library_name = _infer_library_from_model_or_model_class(model=model)
 
         _save_model(
             ov_model,
@@ -535,8 +526,9 @@ def export_from_model(
             f"Compression of the weights to {ov_config.quantization_config} requires nncf, please install it with `pip install nncf`"
         )
 
-    library_name = TasksManager._infer_library_from_model_or_model_class(model=model)
-    TasksManager.standardize_model_attributes(model)
+    library_name = _infer_library_from_model_or_model_class(model)
+    if library_name != "open_clip":
+        TasksManager.standardize_model_attributes(model)
 
     if hasattr(model.config, "export_model_type"):
         model_type = model.config.export_model_type.replace("_", "-")
@@ -597,6 +589,12 @@ def export_from_model(
             kwargs_shapes[input_name] if input_name in kwargs_shapes else DEFAULT_DUMMY_SHAPES[input_name]
         )
 
+    if library_name == "open_clip":
+        custom_architecture = True
+        custom_export_configs, fn_get_submodels = _get_open_clip_submodels_fn_and_export_configs(
+            model, library_name, task, preprocessors, custom_export_configs, fn_get_submodels
+        )
+
     logging.disable(logging.INFO)
     export_config, models_and_export_configs = _get_submodels_and_export_configs(
         model=model,
@@ -614,7 +612,16 @@ def export_from_model(
     )
     logging.disable(logging.NOTSET)
 
-    if library_name != "diffusers":
+    if library_name == "open_clip":
+        if hasattr(model.config, "save_pretrained"):
+            model.config.save_pretrained(output)
+
+        for preprocess in preprocessors:
+            if hasattr(preprocess, "save_pretrained"):
+                preprocess.save_pretrained(output)
+
+        files_subpaths = ["openvino_" + model_name + ".xml" for model_name in models_and_export_configs.keys()]
+    elif library_name != "diffusers":
         # Saving the model config and preprocessor as this is needed sometimes.
         model.config.save_pretrained(output)
         generation_config = getattr(model, "generation_config", None)
@@ -744,6 +751,8 @@ def _add_version_info_to_model(model: Model, library_name: Optional[str] = None)
             model.set_rt_info(_optimum_version, ["optimum", "diffusers_version"])
         elif library_name == "timm":
             model.set_rt_info(_timm_version, ["optimum", "timm_version"])
+        elif library_name == "open_clip":
+            model.set_rt_info(_open_clip_version, ["optimum", "open_clip_version"])
         rt_info = model.get_rt_info()
         if "nncf" in rt_info:
             model.set_rt_info(_nncf_version, ["optimum", "nncf_version"])