Image generation method and system based on diffusion model adaptive reasoning

Abstract

The invention discloses an image generation method and system based on diffusion model adaptive reasoning. The method comprises the following steps: constructing a next noise intensity prediction module TPM; improving an original diffusion model, and adding a TPM; the improved diffusion model is trained by using a reinforcement learning method, the whole denoising process is expanded for multiple times, the next time step is predicted in each iteration, an award is calculated according to the quality of a finally generated image and the number of steps, and TPM parameters are updated by using a PPO loss function; in the actual reasoning process, noise is randomly sampled from Gaussian noise and input into the trained diffusion model, the diffusion model is used for predicting a denoising result, meanwhile, the TPM is used for outputting distribution of time steps based on next denoising, and a distribution mode is used as the next time step used in actual reasoning; and when the predicted time step is smaller than a set threshold value, a final image generation result is obtained based on model output. According to the method, the image generation process can be dynamically adjusted, and a reasoning result with both efficiency and quality is realized.

Description

Technical Field

[0001] The present invention belongs to the field of artificial intelligence and computer vision, and in particular relates to an image generation method and system based on diffusion model adaptive reasoning. Background Art

[0002] Since the advent of ChatGPT, deep neural networks, especially models based on the Transformer architecture, have become increasingly influential in the field of content generation. This architecture has not only promoted the leap forward in natural language processing technology, but has also demonstrated outstanding performance in multimodal generation tasks such as images, audio, and video, becoming an important cornerstone in the field of artificial intelligence content generation (AIGC).

[0003] Among them, the diffusion model based on the Transformer architecture has also achieved remarkable results in the field of image and video generation in recent years. Its powerful modeling ability and flexible generation mechanism have enabled the quality and diversity of generated content to reach an unprecedented level.

[0004] For example, a Chinese patent document with publication number CN117593398A discloses a remote sensing image generation method based on a diffusion model. First, a forward diffusion process is performed to generate a training data set. Secondly, a reverse denoising process is performed to continuously denoise until the original image is obtained. Then, the posterior probability is calculated to obtain the probability distribution of the denoised image given the noisy image and the original image. Then, a deep learning model is built to use an improved UNet model to predict the mean and variance in the reverse denoising process. Finally, image generation is performed to randomly generate a noise image, and a remote sensing image is generated after cyclic denoising.

[0005] A Chinese patent document with publication number CN115908187A discloses an image feature analysis and generation method based on a fast denoising diffusion probability model. An original denoising diffusion probability model is constructed, which includes a score predictor, and a condition generator and a corresponding discriminator are introduced to obtain a fast denoising diffusion probability model. The trained score predictor and the trained condition generator are updated to the fast denoising diffusion probability model, and a Gaussian white noise image is input into the fast denoising diffusion probability model to obtain a true high-quality output image.

[0006] The diffusion process of the diffusion model is to denoise the random noise step by step to generate realistic samples. However, the current diffusion model often relies on a certain denoising plan to plan the noise intensity of each denoising step during the reasoning process of different prompt words. Generally speaking, when generating images, a series of time steps (i.e., noise intensities) {t0,…,t N}, at the beginning, t0 = 1, it is completely noise Final N= 0 to get a clean image The denoising process at each step is:

[0007] However, it is not reasonable to use a unified denoising plan for different prompt words. For less complex generation targets, fewer denoising steps can achieve high-quality results; while for more complex generation targets, more denoising steps are required to ensure the fineness and accuracy of the generation results. In addition, the denoising intensity of each step should be dynamically adjusted according to the current denoising state, rather than pre-fixing the denoising plan for the entire reasoning process. Summary of the invention

[0008] The present invention provides an image generation method and system based on diffusion model adaptive reasoning, which can dynamically adjust the image generation process to achieve reasoning results that emphasize both efficiency and quality; for images with rich content, denoising steps are automatically added to capture more details, and for images with simple content, denoising steps are automatically reduced to improve efficiency.

[0009] An image generation method based on diffusion model adaptive reasoning includes the following steps:

[0010] (1) Construct the next step noise intensity prediction module TPM to dynamically adjust the number of denoising steps and noise intensity according to the current potential features;

[0011] (2) Improve the original diffusion model and add a prediction module TPM;

