An adaptive watermarking method for diffusion models robust to geometric attacks

Abstract

This invention relates to a diffusion model-based adaptive watermarking method robust to geometric attacks and a neural network model, belonging to the field of generative artificial intelligence technology. It addresses the difficulty in achieving a balance between embedding strength and image visual quality in existing diffusion model watermarking techniques. In the watermark embedding stage, the method includes: obtaining an initial latent vector generated by the diffusion model based on user-input prompts; encoding the watermark using a watermark encoder to obtain a watermark feature vector with the same vector space size as the initial latent vector; calculating the embedding strength of the initial latent vector using a two-branch adaptive strength generator based on an Illumination-Adaptive-Transformer network structure to obtain a spatial adaptive embedding strength map of the initial latent vector; embedding the watermark feature vector into the initial latent vector according to the spatial adaptive embedding strength map to obtain a watermarked latent vector; and converting the watermarked latent vector into a watermarked image using the diffusion model.

Description

Technical Field

[0001] This invention relates to the field of generative artificial intelligence technology, and in particular to an adaptive watermarking method for diffusion models that is robust to geometric attacks, and a neural network model. Background Technology

[0002] With the rapid development of deep learning technology, generative artificial intelligence (AIGC) technology, represented by diffusion models, has made breakthrough progress in the field of image synthesis, capable of generating highly realistic digital content based on text descriptions. However, this has also raised serious security issues such as copyright ownership and the spread of misinformation, making digital watermarking technology a key means to ensure content credibility and traceability.

[0003] In the existing technological system, watermarking schemes for diffusion models mainly present four technical paths. The first is the image post-processing embedding method. Technicians usually use discrete cosine transform or deep neural networks to superimpose watermark signals in the pixel domain after the image is generated. Although this method is direct, the watermark signal is completely decoupled from the generation process and is easily removed by conventional denoising, compression, and other editing techniques. The second is the model weight fine-tuning method, which involves retraining the diffusion model by introducing a specific loss function to embed the watermark features into the model parameters, so that each generated image has a fixed identifier. However, this method is not only computationally expensive, but also difficult to flexibly change the bit information contained in the watermark once the model is deployed. The third approach is embedding methods targeting the initial noise of the diffusion model. This is a relatively new approach, which involves pre-setting specific spatial structures or spectral features (such as setting specific distributions in the Fourier ring region) in the random Gaussian noise before the diffusion iteration begins, attempting to allow the watermark information to be naturally coupled to the final image during the denoising process. However, due to the highly random nature of the diffusion model denoising path, the weak signals in the initial noise are prone to uncontrollable attenuation or distortion during long-term iterative evolution, making it difficult to obtain stable signals at the extraction end. The fourth approach is embedding methods based on the latent space generation process. These methods utilize the latent diffusion model to directly modulate information on the VAE latent vector. Because the latent space is rich in high-level semantics and highly integrated with the generation process, it has a natural advantage in terms of invisibility.

[0004] However, the aforementioned existing technologies still reveal two fatal technical bottlenecks when facing complex and ever-changing real-world application scenarios. First, existing solutions generally lack an adaptive balance mechanism between watermark embedding strength and visual quality. Most technologies employ a globally uniform embedding step size, or can only make coarse adjustments based on local pixel variance, failing to establish a deep perception of the global semantic layout and local texture complexity of the image. This results in visible color block artifacts in smooth areas such as the sky and human faces due to excessive embedding; while in high-frequency texture areas such as hair and grass, insufficient embedding strength causes the watermark signal to disappear completely in the slightest post-processing fluctuations, making it impossible to achieve optimal embedding across the entire image.

[0005] Secondly, existing watermarking methods suffer from severe inadequacy in robustness against geometric transformation attacks. Most current extraction networks are based on standard convolutional neural network architectures, whose core mechanisms heavily rely on the strict spatial synchronization and alignment of image pixels or feature points. When a watermarked image is subjected to geometric attacks such as rotation, scaling, cropping, or affine transformations, although the semantic content of the image remains essentially unchanged, the grid positions of its latent features will be resampled and shifted, resulting in a complete disruption of spatial synchronization. In this situation, the extractor cannot read the correct watermark code at the expected coordinates, leading to a sharp increase in the bit error rate and even causing complete extraction failure. Summary of the Invention

[0006] Based on the above analysis, the embodiments of the present invention aim to solve the problem that existing diffusion model watermarking technology is difficult to balance between embedding strength and image visual quality.

[0007] In a first aspect, embodiments of the present invention provide an adaptive watermarking method for diffusion models that is robust to geometric attacks. During the watermark embedding stage, the method includes:

[0008] Obtain the initial latent vector generated by the diffusion model based on the prompts input by the user; The watermark is encoded by a watermark encoder to obtain a watermark feature vector that has the same vector space size as the initial potential vector. The embedding strength of the initial latent vector is calculated using a dual-branch adaptive intensifier based on an Illumination-Adaptive-Transformer network structure, resulting in a spatial adaptive embedding strength map of the initial latent vector. Based on the spatial adaptive embedding intensity map, the watermark feature vector is embedded into the initial latent vector to obtain a watermarked latent vector; and Using the diffusion model, the watermarked latent vector is converted into a watermarked image.

