update: add color inpainting feature

app.py

@@ -66,16 +66,10 @@ def resize_image(input_image, resolution, max_edge=False, edge_limit=False):
     pimg[oH:oH+H, oW:oW+W] = img
     return pimg
 
-def get_canny_filter(image, format='pil', low_threshold=100, high_threshold=200):
-    
-    if not isinstance(image, np.ndarray):
-        image = np.array(image) 
-        
+def get_canny_filter(image, low_threshold=100, high_threshold=200):
     image = cv2.Canny(image, low_threshold, high_threshold)
     image = image[:, :, None]
     image = np.concatenate([image, image, image], axis=2)
-    if format == 'pil':
-        image = Image.fromarray(image)
     return image
 
 def get_color_filter(cond_image, mask_size=64):
@@ -85,16 +79,11 @@ def get_color_filter(cond_image, mask_size=64):
     return color
 
 def get_colorcanny(image, mask_size):
-    
-    if not isinstance(image, np.ndarray):
-        image = np.array(image) 
-     
-    canny_img = get_canny_filter(image, format='np')
+    canny_img = get_canny_filter(image)
 
     color_img = get_color_filter(image, int(mask_size))
 
     color_img[np.where(canny_img > 128)] = 255
-    color_img = Image.fromarray(color_img)
     return color_img
     
 def process(input_image, prompt, n_prompt, strength=1.0, color_mask_size=96, size=512, scale=6.0, ddim_steps=20):
@@ -102,6 +91,35 @@ def process(input_image, prompt, n_prompt, strength=1.0, color_mask_size=96, siz
     input_image = resize_image(input_image, size, max_edge=True, edge_limit=True)
 
     cond_img = get_colorcanny(input_image, color_mask_size)
+    cond_img = Image.fromarray(cond_img)
+    output = pipe(
+        prompt_embeds=prompt_embeds, 
+        negative_prompt_embeds=negative_prompt_embeds,
+        image=cond_img,
+        generator=generator,
+        num_images_per_prompt=1,
+        num_inference_steps=ddim_steps,
+        guidance_scale=scale,
+        controlnet_conditioning_scale=float(strength)
+    )
+    return [output.images[0], cond_img]
+
+
+def inpaint_process(inpaint_image, input_image, prompt, n_prompt, strength=1.0, color_mask_size=96, size=512, scale=6.0, ddim_steps=20):
+    if inpaint_image is None:
+        return process(input_image, prompt, n_prompt, strength, color_mask_size, size, scale, ddim_steps)
+    
+    prompt_embeds, negative_prompt_embeds = _encode_prompt(pipe, prompt, pipe.device, 1, True, n_prompt, 3)
+    input_image = resize_image(input_image, size, max_edge=True, edge_limit=True)
+    inpaint_image = resize_image(inpaint_image, size, max_edge=True, edge_limit=True)
+
+    canny_img = get_canny_filter(input_image)
+
+    color_img = get_color_filter(inpaint_image, int(color_mask_size))
+
+    color_img[np.where(canny_img > 128)] = 255
+    cond_img = Image.fromarray(color_img)
+
     output = pipe(
         prompt_embeds=prompt_embeds, 
         negative_prompt_embeds=negative_prompt_embeds,
@@ -130,9 +148,12 @@ with block:
     with gr.Row():
         with gr.Column():
             input_image = gr.Image(source='upload', type="numpy")
+            color_image = gr.ImagePaint(type="numpy")
             prompt = gr.Textbox(label="Prompt", value='')
             n_prompt = gr.Textbox(label="Negative Prompt", value='')
-            run_button = gr.Button(label="Run")
+            with gr.Row():
+                run_button = gr.Button(label="Run")
+                run_edit_button = gr.Button(value='Run with inpaint color', label="Run with inpaint color")
             with gr.Accordion('Advanced', open=False):
                 strength = gr.Slider(label="Control Strength", minimum=0.0, maximum=2.0, value=1.0, step=0.01)
                 color_mask_size = gr.Slider(label="Color Mask Size", minimum=32, maximum=256, value=96, step=16)
@@ -144,6 +165,8 @@ with block:
             result_gallery = gr.Gallery(label='Output', show_label=False, elem_id="gallery").style(grid=2, height='auto')
     ips = [input_image, prompt, n_prompt, strength, color_mask_size, size, scale, ddim_steps]
     run_button.click(fn=process, inputs=ips, outputs=[result_gallery])
+    run_edit_button.click(fn=inpaint_process, inputs=[color_image] + ips, outputs=[result_gallery])
+
 
     gr.Examples(
         examples=[