Spaces:
Sleeping
Sleeping
| 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 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 process_image_img2img(raw_image,prompt): | |
| 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 = 9222 | |
| # SEED = np.random.randint(low=0, high=10000) | |
| # Play with these for improving generated image quality | |
| STRENGTH = 0.5 | |
| GUIDANCE = 7.5 | |
| NUM_STEPS = 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): | |
| # 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(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 = 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 = 7.5 | |
| # num_inference_steps = 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"]] | |
| 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 safe-guarding images against manipulation by ML-powerd 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> | |
| ''') | |
| 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.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], outputs = [result_gallery]) | |
| examples = gr.Examples(examples=examples_list,inputs = [input_image_img2img,input_prompt_img2img], 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], outputs = [result_gallery]) | |
| demo.launch(debug=True) |