Your Text Encoder Can Be An Object-Level Watermarking Controller

TL;DR

The paper addresses the challenge of watermarking AI-generated images by proposing an in-generation, object-level watermarking approach for latent diffusion models. It introduces a dedicated watermark token $\bm{\mathcal{W}_*}$ added to the text encoder, enabling selective watermarking of image regions via cross-attention while preserving overall image quality. The method optimizes a dual loss with an empirically chosen optimal timestep $\tau^*=8$, achieving high bit accuracy (up to $99\%$ for $48$ bits) with a dramatic $10^5\times$ reduction in trainable parameters, and demonstrates plug-and-play compatibility with Stable Diffusion variants as well as Textual Inversion. Object-level localization, robustness to common attacks, and compatibility with personalized diffusion pipelines highlight the practical impact for provenance, copyright protection, and controllable generation.

Abstract

Invisible watermarking of AI-generated images can help with copyright protection, enabling detection and identification of AI-generated media. In this work, we present a novel approach to watermark images of T2I Latent Diffusion Models (LDMs). By only fine-tuning text token embeddings $W_*$, we enable watermarking in selected objects or parts of the image, offering greater flexibility compared to traditional full-image watermarking. Our method leverages the text encoder's compatibility across various LDMs, allowing plug-and-play integration for different LDMs. Moreover, introducing the watermark early in the encoding stage improves robustness to adversarial perturbations in later stages of the pipeline. Our approach achieves $99\%$ bit accuracy ($48$ bits) with a $10^5 \times$ reduction in model parameters, enabling efficient watermarking.

Figures (16)

• Figure 1: Text prompt controlled Object-level watermarking. Our method takes a text prompt $\mathcal{P}$ and a watermarking token $\mathcal{W}_*$ as input and integrates watermarking into text-to-image generation. (Left Image): Given a watermarking token $\mathcal{W}_*$, the user can pick any subset of tokens $\{\mathcal{P}_i, \mathcal{P}_j, ...\}$ within the prompt $\mathcal{P}$, we then watermark specific objects corresponding to these user-selected tokens within the image using cross-attention maps. The user has the flexibility to watermark full-image or watermark one or more specific objects while perfectly preserving non-watermarked regions. (Right Image): We demonstrate our method's ability to perform object watermarking with personalization using Styled UNet and Textual Inversion.
• Figure 2: $\bm{\mathcal{W}_*}$ training pipeline.(Left) To find $\mathcal{W}_*$ token embeddings, we use an Img2Img generation pipeline. $\mathcal{D}$ and $\mathcal{D}_w$ represent VAE decoder in the LDM, and Watermark Detector respectively. While training, we send the input image through LDM encoder to retrieve the latent $z_0$, we then add a forward diffusion noise of $\tau^*$ timesteps, followed by iteratively denoising $z_{\tau^*}$ using Classifier-Free Guidance ho2022classifierfreediffusionguidance from $[\tau^* \to 0]$ to retrieve $z'_{0,w}$. During the denoising process, we train for $\mathcal{W}_*$ token embeddings. (Middle) We use latent matching loss to control the trajectory of watermarked latents and bit loss to find $\mathcal{W}_*$ token embeddings. (Right) We then use trained $\mathcal{W}_*$ embeddings to generate watermarked images.
• Figure 3: Qualitative results and watermark heatmaps.We present qualitative results of our watermarking approach while watermarking up to three objects within a single image. All the images are watermarked within the T2I generation pipeline. Recall from \ref{['sec:obj_wat_method']}, we control object-level watermarking directly from the text prompt $\mathcal{P}$. (Top row) presents single object watermarking along with the watermark heatmap for each generated image. We see high bit accuracy within each object selected for watermarking and 0 outside the object. (Second row, third row) presents our qualitative results for two and more than three objects, respectively. We see that our method can accurately retrieve watermarked objects with high bit accuracy. We see (bottom row, last column) that when attention maps are imperfectly defined, the watermark is not confined within the objects, but high bit accuracy is still achieved within the image.
• Figure 4: Plug-and-Play ability of our method.We present our method’s ability to be plugged into any combination of personalized T2I model. Above image shows four such combinations where we use Textual Inversion style tokens and styled T2I pipelines can be seen. It can be observed that the object-level control and watermarking ability of our method is preserved across all these pipelines.
• Figure 5: Watermark heatmap enhancement using SAM.Our watermarking is embedded into the generation pipeline using Attention maps. We test the performance of our watermark detection by plugging our method into P2P and SAM.