Add stable diffusion OpenVINO pipeline (#195)

* Add stable diffusion OpenVINO pipeline * Fix style * Add tests * add diffusers import utils * add test * fix style * Enable model caching * Fix optimum version requirements
huggingface · Mar 2, 2023 · 21c24c5 · 21c24c5
1 parent feba68e
commit 21c24c5
Show file tree

Hide file tree

Showing 8 changed files with 685 additions and 38 deletions.
diff --git a/optimum/intel/openvino/__init__.py b/optimum/intel/openvino/__init__.py
@@ -14,7 +14,7 @@
 
 import importlib.util
 
-from ..utils.import_utils import is_nncf_available
+from ..utils.import_utils import is_diffusers_available, is_nncf_available
 from .utils import OV_DECODER_NAME, OV_DECODER_WITH_PAST_NAME, OV_ENCODER_NAME, OV_XML_FILE_NAME
 
 
@@ -39,3 +39,7 @@
     OVModelForTokenClassification,
 )
 from .modeling_seq2seq import OVModelForSeq2SeqLM
+
+
+if is_diffusers_available():
+    from .modeling_diffusion import OVStableDiffusionPipeline
diff --git a/optimum/intel/openvino/modeling_base.py b/optimum/intel/openvino/modeling_base.py
@@ -104,17 +104,18 @@ def __init__(
             self.generation_config = GenerationConfig.from_model_config(config) if self.can_generate() else None
 
     @staticmethod
-    def load_model(file_name: Union[str, Path], bin_file_name: Optional[Union[str, Path]] = None):
+    def load_model(file_name: Union[str, Path]):
         """
         Loads the model.
 
         Arguments:
             file_name (`str` or `Path`):
                 The path of the model ONNX or XML file.
-            bin_file_name (`str` or `Path`, *optional*):
-                The path of the model binary file, for OpenVINO IR the weights file is expected to be in the same
-                directory as the .xml file if not provided.
         """
+        if isinstance(file_name, str):
+            file_name = Path(file_name)
+        bin_file_name = file_name.with_suffix(".bin") if file_name.suffix == ".xml" else None
+
         return core.read_model(file_name, bin_file_name)
 
     def _save_pretrained(self, save_directory: Union[str, Path], file_name: Optional[str] = None, **kwargs):
@@ -184,8 +185,7 @@ def _from_pretrained(
                     "The file names `ov_model.xml` and `ov_model.bin` will be soon deprecated."
                     "Make sure to rename your file to respectively `openvino_model.xml` and `openvino_model.bin`"
                 )
-            bin_file_name = file_name.replace(".xml", ".bin") if not from_onnx else None
-            model = cls.load_model(file_name, bin_file_name)
+            model = cls.load_model(file_name)
             model_save_dir = model_id
         # Download the model from the hub
         else:
@@ -225,8 +225,7 @@ def _from_pretrained(
                     "Make sure to rename your file to respectively `openvino_model.xml` and `openvino_model.bin`"
                 )
             model_save_dir = Path(model_cache_path).parent
-            bin_file_name = file_names[1] if not from_onnx else None
-            model = cls.load_model(file_names[0], bin_file_name=bin_file_name)
+            model = cls.load_model(file_names[0])
         return cls(model, config=config, model_save_dir=model_save_dir, **kwargs)
 
     @classmethod

diff --git a/optimum/intel/openvino/modeling_base_seq2seq.py b/optimum/intel/openvino/modeling_base_seq2seq.py
@@ -190,22 +190,11 @@ def _from_pretrained(
                     "will be soon deprecated. Make sure to rename your file to respectively `openvino_encoder_model.xml`, "
                     "`openvino_decoder_model.xml` and `openvino_decoder_with_past_model.xml`"
                 )
-            encoder_bin_file_name = (
-                os.path.join(model_id, encoder_file_name.replace(".xml", ".bin")) if not from_onnx else None
-            )
-            decoder_bin_file_name = (
-                os.path.join(model_id, decoder_file_name.replace(".xml", ".bin")) if not from_onnx else None
-            )
-            decoder_with_past_bin_file_name = (
-                os.path.join(model_id, decoder_with_past_file_name.replace(".xml", ".bin")) if not from_onnx else None
-            )
 
-            encoder = cls.load_model(os.path.join(model_id, encoder_file_name), encoder_bin_file_name)
-            decoder = cls.load_model(os.path.join(model_id, decoder_file_name), decoder_bin_file_name)
+            encoder = cls.load_model(os.path.join(model_id, encoder_file_name))
+            decoder = cls.load_model(os.path.join(model_id, decoder_file_name))
             decoder_with_past = (
-                cls.load_model(os.path.join(model_id, decoder_with_past_file_name), decoder_with_past_bin_file_name)
-                if use_cache
-                else None
+                cls.load_model(os.path.join(model_id, decoder_with_past_file_name)) if use_cache else None
             )
             model_save_dir = Path(model_id)
 
@@ -257,14 +246,9 @@ def _from_pretrained(
                 )
 
             model_save_dir = Path(model_cache_path).parent
-            encoder = cls.load_model(file_names["encoder"], bin_file_name=file_names.pop("encoder_bin", None))
-            decoder = cls.load_model(file_names["decoder"], bin_file_name=file_names.pop("decoder_bin", None))
-            if use_cache:
-                decoder_with_past = cls.load_model(
-                    file_names["decoder_with_past"], bin_file_name=file_names.pop("decoder_with_past_bin", None)
-                )
-            else:
-                decoder_with_past = None
+            encoder = cls.load_model(file_names["encoder"])
+            decoder = cls.load_model(file_names["decoder"])
+            decoder_with_past = cls.load_model(file_names["decoder_with_past"]) if use_cache else None
 
         return cls(
             encoder=encoder,