Skip to content

Commit 47fa545

Browse files
committed
fix test
1 parent 429e34c commit 47fa545

File tree

1 file changed

+8
-15
lines changed

1 file changed

+8
-15
lines changed

tests/openvino/test_stable_diffusion.py

+8-15
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,6 @@
2828
from diffusers.utils import load_image
2929
from diffusers.utils.testing_utils import floats_tensor
3030
from openvino.runtime.ie_api import CompiledModel
31-
from packaging.version import Version, parse
3231
from parameterized import parameterized
3332
from utils_tests import MODEL_NAMES, SEED
3433

@@ -46,13 +45,8 @@
4645
OVModelVaeDecoder,
4746
OVModelVaeEncoder,
4847
)
49-
from optimum.onnxruntime import (
50-
ORTStableDiffusionImg2ImgPipeline,
51-
ORTStableDiffusionInpaintPipeline,
52-
ORTStableDiffusionXLImg2ImgPipeline,
53-
ORTStableDiffusionXLPipeline,
54-
)
55-
from optimum.utils.import_utils import _diffusers_version
48+
from optimum.intel.utils.import_utils import is_diffusers_version
49+
from optimum.utils.import_utils import is_onnxruntime_available
5650

5751

5852
F32_CONFIG = {"INFERENCE_PRECISION_HINT": "f32"}
@@ -167,7 +161,6 @@ def generate_inputs(self, height=128, width=128, batch_size=1):
167161
class OVStableDiffusionImg2ImgPipelineTest(OVStableDiffusionPipelineBaseTest):
168162
SUPPORTED_ARCHITECTURES = ("stable-diffusion",)
169163
MODEL_CLASS = OVStableDiffusionImg2ImgPipeline
170-
ORT_MODEL_CLASS = ORTStableDiffusionImg2ImgPipeline
171164
TASK = "image-to-image"
172165

173166
@parameterized.expand(SUPPORTED_ARCHITECTURES)
@@ -298,11 +291,13 @@ def test_height_width_properties(self, model_arch: str):
298291
class OVStableDiffusionInpaintPipelineTest(OVStableDiffusionPipelineBaseTest):
299292
SUPPORTED_ARCHITECTURES = ("stable-diffusion",)
300293
MODEL_CLASS = OVStableDiffusionInpaintPipeline
301-
ORT_MODEL_CLASS = ORTStableDiffusionInpaintPipeline
302294
TASK = "inpaint"
303295

304296
@parameterized.expand(SUPPORTED_ARCHITECTURES)
297+
@unittest.skipIf(not is_onnxruntime_available(), "this test requires onnxruntime")
305298
def test_compare_diffusers_pipeline(self, model_arch: str):
299+
from optimum.onnxruntime import ORTStableDiffusionInpaintPipeline
300+
306301
model_id = MODEL_NAMES[model_arch]
307302
pipeline = self.MODEL_CLASS.from_pretrained(model_id, export=True, ov_config=F32_CONFIG)
308303
batch_size, num_images, height, width = 1, 1, 64, 64
@@ -329,7 +324,7 @@ def test_compare_diffusers_pipeline(self, model_arch: str):
329324
outputs = pipeline(**inputs, latents=latents).images
330325
self.assertEqual(outputs.shape, (batch_size * num_images, height, width, 3))
331326

332-
ort_pipeline = self.ORT_MODEL_CLASS.from_pretrained(model_id, export=True)
327+
ort_pipeline = ORTStableDiffusionInpaintPipeline.from_pretrained(model_id, export=True)
333328
ort_outputs = ort_pipeline(**inputs, latents=latents).images
334329
self.assertTrue(np.allclose(outputs, ort_outputs, atol=1e-1))
335330

@@ -358,7 +353,6 @@ def generate_inputs(self, height=128, width=128, batch_size=1):
358353
class OVtableDiffusionXLPipelineTest(unittest.TestCase):
359354
SUPPORTED_ARCHITECTURES = ("stable-diffusion-xl",)
360355
MODEL_CLASS = OVStableDiffusionXLPipeline
361-
ORT_MODEL_CLASS = ORTStableDiffusionXLPipeline
362356
PT_MODEL_CLASS = StableDiffusionXLPipeline
363357
TASK = "text-to-image"
364358

@@ -444,7 +438,6 @@ def test_num_images_per_prompt_static_model(self, model_arch: str):
444438
class OVStableDiffusionXLImg2ImgPipelineTest(unittest.TestCase):
445439
SUPPORTED_ARCHITECTURES = ("stable-diffusion-xl", "stable-diffusion-xl-refiner")
446440
MODEL_CLASS = OVStableDiffusionXLImg2ImgPipeline
447-
ORT_MODEL_CLASS = ORTStableDiffusionXLImg2ImgPipeline
448441
PT_MODEL_CLASS = StableDiffusionXLImg2ImgPipeline
449442
TASK = "image-to-image"
450443

@@ -489,7 +482,7 @@ class OVLatentConsistencyModelPipelineTest(unittest.TestCase):
489482
TASK = "text-to-image"
490483

491484
@parameterized.expand(SUPPORTED_ARCHITECTURES)
492-
@unittest.skipIf(parse(_diffusers_version) <= Version("0.21.4"), "not supported with this diffusers version")
485+
@unittest.skipIf(is_diffusers_version("<=", "0.21.4"), "not supported with this diffusers version")
493486
def test_compare_to_diffusers(self, model_arch: str):
494487
ov_pipeline = self.MODEL_CLASS.from_pretrained(MODEL_NAMES[model_arch], export=True, ov_config=F32_CONFIG)
495488
self.assertIsInstance(ov_pipeline.text_encoder, OVModelTextEncoder)
@@ -532,7 +525,7 @@ def test_compare_to_diffusers(self, model_arch: str):
532525
self.assertEqual(pipeline.device.type, ov_pipeline.device)
533526

534527
@parameterized.expand(SUPPORTED_ARCHITECTURES)
535-
@unittest.skipIf(parse(_diffusers_version) <= Version("0.21.4"), "not supported with this diffusers version")
528+
@unittest.skipIf(is_diffusers_version("<=", "0.21.4"), "not supported with this diffusers version")
536529
def test_num_images_per_prompt_static_model(self, model_arch: str):
537530
model_id = MODEL_NAMES[model_arch]
538531
pipeline = self.MODEL_CLASS.from_pretrained(model_id, export=True, compile=False, dynamic_shapes=False)

0 commit comments

Comments
 (0)