Skip to content

Commit b27e6f7

Browse files
committed
add tests for variant
1 parent 1760da7 commit b27e6f7

File tree

3 files changed

+28
-0
lines changed

3 files changed

+28
-0
lines changed

optimum/intel/__init__.py

+2
Original file line numberDiff line numberDiff line change
@@ -305,6 +305,7 @@
305305
OVPipelineForImage2Image,
306306
OVPipelineForInpainting,
307307
OVPipelineForText2Image,
308+
OVSanaPipeline,
308309
OVStableDiffusion3Img2ImgPipeline,
309310
OVStableDiffusion3InpaintPipeline,
310311
OVStableDiffusion3Pipeline,
@@ -323,6 +324,7 @@
323324
OVPipelineForImage2Image,
324325
OVPipelineForInpainting,
325326
OVPipelineForText2Image,
327+
OVSanaPipeline,
326328
OVStableDiffusion3Img2ImgPipeline,
327329
OVStableDiffusion3InpaintPipeline,
328330
OVStableDiffusion3Pipeline,

tests/openvino/test_diffusion.py

+15
Original file line numberDiff line numberDiff line change
@@ -416,6 +416,21 @@ def test_textual_inversion(self):
416416

417417
np.testing.assert_allclose(ov_output, diffusers_output, atol=1e-4, rtol=1e-2)
418418

419+
@require_diffusers
420+
def test_load_custom_weight_variant(self):
421+
model_id = "katuni4ka/tiny-stable-diffusion-torch-custom-variant"
422+
diffusers_pipeline = self.AUTOMODEL_CLASS.from_pretrained(model_id, variant="custom")
423+
ov_pipeline = self.OVMODEL_CLASS.from_pretrained(model_id, compile=False, variant="custom")
424+
height, width, batch_size = 32, 64, 1
425+
inputs = self.generate_inputs(height=height, width=width, batch_size=batch_size)
426+
427+
ov_output = ov_pipeline(**inputs, generator=get_generator("pt", SEED))
428+
diffusers_output = diffusers_pipeline(**inputs, generator=get_generator("pt", SEED))
429+
ov_images = ov_output.images
430+
diffusers_images = diffusers_output.images
431+
432+
np.testing.assert_allclose(ov_images, diffusers_images, atol=1e-4, rtol=1e-2)
433+
419434

420435
class OVPipelineForImage2ImageTest(unittest.TestCase):
421436
SUPPORTED_ARCHITECTURES = ["stable-diffusion", "stable-diffusion-xl", "latent-consistency"]

tests/openvino/test_exporters_cli.py

+11
Original file line numberDiff line numberDiff line change
@@ -568,3 +568,14 @@ def test_export_openvino_with_missed_weight_format(self):
568568
"Some compression parameters are provided, but the weight format is not specified.",
569569
str(exc_info.exception.stderr),
570570
)
571+
572+
def test_export_openvino_with_custom_variant(self):
573+
with TemporaryDirectory() as tmpdir:
574+
subprocess.run(
575+
f"optimum-cli export openvino --model katuni4ka/tiny-stable-diffusion-torch-custom-variant --variant custom {tmpdir}",
576+
shell=True,
577+
check=True,
578+
)
579+
model = eval(_HEAD_TO_AUTOMODELS["stable-diffusion"]).from_pretrained(tmpdir, compile=False)
580+
for component in ["text_encoder", "tokenizer", "unet", "vae_encoder", "vae_decoder"]:
581+
self.assertIsNotNone(getattr(model, component))

0 commit comments

Comments
 (0)