[0009] Based on a further improvement of the above method, in the watermark extraction stage, the method includes: The watermarked image to be detected is input into an image corrector based on a spatial transformation network to obtain the corrected image of the image to be detected. The corrected image is input into the diffusion model to obtain the latent vector of the corrected image; The latent vector of the corrected image is input into the watermark extractor to obtain the target watermark feature vector, wherein the network structure of the watermark extractor is the same as that of the dual-branch adaptive intensifier. The target watermark feature vector is input into the watermark decoder to obtain the watermark, wherein the watermark decoder adopts an inverse network structure that is symmetrical to the network structure of the watermark encoder.

[0010] Based on a further improvement of the above method, the embedding strength of the initial latent vector is calculated using a dual-branch adaptive intensifier based on an Illumination-Adaptive-Transformer network structure, resulting in a spatial adaptive embedding strength map of the initial latent vector, including: The initial latent vector is input into a local branch of the dual-branch adaptive intensifier to obtain a high-frequency feature map that reflects the local texture complexity. The initial latent vector is input into the global branch of the dual-branch adaptive intensifier to obtain a global modulation vector reflecting the overall semantic sensitivity of the image; and The high-frequency feature map is modulated by affine transformation using the global modulation vector, and the modulated high-frequency feature map is input into a 1×1 convolutional layer and a sigmoid activation function to obtain a spatial adaptive embedding intensity map.

[0011] Based on a further improvement of the above method, according to the spatial adaptive embedding intensity map, the watermark feature vector is embedded into the initial latent vector to obtain a watermarked latent vector, including: The watermarked potential vector is generated according to the following formula. :

[0012] in, Let be the initial potential vector. This is the spatial adaptive embedding intensity map. The watermark feature vector, This is the preset global intensity control factor. This is for element-wise multiplication.

[0013] Based on a further improvement of the above method, the image corrector includes: a texture feature enhancement module, a parameter regression subnetwork, and a differentiable resampling layer, wherein: The texture feature enhancement module is used to extract the high-frequency feature map of the image to be detected and fuse the high-frequency feature map with the image to be detected; The parameter regression subnetwork is used to receive the output of the texture feature enhancement module and regress the affine transformation matrix for counteracting geometric attacks through a convolutional neural network; and The differentiable resampling layer is used to establish a coordinate mapping relationship based on the affine transformation matrix, and to resample the image to be detected to obtain the corrected image of the image to be detected.

[0014] Based on a further improvement of the above method, the texture feature enhancement module includes a fixed Sobel operator layer, and the parameter regression subnetwork is a convolutional neural network structure based on ResNet-18, wherein the last layer of the parameter regression subnetwork is a fully connected layer for outputting a 2×3 affine transformation matrix.

[0015] Based on a further improvement to the above method, the differentiable resampling layer resamples the image to be detected according to the following formula:

[0016] in, and To correct the pixel value of the i-th pixel in the image, and Let be the pixel value of the i-th pixel in the image to be detected. Let be the affine transformation matrix.

[0017] Based on a further improvement of the above method, the latent vector of the corrected image is input into the watermark extractor to obtain the target watermark feature vector, which includes: The target watermark feature vector is generated according to the following formula. :

[0018] in, The latent vector of the corrected image, The spatially adaptive embedding intensity map is the latent vector of the corrected image.

[0019] Secondly, an embodiment of the present invention provides a neural network model for an adaptive watermarking method, comprising: The system consists of a diffusion model, a watermark encoder, a bi-branch adaptive strength generator, a feature fusion layer, an image corrector, a watermark extractor, and a watermark decoder, among which: The diffusion model includes a variational autoencoder and a variational autodecoder, wherein the variational autoencoder is used to generate an initial latent vector based on prompt words input by the user; The watermark encoder is used to encode the watermark to obtain a watermark feature vector that has the same vector space size as the initial potential vector. The dual-branch adaptive intensifier is used to calculate the embedding strength of the initial latent vector to obtain the spatial adaptive embedding strength map of the initial latent vector. The feature fusion layer is used to embed the watermark feature vector into the initial latent vector according to the spatial adaptive embedding intensity map, so as to obtain a watermarked latent vector; The variational autodecoder is used to convert the watermarked latent vector into a watermarked image. The image corrector employs a spatial transformation network, which is used to correct the image to be detected to obtain a corrected image; The variational autoencoder is further configured to generate a latent vector of the corrected image based on the corrected image; The watermark extractor has the same mesh structure as the dual-branch adaptive intensifier, and it is used to convert the latent vector of the corrected image into the target watermark feature vector. The watermark decoder employs a reverse network structure that is symmetrical to the network structure of the watermark encoder, and is used to convert the target watermark feature vector into a watermark.

