[Image generation] added progress bar to C++ samples

openvinotoolkit · Jan 9, 2025 · d254b1e · d254b1e
1 parent 7ef754c
commit d254b1e
Show file tree

Hide file tree

Showing 7 changed files with 62 additions and 16 deletions.
diff --git a/samples/cpp/image_generation/CMakeLists.txt b/samples/cpp/image_generation/CMakeLists.txt
@@ -11,12 +11,23 @@ find_package(OpenVINOGenAI REQUIRED
 file(DOWNLOAD https://raw.githubusercontent.com/nothings/stb/f75e8d1cad7d90d72ef7a4661f1b994ef78b4e31/stb_image.h ${CMAKE_BINARY_DIR}/stb_image.h
      EXPECTED_HASH MD5=27932e6fb3a2f26aee2fc33f2cb4e696)
 
+include(FetchContent)
+
+if(POLICY CMP0135)
+    cmake_policy(SET CMP0135 NEW)
+endif()
+
+FetchContent_Declare(indicators
+    URL https://github.com/p-ranav/indicators/archive/refs/tags/v2.3.tar.gz
+    URL_HASH SHA256=70da7a693ff7a6a283850ab6d62acf628eea17d386488af8918576d0760aef7b)
+FetchContent_MakeAvailable(indicators)
+
 # create main sample executable
 
 add_executable(text2image text2image.cpp imwrite.cpp)
 
 target_include_directories(text2image PRIVATE ${CMAKE_BINARY_DIR} "${CMAKE_CURRENT_SOURCE_DIR}")
-target_link_libraries(text2image PRIVATE openvino::genai)
+target_link_libraries(text2image PRIVATE openvino::genai indicators::indicators)
 
 set_target_properties(text2image PROPERTIES
     COMPILE_PDB_NAME text2image
@@ -33,7 +44,7 @@ install(TARGETS text2image
 add_executable(lora_text2image lora_text2image.cpp imwrite.cpp)
 
 target_include_directories(lora_text2image PRIVATE "${CMAKE_CURRENT_SOURCE_DIR}")
-target_link_libraries(lora_text2image PRIVATE openvino::genai)
+target_link_libraries(lora_text2image PRIVATE openvino::genai indicators::indicators)
 
 set_target_properties(lora_text2image PROPERTIES
     COMPILE_PDB_NAME lora_text2image
@@ -52,7 +63,7 @@ add_executable(heterogeneous_stable_diffusion
     imwrite.cpp)
 
 target_include_directories(heterogeneous_stable_diffusion PRIVATE "${CMAKE_CURRENT_SOURCE_DIR}")
-target_link_libraries(heterogeneous_stable_diffusion PRIVATE openvino::genai)
+target_link_libraries(heterogeneous_stable_diffusion PRIVATE openvino::genai indicators::indicators)
 
 set_target_properties(heterogeneous_stable_diffusion PROPERTIES
     COMPILE_PDB_NAME heterogeneous_stable_diffusion
@@ -69,7 +80,7 @@ install(TARGETS heterogeneous_stable_diffusion
 add_executable(image2image image2image.cpp load_image.cpp imwrite.cpp)
 
 target_include_directories(image2image PRIVATE "${CMAKE_CURRENT_SOURCE_DIR}" "${CMAKE_BINARY_DIR}")
-target_link_libraries(image2image PRIVATE openvino::genai)
+target_link_libraries(image2image PRIVATE openvino::genai indicators::indicators)
 
 set_target_properties(image2image PROPERTIES
     COMPILE_PDB_NAME image2image
@@ -80,13 +91,13 @@ install(TARGETS image2image
         RUNTIME DESTINATION samples_bin/
         COMPONENT samples_bin
         EXCLUDE_FROM_ALL)
-        
-# create LoRA sample executable
+
+# create inpainting executable
 
 add_executable(inpainting inpainting.cpp load_image.cpp imwrite.cpp)
 
 target_include_directories(inpainting PRIVATE "${CMAKE_CURRENT_SOURCE_DIR}" "${CMAKE_BINARY_DIR}")
-target_link_libraries(inpainting PRIVATE openvino::genai)
+target_link_libraries(inpainting PRIVATE openvino::genai indicators::indicators)
 
 set_target_properties(inpainting PROPERTIES
     COMPILE_PDB_NAME inpainting

diff --git a/samples/cpp/image_generation/heterogeneous_stable_diffusion.cpp b/samples/cpp/image_generation/heterogeneous_stable_diffusion.cpp
@@ -1,9 +1,11 @@
 // Copyright (C) 2023-2024 Intel Corporation
 // SPDX-License-Identifier: Apache-2.0
 
-#include "imwrite.hpp"
 #include "openvino/genai/image_generation/text2image_pipeline.hpp"
 
+#include "imwrite.hpp"
+#include "progress_bar.hpp"
+
 int32_t main(int32_t argc, char* argv[]) try {
     OPENVINO_ASSERT(argc >= 3 && argc <= 6,
                     "Usage: ",
@@ -102,7 +104,8 @@ int32_t main(int32_t argc, char* argv[]) try {
                                          ov::genai::width(width),
                                          ov::genai::height(height),
                                          ov::genai::guidance_scale(guidance_scale),
-                                         ov::genai::num_inference_steps(number_of_inference_steps_per_image));
+                                         ov::genai::num_inference_steps(number_of_inference_steps_per_image),
+                                         ov::genai::callback(progress_bar));
 
         imwrite("image_" + std::to_string(imagei) + ".bmp", image, true);
     }

diff --git a/samples/cpp/image_generation/image2image.cpp b/samples/cpp/image_generation/image2image.cpp
@@ -3,8 +3,9 @@
 
 #include "openvino/genai/image_generation/image2image_pipeline.hpp"
 
-#include "load_image.hpp"
 #include "imwrite.hpp"
+#include "load_image.hpp"
+#include "progress_bar.hpp"
 
 int32_t main(int32_t argc, char* argv[]) try {
     OPENVINO_ASSERT(argc == 4, "Usage: ", argv[0], " <MODEL_DIR> '<PROMPT>' <IMAGE>");
@@ -17,7 +18,8 @@ int32_t main(int32_t argc, char* argv[]) try {
     ov::genai::Image2ImagePipeline pipe(models_path, device);
     ov::Tensor generated_image = pipe.generate(prompt, image,
         // controls how initial image is noised after being converted to latent space. `1` means initial image is fully noised
-        ov::genai::strength(0.8f));
+        ov::genai::strength(0.8f),
+        ov::genai::callback(progress_bar));
 
     // writes `num_images_per_prompt` images by pattern name
     imwrite("image_%d.bmp", generated_image, true);

diff --git a/samples/cpp/image_generation/inpainting.cpp b/samples/cpp/image_generation/inpainting.cpp
@@ -3,8 +3,9 @@
 
 #include "openvino/genai/image_generation/inpainting_pipeline.hpp"
 
-#include "load_image.hpp"
 #include "imwrite.hpp"
+#include "load_image.hpp"
+#include "progress_bar.hpp"
 
 int32_t main(int32_t argc, char* argv[]) try {
     OPENVINO_ASSERT(argc == 5, "Usage: ", argv[0], " <MODEL_DIR> '<PROMPT>' <IMAGE> <MASK_IMAGE>");
@@ -16,7 +17,7 @@ int32_t main(int32_t argc, char* argv[]) try {
     ov::Tensor mask_image = utils::load_image(mask_image_path);
 
     ov::genai::InpaintingPipeline pipe(models_path, device);
-    ov::Tensor generated_image = pipe.generate(prompt, image, mask_image);
+    ov::Tensor generated_image = pipe.generate(prompt, image, mask_image, ov::genai::callback(progress_bar));
 
     // writes `num_images_per_prompt` images by pattern name
     imwrite("image_%d.bmp", generated_image, true);

diff --git a/samples/cpp/image_generation/lora_text2image.cpp b/samples/cpp/image_generation/lora_text2image.cpp
@@ -4,6 +4,7 @@
 #include "openvino/genai/image_generation/text2image_pipeline.hpp"
 
 #include "imwrite.hpp"
+#include "progress_bar.hpp"
 
 int32_t main(int32_t argc, char* argv[]) try {
     OPENVINO_ASSERT(argc >= 3 && (argc - 3) % 2 == 0, "Usage: ", argv[0], " <MODEL_DIR> '<PROMPT>' [<LORA_SAFETENSORS> <ALPHA> ...]]");
@@ -27,7 +28,8 @@ int32_t main(int32_t argc, char* argv[]) try {
         ov::genai::width(512),
         ov::genai::height(896),
         ov::genai::num_inference_steps(20),
-        ov::genai::rng_seed(42));
+        ov::genai::rng_seed(42),
+        ov::genai::callback(progress_bar));
     imwrite("lora.bmp", image, true);
 
     std::cout << "Generating image without LoRA adapters applied, resulting image will be in baseline.bmp\n";
@@ -36,7 +38,8 @@ int32_t main(int32_t argc, char* argv[]) try {
         ov::genai::width(512),
         ov::genai::height(896),
         ov::genai::num_inference_steps(20),
-        ov::genai::rng_seed(42));
+        ov::genai::rng_seed(42),
+        ov::genai::callback(progress_bar));
     imwrite("baseline.bmp", image, true);
 
     return EXIT_SUCCESS;

diff --git a/samples/cpp/image_generation/progress_bar.hpp b/samples/cpp/image_generation/progress_bar.hpp
@@ -0,0 +1,24 @@
+// Copyright (C) 2023-2025 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+
+#include "indicators/progress_bar.hpp"
+
+bool progress_bar(size_t step, size_t num_steps, ov::Tensor& /* latent */) {
+    using namespace indicators;
+
+    static ProgressBar bar{
+        option::BarWidth{50},
+        option::ForegroundColor{Color::green},
+        option::FontStyles{std::vector<FontStyle>{FontStyle::bold}},
+        option::ShowElapsedTime{true},
+        option::ShowRemainingTime{true},
+    };
+
+    std::stringstream stream;
+    stream << "Image generation step " << (step + 1) << " / " << num_steps;
+
+    bar.set_option(option::PostfixText{stream.str()});
+    bar.set_progress((100 * (step + 1)) / num_steps);
+
+    return false;
+}
diff --git a/samples/cpp/image_generation/text2image.cpp b/samples/cpp/image_generation/text2image.cpp
@@ -4,6 +4,7 @@
 #include "openvino/genai/image_generation/text2image_pipeline.hpp"
 
 #include "imwrite.hpp"
+#include "progress_bar.hpp"
 
 int32_t main(int32_t argc, char* argv[]) try {
     OPENVINO_ASSERT(argc == 3, "Usage: ", argv[0], " <MODEL_DIR> '<PROMPT>'");
@@ -16,7 +17,8 @@ int32_t main(int32_t argc, char* argv[]) try {
         ov::genai::width(512),
         ov::genai::height(512),
         ov::genai::num_inference_steps(20),
-        ov::genai::num_images_per_prompt(1));
+        ov::genai::num_images_per_prompt(1),
+        ov::genai::callback(progress_bar));
 
     // writes `num_images_per_prompt` images by pattern name
     imwrite("image_%d.bmp", image, true);