retry the ImageEditor

app.py

@@ -51,17 +51,11 @@ class BZHStableSignatureDemo(object):
         self.pipe.vae.load_state_dict(sd, strict=False)
 
         output = self.pipe(prompt, num_inference_steps=4, guidance_scale=0.0, output_type="pil")
-        return output.images[0]
+        return { "background": output.images[0] }
 
-    @staticmethod
-    def pad(img, padding, mode="edge"):
-        npimg = np.asarray(img)
-        nppad = ((padding[1], padding[3]), (padding[0], padding[2]), (0,0))
-        npimg = np.pad(npimg, nppad, mode=mode)
-        return Image.fromarray(npimg)
-
-    def attack_detect(self, img, jpeg_compression, downscale, saturation):
+    def attack_detect(self, img_edit, jpeg_compression, downscale, saturation):
 
+        img = img_edit["composite"]
         # attack
         if downscale != 1:
             size = img.size
@@ -121,7 +115,7 @@ def interface():
         with gr.Row():
             btn1 = gr.Button("Generate")
         with gr.Row():
-            watermarked_image = gr.Image(type="pil", width=512, height=512)
+            watermarked_image = gr.ImageEditor(type="pil", width=512, height=512)
             with gr.Column():
                 downscale = gr.Slider(1, 3, value=1, step=0.1, label="Downscale ratio")
                 saturation = gr.Slider(0, 2, value=1, step=0.1, label="Color saturation")