update

Author: Yi Zhang

config.py

@@ -3,6 +3,7 @@
 num_inference_steps = 50
 num_batch = 5  
 batch_size = 12
+sample_num = 800
 model_id = "runwayml/stable-diffusion-v1-5"
 
 e_threshold = 0.08

version_defence.py

@@ -15,7 +15,7 @@
 import pandas as pd
 torch.cuda.empty_cache()
 from config import (
-    e_threshold, origin_prompt_path, sigma,
+    sample_num, origin_prompt_path, sigma,
     num_inference_steps, num_batch, batch_size,
     model_id
 )
@@ -184,7 +184,7 @@ def defence_autocorrect(influent_sentence):
                 logger.info(f"disturb_num: {sample_data[0]}")
                 n = 1
                 epsilon = 1000
-                for count in range(400): 
+                for count in range(sample_num): 
                     selected = random.choices(strings, k=1)[0]
                     disturb_prompt = defence_spellchecker(selected)
                     if disturb_prompt == ori_prompt:

version_sample.py

@@ -16,7 +16,7 @@
 import pandas as pd
 torch.cuda.empty_cache()
 from config import (
-    e_threshold, origin_prompt_path, sigma,
+    sample_num, origin_prompt_path, sigma,
     num_inference_steps, num_batch, batch_size,
     model_id
 )
@@ -165,7 +165,7 @@ def get_origin_prompt(origin_prompt_path):
                 logger.info(f"disturb_num: {sample_data[0]}")
                 n = 1
                 epsilon = 1000
-                for count in range(400):
+                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)