Update app.py

Add inpainting model

app.py

@@ -7,10 +7,10 @@ import torch
 import requests
 from tqdm import tqdm
 
-from diffusers import StableDiffusionImg2ImgPipeline
+from diffusers import StableDiffusionImg2ImgPipeline, StableDiffusionInpaintPipeline
 import torchvision.transforms as T
 
-from utils import preprocess, recover_image
+from utils import preprocess,prepare_mask_and_masked_image, recover_image
 
 to_pil = T.ToPILImage()
 
@@ -28,15 +28,22 @@ pipe_img2img = StableDiffusionImg2ImgPipeline.from_pretrained(
     torch_dtype=torch.float16,
 )
 pipe_img2img = pipe_img2img.to("cuda")
-def pgd(X, model, eps=0.1, step_size=0.015, iters=40, clamp_min=0, clamp_max=1, mask=None):
+
+pipe_inpaint = StableDiffusionInpaintPipeline.from_pretrained(
+    "runwayml/stable-diffusion-inpainting",
+    revision="fp16",
+    torch_dtype=torch.float16,
+)
+pipe_inpaint = pipe_inpaint.to("cuda")
+
+def pgd(X, target, model, criterion, 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()
     pbar = tqdm(range(iters))
     for i in pbar:
         actual_step_size = step_size - (step_size - step_size / 100) / iters * i  
+        X_adv.requires_grad_(True)                    
 
-        X_adv.requires_grad_(True)
-
-        loss = (model(X_adv).latent_dist.mean).norm()
+        loss = (model(X_adv).latent_dist.mean - target).norm()
 
         pbar.set_description(f"[Running attack]: Loss {loss.item():.5f} | step size: {actual_step_size:.4}")