[0020] Based on the further improvement of the above model, its training method is as follows: In the first stage, all parameters of the diffusion model, the bi-branch adaptive intensifier, and the image corrector are frozen, and only the watermark encoder and the watermark decoder are trained. In the second stage, firstly, all parameters of the watermark decoder and the image corrector are frozen, the watermark encoder and the dual-branch adaptive strength generator are activated, and the dual-branch adaptive strength generator is trained. Secondly, all parameters of the dual-branch adaptive strength generator are frozen, the image corrector is activated, and the image corrector is trained. In the third stage, the bi-branch adaptive intensifier, the image corrector, the watermark extractor, and the watermark decoder are activated, the diffusion model is frozen, and the neural network model is trained. The total loss function is... for:

[0021] in, The message cross-entropy loss is used to ensure the readability of the watermark under attack. Image perceptual loss constrains image visual quality; Geometric alignment constraints are used to assist the convergence of the image corrector; The weights are the message cross-entropy loss values; The weights for image perception loss; The weights of the geometric alignment constraints.

[0022] Compared with the prior art, the present invention can achieve at least one of the following beneficial effects: 1. Significantly improves robustness against geometric transformation attacks: The differentiable STN image corrector introduced in this invention achieves automatic spatial alignment of the attacked image, effectively restoring the spatial synchronization of the watermark signal, so that the model can still maintain extremely high watermark extraction accuracy when subjected to arbitrary angle rotation or large-scale scaling.

[0023] 2. Achieving an adaptive balance between image fidelity and watermark embedding depth: Based on an improved IAT architecture, a dual-branch adaptive strength generator combines the pixel-level texture perception capability of the local branch with the overall semantic control capability of the global branch. This ensures that the watermark is mainly embedded in high-frequency texture areas, while effectively avoiding smooth backgrounds, thus resolving the contradiction between invisibility and robustness in generative watermarking.

[0024] 3. Excellent training stability and modular compatibility: The phased, progressive training strategy effectively solves the convergence problem of complex multi-module networks. Furthermore, since the core operations are performed in the latent space, there is no need to retrain the main diffusion model, resulting in advantages such as low cost and easy adaptation.

[0025] In this invention, the above-described technical solutions can be combined with each other to achieve more preferred combinations. Other features and advantages of this invention will be set forth in the following description, and some advantages may become apparent from the description or be learned by practicing the invention. The objects and other advantages of this invention can be realized and obtained from what is particularly pointed out in the description and drawings. Attached Figure Description

[0026] The accompanying drawings are for illustrative purposes only and are not intended to limit the invention. Throughout the drawings, the same reference numerals denote the same parts. Figure 1 A flowchart illustrating an adaptive watermarking method for a diffusion model robust to geometric attacks, according to an embodiment of the present invention, is shown.

[0027] Figure 2 A flowchart illustrating the watermark extraction stage of the adaptive watermarking method for a diffusion model robust to geometric attacks according to an embodiment of the present invention is shown.

[0028] Figure 3The diagram shows the robustness of the neural network model proposed in this embodiment to various types of noise. Detailed Implementation

[0029] Preferred embodiments of the present invention will now be described in detail with reference to the accompanying drawings, which form part of this application and are used together with the embodiments of the present invention to illustrate the principles of the present invention, but are not intended to limit the scope of the present invention.

[0030] According to one aspect of the present invention, an adaptive watermarking method for diffusion models that is robust to geometric attacks is provided. Figure 1 A schematic diagram of the watermark embedding stage of the adaptive watermarking method for a diffusion model robust to geometric attacks according to an embodiment of the present invention is shown. Figure 1 As shown, this adaptive watermarking method for a diffusion model robust to geometric attacks includes the following steps in the watermark embedding stage: Step 101: Obtain the initial latent vector generated by the diffusion model based on the prompt words input by the user.

[0031] In this embodiment, after the user inputs a text prompt, a pre-trained diffusion model (e.g., Stable Diffusion v2.1) can be invoked. The diffusion model includes a variational autoencoder and a variational autodecoder. The variational autoencoder first performs multi-step denoising sampling on the tensors in the latent space (e.g., using a DDIM sampler) to generate an initial latent vector. In this embodiment, the dimension of the initial potential vector is... (e.g. 4) ).

[0032] Step 102: Encode the watermark using a watermark encoder to obtain a watermark feature vector that has the same vector space size as the initial potential vector.

[0033] In this embodiment, after obtaining the binary watermark M to be embedded (e.g., 64 bits), M can be expanded into a watermark feature vector using a watermark encoder. Specifically, the watermark encoder may contain fully connected layers and a reshape operation to map and reshape the bitstream to the initial latent vector. Tensors of uniform size in vector space (e.g., 4) The Reshape operation is used to change the shape of a tensor, but does not change its underlying data content.

[0034] Step 103: Calculate the embedding strength of the initial latent vector using a dual-branch adaptive intensifier based on the Illumination-Adaptive-Transformer network structure to obtain the spatial adaptive embedding strength map of the initial latent vector.

