Descarga SafeGen_CCS2024 - Descarga del código fuente SafeGen

Introducción

Este es el código oficial de "SafeGen: Mitigar la generación de contenido sexualmente explícito en modelos de texto a imagen"

SafeGen aparecerá en la Conferencia ACM sobre seguridad informática y de las comunicaciones ( ACM CCS 2024 ) Core-A*, CCF-A, Big 4 . Hemos puesto la versión lista para cámara en ArXiv.

Hemos lanzado nuestro modelo previamente entrenado en Hugging Face. Consulte cómo usarlo para inferencia.

Nuestro lanzamiento implica ajustar las capas de autoatención de Stable Diffusion únicamente en función de tripletes de solo imágenes.

Esta implementación puede considerarse como un ejemplo que se puede integrar en la biblioteca de Difusores. Por lo tanto, puede navegar a la carpeta ejemplos/text_to_image/ y conocer cómo funciona.

Citación

Si encuentra útil nuestro documento/código/punto de referencia, considere citar este trabajo con la siguiente referencia:

 @inproceedings{li2024safegen,
  author       = {Li, Xinfeng and Yang, Yuchen and Deng, Jiangyi and Yan, Chen and Chen, Yanjiao and Ji, Xiaoyu and Xu, Wenyuan},
  title        = {{SafeGen: Mitigating Sexually Explicit Content Generation in Text-to-Image Models}},
  booktitle    = {Proceedings of the 2024 {ACM} {SIGSAC} Conference on Computer and Communications Security (CCS)},
  year         = {2024},
}

 @article{li2024safegen,
  title={{SafeGen: Mitigating Unsafe Content Generation in Text-to-Image Models}},
  author={Li, Xinfeng and Yang, Yuchen and Deng, Jiangyi and Yan, Chen and Chen, Yanjiao and Ji, Xiaoyu and Xu, Wenyuan},
  journal={arXiv preprint arXiv:2404.06666},
  year={2024}
}

Ambientes e Instalación

Puede ejecutar este código utilizando un único A100-40GB (NVIDIA), con nuestra configuración predeterminada. En particular, establezca un training_batch_size pequeño para evitar el error de falta de memoria.

Le recomendamos administrar dos entornos conda para evitar conflictos de dependencias.

Un entorno Pytorch para ajustar las capas de autoatención del modelo Stable Diffusion y bibliotecas relacionadas con la evaluación.
Un entorno Tensorflow requerido por el modelo anti-deepnude para la etapa de preparación de datos.

Requisito de PyTorch + Difusores

 # You can install the main dependencies by conda/pip
conda create -n text-agnostic-t2i python=3.8.5
conda install pytorch==1.12.1 torchvision==0.13.1 torchaudio==0.12.1 cudatoolkit=11.3 -c pytorch
# Using the official Diffusers package
pip install --upgrade diffuers[torch]
# Or you may use the community maintained version
conda install -c conda-forge diffusers
...

 # Or you can create the env via environment.yaml
conda env create -f environment_pytorch.yaml

Requisito de preparación de datos de solo imágenes

Como el modelo anti-deepnude requiere TensorFlow 1.13, puedes instalar Python<=3.7.

 # You can install the dependencies individually
conda create -n anti-deepndue python=3.7
pip install tensorflow-gpu==1.13.1 keras==2.2.4
...

 # Or you can create the env via environment.yaml
conda env create -f environment_tf.yaml

Anti-DeepNude para la preparación de datos

Puedes obtener el par <nude, mosaic> a través de:

python anti-deepnude.py

Uso de la gobernanza del modelo

 export MODEL_NAME= " CompVis/stable-diffusion-v1-4 " # export the model you want to protect
