Skip to content

Commit eb1d06d

Browse files
committed
add test
1 parent b963c54 commit eb1d06d

File tree

2 files changed

+18
-1
lines changed

2 files changed

+18
-1
lines changed

optimum/intel/openvino/modeling_diffusion.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -972,7 +972,7 @@ def __call__(self, *args, height=None, width=None, **kwargs):
972972

973973
for k, v in kwargs.items():
974974
kwargs[k] = np_to_pt_generators(v, self.device)
975-
975+
976976
if self.height != -1:
977977
if height is not None and height != self.height:
978978
logger.warning(f"Incompatible height argument provided {height}. Pipeline only support {self.height}.")

tests/openvino/test_diffusion.py

+17
Original file line numberDiff line numberDiff line change
@@ -438,6 +438,23 @@ def test_load_custom_weight_variant(self):
438438

439439
np.testing.assert_allclose(ov_images, diffusers_images, atol=1e-4, rtol=1e-2)
440440

441+
@parameterized.expand(SUPPORTED_ARCHITECTURES)
442+
@require_diffusers
443+
def test_static_shape_image_generation(self, model_arch):
444+
pipeline = self.OVMODEL_CLASS.from_pretrained(MODEL_NAMES[model_arch], compile=False)
445+
pipeline.reshape(height=40, width=32)
446+
pipeline.compile()
447+
# generation with incompatible size
448+
height, width, batch_size = 64, 64, 1
449+
inputs = self.generate_inputs(height=height, width=width, batch_size=batch_size)
450+
image = pipeline(**inputs, num_inference_steps=2).images[0]
451+
self.assertTupleEqual(image.size, (32, 40))
452+
# generation without height / width provided
453+
inputs.pop("height")
454+
inputs.pop("width")
455+
image = pipeline(**inputs, num_inference_steps=2).images[0]
456+
self.assertTupleEqual(image.size, (32, 40))
457+
441458

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

0 commit comments

Comments
 (0)