sd3 notebook migrate to genai

notebooks/stable-diffusion-v3/README.md

@@ -31,12 +31,9 @@ This folder contains notebooks that demonstrate the use of the Stable Diffusion
 The OpenVINO tutorial consists of the following steps:
 
 - Install prerequisites
-- Collect Pytorch model pipeline
 - Convert model to OpenVINO intermediate representation (IR) format and compress weights using NNCF
 - Prepare OpenVINO Inference pipeline
 - Run Text-to-Image generation
-- Optimize pipeline with [NNCF](https://github.com/openvinotoolkit/nncf/)
-- Compare the results of the original and optimized pipelines
 - Launch interactive demo
 
 The Torch FX tutorial consists of the following steps:

notebooks/stable-diffusion-v3/gradio_helper.py

@@ -1,6 +1,10 @@
 import torch
 import gradio as gr
 import numpy as np
+import sys
+import openvino_genai as ov_genai
+from tqdm.auto import tqdm
+from PIL import Image
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1344
@@ -24,30 +28,43 @@
 """
 
 
-def make_demo(pipeline, use_flash_lora):
+def make_demo(pipeline, turbo):
     def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, progress=gr.Progress(track_tqdm=True)):
         if randomize_seed:
             seed = np.random.randint(0, MAX_SEED)
 
-        generator = torch.Generator().manual_seed(seed)
+        generator = ov_genai.TorchGenerator(seed)
+        pbar = tqdm(total=num_inference_steps)
+        def callback(step, num_steps, latent):
+            if num_steps != pbar.total:
+                pbar.reset(num_steps)
+            pbar.update(1)
+            sys.stdout.flush()
+            return False
 
-        image = pipeline(
+        generate_kwargs = {}
+
+        if guidance_scale > 1:
+            generate_kwargs["negative_prompt"] = negative_prompt
+
+        image_tensor = pipeline.generate(
             prompt=prompt,
-            negative_prompt=negative_prompt,
             guidance_scale=guidance_scale,
             num_inference_steps=num_inference_steps,
             width=width,
             height=height,
             generator=generator,
-        ).images[0]
-
+            callback=callback,
+            **generate_kwargs
+        )
+        image = Image.fromarray(image_tensor.data[0])
         return image, seed
 
     with gr.Blocks(css=css) as demo:
         with gr.Column(elem_id="col-container"):
             gr.Markdown(
                 """
-            # Demo [Stable Diffusion 3 Medium](https://huggingface.co/stabilityai/stable-diffusion-3-medium) with OpenVINO
+            # Demo Stable Diffusion 3  with OpenVINO
             """
             )
 
@@ -102,17 +119,17 @@ def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance
                     guidance_scale = gr.Slider(
                         label="Guidance scale",
                         minimum=0.0,
-                        maximum=10.0 if not use_flash_lora else 2,
+                        maximum=10.0 if not turbo else 2,
                         step=0.1,
-                        value=5.0 if not use_flash_lora else 0,
+                        value=5.0 if not turbo else 1.5,
                     )
 
                     num_inference_steps = gr.Slider(
                         label="Number of inference steps",
                         minimum=1,
                         maximum=50,
                         step=1,
-                        value=28 if not use_flash_lora else 4,
+                        value=28 if not turbo else 8,
                     )
 
             gr.Examples(examples=examples, inputs=[prompt])