Update app.py

Remove inpaint portion temporarily

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 = [["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(
@@ -197,7 +197,7 @@ with gr.Blocks() as demo:
         ).style(grid=[2], height="auto")
             
     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])   
+    # run_btn_inpaint.click(process_image_inpaint, inputs = [input_image_inpaint,mask_image_inpaint,input_prompt_inpaint], outputs = [result_gallery])   
     
                      
 demo.launch(debug=True)