[0012] (3) Obtain a text training set and use the reinforcement learning method to train the improved diffusion model; during the training process, the entire denoising process is carried out multiple times, and the next time step is predicted in each iteration. The reward is calculated based on the quality of the final generated image and the number of steps, and the PPO loss function is used to update the parameters of the prediction module TPM until the entire model converges;

[0013] (4) In the actual inference process, noise is randomly sampled from Gaussian noise and input into the trained diffusion model. The initial time step is t n＝0 =1;

[0014] (5) At time step t n When the diffusion model is used to predict the denoising result, TPM is used to output the time step t based on which the next denoising is based. n+1 The distribution belongs to it, and the distribution mode is used as the next time step used in actual inference;

[0015] (6) Repeat step (5) until the predicted time step t n+1 Less than the set threshold, the final generated image result is obtained based on the model output.

[0016] In step (1), the prediction module TPM includes an input layer, a convolution layer, a pooling layer and a linear layer; wherein the input of the input layer is the potential feature of the current denoising step, and after being processed by the convolution layer, the pooling layer and the linear layer, the parameters of the Beta distribution are output to predict the next time step, i.e., the noise intensity of the next step.

[0017] The latent features of the current denoising step are derived from the intermediate layers of the diffusion model, which contain the noise information of the current image and the intermediate results in the denoising process.

[0018] The prediction module TPM predicts two real numbers a and b, and determines the corresponding Beta distribution through the following formula:

[0019] α=1+e a

[0020] β=1+e b

[0021] Where α and β are the two parameters of Beta distribution.

[0022] The formula for predicting the next time step is:

[0023] t n =r n ·t n-1

[0024] Where, t n represents the time step t n-1 The next time step; r n Represents a value sampled from a Beta distribution.

[0025] In step (3), during the training process, the original diffusion model is frozen and only the parameters of the added prediction module TPM are updated.

[0026] In step (3), when calculating the reward based on the quality of the final generated image and the number of steps, a discount factor is applied for discount. The specific formula is as follows:

[0027]

[0028] In the formula, R(s,y) represents the reward for this denoising, γ represents the discount factor, N represents the total number of denoising steps, and IR(y,s) represents the reward of the reward model.

[0029] In step (3), the PPO loss function is as follows:

[0030]

[0031] Where s = (c, ∈) represents the initial state, including the input prompt c and the standard normal distribution noise πθ represents the policy network, i.e. the prediction module TPM; π θ (y|s) represents the current strategy for sampling trajectories, π θ (·|s) represents the action distribution of the current strategy in state s. ;

[0032] π old (y|s) represents the old strategy used to sample the trajectory; π ref represents the reference strategy for regularization, π ref (·|s) is the action distribution of the reference strategy in state s; y represents the final image obtained by the strategy, that is, the scheduling time; represents the advantage of taking action y in state s; λ represents the coefficient of KL; KL represents the KL divergence; θ represents the parameters of the current policy network.

[0033] Based on the same inventive principle, the present invention also provides an image generation system based on diffusion model adaptive reasoning, including a memory and one or more processors, wherein the memory stores executable code, and when the one or more processors execute the executable code, they are used to implement the above-mentioned image generation method.

[0034] Compared with the prior art, the present invention has the following beneficial effects:

[0035] 1. The present invention proposes an adaptive reasoning mechanism that can dynamically optimize the reasoning process according to different text guides. By designing a plug-and-play prediction module TPM, it can analyze the current denoising potential features in real time and predict the noise intensity of the next reasoning, thereby realizing dynamic adjustment of the entire reasoning process. For images with rich content, the denoising steps are automatically increased to capture more details; for images with simple content, the denoising steps are automatically reduced to improve efficiency.

[0036] 2. Different from the traditional diffusion model single-step denoising training method, the present invention uses reinforcement learning method to supervise the denoising process throughout the whole process to ensure the high consistency of training and reasoning process. In addition, unlike model distillation or other methods to optimize the reasoning process, the present invention only needs very few training resources to achieve the expected goal, which greatly improves the training efficiency and generation quality of the model.

[0037] 3. The adaptive reasoning mechanism adopted in the present invention not only optimizes the generation effect, but also significantly reduces the training and reasoning costs, providing a new idea for the wide application of diffusion models. At the same time, the TPM module can be easily integrated into the existing diffusion model without major modifications to the original model. BRIEF DESCRIPTION OF THE DRAWINGS

