Update app.py

app.py

@@ -27,12 +27,12 @@ pipe_img2img = StableDiffusionImg2ImgPipeline.from_pretrained(
 )
 pipe_img2img = pipe_img2img.to("cuda")
 
-# pipe_inpaint = StableDiffusionInpaintPipeline.from_pretrained(
-#     "runwayml/stable-diffusion-inpainting",
-#     revision="fp16",
-#     torch_dtype=torch.float16,
-# )
-# pipe_inpaint = pipe_inpaint.to("cuda")
+pipe_inpaint = StableDiffusionInpaintPipeline.from_pretrained(
+    "runwayml/stable-diffusion-inpainting",
+    revision="fp16",
+    torch_dtype=torch.float16,
+)
+pipe_inpaint = pipe_inpaint.to("cuda")
 
 def pgd(X, model, eps=0.1, step_size=0.015, iters=40, clamp_min=0, clamp_max=1, mask=None):
     X_adv = X.clone().detach() + (torch.rand(*X.shape)*2*eps-eps).cuda()
@@ -95,77 +95,77 @@ def process_image_img2img(raw_image,prompt):
 
     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()
+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
+        # 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 = 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_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)
+    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
+    # A good seed
+    SEED = 9209
     
-#     # Uncomment the below to generated other images
-#     # SEED = np.random.randint(low=0, high=100000)
+    # Uncomment the below to generated other images
+    # SEED = np.random.randint(low=0, high=100000)
     
-#     torch.manual_seed(SEED)
-#     print(SEED)
+    torch.manual_seed(SEED)
+    print(SEED)
     
-#     strength = 0.7
-#     guidance_scale = 7.5
-#     num_inference_steps = 100
+    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)
+    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)
+    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")]
+    return [(init_image,"Source Image"), (adv_image, "Adv Image"), (image_nat,"Gen. Image Nat"), (image_adv, "Gen. Image Adv")]
 
 
 examples_list = [["dog.png", "dog under heavy rain and muddy ground real"]]
@@ -185,11 +185,11 @@ with gr.Blocks() as demo:
             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.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(
@@ -198,7 +198,7 @@ with gr.Blocks() as demo:
             
     run_btn_img2img.click(process_image_img2img, inputs = [input_image_img2img,input_prompt_img2img], outputs = [result_gallery])
     examples = gr.Examples(examples=examples_list,inputs = [input_image_img2img,input_prompt_img2img],  outputs = [result_gallery], cache_examples = True, fn = process_image_img2img)
-    # run_btn_inpaint.click(process_image_inpaint, inputs = [input_image_inpaint,mask_image_inpaint,input_prompt_inpaint], outputs = [result_gallery])   
+    run_btn_inpaint.click(process_image_inpaint, inputs = [input_image_inpaint,mask_image_inpaint,input_prompt_inpaint], outputs = [result_gallery])   
     
                      
 demo.launch(debug=True)