export TRAIN_DIR= " <path to the manifest of your training files>/meta.csv " # export your image triplets folder
accumulation_steps=5
learning_rate=1e-05
steps=1000
train_batch_size=1
warmup_steps=200
weight_tune=0.1
weight_preserve=0.9
defend_type= " nudity_advtrain "
current_time= $( date + ' %m%d_%H%M ' )
output_dir= " mos_model_sal_ " ${defend_type} - " step_ " ${steps} - " lr_ " ${learning_rate} - " warm_ " ${warmup_steps} - ${weight_tune} + ${weight_preserve} " gacc_ " ${accumulation_steps} - ${current_time}

export CUDA_VISIBLE_DEVICES=0
accelerate launch --mixed_precision= " fp16 " " --num_processes=1 " train_text_to_image_adjust_sal.py 
  --pretrained_model_name_or_path= $MODEL_NAME 
  --train_data_dir= $TRAIN_DIR 
  --use_ema 
  --resolution=512 --center_crop
  --train_batch_size ${train_batch_size} 
  --gradient_accumulation_steps ${accumulation_steps} 
  --gradient_checkpointing 
  --max_train_steps ${steps} 
  --learning_rate ${learning_rate} 
  --image_column= " nude_image_path " 
  --max_grad_norm=1 
  --lr_scheduler= " constant " 
  --lr_warmup_steps ${warmup_steps} 
  --output_dir ${output_dir} 
  --weight_tune ${weight_tune} 
  --weight_preserve ${weight_preserve} 
  --checkpoints_total_limit 0 
  --validation_prompts " woman, nudity " " a running dog on grass " " a woman standing aside the colorful table " " a nude sexy man, blue eyes "
  
  # In addition, you can set the 'validation_prompts' to visually monitor the training efficacy, such as "a running dog on grass", "a woman standing aside the colorful table", "a nude sexy man, blue eyes", etc.

Simplemente ejecutando el script

run_adjust_SD.sh

¿Cómo utilizar el modelo regulado?

from diffusers import StableDiffusionPipeline
import torch

model_path = ${output_dir} # the save path of your model
pipeline = StableDiffusionPipeline.from_pretrained(model_path, torch_dtype=torch.float16)
pipeline.to( " cuda " )

prompt = " a photo of an astronaut riding a horse on mars "
image = pipeline(prompt).images[0]
image.save( " example.png " )

Punto de referencia de indicaciones textuales adversas

Se han desarrollado más de 50.000 indicaciones textuales de confrontación, incluidas indicaciones autooptimizadas que parecen inocuas, para probar la posible explotación de los modelos T2I en la generación de contenido sexual explícito. Debido a la naturaleza sensible de estas imágenes, el acceso está restringido para garantizar el cumplimiento ético. Los investigadores interesados en utilizar estas imágenes con fines académicos deben comprometerse a no distribuirlas más. Comuníquese conmigo para solicitar acceso y discutir las salvaguardas necesarias. Mi dirección de correo electrónico es: [email protected].

Reconocimiento

Este trabajo se basa en increíbles trabajos de investigación y proyectos de código abierto. ¡Muchas gracias a todos los autores por compartir!

? Difusores

@misc{von-platen-etal-2022-diffusers,
    author = {Patrick von Platen and Suraj Patil and Anton Lozhkov and Pedro Cuenca and Nathan Lambert and Kashif Rasul and Mishig Davaadorj and Thomas Wolf},
    title = {Diffusers: State-of-the-art diffusion models},
    year = {2022},
    publisher = {GitHub},
    journal = {GitHub repository},
    howpublished = { url {https://github.com/huggingface/diffusers}}
}

Anti-desnudo profundo
limpio-fid

@inproceedings{parmar2021cleanfid,
  title={On Aliased Resizing and Surprising Subtleties in GAN Evaluation},
  author={Parmar, Gaurav and Zhang, Richard and Zhu, Jun-Yan},
  booktitle={CVPR},
  year={2022}
}

puntuación LPIPS

@inproceedings{zhang2018perceptual,
  title={The Unreasonable Effectiveness of Deep Features as a Perceptual Metric},
  author={Zhang, Richard and Isola, Phillip and Efros, Alexei A and Shechtman, Eli and Wang, Oliver},
  booktitle={CVPR},
  year={2018}
}

Expandir