compare

app.py

@@ -64,31 +64,45 @@ def mediapipe_segmentation(image_file, mask_file):
         cv2.imwrite(mask_file, output_image)
 
 
-def image_inpainting(prompt, image_path, mask_image_path, num_samples=4):
+def image_inpainting(prompt, image_path, mask_image_path, num_samples=4, is_origin=False):
     image = PIL.Image.open(image_path).convert("RGB").resize((512, 512))
     mask_image = PIL.Image.open(mask_image_path).convert("RGB").resize((512, 512))
+    if not is_origin:
+        guidance_scale = 7.5
+        generator = torch.Generator(device=device).manual_seed(0)  # change the seed to get different results
 
-    guidance_scale = 7.5
-    generator = torch.Generator(device=device).manual_seed(0)  # change the seed to get different results
-
-    images = pipe(prompt=prompt, image=image, mask_image=mask_image, guidance_scale=guidance_scale, generator=generator,
-                  num_images_per_prompt=num_samples).images
+        images = pipe(prompt=prompt, image=image, mask_image=mask_image, guidance_scale=guidance_scale, generator=generator,
+                      num_images_per_prompt=num_samples).images
+    else:
+        images = pipe(prompt=prompt, image=image, mask_image=mask_image, num_images_per_prompt=num_samples).images
 
     # insert initial image in the list so we can compare side by side
     # images.insert(0, image)
     return image_grid(images, 2, math.ceil(num_samples/2))
 
 
+title = "Person Matting & Stable Diffusion In-Painting"
+description = "Inpainting Stable Diffusion <br/><b>mediapipe + Stable Diffusion<b/><br/>"
+
+
 def predict1(dict, prompt):
     dict['image'].save('image.png')
-    dict['mask'].save('mask.png')
-    mediapipe_segmentation('image.png', 'm_mask.png')
+    # dict['mask'].save('mask.png')
+    mediapipe_segmentation('image.png')
+    image = image_inpainting(prompt, image_path='image.png', mask_image_path='m_mask.png', is_origin=False)
+    return image
 
-    image = image_inpainting(prompt, image_path='image.png', mask_image_path='m_mask.png')
+
+def predict2(dict, prompt):
+    dict['image'].save('image.png')
+    dict['mask'].save('mask.png')
+    image = image_inpainting(prompt, image_path='image.png', mask_image_path='mask.png', is_origin=True)
     return image
 
 
-title = "Person Matting & Stable Diffusion In-Painting"
-description = "Inpainting Stable Diffusion <br/><b>mediapipe + Stable Diffusion<b/><br/>"
-gr.Interface(predict1, inputs=[gr.Image(source='upload', tool='sketch', type='pil'), gr.Textbox(label='prompt')],
-             outputs='image', title=title, description=description).launch(max_threads=True)
+image_input = gr.Image(source='upload', tool='sketch', type='pil')
+prompt = gr.Textbox(label='prompt')
+
+greeter_1 = gr.Interface(predict1, inputs=[image_input, prompt], outputs=gr.Image(label='auto'))
+greeter_2 = gr.Interface(predict2, inputs=[image_input, prompt], outputs=gr.Image(label='paint'))
+demo = gr.Parallel(greeter_1, greeter_2).launch(max_threads=True)