Official implementation of the paper accepted at NeurIPS 2025
Authors: Jaesin Ahn, Heechul Jung (Kyungpook National University)
Diffusion models show remarkable image generation performance following text prompts, but risk generating sexual contents. Existing approaches struggle to defend against adversarial attacks while maintaining benign image quality. We propose DES (Distorting Embedding Space), a text encoder-based defense mechanism that effectively tackles these issues through innovative embedding space control.
DES transforms unsafe embeddings extracted from text encoders toward carefully calculated safe embedding regions to prevent unsafe content generation, while reproducing the original safe embeddings. DES also neutralizes the "nudity" embedding by aligning it with neutral embedding to enhance robustness against adversarial attacks.
- State-of-the-Art Defense: ASR of 9.47% on FLUX.1 and 0.52% on SD v1.5
- Superior Performance: 76.5% and 63.9% ASR reduction compared to previous SOTA methods
- Quality Preservation: Maintains FID and CLIP scores comparable to original models
- Practical Efficiency: 90 seconds training time with zero inference overhead
- Stable Diffusion: v1.x, v2.x, XL, v3.x
- FLUX.1: FLUX.1-dev
neurips/
├── train_des.py # Main DES training script
├── save_codebook.py # Safe embedding codebook generation
├── generate.py # Image generation
├── fid.py # FID calculation
├── clipscore.py # CLIP Score calculation
│
├── train.sh # Training example script
├── generate.sh # Generation example script
├── evaluate.sh # Evaluation example script
│
├── datasets/ # Prompt datasets (safe/unsafe prompts, attack datasets)
│
├── tasks/ # Evaluation utilities
│ ├── img_batch_classify.py # NudeNet classification
│ ├── img_batch_classify_q16.py # Q16 classification
│ └── utils/ # Evaluation metrics and utilities
│
└── README.md # This file
# Clone the repository
git clone https://github.com/amoeba04/des-neurips2025.git
cd des-neurips2025/neurips
# Create conda environment
conda create -n des python=3.8
conda activate des
# Install PyTorch 2.2.1 with CUDA 11.8 support
pip install torch==2.2.1 torchvision==0.17.1 --index-url https://download.pytorch.org/whl/cu118
# Install core dependencies
pip install -r requirements.txt
# Install evaluation dependencies
pip install git+https://github.com/openai/CLIP.git
pip install git+https://github.com/boomb0om/text2image-benchmarkpython save_codebook.py \
--model_path stable-diffusion-v1-5/stable-diffusion-v1-5 \
--device cuda:0 \
--csv_path datasets/safe_prompts_copro_sexual.csv \
--save_dir codebook_copro_sexualpython train_des.py \
--model_path stable-diffusion-v1-5/stable-diffusion-v1-5 \
--device cuda:0 \
--codebook_dir codebook_copro_sexual \
--unsafe_csv_path datasets/unsafe_prompts_copro_sexual.csv \
--safe_csv_path datasets/safe_prompts_copro_sexual.csv \
--output_dir checkpoints/des_copro_sexual \
--num_epochs 2 \
--learning_rate 1e-5 \
--batch_size 128 \
--lambda_safe 0.3 \
--save_every 1 \
--ablation 1 2 3 \
--concept_prompt "nudity" \
--concept_guidance_scale 200.0 \
--safe_embedding_path checkpoints/des_copro_sexual/safe_embeddings.pthTraining completes in ~90 seconds for CLIP-L/14.
# Unsafe images
python generate.py \
--model_path stable-diffusion-v1-5/stable-diffusion-v1-5 \
--device cuda:0 \
--prompts_csv "datasets/mma_prompts.csv" \
--output_path "t2i_mma" \
--start_idx 0 \
--training_method des \
--text_encoder_path "checkpoints/des_copro_sexual/1e-05_0.3/checkpoint-2.pt"
# Safe images
python generate.py \
--model_path stable-diffusion-v1-5/stable-diffusion-v1-5 \
--device cuda:0 \
--prompts_csv "datasets/coco_prompts.csv" \
--output_path "t2i_coco" \
--start_idx 0 \
--end_idx 10000 \
--training_method des \
--text_encoder_path "checkpoints/des_copro_sexual/1e-05_0.3/checkpoint-2.pt"
# NudeNet-based evaluation
python tasks/img_batch_classify.py \
--job nudity \
--cls_class nudity \
--folder_dir results/des/stable-diffusion-v1-5/t2i_sneaky/ \
--devices 0
# Q16-based evaluation
python tasks/img_batch_classify_q16.py \
--job nudity \
--cls_class nudity \
--folder_dir results/des/stable-diffusion-v1-5/t2i_sneaky/ \
--devices 0# FID (Fréchet Inception Distance)
python fid.py \
--gen_imgs_path results/des/stable-diffusion-v1-5/t2i_coco/ \
--coco_imgs_path datasets/coco_10k/ \
--device cuda:0
# CLIP Score
python clipscore.py \
--image_folder results/des/stable-diffusion-v1-5/t2i_coco/ \
--csv_file datasets/coco_prompts.csv \
--device cuda:0Core Arguments:
--model_path: Hugging Face model ID or local path--codebook_dir: Safe embedding codebook directory--unsafe_csv_path: CSV file with unsafe prompts (for Unsafe loss)--safe_csv_path: CSV file with safe prompts (for Safe loss)--output_dir: Checkpoint save directory
Hyperparameters (Recommended):
--lambda_safe: Balance parameter (default: 0.3)--concept_guidance_scale: Nudity subtraction scale (default: 200.0)
Training Configuration:
--num_epochs: Training epochs (default: 2, typically sufficient)--learning_rate: Learning rate (default: 1e-5)--batch_size: 128 for SD v1.x, 16 for FLUX.1
SD v3.5 (Multiple Text Encoders):
SD v3.5 has 3 text encoders:
- Encoder 1: CLIP-L/14 (768 dim)
- Encoder 2: CLIP-G/14 (1280 dim)
- Encoder 3: T5-XXL (4096 dim)
Step 1: Generate codebook for specific encoder
python save_codebook.py \
--model_path stabilityai/stable-diffusion-3.5-medium \
--device cuda:0 \
--csv_path datasets/safe_prompts_copro_sexual.csv \
--save_dir codebook_sdv3_copro_sexual \
--encoder_idx 3 # 1, 2, or 3Step 2: Train DES on specific encoder
For CLIP encoders (1 or 2) - Single GPU is sufficient:
python train_des.py \
--model_path stabilityai/stable-diffusion-3.5-medium \
--codebook_dir codebook_sdv3_copro_sexual \
--unsafe_csv_path datasets/unsafe_prompts_copro_sexual.csv \
--safe_csv_path datasets/safe_prompts_copro_sexual.csv \
--output_dir checkpoints/des_sdv3_encoder1 \
--text_encoder_idx 1 \
--batch_size 128 \
--learning_rate 1e-5 \
--lambda_safe 0.3 \
...For T5 encoder (3) - Multi-GPU recommended for 40GB GPUs:
# T5-XXL requires multi-GPU setup for 40GB GPUs
CUDA_VISIBLE_DEVICES=0,1,2,3 python train_des.py \
--model_path stabilityai/stable-diffusion-3.5-medium \
--codebook_dir codebook_sdv3_copro_sexual \
--unsafe_csv_path datasets/unsafe_prompts_copro_sexual.csv \
--safe_csv_path datasets/safe_prompts_copro_sexual.csv \
--output_dir checkpoints/des_sdv3_encoder3 \
--text_encoder_idx 3 \
--multi_gpu \
--gradient_checkpointing \
--batch_size 4 \
--learning_rate 1e-5 \
--lambda_safe 0.3 \
--concept_prompt "nudity" \
--concept_guidance_scale 200.0 \
--safe_embedding_path checkpoints/des_sdv3_encoder3/safe_embeddings.pthNote: T5-XXL requires significant GPU memory. With 40GB GPUs, use
--multi_gpuflag to distribute the model across multiple GPUs automatically.
- CLIP-L/14 (Stable Diffusion v1.x, Stable Diffusion v3.x, FLUX.1): link
- CLIP-G/14 (Stable Diffusion v3.x): link
- T5-XXL (Stable Diffusion v3.x, FLUX.1): link
- CLIP-L/14 (Stable Diffusion v1.x): link
- CLIP-L/14 (Stable Diffusion v1.x): link
This project is licensed under the CC BY-NC 4.0 (Creative Commons Attribution-NonCommercial 4.0 International License).
For detailed terms and conditions, see the LICENSE file.
If you use this code or find our work helpful, please cite:
@inproceedings{ahn2025des,
title={Mitigating Sexual Content Generation via Embedding Distortion in Text-conditioned Diffusion Models},
author={Ahn, Jaesin and Jung, Heechul},
booktitle={The Thirty-ninth Annual Conference on Neural Information Processing Systems},
year={2025}
}We thank the authors of the following repository for their code, which we partially referenced in this project: