diff --git a/tests/openvino/test_diffusion.py b/tests/openvino/test_diffusion.py index bb808ecc0..477799345 100644 --- a/tests/openvino/test_diffusion.py +++ b/tests/openvino/test_diffusion.py @@ -667,13 +667,14 @@ class OVPipelineForInpaintingTest(unittest.TestCase): if is_transformers_version(">=", "4.40.0"): SUPPORTED_ARCHITECTURES.append("stable-diffusion-3") SUPPORTED_ARCHITECTURES.append("flux") + SUPPORTED_ARCHITECTURES.append("flux-fill") AUTOMODEL_CLASS = AutoPipelineForInpainting OVMODEL_CLASS = OVPipelineForInpainting TASK = "inpainting" - def generate_inputs(self, height=128, width=128, batch_size=1, channel=3, input_type="pil"): + def generate_inputs(self, height=128, width=128, batch_size=1, channel=3, input_type="pil", model_arch=""): inputs = _generate_prompts(batch_size=batch_size) inputs["image"] = _generate_images( @@ -683,7 +684,8 @@ def generate_inputs(self, height=128, width=128, batch_size=1, channel=3, input_ height=height, width=width, batch_size=batch_size, channel=1, input_type=input_type ) - inputs["strength"] = 0.75 + if model_arch != "flux-fill": + inputs["strength"] = 0.75 inputs["height"] = height inputs["width"] = width @@ -699,7 +701,12 @@ def test_load_vanilla_model_which_is_not_supported(self): @parameterized.expand(SUPPORTED_ARCHITECTURES) @require_diffusers def test_ov_pipeline_class_dispatch(self, model_arch: str): - auto_pipeline = self.AUTOMODEL_CLASS.from_pretrained(MODEL_NAMES[model_arch]) + if model_arch != "flux-fill": + auto_pipeline = self.AUTOMODEL_CLASS.from_pretrained(MODEL_NAMES[model_arch]) + else: + from diffusers import FluxFillPipeline + + auto_pipeline = FluxFillPipeline.from_pretrained(MODEL_NAMES[model_arch]) ov_pipeline = self.OVMODEL_CLASS.from_pretrained(MODEL_NAMES[model_arch]) self.assertEqual(ov_pipeline.auto_model_class, auto_pipeline.__class__) @@ -713,7 +720,9 @@ def test_num_images_per_prompt(self, model_arch: str): for height in [64, 128]: for width in [64, 128]: for num_images_per_prompt in [1, 3]: - inputs = self.generate_inputs(height=height, width=width, batch_size=batch_size) + inputs = self.generate_inputs( + height=height, width=width, batch_size=batch_size, model_arch=model_arch + ) outputs = pipeline(**inputs, num_images_per_prompt=num_images_per_prompt).images self.assertEqual(outputs.shape, (batch_size * num_images_per_prompt, height, width, 3)) @@ -752,7 +761,9 @@ def test_shape(self, model_arch: str): height, width, batch_size = 128, 64, 1 for input_type in ["pil", "np", "pt"]: - inputs = self.generate_inputs(height=height, width=width, batch_size=batch_size, input_type=input_type) + inputs = self.generate_inputs( + height=height, width=width, batch_size=batch_size, input_type=input_type, model_arch=model_arch + ) for output_type in ["pil", "np", "pt", "latent"]: inputs["output_type"] = output_type @@ -764,7 +775,7 @@ def test_shape(self, model_arch: str): elif output_type == "pt": self.assertEqual(outputs.shape, (batch_size, 3, height, width)) else: - if model_arch != "flux": + if not model_arch.startswith("flux"): out_channels = ( pipeline.unet.config.out_channels if pipeline.unet is not None @@ -782,17 +793,26 @@ def test_shape(self, model_arch: str): else: packed_height = height // pipeline.vae_scale_factor // 2 packed_width = width // pipeline.vae_scale_factor // 2 - channels = pipeline.transformer.config.in_channels + channels = ( + pipeline.transformer.config.in_channels + if model_arch != "flux-fill" + else pipeline.transformer.out_channels + ) self.assertEqual(outputs.shape, (batch_size, packed_height * packed_width, channels)) @parameterized.expand(SUPPORTED_ARCHITECTURES) @require_diffusers def test_compare_to_diffusers_pipeline(self, model_arch: str): ov_pipeline = self.OVMODEL_CLASS.from_pretrained(MODEL_NAMES[model_arch]) - diffusers_pipeline = self.AUTOMODEL_CLASS.from_pretrained(MODEL_NAMES[model_arch]) + if model_arch != "flux-fill": + diffusers_pipeline = self.AUTOMODEL_CLASS.from_pretrained(MODEL_NAMES[model_arch]) + else: + from diffusers import FluxFillPipeline + + diffusers_pipeline = FluxFillPipeline.from_pretrained(MODEL_NAMES[model_arch]) height, width, batch_size = 64, 64, 1 - inputs = self.generate_inputs(height=height, width=width, batch_size=batch_size) + inputs = self.generate_inputs(height=height, width=width, batch_size=batch_size, model_arch=model_arch) for output_type in ["latent", "np", "pt"]: inputs["output_type"] = output_type @@ -804,7 +824,7 @@ def test_compare_to_diffusers_pipeline(self, model_arch: str): # test generation when input resolution nondevisible on 64 height, width, batch_size = 96, 96, 1 - inputs = self.generate_inputs(height=height, width=width, batch_size=batch_size) + inputs = self.generate_inputs(height=height, width=width, batch_size=batch_size, model_arch=model_arch) for output_type in ["latent", "np", "pt"]: inputs["output_type"] = output_type @@ -820,7 +840,7 @@ def test_image_reproducibility(self, model_arch: str): pipeline = self.OVMODEL_CLASS.from_pretrained(MODEL_NAMES[model_arch]) height, width, batch_size = 64, 64, 1 - inputs = self.generate_inputs(height=height, width=width, batch_size=batch_size) + inputs = self.generate_inputs(height=height, width=width, batch_size=batch_size, model_arch=model_arch) for generator_framework in ["np", "pt"]: ov_outputs_1 = pipeline(**inputs, generator=get_generator(generator_framework, SEED)) diff --git a/tests/openvino/utils_tests.py b/tests/openvino/utils_tests.py index 02b81bfdc..9373e8afa 100644 --- a/tests/openvino/utils_tests.py +++ b/tests/openvino/utils_tests.py @@ -66,6 +66,7 @@ "falcon-40b": "katuni4ka/tiny-random-falcon-40b", "flaubert": "hf-internal-testing/tiny-random-flaubert", "flux": "katuni4ka/tiny-random-flux", + "flux-fill": "katuni4ka/tiny-random-flux-fill", "gpt_bigcode": "hf-internal-testing/tiny-random-GPTBigCodeModel", "gpt2": "hf-internal-testing/tiny-random-gpt2", "gpt_neo": "hf-internal-testing/tiny-random-GPTNeoModel",