binzhwang/tnc-defenseCódigo de investigación para detectar y desintoxicar backdoors en modelos de difusión de texto a imagen, con pipelines para Stable Diffusion v1.4, v1.5, XL y 3 Medium.
TNC-Defense proporciona código de investigación para detectar y desintoxicar puertas traseras en modelos de difusión de texto a imagen. El repositorio incluye actualmente pipelines de detección para Stable Diffusion v1.4, Stable Diffusion v1.5, Stable Diffusion XL y Stable Diffusion 3 Medium, junto con pipelines de desintoxicación para Stable Diffusion v1.4 y Stable Diffusion XL.
Solo para uso en investigación. Este código está pensado para estudiar la robustez y la seguridad de los modelos de difusión de texto a imagen. Los usuarios son responsables de cumplir con las licencias y las políticas de uso de todos los modelos, conjuntos de datos e implementaciones upstream.
TNC-Defense/
├── TNC-Detect/
│ ├── detect_sd14.py # Detection on Stable Diffusion v1.4
│ ├── detect_sd15.py # Detection on Stable Diffusion v1.5
│ ├── detect_sdxl.py # Detection on Stable Diffusion XL
│ ├── detect_sd3.py # Detection on Stable Diffusion 3 Medium
│ ├── ana.py # Detection metrics and visualization
│ ├── run_detect_sd14.sh # Multi-GPU SD1.4 detection
│ └── ana.sh # Analysis entry point
└── TNC-Detox/
├── data/ # Paired clean/backdoored training data
├── common.py # Shared data and reproducibility utilities
├── detox_sd14.py # Detoxification for SD1.4
├── detox_sdxl.py # Detoxification for SDXL
├── evaluate_sd14.py # SD1.4 generation-based evaluation
├── compute_fid.py # FID evaluation
├── run_detox_sd14.sh # Multi-GPU SD1.4 detoxification
└── run_detox_sdxl.sh # SDXL detoxification
El código ha sido probado con pruebas de humo en el entorno Conda TCN con:
Las utilidades de análisis y evaluación requieren además paquetes como NumPy, pandas, SciPy, scikit-learn, Matplotlib, Pillow y tqdm.
Active el entorno antes de ejecutar un experimento:
cd TNC-Defense
conda create -n TCN python=3.10 -y
conda activate TCN
python -m pip install --upgrade pip
python -m pip install -r requirements.txt
Los puntos de entrada Bash utilizan el python del entorno activo. Se puede proporcionar un intérprete diferente mediante la variable de entorno PYTHON:
PYTHON=python bash TNC-Detect/run_detect_sd14.sh
Todas las rutas del repositorio se resuelven en relación con la ubicación de cada script, por lo que los comandos no dependen del directorio de trabajo actual. Las rutas de modelo predeterminadas siguen el diseño local del proyecto utilizado para los experimentos. Si los modelos están almacenados en otro lugar, sobrescriba los argumentos de línea de comandos correspondientes, como --base-model-path, --backdoor-model-path, --base-model o --tox-path.
Los conjuntos de datos de desintoxicación se incluyen en:
TNC-Detox/data/
├── sd14/
│ ├── pixel/
│ ├── pixel_mul_tokens/
│ ├── eviledit/
│ ├── personal_bkd/
│ └── villain_mul/
└── sdxl/
├── pixel/
└── pixel_mul_tokens/
Cada conjunto de datos contiene clean/, poison/ y metadata.json.
El ejecutor predeterminado evalúa todos los métodos de puerta trasera SD1.4 configurados. Asigna los métodos a los IDs de GPU definidos en run_detect_sd14.sh y los lanza en paralelo.
cd TNC-Defense/TNC-Detect
bash run_detect_sd14.sh
Para ejecutar un método directamente:
python detect_sd14.py \
--backdoor-method pixel \
--device cuda:0 \
--seed 0 \
--num-inference-steps 50
Los argumentos disponibles para cada familia de modelos se pueden consultar con:
python detect_sd14.py --help
python detect_sd15.py --help
python detect_sdxl.py --help
python detect_sd3.py --help
Después de generar los archivos CSV de MSE, configure los archivos de entrada en ana.sh y ejecute:
bash ana.sh
El detector admite tanto la regla de decisión ingenua como la de k dinámico. Las métricas y los gráficos los produce ana.py.
El ejecutor SD1.4 predeterminado lanza los cinco métodos configurados en los cinco IDs de GPU especificados en el script:
cd TNC-Defense/TNC-Detox
bash run_detox_sd14.sh
Para ejecutar un solo método:
python detox_sd14.py \
--method pixel \
--device cuda:0 \
--seed 42
Configure METHOD y GPU_ID en run_detox_sdxl.sh y, a continuación, ejecute:
bash run_detox_sdxl.sh
Los valores predeterminados específicos del método, incluidos los pasos de entrenamiento, las proporciones de muestreo de timestep, los pesos de pérdida, los tamaños de lote y los intervalos de checkpoint, se definen en la sección METHOD_PROFILES de cada script de entrenamiento. Los argumentos de línea de comandos pueden sobrescribir estos valores predeterminados.
Genere imágenes a partir de uno o más UNets SD1.4 desintoxicados:
python evaluate_sd14.py \
--detox-model checkpoints/sd14/<checkpoint> \
--prompt-file <prompt-file>
Calcule el FID entre los directorios de imágenes de referencia y desintoxicadas:
python compute_fid.py \
--reference-dir <reference-images> \
--candidate-dir <detoxified-images>
42 de forma predeterminada y habilita el comportamiento determinista de cuDNN.Los modelos con puertas traseras y las configuraciones de ataque utilizados por este proyecto provienen de los siguientes recursos públicos:
Cite los artículos originales correspondientes cuando utilice estos ataques o checkpoints: BadT2I (badt2izhai2023), EvilEdit (evileditwang2024), VillanDiffusion (villandiffusionchou2023) y PersonalBKD (personalizationhuang2024).
Si este proyecto le resulta útil en su investigación, cite:
@article{wang2026backdoor,
title={Backdoor sentinel: Detecting and detoxifying backdoors in diffusion models via temporal noise consistency},
author={Wang, Bingzheng and Gu, Xiaoyan and Xu, Hongbo and Li, Hongcheng and Yu, Zimo and Zhou, Jiang and Wang, Weiping},
journal={arXiv preprint arXiv:2602.01765},
year={2026}
}
Este proyecto se ha construido con una referencia sustancial al excelente codebase NaviT2I. Agradecemos sinceramente a los autores de NaviT2I por publicar su implementación y los recursos de apoyo.
También agradecemos a los autores de BadT2I, EvilEdit, VillanDiffusion, PersonalBKD, Hugging Face Diffusers y los modelos de difusión preentrenados asociados por poner a disposición del público su código y sus checkpoints. Todos los componentes de terceros siguen sujetos a sus respectivas licencias y condiciones de uso.