update

Author: Yi Zhang

config.py

@@ -5,6 +5,7 @@
 batch_size = 12
 sample_num = 800
 model_id = "runwayml/stable-diffusion-v1-5"
+clip_version = "openai/clip-vit-large-patch14-336"
 
 e_threshold = 0.08
 sigma = 0.3

docs/perturb_example.PNG

@@ -1,5 +1,5 @@
 import os
-os.environ["CUDA_VISIBLE_DEVICES"] = "1"
+os.environ["CUDA_VISIBLE_DEVICES"] = "0"
 import torch
 device = "cuda" if torch.cuda.is_available() else "cpu"
 import logging
@@ -16,12 +16,12 @@
 import pandas as pd
 torch.cuda.empty_cache()
 from config import (
-    sample_num, origin_prompt_path, sigma,
+    e_threshold, origin_prompt_path, sigma,clip_version,
     num_inference_steps, num_batch, batch_size,
     model_id
 )
 import random
-# random.seed(42) 
+
 
 def setup_logger(file_name):
     logger = logging.getLogger(file_name)
@@ -35,12 +35,14 @@ def setup_logger(file_name):
     return logger  
 
 pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
+# pipe = StableDiffusionPipeline.from_pretrained('CompVis/stable-diffusion-v1-4', torch_dtype=torch.float16)
 pipe = pipe.to(device)
 
 transform = transforms.Compose([transforms.ToTensor()])
-metric = CLIPScore().to(device)
+metric = CLIPScore(model_name_or_path=clip_version).to(device)
 
 def calculate_text_image_distance(text, image):
+    # print(transform(image).shape)  # torch.Size([1, 512, 512])
     img = transform(image)*255
     score = metric(img.to(device), text)
     return score.detach().cpu().numpy().item()
@@ -62,6 +64,17 @@ def calculate_t_zValue(data1, data2):
     return p_val, z_value
 
 
+def generate_func(pipe, prompt, seed):
+    while True:
+        generator = torch.Generator(device="cuda").manual_seed(seed)
+        output = pipe([prompt] * batch_size, num_inference_steps = num_inference_steps, generator = generator)
+        print(output.nsfw_content_detected, seed)
+        if any(output.nsfw_content_detected):
+            print(f"Potential NSFW content was detected with seed {seed}. Trying a new seed.")
+            seed = seed + 100 
+        else:
+            return output
+
 def cal_loss(ori_loss, disturb_prompt, ori_prompt):
     print("dis_prompt", disturb_prompt)
     print("ori_prompt", ori_prompt)
@@ -71,18 +84,20 @@ def cal_loss(ori_loss, disturb_prompt, ori_prompt):
     dis_interim_loss = []
     alpha = [0.0148, 0.0262, 0.0354, 0.0432, 0.05]
     futility_boundary = [-0.145, 0.511, 1.027, 1.497, float('inf')]
-    efficacy_boundary = [2.176, 2.144, 2.113, 2.090, 2.071] 
-    for i in range(num_batch): 
+    efficacy_boundary = [2.176, 2.144, 2.113, 2.090, 2.071]  # critical value when Z-score > critical value, the reject the H0
+    for i in range(num_batch):  # 5
         ori_interim_loss = ori_loss[0: batch_size*(i+1)]
-        generator = torch.Generator(device).manual_seed(1023+i)
-        images = pipe([disturb_prompt] * batch_size, num_inference_steps = num_inference_steps, generator = generator)
-        for j in range(batch_size):  
+        # generator = torch.Generator(device).manual_seed(1023+i)
+        # images = pipe([disturb_prompt] * batch_size, num_inference_steps = num_inference_steps, generator = generator)
+        images = generate_func(pipe, disturb_prompt, seed=1023+i)
+        for j in range(batch_size):  # 5
             dis_interim_loss.append(calculate_text_image_distance(ori_prompt, images.images[j]))
+        # _, p_val = mannwhitneyu(dis_interim_loss, ori_interim_loss)
         logger.info(f"dis_interim_loss: {len(dis_interim_loss)}; {dis_interim_loss}")
         logger.info(f"ori_interim_loss: {len(ori_interim_loss)}; {ori_interim_loss}")
         _, p_1 = shapiro(ori_interim_loss[0: 12*(i+1)])
         _, p_2 = shapiro(dis_interim_loss[0: 12*(i+1)]) 
-        if p_1 > 0.05 and p_2 > 0.05: 
+        if p_1 > 0.05 and p_2 > 0.05:  # normal distr
             p_val, z_val = calculate_t_zValue(ori_interim_loss[0: 12*(i+1)], dis_interim_loss[0: 12*(i+1)])
         else:
             p_val, z_val = calculate_u_zValue(ori_interim_loss[0: 12*(i+1)], dis_interim_loss[0: 12*(i+1)])
@@ -107,6 +122,7 @@ def get_AE(sample_data):
     import random
     random.seed(42) 
     strings = [line.split(':')[0].strip() for line in sample_data[1:]]
+    # sampled_strings = random.sample(strings, len(strings))
     sampled_strings = random.choices(strings, k=1)
     return sampled_strings
 
@@ -130,80 +146,73 @@ def get_origin_prompt(origin_prompt_path):
     return origin_prompt
 
 
+
 if __name__ == "__main__":
     start_time = time.time()
     robust_left, robust_right = 0, 0
-    origin_prompts = get_origin_prompt(origin_prompt_path) 
-    for index, ori_prompt in origin_prompts.items():
+    test_path = "dataset/"
+    for index in range(1, 100):
+        AEdata_path = f"{test_path}/{index}.txt"
+        with open(AEdata_path, 'r', encoding = 'utf-8') as file:
+            content = file.readlines()
+        ori_prompt = content[0].strip()
         efficient_m, efficient_n = 0, 0
-        AEdata_path = f"./generate_AE/coco/char_AE/result_{index}.csv"
-        logger = setup_logger(f"adaptive_log/log_char_{index}.log")
+        logger = setup_logger(f"adaptive_log/origin_result/10_rate/log_char_{index}.log")
         logger.info(f"sigma: {sigma}")
         logger.info(f"num_inference_steps: {num_inference_steps}")
         logger.info(f"num_batch: {num_batch}")
         logger.info(f"batch_size: {batch_size}")
         logger.info(AEdata_path)
         logger.info(f"ori_prompt: {ori_prompt}")
-        df = pd.read_csv(AEdata_path)
         ori_loss = []
         for i in range(num_batch):
-            generator = torch.Generator(device).manual_seed(1023+i)
-            images = pipe([ori_prompt] * batch_size, num_inference_steps = num_inference_steps, generator = generator)
+            # generator = torch.Generator(device).manual_seed(1023+i)
+            # images = pipe([ori_prompt] * batch_size, num_inference_steps = num_inference_steps, generator = generator)
+            images = generate_func(pipe, ori_prompt, seed=1023+i)
             for j in range(batch_size):
                 ori_loss.append(calculate_text_image_distance(ori_prompt, images.images[j]))
         logger.info(f"ori_loss: {len(ori_loss)} {ori_loss}")
         logger.info(f"*" * 120)