[0038] Figure 1 This is an overall structural diagram of adding a TPM module to a diffusion model in an embodiment of the present invention;

[0039] Figure 2 Generate an effect diagram for the picture in the embodiment of the present invention. DETAILED DESCRIPTION

[0040] The present invention is further described in detail below in conjunction with the accompanying drawings and embodiments. It should be pointed out that the embodiments described below are intended to facilitate the understanding of the present invention and do not have any limiting effect on the present invention.

[0041] An image generation method based on adaptive reasoning of diffusion model mainly includes two core parts:

[0042] (1) The present invention designs a lightweight module (TPM) consisting of only a few layers of neural networks, which is used to predict the noise intensity of the next step in real time during the denoising process. The module can dynamically adjust the number of denoising steps and the noise intensity according to the current latent features, thereby significantly improving the efficiency of the model while ensuring image quality. Specifically, the module processes the input latent features through convolutional layers and linear layers, and outputs a Beta distribution parameter for predicting the noise intensity of the next step.

[0043] (2) In order to optimize the prediction module TPM, the present invention proposes a training algorithm based on reinforcement learning. Specifically, the entire denoising process is regarded as a complete trajectory, and a reward model aligned with human preferences is used to evaluate the quality of the generated image. By maximizing the quality of the generated image and discounting the number of denoising steps, the present invention trains the prediction module to achieve accurate time prediction. The core of the algorithm is to regard the entire generation process as a single action and optimize the parameters of the prediction module through the policy gradient method, thereby reducing the denoising steps by about 50% on average while generating high-quality images.

[0044] Specifically, the embodiment of the present invention includes the following steps:

[0045] Step 1: Extract the text annotation part from the image-text dataset (such as LAION-400M) as the training set, and clean and preprocess the dataset to ensure data quality and diversity.

[0046] Step 2: Design and implement the next step noise intensity prediction module (TPM), which consists of a lightweight neural network that can dynamically adjust the denoising step and noise intensity according to the current potential features. Integrate the TPM module into the diffusion model that needs to be tuned to ensure that it can be seamlessly integrated and work together.

[0047] The next noise intensity prediction module (TPM) is an important part of the present invention. Its main function is to dynamically predict the noise intensity of the next step during the denoising process. This module uses a lightweight neural network architecture to flexibly adjust the denoising step and noise intensity according to the current potential features, which not only ensures the quality of the generated image, but also significantly improves the model operation efficiency.

[0048] The overall structure of adding TPM module in the diffusion model is as follows Figure 1 As shown, the TPM module includes an input layer, a convolutional layer, a pooling layer and a linear layer.

[0049] Input layer: The input of TPM is the potential features of the current denoising step. These features usually come from the intermediate layers of the diffusion model, which contain the noise information of the current image and the intermediate results of the denoising process.

[0050] Convolutional layer: In order to efficiently process the potential features of the input, TPM uses a multi-layer convolutional network. The convolutional layer can extract the spatial information in the input features and convert it into a feature representation suitable for further processing.

[0051] Pooling layer: After the convolution layer, the pooling layer is used to reduce the dimension of the features. On the one hand, it reduces the computational complexity, and on the other hand, it retains key information and avoids unnecessary data redundancy.

[0052] Linear layer: The pooled feature vector is fed into two linear layers, outputting a and b.

[0053] The purpose of letting the model predict Beta distribution in the present invention is to obtain the probability distribution of each action in reinforcement learning training, so as to optimize the entire action trajectory. It is worth noting that when α>1 and β>1, Beta distribution presents a unimodal characteristic. Based on this characteristic, the present invention reparameterizes TPM to make it predict two real numbers a and b, and determines the corresponding Beta distribution through the following formula.

[0054] α=1+e a ,β=1+e b

[0055] r~Beta(α,β)

[0056] To enable the module to dynamically adjust the noise schedule, the TPM predicts how the noise intensity t decays in adjacent steps. Therefore, the next noise level t n r can be obtained by sampling from the distribution n , and calculated using the following formula.

[0057] t n =r n ·t n-1

[0058] During training, the original diffusion model is frozen and only the newly introduced TPM is updated. Therefore, the model learns to predict the next time while retaining the original ability of image generation.