[0035] To enhance robustness while maintaining image quality, this embodiment designs a dual-branch adaptive intensifier based on an Illumination-Adaptive-Transformer network structure to refine the initial latent vector. Analysis is performed. This dual-branch adaptive intensifier is based on a parallel two-stream architecture, analyzing the initial latent vector from two dimensions: local texture details and global semantic layout, and automatically calculating a spatial adaptive embedding intensity map.

[0036] The dual-branch adaptive intensifier in this embodiment includes: 1. Local Branch: This local branch is used to generate a high-frequency feature map reflecting the local texture complexity. This high-frequency feature map contains pixel-level basic texture features. For example, the input initial latent vector... A high-frequency feature map reflecting local texture complexity is extracted using a network consisting of two 3×3 depthwise separable convolutional layers. This high-frequency feature map It not only includes edge information, but also preserves the initial latent vector. Spatial structure.

[0037] 2. Global Branch: This global branch is used to generate a global modulation vector that reflects the overall semantic sensitivity of the image. For example, the initial latent vector is input... After downsampling, the data is input into the Transformer module, where a self-attention mechanism is used to aggregate global information. Unlike the conventional output feature map, the end of this global branch is connected to a global average pooling layer and a fully connected layer, regressing a set of global modulation vectors γ and β. These parameters reflect the overall semantic sensitivity of the image.

[0038] 3. Adaptive Modulation and Generation: Utilizing the global modulation vector to generate high-frequency feature maps. Perform affine transformation modulation and then convert the modulated high-frequency feature map. The input is fed into a 1×1 convolutional layer and a sigmoid activation function to obtain a spatially adaptive embedding intensity map. For example, affine transform modulation is performed on the high-frequency feature map of the local branch output using global modulation vectors γ and β. The specific calculation formula is as follows:

[0039] in, Responsible for scaling the response amplitude of local texture features. Responsible for adjusting the reference bias and modulated features Finally, a 1×1 convolutional layer and a sigmoid activation function are used to generate the final spatially adaptive embedding intensity map. .

[0040] Step 104: Based on the spatial adaptive embedding intensity map, embed the watermark feature vector into the initial latent vector to obtain the watermarked latent vector.

[0041] In this embodiment, the watermark embedding intensity can be enhanced in areas with complex textures and suppressed in areas with smooth or semantically sensitive textures, based on the spatially adaptive embedding intensity map, and the watermark feature vector can be superimposed on the initial latent vector.

[0042] In some embodiments, a watermarked potential vector can be generated according to the following formula. :

[0043] in, Let be the initial potential vector. This is the spatial adaptive embedding intensity map. The watermark feature vector, This is the preset global intensity control factor. This is element-wise multiplication. For example, The value is 1.

[0044] Step 105: Using the diffusion model, the watermarked latent vector is converted into a watermarked image.

[0045] In this embodiment, the watermarked potential vector (such as the one mentioned above) can be used. The input is fed into the variational self-decoder of the diffusion model to generate the final watermarked image. .

[0046] According to another aspect of the present invention, a diffusion model adaptive watermarking method robust to geometric attacks is provided. Figure 2 A flowchart illustrating the watermark extraction stage of the adaptive watermarking method for a diffusion model robust to geometric attacks according to an embodiment of the present invention is shown. Figure 2 As shown, this adaptive watermarking method for diffusion models robust to geometric attacks includes the following steps in the watermark extraction stage: Step 201: Input the watermarked image to be detected into an image corrector based on a spatial transformation network to obtain the corrected image of the image to be detected.

[0047] In this embodiment, the image to be detected is acquired. The image may have undergone geometric attacks such as rotation, scaling, and cropping during transmission. To address the issue of watermark extraction failing after image rotation or scaling, this invention incorporates an image corrector based on a spatial transformation network before extraction. (Image to be detected) The image is first input into the image corrector, which automatically predicts what deformations have occurred in the image and uses a differentiable resampling mechanism to detect the image. Reverse transformation back to standard form.

[0048] In some embodiments, the image corrector includes: a texture feature enhancement module, a parametric regression subnetwork, and a differentiable resampling layer. The texture feature enhancement module is used to extract the image to be detected. The high-frequency feature map is obtained and compared with the image to be detected. The data is then fused. A parameter regression subnetwork receives the output of the texture feature enhancement module and regresses the inverse transform parameters used to counteract geometric attacks via a convolutional neural network. A differentiable resampling layer is used to establish a coordinate mapping relationship based on the inverse transform parameters for the image to be detected. Perform resampling.

[0049] Unlike traditional methods that directly use RGB pixels for parameter prediction, this embodiment considers that the watermark signal is mainly hidden in high-frequency texture areas, and therefore designs a texture enhancement mechanism to assist in geometric correction. Firstly, the image to be detected can be... The input is fed into a fixed Sobel operator layer to extract high-frequency feature maps containing only edges and textures. Then the original image can be... High-frequency feature map The texture enhancement tensor is obtained by splicing and blending along the channel dimension. Finally, the texture enhancement tensor... Input to the parameter regression subnet.

