Merge branch 'main' into ea/fix_win_compression

huggingface · Oct 9, 2024 · 3d521c0 · 3d521c0
2 parents c66f866 + 41d93a1
commit 3d521c0
Show file tree

Hide file tree

Showing 10 changed files with 1,469 additions and 1,166 deletions.
diff --git a/optimum/exporters/openvino/model_configs.py b/optimum/exporters/openvino/model_configs.py
@@ -33,9 +33,6 @@
     MistralOnnxConfig,
     MPTOnnxConfig,
     PhiOnnxConfig,
-    UNetOnnxConfig,
-    VaeDecoderOnnxConfig,
-    VaeEncoderOnnxConfig,
     VisionOnnxConfig,
 )
 from optimum.exporters.tasks import TasksManager
@@ -119,12 +116,12 @@ def init_model_configs():
 
 
 if TYPE_CHECKING:
-    from transformers.modeling_utils import PreTrainedModel
+    from transformers.modeling_utils import PreTrainedModel  # noqa: F811
 
-    from optimum.exporters.onnx.model_patcher import ModelPatcher
+    from optimum.exporters.onnx.model_patcher import ModelPatcher  # noqa: F811
 
     if is_tf_available():
-        from transformers.modeling_tf_utils import TFPreTrainedModel
+        from transformers.modeling_tf_utils import TFPreTrainedModel  # noqa: F811
 
 
 register_in_tasks_manager = TasksManager.create_register("openvino", overwrite_existing=True)
@@ -675,62 +672,6 @@ def patch_model_for_export(
         return FalconModelPatcher(self, model, model_kwargs=model_kwargs)
 
 
-@register_in_tasks_manager("unet", *["semantic-segmentation"], library_name="diffusers")
-class UNetOpenVINOConfig(UNetOnnxConfig):
-    @property
-    def inputs(self) -> Dict[str, Dict[int, str]]:
-        common_inputs = {
-            "sample": {0: "batch_size", 2: "height", 3: "width"},
-            "timestep": {0: "steps"},
-            "encoder_hidden_states": {0: "batch_size", 1: "sequence_length"},
-        }
-
-        # TODO : add text_image, image and image_embeds
-        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"},
-        }
-
-
-@register_in_tasks_manager("vae-encoder", *["semantic-segmentation"], library_name="diffusers")
-class VaeEncoderOpenVINOConfig(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]]:
-        return {
-            "latent_sample": {0: "batch_size", 2: "height_latent", 3: "width_latent"},
-        }
-
-
-@register_in_tasks_manager("vae-decoder", *["semantic-segmentation"], library_name="diffusers")
-class VaeDecoderOpenVINOConfig(VaeDecoderOnnxConfig):
-    @property
-    def inputs(self) -> Dict[str, Dict[int, str]]:
-        return {
-            "latent_sample": {0: "batch_size", 2: "height_latent", 3: "width_latent"},
-        }
-
-    @property
-    def outputs(self) -> Dict[str, Dict[int, str]]:
-        return {
-            "sample": {0: "batch_size", 2: "height", 3: "width"},
-        }
-
-
 @register_in_tasks_manager(
     "persimmon",
     *[

diff --git a/optimum/intel/__init__.py b/optimum/intel/__init__.py
@@ -99,7 +99,13 @@
         "OVStableDiffusionInpaintPipeline",
         "OVStableDiffusionXLPipeline",
         "OVStableDiffusionXLImg2ImgPipeline",
+        "OVStableDiffusionXLInpaintPipeline",
         "OVLatentConsistencyModelPipeline",
+        "OVLatentConsistencyModelImg2ImgPipeline",
+        "OVPipelineForImage2Image",
+        "OVPipelineForText2Image",
+        "OVPipelineForInpainting",
+        "OVDiffusionPipeline",
     ]
 else:
     _import_structure["openvino"].extend(
@@ -109,7 +115,13 @@
             "OVStableDiffusionInpaintPipeline",
             "OVStableDiffusionXLPipeline",
             "OVStableDiffusionXLImg2ImgPipeline",
+            "OVStableDiffusionXLInpaintPipeline",
             "OVLatentConsistencyModelPipeline",
+            "OVLatentConsistencyModelImg2ImgPipeline",
+            "OVPipelineForImage2Image",
+            "OVPipelineForText2Image",
+            "OVPipelineForInpainting",
+            "OVDiffusionPipeline",
         ]
     )
 
@@ -250,7 +262,11 @@
             raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
         from .utils.dummy_openvino_and_diffusers_objects import (
+            OVDiffusionPipeline,
             OVLatentConsistencyModelPipeline,
+            OVPipelineForImage2Image,
+            OVPipelineForInpainting,
+            OVPipelineForText2Image,
             OVStableDiffusionImg2ImgPipeline,
             OVStableDiffusionInpaintPipeline,
             OVStableDiffusionPipeline,
@@ -259,11 +275,17 @@
         )
     else:
         from .openvino import (
+            OVDiffusionPipeline,
+            OVLatentConsistencyModelImg2ImgPipeline,
             OVLatentConsistencyModelPipeline,
+            OVPipelineForImage2Image,
+            OVPipelineForInpainting,
+            OVPipelineForText2Image,
             OVStableDiffusionImg2ImgPipeline,
             OVStableDiffusionInpaintPipeline,
             OVStableDiffusionPipeline,
             OVStableDiffusionXLImg2ImgPipeline,
+            OVStableDiffusionXLInpaintPipeline,
             OVStableDiffusionXLPipeline,
         )
 

diff --git a/optimum/intel/openvino/__init__.py b/optimum/intel/openvino/__init__.py
@@ -81,11 +81,17 @@
 
 if is_diffusers_available():
     from .modeling_diffusion import (
+        OVDiffusionPipeline,
+        OVLatentConsistencyModelImg2ImgPipeline,
         OVLatentConsistencyModelPipeline,
+        OVPipelineForImage2Image,
+        OVPipelineForInpainting,
+        OVPipelineForText2Image,
         OVStableDiffusionImg2ImgPipeline,
         OVStableDiffusionInpaintPipeline,
         OVStableDiffusionPipeline,
         OVStableDiffusionXLImg2ImgPipeline,
+        OVStableDiffusionXLInpaintPipeline,
         OVStableDiffusionXLPipeline,
     )