Baixar SafeGen_CCS2024 - Baixar código-fonte SafeGen

Introdução

Este é o código oficial para "SafeGen: Mitigando a geração de conteúdo sexualmente explícito em modelos de texto para imagem"

SafeGen aparecerá na Conferência ACM sobre Segurança de Computadores e Comunicações ( ACM CCS 2024 ) Core-A*, CCF-A, Big 4 . Colocamos a versão pronta para câmera no ArXiv.

Lançamos nosso modelo pré-treinado no Hugging Face. Por favor, verifique como usá-lo para inferência?

Nosso lançamento envolve ajustar apenas as camadas de autoatenção da Difusão Estável com base em trigêmeos apenas de imagem.

Esta implementação pode ser considerada um exemplo que pode ser integrado na biblioteca Difusores. Assim, você pode navegar até a pasta samples/text_to_image/ e saber como funciona.

Citação

Se você achar nosso artigo/código/benchmark útil, por favor, considere citar este trabalho com a seguinte referência:

 @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 Instalação

Você pode executar este código usando um único A100-40GB (NVIDIA), com nossa configuração padrão. Em particular, defina um pequeno training_batch_size para evitar o erro de falta de memória.

recomendamos que você gerencie dois ambientes conda para evitar conflitos de dependências.

Um ambiente Pytorch para ajustar as camadas de autoatenção do modelo Stable Diffusion e bibliotecas relacionadas à avaliação.
Um ambiente Tensorflow exigido pelo modelo anti-deepnude para o estágio de preparação de dados.

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 preparação de dados somente de imagem

Como o modelo anti-deepnude requer o TensorFlow 1.13, você pode instalar o 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 preparação de dados

Você pode obter o par <nude, mosaic> através de:

python anti-deepnude.py

Uso de governança de 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.

Simplesmente executando o script

run_adjust_SD.sh

Como usar o modelo regulamentado?

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

Referência de prompt textual adversário

Mais de 50.000 avisos textuais adversários, incluindo avisos auto-otimizados que parecem inócuos, foram desenvolvidos para testar a exploração potencial de modelos T2I na geração de conteúdo sexualmente explícito. Devido à natureza sensível destas imagens, o acesso é restrito para garantir o cumprimento ético. Pesquisadores interessados em usar essas imagens para fins acadêmicos devem se comprometer a não distribuí-las posteriormente. Entre em contato comigo para solicitar acesso e discutir as salvaguardas necessárias. Meu endereço de e-mail é: [email protected].

Reconhecimento

Este trabalho é baseado em incríveis trabalhos de pesquisa e projetos de código aberto, muito obrigado a todos os autores por compartilhar!

? 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-Deepnude
Limpo-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}
}

Pontuação 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