binzhwang/tnc-defenseИсследовательский код для обнаружения и обезвреживания бэкдоров в диффузионных моделях генерации изображений по текстовому описанию, с пайплайнами для Stable Diffusion v1.4, v1.5, XL и 3 Medium.
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
Код прошёл смоук-тест в Conda-окружении TCN со следующим набором зависимостей:
Инструменты анализа и оценки дополнительно требуют такие пакеты, как 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. Он распределяет методы по GPU ID, определённым в run_detect_sd14.sh, и запускает их параллельно.
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
После формирования CSV-файлов MSE настройте входные файлы в ana.sh и выполните:
bash ana.sh
Детектор поддерживает как наивное правило принятия решений, так и правило с динамическим 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
Настройте METHOD и GPU_ID в run_detox_sdxl.sh, а затем выполните:
bash run_detox_sdxl.sh
Значения по умолчанию для конкретных методов, включая количество шагов обучения, коэффициенты сэмплирования временных шагов, веса функций потерь, размеры батчей и интервалы сохранения контрольных точек, определены в разделе METHOD_PROFILES каждого обучающего скрипта. Аргументы командной строки могут переопределять эти значения.
Сгенерируйте изображения с помощью одной или нескольких обезвреженных UNet-моделей SD1.4:
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 и связанных предобученных диффузионных моделей за то, что они сделали свой код и контрольные точки общедоступными. Все сторонние компоненты остаются под действием соответствующих лицензий и условий использования.