add export tests

Author: eaidova

optimum/exporters/openvino/convert.py

@@ -919,7 +919,12 @@ def get_diffusion_models_for_export_ext(
 
     if not is_sd3:
         return None, get_diffusion_models_for_export(pipeline, int_dtype, float_dtype, exporter)
+    models_for_export = get_sd3_models_for_export(pipeline, exporter, int_dtype, float_dtype)
 
+    return None, models_for_export
+
+
+def get_sd3_models_for_export(pipeline, exporter, int_dtype, float_dtype):
     models_for_export = {}
 
     # Text encoder
@@ -948,7 +953,7 @@ def get_diffusion_models_for_export_ext(
         exporter=exporter,
         library_name="diffusers",
         task="semantic-segmentation",
-        model_type="transformer",
+        model_type="sd3-transformer",
     )
     transformer_export_config = export_config_constructor(
         pipeline.transformer.config, int_dtype=int_dtype, float_dtype=float_dtype
@@ -1015,4 +1020,4 @@ def get_diffusion_models_for_export_ext(
         )
         models_for_export["text_encoder_3"] = (text_encoder_3, export_config)
 
-    return None, models_for_export
+    return models_for_export

optimum/exporters/openvino/model_configs.py

@@ -1550,7 +1550,7 @@ def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int
         return self.random_float_tensor(shape, framework=framework, dtype=float_dtype)
 
 
-@register_in_tasks_manager("transformer", *["semantic-segmentation"], library_name="diffusers")
+@register_in_tasks_manager("sd3-transformer", *["semantic-segmentation"], library_name="diffusers")
 class TransformerOpenVINOConfig(UNetOnnxConfig):
     DUMMY_INPUT_GENERATOR_CLASSES = UNetOnnxConfig.DUMMY_INPUT_GENERATOR_CLASSES + (
         PooledProjectionsDummyInputGenerator,

optimum/intel/__init__.py

@@ -100,6 +100,9 @@
         "OVStableDiffusionXLPipeline",
         "OVStableDiffusionXLImg2ImgPipeline",
         "OVStableDiffusionXLInpaintPipeline",
+        "OVStableDiffusion3Pipeline",
+        "OVStableDiffusion3Image2ImagePipeline",
+        "OVStableDiffusion3InpaintPipeline",
         "OVLatentConsistencyModelPipeline",
         "OVLatentConsistencyModelImg2ImgPipeline",
         "OVPipelineForImage2Image",
@@ -116,6 +119,9 @@
             "OVStableDiffusionXLPipeline",
             "OVStableDiffusionXLImg2ImgPipeline",
             "OVStableDiffusionXLInpaintPipeline",
+            "OVStableDiffusion3Pipeline",
+            "OVStableDiffusion3Image2ImagePipeline",
+            "OVStableDiffusion3InpaintPipeline",
             "OVLatentConsistencyModelPipeline",
             "OVLatentConsistencyModelImg2ImgPipeline",
             "OVPipelineForImage2Image",

tests/openvino/test_export.py

@@ -41,6 +41,7 @@
     OVModelForSequenceClassification,
     OVModelForSpeechSeq2Seq,
     OVModelForTokenClassification,
+    OVStableDiffusion3Pipeline,
     OVStableDiffusionPipeline,
     OVStableDiffusionXLImg2ImgPipeline,
     OVStableDiffusionXLPipeline,
@@ -68,6 +69,7 @@ class ExportModelTest(unittest.TestCase):
         "stable-diffusion-xl": OVStableDiffusionXLPipeline,
         "stable-diffusion-xl-refiner": OVStableDiffusionXLImg2ImgPipeline,
         "latent-consistency": OVLatentConsistencyModelPipeline,
+        "stable-diffusion-3": OVStableDiffusion3Pipeline,
     }
 
     GENERATIVE_MODELS = ("pix2struct", "t5", "bart", "gpt2", "whisper")

tests/openvino/test_exporters_cli.py

@@ -71,6 +71,7 @@ class OVCLIExportTestCase(unittest.TestCase):
         ("feature-extraction", "blenderbot"),
         ("text-to-image", "stable-diffusion"),
         ("text-to-image", "stable-diffusion-xl"),
+        ("text-to-image", "stable-diffusion-3"),
         ("image-to-image", "stable-diffusion-xl-refiner"),
     )
     EXPECTED_NUMBER_OF_TOKENIZER_MODELS = {
@@ -85,6 +86,7 @@ class OVCLIExportTestCase(unittest.TestCase):
         "blenderbot": 2 if is_tokenizers_version("<", "0.20") else 0,
         "stable-diffusion": 2 if is_tokenizers_version("<", "0.20") else 0,
         "stable-diffusion-xl": 4 if is_tokenizers_version("<", "0.20") else 0,
+        "stable-diffusion-3": 6 if is_tokenizers_version("<", "0.20") else 0,
     }
 
     SUPPORTED_SD_HYBRID_ARCHITECTURES = (