Create app.py

app.py

@@ -0,0 +1,485 @@
+import gradio as gr
+
+import os
+
+from PIL import Image, ImageOps
+
+import matplotlib.pyplot as plt
+
+import numpy as np
+
+import torch
+
+import requests
+
+from tqdm import tqdm
+
+from io import BytesIO
+
+
+
+from diffusers import StableDiffusionImg2ImgPipeline, StableDiffusionInpaintPipeline
+
+import torchvision.transforms as T
+
+
+
+from utils import preprocess,prepare_mask_and_masked_image, recover_image
+
+
+
+to_pil = T.ToPILImage()
+
+
+
+model_id_or_path = "runwayml/stable-diffusion-v1-5"
+
+# model_id_or_path = "CompVis/stable-diffusion-v1-4"
+
+# model_id_or_path = "CompVis/stable-diffusion-v1-3"
+
+# model_id_or_path = "CompVis/stable-diffusion-v1-2"
+
+# model_id_or_path = "CompVis/stable-diffusion-v1-1"
+
+
+
+pipe_img2img = StableDiffusionImg2ImgPipeline.from_pretrained(
+
+model_id_or_path,
+
+revision="fp16", 
+
+torch_dtype=torch.float16,
+
+)
+
+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")
+
+
+
+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).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 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, scale, num_steps, seed):
+
+resize = T.transforms.Resize(512)
+
+center_crop = T.transforms.CenterCrop(512)
+
+init_image = center_crop(resize(raw_image))
+
+with torch.autocast('cuda'):
+
+X = preprocess(init_image).half().cuda()
+
+adv_X = pgd(X, 
+
+model=pipe_img2img.vae.encode, 
+
+clamp_min=-1, 
+
+clamp_max=1,
+
+eps=0.06, # The higher, the less imperceptible the attack is 
+
+step_size=0.02, # Set smaller than eps
+
+iters=100, # The higher, the stronger your attack will be
+
+)
+
+
+
+# convert pixels back to [0,1] range
+
+adv_X = (adv_X / 2 + 0.5).clamp(0, 1)
+
+
+
+adv_image = to_pil(adv_X[0]).convert("RGB")
+
+
+
+# a good seed (uncomment the line below to generate new images)
+
+SEED = seed# Default is 9222
+
+# SEED = np.random.randint(low=0, high=10000)
+
+
+
+# Play with these for improving generated image quality
+
+STRENGTH = 0.5
+
+GUIDANCE = scale # Default is 7.5
+
+NUM_STEPS = num_steps # Default is 50
+
+
+
+with torch.autocast('cuda'):
+
+torch.manual_seed(SEED)
+
+image_nat = pipe_img2img(prompt=prompt, image=init_image, strength=STRENGTH, guidance_scale=GUIDANCE, num_inference_steps=NUM_STEPS).images[0]
+
+torch.manual_seed(SEED)
+
+image_adv = pipe_img2img(prompt=prompt, image=adv_image, strength=STRENGTH, guidance_scale=GUIDANCE, num_inference_steps=NUM_STEPS).images[0]
+
+
+
+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,scale, num_steps, seed):
+
+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
+
+
+
+adv_X = pgd_inpaint(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_image = to_pil(adv_X[0]).convert("RGB")
+
+adv_image = recover_image(adv_image, init_image, mask_image, background=True)
+
+
+
+# A good seed
+
+SEED = seed #Default is 9209
+
+
+
+# Uncomment the below to generated other images
+
+# SEED = np.random.randint(low=0, high=100000)
+
+
+
+torch.manual_seed(SEED)
+
+print(SEED)
+
+
+
+#strength = 0.7
+
+guidance_scale = scale# Default is 7.5
+
+num_inference_steps = num_steps # Default is 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)
+
+
+
+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")]
+
+
+
+
+
+examples_list = [["dog.png", "dog under heavy rain and muddy ground real", 7.5, 50, 9222]]
+
+
+
+
+
+with gr.Blocks() as demo:
+
+gr.Markdown("""
+
+## Interactive demo: Raising the Cost of Malicious AI-Powered Image Editing
+
+""")
+
+gr.HTML('''
+
+<p style="margin-bottom: 10px; font-size: 94%">This is an unofficial demo for Photoguard, which is an approach to safeguarding images against manipulation by ML-powered photo-editing models such as stable diffusion through immunization of images. The demo is based on the <a href='https://github.com/MadryLab/photoguard' style='text-decoration: underline;' target='_blank'> Github </a> implementation provided by the authors.</p>
+
+''')
+
+gr.HTML('''
+
+<p align="center"><img src="https://raw.githubusercontent.com/MadryLab/photoguard/main/assets/hero_fig.PNG" style="width:60%"/></p>
+
+''')
+
+gr.HTML('''
+
+<p style="margin-bottom: 10px; font-size: 94%"> A malevolent actor might download
+
+photos of people posted online and edit them maliciously using an off-the-shelf diffusion model. The adversary
+
+describes via a textual prompt the desired changes and then uses a diffusion model to generate a realistic
+
+image that matches the prompt (similar to the top row in the image). By immunizing the original image before the adversary can access it,
+
+we disrupt their ability to successfully perform such edits forcing them to generate unrealistic images (similar to the bottom row in the image). For a more detailed explanation, please read the accompanying <a href='https://arxiv.org/abs/2302.06588' style='text-decoration: underline;' target='_blank'> Paper </a> or <a href='https://gradientscience.org/photoguard/' style='text-decoration: underline;' target='_blank'> Blogpost </a>
+
+''')
+
+
+
+with gr.Column(): 
+
+with gr.Tab("Simple Image to Image"):
+
+input_image_img2img = gr.Image(type="pil", label = "Source Image")
+
+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.Accordion("Advanced options", open=False):
+
+scale = gr.Slider(label="Guidance Scale", minimum=0.1, maximum=30.0, value=7.5, step=0.1)
+
+num_steps = gr.Slider(label="Number of Inference Steps", minimum=5, maximum=125, value=100, step=5)
+
+seed = gr.Slider(label="Seed", minimum=0, maximum=2147483647, step=1, randomize=True)
+
+
+
+with gr.Row():
+
+result_gallery = gr.Gallery(
+
+label="Generated images", show_label=False, elem_id="gallery"
+
+).style(grid=[2], height="auto")
+
+
+
+run_btn_img2img.click(process_image_img2img, inputs = [input_image_img2img,input_prompt_img2img, scale, num_steps, seed], outputs = [result_gallery])
+
+examples = gr.Examples(examples=examples_list,inputs = [input_image_img2img,input_prompt_img2img,scale, num_steps, seed], 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,scale, num_steps, seed], outputs = [result_gallery]) 
+
+
+
+
+
+demo.launch(debug=True)