Wukong Framework for Not Safe For Work Detection in Text-to-Image systems

TL;DR

Wukong is a transformer-based framework integrated into diffusion models for early, efficient, and accurate NSFW detection in text-to-image generation, leveraging intermediate denoising outputs and cross-attention features.

Contribution

The paper introduces Wukong, a novel method that detects NSFW content during the diffusion process by utilizing intermediate outputs and shared attention parameters, improving efficiency and accuracy.

Findings

Wukong outperforms text-based safeguards in accuracy.

Wukong achieves comparable results to image filters.

Wukong enables early NSFW detection during image generation.

Abstract

Text-to-Image (T2I) generation is a popular AI-generated content (AIGC) technology enabling diverse and creative image synthesis. However, some outputs may contain Not Safe For Work (NSFW) content (e.g., violence), violating community guidelines. Detecting NSFW content efficiently and accurately, known as external safeguarding, is essential. Existing external safeguards fall into two types: text filters, which analyze user prompts but overlook T2I model-specific variations and are prone to adversarial attacks; and image filters, which analyze final generated images but are computationally costly and introduce latency. Diffusion models, the foundation of modern T2I systems like Stable Diffusion, generate images through iterative denoising using a U-Net architecture with ResNet and Transformer blocks. We observe that: (1) early denoising steps define the semantic layout of the image, and (2) cross-attention layers in U-Net are crucial for aligning text and image regions. Based on these insights, we propose Wukong, a transformer-based NSFW detection framework that leverages intermediate outputs from early denoising steps and reuses U-Net's pre-trained cross-attention parameters. Wukong operates within the diffusion process, enabling early detection without waiting for full image generation. We also introduce a new dataset containing prompts, seeds, and image-specific NSFW labels, and evaluate Wukong on this and two public benchmarks. Results show that Wukong significantly outperforms text-based safeguards and achieves comparable accuracy of image filters, while offering much greater efficiency.