Modify pgd to handle inpainting

app.py

@@ -58,6 +58,29 @@ def pgd(X, model, eps=0.1, step_size=0.015, iters=40, clamp_min=0, clamp_max=1,
             X_adv.data *= mask
             
     return X_adv
+
+def pgd_inpaint(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)                    
+
+        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}")
+
+        grad, = torch.autograd.grad(loss, [X_adv])
+        
+        X_adv = X_adv - grad.detach().sign() * actual_step_size
+        X_adv = torch.minimum(torch.maximum(X_adv, X - eps), X + eps)
+        X_adv.data = torch.clamp(X_adv, min=clamp_min, max=clamp_max)
+        X_adv.grad = None    
+        
+        if mask is not None:
+            X_adv.data *= mask
+            
+    return X_adv
     
 def process_image_img2img(raw_image,prompt):
     resize = T.transforms.Resize(512)
@@ -115,7 +138,7 @@ def process_image_inpaint(raw_image,mask, prompt):
         # 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, 
+        adv_X = pgd_inpaint(X, 
                     target = target,
                     model=pipe_inpaint.vae.encode, 
                     criterion=torch.nn.MSELoss(),