-        for id in range(2, 3): 
-            efficient_n = 0
-            Non_AE, n = 0, 0
-            L_distance, AdvSt2i = [], []
-            robust_re, epsilon_re = [], []
-            sample_data = list(df[f"Column {id}"].dropna())
-            strings = [line.split(':')[0].strip() for line in sample_data[1:]]
-            logger.info(f"disturb rate: {id}")
-            logger.info(f"disturb_num: {sample_data[0]}")
-            n = 1
-            epsilon = 1000
-            for count in range(sample_num):
-                disturb_prompt = random.choices(strings, k=1)[0]
-                L_distance.append(Levenshtein.distance(ori_prompt, disturb_prompt))
-                whether_robust = cal_loss(ori_loss, disturb_prompt, ori_prompt)
-                Non_AE += 1 if whether_robust else 0
-                robust_left, robust_right, epsilon = calculate_R(Non_AE, n)
-                robust_re.append((robust_left, robust_right))
-                epsilon_re.append(epsilon)
-                logger.info(f"stop_early: {efficient_n}")
-                logger.info(f"futility: {efficient_m}")
-                logger.info(f"Non_AE: {Non_AE}")
-                logger.info(f"n: {n}")
-                logger.info(f"robust reach: {robust_left} , {robust_right}")
-                logger.info(f"epsilon reach: {epsilon}")
-                print("*" * 120)
-                logger.info(f"*" * 120)
-                n += 1
+        efficient_n = 0
+        Non_AE, n = 0, 0
+        L_distance, AdvSt2i = [], []
+        robust_re, epsilon_re = [], []
+        n = 1
+        epsilon = 1000
+        for item in content[1:]:
+            disturb_prompt = item.strip()
+            whether_robust = cal_loss(ori_loss, disturb_prompt, ori_prompt)
+            Non_AE += 1 if whether_robust else 0
+            robust_left, robust_right, epsilon = calculate_R(Non_AE, n)
+            robust_re.append((robust_left, robust_right))
+            epsilon_re.append(epsilon)
+            logger.info(f"stop_early: {efficient_n}")
+            logger.info(f"futility: {efficient_m}")
+            logger.info(f"Non_AE: {Non_AE}")
+            logger.info(f"n: {n}")
+            logger.info(f"robust reach: {robust_left} , {robust_right}")
+            logger.info(f"epsilon reach: {epsilon}")
             print("*" * 120)
             logger.info(f"*" * 120)
-            logger.info(f"robust = {robust_re}")
-            logger.info(f"epsilon = {epsilon_re}")
-            logger.info(f"stop_early = {efficient_n}")
-            logger.info(f"futility = {efficient_m}")
-            logger.info(f"Non_AE = {Non_AE}")
-            logger.info(f"n = {n}")
-            logger.info(f"AdvSt2i = {round(np.mean(AdvSt2i), 2)}")
-            logger.info(f"OriSt2i = {round(np.mean(ori_loss), 2)}")
-            logger.info(f"Levenshtein = {round(np.mean(L_distance), 2)}")
-            logger.info(f"robust = {robust_left} , {robust_right}")
-            logger.info(f"epsilon = {epsilon}")
-
-            end_time = time.time()
-            elapsed_time = end_time - start_time
-            hours, remainder = divmod(elapsed_time, 3600)
-            minutes, seconds = divmod(remainder, 60)
-            print(f"time cost: {int(hours)} hours, {int(minutes)} minutes, {int(seconds)} seconds")
-            logger.info(f"time cost: {int(hours)} hours, {int(minutes)} minutes, {int(seconds)} seconds")
-            logger.info(f"&" * 150)
-
-
-
-
-
+            n += 1
+        print("*" * 120)
+        logger.info(f"*" * 120)
+        logger.info(f"robust = {robust_re}")
+        logger.info(f"epsilon = {epsilon_re}")
+        logger.info(f"stop_early = {efficient_n}")
+        logger.info(f"futility = {efficient_m}")
+        logger.info(f"Non_AE = {Non_AE}")
+        logger.info(f"n = {n}")
+        logger.info(f"AdvSt2i = {round(np.mean(AdvSt2i), 2)}")
+        logger.info(f"OriSt2i = {round(np.mean(ori_loss), 2)}")
+        logger.info(f"Levenshtein = {round(np.mean(L_distance), 2)}")
+        logger.info(f"robust = {robust_left} , {robust_right}")
+        logger.info(f"epsilon = {epsilon}")
+
+        end_time = time.time()
+        elapsed_time = end_time - start_time
+        hours, remainder = divmod(elapsed_time, 3600)
+        minutes, seconds = divmod(remainder, 60)
+        print(f"time cost: {int(hours)} hours, {int(minutes)} minutes, {int(seconds)} seconds")
+        logger.info(f"time cost: {int(hours)} hours, {int(minutes)} minutes, {int(seconds)} seconds")
+        logger.info(f"&" * 150)