src/backend/app/algorithms/perturbation/caat.py

@@ -545,12 +545,14 @@ def _save_final_images(
     os.makedirs(save_folder, exist_ok=True)
 
     noised_imgs = final_perturbed_images.detach().float().cpu()
-    img_names = [str(instance_path[0]).split("/")[-1] for instance_path in train_dataset.instance_images_path]
+
+    from pathlib import Path
+    img_names = [Path(instance_path[0]).stem for instance_path in train_dataset.instance_images_path]
 
     for i in range(len(img_names)):
         img_pixel = noised_imgs[i]
         img_name = img_names[i]
-        save_path = os.path.join(save_folder, f"final_noise_{img_name}")
+        save_path = os.path.join(save_folder, f"perturbed_{img_name}.png")
         Image.fromarray(
             (img_pixel * 127.5 + 128).clamp(0, 255).to(torch.uint8).permute(1, 2, 0).numpy()
         ).save(save_path)

src/backend/app/algorithms/perturbation/pid.py

@@ -11,6 +11,7 @@ from torch.utils.data import Dataset
 from torchvision import transforms
 from tqdm.auto import tqdm
 from diffusers import AutoencoderKL
+from pathlib import Path
 
 
 def parse_args(input_args=None):
@@ -225,7 +226,8 @@ def main(args):
                 for i in range(0, len(dataset.instance_images_path)):
                     img = dataset[i]['pixel_values']
                     img = to_image(img + attackmodel.delta[i])
-                    img.save(os.path.join(args.output_dir, f"{i}.png"))
+                    original_name = Path(dataset.instance_images_path[i]).stem
+                    img.save(os.path.join(args.output_dir, f"perturbed_{original_name}.png"))
 
             # 分别计算原始图像和中毒（添加扰动）图像在潜空间的分布
             clean_embedding = attackmodel(vae, batch['pixel_values'], batch['index'], False)

src/backend/app/scripts/finetune_db.sh

@@ -37,22 +37,21 @@ CUDA_VISIBLE_DEVICES=0 accelerate launch ../finetune_infras/train_db_gen_trace.p
   --output_dir=$DREAMBOOTH_OUTPUT_DIR \
   --with_prior_preservation \
   --train_text_encoder \
-  --prior_loss_weight=0.4 \
+  --prior_loss_weight=0.15 \
   --instance_prompt="a selfie photo of <sks> person" \
   --class_prompt="a selfie photo of person" \
   --resolution=512 \
   --train_batch_size=1 \
-  --gradient_accumulation_steps=1 \
-  --learning_rate=5e-7 \
-  --lr_scheduler="constant_with_warmup" \
-  --lr_warmup_steps=50 \
-  --num_class_images=100 \
-  --max_train_steps=800 \
+  --gradient_accumulation_steps=4 \
+  --learning_rate=1e-6 \
+  --lr_scheduler="cosine_with_restarts" \
+  --num_class_images=80 \
+  --max_train_steps=1200 \
   --checkpointing_steps=400 \
   --mixed_precision=bf16 \
   --prior_generation_precision=bf16 \
   --sample_batch_size=5 \
-  --validation_prompt="a selfie photo of <sks> person, head-and-shoulders, face looking at the camera, Eiffel Tower clearly visible behind, outdoor daytime, realistic" \
+  --validation_prompt="a selfie photo of <sks> person" \
   --num_validation_images=5 \
   --validation_num_inference_steps=120 \
   --validation_guidance_scale=7.0 \

src/backend/app/scripts/finetune_lora.sh

@@ -53,8 +53,8 @@ CUDA_VISIBLE_DEVICES=0 accelerate launch ../finetune_infras/train_lora_gen_trace
   --output_dir=$DREAMBOOTH_OUTPUT_DIR \
   --validation_image_output_dir=$OUTPUT_INFER_DIR \
   --with_prior_preservation --prior_loss_weight=1.0 \
-  --instance_prompt="a photo of sks person" \
-  --class_prompt="a photo of person" \
+  --instance_prompt="a selfie photo of <sks> person" \
+  --class_prompt="a selfie photo of person" \
   --resolution=512 \
   --train_batch_size=1 \
   --gradient_accumulation_steps=1 \
@@ -67,7 +67,7 @@ CUDA_VISIBLE_DEVICES=0 accelerate launch ../finetune_infras/train_lora_gen_trace
   --seed=0 \
   --mixed_precision=fp16 \
   --rank=4 \
-  --validation_prompt="a photo of sks person" \
+  --validation_prompt="a selfie photo of <sks> person" \
   --num_validation_images 10 \
   --positions_save_path="$POSITION_DIR" \
   --coords_log_interval 10 \

src/backend/app/scripts/finetune_ti.sh

@@ -47,9 +47,9 @@ CUDA_VISIBLE_DEVICES=0 accelerate launch ../finetune_infras/train_ti_gen_trace.p
   --instance_data_dir=$INSTANCE_DIR \
   --output_dir=$TI_OUTPUT_DIR \
   --validation_image_output_dir=$OUTPUT_INFER_DIR \
-  --placeholder_token="<sks-concept>" \
+  --placeholder_token="<sks>" \
   --initializer_token="person" \
-  --instance_prompt="a photo of <sks-concept> person" \
+  --instance_prompt="a selfie photo of <sks> person" \
   --resolution=512 \
   --train_batch_size=1 \
   --gradient_accumulation_steps=1 \
@@ -60,7 +60,7 @@ CUDA_VISIBLE_DEVICES=0 accelerate launch ../finetune_infras/train_ti_gen_trace.p
   --checkpointing_steps=500 \
   --seed=0 \
   --mixed_precision=fp16 \
-  --validation_prompt="a close-up photo of <sks-concept> person" \
+  --validation_prompt="a selfie photo of <sks> person" \
   --num_validation_images 4 \
   --validation_epochs 50 \
   --coords_save_path="$COORD_DIR" \

src/backend/config/algorithm_config.py

@@ -318,17 +318,16 @@ class AlgorithmConfig:
             'default_params': {
                 'pretrained_model_name_or_path': MODELS_DIR['model2'],
                 'with_prior_preservation': True,
-                'prior_loss_weight': 0.4,
+                'prior_loss_weight': 0.15,
                 'instance_prompt': 'a selfie photo of <sks> person',
                 'class_prompt': 'a selfie photo of person',
                 'resolution': 512,
                 'train_batch_size': 1,
-                'gradient_accumulation_steps': 1,
+                'gradient_accumulation_steps': 4,
                 'learning_rate': 1e-6,
-                'lr_scheduler': 'constant_with_warmup',
-                'lr_warmup_steps': 50,
-                'num_class_images': 100,
-                'max_train_steps': 800,
+                'lr_scheduler': 'cosine_with_restarts',
+                'num_class_images': 80,
+                'max_train_steps': 1200,
                 'checkpointing_steps': 400,
                 'center_crop': True,
                 'mixed_precision': 'bf16',

src/backend/init_db.py

@@ -105,8 +105,8 @@ def init_database():
                 'data_type_code': 'art', 
                 'instance_prompt': 'a painting in <sks> style', 
                 'class_prompt': 'a painting',
-                'validation_prompt_prefix_db_lora': '((a painting in <sks> style)), distinct <sks> feature, coherent aesthetic, balanced lighting and color, fine art texture',
-                'validation_prompt_prefix_ti': 'a painting in <sks> style',
+                'validation_prompt_prefix_db_lora': 'a painting in strong <sks> style',
+                'validation_prompt_prefix_ti': 'a painting in strong <sks> style',
                 'placeholder_token': '<sks>',
                 'initializer_token': 'style',
                 'description': '艺术品类型的数据集'