[0050] In this embodiment, the parameter regression subnetwork can employ a lightweight convolutional neural network structure, such as a variant based on ResNet-18, responsible for extracting geometric features from the image. The last layer of the parameter regression subnetwork is a fully connected layer, outputting a 2×3 affine transformation matrix. :

[0051] The affine transformation matrix Characterizes the image to be detected The inverse transform parameters required to restore the standard form.

[0052] To accelerate the convergence of the training parametric regression subnet, the weights of the last layer of the parametric regression subnet can be initialized to 0 and the biases to the identity matrix during training, so that the output image does not change in the initial stage of training.

[0053] In some embodiments, the differential resampling layer can be configured according to the following formula for the image to be detected. Perform resampling:

[0054] in, and To correct the pixel value of the i-th pixel in the image, and To detect the pixel value of the i-th pixel in an image, Let be the affine transformation matrix.

[0055] Due to the calculated sampling coordinates ( , The coordinates are typically non-integer. In this embodiment, bilinear interpolation can be used to calculate the pixel value at these coordinates, ensuring that subsequent watermark decoding errors can be backpropagated to the parameter regression subnet through this layer, thereby driving the network to learn the correct geometric correction strategy. The final output is a spatially aligned corrected image. This restored the spatial synchronization of the watermark signal.

[0056] Step 202: Input the corrected image into the diffusion model to obtain the latent vector of the corrected image.

[0057] In this embodiment, the corrected image can be... The input is fed into the variational autoencoder of the diffusion model, which maps it from the pixel space back to the latent space to obtain the extracted latent vector of the corrected image. .

[0058] Step 203: Input the latent vector of the corrected image into the watermark extractor to obtain the target watermark feature vector, wherein the network structure of the watermark extractor is the same as that of the dual-branch adaptive intensifier.

[0059] In order to accurately separate the weak watermark signal from the background image features, this step will extract the latent vector. Input to the watermark extractor.

[0060] In this embodiment, the watermark extractor adopts the same parallel dual-stream network architecture as the dual-branch adaptive intensifier, which also includes a local texture-aware branch that does not perform downsampling and a global semantic attention branch that captures long-range dependencies.

[0061] The watermark extractor can process the input... Analysis was performed. Because the watermark signal is relatively weak compared to image features, Semantic and texture distribution and initial latent vector The high degree of consistency allows the watermark extractor to reconstruct an approximate spatial adaptive intensity map through a dual-branch network. This spatial adaptive intensity map Precisely indicates the latent vector Which areas contain strong watermark signals, and which areas have almost no watermark?

[0062] This embodiment can reconstruct the spatial adaptive intensity map. As a soft attention mask, for latent vectors Feature decoupling is performed, specifically through the following operations:

[0063] Through this operation, the watermark extractor can automatically "focus" on areas with high watermark embedding intensity, while simultaneously using a spatially adaptive intensity map. The feature response of regions with median values ​​close to 0 is suppressed to 0, ultimately resulting in a watermark feature vector with a significantly improved signal-to-noise ratio. .

[0064] Step 204: Input the target watermark feature vector into the watermark decoder to obtain the watermark, wherein the watermark decoder adopts an inverse network structure that is symmetrical to the network structure of the watermark encoder. In this embodiment, the watermark feature vector can be... The input is fed into the watermark decoder and mapped back to a binary bitstream. Calculate bit accuracy.

[0065] Specifically, the watermark decoder employs a reverse network structure symmetrical to the watermark encoder. For example, firstly, a flattening layer stretches the input four-dimensional feature tensor (dimensions 1×4×64×64) into a one-dimensional feature vector. Subsequently, this vector undergoes dimensionality reduction mapping through multiple fully connected layers, outputting a prediction vector of dimension N (64). Finally, the output value is mapped to a sigmoid activation function. The interval is used to obtain the predicted binary bit stream. The predicted bitstream The watermark information is compared bit by bit with the original watermark information M to calculate the bit accuracy, thereby determining whether the watermark extraction was successful or to conduct copyright ownership authentication.

[0066] According to another aspect of the present invention, a neural network model for implementing any of the above-described adaptive watermarking methods is provided, comprising: The system consists of a diffusion model, a watermark encoder, a bi-branch adaptive strength generator, a feature fusion layer, an image corrector, a watermark extractor, and a watermark decoder, among which: The diffusion model includes a variational autoencoder and a variational autodecoder, wherein the variational autoencoder is used to generate an initial latent vector based on prompt words input by the user; The watermark encoder is used to encode the watermark to obtain a watermark feature vector that has the same vector space size as the initial potential vector. The dual-branch adaptive intensifier is used to calculate the embedding strength of the initial latent vector to obtain the spatial adaptive embedding strength map of the initial latent vector. The feature fusion layer is used to embed the watermark feature vector into the initial latent vector according to the spatial adaptive embedding intensity map, so as to obtain a watermarked latent vector; The variational autodecoder is used to convert the watermarked latent vector into a watermarked image. The image corrector employs a spatial transformation network, which is used to correct the image to be detected to obtain a corrected image; The variational autoencoder is further configured to generate a latent vector of the corrected image based on the corrected image; The watermark extractor has the same mesh structure as the dual-branch adaptive intensifier, and it is used to convert the latent vector of the corrected image into the target watermark feature vector. The watermark decoder employs a reverse network structure that is symmetrical to the network structure of the watermark encoder, and is used to convert the target watermark feature vector into a watermark.

[0067] The functions of each module in the above model are combined with the above text. Figure 1 and Figure 2 The adaptive watermarking method for diffusion models that is robust to geometric attacks is similar and will not be described in detail here.

[0068] In some embodiments, the training method for the above neural network model is as follows: In the first stage, all parameters of the diffusion model, the bi-branch adaptive intensifier, and the image corrector are frozen, and only the watermark encoder and the watermark decoder are trained.

[0069] In the second stage, firstly, all parameters of the watermark decoder and the image corrector are frozen, the watermark encoder and the bi-branch adaptive strength generator are activated, and the bi-branch adaptive strength generator is trained. Secondly, all parameters of the bi-branch adaptive strength generator are frozen, the image corrector is activated, and the image corrector is trained.

[0070] In the third stage, the bi-branch adaptive intensifier, the image corrector, the watermark extractor, and the watermark decoder are activated, the diffusion model is frozen, and the neural network model is trained. The total loss function is... for:

[0071] in, The message cross-entropy loss is used to ensure the readability of the watermark under attack. Image perceptual loss constrains image visual quality; Geometric alignment constraints are used to assist the convergence of the image corrector; The weights are the message cross-entropy loss values; The weights for image perception loss; The weights of the geometric alignment constraints.

[0072] In this embodiment, the first stage is the payload channel construction stage. The goal of this stage is to establish a basic mapping relationship between the latent spatial tensor and binary information, and it does not involve the generation process of the diffusion model or the embedding and superposition of watermarks.

[0073] The specific operation is as follows: Freeze all parameters of the diffusion model, the two-branch adaptive intensifier, and the image corrector, and train only the watermark encoder and watermark decoder. During training, a random binary watermark sequence is first generated and input into the watermark encoder, which maps it to a watermark feature tensor with the same dimension as the latent vector of the diffusion model. Subsequently, this watermark feature tensor is directly input into the watermark decoder (i.e., skipping the step of embedding into the latent vector), and the decoder restores it to the predicted binary sequence. By calculating the reconstruction loss (mean squared error loss and cross-entropy loss) between the input sequence and the predicted sequence, the weights of the watermark encoder and decoder are updated using the backpropagation algorithm until the watermark reading and writing channels have stable reconstruction capabilities.

[0074] The second stage is the functional module adaptation stage. After the first stage of training is completed, this stage is specifically divided into two sequential sub-steps. First, the dual-branch adaptive intensifier is trained: the watermark encoder and the dual-branch adaptive intensifier are activated, while the watermark decoder and the image corrector are frozen. The feature tensor generated by the watermark encoder is input into the intensifier to calculate the adaptive intensity map, which is then superimposed on the initial latent vector generated by the diffusion model. Subsequently, the watermarked image is generated through the VAE decoder, and the feature difference (MSE Loss) in the latent space and the perceptual difference (LPIPS Loss) in the pixel space between the watermarked image and the original image are calculated. This updates the intensifier parameters, enabling it to learn to allocate the watermark intensity using local and global branches to ensure invisibility.

[0075] The STN image corrector is then trained: keeping the intensifier parameters fixed, the image corrector is activated. A geometric attack simulation layer is introduced into the training data stream, applying random rotation, scaling, and cropping transformations to the generated watermarked image. The attacked image is input into the corrector for processing, resulting in a corrected image. Then, the pixel-level geometric alignment loss (mean squared error loss, MSE) between the corrected image and the original image is calculated. The gradient of this alignment loss is backpropagated to the parameter regression subnetwork of the corrector, forcing the corrector to quickly learn the geometric transformation parameters to restore the deformed image to its standard shape under supervised conditions.

[0076] The third stage is the global joint fine-tuning stage. This stage aims to achieve collaborative optimization of all modules. Specifically, it involves unfreezing the dual-branch adaptive intensifier, image corrector, watermark extractor, and watermark decoder (while keeping the main diffusion model frozen), simulating the entire "generation-embedding-attack-correction-extraction" process. A weighted total loss function is used. :

[0077] in, The message cross-entropy loss is used to ensure the readability of the watermark under attack. Image perceptual loss constrains image visual quality; Geometric alignment constraints are used to assist the convergence of the image corrector; The weights are the message cross-entropy loss values; The weights for image perception loss; The weights of the geometric alignment constraints.

[0078] To verify the effectiveness of the proposed adaptive watermarking method for geometrically robust diffusion models, this embodiment performs robustness tests on a subset of the MS-COCO dataset. The GPU used is an RTX 4090 (24GB). 1. Python version 3.12.3 and CUDA version 12.3 were used. The AdamW optimizer was used during the training phase. The learning rate was set to 1e-4 in stages 1 and 2, and 1e-5 in stage 3. The exponential moving average threshold for bit precision at each stage was 0.990, 0.045, and 0.990. The batch size was 4. The latent diffusion model used in the experiments was stable-diffusion-2-1-base. The DDIM sampler was used during the generation process, with 50 sampling steps.

[0079] The test selected five common attack types: rotation, center crop, resize, JPEG compression, and Gaussian blur.

[0080] Specifically, to quantitatively evaluate the model's performance under different attack intensities, this embodiment sets minimum and maximum values ​​for attack parameters for each of the five attack types, and evenly divides them into five levels (Level 1 to Level 5) within this range, where Level 1 represents the weakest attack intensity and Level 5 represents the strongest attack intensity. The specific parameter settings for each attack type are as follows: Rotation Attack: The rotation angle range is set from ±5 degrees to ±20 degrees. That is, Level 1 corresponds to a rotation of ±5 degrees, Level 5 corresponds to a rotation of ±20 degrees, and intermediate levels are distributed by linear interpolation.

[0081] Center Crop Attack: The cropping retention ratio is set between 90% and 50%. That is, Level 1 corresponds to retaining 90% of the image area, and Level 5 corresponds to retaining only 50% of the image area. The attack strength increases as the retained area decreases.

[0082] Resize attack: The image scaling ratio is set to 0.9x to 0.5x. That is, Level 1 corresponds to scaling to 0.9x of the original image, and Level 5 corresponds to scaling to 0.5x of the original image.

[0083] JPEG compression attack: The compression quality factor is set to a range of 90 to 10. That is, Level 1 corresponds to a quality factor of 90, Level 5 corresponds to a quality factor of 10, and the lower the quality factor, the greater the attack strength caused by compression.

[0084] Gaussian Blur Attack: The standard deviation of the Gaussian kernel is set to range from 0.5 to 3.0. That is, Level 1 corresponds to Sigma=0.5, and Level 5 corresponds to Sigma=3.0.

[0085] Figure 3 The diagram shows the robustness of the neural network model proposed in this embodiment of the invention to various types of noise. Figure 3 As can be seen, the accuracy of watermark extraction bits of the present invention remains at a high level under the different levels of attacks mentioned above. In particular, when facing strong geometric deformations, thanks to the introduction of the STN image corrector, the accuracy drops by very little, demonstrating the excellent robustness of the present invention in geometric attack scenarios.

[0086] Compared with the prior art, the embodiments of the present invention can achieve at least one of the following beneficial effects: 1. Significantly improves robustness against geometric transformation attacks: The differentiable STN image corrector introduced in this invention achieves automatic spatial alignment of the attacked image, effectively restoring the spatial synchronization of the watermark signal, so that the model can still maintain extremely high watermark extraction accuracy when subjected to arbitrary angle rotation or large-scale scaling.

[0087] 2. Achieving an adaptive balance between image fidelity and watermark embedding depth: Based on an improved IAT architecture, a dual-branch adaptive strength generator combines the pixel-level texture perception capability of the local branch with the overall semantic control capability of the global branch. This ensures that the watermark is mainly embedded in high-frequency texture areas, while effectively avoiding smooth backgrounds, thus resolving the contradiction between invisibility and robustness in generative watermarking.

[0088] 3. Excellent training stability and modular compatibility: The phased, progressive training strategy effectively solves the convergence problem of complex multi-module networks. Furthermore, since the core operations are performed in the latent space, there is no need to retrain the main diffusion model, resulting in advantages such as low cost and easy adaptation.

[0089] The above description is only a preferred embodiment of the present invention, but the scope of protection of the present invention is not limited thereto. Any changes or substitutions that can be easily conceived by those skilled in the art within the scope of the technology disclosed in the present invention should be included within the scope of protection of the present invention.

Claims

1. An adaptive watermarking method for a diffusion model robust to geometric attacks, characterized in that, In the watermark embedding stage, the method includes: Obtain the initial latent vector generated by the diffusion model based on the prompts input by the user; The watermark is encoded by a watermark encoder to obtain a watermark feature vector that has the same vector space size as the initial potential vector. The embedding strength of the initial latent vector is calculated using a dual-branch adaptive intensifier based on an Illumination-Adaptive-Transformer network structure, resulting in a spatial adaptive embedding strength map of the initial latent vector. Based on the spatial adaptive embedding intensity map, the watermark feature vector is embedded into the initial latent vector to obtain a watermarked latent vector; and Using the diffusion model, the watermarked latent vector is converted into a watermarked image.

2. The method according to claim 1, characterized in that, In the watermark extraction stage, the method includes: The watermarked image to be detected is input into an image corrector based on a spatial transformation network to obtain the corrected image of the image to be detected. The corrected image is input into the diffusion model to obtain the latent vector of the corrected image; The latent vector of the corrected image is input into the watermark extractor to obtain the target watermark feature vector, wherein the network structure of the watermark extractor is the same as that of the dual-branch adaptive intensifier. The target watermark feature vector is input into the watermark decoder to obtain the watermark, wherein the watermark decoder adopts an inverse network structure that is symmetrical to the network structure of the watermark encoder.

3. The method according to claim 2, characterized in that, The embedding strength of the initial latent vector is calculated using a dual-branch adaptive intensifier based on an Illumination-Adaptive-Transformer network structure, resulting in a spatial adaptive embedding strength map of the initial latent vector, including: The initial latent vector is input into a local branch of the dual-branch adaptive intensifier to obtain a high-frequency feature map that reflects the local texture complexity. The initial latent vector is input into the global branch of the dual-branch adaptive intensifier to obtain a global modulation vector reflecting the overall semantic sensitivity of the image; and The high-frequency feature map is modulated by affine transformation using the global modulation vector, and the modulated high-frequency feature map is input into a 1×1 convolutional layer and a sigmoid activation function to obtain a spatial adaptive embedding intensity map.

4. The method according to claim 3, characterized in that, Based on the spatial adaptive embedding intensity map, the watermark feature vector is embedded into the initial latent vector to obtain the watermarked latent vector, including: The watermarked potential vector is generated according to the following formula. : in, Let be the initial potential vector. This is the spatial adaptive embedding intensity map. The watermark feature vector is... This is the preset global intensity control factor. This is for element-wise multiplication.

5. The method according to claim 2, characterized in that, The image corrector includes: a texture feature enhancement module, a parameter regression subnetwork, and a differentiable resampling layer, wherein: The texture feature enhancement module is used to extract the high-frequency feature map of the image to be detected and fuse the high-frequency feature map with the image to be detected; The parameter regression subnetwork is used to receive the output of the texture feature enhancement module and regress the affine transformation matrix used to counteract geometric attacks through a convolutional neural network; and The differentiable resampling layer is used to establish a coordinate mapping relationship based on the affine transformation matrix, and to resample the image to be detected to obtain the corrected image of the image to be detected.

6. The method according to claim 5, characterized in that, The texture feature enhancement module includes a fixed Sobel operator layer, and the parametric regression subnetwork is a convolutional neural network structure based on ResNet-18, wherein the last layer of the parametric regression subnetwork is a fully connected layer for outputting a 2×3 affine transformation matrix.

7. The method according to claim 6, characterized in that, The differentiable resampling layer resamples the image to be detected according to the following formula: in, and To correct the pixel value of the i-th pixel in the image, and Let be the pixel value of the i-th pixel in the image to be detected. Let be the affine transformation matrix.

8. The method according to claim 3, characterized in that, The latent vector of the corrected image is input into the watermark extractor to obtain the target watermark feature vector, including: The target watermark feature vector is generated according to the following formula. : in, The latent vector of the corrected image, The spatially adaptive embedding intensity map is the latent vector of the corrected image.

9. A neural network model for an adaptive watermarking method, characterized in that, include: The system consists of a diffusion model, a watermark encoder, a bi-branch adaptive strength generator, a feature fusion layer, an image corrector, a watermark extractor, and a watermark decoder, among which: The diffusion model includes a variational autoencoder and a variational autodecoder, wherein the variational autoencoder is used to generate an initial latent vector based on prompt words input by the user; The watermark encoder is used to encode the watermark to obtain a watermark feature vector that has the same vector space size as the initial potential vector. The dual-branch adaptive intensifier is used to calculate the embedding strength of the initial latent vector to obtain the spatial adaptive embedding strength map of the initial latent vector. The feature fusion layer is used to embed the watermark feature vector into the initial latent vector according to the spatial adaptive embedding intensity map, so as to obtain a watermarked latent vector; The variational autodecoder is used to convert the watermarked latent vector into a watermarked image. The image corrector employs a spatial transformation network, which is used to correct the image to be detected to obtain a corrected image; The variational autoencoder is further configured to generate a latent vector of the corrected image based on the corrected image; The watermark extractor has the same mesh structure as the dual-branch adaptive intensifier, and it is used to convert the latent vector of the corrected image into the target watermark feature vector. The watermark decoder employs a reverse network structure that is symmetrical to the network structure of the watermark encoder, and is used to convert the target watermark feature vector into a watermark.

10. The neural network model according to claim 9, characterized in that, The training method is as follows: In the first stage, all parameters of the diffusion model, the bi-branch adaptive intensifier, and the image corrector are frozen, and only the watermark encoder and the watermark decoder are trained. In the second stage, firstly, all parameters of the watermark decoder and the image corrector are frozen, the watermark encoder and the dual-branch adaptive strength generator are activated, and the dual-branch adaptive strength generator is trained. Secondly, all parameters of the dual-branch adaptive strength generator are frozen, the image corrector is activated, and the image corrector is trained. In the third stage, the bi-branch adaptive intensifier, the image corrector, the watermark extractor, and the watermark decoder are activated, the diffusion model is frozen, and the neural network model is trained. The total loss function is... for: in, The message cross-entropy loss is used to ensure the readability of the watermark under attack. Image perceptual loss constrains image visual quality; Geometric alignment constraints are used to assist the convergence of the image corrector; The weights are the message cross-entropy loss values; The weights for image perception loss; The weights of the geometric alignment constraints.

Images