Update text2image genai notebook (openvinotoolkit#2660)

sgolebiewski-intel · Jan 16, 2025 · 231f963 · 231f963
1 parent 383cfbd
commit 231f963
Show file tree

Hide file tree

Showing 2 changed files with 406 additions and 35 deletions.
diff --git a/notebooks/text-to-image-genai/gradio_helper.py b/notebooks/text-to-image-genai/gradio_helper.py
@@ -1,5 +1,7 @@
 import gradio as gr
 import numpy as np
+from tqdm.auto import tqdm
+import sys
 
 import openvino_genai as ov_genai
 
@@ -30,6 +32,13 @@ def infer(prompt, negative_prompt, seed, randomize_seed, width, height, num_infe
 
         generator = generator_cls(seed)
 
+        pbar = tqdm(total=num_inference_steps)
+
+        def callback(step, num_steps, latent):
+            pbar.update(1)
+            sys.stdout.flush()
+            return False
+
         image_tensor = pipeline.generate(
             prompt=prompt,
             negative_prompt=negative_prompt,
@@ -38,6 +47,7 @@ def infer(prompt, negative_prompt, seed, randomize_seed, width, height, num_infe
             height=height,
             generator=generator,
             adapters=adapter_config if use_lora else ov_genai.AdapterConfig(),
+            callback=callback,
         )
 
         return image_tensor.data[0], seed