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

TL;DR

This work tackles NSFW content in text-to-image generation by introducing Wukong, a model-specific external safeguard that detects unsafe outputs within the diffusion process. By using intermediate latent representations from early denoising steps and reusing pre-trained cross-attention parameters from the U-Net, Wukong achieves fast, accurate NSFW detection with low latency. The approach is trained as a multi-label classifier over seven NSFW categories and can halt generation at an early step $T_C$ if unsafe content is detected, offering substantial efficiency gains over post-hoc image filtering. A new Wukong-Demons dataset provides image-grounded, model/seed-sensitive NSFW annotations, and experiments show Wukong outperforms text-based safeguards while matching image-based methods in accuracy, with robust performance against adversarial prompts. Overall, the framework advances practical deployment of safe T2I systems by combining model-aware signals with efficient, early-stage detection.

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.

Figures (6)

• Figure 1: An illustrative example of modifying the textual condition during the early denoising steps in the Stable Diffusion process. The initial latent noise is conditioned with an unsafe prompt (prompt_1) in the first step. From the second step onward, the first row remains conditioned on prompt_1, while the second row switches to a safe prompt (prompt_2).
• Figure 2: Visualization of Attention Maps. (a) shows an image generated with the prompt "car" using Stable Diffusion XL. (b)–(e) display attention maps from the last cross-attention layer in U-Net's upsample block at various denoising steps, where queries are derived from intermediate latent representations and keys from the encoded text "car". (f) shows an image generated with the prompt "a naked woman", while (g)–(j) show corresponding attention maps using the NSFW concept "sexual" as the key. All attention maps are reshaped from $4096\times 1$ to $64\times 64$, then resized to $1024\times 1024$ for visualisation.
• Figure 3: Overview of the Wukong Framework . The left portion illustrates the full denoising process of the Stable Diffusion model, where $x_{T-t}$ denotes the latent representation at denoising step $t$. In the Wukong framework, the diffusion proceeds through $T_C-1$ full denoising steps and up to, but not including, the final cross-attention layer in the upsample block at step $T_C$, producing an intermediate latent representation $\varphi\left(x_{\left(T-T_C\right)}\right)$. The remaining $T-T_C$ steps are skipped (typically $T_C\ll T$). The right portion shows the transformer-based decoder used for NSFW detection. The query and key weight $W_Q$ and $W_K$ are reused from the final cross-attention layer of the upsample block in the U-Net (highlighted with a red dashed rectangle, and defined in Equations \ref{['E_NSFW']} and \ref{['new_KV']}). For clarity, auxiliary layers such as LayerNorm are omitted.
• Figure 4: The performances comparison varying $T_C$.
• Figure 5: F1-score under different values of $\delta$.