Diffusion tasks (huggingface#827)

* adapt to the new tasks * Update optimum/exporters/openvino/convert.py Co-authored-by: Ella Charlaix <[email protected]> * update tasks manager * fix * explicit kwargs * fix * fix model saving * fix library name overriding * fix _save_model call assesrtions * fix tasks vs model types in tests * last fix --------- Co-authored-by: Ella Charlaix <[email protected]>
apaniukov · Jul 18, 2024 · b362180 · b362180
1 parent d35ced8
commit b362180
Show file tree

Hide file tree

Showing 8 changed files with 151 additions and 65 deletions.
diff --git a/optimum/commands/export/openvino.py b/optimum/commands/export/openvino.py
@@ -227,13 +227,19 @@ def _get_default_int4_config(model_id_or_path, library_name):
 
             return _DEFAULT_4BIT_CONFIG
 
-        library_name = TasksManager.infer_library_from_model(self.args.model, library_name=self.args.library)
-        if library_name == "sentence_transformers" and self.args.library is None:
-            logger.warning(
-                "Library name is not specified. There are multiple possible variants: `sentence_transformers`, `transformers`."
-                "`transformers` will be selected. If you want to load your model with the `sentence-transformers` library instead, please set --library sentence_transformers"
+        if self.args.library is None:
+            # TODO: add revision, subfolder and token to args
+            library_name = TasksManager._infer_library_from_model_name_or_path(
+                model_name_or_path=self.args.model, cache_dir=self.args.cache_dir
             )
-            library_name = "transformers"
+            if library_name == "sentence_transformers":
+                logger.warning(
+                    "Library name is not specified. There are multiple possible variants: `sentence_transformers`, `transformers`."
+                    "`transformers` will be selected. If you want to load your model with the `sentence-transformers` library instead, please set --library sentence_transformers"
+                )
+                library_name = "transformers"
+        else:
+            library_name = self.args.library
 
         if self.args.weight_format is None:
             ov_config = None

diff --git a/optimum/exporters/openvino/__main__.py b/optimum/exporters/openvino/__main__.py
@@ -47,11 +47,24 @@
 logger = logging.getLogger(__name__)
 
 
-def infer_task(task, model_name_or_path):
+def infer_task(
+    task,
+    model_name_or_path,
+    subfolder: str = "",
+    revision: Optional[str] = None,
+    cache_dir: str = HUGGINGFACE_HUB_CACHE,
+    token: Optional[Union[bool, str]] = None,
+):
     task = TasksManager.map_from_synonym(task)
     if task == "auto":
         try:
-            task = TasksManager.infer_task_from_model(model_name_or_path)
+            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}"
@@ -193,19 +206,27 @@ def main_export(
             ov_config = OVConfig(quantization_config=q_config)
 
     original_task = task
-    task = infer_task(task, model_name_or_path)
-    framework = TasksManager.determine_framework(model_name_or_path, subfolder=subfolder, framework=framework)
-    library_name_is_not_provided = library_name is None
-    library_name = TasksManager.infer_library_from_model(
-        model_name_or_path, subfolder=subfolder, library_name=library_name
+    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 library_name == "sentence_transformers" and library_name_is_not_provided:
-        logger.warning(
-            "Library name is not specified. There are multiple possible variants: `sentence_tenasformers`, `transformers`."
-            "`transformers` will be selected. If you want to load your model with the `sentence-transformers` library instead, please set --library sentence_transformers"
+    if library_name is None:
+        library_name = TasksManager._infer_library_from_model_name_or_path(
+            model_name_or_path=model_name_or_path,
+            subfolder=subfolder,
+            revision=revision,
+            cache_dir=cache_dir,
+            token=token,
         )
-        library_name = "transformers"
+        if library_name == "sentence_transformers":
+            logger.warning(
+                "Library name is not specified. There are multiple possible variants: `sentence_tenasformers`, `transformers`."
+                "`transformers` will be selected. If you want to load your model with the `sentence-transformers` library instead, please set --library sentence_transformers"
+            )
+            library_name = "transformers"
 
     do_gptq_patching = False
     custom_architecture = False
@@ -317,9 +338,7 @@ class StoreAttr(object):
                 )
             model.config.pad_token_id = pad_token_id
 
-    if "stable-diffusion" in task:
-        model_type = "stable-diffusion"
-    elif hasattr(model.config, "export_model_type"):
+    if hasattr(model.config, "export_model_type"):
         model_type = model.config.export_model_type.replace("_", "-")
     else:
         model_type = model.config.model_type.replace("_", "-")

diff --git a/optimum/exporters/openvino/convert.py b/optimum/exporters/openvino/convert.py
@@ -64,7 +64,7 @@
     from transformers.modeling_utils import PreTrainedModel
 
 if is_diffusers_available():
-    from diffusers import ModelMixin
+    from diffusers import DiffusionPipeline, ModelMixin
 
 if is_tf_available():
     from transformers.modeling_tf_utils import TFPreTrainedModel
@@ -74,7 +74,7 @@
     from optimum.intel.openvino.configuration import OVConfig
 
 
-def _save_model(model, path: str, ov_config: Optional["OVConfig"] = None):
+def _save_model(model, path: str, ov_config: Optional["OVConfig"] = None, library_name: Optional[str] = None):
     compress_to_fp16 = False
 
     if ov_config is not None:
@@ -90,13 +90,12 @@ def _save_model(model, path: str, ov_config: Optional["OVConfig"] = None):
 
         compress_to_fp16 = ov_config.dtype == "fp16"
 
-    library_name = TasksManager.infer_library_from_model(Path(path).parent)
     model = _add_version_info_to_model(model, library_name)
     save_model(model, path, compress_to_fp16)
 
 
 def export(
-    model: Union["PreTrainedModel", "TFPreTrainedModel", "ModelMixin"],
+    model: Union["PreTrainedModel", "TFPreTrainedModel", "ModelMixin", "DiffusionPipeline"],
     config: OnnxConfig,
     output: Path,
     opset: Optional[int] = None,
@@ -139,7 +138,7 @@ def export(
         )
 
     if "diffusers" in str(model.__class__) and not is_diffusers_available():
-        raise ImportError("The pip package `diffusers` is required to export stable diffusion models to ONNX.")
+        raise ImportError("The package `diffusers` is required to export diffusion models to OpenVINO.")
 
     if stateful:
         # This will be checked anyway after the model conversion, but checking it earlier will save time for a user if not suitable version is used
@@ -198,7 +197,19 @@ def export_tensorflow(
     onnx_path = Path(output).with_suffix(".onnx")
     input_names, output_names = export_tensorflow_onnx(model, config, opset, onnx_path)
     ov_model = convert_model(str(onnx_path))
-    _save_model(ov_model, output.parent / output, ov_config=ov_config)
+
+    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)
+
+    _save_model(
+        ov_model,
+        output.parent / output,
+        ov_config=ov_config,
+        library_name=library_name,
+    )
     return input_names, output_names, True
 
 
@@ -251,7 +262,19 @@ def export_pytorch_via_onnx(
     )
     torch.onnx.export = orig_torch_onnx_export
     ov_model = convert_model(str(onnx_output))
-    _save_model(ov_model, output.parent / OV_XML_FILE_NAME if output.suffix != ".xml" else output, ov_config=ov_config)
+
+    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)
+
+    _save_model(
+        ov_model,
+        output.parent / OV_XML_FILE_NAME if output.suffix != ".xml" else output,
+        ov_config=ov_config,
+        library_name=library_name,
+    )
     return input_names, output_names, True
 
 
@@ -413,7 +436,18 @@ def ts_patched_forward(*args, **kwargs):
         if stateful:
             patch_stateful(model.config, ov_model)
 
-        _save_model(ov_model, output, ov_config=ov_config)
+        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)
+
+        _save_model(
+            ov_model,
+            output,
+            ov_config=ov_config,
+            library_name=library_name,
+        )
         clear_class_registry()
         del model
         gc.collect()
@@ -422,7 +456,7 @@ def ts_patched_forward(*args, **kwargs):
 
 def export_models(
     models_and_export_configs: Dict[
-        str, Tuple[Union["PreTrainedModel", "TFPreTrainedModel", "ModelMixin"], "OnnxConfig"]
+        str, Tuple[Union["PreTrainedModel", "TFPreTrainedModel", "ModelMixin", "DiffusionPipeline"], "OnnxConfig"]
     ],
     output_dir: Path,
     opset: Optional[int] = None,
@@ -491,7 +525,7 @@ def export_models(
 
 
 def export_from_model(
-    model: Union["PreTrainedModel", "TFPreTrainedModel"],
+    model: Union["PreTrainedModel", "TFPreTrainedModel", "ModelMixin", "DiffusionPipeline"],
     output: Union[str, Path],
     task: Optional[str] = None,
     ov_config: Optional["OVConfig"] = None,
@@ -505,14 +539,15 @@ def export_from_model(
     trust_remote_code: bool = False,
     **kwargs_shapes,
 ):
+    model_kwargs = model_kwargs or {}
+
     if ov_config is not None and ov_config.quantization_config and not is_nncf_available():
         raise ImportError(
             f"Compression of the weights to {ov_config.quantization_config} requires nncf, please install it with `pip install nncf`"
         )
 
-    model_kwargs = model_kwargs or {}
-    library_name = TasksManager._infer_library_from_model(model)
-    TasksManager.standardize_model_attributes(model, library_name)
+    library_name = TasksManager._infer_library_from_model_or_model_class(model=model)
+    TasksManager.standardize_model_attributes(model)
 
     if hasattr(model.config, "export_model_type"):
         model_type = model.config.export_model_type.replace("_", "-")
@@ -521,7 +556,7 @@ def export_from_model(
 
     custom_architecture = library_name == "transformers" and model_type not in TasksManager._SUPPORTED_MODEL_TYPE
 
-    if task is not None:
+    if task is not None and task != "auto":
         task = TasksManager.map_from_synonym(task)
     else:
         try:

diff --git a/optimum/intel/openvino/utils.py b/optimum/intel/openvino/utils.py
@@ -18,11 +18,12 @@
 import os
 from glob import glob
 from pathlib import Path
-from typing import Tuple, Union
+from typing import Tuple, Type, Union
 
 import numpy as np
 from huggingface_hub import model_info
-from openvino.runtime import Core, Type, properties
+from openvino.runtime import Core, properties
+from openvino.runtime import Type as OVType
 from transformers import AutoTokenizer, CLIPTokenizer, PreTrainedTokenizer, PreTrainedTokenizerFast
 from transformers.onnx.utils import ParameterFormat, compute_serialized_parameters_size
 
@@ -70,19 +71,19 @@
 
 
 STR_TO_OV_TYPE = {
-    "boolean": Type.boolean,
-    "f16": Type.f16,
-    "f32": Type.f32,
-    "f64": Type.f64,
-    "i8": Type.i8,
-    "i16": Type.i16,
-    "i32": Type.i32,
-    "i64": Type.i64,
-    "u8": Type.u8,
-    "u16": Type.u16,
-    "u32": Type.u32,
-    "u64": Type.u64,
-    "bf16": Type.bf16,
+    "boolean": OVType.boolean,
+    "f16": OVType.f16,
+    "f32": OVType.f32,
+    "f64": OVType.f64,
+    "i8": OVType.i8,
+    "i16": OVType.i16,
+    "i32": OVType.i32,
+    "i64": OVType.i64,
+    "u8": OVType.u8,
+    "u16": OVType.u16,
+    "u32": OVType.u32,
+    "u64": OVType.u64,
+    "bf16": OVType.bf16,
 }
 
 
@@ -110,7 +111,7 @@
 }
 
 
-NEED_CONVERT_TO_FAST_TOKENIZER: Tuple[type(PreTrainedTokenizer)] = (CLIPTokenizer,)
+NEED_CONVERT_TO_FAST_TOKENIZER: Tuple[Type[PreTrainedTokenizer]] = (CLIPTokenizer,)
 
 
 def maybe_convert_tokenizer_to_fast(

diff --git a/setup.py b/setup.py
@@ -29,7 +29,7 @@
 INSTALL_REQUIRE = [
     "torch>=1.11",
     "transformers>=4.36.0,<4.43.0",
-    "optimum>=1.21.2,<1.22.0",
+    "optimum@git+https://github.com/huggingface/optimum.git",
     "datasets>=1.4.0",
     "sentencepiece",
     "setuptools",

diff --git a/tests/openvino/test_stable_diffusion.py → tests/openvino/test_diffusion.py b/tests/openvino/test_stable_diffusion.py → tests/openvino/test_diffusion.py