[0059] Step 3: Select a pre-trained reward model (such as ImageReward) to evaluate the quality of the generated image. Use the reinforcement learning method to train the improved diffusion model; during the training process, the entire denoising process is carried out multiple times, and the next time step is predicted in each iteration. The reward is calculated based on the quality of the final generated image and the number of steps, and the PPO loss function is used to update the parameters of the policy network until the model converges.

[0060] To train the TPM, the present invention uses a reinforcement learning method to directly optimize the final image quality after the entire denoising process. This method ensures that the TPM learns to predict the optimal denoising process to maximize the quality of the generated image and minimize the number of steps required. Since the present invention uses a reward model with reinforcement learning to optimize the entire process, there is no need for image-text pairs, only a pre-trained reward model and raw image prompt words are needed for optimization.

[0061] During training, the invention supervises the entire denoising process, predicting the next time step at each iteration. The quality of the final generated image is measured by a reward model aligned with human preferences. This reward is discounted according to the number of steps required to generate the image, encouraging the model to find a balance between image quality and efficiency.

[0062] The present invention uses proximal policy optimization (PPO) as the optimization algorithm. The loss function of PPO is defined as:

[0063]

[0064] In the problem definition of the present invention, s = (c, ∈) represents the initial state, including the input prompt c and the standard normal distribution noise π θ Represents a policy network, such as the TPM of the present invention. old represents the old strategy used to sample trajectories; π ref represents the reference strategy used for regularization; y represents the action taken by the strategy of the present invention, that is, the scheduling time; represents the advantage of taking action y in state s.

[0065] The present invention treats the entire denoising plan as the target to be optimized. Typically, when a model makes a series of predictions, PPO treats each prediction as an action and optimizes them as a batch. The RLOO algorithm proposes that when the reward signal only appears at the end of the sequence and the environment dynamics are fully determined, the entire sequence can be treated as a single action without compromising performance. Therefore, for simplicity, the present invention treats the entire generation process (including all prediction times in the schedule) as a single optimization action. Therefore, when calculating the PPO loss, the present invention optimizes the entire trajectory as a training sample. However, TPM only outputs the distribution of a single prediction, and the present invention uses It means that according to the chain rule, the probability of the entire generation can be calculated as the product of each prediction, as shown below:

[0066]

[0067] Wherein, N represents the total number of generation steps. In the model of the present invention, the probability strategy network Each factor of is calculated by TPM, and TPM outputs the possible diffusion time t for the next step i The distribution on .

[0068] In the present invention, the image reward decays by the number of steps. When generating samples for policy gradient training, the trajectory is obtained by generating images from Gaussian noise using the current TPM policy. The present invention selects the pre-trained ImageReward model to allocate reward scores only based on the last final image. Considering that too many steps are not desired in the generation process, the present invention applies a discount factor γ<1 to discount the reward to the intermediate diffusion step and calculates the average of the total number of denoising steps N. The final reward of the trajectory is shown as follows:

[0069]

[0070] The present invention directly calculates the advantage from R(s,y) of the training batch Without optimizing the value model. This reward function will encourage TPM to help the original diffusion model generate high-quality images and reduce the number of diffusion steps to improve model efficiency. By adjusting the decay factor γ, you can control how many diffusion steps are applied to generate high-quality images. The smaller the value of γ, the fewer the average number of inference denoising times for the model that finally converges with training.

[0071] Step 5: In the inference model, the mode of the TPM output distribution is used as the module output value to avoid the sampling uncertainty of reinforcement learning exploration during training, thereby stabilizing the inference process and achieving a balance between efficiency and quality.

[0072] The core advantage of the present invention is that it can dynamically adjust the number of denoising steps and noise intensity according to the current potential features, thereby significantly improving the efficiency of the model while ensuring image quality. Specifically, the method of the present invention can reduce the denoising steps by about 50% on average while generating high-quality images, thus achieving a good balance between performance and efficiency.

[0073] In addition, the method of the present invention can be easily integrated into the existing diffusion model without major modifications to the original model, and only a lightweight prediction module needs to be added. This makes the present invention highly compatible and scalable, and can be widely used in various generation tasks based on diffusion models, such as image synthesis, video generation, etc.

[0074] Another important feature of the present invention is that it can adaptively adjust the denoising plan according to different input prompts. For example, for images with rich content, the model will automatically increase the denoising steps to capture more details; while for images with simple content, the model will reduce the denoising steps to improve efficiency. This adaptability not only improves the quality of the generated images, but also significantly reduces the computational cost during the inference process.

