Update app.py

Add inpainting

app.py

@@ -14,8 +14,6 @@ from utils import preprocess,prepare_mask_and_masked_image, recover_image
 
 to_pil = T.ToPILImage()
 
-title = "Interactive demo: Raising the Cost of Malicious AI-Powered Image Editing"
-
 model_id_or_path = "runwayml/stable-diffusion-v1-5"
 # model_id_or_path = "CompVis/stable-diffusion-v1-4"
 # model_id_or_path = "CompVis/stable-diffusion-v1-3"
@@ -60,7 +58,7 @@ def pgd(X, model, eps=0.1, step_size=0.015, iters=40, clamp_min=0, clamp_max=1,
             
     return X_adv
     
-def process_image(raw_image,prompt):
+def process_image_img2img(raw_image,prompt):
     resize = T.transforms.Resize(512)
     center_crop = T.transforms.CenterCrop(512)
     init_image = center_crop(resize(raw_image))
@@ -96,19 +94,110 @@ def process_image(raw_image,prompt):
         image_adv = pipe_img2img(prompt=prompt, image=adv_image, strength=STRENGTH, guidance_scale=GUIDANCE, num_inference_steps=NUM_STEPS).images[0]
 
     return [(init_image,"Source Image"), (adv_image, "Adv Image"), (image_nat,"Gen. Image Nat"), (image_adv, "Gen. Image Adv")]
+
+def process_image_inpaint(raw_image,mask, prompt):
+    init_image = raw_image.convert('RGB').resize((512,512))
+    mask_image = mask.convert('RGB')
+    mask_image = ImageOps.invert(mask_image).resize((512,512))
+
+    # Attack using embedding of random image from internet
+    target_url = "https://bostonglobe-prod.cdn.arcpublishing.com/resizer/2-ZvyQ3aRNl_VNo7ja51BM5-Kpk=/960x0/cloudfront-us-east-1.images.arcpublishing.com/bostonglobe/CZOXE32LQQX5UNAB42AOA3SUY4.jpg"
+    response = requests.get(target_url)
+    target_image = Image.open(BytesIO(response.content)).convert("RGB")
+    target_image = target_image.resize((512, 512))
+
+    with torch.autocast('cuda'):
+    mask, X = prepare_mask_and_masked_image(init_image, mask_image)
+    X = X.half().cuda()
+    mask = mask.half().cuda()
+    
+    # Here we attack towards the embedding of a random target image. You can also simply attack towards an embedding of zeros!
+    target = pipe_inpaint.vae.encode(preprocess(target_image).half().cuda()).latent_dist.mean
+
+    adv_X = pgd(X, 
+                target = target,
+                model=pipe_inpaint.vae.encode, 
+                criterion=torch.nn.MSELoss(), 
+                clamp_min=-1, 
+                clamp_max=1,
+                eps=0.06, 
+                step_size=0.01, 
+                iters=1000,
+                mask=1-mask
+               )
+
+    adv_X = (adv_X / 2 + 0.5).clamp(0, 1)
+
+    adv_image = to_pil(adv_X[0]).convert("RGB")
+    adv_image = recover_image(adv_image, init_image, mask_image, background=True)
+
+    # A good seed
+    SEED = 9209
+    
+    # Uncomment the below to generated other images
+    # SEED = np.random.randint(low=0, high=100000)
+    
+    torch.manual_seed(SEED)
+    print(SEED)
     
-description = "This is an unofficial demo for Photoguard, which is an approach to safe-guarding images against manipulation by ML-powerd photo-editing models such as stable diffusion through immunization of images. The demo is based on the <a href='https://github.com/MadryLab/photoguard' style='text-decoration: underline;' target='_blank'> Github </a> implementation provided by the authors." 
+    strength = 0.7
+    guidance_scale = 7.5
+    num_inference_steps = 100
+    
+    image_nat = pipe_inpaint(prompt=prompt, 
+                         image=init_image, 
+                         mask_image=mask_image, 
+                         eta=1,
+                         num_inference_steps=num_inference_steps,
+                         guidance_scale=guidance_scale,
+                         strength=strength
+                        ).images[0]
+    image_nat = recover_image(image_nat, init_image, mask_image)
+    
+    torch.manual_seed(SEED)
+    image_adv = pipe_inpaint(prompt=prompt, 
+                         image=adv_image, 
+                         mask_image=mask_image, 
+                         eta=1,
+                         num_inference_steps=num_inference_steps,
+                         guidance_scale=guidance_scale,
+                         strength=strength
+                        ).images[0]
+    image_adv = recover_image(image_adv, init_image, mask_image)
+
+    return [(init_image,"Source Image"), (adv_image, "Adv Image"), (image_nat,"Gen. Image Nat"), (image_adv, "Gen. Image Adv")]
+
+
 examples = [["dog.png", "dog under heavy rain and muddy ground real"]]
 
-interface = gr.Interface(fn=process_image, 
-                     inputs=[gr.Image(type="pil"), gr.Textbox(label="Prompt")],
-                     outputs=[gr.Gallery(
+
+with gr.Blocks() as demo:
+    gr.Markdown("""
+                ## Interactive demo: Raising the Cost of Malicious AI-Powered Image Editing
+            """)
+    gr.HTML('''
+     <p style="margin-bottom: 10px; font-size: 94%">This is an unofficial demo for Photoguard, which is an approach to safe-guarding images against manipulation by ML-powerd photo-editing models such as stable diffusion through immunization of images. The demo is based on the <a href='https://github.com/MadryLab/photoguard' style='text-decoration: underline;' target='_blank'> Github </a> implementation provided by the authors.</p>
+              ''')
+    
+    with gr.Column(): 
+        with gr.Tab("Simple Image to Image"):
+            input_image_img2img = gr.Image(type="pil", label = "Source Image")
+            input_prompt_img2img = gr.Textbox(label="Prompt")
+            run_btn_img2img = gr.Button('Run')
+        
+        with gr.Tab("Simple Inpainting"):
+            input_image_inpaint = gr.Image(type="pil", label = "Source Image")
+            mask_image_inpaint = gr.Image(type="pil", label = "Mask")
+            input_prompt_inpaint = gr.Textbox(label="Prompt")
+            run_btn_inpaint = gr.Button('Run')
+        
+        with gr.Row():
+            result_gallery = gr.Gallery(
             label="Generated images", show_label=False, elem_id="gallery"
         ).style(grid=[2], height="auto")
-                              ],
-                     title=title,
-                     description = description,
-                     examples = examples
-                     )
+            
+    run_btn_img2img.click(process_image_img2img, inputs = [input_image_img2img,input_prompt_img2img], outputs = [result_gallery])
+    run_btn_inpaint.click(process_image_inpaint, inputs = [input_image_inpaint,mask_image_inpaint,input_prompt_inpaint], outputs = [result_gallery])   
+    
                      
-interface.launch(debug=True)
+demo.launch(debug=True)