Diffusion-Based Dataset Augmentation with Safety in Mind

Usage

1. For building the running environment:

You can use SGLang Docker and then in your container terminal:

git clone https://github.com/alikhalajii/diffBased-DSAug-Safety.git
pip install git+https://github.com/ml-research/ledits_pp.git
pip install git+https://github.com/openai/CLIP.git

Alternatively, create the Conda environment using the provided YAML file:

git clone https://github.com/alikhalajii/diffBased-DSAug-Safety.git
cd diffBased-DSAug-Safety
conda env create -f sgl-lpp.yml
conda activate sgl-lpp

2. Set the Hugging Face token with access to the following repositories:

AIML-TUDA/LlavaGuard-7B, AIML-TUDA/LlavaGuard-13B, llava-hf/llama3-llava-next-8b-hf, llava-v1.6-vicuna-13b-hf

and activate your access token:

huggingface-cli login

3. Reinstall the flashinfer package:

pip uninstall flashinfer
pip install flashinfer -i https://flashinfer.ai/whl/cu121/torch2.3/

4. Launch the SGLang server as backend with one of the below LlavaGuard checkpoints.

CUDA_VISIBLE_DEVICES=0 python3 -m sglang.launch_server --model-path AIML-TUDA/LlavaGuard-7B --tokenizer-path llava-hf/llava-1.5-7b-hf --port 10000
CUDA_VISIBLE_DEVICES=0 python3 -m sglang.launch_server --model-path AIML-TUDA/LlavaGuard-13B --tokenizer-path llava-hf/llava-1.5-13b-hf --port 10000

5. Once the server is set up, run the safeguarding pipeline:

CUDA_VISIBLE_DEVICES=1 python3 src/main_pipe.py --input_dir INPUT_DIR --output_dir OUTPUT_DIR --llava_model "llava-hf/llama3-llava-next-8b-hf"
CUDA_VISIBLE_DEVICES=1 python3 src/main_pipe.py --input_dir INPUT_DIR --output_dir OUTPUT_DIR --llava_model "llava-hf/llava-v1.6-vicuna-13b-hf"

To limit the number of images in the input directory to a specific number (e.g., 100):

CUDA_VISIBLE_DEVICES=1 python src/main_pipe.py --input_dir INPUT_DIR --output_dir OUTPUT_DIR --llava_model LLAVA_MODEL_ID --num_images 100

Abstract

The rapid rise of large-scale image datasets, often sourced without ethical or cultural considerations, has raised significant safety concerns. These datasets frequently include unsafe or explicit content, such as depictions of violence, self-harm, or nudity, which may perpetuate harmful societal narratives when incorporated into model training. Addressing this issue poses a dual challenge: ensuring the datasets remain ethically compliant while preserving their diversity and utility for real-world applications. This thesis proposes a diffusion-based dataset augmentation approach leveraging pre-trained denoising diffusion models and VLMs to transform unsafe elements within images. The method focuses on refining inappropriate content by selectively altering high-level visual features, guided by dynamically generated textual terms. By doing so, the pipeline preserves the contextual and semantic coherence of the original images, offering an alternative to outright content removal. Experimental results demonstrate the pipeline's efficacy in addressing safety concerns, particularly in localized categories like nudity. However, challenges persist in refining complex safety categories, where unsafe elements are often intertwined with broader visual narratives. This work underscores the potential of integrating diffusion models and VLMs to address safety challenges in image datasets.

Zusammenfassung

Die rasche Zunahme großer Bilddatensätze, die oft ohne Berücksichtigung ethischer oder kultureller Aspekte erstellt werden, wirft erhebliche Sicherheitsbedenken auf. Solche Datensätze enthalten häufig unsichere oder explizite Inhalte, wie Darstellungen von Gewalt, Selbstverletzung oder Nacktheit, die schädliche gesellschaftliche Narrative verstärken können, wenn sie in das Modelltraining einfließen. Die Herausforderung besteht darin, Datensätze sowohl ethisch vertretbar zu gestalten als auch ihre Vielfalt und ihren Nutzen für reale Anwendungen zu bewahren.Diese Arbeit stellt einen diffusionsbasierten Ansatz zur Datensatzerweiterung vor, der vortrainierte Diffusionsmodelle und Vision-Sprachmodelle nutzt, um unsichere Bildelemente zu transformieren. Durch die selektive Änderung hochdimensionaler visueller Merkmale, gesteuert durch dynamisch generierte Textbegriffe, bewahrt die Methode die kontextuelle und semantische Kohärenz der Originalbilder und bietet eine Alternative zur vollständigen Entfernung sensibler Inhalte. Die experimentellen Ergebnisse zeigen, dass die Pipeline Sicherheitsbedenken erfolgreich adressiert, insbesondere in Kategorien wie Nacktheit mit klar definierten semantischen Grenzen. In komplexeren Sicherheitskategorien wie Gewalt oder Selbstverletzung, bei denen unsichere Elemente oft mit umfassenderen visuellen Kontexten verknüpft sind, bestehen jedoch weiterhin Herausforderungen. Die Arbeit hebt das Potenzial hervor, Diffusionsmodelle und Vision-Sprachmodelle zu integrieren, um Sicherheitsprobleme in Bilddatensätzen effektiver zu bewältigen.