Skip to content

Commit 6fe72b5

Browse files
eaidovaeaidova
committed
apply review comments
1 parent 12abb0e commit 6fe72b5

File tree

6 files changed

+32
-43
lines changed

6 files changed

+32
-43
lines changed

‎optimum/exporters/openvino/model_configs.py

+14-36
Original file line numberDiff line numberDiff line change
@@ -1678,43 +1678,17 @@ def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int
16781678
shape = [self.batch_size, (self.height // 2) * (self.width // 2), self.num_channels * 4]
16791679
return self.random_float_tensor(shape, framework=framework, dtype=float_dtype)
16801680
if input_name == "img_ids":
1681-
return self.prepare_image_ids(framework, int_dtype, float_dtype)
1682-
1683-
return super().generate(input_name, framework, int_dtype, float_dtype)
1684-
1685-
def prepare_image_ids(self, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"):
1686-
img_ids_height = self.height // 2
1687-
img_ids_width = self.width // 2
1688-
if framework == "pt":
1689-
import torch
1690-
1691-
latent_image_ids = torch.zeros(img_ids_height, img_ids_width, 3)
1692-
latent_image_ids[..., 1] = latent_image_ids[..., 1] + torch.arange(img_ids_height)[:, None]
1693-
latent_image_ids[..., 2] = latent_image_ids[..., 2] + torch.arange(img_ids_width)[None, :]
1694-
1695-
latent_image_id_height, latent_image_id_width, latent_image_id_channels = latent_image_ids.shape
1696-
1697-
latent_image_ids = latent_image_ids[None, :].repeat(self.batch_size, 1, 1, 1)
1698-
latent_image_ids = latent_image_ids.reshape(
1699-
self.batch_size, latent_image_id_height * latent_image_id_width, latent_image_id_channels
1681+
img_ids_height = self.height // 2
1682+
img_ids_width = self.width // 2
1683+
return self.random_int_tensor(
1684+
[self.batch_size, img_ids_height * img_ids_width, 3],
1685+
min_value=0,
1686+
max_value=min(img_ids_height, img_ids_width),
1687+
framework=framework,
1688+
dtype=float_dtype,
17001689
)
1701-
latent_image_ids.to(DTYPE_MAPPER.pt(float_dtype))
1702-
return latent_image_ids
1703-
if framework == "np":
1704-
import numpy as np
17051690

1706-
latent_image_ids = np.zeros(img_ids_height, img_ids_width, 3)
1707-
latent_image_ids[..., 1] = latent_image_ids[..., 1] + np.arange(img_ids_height)[:, None]
1708-
latent_image_ids[..., 2] = latent_image_ids[..., 2] + np.arange(img_ids_width)[None, :]
1709-
1710-
latent_image_id_height, latent_image_id_width, latent_image_id_channels = latent_image_ids.shape
1711-
1712-
latent_image_ids = np.tile(latent_image_ids[None, :], (self.batch_size, 1, 1, 1))
1713-
latent_image_ids = latent_image_ids.reshape(
1714-
self.batch_size, latent_image_id_height * latent_image_id_width, latent_image_id_channels
1715-
)
1716-
latent_image_ids.astype(DTYPE_MAPPER.np[float_dtype])
1717-
return latent_image_ids
1691+
return super().generate(input_name, framework, int_dtype, float_dtype)
17181692

17191693

17201694
class DummyFluxTextInputGenerator(DummySeq2SeqDecoderTextInputGenerator):
@@ -1728,7 +1702,11 @@ class DummyFluxTextInputGenerator(DummySeq2SeqDecoderTextInputGenerator):
17281702

17291703
def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"):
17301704
if input_name == "txt_ids":
1731-
return self.constant_tensor([self.batch_size, self.sequence_length, 3], 0, DTYPE_MAPPER.pt(float_dtype))
1705+
import torch
1706+
1707+
shape = [self.batch_size, self.sequence_length, 3]
1708+
dtype = DTYPE_MAPPER.pt(float_dtype)
1709+
return torch.full(shape, 0, dtype=dtype)
17321710
return super().generate(input_name, framework, int_dtype, float_dtype)
17331711

17341712

‎optimum/intel/openvino/utils.py

+1
Original file line numberDiff line numberDiff line change
@@ -116,6 +116,7 @@
116116
"stable-diffusion": "OVStableDiffusionPipeline",
117117
"stable-diffusion-xl": "OVStableDiffusionXLPipeline",
118118
"stable-diffusion-3": "OVStableDiffusion3Pipeline",
119+
"flux": "OVFluxPipeline",
119120
"pix2struct": "OVModelForPix2Struct",
120121
"latent-consistency": "OVLatentConsistencyModelPipeline",
121122
"open_clip_text": "OVModelOpenCLIPText",

‎tests/openvino/test_export.py

+5-2
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
from optimum.exporters.openvino import export_from_model, main_export
2929
from optimum.exporters.tasks import TasksManager
3030
from optimum.intel import (
31+
OVFluxPipeline,
3132
OVLatentConsistencyModelPipeline,
3233
OVModelForAudioClassification,
3334
OVModelForCausalLM,
@@ -47,7 +48,7 @@
4748
OVStableDiffusionXLPipeline,
4849
)
4950
from optimum.intel.openvino.modeling_base import OVBaseModel
50-
from optimum.intel.utils.import_utils import _transformers_version
51+
from optimum.intel.utils.import_utils import _transformers_version, is_transformers_version
5152
from optimum.utils.save_utils import maybe_load_preprocessors
5253

5354

@@ -69,9 +70,11 @@ class ExportModelTest(unittest.TestCase):
6970
"stable-diffusion-xl": OVStableDiffusionXLPipeline,
7071
"stable-diffusion-xl-refiner": OVStableDiffusionXLImg2ImgPipeline,
7172
"latent-consistency": OVLatentConsistencyModelPipeline,
72-
"stable-diffusion-3": OVStableDiffusion3Pipeline,
7373
}
7474

75+
if is_transformers_version(">=", "4.45"):
76+
SUPPORTED_ARCHITECTURES.update({"stable-diffusion-3": OVStableDiffusion3Pipeline, "flux": OVFluxPipeline})
77+
7578
GENERATIVE_MODELS = ("pix2struct", "t5", "bart", "gpt2", "whisper")
7679

7780
def _openvino_export(

‎tests/openvino/test_exporters_cli.py

+5-2
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626

2727
from optimum.exporters.openvino.__main__ import main_export
2828
from optimum.intel import ( # noqa
29+
OVFluxPipeline,
2930
OVLatentConsistencyModelPipeline,
3031
OVModelForAudioClassification,
3132
OVModelForCausalLM,
@@ -77,7 +78,7 @@ class OVCLIExportTestCase(unittest.TestCase):
7778
]
7879

7980
if is_transformers_version(">=", "4.45"):
80-
SUPPORTED_ARCHITECTURES.append(("text-to-image", "stable-diffusion-3"))
81+
SUPPORTED_ARCHITECTURES.extend([("text-to-image", "stable-diffusion-3"), ("text-to-image", "flux")])
8182
EXPECTED_NUMBER_OF_TOKENIZER_MODELS = {
8283
"gpt2": 2 if is_tokenizers_version("<", "0.20") else 0,
8384
"t5": 0, # no .model file in the repository
@@ -91,6 +92,7 @@ class OVCLIExportTestCase(unittest.TestCase):
9192
"stable-diffusion": 2 if is_tokenizers_version("<", "0.20") else 0,
9293
"stable-diffusion-xl": 4 if is_tokenizers_version("<", "0.20") else 0,
9394
"stable-diffusion-3": 6 if is_tokenizers_version("<", "0.20") else 0,
95+
"flux": 4 if is_tokenizers_version("<", "0.20") else 0,
9496
}
9597

9698
SUPPORTED_SD_HYBRID_ARCHITECTURES = [
@@ -101,6 +103,7 @@ class OVCLIExportTestCase(unittest.TestCase):
101103

102104
if is_transformers_version(">=", "4.45"):
103105
SUPPORTED_SD_HYBRID_ARCHITECTURES.append(("stable-diffusion-3", 9, 65))
106+
SUPPORTED_SD_HYBRID_ARCHITECTURES.append(("flux", 9, 65))
104107

105108
TEST_4BIT_CONFIGURATONS = [
106109
("text-generation-with-past", "opt125m", "int4 --sym --group-size 128", {"int8": 4, "int4": 72}),
@@ -218,7 +221,7 @@ def test_exporters_cli_int8(self, task: str, model_type: str):
218221
models = [model.encoder, model.decoder]
219222
if task.endswith("with-past"):
220223
models.append(model.decoder_with_past)
221-
elif model_type.startswith("stable-diffusion"):
224+
elif model_type.startswith("stable-diffusion") or model_type.startswith("flux"):
222225
models = [model.unet or model.transformer, model.vae_encoder, model.vae_decoder]
223226
models.append(model.text_encoder if model_type == "stable-diffusion" else model.text_encoder_2)
224227
else:

‎tests/openvino/test_quantization.py

+5-2
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,7 @@
5858
OVStableDiffusionPipeline,
5959
OVStableDiffusionXLPipeline,
6060
OVStableDiffusion3Pipeline,
61+
OVFluxPipeline,
6162
OVQuantizer,
6263
OVTrainer,
6364
OVQuantizationConfig,
@@ -308,8 +309,10 @@ class OVWeightCompressionTest(unittest.TestCase):
308309
]
309310

310311
if is_transformers_version(">=", "4.45.0"):
311-
SUPPORTED_ARCHITECTURES_WITH_HYBRID_QUANTIZATION.append(
312-
(OVStableDiffusion3Pipeline, "stable-diffusion-3", 9, 65)
312+
SUPPORTED_ARCHITECTURES_WITH_HYBRID_QUANTIZATION.extend(
313+
[
314+
(OVStableDiffusion3Pipeline, "stable-diffusion-3", 9, 65),
315+
]
313316
)
314317

315318
IS_SUPPORT_STATEFUL = is_openvino_version(">=", "2023.3")

‎tests/openvino/utils_tests.py

+2-1
Original file line numberDiff line numberDiff line change
@@ -172,7 +172,8 @@
172172
"stable-diffusion-xl": (366, 34, 42, 66),
173173
"stable-diffusion-xl-refiner": (366, 34, 42, 66),
174174
"open-clip": (20, 28),
175-
"stable-diffusion-3": (366, 34, 42, 66),
175+
"stable-diffusion-3": (66, 42, 58, 30),
176+
"flux": (56, 24, 28, 64),
176177
}
177178

178179

0 commit comments

Comments
 (0)