Added inpainting pipeline for paint your dreams demo (#229)

Author: adrianboguszewski

demos/paint_your_dreams_demo/main.py

@@ -30,8 +30,7 @@
 
 safety_checker: Optional[Pipeline] = None
 
-ov_pipelines_t2i = {}
-ov_pipelines_i2i = {}
+ov_pipelines = {}
 
 stop_generating: bool = True
 hf_model_name: Optional[str] = None
@@ -83,6 +82,8 @@ async def load_static_pipeline(model_dir: Path,  device: str, size: int, pipelin
         ov_pipeline = genai.Text2ImagePipeline.latent_consistency_model(scheduler, text_encoder, unet, vae)
     elif pipeline == "image2image":
         ov_pipeline = genai.Image2ImagePipeline.latent_consistency_model(scheduler, text_encoder, unet, vae)
+    elif pipeline == "inpainting":
+        ov_pipeline = genai.InpaintingPipeline.latent_consistency_model(scheduler, text_encoder, unet, vae)
     else:
         raise ValueError(f"Unknown pipeline: {pipeline}")
 
@@ -92,17 +93,10 @@ async def load_static_pipeline(model_dir: Path,  device: str, size: int, pipelin
 async def load_pipeline(model_name: str, device: str, size: int, pipeline: str):
     model_dir = MODEL_DIR / model_name
 
-    if pipeline == "text2image":
-        if device not in ov_pipelines_t2i:
-            ov_pipelines_t2i[device] = await load_static_pipeline(model_dir, device, size, pipeline)
-
-        return ov_pipelines_t2i[device]
+    if (device, pipeline) not in ov_pipelines:
+        ov_pipelines[(device, pipeline)] = await load_static_pipeline(model_dir, device, size, pipeline)
 
-    if pipeline == "image2image":
-        if device not in ov_pipelines_i2i:
-            ov_pipelines_i2i[device] = await load_static_pipeline(model_dir, device, size, pipeline)
-
-        return ov_pipelines_i2i[device]
+    return ov_pipelines[(device, pipeline)]
 
 
 async def stop():
@@ -124,27 +118,42 @@ def progress(step, num_steps, latent):
     return False
 
 
-async def generate_images(input_image: np.ndarray, prompt: str, seed: int, guidance_scale: float, num_inference_steps: int,
+async def generate_images(input_image_mask: np.ndarray, prompt: str, seed: int, guidance_scale: float, num_inference_steps: int,
                           strength: float, randomize_seed: bool, device: str, endless_generation: bool, size: int) -> tuple[np.ndarray, float]:
     global stop_generating
     stop_generating = not endless_generation
 
+    input_image = input_image_mask["background"][:, :, :3]
+    image_mask = input_image_mask["layers"][0][:, :, 3:]
+
+    # ensure image is square
+    input_image = utils.crop_center(input_image)
+    input_image = cv2.resize(input_image, (size, size))
+    image_mask = cv2.resize(image_mask, (size, size), interpolation=cv2.INTER_NEAREST)
+    image_mask = cv2.cvtColor(image_mask, cv2.COLOR_GRAY2BGR)
+
     while True:
         if randomize_seed:
             seed = random.randint(0, MAX_SEED)
 
         start_time = time.time()
-        if input_image is None:
+        if input_image.any():
+
+            # inpainting pipeline
+            if image_mask.any():
+                ov_pipeline = await load_pipeline(hf_model_name, device, size, "inpainting")
+                result = ov_pipeline.generate(prompt=prompt, image=ov.Tensor(input_image[None]), mask_image=ov.Tensor(image_mask[None]), num_inference_steps=num_inference_steps,
+                                              width=size, height=size, guidance_scale=guidance_scale, strength=1.0 - strength, rng_seed=seed, callback=progress).data[0]
+            # image2image pipeline
+            else:
+                ov_pipeline = await load_pipeline(hf_model_name, device, size,"image2image")
+                result = ov_pipeline.generate(prompt=prompt, image=ov.Tensor(input_image[None]), num_inference_steps=num_inference_steps, width=size, height=size,
+                                 guidance_scale=guidance_scale, strength=1.0 - strength, rng_seed=seed, callback=progress).data[0]
+        # text2image pipeline
+        else:
             ov_pipeline = await load_pipeline(hf_model_name, device, size, "text2image")
             result = ov_pipeline.generate(prompt=prompt, num_inference_steps=num_inference_steps, width=size, height=size,
                                           guidance_scale=guidance_scale, rng_seed=seed, callback=progress).data[0]
-        else:
-            ov_pipeline = await load_pipeline(hf_model_name, device, size,"image2image")
-            # ensure image is square
-            input_image = utils.crop_center(input_image)
-            input_image = cv2.resize(input_image, (size, size))
-            result = ov_pipeline.generate(prompt=prompt, image=ov.Tensor(input_image[None]), num_inference_steps=num_inference_steps, width=size, height=size,
-                             guidance_scale=guidance_scale, strength=1.0 - strength, rng_seed=seed, callback=progress).data[0]
         end_time = time.time()
 
         label = safety_checker(Image.fromarray(result), top_k=1)
@@ -191,7 +200,7 @@ def build_ui(image_size: int) -> gr.Interface:
             with gr.Row():
                 with gr.Column():
                     with gr.Row():
-                        input_image = gr.Image(label="Input image (leave blank for text2image generation)", sources=["webcam", "clipboard", "upload"])
+                        input_image = gr.ImageMask(label="Input image (leave blank for text2image generation)", sources=["webcam", "clipboard", "upload"])
                         result_img = gr.Image(label="Generated image", elem_id="output_image", format="png")
                     with gr.Row():
                         result_time_label = gr.Text("", label="Inference time", type="text")

demos/paint_your_dreams_demo/requirements.txt

@@ -7,13 +7,13 @@ openvino-genai==2025.1.0.0.dev20250305
 optimum-intel==1.22.0
 optimum==1.24.0
 onnx==1.17.0
-huggingface-hub==0.27.0
+huggingface-hub==0.29.3
 diffusers==0.32.1
 transformers==4.48.3
 torch==2.5.1
 accelerate==1.2.1
 pillow==11.1.0
 opencv-python==4.10.0.84
 numpy==2.1.3
-gradio==5.12.0
+gradio==5.22.0
 tqdm==4.67.1