[0075] like Figure 2 As shown, from left to right, the images generated by the FLUX-1.0dev model embedded in the TPM of the present invention gradually change from simple to complex. The TPM of the present invention can adaptively adjust the reasoning plan to match the complexity of each generated target.

[0076] The embodiments described above provide a detailed description of the technical solutions and beneficial effects of the present invention. It should be understood that the above are only specific embodiments of the present invention and are not intended to limit the present invention. Any modifications, supplements and equivalent substitutions made within the scope of the principles of the present invention should be included in the protection scope of the present invention.

Claims

1. An image generation method based on adaptive reasoning of a diffusion model, characterized in that: The following steps are involved: (1) Construct the next step noise intensity prediction module TPM to dynamically adjust the number of denoising steps and noise intensity according to the current potential features; (2) Improve the original diffusion model and add a prediction module TPM; (3) Obtain text training sets and use reinforcement learning methods to train the improved diffusion model; During the training process, the entire denoising process is carried out multiple times, and the next time step is predicted in each iteration. The reward is calculated based on the quality of the final generated image and the number of steps, and the parameters of the prediction module TPM are updated using the PPO loss function until the entire model converges; (4) In the actual inference process, noise is randomly sampled from Gaussian noise and input into the trained diffusion model. The initial time step is t n＝0 =1; (5) At time step t n When the diffusion model is used to predict the denoising result, the TPM is used to output the time step t based on which the next denoising is based. n+1 The distribution belongs to it, and the distribution mode is used as the next time step used in actual inference; (6) Repeat step (5) until the predicted time step t n+1 Less than the set threshold, the final generated image result is obtained based on the model output.

2. The image generation method based on diffusion model adaptive reasoning according to claim 1 is characterized in that: In step (1), the prediction module TPM includes an input layer, a convolution layer, a pooling layer and a linear layer; wherein the input of the input layer is the potential feature of the current denoising step, and after being processed by the convolution layer, the pooling layer and the linear layer, the parameters of the Beta distribution are output to predict the next time step, i.e., the noise intensity of the next step.

3. The image generation method based on diffusion model adaptive reasoning according to claim 2 is characterized in that: The latent features of the current denoising step are derived from the intermediate layers of the diffusion model, which contain the noise information of the current image and the intermediate results in the denoising process.

4. The image generation method based on diffusion model adaptive reasoning according to claim 2 is characterized in that: The prediction module TPM predicts two real numbers a and b, and determines the corresponding Beta distribution through the following formula: α＝1+e a β＝1+e b Where α and β are the two parameters of Beta distribution.

5. The image generation method based on diffusion model adaptive reasoning according to claim 2 is characterized in that: The formula for predicting the next time step is: t n =r n ·t n-1 Where, t n represents the time step t n-1 The next time step; r n Represents a value sampled from a Beta distribution.

6. The image generation method based on diffusion model adaptive reasoning according to claim 1, characterized in that: In step (3), during the training process, the original diffusion model is frozen and only the parameters of the added prediction module TPM are updated.

7. The image generation method based on diffusion model adaptive reasoning according to claim 1, characterized in that: In step (3), when calculating the reward based on the quality of the final generated image and the number of steps, a discount factor is applied for discount. The specific formula is as follows: Where R(s, y) represents the reward for this denoising, γ represents the discount factor, N represents the total number of denoising steps, and IR(y, ,s) represents the reward of the reward model.

8. The image generation method based on diffusion model adaptive reasoning according to claim 1, characterized in that: In step (3), the PPO loss function is as follows: Where s = (c, ∈) represents the initial state, including the input prompt c and the standard normal distribution noise π θ represents the policy network, i.e. the prediction module TPM; π θ (y|s) represents the current strategy for sampling trajectories, π θ (·|s) represents the action distribution of the current strategy in state s; π old (y|s) represents the old strategy used to sample the trajectory; π ref (·|s) is the action distribution of the reference strategy in state s; y is the final image obtained by the strategy; represents the advantage of taking action y in state s; λ represents the coefficient of KL; KL represents the KL divergence; θ represents the parameters of the current policy network.

9. An image generation system based on diffusion model adaptive reasoning, characterized in that: The invention comprises a memory and one or more processors, wherein the memory stores executable codes, and when the one or more processors execute the executable codes, the image generating method according to any one of claims 1 to 8 is implemented.

Images