├── assets └── teaser.png ├── .gitignore ├── main_edit.py ├── evaluation ├── eval_psnr.py ├── eval_lpips.py ├── eval_ssim.py ├── eval_clip_i.py ├── eval_facial.py └── eval_clip_s.py ├── edit.py ├── utils.py ├── methods.py ├── defend.py ├── main_defend.py ├── README.md └── environment.yml /assets/teaser.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/taco-group/FaceLock/HEAD/assets/teaser.png -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | __pycache__ 2 | celeba 3 | edit_*/ 4 | facelock/ 5 | *.png 6 | !assets/*.png 7 | *.jpg 8 | *.csv 9 | *.out -------------------------------------------------------------------------------- /main_edit.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | from PIL import Image 3 | import torch 4 | from diffusers import StableDiffusionInstructPix2PixPipeline, EulerAncestralDiscreteScheduler 5 | from utils import edit_prompts 6 | import os 7 | 8 | if __name__ == "__main__": 9 | parser = argparse.ArgumentParser() 10 | parser.add_argument("--seed", type=int, default=-1) 11 | parser.add_argument("--src_dir", required=True, help="path to the directory of the src images") 12 | parser.add_argument("--edit_dir", required=True, help="path to the directory of the edited images") 13 | # edit configuration 14 | parser.add_argument("--num_inference_steps", type=int, default=50) 15 | parser.add_argument("--image_guidance_scale", type=float, default=1.5) 16 | parser.add_argument("--guidance_scale", type=float, default=7.5) 17 | args = parser.parse_args() 18 | 19 | guidance_scale = args.guidance_scale 20 | image_guidance_scale = args.image_guidance_scale 21 | num_inference_steps = args.num_inference_steps 22 | if args.seed == -1: 23 | import random 24 | seed = random.randint(0, 2**32 - 1) 25 | else: 26 | seed = args.seed 27 | model = StableDiffusionInstructPix2PixPipeline.from_pretrained( 28 | "timbrooks/instruct-pix2pix", 29 | torch_dtype=torch.float16, 30 | safety_checker=None, 31 | ).to("cuda") 32 | model.scheduler = EulerAncestralDiscreteScheduler.from_config(model.scheduler.config) 33 | 34 | src_dir = args.src_dir 35 | edit_dir = args.edit_dir 36 | image_files = sorted(os.listdir(src_dir)) 37 | 38 | for i, image_file in enumerate(image_files): 39 | image_path = os.path.join(src_dir, image_file) 40 | src_image = Image.open(image_path).convert("RGB") 41 | for idx, prompt in edit_prompts.items(): 42 | save_dir = os.path.join(edit_dir, f"seed{seed}", f"prompt{idx}") 43 | if not os.path.exists(save_dir): 44 | os.makedirs(save_dir) 45 | torch.manual_seed(seed) 46 | edit_image = model( 47 | prompt=prompt, 48 | image=src_image, 49 | num_inference_steps=num_inference_steps, 50 | image_guidance_scale=image_guidance_scale, 51 | guidance_scale=guidance_scale 52 | ).images[0] 53 | edit_image.save(os.path.join(save_dir, image_file)) 54 | if (i + 1) % 100 == 0: 55 | print(f"Edited [{i + 1}/{len(image_files)}]") -------------------------------------------------------------------------------- /evaluation/eval_psnr.py: -------------------------------------------------------------------------------- 1 | # Compute PSNR between the edits on protected and unprotected images 2 | import argparse 3 | import pandas as pd 4 | import os 5 | import numpy as np 6 | from PIL import Image 7 | import torchvision.transforms as T 8 | from torchmetrics.image import PeakSignalNoiseRatio 9 | from tqdm import trange 10 | import sys 11 | sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) 12 | from utils import edit_prompts 13 | 14 | psnr = PeakSignalNoiseRatio(data_range=1.0) 15 | 16 | if __name__ == "__main__": 17 | pd.set_option("display.max_rows", None) 18 | pd.set_option("display.max_columns", None) 19 | 20 | parser = argparse.ArgumentParser() 21 | parser.add_argument("--clean_edit_dir", required=True, type=str, help="path to the directory containing edits on the unprotected images") 22 | parser.add_argument("--defend_edit_dirs", required=True, nargs="+", help="path to the directory containing different attack budget subdirectories") 23 | parser.add_argument("--seed", required=True, type=int, help="the seed to evaluate on") 24 | args = parser.parse_args() 25 | 26 | clean_edit_dir = args.clean_edit_dir 27 | defend_edit_dirs = args.defend_edit_dirs 28 | prompt_num = len(edit_prompts) 29 | seed = args.seed 30 | for x in defend_edit_dirs: 31 | assert os.path.exists(x) 32 | 33 | result = [] 34 | 35 | for edit_dir in defend_edit_dirs: 36 | eps_dirs = sorted(os.listdir(edit_dir)) 37 | for eps_dir in eps_dirs: 38 | cur_method = os.path.join(edit_dir, eps_dir) 39 | psnr_dict = {"method": cur_method} 40 | psnr_scores = [] 41 | print(f"Processing {cur_method}") 42 | seed_dir = os.path.join(cur_method, f"seed{seed}") 43 | clean_seed_dir = os.path.join(clean_edit_dir, f"seed{seed}") 44 | for i in range(prompt_num): 45 | prompt_dir = os.path.join(seed_dir, f"prompt{i}") 46 | clean_prompt_dir = os.path.join(clean_seed_dir, f"prompt{i}") 47 | assert os.path.exists(prompt_dir) and os.path.exists(clean_prompt_dir) 48 | psnr_i = 0 49 | image_files = sorted(os.listdir(prompt_dir)) 50 | clean_image_files = sorted(os.listdir(clean_prompt_dir)) 51 | num = len(image_files) 52 | assert num == len(clean_image_files) 53 | for k in trange(num): 54 | edit_image = Image.open(os.path.join(prompt_dir, image_files[k])).convert("RGB") 55 | clean_edit_image = Image.open(os.path.join(clean_prompt_dir, clean_image_files[k])).convert("RGB") 56 | edit_tensor = T.ToTensor()(edit_image) 57 | clean_edit_tensor = T.ToTensor()(clean_edit_image) 58 | psnr_i += psnr(edit_tensor, clean_edit_tensor).item() 59 | psnr_i /= num 60 | psnr_scores.append(psnr_i) 61 | psnr_dict[f"prompt{i}"] = psnr_i 62 | 63 | psnr_dict["mean"] = np.mean(psnr_scores) 64 | result.append(psnr_dict) 65 | 66 | df = pd.DataFrame(result) 67 | print(df) 68 | df.to_csv("psnr_metric.csv", index=False) 69 | -------------------------------------------------------------------------------- /evaluation/eval_lpips.py: -------------------------------------------------------------------------------- 1 | # Compute LPIPS between the edits on protected and unprotected images 2 | import argparse 3 | import pandas as pd 4 | import torch 5 | import os 6 | import numpy as np 7 | from PIL import Image 8 | import torchvision.transforms as T 9 | import lpips 10 | from tqdm import trange 11 | import sys 12 | sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) 13 | from utils import edit_prompts 14 | 15 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 16 | lpips_fn = lpips.LPIPS(net='vgg').to(device) 17 | 18 | if __name__ == "__main__": 19 | pd.set_option("display.max_rows", None) 20 | pd.set_option("display.max_columns", None) 21 | 22 | parser = argparse.ArgumentParser() 23 | parser.add_argument("--clean_edit_dir", required=True, type=str, help="path to the directory containing edits on the unprotected images") 24 | parser.add_argument("--defend_edit_dirs", required=True, nargs="+", help="path to the directory containing different attack budget subdirectories") 25 | parser.add_argument("--seed", required=True, type=int, help="the seed to evaluate on") 26 | args = parser.parse_args() 27 | 28 | clean_edit_dir = args.clean_edit_dir 29 | defend_edit_dirs = args.defend_edit_dirs 30 | prompt_num = len(edit_prompts) 31 | seed = args.seed 32 | for x in defend_edit_dirs: 33 | assert os.path.exists(x) 34 | 35 | result = [] 36 | 37 | for edit_dir in defend_edit_dirs: 38 | eps_dirs = sorted(os.listdir(edit_dir)) 39 | for eps_dir in eps_dirs: 40 | cur_method = os.path.join(edit_dir, eps_dir) 41 | lpips_dict = {"method": cur_method} 42 | lpips_scores = [] 43 | print(f"Processing {cur_method}") 44 | seed_dir = os.path.join(cur_method, f"seed{seed}") 45 | clean_seed_dir = os.path.join(clean_edit_dir, f"seed{seed}") 46 | for i in range(prompt_num): 47 | prompt_dir = os.path.join(seed_dir, f"prompt{i}") 48 | clean_prompt_dir = os.path.join(clean_seed_dir, f"prompt{i}") 49 | assert os.path.exists(prompt_dir) and os.path.exists(clean_prompt_dir) 50 | lpips_i = 0 51 | image_files = sorted(os.listdir(prompt_dir)) 52 | clean_image_files = sorted(os.listdir(clean_prompt_dir)) 53 | num = len(image_files) 54 | assert num == len(clean_image_files) 55 | for k in trange(num): 56 | clean_edit_image = lpips.im2tensor(lpips.load_image(os.path.join(clean_prompt_dir, clean_image_files[k]))).to(device) 57 | edit_image = lpips.im2tensor(lpips.load_image(os.path.join(prompt_dir,image_files[k]))).to(device) 58 | cur_score = lpips_fn(clean_edit_image, edit_image).item() 59 | lpips_i += cur_score 60 | lpips_i /= num 61 | lpips_scores.append(lpips_i) 62 | lpips_dict[f"prompt{i}"] = lpips_i 63 | 64 | lpips_dict["mean"] = np.mean(lpips_scores) 65 | result.append(lpips_dict) 66 | 67 | df = pd.DataFrame(result) 68 | print(df) 69 | df.to_csv("lpips_metric.csv", index=False) 70 | -------------------------------------------------------------------------------- /evaluation/eval_ssim.py: -------------------------------------------------------------------------------- 1 | # Compute ssim between the edits on protected and unprotected images 2 | import argparse 3 | import pandas as pd 4 | import os 5 | import numpy as np 6 | from PIL import Image 7 | import torchvision.transforms as T 8 | from torchmetrics.image import StructuralSimilarityIndexMeasure 9 | from tqdm import trange 10 | import sys 11 | sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) 12 | from utils import edit_prompts 13 | 14 | ssim = StructuralSimilarityIndexMeasure(data_range=1.0) 15 | 16 | if __name__ == "__main__": 17 | pd.set_option("display.max_rows", None) 18 | pd.set_option("display.max_columns", None) 19 | 20 | parser = argparse.ArgumentParser() 21 | parser.add_argument("--clean_edit_dir", required=True, type=str, help="path to the directory containing edits on the unprotected images") 22 | parser.add_argument("--defend_edit_dirs", required=True, nargs="+", help="path to the directory containing different attack budget subdirectories") 23 | parser.add_argument("--seed", required=True, type=int, help="the seed to evaluate on") 24 | args = parser.parse_args() 25 | 26 | clean_edit_dir = args.clean_edit_dir 27 | defend_edit_dirs = args.defend_edit_dirs 28 | prompt_num = len(edit_prompts) 29 | seed = args.seed 30 | for x in defend_edit_dirs: 31 | assert os.path.exists(x) 32 | 33 | result = [] 34 | 35 | for edit_dir in defend_edit_dirs: 36 | eps_dirs = sorted(os.listdir(edit_dir)) 37 | for eps_dir in eps_dirs: 38 | cur_method = os.path.join(edit_dir, eps_dir) 39 | ssim_dict = {"method": cur_method} 40 | ssim_scores = [] 41 | print(f"Processing {cur_method}") 42 | seed_dir = os.path.join(cur_method, f"seed{seed}") 43 | clean_seed_dir = os.path.join(clean_edit_dir, f"seed{seed}") 44 | for i in range(prompt_num): 45 | prompt_dir = os.path.join(seed_dir, f"prompt{i}") 46 | clean_prompt_dir = os.path.join(clean_seed_dir, f"prompt{i}") 47 | assert os.path.exists(prompt_dir) and os.path.exists(clean_prompt_dir) 48 | ssim_i = 0 49 | image_files = sorted(os.listdir(prompt_dir)) 50 | clean_image_files = sorted(os.listdir(clean_prompt_dir)) 51 | num = len(image_files) 52 | assert num == len(clean_image_files) 53 | for k in trange(num): 54 | edit_image = Image.open(os.path.join(prompt_dir, image_files[k])).convert("RGB") 55 | clean_edit_image = Image.open(os.path.join(clean_prompt_dir, clean_image_files[k])).convert("RGB") 56 | edit_tensor = T.ToTensor()(edit_image).unsqueeze(0) 57 | clean_edit_tensor = T.ToTensor()(clean_edit_image).unsqueeze(0) 58 | ssim_i += ssim(edit_tensor, clean_edit_tensor).item() 59 | ssim_i /= num 60 | ssim_scores.append(ssim_i) 61 | ssim_dict[f"prompt{i}"] = ssim_i 62 | 63 | ssim_dict["mean"] = np.mean(ssim_scores) 64 | result.append(ssim_dict) 65 | 66 | df = pd.DataFrame(result) 67 | print(df) 68 | df.to_csv("ssim_metric.csv", index=False) 69 | -------------------------------------------------------------------------------- /edit.py: -------------------------------------------------------------------------------- 1 | from PIL import Image 2 | import torch 3 | import argparse 4 | from diffusers import StableDiffusionInstructPix2PixPipeline, EulerAncestralDiscreteScheduler, AutoPipelineForImage2Image 5 | import os 6 | 7 | def get_args_parser(): 8 | parser = argparse.ArgumentParser() 9 | 10 | # 1. image arguments 11 | parser.add_argument("--input_path", required=True, type=str, help="the path to the input image") 12 | 13 | # 2. edit model arguments 14 | parser.add_argument("--model", default="instruct-pix2pix", type=str, help="the model used to edit images [instruct-pix2pix/stable-diffusion]") 15 | parser.add_argument("--model_id", default="timbrooks/instruct-pix2pix", type=str, help="model id from hugging face for the model") 16 | parser.add_argument("--seed", default=-1, type=int, help="the seed used to control generation") 17 | parser.add_argument("--guidance_scale", default=7.5, type=float, help="the guidance scale for the textual prompt") 18 | parser.add_argument("--image_guidance_scale", default=1.5, type=float, help="the image guidance scale for the image in instruct-pix2pix") 19 | parser.add_argument("--strength", default=0.5, type=float, help="the strength value for the image in stable diffusion") 20 | parser.add_argument("--num_inference_steps", default=50, type=int, help="the number of denoising steps for image generation") 21 | parser.add_argument("--prompt", default="", type=str, help="the prompt used to edit the image") 22 | 23 | # 3. output 24 | parser.add_argument("--output_path", default=None, help="the path used to save the generated image") 25 | 26 | return parser 27 | 28 | def main(args): 29 | # 1. prepare the image 30 | src_image = Image.open(args.input_path).convert("RGB") 31 | 32 | # 2. prepare the edit model 33 | model = None 34 | if args.model == "stable-diffusion": 35 | model = AutoPipelineForImage2Image.from_pretrained( 36 | args.model_id, 37 | torch_dtype=torch.float16, 38 | safety_checker=None, 39 | ) 40 | elif args.model == "instruct-pix2pix": 41 | model = StableDiffusionInstructPix2PixPipeline.from_pretrained( 42 | args.model_id, 43 | torch_dtype=torch.float16, 44 | safety_checker=None, 45 | ) 46 | model.scheduler = EulerAncestralDiscreteScheduler.from_config(model.scheduler.config) 47 | else: 48 | raise ValueError(f"Invalid model '{args.model}'. Valid options are 'stable-diffusion' or 'instruct-pix2pix'.") 49 | 50 | model.to("cuda") 51 | 52 | # 3. edit the image 53 | if args.seed == -1: 54 | import random 55 | seed = random.randint(0, 2**32 - 1) 56 | else: 57 | seed = args.seed 58 | torch.manual_seed(seed) 59 | prompt = args.prompt 60 | if args.model == "stable-diffusion": 61 | edit_image = model( 62 | prompt=prompt, 63 | image=src_image, 64 | num_inference_steps=args.num_inference_steps, 65 | strength=args.strength, 66 | guidance_scale=args.guidance_scale, 67 | ).images[0] 68 | elif args.model == "instruct-pix2pix": 69 | edit_image = model( 70 | prompt=prompt, 71 | image=src_image, 72 | num_inference_steps=args.num_inference_steps, 73 | image_guidance_scale=args.image_guidance_scale, 74 | guidance_scale=args.guidance_scale, 75 | ).images[0] 76 | 77 | # 4. store the edited image 78 | edit_image.save(args.output_path) 79 | 80 | if __name__ == "__main__": 81 | parser = get_args_parser() 82 | args = parser.parse_args() 83 | 84 | if args.output_path is None: 85 | init_path = args.input_path 86 | directory, filename = os.path.split(init_path) 87 | new_path = os.path.join(directory, "edit_" + filename) 88 | args.output_path = new_path 89 | 90 | main(args) -------------------------------------------------------------------------------- /utils.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | import shutil 4 | from PIL import Image 5 | import numpy as np 6 | import torch 7 | import torchvision.transforms as T 8 | from torchvision.transforms import Compose, ToTensor, Normalize, ToPILImage 9 | from huggingface_hub import hf_hub_download 10 | from transformers import AutoModel 11 | import inspect 12 | 13 | # helpfer function to download huggingface repo and use model 14 | def download(repo_id, path, HF_TOKEN=None): 15 | os.makedirs(path, exist_ok=True) 16 | files_path = os.path.join(path, 'files.txt') 17 | if not os.path.exists(files_path): 18 | hf_hub_download(repo_id, 'files.txt', token=HF_TOKEN, local_dir=path, local_dir_use_symlinks=False) 19 | with open(os.path.join(path, 'files.txt'), 'r') as f: 20 | files = f.read().split('\n') 21 | for file in [f for f in files if f] + ['config.json', 'wrapper.py', 'model.safetensors']: 22 | full_path = os.path.join(path, file) 23 | if not os.path.exists(full_path): 24 | hf_hub_download(repo_id, file, token=HF_TOKEN, local_dir=path, local_dir_use_symlinks=False) 25 | 26 | 27 | # helpfer function to download huggingface repo and use model 28 | def load_model_from_local_path(path, HF_TOKEN=None): 29 | cwd = os.getcwd() 30 | os.chdir(path) 31 | sys.path.insert(0, path) 32 | model = AutoModel.from_pretrained(path, trust_remote_code=True, token=HF_TOKEN) 33 | os.chdir(cwd) 34 | sys.path.pop(0) 35 | return model 36 | 37 | 38 | # helpfer function to download huggingface repo and use model 39 | def load_model_by_repo_id(repo_id, save_path, HF_TOKEN=None, force_download=False): 40 | if force_download: 41 | if os.path.exists(save_path): 42 | shutil.rmtree(save_path) 43 | download(repo_id, save_path, HF_TOKEN) 44 | return load_model_from_local_path(save_path, HF_TOKEN) 45 | 46 | def pil_to_input(pil_image): 47 | trans = Compose([ToTensor(), Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])]) 48 | input = trans(pil_image).unsqueeze(0) 49 | return input 50 | 51 | def compute_score(input1, input2, aligner, fr_model, return_bbox=False): 52 | aligned_x1, orig_pred_ldmks1, aligned_ldmks1, score1, thetas1, normalized_bbox1 = aligner(input1) 53 | aligned_x2, orig_pred_ldmks2, aligned_ldmks2, score2, thetas2, normalized_bbox2 = aligner(input2) 54 | 55 | # recognize 56 | input_signature = inspect.signature(fr_model.model.net.forward) 57 | if input_signature.parameters.get('keypoints') is not None: 58 | feat1 = fr_model(aligned_x1, aligned_ldmks1) 59 | feat2 = fr_model(aligned_x2, aligned_ldmks2) 60 | else: 61 | feat1 = fr_model(aligned_x1) 62 | feat2 = fr_model(aligned_x2) 63 | 64 | # compute cosine similarity 65 | cossim = torch.nn.functional.cosine_similarity(feat1, feat2) 66 | if return_bbox: 67 | return cossim, normalized_bbox1, normalized_bbox2 68 | return cossim 69 | 70 | edit_prompts = { 71 | 0: "Turn the person's hair pink", 72 | 1: "Let the person turn bald", 73 | 2: "Let the person have a tatto", 74 | 3: "Let the person wear purple makeup", 75 | 4: "Let the person grow a mustach", 76 | 5: "Turn the person into a zombie", 77 | 6: "Change the skin color to Avatar blue", 78 | 7: "Add elf-like ears", 79 | 8: "Add large vampire fangs", 80 | 9: "Apply Goth style makeup", 81 | 10: "Let the person wear a police suit", 82 | 11: "Let the person wear a bowtie", 83 | 12: "Let the person wear a helmet", 84 | 13: "Let the person wear sunglasses", 85 | 14: "Let the person wear earrings", 86 | 15: "Let the person smoke a cigar", 87 | 16: "Place a headband in the hair", 88 | 17: "Place a tiara on the top of the head", 89 | 18: "Let it be snowy", 90 | 19: "Change the background to a beach", 91 | 20: "Add a city skyline background", 92 | 21: "Add a forest background", 93 | 22: "Change the background to a desert", 94 | 23: "Set the background in a library", 95 | 24: "Let the person stand under the moon" 96 | } 97 | 98 | def get_image_embeddings(images, clip_processor, clip_model, device="cuda"): 99 | inputs = clip_processor(images=images, return_tensors="pt").to(device) 100 | with torch.no_grad(): 101 | image_features = clip_model.get_image_features(**inputs) 102 | return image_features 103 | 104 | def get_text_embeddings(prompts, clip_processor, clip_model, device="cuda"): 105 | inputs = clip_processor(text=prompts, return_tensors="pt", padding=True, truncation=True).to(device) 106 | with torch.no_grad(): 107 | text_features = clip_model.get_text_features(**inputs) 108 | return text_features -------------------------------------------------------------------------------- /evaluation/eval_clip_i.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import pandas as pd 3 | import torch 4 | import os 5 | import numpy as np 6 | from PIL import Image 7 | import torchvision.transforms as T 8 | import torch.nn.functional as F 9 | from transformers import CLIPProcessor, CLIPModel 10 | from tqdm import trange 11 | import sys 12 | sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) 13 | from utils import edit_prompts, get_image_embeddings 14 | 15 | device = 'cuda' if torch.cuda.is_available() else 'cpu' 16 | clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32").to(device) 17 | clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32") 18 | 19 | if __name__ == "__main__": 20 | pd.set_option("display.max_rows", None) 21 | pd.set_option("display.max_columns", None) 22 | 23 | parser = argparse.ArgumentParser() 24 | parser.add_argument("--src_dir", required=True, type=str, help="path to the directory containing the source images") 25 | parser.add_argument("--clean_edit_dir", default=None, help="path to the directory containing edits on the unprotected images") 26 | parser.add_argument("--defend_edit_dirs", required=True, nargs="+", help="path to the directory containing different attack budget subdirectories") 27 | parser.add_argument("--seed", required=True, type=int, help="the seed to evaluate on") 28 | args = parser.parse_args() 29 | 30 | src_dir = args.src_dir 31 | src_image_files = sorted(os.listdir(src_dir)) 32 | num = len(src_image_files) 33 | defend_edit_dirs = args.defend_edit_dirs 34 | prompt_num = len(edit_prompts) 35 | seed = args.seed 36 | for x in defend_edit_dirs: 37 | assert os.path.exists(x) 38 | 39 | result = [] 40 | 41 | if args.clean_edit_dir is not None: 42 | clean_edit_dir = args.clean_edit_dir 43 | print("Processing clean") 44 | clip_dict = {"method": "clean"} 45 | clip_scores = [] 46 | seed_dir = os.path.join(clean_edit_dir, f"seed{seed}") 47 | for i in range(prompt_num): 48 | prompt_dir = os.path.join(seed_dir, f"prompt{i}") 49 | assert os.path.exists(prompt_dir) 50 | clip_i = 0 51 | edit_image_files = sorted(os.listdir(prompt_dir)) 52 | for k in trange(num): 53 | src_image = Image.open(os.path.join(src_dir, src_image_files[k])).convert("RGB") 54 | edit_image = Image.open(os.path.join(prompt_dir, edit_image_files[k])).convert("RGB") 55 | src_embeddings = get_image_embeddings(src_image, clip_processor, clip_model, device) 56 | edit_embeddings = get_image_embeddings(edit_image, clip_processor, clip_model, device) 57 | similarity_score = F.cosine_similarity(edit_embeddings, src_embeddings, dim=-1).mean().item() 58 | clip_i += similarity_score 59 | clip_i /= num 60 | clip_scores.append(clip_i) 61 | clip_dict[f"prompt{i}"] = clip_i 62 | 63 | clip_dict["mean"] = np.mean(clip_scores) 64 | result.append(clip_dict) 65 | 66 | df = pd.DataFrame(result) 67 | print(df) 68 | df.to_csv("clip_i_metric.csv", index=False) 69 | 70 | for edit_dir in defend_edit_dirs: 71 | eps_dirs = sorted(os.listdir(edit_dir)) 72 | for eps_dir in eps_dirs: 73 | cur_method = os.path.join(edit_dir, eps_dir) 74 | clip_dict = {"method": cur_method} 75 | clip_scores = [] 76 | print(f"Processing {cur_method}") 77 | seed_dir = os.path.join(cur_method, f"seed{seed}") 78 | for i in range(prompt_num): 79 | prompt_dir = os.path.join(seed_dir, f"prompt{i}") 80 | assert os.path.exists(prompt_dir) 81 | clip_i = 0 82 | edit_image_files = sorted(os.listdir(prompt_dir)) 83 | for k in trange(num): 84 | src_image = Image.open(os.path.join(src_dir, src_image_files[k])).convert("RGB") 85 | edit_image = Image.open(os.path.join(prompt_dir, edit_image_files[k])).convert("RGB") 86 | src_embeddings = get_image_embeddings(src_image, clip_processor, clip_model, device) 87 | edit_embeddings = get_image_embeddings(edit_image, clip_processor, clip_model, device) 88 | similarity_score = F.cosine_similarity(edit_embeddings, src_embeddings, dim=-1).mean().item() 89 | clip_i += similarity_score 90 | clip_i /= num 91 | clip_scores.append(clip_i) 92 | clip_dict[f"prompt{i}"] = clip_i 93 | 94 | clip_dict["mean"] = np.mean(clip_scores) 95 | result.append(clip_dict) 96 | 97 | df = pd.DataFrame(result) 98 | print(df) 99 | df.to_csv("clip_i_metric.csv", index=False) 100 | -------------------------------------------------------------------------------- /evaluation/eval_facial.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import pandas as pd 3 | import torch 4 | import os 5 | import numpy as np 6 | from PIL import Image 7 | import torchvision.transforms as T 8 | import torch.nn.functional as F 9 | from tqdm import trange 10 | import sys 11 | sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) 12 | from utils import edit_prompts, load_model_by_repo_id, compute_score, pil_to_input 13 | 14 | repo_id = 'minchul/cvlface_adaface_vit_base_kprpe_webface4m' 15 | aligner_id = 'minchul/cvlface_DFA_mobilenet' 16 | # load model 17 | device = 'cuda' if torch.cuda.is_available() else 'cpu' 18 | fr_model = load_model_by_repo_id(repo_id=repo_id, 19 | save_path=f'{os.environ["HF_HOME"]}/{repo_id}', 20 | HF_TOKEN=os.environ['HUGGINGFACE_HUB_TOKEN']).to(device) 21 | aligner = load_model_by_repo_id(repo_id=aligner_id, 22 | save_path=f'{os.environ["HF_HOME"]}/{aligner_id}', 23 | HF_TOKEN=os.environ['HUGGINGFACE_HUB_TOKEN']).to(device) 24 | 25 | if __name__ == "__main__": 26 | pd.set_option("display.max_rows", None) 27 | pd.set_option("display.max_columns", None) 28 | 29 | parser = argparse.ArgumentParser() 30 | parser.add_argument("--src_dir", required=True, type=str, help="path to the directory containing the source images") 31 | parser.add_argument("--clean_edit_dir", default=None, help="path to the directory containing edits on the unprotected images") 32 | parser.add_argument("--defend_edit_dirs", required=True, nargs="+", help="path to the directory containing different attack budget subdirectories") 33 | parser.add_argument("--seed", required=True, type=int, help="the seed to evaluate on") 34 | args = parser.parse_args() 35 | 36 | src_dir = args.src_dir 37 | src_image_files = sorted(os.listdir(src_dir)) 38 | num = len(src_image_files) 39 | defend_edit_dirs = args.defend_edit_dirs 40 | prompt_num = len(edit_prompts) 41 | seed = args.seed 42 | for x in defend_edit_dirs: 43 | assert os.path.exists(x) 44 | 45 | result = [] 46 | 47 | if args.clean_edit_dir is not None: 48 | clean_edit_dir = args.clean_edit_dir 49 | print("Processing clean") 50 | facial_dict = {"method": "clean"} 51 | facial_scores = [] 52 | seed_dir = os.path.join(clean_edit_dir, f"seed{seed}") 53 | for i in range(prompt_num): 54 | prompt_dir = os.path.join(seed_dir, f"prompt{i}") 55 | assert os.path.exists(prompt_dir) 56 | facial_i = 0 57 | edit_image_files = sorted(os.listdir(prompt_dir)) 58 | for k in trange(num): 59 | src_image = Image.open(os.path.join(src_dir, src_image_files[k])).convert("RGB") 60 | edit_image = Image.open(os.path.join(prompt_dir, edit_image_files[k])).convert("RGB") 61 | similarity_score = compute_score(pil_to_input(src_image).cuda(), pil_to_input(edit_image).cuda(), aligner, fr_model).item() 62 | facial_i += similarity_score 63 | facial_i /= num 64 | facial_scores.append(facial_i) 65 | facial_dict[f"prompt{i}"] = facial_i 66 | 67 | facial_dict["mean"] = np.mean(facial_scores) 68 | result.append(facial_dict) 69 | 70 | df = pd.DataFrame(result) 71 | print(df) 72 | df.to_csv("facial_metric.csv", index=False) 73 | 74 | for edit_dir in defend_edit_dirs: 75 | eps_dirs = sorted(os.listdir(edit_dir)) 76 | for eps_dir in eps_dirs: 77 | cur_method = os.path.join(edit_dir, eps_dir) 78 | facial_dict = {"method": cur_method} 79 | facial_scores = [] 80 | print(f"Processing {cur_method}") 81 | seed_dir = os.path.join(cur_method, f"seed{seed}") 82 | for i in range(prompt_num): 83 | prompt_dir = os.path.join(seed_dir, f"prompt{i}") 84 | assert os.path.exists(prompt_dir) 85 | facial_i = 0 86 | edit_image_files = sorted(os.listdir(prompt_dir)) 87 | for k in trange(num): 88 | src_image = Image.open(os.path.join(src_dir, src_image_files[k])).convert("RGB") 89 | edit_image = Image.open(os.path.join(prompt_dir, edit_image_files[k])).convert("RGB") 90 | similarity_score = compute_score(pil_to_input(src_image).cuda(), pil_to_input(edit_image).cuda(), aligner, fr_model).item() 91 | facial_i += similarity_score 92 | facial_i /= num 93 | facial_scores.append(facial_i) 94 | facial_dict[f"prompt{i}"] = facial_i 95 | 96 | facial_dict["mean"] = np.mean(facial_scores) 97 | result.append(facial_dict) 98 | 99 | df = pd.DataFrame(result) 100 | print(df) 101 | df.to_csv("facial_metric.csv", index=False) 102 | -------------------------------------------------------------------------------- /evaluation/eval_clip_s.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import pandas as pd 3 | import torch 4 | import os 5 | import numpy as np 6 | from PIL import Image 7 | import torchvision.transforms as T 8 | import torch.nn.functional as F 9 | from transformers import CLIPProcessor, CLIPModel 10 | from tqdm import trange 11 | import sys 12 | sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) 13 | from utils import edit_prompts, get_image_embeddings, get_text_embeddings 14 | 15 | def get_image_embeddings(images, clip_processor, clip_model, device="cuda"): 16 | inputs = clip_processor(images=images, return_tensors="pt").to(device) 17 | with torch.no_grad(): 18 | image_features = clip_model.get_image_features(**inputs) 19 | return image_features 20 | 21 | device = 'cuda' if torch.cuda.is_available() else 'cpu' 22 | clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32").to(device) 23 | clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32") 24 | 25 | if __name__ == "__main__": 26 | pd.set_option("display.max_rows", None) 27 | pd.set_option("display.max_columns", None) 28 | 29 | parser = argparse.ArgumentParser() 30 | parser.add_argument("--src_dir", required=True, type=str, help="path to the directory containing the source images") 31 | parser.add_argument("--clean_edit_dir", default=None, help="path to the directory containing edits on the unprotected images") 32 | parser.add_argument("--defend_edit_dirs", required=True, nargs="+", help="path to the directory containing different attack budget subdirectories") 33 | parser.add_argument("--seed", required=True, type=int, help="the seed to evaluate on") 34 | args = parser.parse_args() 35 | 36 | src_dir = args.src_dir 37 | src_image_files = sorted(os.listdir(src_dir)) 38 | num = len(src_image_files) 39 | defend_edit_dirs = args.defend_edit_dirs 40 | prompt_num = len(edit_prompts) 41 | seed = args.seed 42 | for x in defend_edit_dirs: 43 | assert os.path.exists(x) 44 | 45 | result = [] 46 | 47 | if args.clean_edit_dir is not None: 48 | clean_edit_dir = args.clean_edit_dir 49 | print("Processing clean") 50 | clip_dict = {"method": "clean"} 51 | clip_scores = [] 52 | seed_dir = os.path.join(clean_edit_dir, f"seed{seed}") 53 | for i in range(prompt_num): 54 | prompt_dir = os.path.join(seed_dir, f"prompt{i}") 55 | cur_prompt = edit_prompts[i] 56 | text_embeddings = get_text_embeddings(cur_prompt, clip_processor, clip_model, device) 57 | assert os.path.exists(prompt_dir) 58 | clip_i = 0 59 | edit_image_files = sorted(os.listdir(prompt_dir)) 60 | for k in trange(num): 61 | src_image = Image.open(os.path.join(src_dir, src_image_files[k])).convert("RGB") 62 | edit_image = Image.open(os.path.join(prompt_dir, edit_image_files[k])).convert("RGB") 63 | src_embeddings = get_image_embeddings(src_image, clip_processor, clip_model, device) 64 | edit_embeddings = get_image_embeddings(edit_image, clip_processor, clip_model, device) 65 | delta_embeddings = edit_embeddings - src_embeddings 66 | similarity_score = F.cosine_similarity(delta_embeddings, text_embeddings).item() 67 | clip_i += similarity_score 68 | clip_i /= num 69 | clip_scores.append(clip_i) 70 | clip_dict[f"prompt{i}"] = clip_i 71 | 72 | clip_dict["mean"] = np.mean(clip_scores) 73 | result.append(clip_dict) 74 | 75 | df = pd.DataFrame(result) 76 | print(df) 77 | df.to_csv("clip_s_metric.csv", index=False) 78 | 79 | for edit_dir in defend_edit_dirs: 80 | eps_dirs = sorted(os.listdir(edit_dir)) 81 | for eps_dir in eps_dirs: 82 | cur_method = os.path.join(edit_dir, eps_dir) 83 | clip_dict = {"method": cur_method} 84 | clip_scores = [] 85 | print(f"Processing {cur_method}") 86 | seed_dir = os.path.join(cur_method, f"seed{seed}") 87 | for i in range(prompt_num): 88 | prompt_dir = os.path.join(seed_dir, f"prompt{i}") 89 | cur_prompt = edit_prompts[i] 90 | text_embeddings = get_text_embeddings(cur_prompt, clip_processor, clip_model, device) 91 | assert os.path.exists(prompt_dir) 92 | clip_i = 0 93 | edit_image_files = sorted(os.listdir(prompt_dir)) 94 | for k in trange(num): 95 | src_image = Image.open(os.path.join(src_dir, src_image_files[k])).convert("RGB") 96 | edit_image = Image.open(os.path.join(prompt_dir, edit_image_files[k])).convert("RGB") 97 | src_embeddings = get_image_embeddings(src_image, clip_processor, clip_model, device) 98 | edit_embeddings = get_image_embeddings(edit_image, clip_processor, clip_model, device) 99 | delta_embeddings = edit_embeddings - src_embeddings 100 | similarity_score = F.cosine_similarity(delta_embeddings, text_embeddings).item() 101 | clip_i += similarity_score 102 | clip_i /= num 103 | clip_scores.append(clip_i) 104 | clip_dict[f"prompt{i}"] = clip_i 105 | 106 | clip_dict["mean"] = np.mean(clip_scores) 107 | result.append(clip_dict) 108 | 109 | df = pd.DataFrame(result) 110 | print(df) 111 | df.to_csv("clip_s_metric.csv", index=False) 112 | -------------------------------------------------------------------------------- /methods.py: -------------------------------------------------------------------------------- 1 | import os 2 | import torch 3 | import torch.nn as nn 4 | from tqdm import tqdm 5 | import torch.optim as optim 6 | import torchvision.transforms as T 7 | import torch.nn.functional as F 8 | import numpy as np 9 | from utils import compute_score 10 | import pdb 11 | 12 | # CW L2 attack 13 | def cw_l2_attack(X, model, c=0.1, lr=0.01, iters=100, targeted=False): 14 | encoder = model.vae.encode 15 | clean_latents = encoder(X).latent_dist.mean 16 | 17 | def f(x): 18 | latents = encoder(x).latent_dist.mean 19 | if targeted: 20 | return latents.norm() 21 | else: 22 | return -torch.norm(latents - clean_latents.detach(), p=2, dim=-1) 23 | 24 | w = torch.zeros_like(X, requires_grad=True).cuda() 25 | pbar = tqdm(range(iters)) 26 | optimizer = optim.Adam([w], lr=lr) 27 | 28 | for step in pbar: 29 | a = 1/2*(nn.Tanh()(w) + 1) 30 | 31 | loss1 = nn.MSELoss(reduction='sum')(a, X) 32 | loss2 = torch.sum(c*f(a)) 33 | 34 | cost = loss1 + loss2 35 | pbar.set_description(f"Loss: {cost.item():.5f} | loss1: {loss1.item():.5f} | loss2: {loss2.item():.5f}") 36 | # pdb.set_trace() 37 | 38 | optimizer.zero_grad() 39 | cost.backward() 40 | optimizer.step() 41 | 42 | X_adv = 1/2*(nn.Tanh()(w) + 1) 43 | return X_adv 44 | 45 | # Encoder attack - Targeted / Untargeted 46 | def encoder_attack(X, model, eps=0.03, step_size=0.01, iters=100, clamp_min=-1, clamp_max=1, targeted=False): 47 | """ 48 | Processing encoder attack using l_inf norm 49 | Params: 50 | X - image tensor we hope to protect 51 | model - the targeted edit model 52 | eps - attack budget 53 | step_size - attack step size 54 | iters - attack iterations 55 | clamp_min - min value for the image pixels 56 | clamp_max - max value for the image pixels 57 | Return: 58 | X_adv - image tensor for the protected image 59 | """ 60 | encoder = model.vae.encode 61 | X_adv = torch.clamp(X.clone().detach() + (torch.rand(*X.shape)*2*eps-eps).half().cuda(), min=clamp_min, max=clamp_max) 62 | if not targeted: 63 | loss_fn = nn.MSELoss() 64 | clean_latent = encoder(X).latent_dist.mean 65 | pbar = tqdm(range(iters)) 66 | for i in pbar: 67 | actual_step_size = step_size - (step_size - step_size / 100) / iters * i 68 | 69 | X_adv.requires_grad_(True) 70 | latent = encoder(X_adv).latent_dist.mean 71 | if targeted: 72 | loss = latent.norm() 73 | grad, = torch.autograd.grad(loss, [X_adv]) 74 | X_adv = X_adv - grad.detach().sign() * actual_step_size 75 | else: 76 | loss = loss_fn(latent, clean_latent) 77 | grad, = torch.autograd.grad(loss, [X_adv]) 78 | X_adv = X_adv + grad.detach().sign() * actual_step_size 79 | 80 | pbar.set_description(f"[Running attack]: Loss {loss.item():.5f} | step size: {actual_step_size:.4}") 81 | 82 | X_adv = torch.minimum(torch.maximum(X_adv, X - eps), X + eps) 83 | X_adv.data = torch.clamp(X_adv, min=clamp_min, max=clamp_max) 84 | X_adv.grad = None 85 | 86 | pbar.set_postfix(norm_2=(X_adv - X).norm().item(), norm_inf=(X_adv - X).abs().max().item()) 87 | 88 | return X_adv 89 | 90 | def vae_attack(X, model, eps=0.03, step_size=0.01, iters=100, clamp_min=-1, clamp_max=1): 91 | """ 92 | Processing encoder attack using l_inf norm 93 | Params: 94 | X - image tensor we hope to protect 95 | model - the targeted edit model 96 | eps - attack budget 97 | step_size - attack step size 98 | iters - attack iterations 99 | clamp_min - min value for the image pixels 100 | clamp_max - max value for the image pixels 101 | Return: 102 | X_adv - image tensor for the protected image 103 | """ 104 | vae = model.vae 105 | X_adv = torch.clamp(X.clone().detach() + (torch.rand(*X.shape)*2*eps-eps).half().cuda(), min=clamp_min, max=clamp_max) 106 | pbar = tqdm(range(iters)) 107 | for i in pbar: 108 | actual_step_size = step_size - (step_size - step_size / 100) / iters * i 109 | 110 | X_adv.requires_grad_() 111 | image = vae(X_adv).sample 112 | 113 | loss = (image).norm() 114 | grad, = torch.autograd.grad(loss, [X_adv]) 115 | X_adv = X_adv - grad.detach().sign() * actual_step_size 116 | 117 | pbar.set_description(f"[Running attack]: Loss {loss.item():.5f} | step size: {actual_step_size:.4}") 118 | 119 | X_adv = torch.minimum(torch.maximum(X_adv, X - eps), X + eps) 120 | X_adv.data = torch.clamp(X_adv, min=clamp_min, max=clamp_max) 121 | X_adv.grad = None 122 | 123 | return X_adv 124 | 125 | def facelock(X, model, aligner, fr_model, lpips_fn, eps=0.03, step_size=0.01, iters=100, clamp_min=-1, clamp_max=1): 126 | X_adv = torch.clamp(X.clone().detach() + (torch.rand(*X.shape)*2*eps-eps).to(X.device), min=clamp_min, max=clamp_max).half() 127 | pbar = tqdm(range(iters)) 128 | 129 | vae = model.vae 130 | X_adv.requires_grad_(True) 131 | clean_latent = vae.encode(X).latent_dist.mean 132 | 133 | for i in pbar: 134 | # actual_step_size = step_size 135 | actual_step_size = step_size - (step_size - step_size / 100) / iters * i 136 | 137 | latent = vae.encode(X_adv).latent_dist.mean 138 | image = vae.decode(latent).sample.clip(-1, 1) 139 | 140 | loss_cvl = compute_score(image.float(), X.float(), aligner=aligner, fr_model=fr_model) 141 | loss_encoder = F.mse_loss(latent, clean_latent) 142 | loss_lpips = lpips_fn(image, X) 143 | loss = -loss_cvl * (1 if i >= iters * 0.35 else 0.0) + loss_encoder * 0.2 + loss_lpips * (1 if i > iters * 0.25 else 0.0) 144 | grad, = torch.autograd.grad(loss, [X_adv]) 145 | X_adv = X_adv + grad.detach().sign() * actual_step_size 146 | 147 | X_adv = torch.minimum(torch.maximum(X_adv, X - eps), X + eps) 148 | X_adv.data = torch.clamp(X_adv, min=clamp_min, max=clamp_max) 149 | X_adv.grad = None 150 | 151 | pbar.set_postfix(loss_cvl=loss_cvl.item(), loss_encoder=loss_encoder.item(), loss_lpips=loss_lpips.item(), loss=loss.item()) 152 | 153 | return X_adv -------------------------------------------------------------------------------- /defend.py: -------------------------------------------------------------------------------- 1 | import os 2 | from PIL import Image 3 | import numpy as np 4 | import torch 5 | import torchvision 6 | import lpips 7 | from utils import load_model_by_repo_id, pil_to_input 8 | from methods import cw_l2_attack, vae_attack, encoder_attack, facelock 9 | import argparse 10 | from diffusers import StableDiffusionInstructPix2PixPipeline, StableDiffusionImg2ImgPipeline, EulerAncestralDiscreteScheduler 11 | import pdb 12 | 13 | def get_args_parser(): 14 | parser = argparse.ArgumentParser() 15 | 16 | # 1. image arguments 17 | parser.add_argument("--input_path", required=True, type=str, help="the path to the image you hope to protect") 18 | 19 | # 2. target model arguments 20 | parser.add_argument("--target_model", default="instruct-pix2pix", type=str, help="the target image editing model [instruct-pix2pix/stable-diffusion]") 21 | parser.add_argument("--model_id", default="timbrooks/instruct-pix2pix", type=str, help="model id from hugging face for the model") 22 | 23 | # 3. attack arguments 24 | parser.add_argument("--defend_method", required=True, type=str, help="the chosen attack method between [encoder/vae/cw/facelock]") 25 | parser.add_argument("--attack_budget", default=0.03, type=float, help="the attack budget") 26 | parser.add_argument("--step_size", default=0.01, type=float, help="the attack step size") 27 | parser.add_argument("--num_iters", default=100, type=int, help="the number of attack iterations") 28 | parser.add_argument("--targeted", default=True, action='store_true', help="targeted (towards 0 tensor) attack") 29 | parser.add_argument("--untargeted", action='store_false', dest='targeted', help="untargeted attack") 30 | 31 | # 3.1 cw attack other arguments 32 | parser.add_argument("--c", default=0.03, type=float, help="the constant ratio used in cw attack") 33 | parser.add_argument("--lr", default=0.03, type=float, help="the learning rate for the optimizer used in cw attack") 34 | 35 | # 4. output arguments 36 | parser.add_argument("--output_path", default=None, type=str, help="the output path the protected images") 37 | return parser 38 | 39 | def process_encoder_attack(X, model, args): 40 | with torch.autocast("cuda"): 41 | X_adv = encoder_attack( 42 | X=X, 43 | model=model, 44 | eps=args.attack_budget, 45 | step_size=args.step_size, 46 | iters=args.num_iters, 47 | clamp_min=-1, 48 | clamp_max=1, 49 | targeted=args.targeted, 50 | ) 51 | return X_adv 52 | 53 | def process_vae_attack(X, model, args): 54 | with torch.autocast("cuda"): 55 | X_adv = vae_attack( 56 | X=X, 57 | model=model, 58 | eps=args.attack_budget, 59 | step_size=args.step_size, 60 | iters=args.num_iters, 61 | clamp_min=-1, 62 | clamp_max=1, 63 | ) 64 | return X_adv 65 | 66 | def process_cw_attack(X, model, args): 67 | X_adv = cw_l2_attack( 68 | X=X, 69 | model=model, 70 | c=args.c, 71 | lr=args.lr, 72 | iters=args.num_iters, 73 | ) 74 | delta = X_adv - X 75 | delta_clip = delta.clip(-args.attack_budget, args.attack_budget) 76 | X_adv = (X + delta_clip).clip(0, 1) 77 | return X_adv 78 | 79 | def process_facelock(X, model, args): 80 | fr_id = 'minchul/cvlface_adaface_vit_base_kprpe_webface4m' 81 | aligner_id = 'minchul/cvlface_DFA_mobilenet' 82 | device = 'cuda' 83 | fr_model = load_model_by_repo_id(repo_id=fr_id, 84 | save_path=f'{os.environ["HF_HOME"]}/{fr_id}', 85 | HF_TOKEN=os.environ['HUGGINGFACE_HUB_TOKEN']).to(device) 86 | aligner = load_model_by_repo_id(repo_id=aligner_id, 87 | save_path=f'{os.environ["HF_HOME"]}/{aligner_id}', 88 | HF_TOKEN=os.environ['HUGGINGFACE_HUB_TOKEN']).to(device) 89 | lpips_fn = lpips.LPIPS(net="vgg").to(device) 90 | 91 | with torch.autocast("cuda"): 92 | X_adv = facelock( 93 | X=X, 94 | model=model, 95 | aligner=aligner, 96 | fr_model=fr_model, 97 | lpips_fn=lpips_fn, 98 | eps=args.attack_budget, 99 | step_size=args.step_size, 100 | iters=args.num_iters, 101 | clamp_min=-1, 102 | clamp_max=1, 103 | ) 104 | return X_adv 105 | 106 | def main(args): 107 | # 1. prepare the image 108 | init_image = Image.open(args.input_path).convert("RGB") 109 | to_tensor = torchvision.transforms.ToTensor() 110 | if args.defend_method != "cw": 111 | X = pil_to_input(init_image).cuda().half() 112 | else: 113 | X = to_tensor(init_image).cuda().unsqueeze(0) # perform cw attack using torch.float32 114 | 115 | # 2. prepare the targeted model 116 | model = None 117 | if args.target_model == "stable-diffusion": 118 | model = StableDiffusionImg2ImgPipeline.from_pretrained( 119 | pretrained_model_name_or_path=args.model_id, 120 | torch_dtype=torch.float16 if args.defend_method != "cw" else torch.float32, 121 | safety_checker=None, 122 | ) 123 | elif args.target_model == "instruct-pix2pix": 124 | model = StableDiffusionInstructPix2PixPipeline.from_pretrained( 125 | pretrained_model_name_or_path=args.model_id, 126 | torch_dtype=torch.float16 if args.defend_method != "cw" else torch.float32, 127 | safety_checker=None, 128 | ) 129 | model.scheduler = EulerAncestralDiscreteScheduler.from_config(model.scheduler.config) 130 | else: 131 | raise ValueError(f"Invalid target_model '{args.target_model}'. Valid options are 'stable-diffusion' or 'instruct-pix2pix'.") 132 | 133 | model.to("cuda") 134 | 135 | # 3. set up defend 136 | defend_fn = None 137 | if args.defend_method == "encoder": 138 | defend_fn = process_encoder_attack 139 | elif args.defend_method == "vae": 140 | defend_fn = process_vae_attack 141 | elif args.defend_method == "cw": 142 | defend_fn = process_cw_attack 143 | elif args.defend_method == "facelock": 144 | defend_fn = process_facelock 145 | else: 146 | raise ValueError(f"Invalid defend_method '{args.defend_method}'. Valid options are 'encoder', 'vae', 'cw', or 'facelock'.") 147 | 148 | # 4. process defend 149 | X_adv = defend_fn(X, model, args) 150 | 151 | # 5. convert back to image and store it 152 | to_pil = torchvision.transforms.ToPILImage() 153 | if args.defend_method != "cw": 154 | X_adv = (X_adv / 2 + 0.5).clamp(0, 1) 155 | protected_image = to_pil(X_adv[0]).convert("RGB") 156 | protected_image.save(args.output_path) 157 | 158 | if __name__ == "__main__": 159 | parser = get_args_parser() 160 | args = parser.parse_args() 161 | 162 | main(args) -------------------------------------------------------------------------------- /main_defend.py: -------------------------------------------------------------------------------- 1 | import os 2 | from PIL import Image 3 | import numpy as np 4 | import torch 5 | import torchvision 6 | import lpips 7 | from utils import load_model_by_repo_id, pil_to_input 8 | from methods import cw_l2_attack, vae_attack, encoder_attack, facelock 9 | import argparse 10 | from diffusers import StableDiffusionInstructPix2PixPipeline, StableDiffusionImg2ImgPipeline, EulerAncestralDiscreteScheduler 11 | import pdb 12 | 13 | def get_args_parser(): 14 | parser = argparse.ArgumentParser() 15 | 16 | # 1. image arguments 17 | parser.add_argument("--image_dir", required=True, type=str, help="the path to the images you hope to protect") 18 | 19 | # 2. target model arguments 20 | parser.add_argument("--target_model", default="instruct-pix2pix", type=str, help="the target image editing model [instruct-pix2pix/stable-diffusion]") 21 | parser.add_argument("--model_id", default="timbrooks/instruct-pix2pix", type=str, help="model id from hugging face for the model") 22 | 23 | # 3. attack arguments 24 | parser.add_argument("--defend_method", required=True, type=str, help="the chosen attack method between [encoder/vae/cw/facelock]") 25 | parser.add_argument("--attack_budget", default=0.03, type=float, help="the attack budget") 26 | parser.add_argument("--step_size", default=0.01, type=float, help="the attack step size") 27 | parser.add_argument("--clamp_min", default=-1, type=float, help="min value for the image pixels") 28 | parser.add_argument("--clamp_max", default=1, type=float, help="max value for the image pixels") 29 | parser.add_argument("--num_iters", default=100, type=int, help="the number of attack iterations") 30 | parser.add_argument("--targeted", default=True, action='store_true', help="targeted (towards 0 tensor) attack") 31 | parser.add_argument("--untargeted", action='store_false', dest='targeted', help="untargeted attack") 32 | 33 | # 3.1 cw attack other arguments 34 | parser.add_argument("--c", default=0.03, type=float, help="the constant ratio used in cw attack") 35 | parser.add_argument("--lr", default=0.03, type=float, help="the learning rate for the optimizer used in cw attack") 36 | 37 | # 4. output arguments 38 | parser.add_argument("--output_dir", required=True, type=str, help="the to the output directory") 39 | return parser 40 | 41 | def main(args): 42 | # 1. prepare the images 43 | image_dir = args.image_dir 44 | image_files = sorted(os.listdir(image_dir)) 45 | 46 | # 2. prepare the targeted model 47 | model = None 48 | if args.target_model == "stable-diffusion": 49 | model = StableDiffusionImg2ImgPipeline.from_pretrained( 50 | pretrained_model_name_or_path=args.model_id, 51 | torch_dtype=torch.float16 if args.defend_method != "cw" else torch.float32, 52 | safety_checker=None, 53 | ) 54 | elif args.target_model == "instruct-pix2pix": 55 | model = StableDiffusionInstructPix2PixPipeline.from_pretrained( 56 | pretrained_model_name_or_path=args.model_id, 57 | torch_dtype=torch.float16 if args.defend_method != "cw" else torch.float32, 58 | safety_checker=None, 59 | ) 60 | model.scheduler = EulerAncestralDiscreteScheduler.from_config(model.scheduler.config) 61 | else: 62 | raise ValueError(f"Invalid target_model '{args.target_model}'. Valid options are 'stable-diffusion' or 'instruct-pix2pix'.") 63 | 64 | model.to("cuda") 65 | 66 | if args.defend_method == "facelock": 67 | fr_id = 'minchul/cvlface_adaface_vit_base_kprpe_webface4m' 68 | aligner_id = 'minchul/cvlface_DFA_mobilenet' 69 | device = 'cuda' 70 | fr_model = load_model_by_repo_id(repo_id=fr_id, 71 | save_path=f'{os.environ["HF_HOME"]}/{fr_id}', 72 | HF_TOKEN=os.environ['HUGGINGFACE_HUB_TOKEN']).to(device) 73 | aligner = load_model_by_repo_id(repo_id=aligner_id, 74 | save_path=f'{os.environ["HF_HOME"]}/{aligner_id}', 75 | HF_TOKEN=os.environ['HUGGINGFACE_HUB_TOKEN']).to(device) 76 | lpips_fn = lpips.LPIPS(net="vgg").to(device) 77 | 78 | # 3. set up and process defend 79 | to_pil = torchvision.transforms.ToPILImage() 80 | to_tensor = torchvision.transforms.ToTensor() 81 | total_num = len(image_files) 82 | for idx, image_file in enumerate(image_files): 83 | if (idx + 1) % 100 == 0: 84 | print(f"Processing {idx + 1}/{total_num}") 85 | init_image = Image.open(os.path.join(image_dir, image_file)).convert("RGB") 86 | 87 | # process the input 88 | if args.defend_method == "cw": 89 | X = to_tensor(init_image).cuda().unsqueeze(0) 90 | else: 91 | if args.clamp_min == -1: 92 | X = pil_to_input(init_image).cuda().half() 93 | else: 94 | X = to_tensor(init_image).half().cuda().unsqueeze(0) 95 | 96 | # defend 97 | if args.defend_method == "encoder": 98 | with torch.autocast("cuda"): 99 | X_adv = encoder_attack( 100 | X=X, 101 | model=model, 102 | eps=args.attack_budget, 103 | step_size=args.step_size, 104 | iters=args.num_iters, 105 | clamp_min=-1, 106 | clamp_max=1, 107 | targeted=args.targeted, 108 | ) 109 | elif args.defend_method == "vae": 110 | with torch.autocast("cuda"): 111 | X_adv = vae_attack( 112 | X=X, 113 | model=model, 114 | eps=args.attack_budget, 115 | step_size=args.step_size, 116 | iters=args.num_iters, 117 | clamp_min=-1, 118 | clamp_max=1, 119 | ) 120 | elif args.defend_method == "cw": 121 | X_adv = cw_l2_attack( 122 | X=X, 123 | model=model, 124 | c=args.c, 125 | lr=args.lr, 126 | iters=args.num_iters, 127 | ) 128 | delta = X_adv - X 129 | delta_clip = delta.clip(-args.attack_budget, args.attack_budget) 130 | X_adv = (X + delta_clip).clip(0, 1) 131 | elif args.defend_method == "facelock": 132 | with torch.autocast("cuda"): 133 | X_adv = facelock( 134 | X=X, 135 | model=model, 136 | aligner=aligner, 137 | fr_model=fr_model, 138 | lpips_fn=lpips_fn, 139 | eps=args.attack_budget, 140 | step_size=args.step_size, 141 | iters=args.num_iters, 142 | clamp_min=-1, 143 | clamp_max=1, 144 | ) 145 | 146 | if args.clamp_min == -1: 147 | X_adv = (X_adv / 2 + 0.5).clamp(0, 1) 148 | protected_image = to_pil(X_adv[0]).convert("RGB") 149 | protected_image.save(os.path.join(args.output_dir, image_file)) 150 | 151 | if __name__ == "__main__": 152 | parser = get_args_parser() 153 | args = parser.parse_args() 154 | 155 | args.output_dir = os.path.join(args.output_dir, f"budget_{args.attack_budget}") 156 | if not os.path.exists(args.output_dir): 157 | os.makedirs(args.output_dir) 158 | 159 | main(args) -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # FaceLock 2 | 3 | [](https://arxiv.org/abs/2411.16832) [](https://opensource.org/licenses/MIT) 4 | 5 |
8 |
9 | 10 | Figure 1: An illustration of adversarial perturbation generation for safeguarding personal images. 11 | |
12 |