Unified Safe In-context Image Generation in Multimodal Diffusion Transformers via Restricting Unsafe Information Flows

Diffusion transformers (DiTs) equipped with multimodal attention (MM-Attn) have become a dominant paradigm for image generation. However, preventing the generation of harmful content remains a critical challenge, particularly in image-to-image (I2I) editing tasks. Existing safety mechanisms are primarily designed for text-to-image (T2I) synthesis or U-Net-based architectures, which limits their effectiveness for unified safety mitigation in DiT-based frameworks. To bridge this gap, we propose Unified Visual Safety Regulator (UVR), a training-free safe generation framework that regulates unsafe semantics in generated images. UVR is grounded in an analysis of attention dynamics from the perspective of information flow in MM-Attn. We identify a task-independent start-up stage, during which unsafe semantics in output patches rapidly emerge and can be accurately localized, followed by task-specific semantic amplification and interference stages, where harmful signals are further propagated and entangled with benign content. Based on these observations, UVR mitigates unsafe generation through unified, targeted attention modulation and explicit restriction of harmful information flow over the identified unsafe output patches. Experiments across various concepts show that UVR achieves state-of-the-art safety performance by achieving 91% and 77% erase rate in image synthesis and editing tasks, while preserving visual quality and fidelity with minimal degradation. Code is available at https://github.com/deng12yx/UVR.

Modern image generation models can now create images from text prompts and edit images based on reference pictures. However, this also creates new safety risks: harmful content may come not only from unsafe text prompts, but also from unsafe reference images used during editing. Existing safety methods often focus on blocking prompts or checking the final image, which may miss harmful information that appears and spreads during the generation process. We propose Unified Visual Safety Regulator (UVR), a training-free method that monitors the internal generation process of multimodal diffusion transformers. Instead of treating text-to-image generation and image editing as separate problems, UVR looks at how unsafe information flows through multimodal attention and affects the output image tokens. It identifies risky visual regions as they emerge, then blocks the unsafe information flow before it is further propagated into the final image. This unified process-level safety mechanism can reduce unsafe outputs in both text-to-image generation and image-to-image editing while preserving visual quality and editing ability. Our work helps make modern image generation systems safer without requiring model retraining or separate safety designs for different tasks.