Update

app.py

@@ -13,7 +13,7 @@ import cv2
 from matplotlib import pyplot as plt
 from torchvision import transforms
 from diffusers import DiffusionPipeline
-
+from diffusers.utils import torch_device
 pipe = DiffusionPipeline.from_pretrained(
     "patrickvonplaten/new_inpaint_test",
     torch_dtype=torch.float16,
@@ -30,11 +30,19 @@ def read_content(file_path: str) -> str:
 
     return content
 
-def predict(dict, example_image):
-    init_image = dict["image"].convert("RGB")
-    mask = dict["mask"].convert("RGB")
-    image = pipe(image=init_image, mask_image=mask, example_image=example_image).images[0]
-    return image, gr.update(visible=True), gr.update(visible=True), gr.update(visible=True)
+def predict(dict, reference, scale, seed, step):
+    init_image = dict["image"].convert("RGB").resize((512,512))
+    mask = dict["mask"].convert("RGB").resize((512,512))
+    generator = torch.Generator('cuda').manual_seed(seed) if seed != 0 else None
+    output = pipe(
+        image=init_image,
+        mask_image=mask,
+        example_image=reference,
+        generator=generator,
+        guidance_scale=scale,
+        num_inference_steps=step,
+    ).images[0]
+    return output, gr.update(visible=True), gr.update(visible=True), gr.update(visible=True)
 
 
 css = '''
@@ -75,6 +83,16 @@ css = '''
     display: none !important;
 }
 '''
+example={}
+for i in range(1,4):
+    ex_image_path='examples/image/example_'+str(i)+'.png'
+    ex_mask_path='examples/mask/example_'+str(i)+'.png'
+    ex_reference_path='examples/reference/example_'+str(i)+'.jpg'
+    ex_image=Image.open(ex_image_path)
+    ex_mask=Image.open(ex_mask_path)
+    ex_reference=Image.open(ex_reference_path)
+    example[i]={'image':{'image':ex_image,'mask':ex_mask},'reference':ex_reference}
+
 
 image_blocks = gr.Blocks(css=css)
 with image_blocks as demo:
@@ -83,23 +101,39 @@ with image_blocks as demo:
         with gr.Box():
             with gr.Row():
                 with gr.Column():
-                    image = gr.Image(source='upload', tool='sketch', elem_id="image_upload", type="pil", label="Upload")
-                    example = gr.Image(source='upload', elem_id="image_upload", type="pil", label="Upload")
+                    image = gr.Image(source='upload', tool='sketch', elem_id="image_upload", type="pil", label="Source Image")
+                    reference = gr.Image(source='upload', elem_id="image_upload", type="pil", label="Reference Image")
+
+                with gr.Column():
+                    image_out = gr.Image(label="Output", elem_id="output-img").style(height=400)
+                    guidance = gr.Slider(label="Guidance scale", value=5, maximum=15,interactive=True)
+                    steps = gr.Slider(label="Steps", value=50, minimum=2, maximum=75, step=1,interactive=True)
+
+                    seed = gr.Slider(0, 10000, label='Seed (0 = random)', value=0, step=1)
+
                     with gr.Row(elem_id="prompt-container").style(mobile_collapse=False, equal_height=True):
-                        btn = gr.Button("Inpaint!").style(
+                        btn = gr.Button("Paint!").style(
                             margin=False,
                             rounded=(False, True, True, False),
-                            full_width=False,
+                            full_width=True,
                         )
-                with gr.Column():
-                    image_out = gr.Image(label="Output", elem_id="output-img").style(height=400)
                     with gr.Group(elem_id="share-btn-container"):
                         community_icon = gr.HTML(community_icon_html, visible=False)
                         loading_icon = gr.HTML(loading_icon_html, visible=False)
-                        share_button = gr.Button("Share to community", elem_id="share-btn", visible=False)
+                        share_button = gr.Button("Share to community", elem_id="share-btn", visible=False).style(
+                            margin=False,
+                            rounded=(False, True, True, False),
+                            full_width=True,
+                        )
             
-
-            btn.click(fn=predict, inputs=[image, example], outputs=[image_out, community_icon, loading_icon, share_button])
+            with gr.Row():
+                gr.Examples([
+                    ['examples/image/example_2.png', 'examples/reference/example_2.jpg',5,50],
+                    ['examples/image/example_3.png', 'examples/reference/example_3.jpg',5,50],
+                    ['examples/image/example_1.png', 'examples/reference/example_1.jpg',5,50],
+                ], inputs=[image, reference, guidance, steps])
+            
+            btn.click(fn=predict, inputs=[image, reference, guidance, seed, steps], outputs=[image_out, community_icon, loading_icon, share_button])
             share_button.click(None, [], [], _js=share_js)
 
 
@@ -113,10 +147,7 @@ with image_blocks as demo:
                     <div class="acknowledgments">
                         <p><h4>LICENSE</h4>
         The model is licensed with a <a href="https://huggingface.co/spaces/CompVis/stable-diffusion-license" style="text-decoration: underline;" target="_blank">CreativeML Open RAIL-M</a> license. The authors claim no rights on the outputs you generate, you are free to use them and are accountable for their use which must not go against the provisions set in this license. The license forbids you from sharing any content that violates any laws, produce any harm to a person, disseminate any personal information that would be meant for harm, spread misinformation and target vulnerable groups. For the full list of restrictions please <a href="https://huggingface.co/spaces/CompVis/stable-diffusion-license" target="_blank" style="text-decoration: underline;" target="_blank">read the license</a></p>
-                        <p><h4>Biases and content acknowledgment</h4>
-        Despite how impressive being able to turn text into image is, beware to the fact that this model may output content that reinforces or exacerbates societal biases, as well as realistic faces, pornography and violence. The model was trained on the <a href="https://laion.ai/blog/laion-5b/" style="text-decoration: underline;" target="_blank">LAION-5B dataset</a>, which scraped non-curated image-text-pairs from the internet (the exception being the removal of illegal content) and is meant for research purposes. You can read more in the <a href="https://huggingface.co/CompVis/stable-diffusion-v1-4" style="text-decoration: underline;" target="_blank">model card</a></p>
-                    </div>
                 """
             )
 
-image_blocks.launch()
+image_blocks.launch(server_name='0.0.0.0')

header.html

@@ -12,7 +12,10 @@
   </div>
   <div>
     <p style="align-items: center; margin-bottom: 7px;">
-      Paint by Example, add a input and draw a mask for what you want to replace with a example image 
+      Paint by Example, upload a source image and draw a mask for what you want to replace with a example image. 
+    </p>
+    <p style="align-items: center; margin-bottom: 7px;">
+      This demo currently only supports square source image. It will raise an error if the souce image is not square.
     </p>
   </div>
 </div>