apply review comments

optimum/commands/export/openvino.py

@@ -322,10 +322,6 @@ def run(self):
                 from optimum.intel import OVStableDiffusion3Pipeline
 
                 model_cls = OVStableDiffusion3Pipeline
-            elif class_name == "FluxPipeline":
-                from optimum.intel import OVFluxPipeline
-
-                model_cls = OVFluxPipeline
             else:
                 raise NotImplementedError(f"Quantization in hybrid mode isn't supported for class {class_name}.")
 

optimum/exporters/openvino/model_configs.py

@@ -1678,43 +1678,17 @@ def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int
             shape = [self.batch_size, (self.height // 2) * (self.width // 2), self.num_channels * 4]
             return self.random_float_tensor(shape, framework=framework, dtype=float_dtype)
         if input_name == "img_ids":
-            return self.prepare_image_ids(framework, int_dtype, float_dtype)
-
-        return super().generate(input_name, framework, int_dtype, float_dtype)
-
-    def prepare_image_ids(self, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"):
-        img_ids_height = self.height // 2
-        img_ids_width = self.width // 2
-        if framework == "pt":
-            import torch
-
-            latent_image_ids = torch.zeros(img_ids_height, img_ids_width, 3)
-            latent_image_ids[..., 1] = latent_image_ids[..., 1] + torch.arange(img_ids_height)[:, None]
-            latent_image_ids[..., 2] = latent_image_ids[..., 2] + torch.arange(img_ids_width)[None, :]
-
-            latent_image_id_height, latent_image_id_width, latent_image_id_channels = latent_image_ids.shape
-
-            latent_image_ids = latent_image_ids[None, :].repeat(self.batch_size, 1, 1, 1)
-            latent_image_ids = latent_image_ids.reshape(
-                self.batch_size, latent_image_id_height * latent_image_id_width, latent_image_id_channels
+            img_ids_height = self.height // 2
+            img_ids_width = self.width // 2
+            return self.random_int_tensor(
+                [self.batch_size, img_ids_height * img_ids_width, 3],
+                min_value=0,
+                max_value=min(img_ids_height, img_ids_width),
+                framework=framework,
+                dtype=float_dtype,
             )
-            latent_image_ids.to(DTYPE_MAPPER.pt(float_dtype))
-            return latent_image_ids
-        if framework == "np":
-            import numpy as np
 
-            latent_image_ids = np.zeros(img_ids_height, img_ids_width, 3)
-            latent_image_ids[..., 1] = latent_image_ids[..., 1] + np.arange(img_ids_height)[:, None]
-            latent_image_ids[..., 2] = latent_image_ids[..., 2] + np.arange(img_ids_width)[None, :]
-
-            latent_image_id_height, latent_image_id_width, latent_image_id_channels = latent_image_ids.shape
-
-            latent_image_ids = np.tile(latent_image_ids[None, :], (self.batch_size, 1, 1, 1))
-            latent_image_ids = latent_image_ids.reshape(
-                self.batch_size, latent_image_id_height * latent_image_id_width, latent_image_id_channels
-            )
-            latent_image_ids.astype(DTYPE_MAPPER.np[float_dtype])
-            return latent_image_ids
+        return super().generate(input_name, framework, int_dtype, float_dtype)
 
 
 class DummyFluxTextInputGenerator(DummySeq2SeqDecoderTextInputGenerator):
@@ -1728,7 +1702,11 @@ class DummyFluxTextInputGenerator(DummySeq2SeqDecoderTextInputGenerator):
 
     def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"):
         if input_name == "txt_ids":
-            return self.constant_tensor([self.batch_size, self.sequence_length, 3], 0, DTYPE_MAPPER.pt(float_dtype))
+            import torch
+
+            shape = [self.batch_size, self.sequence_length, 3]
+            dtype = DTYPE_MAPPER.pt(float_dtype)
+            return torch.full(shape, 0, dtype=dtype)
         return super().generate(input_name, framework, int_dtype, float_dtype)
 
 

optimum/intel/openvino/utils.py

@@ -116,6 +116,7 @@
     "stable-diffusion": "OVStableDiffusionPipeline",
     "stable-diffusion-xl": "OVStableDiffusionXLPipeline",
     "stable-diffusion-3": "OVStableDiffusion3Pipeline",
+    "flux": "OVFluxPipeline",
     "pix2struct": "OVModelForPix2Struct",
     "latent-consistency": "OVLatentConsistencyModelPipeline",
     "open_clip_text": "OVModelOpenCLIPText",

tests/openvino/test_export.py

@@ -28,6 +28,7 @@
 from optimum.exporters.openvino import export_from_model, main_export
 from optimum.exporters.tasks import TasksManager
 from optimum.intel import (
+    OVFluxPipeline,
     OVLatentConsistencyModelPipeline,
     OVModelForAudioClassification,
     OVModelForCausalLM,
@@ -47,7 +48,7 @@
     OVStableDiffusionXLPipeline,
 )
 from optimum.intel.openvino.modeling_base import OVBaseModel
-from optimum.intel.utils.import_utils import _transformers_version
+from optimum.intel.utils.import_utils import _transformers_version, is_transformers_version
 from optimum.utils.save_utils import maybe_load_preprocessors
 
 
@@ -69,9 +70,11 @@ class ExportModelTest(unittest.TestCase):
         "stable-diffusion-xl": OVStableDiffusionXLPipeline,
         "stable-diffusion-xl-refiner": OVStableDiffusionXLImg2ImgPipeline,
         "latent-consistency": OVLatentConsistencyModelPipeline,
-        "stable-diffusion-3": OVStableDiffusion3Pipeline,
     }
 
+    if is_transformers_version(">=", "4.45"):
+        SUPPORTED_ARCHITECTURES.update({"stable-diffusion-3": OVStableDiffusion3Pipeline, "flux": OVFluxPipeline})
+
     GENERATIVE_MODELS = ("pix2struct", "t5", "bart", "gpt2", "whisper")
 
     def _openvino_export(

tests/openvino/test_exporters_cli.py

@@ -26,6 +26,7 @@
 
 from optimum.exporters.openvino.__main__ import main_export
 from optimum.intel import (  # noqa
+    OVFluxPipeline,
     OVLatentConsistencyModelPipeline,
     OVModelForAudioClassification,
     OVModelForCausalLM,
@@ -77,7 +78,7 @@ class OVCLIExportTestCase(unittest.TestCase):
     ]
 
     if is_transformers_version(">=", "4.45"):
-        SUPPORTED_ARCHITECTURES.append(("text-to-image", "stable-diffusion-3"))
+        SUPPORTED_ARCHITECTURES.extend([("text-to-image", "stable-diffusion-3"), ("text-to-image", "flux")])
     EXPECTED_NUMBER_OF_TOKENIZER_MODELS = {
         "gpt2": 2 if is_tokenizers_version("<", "0.20") else 0,
         "t5": 0,  # no .model file in the repository
@@ -90,7 +91,8 @@ class OVCLIExportTestCase(unittest.TestCase):
         "blenderbot": 2 if is_tokenizers_version("<", "0.20") else 0,
         "stable-diffusion": 2 if is_tokenizers_version("<", "0.20") else 0,
         "stable-diffusion-xl": 4 if is_tokenizers_version("<", "0.20") else 0,
-        "stable-diffusion-3": 6 if is_tokenizers_version("<", "0.20") else 0,
+        "stable-diffusion-3": 6 if is_tokenizers_version("<", "0.20") else 2,
+        "flux": 4 if is_tokenizers_version("<", "0.20") else 0,
     }
 
     SUPPORTED_SD_HYBRID_ARCHITECTURES = [
@@ -218,7 +220,7 @@ def test_exporters_cli_int8(self, task: str, model_type: str):
                 models = [model.encoder, model.decoder]
                 if task.endswith("with-past"):
                     models.append(model.decoder_with_past)
-            elif model_type.startswith("stable-diffusion"):
+            elif model_type.startswith("stable-diffusion") or model_type.startswith("flux"):
                 models = [model.unet or model.transformer, model.vae_encoder, model.vae_decoder]
                 models.append(model.text_encoder if model_type == "stable-diffusion" else model.text_encoder_2)
             else:

tests/openvino/test_quantization.py

@@ -58,6 +58,7 @@
     OVStableDiffusionPipeline,
     OVStableDiffusionXLPipeline,
     OVStableDiffusion3Pipeline,
+    OVFluxPipeline,
     OVQuantizer,
     OVTrainer,
     OVQuantizationConfig,
@@ -308,8 +309,10 @@ class OVWeightCompressionTest(unittest.TestCase):
     ]
 
     if is_transformers_version(">=", "4.45.0"):
-        SUPPORTED_ARCHITECTURES_WITH_HYBRID_QUANTIZATION.append(
-            (OVStableDiffusion3Pipeline, "stable-diffusion-3", 9, 65)
+        SUPPORTED_ARCHITECTURES_WITH_HYBRID_QUANTIZATION.extend(
+            [
+                (OVStableDiffusion3Pipeline, "stable-diffusion-3", 9, 65),
+            ]
         )
 
     IS_SUPPORT_STATEFUL = is_openvino_version(">=", "2023.3")

tests/openvino/utils_tests.py

@@ -172,7 +172,8 @@
     "stable-diffusion-xl": (366, 34, 42, 66),
     "stable-diffusion-xl-refiner": (366, 34, 42, 66),
     "open-clip": (20, 28),
-    "stable-diffusion-3": (366, 34, 42, 66),
+    "stable-diffusion-3": (66, 42, 58, 30),
+    "flux": (56, 24, 28, 64),
 }