Téléchargement SafeGen_CCS2024 - Téléchargement du code source SafeGen

Introduction

Ceci est le code officiel de « SafeGen : Atténuation de la génération de contenu sexuellement explicite dans les modèles texte-image »

SafeGen apparaîtra à la conférence ACM sur la sécurité informatique et des communications ( ACM CCS 2024 ) Core-A*, CCF-A, Big 4 . Nous avons mis en ligne la version prête à photographier sur ArXiv.

Nous avons publié notre modèle pré-entraîné sur Hugging Face. Veuillez vérifier comment l'utiliser pour l'inférence ?.

Notre version implique l'ajustement des couches d'auto-attention de Stable Diffusion uniquement sur la base de triplets d'images uniquement.

Cette implémentation peut être considérée comme un exemple pouvant être intégré à la bibliothèque Diffusers. Ainsi, vous pouvez accéder au dossier examples/text_to_image/ et découvrir son fonctionnement.

Citation

Si vous trouvez notre article/code/référence utile, veuillez envisager de citer ce travail avec la référence suivante :

 @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}
}

Environnements et installation

Vous pouvez exécuter ce code en utilisant un seul A100-40GB (NVIDIA), avec notre configuration par défaut. En particulier, définissez un petit training_batch_size pour éviter l'erreur de manque de mémoire.

nous vous recommandons de gérer deux environnements conda pour éviter les conflits de dépendances.

Un environnement Pytorch pour ajuster les couches d'auto-attention du modèle Stable Diffusion et les bibliothèques liées à l'évaluation.
Un environnement Tensorflow requis par le modèle anti-deepnude pour l'étape de préparation des données.

Exigence de PyTorch + Diffuseurs

 # 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

Exigence de préparation des données d'image uniquement

Comme le modèle anti-deepnude nécessite TensorFlow 1.13, vous pouvez installer 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 pour la préparation des données

Vous pouvez obtenir la paire <nu, mosaïque> via :

python anti-deepnude.py

Utilisation du modèle de gouvernance

 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.

Exécuter simplement le script

run_adjust_SD.sh

Comment utiliser le modèle régulé ?

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 " )

Évaluation des invites textuelles contradictoires

Plus de 50 000 invites textuelles contradictoires, y compris des invites auto-optimisées qui semblent inoffensives, ont été développées pour tester l'exploitation potentielle des modèles T2I dans la génération de contenu sexuellement explicite. En raison du caractère sensible de ces images, leur accès est restreint pour garantir le respect de l’éthique. Les chercheurs souhaitant utiliser ces images à des fins scientifiques doivent s’engager à ne plus les diffuser. Veuillez me contacter pour demander l’accès et discuter des garanties nécessaires. Mon adresse e-mail est : [email protected].

Reconnaissance

Ce travail est basé sur des travaux de recherche incroyables et des projets open source, merci beaucoup à tous les auteurs pour le partage !

? Diffuseurs

@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-nu profond
Propre-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}
}

Score 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}
}

Développer