review fixes

openvinotoolkit · Jan 23, 2025 · 641b090 · 641b090
1 parent a7a60e8
commit 641b090
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 13 deletions.
diff --git a/src/cpp/src/image_generation/flux_pipeline.hpp b/src/cpp/src/image_generation/flux_pipeline.hpp
@@ -429,15 +429,6 @@ class FluxPipeline : public DiffusionPipeline {
         const size_t vae_scale_factor = m_vae->get_vae_scale_factor();
         const auto& transformer_config = m_transformer->get_config();
 
-        // in case of image to image generation_config_value is just ignored and computed based on initial image
-        if (m_pipeline_type == PipelineType::IMAGE_2_IMAGE) {
-            OPENVINO_ASSERT(initial_image, "Initial image is empty for image to image pipeline");
-            ov::Shape shape = initial_image.get_shape();
-            int64_t dim_val = shape[dim_idx];
-
-            generation_config_value = dim_val - (dim_val % vae_scale_factor);
-        }
-
         if (generation_config_value < 0)
             generation_config_value = transformer_config.m_default_sample_size * vae_scale_factor;
     }
@@ -460,8 +451,6 @@ class FluxPipeline : public DiffusionPipeline {
                 m_generation_config.guidance_scale = 7.0f;
                 m_generation_config.num_inference_steps = 28;
                 m_generation_config.strength = 0.6f;
-                m_generation_config.height = 1024;
-                m_generation_config.width = 1024;
             }
             m_generation_config.max_sequence_length = 512;
         } else {

diff --git a/src/python/py_image_generation_pipelines.cpp b/src/python/py_image_generation_pipelines.cpp
@@ -330,6 +330,7 @@ void init_image_generation_pipelines(py::module_& m) {
         .def_static("stable_diffusion", &ov::genai::Image2ImagePipeline::stable_diffusion, py::arg("scheduler"), py::arg("clip_text_model"), py::arg("unet"), py::arg("vae"))
         .def_static("latent_consistency_model", &ov::genai::Image2ImagePipeline::latent_consistency_model, py::arg("scheduler"), py::arg("clip_text_model"), py::arg("unet"), py::arg("vae"))
         .def_static("stable_diffusion_xl", &ov::genai::Image2ImagePipeline::stable_diffusion_xl, py::arg("scheduler"), py::arg("clip_text_model"), py::arg("clip_text_model_with_projection"), py::arg("unet"), py::arg("vae"))
+        .def_static("flux", &ov::genai::Image2ImagePipeline::flux, py::arg("scheduler"), py::arg("clip_text_model"), py::arg("t5_encoder_model"), py::arg("transformer"), py::arg("vae"))
         .def(
             "compile",
             [](ov::genai::Image2ImagePipeline& pipe,

diff --git a/tools/who_what_benchmark/tests/test_cli_image.py b/tools/who_what_benchmark/tests/test_cli_image.py
@@ -103,8 +103,8 @@ def test_image_model_types(model_id, model_type, backend):
                             ])),
 )
 def test_image_model_genai(model_id, model_type):
-    if ("flux" in model_id or "stable-diffusion-3" in model_id) and model_type != "text-to-image":
-        pytest.skip(reason="FLUX or SD3 are supported as text to image only")
+    if ("stable-diffusion-3" in model_id) and model_type != "text-to-image":
+        pytest.skip(reason="SD3 is supported as text to image only")
 
     with tempfile.TemporaryDirectory() as temp_dir:
         GT_FILE = os.path.join(temp_dir, "gt.csv")