add tests for variant

huggingface · Jan 30, 2025 · b27e6f7 · b27e6f7
1 parent 1760da7
commit b27e6f7
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 0 deletions.
diff --git a/optimum/intel/__init__.py b/optimum/intel/__init__.py
@@ -305,6 +305,7 @@
             OVPipelineForImage2Image,
             OVPipelineForInpainting,
             OVPipelineForText2Image,
+            OVSanaPipeline,
             OVStableDiffusion3Img2ImgPipeline,
             OVStableDiffusion3InpaintPipeline,
             OVStableDiffusion3Pipeline,
@@ -323,6 +324,7 @@
             OVPipelineForImage2Image,
             OVPipelineForInpainting,
             OVPipelineForText2Image,
+            OVSanaPipeline,
             OVStableDiffusion3Img2ImgPipeline,
             OVStableDiffusion3InpaintPipeline,
             OVStableDiffusion3Pipeline,

diff --git a/tests/openvino/test_diffusion.py b/tests/openvino/test_diffusion.py
@@ -416,6 +416,21 @@ def test_textual_inversion(self):
 
         np.testing.assert_allclose(ov_output, diffusers_output, atol=1e-4, rtol=1e-2)
 
+    @require_diffusers
+    def test_load_custom_weight_variant(self):
+        model_id = "katuni4ka/tiny-stable-diffusion-torch-custom-variant"
+        diffusers_pipeline = self.AUTOMODEL_CLASS.from_pretrained(model_id, variant="custom")
+        ov_pipeline = self.OVMODEL_CLASS.from_pretrained(model_id, compile=False, variant="custom")
+        height, width, batch_size = 32, 64, 1
+        inputs = self.generate_inputs(height=height, width=width, batch_size=batch_size)
+
+        ov_output = ov_pipeline(**inputs, generator=get_generator("pt", SEED))
+        diffusers_output = diffusers_pipeline(**inputs, generator=get_generator("pt", SEED))
+        ov_images = ov_output.images
+        diffusers_images = diffusers_output.images
+
+        np.testing.assert_allclose(ov_images, diffusers_images, atol=1e-4, rtol=1e-2)
+
 
 class OVPipelineForImage2ImageTest(unittest.TestCase):
     SUPPORTED_ARCHITECTURES = ["stable-diffusion", "stable-diffusion-xl", "latent-consistency"]

diff --git a/tests/openvino/test_exporters_cli.py b/tests/openvino/test_exporters_cli.py
@@ -568,3 +568,14 @@ def test_export_openvino_with_missed_weight_format(self):
                 "Some compression parameters are provided, but the weight format is not specified.",
                 str(exc_info.exception.stderr),
             )
+
+    def test_export_openvino_with_custom_variant(self):
+        with TemporaryDirectory() as tmpdir:
+            subprocess.run(
+                f"optimum-cli export openvino --model katuni4ka/tiny-stable-diffusion-torch-custom-variant --variant custom {tmpdir}",
+                shell=True,
+                check=True,
+            )
+            model = eval(_HEAD_TO_AUTOMODELS["stable-diffusion"]).from_pretrained(tmpdir, compile=False)
+            for component in ["text_encoder", "tokenizer", "unet", "vae_encoder", "vae_decoder"]:
+                self.assertIsNotNone(getattr(model, component))