add tests

Author: eaidova

optimum/intel/openvino/modeling_diffusion.py

@@ -511,6 +511,7 @@ def _from_transformers(
             no_post_process=True,
             revision=revision,
             cache_dir=cache_dir,
+            task=cls.export_feature,
             token=token,
             local_files_only=local_files_only,
             force_download=force_download,
@@ -1479,13 +1480,16 @@ def from_pretrained(cls, pretrained_model_or_path, **kwargs):
 class OVPipelineForText2Image(OVPipelineForTask):
     auto_model_class = AutoPipelineForText2Image
     ov_pipelines_mapping = OV_TEXT2IMAGE_PIPELINES_MAPPING
+    export_feature = "text-to-image"
 
 
 class OVPipelineForImage2Image(OVPipelineForTask):
     auto_model_class = AutoPipelineForImage2Image
     ov_pipelines_mapping = OV_IMAGE2IMAGE_PIPELINES_MAPPING
+    export_feature = "image-to-image"
 
 
 class OVPipelineForInpainting(OVPipelineForTask):
     auto_model_class = AutoPipelineForInpainting
     ov_pipelines_mapping = OV_INPAINT_PIPELINES_MAPPING
+    export_feature = "inpainting"

tests/openvino/test_diffusion.py

@@ -72,7 +72,7 @@ def _generate_images(height=128, width=128, batch_size=1, channel=3, input_type=
 
 
 class OVPipelineForText2ImageTest(unittest.TestCase):
-    SUPPORTED_ARCHITECTURES = ["stable-diffusion", "stable-diffusion-xl", "latent-consistency"]
+    SUPPORTED_ARCHITECTURES = ["stable-diffusion", "stable-diffusion-xl", "latent-consistency", "stable-diffusion-3"]
 
     OVMODEL_CLASS = OVPipelineForText2Image
     AUTOMODEL_CLASS = AutoPipelineForText2Image
@@ -323,7 +323,7 @@ def test_textual_inversion(self):
 
 
 class OVPipelineForImage2ImageTest(unittest.TestCase):
-    SUPPORTED_ARCHITECTURES = ["stable-diffusion", "stable-diffusion-xl", "latent-consistency"]
+    SUPPORTED_ARCHITECTURES = ["stable-diffusion", "stable-diffusion-xl", "latent-consistency", "stable-diffusion-3"]
 
     AUTOMODEL_CLASS = AutoPipelineForImage2Image
     OVMODEL_CLASS = OVPipelineForImage2Image
@@ -535,7 +535,7 @@ def test_textual_inversion(self):
 
 
 class OVPipelineForInpaintingTest(unittest.TestCase):
-    SUPPORTED_ARCHITECTURES = ["stable-diffusion", "stable-diffusion-xl"]
+    SUPPORTED_ARCHITECTURES = ["stable-diffusion", "stable-diffusion-xl", "stable-diffusion-3"]
 
     AUTOMODEL_CLASS = AutoPipelineForInpainting
     OVMODEL_CLASS = OVPipelineForInpainting

tests/openvino/utils_tests.py

@@ -118,6 +118,7 @@
     "stable-diffusion-openvino": "hf-internal-testing/tiny-stable-diffusion-openvino",
     "stable-diffusion-xl": "echarlaix/tiny-random-stable-diffusion-xl",
     "stable-diffusion-xl-refiner": "echarlaix/tiny-random-stable-diffusion-xl-refiner",
+    "stable-diffusion-3": "yujiepan/stable-diffusion-3-tiny-random",
     "stablelm": "hf-internal-testing/tiny-random-StableLmForCausalLM",
     "starcoder2": "hf-internal-testing/tiny-random-Starcoder2ForCausalLM",
     "latent-consistency": "echarlaix/tiny-random-latent-consistency",