Duplicate from hadisalman/photoguard

Co-authored-by: Hadi Salman <[email protected]>

.gitattributes

@@ -0,0 +1,36 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+*jpg filter=lfs diff=lfs merge=lfs -text
+*.jpg filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,14 @@
+---
+title: Photoguard
+emoji: 🛡
+colorFrom: gray
+colorTo: green
+sdk: gradio
+sdk_version: 3.19.1
+app_file: app.py
+pinned: false
+license: apache-2.0
+duplicated_from: hadisalman/photoguard
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,162 @@
+from io import BytesIO
+import requests
+import gradio as gr
+import requests
+import torch
+from tqdm import tqdm
+from PIL import Image, ImageOps
+from diffusers import StableDiffusionInpaintPipeline
+from torchvision.transforms import ToPILImage
+from utils import preprocess, prepare_mask_and_masked_image, recover_image, resize_and_crop
+
+gr.close_all()
+topil = ToPILImage()
+
+pipe_inpaint = StableDiffusionInpaintPipeline.from_pretrained(
+    "runwayml/stable-diffusion-inpainting",
+    revision="fp16",
+    torch_dtype=torch.float16,
+    safety_checker=None,
+)
+pipe_inpaint = pipe_inpaint.to("cuda")
+
+## Good params for editing that we used all over the paper --> decent quality and speed   
+GUIDANCE_SCALE = 7.5
+NUM_INFERENCE_STEPS = 100
+DEFAULT_SEED = 1234
+
+def pgd(X, targets, 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 - targets).norm()
+        pbar.set_description(f"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 get_target():
+    target_url = 'https://www.rtings.com/images/test-materials/2015/204_Gray_Uniformity.png'
+    response = requests.get(target_url)
+    target_image = Image.open(BytesIO(response.content)).convert("RGB")
+    target_image = target_image.resize((512, 512))
+    return target_image
+
+def immunize_fn(init_image, mask_image):
+    with torch.autocast('cuda'):
+        mask, X = prepare_mask_and_masked_image(init_image, mask_image)
+        X = X.half().cuda()
+        mask = mask.half().cuda()
+
+        targets = pipe_inpaint.vae.encode(preprocess(get_target()).half().cuda()).latent_dist.mean
+
+        adv_X = pgd(X, 
+                    targets = targets,
+                    model=pipe_inpaint.vae.encode, 
+                    criterion=torch.nn.MSELoss(), 
+                    clamp_min=-1, 
+                    clamp_max=1,
+                    eps=0.12, 
+                    step_size=0.01, 
+                    iters=200,
+                    mask=1-mask
+                   )
+
+        adv_X = (adv_X / 2 + 0.5).clamp(0, 1)
+        
+        adv_image = topil(adv_X[0]).convert("RGB")
+        adv_image = recover_image(adv_image, init_image, mask_image, background=True)
+        return adv_image        
+
+def run(image, prompt, seed, immunize=False):
+    if seed == '':
+        seed = DEFAULT_SEED
+    else:
+        seed = int(seed)
+    torch.manual_seed(seed)
+
+    init_image = Image.fromarray(image['image'])
+    init_image = resize_and_crop(init_image, (512,512))
+    mask_image = ImageOps.invert(Image.fromarray(image['mask']).convert('RGB'))
+    mask_image = resize_and_crop(mask_image, init_image.size)
+    
+    if immunize:
+        immunized_image = immunize_fn(init_image, mask_image)
+        
+    image_edited = pipe_inpaint(prompt=prompt, 
+                         image=init_image if not immunize else immunized_image, 
+                         mask_image=mask_image, 
+                         height = init_image.size[0],
+                         width = init_image.size[1],
+                         eta=1,
+                         guidance_scale=GUIDANCE_SCALE,
+                         num_inference_steps=NUM_INFERENCE_STEPS,
+                        ).images[0]
+        
+    image_edited = recover_image(image_edited, init_image, mask_image)
+    
+    if immunize:
+        return [(immunized_image, 'Immunized Image'), (image_edited, 'Edited After Immunization')]
+    else:
+        return [(image_edited, 'Edited Image')]
+
+
+demo = gr.Interface(fn=run, 
+                    inputs=[
+                        gr.ImageMask(label='Drawing tool to mask regions you want to keep, e.g. faces'),
+                        gr.Textbox(label='Prompt', placeholder='A photo of a man in a wedding'),
+                        gr.Textbox(label='Seed (Change to get different edits!)', placeholder=str(DEFAULT_SEED), visible=True),
+                        gr.Checkbox(label='Immunize', value=False),
+                    ], 
+                    cache_examples=False,
+                    outputs=[gr.Gallery(
+                            label="Generated images", 
+                            show_label=False, 
+                            elem_id="gallery").style(grid=[1,2], height="auto")],
+                    examples=[
+                    ['./images/hadi_and_trevor.jpg', 'man attending a wedding', '329357'],
+                    ['./images/trevor_2.jpg', 'two men in prison', '329357'],
+                    ['./images/elon_2.jpg', 'man in a metro station', '214213'],
+                    ],
+                    examples_per_page=20,
+                    allow_flagging='never',
+                    title="Interactive Demo: Immunize your Photos Against AI-powered Malicious Manipulation",
+                    description='''<u>Official</u> demo of our paper: <br>
+                    **Raising the Cost of Malicious AI-Powered Image Editing** <br>
+                    *[Hadi Salman](https://twitter.com/hadisalmanX)\*, [Alaa Khaddaj](https://twitter.com/Alaa_Khaddaj)\*, [Guillaume Leclerc](https://twitter.com/gpoleclerc)\*, [Andrew Ilyas](https://twitter.com/andrew_ilyas), [Aleksander Madry](https://twitter.com/aleks_madry)* <br>
+                    MIT &nbsp;&nbsp;[Paper](https://arxiv.org/abs/2302.06588) 
+                    &nbsp;&nbsp;[Blog post](https://gradientscience.org/photoguard/) 
+                    &nbsp;&nbsp;[![](https://badgen.net/badge/icon/GitHub?icon=github&label)](https://github.com/MadryLab/photoguard)
+                    <br />
+                    Below you can test our (encoder attack) immunization method for making images resistant to manipulation by Stable Diffusion. This immunization process forces the model to perform unrealistic edits.
+                    <br />
+**This is a research project and is not production-ready. See Section 5 in our paper for discussion on its limitations.**
+<details closed>
+<summary>Click for demo steps:</summary>
+
++ Upload an image (or select from the below examples!)
++ Mask (using the drawing tool) the parts of the image you want to maintain unedited (e.g., faces of people)
++ Add a prompt to edit the image accordingly (see examples below)
++ Play with the seed and click submit until you get a realistic edit that you are happy with (or use default seeds below)
+
+Now let's immunize your image and try again! 
++ Click on the "immunize" button, then submit.
++ You will get the immunized image (which looks identical to the original one) and the edited image, which is now hopefully unrealistic!                   
+</details>
+                    ''',
+                   )
+
+# demo.launch()
+demo.launch(server_name='0.0.0.0', share=False, server_port=7860, inline=False)

requirements.txt

@@ -0,0 +1,6 @@
+diffusers==0.10.2
+transformers
+scipy
+accelerate
+torchvision
+--extra-index-url https://download.pytorch.org/whl/cu113

utils.py

@@ -0,0 +1,65 @@
+import torch
+import numpy as np
+import torch
+from PIL import Image, ImageOps
+from torchvision.transforms import ToPILImage, ToTensor
+totensor = ToTensor()
+topil = ToPILImage()
+
+
+
+def resize_and_crop(img, size, crop_type="center"):
+    '''Resize and crop the image to the given size.'''
+    if crop_type == "top":
+        center = (0, 0)
+    elif crop_type == "center":
+        center = (0.5, 0.5)
+    else:
+        raise ValueError
+
+    resize = list(size)
+    if size[0] is None:
+        resize[0] = img.size[0]
+    if size[1] is None:
+        resize[1] = img.size[1]
+    return ImageOps.fit(img, resize, centering=center)
+
+
+def recover_image(image, init_image, mask, background=False):
+    image = totensor(image)
+    mask = totensor(mask)[0]
+    init_image = totensor(init_image)
+    
+    if background:
+        result = mask * init_image + (1 - mask) * image
+    else:
+        result = mask * image + (1 - mask) * init_image
+    return topil(result)
+
+
+def preprocess(image):
+    w, h = image.size
+    w, h = map(lambda x: x - x % 32, (w, h))  # resize to integer multiple of 32
+    image = image.resize((w, h), resample=Image.LANCZOS)
+    image = np.array(image).astype(np.float32) / 255.0
+    image = image[None].transpose(0, 3, 1, 2)
+    image = torch.from_numpy(image)
+    return 2.0 * image - 1.0
+
+
+def prepare_mask_and_masked_image(image, mask):
+    
+    image = np.array(image.convert("RGB"))
+    image = image[None].transpose(0, 3, 1, 2)
+    image = torch.from_numpy(image).to(dtype=torch.float32) / 127.5 - 1.0
+
+    mask = np.array(mask.convert("L"))
+    mask = mask.astype(np.float32) / 255.0
+    mask = mask[None, None]
+    mask[mask < 0.5] = 0
+    mask[mask >= 0.5] = 1
+    mask = torch.from_numpy(mask)
+
+    masked_image = image * (mask < 0.5)
+
+    return mask, masked_image