Update app.py

app.py

@@ -36,14 +36,15 @@ pipe_inpaint = StableDiffusionInpaintPipeline.from_pretrained(
 )
 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):
+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()
     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()
+        X_adv.requires_grad_(True)
+
+        loss = (model(X_adv).latent_dist.mean).norm()
 
         pbar.set_description(f"[Running attack]: Loss {loss.item():.5f} | step size: {actual_step_size:.4}")