TNC-Defense 提供用于检测和去除文本到图像扩散模型中后门的研究代码。该仓库目前包含针对 Stable Diffusion v1.4、Stable Diffusion v1.5、Stable Diffusion XL 和 Stable Diffusion 3 Medium 的检测流程,以及针对 Stable Diffusion v1.4 和 Stable Diffusion XL 的去毒流程。
仅限研究使用。 此代码旨在研究文本到图像扩散模型的鲁棒性与安全性。用户有责任遵守所有上游模型、数据集和实现的相关许可与使用政策。
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
该代码已在 TCN Conda 环境中通过冒烟测试,所用版本如下:
分析和评估工具还额外需要 NumPy、pandas、SciPy、scikit-learn、Matplotlib、Pillow 和 tqdm 等软件包。
运行实验前,请先激活环境:
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
Bash 入口使用当前激活环境中的 python。也可以通过 PYTHON 环境变量指定其他解释器:
PYTHON=python bash TNC-Detect/run_detect_sd14.sh
所有仓库路径均相对于各脚本所在位置进行解析,因此命令不依赖当前工作目录。默认模型路径遵循实验所用的本地项目布局。如果模型存储在其他位置,请覆盖相应的命令行参数,如 --base-model-path、--backdoor-model-path、--base-model 或 --tox-path。
去毒数据集位于:
TNC-Detox/data/
├── sd14/
│ ├── pixel/
│ ├── pixel_mul_tokens/
│ ├── eviledit/
│ ├── personal_bkd/
│ └── villain_mul/
└── sdxl/
├── pixel/
└── pixel_mul_tokens/
每个数据集均包含 clean/、poison/ 和 metadata.json。
默认运行脚本会评估所有已配置的 SD1.4 后门方法。它会将各方法分配到 run_detect_sd14.sh 中定义的 GPU ID,并并行启动它们。
cd TNC-Defense/TNC-Detect
bash run_detect_sd14.sh
直接运行某个方法:
python detect_sd14.py \
--backdoor-method pixel \
--device cuda:0 \
--seed 0 \
--num-inference-steps 50
每个模型系列可用的参数可通过以下命令查看:
python detect_sd14.py --help
python detect_sd15.py --help
python detect_sdxl.py --help
python detect_sd3.py --help
生成 MSE CSV 文件后,在 ana.sh 中配置输入文件并运行:
bash ana.sh
检测器支持朴素(naive)和动态 k(dynamic-k)两种决策规则。指标和图表由 ana.py 生成。
默认 SD1.4 运行脚本会在脚本中指定的五个 GPU ID 上启动五种已配置的方法:
cd TNC-Defense/TNC-Detox
bash run_detox_sd14.sh
运行单个方法:
python detox_sd14.py \
--method pixel \
--device cuda:0 \
--seed 42
在 run_detox_sdxl.sh 中配置 METHOD 和 GPU_ID,然后运行:
bash run_detox_sdxl.sh
各方法特有的默认设置(包括训练步数、时间步采样比例、损失权重、批大小和检查点保存间隔)在每个训练脚本的 METHOD_PROFILES 部分定义。命令行参数可以覆盖这些默认值。
从一个或多个去毒后的 SD1.4 UNet 生成图像:
python evaluate_sd14.py \
--detox-model checkpoints/sd14/<checkpoint> \
--prompt-file <prompt-file>
计算参考图像目录与去毒图像目录之间的 FID:
python compute_fid.py \
--reference-dir <reference-images> \
--candidate-dir <detoxified-images>
42,并启用确定性 cuDNN 行为。本项目使用的带后门模型和攻击配置来自以下公开资源:
使用这些攻击或检查点时,请引用相应的原论文:BadT2I(badt2izhai2023)、EvilEdit(evileditwang2024)、VillanDiffusion(villandiffusionchou2023)和 PersonalBKD(personalizationhuang2024)。
如果您觉得本项目对您的研究有帮助,请引用:
@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}
}
本项目的构建大量参考了优秀的 NaviT2I 代码库。我们衷心感谢 NaviT2I 作者发布其实现和支持资源。
我们还要感谢 BadT2I、EvilEdit、VillanDiffusion、PersonalBKD、Hugging Face Diffusers 及相关预训练扩散模型的作者公开其代码和检查点。所有第三方组件均受其各自许可和使用条款的约束。