Training method of image generation model, image generation method and related equipment

Abstract

The invention relates to a training method of an image generation model, an image generation method and related equipment. The method comprises the following steps: acquiring noise image sample tokens of a plurality of noise sample images and condition sample feature information of corresponding condition sample information; the plurality of noise sample images are sample images with a plurality of different image object categories and / or different noise intensities; obtaining a sample token pool corresponding to each noise image sample token in the same training batch; calling an image generation model realized based on a diffusion model, and carrying out denoising processing on the sample token pool and the condition sample feature information to obtain predicted image information; in the de-noising processing process, based on a preset expert capacity, determining target sample tokens which need to be processed by each expert sub-network in the hybrid expert network in the image generation model from a sample token pool; and performing model training based on the difference between the predicted image information and the target image information. The performance of the trained model can be improved.

Description

Technical Field

[0001] The present disclosure relates to the field of computer technology, and in particular, to a training method for an image generation model, an image generation method, and related devices. Background Art

[0002] In recent years, Diffusion Models (DM) are one of the important models in the field of visual generation. It can perform tasks such as image generation based on text instructions and can generate realistic visual effects. Mixture of Experts (MoE) is a machine learning model architecture designed to improve the performance and flexibility of the overall model by combining multiple "expert" models (sub-models).

[0003] In related technologies, if the Mixture of Experts model is directly applied to the Diffusion Model to perform visual generation tasks, the performance of the Diffusion Model will be damaged to a certain extent. Summary of the Invention

[0004] The present disclosure provides a training method for an image generation model, an image generation method, and related devices to solve at least one problem in related technologies. The technical solution of the present disclosure is as follows:

[0005] According to an aspect of an embodiment of the present disclosure, a training method for an image generation model is provided, including:

[0006] Obtain training data, where the training data includes a plurality of noise sample images and corresponding conditional sample information; the plurality of noise sample images are sample images with various different image object categories and / or different noise intensities;

[0007] Obtain noise image sample tokens corresponding to each of the noise sample images, and conditional sample feature information corresponding to the conditional sample information;

[0008] Call a token processing module in the image generation model to flatten each of the noise image sample tokens in the same training batch to obtain a sample token pool at the training batch level; the image generation model is implemented based on a diffusion model;

[0009] Call the image generation model to perform denoising processing on the sample token pool and the conditional sample feature information to obtain predicted image information; during the denoising processing, based on a preset expert capacity, determine target sample tokens required for each expert sub-network in the mixture of experts network in the image generation model from the sample token pool;

[0010] Train the image generation model based on the difference between the predicted image information and the target image information to obtain a trained image generation model.

[0011] In an alternative embodiment, after obtaining the sample token pool at the training batch level, the method further includes:

[0012] Determine the predicted expert capacity of each expert sub-network at each denoising time step through a capacity prediction network; the predicted expert capacity is used to dynamically determine the number of tokens that each expert sub-network needs to process respectively;

[0013] Training the image generation model based on the difference between the predicted image information and the target image information to obtain a trained image generation model includes:

[0014] Train the image generation model based on the difference between the predicted image information and the target image information, and the difference between the trained expert capacity and the predicted expert capacity, to obtain a trained image generation model; the mixture-of-experts network in the trained image generation model includes the trained capacity prediction network.

[0015] In an alternative embodiment, training the image generation model based on the difference between the predicted image information and the target image information, and the difference between the trained expert capacity and the predicted expert capacity to obtain a trained image generation model includes:

[0016] Determine first training loss information based on the difference between the predicted image information and the target image information;

[0017] Determine second training loss information based on the difference between the trained expert capacity and the predicted expert capacity;

[0018] Independently train the diffusion model and the capacity prediction network respectively with the goal of minimizing the first training loss information and minimizing the second training loss information to obtain a trained image generation model.

[0019] In an alternative embodiment, determining the predicted expert capacity of each expert sub-network at each denoising time step through a capacity prediction network includes:

[0020] For each expert sub-network at each denoising time step, input each sample token in the sample token pool into the capacity prediction network to output the capacity prediction probability of each sample token;

[0021] For each expert sub-network at each denoising time step, determine the total number of sample tokens with a capacity prediction probability greater than a preset probability threshold, and use the determined total number of the first sample tokens as the predicted expert capacity of each expert sub-network at each denoising time step.

[0022] In an alternative embodiment, the mixture-of-experts network includes a dynamic threshold module; the method further includes:

[0023] During the training process, the dynamic threshold module uses the exponential moving average technique to determine the prediction probability threshold of each expert subnet at each denoising time step; during the inference process, based on the prediction probability threshold determined by the capacity prediction network and the dynamic threshold module, determine the predicted expert capacity of each expert subnet at each denoising time step.

[0024] In an alternative embodiment, the mixture-of-experts network includes a routing subnet; during the denoising process, the following operations are respectively performed on the mixture-of-experts network at each denoising time step:

[0025] Use the routing subnet to determine the sample affinity matrix between each sample token in the sample token pool and each expert subnet in the mixture-of-experts network;

[0026] Based on the sample affinity matrix and the preset expert capacity, determine the target sample tokens required to be processed by each expert subnet from the sample token pool;

[0027] Process the corresponding target sample tokens through each expert subnet respectively, and perform a fusion process on the processing results of each expert, so as to use the expert fusion processing result as the output of the mixture-of-experts network.

[0028] In an alternative embodiment, the determining the target sample tokens required to be processed by each expert subnet from the sample token pool based on the sample affinity matrix and the preset expert capacity includes:

[0029] Based on the preset expert capacity, determine the token quantity threshold required to be processed by each expert subnet;

[0030] Sort in descending order according to the values of the sample affinity matrix corresponding to each sample token, and determine the filtered sample tokens with the sorting serial number less than or equal to the token quantity threshold from the sample token pool as the target sample tokens required to be processed by each expert subnet.

[0031] According to another aspect of the embodiments of the present disclosure, there is provided an image generation method, including:

[0032] Obtain a plurality of noise images and conditional information;

[0033] Obtain the noise image tokens corresponding to each of the noise images, and the conditional feature information corresponding to the conditional information;

[0034] Call the token processing module in the trained image generation model to flatten each of the noise image tokens to obtain an image token pool; the trained image generation model is trained based on any of the above methods;

[0035] Call the trained image generation model to denoise the image token pool and the conditional feature information to generate a target image; during the denoising process, based on a preset expert capacity, determine the target tokens required for each expert sub-network in the mixture-of-experts network in the trained image generation model from the image token pool.

[0036] In an alternative embodiment, the trained image generation model includes a capacity prediction network. After obtaining the image token pool, the method further includes:

[0037] Determine the predicted expert capacity of each expert sub-network at each denoising time step through the capacity prediction network; the predicted expert capacity is used to dynamically determine the number of tokens required for each expert sub-network; during the denoising process, based on the predicted expert capacity, determine the target tokens required for each expert sub-network in the mixture-of-experts network in the image generation model from the image token pool.

[0038] In an alternative embodiment, the determining the predicted expert capacity of each expert sub-network at each denoising time step through the capacity prediction network includes:

[0039] For each expert sub-network at each denoising time step, input each image token in the image token pool into the capacity prediction network to output the capacity prediction probability of each image token;

[0040] For each expert sub-network at each denoising time step, determine the total number of the first tokens of the sample tokens whose capacity prediction probability is greater than a preset probability threshold, or determine the total number of the second tokens whose capacity prediction probability is greater than the prediction probability threshold as the predicted expert capacity of each expert sub-network at each denoising time step;

[0041] Wherein, the prediction probability threshold is determined by a dynamic threshold module in the trained image generation model.

[0042] In an alternative embodiment, the mixture-of-experts network includes a routing sub-network; during the denoising process, the following operations are respectively performed on the mixture-of-experts network at each denoising time step:

[0043] Use the routing sub-network to determine the affinity matrix between each image token in the image token pool and each expert sub-network in the mixture-of-experts network;

[0044] Based on the affinity matrix and the preset expert capacity or the predicted expert capacity, determine the target tokens to be processed by each of the expert sub-networks from the image token pool;

[0045] Process the corresponding target tokens through each of the expert sub-networks respectively, and perform a fusion process on the processing results of each expert, so as to use the expert fusion processing result as the output of the mixture-of-experts network.

[0046] According to another aspect of the embodiments of the present disclosure, there is provided a training device for an image generation model, including:

[0047] A training data acquisition module, configured to execute acquiring training data, where the training data includes a plurality of noise sample images and corresponding conditional sample information; the plurality of noise sample images are sample images with multiple different image object categories and / or different noise intensities;

[0048] An input acquisition module, configured to execute acquiring noise image sample tokens corresponding to each of the noise sample images, and conditional sample feature information corresponding to the conditional sample information;

[0049] A sample token processing module, configured to execute calling a token processing module in the image generation model to flatten each of the noise image sample tokens in the same training batch to obtain a sample token pool at the training batch level; the image generation model is implemented based on a diffusion model;

[0050] A sample denoising module, configured to execute calling the image generation model to perform denoising processing on the sample token pool and the conditional sample feature information to obtain predicted image information; during the denoising process, based on a preset expert capacity, determine the target sample tokens to be processed by each of the expert sub-networks in the mixture-of-experts network in the image generation model from the sample token pool;

[0051] A training module, configured to execute training the image generation model based on the difference between the predicted image information and the target image information to obtain a trained image generation model.

[0052] In an optional implementation manner, the device further includes:

[0053] A capacity prediction module, configured to execute determining the predicted expert capacity of each expert sub-network at each denoising time step through a capacity prediction network; the predicted expert capacity is used to dynamically determine the number of tokens to be processed by each expert sub-network respectively;

[0054] The training module is further configured to execute:

[0055] Based on the difference between the predicted image information and the target image information, as well as the difference between the trained expert capacity and the predicted expert capacity, train the image generation model to obtain a trained image generation model; the hybrid expert network in the trained image generation model includes the trained capacity prediction network.

[0056] In an alternative embodiment, the training module includes:

[0057] A first loss determination unit, configured to execute to determine first training loss information based on the difference between the predicted image information and the target image information;

[0058] A second loss determination unit, configured to execute to determine second training loss information based on the difference between the trained expert capacity and the predicted expert capacity;

[0059] A training unit, configured to execute to independently train the diffusion model and the capacity prediction network respectively with the goal of minimizing the first training loss information and minimizing the second training loss information, and obtain a trained image generation model.

[0060] In an alternative embodiment, the capacity prediction module includes:

[0061] A capacity prediction probability determination unit, configured to execute to input each sample token in the sample token pool to the capacity prediction network for each expert sub-network at each denoising time step, and output the capacity prediction probability of each sample token;

[0062] A first capacity prediction unit, configured to execute to determine the total number of sample tokens whose capacity prediction probability is greater than a preset probability threshold for each expert sub-network at each denoising time step, and use the determined total number of first sample tokens as the predicted expert capacity of each expert sub-network at each denoising time step.

[0063] In an alternative embodiment, the hybrid expert network includes a dynamic threshold module; the apparatus further includes:

[0064] A threshold determination module, configured to execute to determine the prediction probability threshold of each expert sub-network at each denoising time step through the dynamic threshold module using the exponential moving average technique during the training process; during the inference process, determine the predicted expert capacity of each expert sub-network at each denoising time step based on the capacity prediction network and the prediction probability threshold determined by the dynamic threshold module.

[0065] In an alternative embodiment, the hybrid expert network includes a routing sub-network; during the denoising process, the following operations are respectively performed on the hybrid expert network at each denoising time step:

[0066] Using the routing sub-network, determine the sample affinity matrix between each sample token in the sample token pool and each expert sub-network in the mixture-of-experts network;

[0067] Based on the sample affinity matrix and the preset expert capacity, determine the target sample tokens that each of the expert sub-networks needs to process from the sample token pool;

[0068] Have each of the expert sub-networks process the corresponding target sample tokens respectively, and perform a fusion process on the processing results of each expert, so as to use the expert fusion processing result as the output of the mixture-of-experts network.

[0069] In an alternative embodiment, the determining the target sample tokens that each of the expert sub-networks needs to process from the sample token pool based on the sample affinity matrix and the preset expert capacity includes:

[0070] Based on the preset expert capacity or the predicted expert capacity, determine the token quantity threshold that each of the expert sub-networks needs to process;

[0071] Sort in descending order according to the values of the sample affinity matrix corresponding to each sample token, and determine the screened sample tokens whose sorting serial numbers are less than or equal to the token quantity threshold from the sample token pool as the target sample tokens that each of the expert sub-networks needs to process.

[0072] According to another aspect of the embodiments of the present disclosure, there is provided an image generation device, including:

[0073] A first acquisition module configured to acquire a plurality of noise images and conditional information;

[0074] A second acquisition module configured to acquire the noise image tokens corresponding to each of the noise images and the conditional feature information corresponding to the conditional information;

[0075] A token processing module configured to call the token processing module in the trained image generation model to perform a flattening process on each of the noise image tokens to obtain an image token pool; the trained image generation model is trained by any of the above methods;

[0076] An image generation module configured to call the trained image generation model to perform a denoising process on the image token pool and the conditional feature information to generate a target image; during the denoising process, based on the preset expert capacity, determine the target tokens that each of the expert sub-networks in the mixture-of-experts network of the trained image generation model needs to process from the image token pool.

[0077] In an alternative embodiment, the apparatus further includes:

[0078] A capacity prediction module configured to execute, through the capacity prediction network, to determine the predicted expert capacity of each expert sub-network at each denoising time step; the predicted expert capacity is used to dynamically determine the number of tokens that each expert sub-network needs to process respectively; during the denoising process, based on the predicted expert capacity, determine the target tokens that each expert sub-network in the mixture-of-experts network in the image generation model needs to process respectively from the image token pool.

[0079] In an alternative embodiment, the capacity prediction module includes:

[0080] A capacity prediction probability determination unit configured to execute, for each expert sub-network at each denoising time step, input each image token in the image token pool into the capacity prediction network, and output the capacity prediction probability of each image token;

[0081] A capacity prediction unit configured to execute, for each expert sub-network at each denoising time step, determine the total number of the first tokens of the sample tokens whose capacity prediction probability is greater than a preset probability threshold, or determine the total number of the second tokens whose capacity prediction probability is greater than the prediction probability threshold, as the predicted expert capacity of each expert sub-network at each denoising time step;

[0082] Wherein, the prediction probability threshold is determined by a dynamic threshold module in the trained image generation model.

[0083] In an alternative embodiment, the mixture-of-experts network includes a routing sub-network; during the denoising process, the following operations are respectively performed on the mixture-of-experts network at each denoising time step:

[0084] Use the routing sub-network to determine the affinity matrix between each image token in the image token pool and each expert sub-network in the mixture-of-experts network;

[0085] Based on the affinity matrix and the preset expert capacity or the predicted expert capacity, determine the target tokens that each of the expert sub-networks needs to process respectively from the image token pool;

[0086] Process the corresponding target tokens through each of the expert sub-networks respectively, and perform a fusion process on the processing results of each expert, so as to use the expert fusion processing result as the output of the mixture-of-experts network.

[0087] According to another aspect of the embodiments of the present disclosure, there is provided a computer-readable storage medium, when the instructions in the computer-readable storage medium are executed by a processor of an electronic device, enabling the electronic device to execute the training method and / or the image generation method of the image generation model as described in any one of the above embodiments.

[0088] According to another aspect of the embodiments of the present disclosure, there is provided an electronic device, including:

[0089] A processor;

[0090] A memory for storing executable instructions of the processor;

[0091] Wherein, the processor is configured to execute the instructions to implement the training method and / or the image generation method of the image generation model as described in any of the above embodiments.

[0092] According to another aspect of the embodiments of the present disclosure, there is provided a computer program product, which includes a computer program that, when executed by a processor, implements the training method and / or the image generation method of the image generation model provided in any of the above embodiments.

[0093] The technical solutions provided by the embodiments of the present disclosure at least bring the following beneficial effects:

[0094] In the embodiments of the present disclosure, training data is obtained, where the training data includes a plurality of noise sample images and corresponding conditional sample information; the plurality of noise sample images are sample images with various different image object categories and / or different noise intensities; noise image sample tokens corresponding to each of the noise sample images, and conditional sample feature information corresponding to the conditional sample information are obtained; the token processing module in the image generation model is called to flatten each of the noise image sample tokens in the same training batch to obtain a sample token pool at the training batch level; the image generation model is implemented based on a diffusion model; the image generation model is called to perform denoising processing on the sample token pool and the conditional sample feature information to obtain predicted image information; during the denoising processing, based on a preset expert capacity, target sample tokens to be processed by each expert sub-network in the mixture of experts network in the image generation model are determined from the sample token pool; based on the difference between the predicted image information and the target image information, the image generation model is trained to obtain a trained image generation model.

[0095] Thus, by introducing a token processing module, the tokenization of each of the noise image samples in the same training batch is flattened to obtain a sample token pool at the training batch level. Subsequently, during the diffusion process, from the sample token pool, the target sample tokens to be processed by each expert sub-network in the mixture-of-experts network are determined. Since the noise image sample tokens in the same training batch are from sample images of different image object categories and / or different noise intensities, the sample token pool contains tokens of sample images from different image object categories and / or different noise intensities, realizing the interaction of cross-samples and enabling the experts to access the global token distribution information across different noise levels and image object category conditions, which helps the model capture the inherent heterogeneity among samples during the diffusion process, can improve the performance of the mixture-of-experts model in the diffusion model, and thus is conducive to improving the performance of the trained model.

[0096] It should be understood that the above general description and the following detailed description are merely exemplary and explanatory, and do not limit the present disclosure. BRIEF DESCRIPTION OF THE DRAWINGS

[0097] The accompanying drawings herein are incorporated into the specification and form a part of the present disclosure, showing embodiments consistent with the present disclosure, and are used together with the specification to explain the principles of the present disclosure and do not constitute an undue limitation to the present disclosure.

[0098] Figure 1 is a flowchart of a method for training an image generation model shown according to an exemplary embodiment.

[0099] Figure 2 is a schematic diagram of a partial processing procedure of an image generation model during the training phase shown according to an exemplary embodiment.

[0100] Figure 3 is a schematic diagram of a token selection process of a token selection mechanism and an expert selection mechanism shown.

[0101] Figure 4 is a flowchart of another method for training an image generation model shown according to an exemplary embodiment.

[0102] Figure 5 is a flowchart of an image generation method shown according to an exemplary embodiment.

[0103] Figure 6 is a schematic diagram of a partial processing procedure of an image generation model during the inference phase shown according to an exemplary embodiment.

[0104] Figure 7 In (a), it shows a schematic diagram of the selection result of experts - tokens executed by using the traditional Excpert Choice gating selection mechanism.

[0105] Figure 7 Figure (b) shows a schematic diagram of the selection result of expert - tokens executed by using the traditional Toke Choice gating selection mechanism.

[0106] Figure 8 Figures (a) and (b) show schematic diagrams of the selection results of expert - tokens executed by using the dynamic token selection mechanism of the present disclosure under different noise intensities.

[0107] Figure 9 It is a comparison chart of the computational amounts of the traditional Excpert Choice, Toke Choice gating selection mechanisms, and the dynamic token selection mechanism of the present disclosure at different time steps.

[0108] Figure 10 It is a block diagram of a training device for an image generation model shown according to an exemplary embodiment.

[0109] Figure 11 It is a block diagram of an image generation device shown according to an exemplary embodiment.

[0110] Figure 12 It is a block diagram of an electronic device shown according to an exemplary embodiment. Detailed implementation manners

[0111] In order to enable those of ordinary skill in the art to better understand the technical solutions of the present disclosure, the technical solutions in the embodiments of the present disclosure will be clearly and completely described below with reference to the accompanying drawings.

[0112] It should be noted that the terms "first", "second", etc. in the specification and claims of the present disclosure and the above - mentioned drawings are used to distinguish similar objects, and do not necessarily need to describe a specific order or sequence. It should be understood that such used data can be interchanged under appropriate circumstances so that the embodiments of the present disclosure described here can be implemented in an order different from those illustrated or described here. The implementation manners described in the following exemplary embodiments do not represent all implementation manners consistent with the present disclosure. On the contrary, they are merely examples of devices and methods consistent with some aspects of the present disclosure as detailed in the appended claims.

[0113] It should be noted that the user information (including but not limited to user device information, user personal information, etc.) and data (including but not limited to data for display, data for analysis, etc.) involved in the present disclosure are all information and data authorized by the user or fully authorized by all parties.

[0114] The following explains and describes some technical terms related to the present disclosure.

[0115] Diffusion Models (DM): It is a powerful multi-task learning framework that can model complex data distributions. It iteratively transforms the initial distribution into the target data distribution by sharing parameters. There are mainly two training paradigms for diffusion models: one is the method based on Denoising Diffusion Probabilistic Models (DDPM), which is trained by predicting the noise added at each step with a network. Its core is to learn to estimate and reverse the noise addition process; the other is the flow model method based on Flow-Matching (FM), which directly learns the continuous transformation between distributions.

[0116] Mixture of Experts (MoE): It is a machine learning model architecture designed to improve the performance and flexibility of the overall model by combining multiple "expert" models (sub-models). Each expert model specializes in handling different parts or patterns of the input data, and a "gating network" is responsible for selecting or weighting the outputs of these expert models. The core idea of MoE is division of labor and cooperation. By allowing different expert models to focus on different tasks or data distributions, the expressive power and efficiency of the model are improved. MoE is widely used in large language models to reduce computational overhead through sparse activation. Different expert models can handle different NLP tasks (such as translation, classification, generation, etc.), and the gating network selects experts according to task requirements.

[0117] Expert capacity: It is one of the most important concepts in the mixture of experts model, which is used to define how many tokens an expert in the mixture of experts model can handle. If the capacities of multiple experts reach the upper limit, the tokens will overflow and be passed to the next layer through residual connections or be completely discarded in some cases. Generally, the expert capacity is a pre-set fixed threshold.

[0118] Figure 1 It is a flowchart of a training method for an image generation model shown according to an exemplary embodiment. Figure 2 It is a schematic diagram of a partial processing process of an image generation model in the training stage as shown according to an exemplary embodiment. As Figure 1As shown, the training method of the image generation model can be applied to an electronic device, which can be a terminal, a server, or a similar computing device. Taking the electronic device as a server as an example for illustration, where the server includes, but is not limited to, an independent server, and can also be a server cluster or a distributed system composed of multiple physical servers, or can also be a cloud server that provides basic cloud computing services such as cloud services, cloud databases, cloud computing, cloud functions, cloud storage, network services, cloud communications, intermediate services, domain name services, security services, and big data and artificial intelligence platforms, etc., one or more of them. As Figure 1 and Figure 2 shown, the method includes the following steps.

[0119] In step S101, training data is obtained, and the training data includes a plurality of noise sample images and corresponding conditional sample information; the plurality of noise sample images are sample images with various different image object categories and / or different noise intensities.

[0120] In the embodiments of the present disclosure, the electronic device obtains an original data set (such as the ImageNet data set) for training an image generation model implemented based on a diffusion model, and performs data augmentation processing such as horizontal flipping on the original data set to obtain the training data finally used for training the image generation model. The training data includes a plurality of training batches of sample data sets, and each training batch (batch) of sample data sets includes a plurality of noise sample images and corresponding conditional sample information. Each noise sample image can correspond to one conditional sample information, or multiple noise sample images correspond to the same conditional sample information. The plurality of noise sample images are sample images with various different image object categories and / or different noise intensities.

[0121] Among them, the plurality of noise sample images are sample images with various different image object categories and / or different noise intensities. Optionally, the image object category refers to the category to which the object included in the noise sample image belongs. For example, for a noise sample image of a cat, its corresponding image object category is "cat"; and for a noise sample image of a dog, its corresponding image object category is "dog". The noise sample image can be a randomly generated noise image, and the noise intensity refers to the noise level of the noise sample image. The noise intensity can be represented by ti, where the larger the value of i, the more the noise sample image deviates from the real image.

[0122] The conditional sample information refers to the guiding conditions for constraining image generation. Optionally, the conditional sample information can include, but is not limited to, at least one of image object category, text, label, etc.

[0123] In step S103, the noise image sample tokens corresponding to each of the noise sample images, and the conditional sample feature information corresponding to the conditional sample information are obtained.

[0124] In an embodiment of the present disclosure, for the sample data set of each training batch, the electronic device may obtain the noise image sample tokens corresponding to each of the noise sample images, and obtain the conditional sample feature information corresponding to the conditional sample information.

[0125] Among them, the noise image sample token is an embedded representation used to characterize the noise sample image. The noise image sample token can be obtained through encoding processing by a first encoder. Optionally, the first encoder may include, but is not limited to, a variational autoencoder (VAE Encoder). The conditional sample feature information may be an encoded feature used to characterize the conditional sample information. The conditional sample information can be obtained through encoding processing by a second encoder, and the second encoder may be an encoder that matches the conditional content included in the conditional sample information, and specific limitations are not made here.

[0126] In step S105, the token processing module in the image generation model is called to flatten each of the noise image sample tokens in the same training batch to obtain a sample token pool at the training batch level; the image generation model is implemented based on a diffusion model.

[0127] In an embodiment of the present disclosure, after the electronic device obtains the noise image sample tokens of a training batch, it can use the token processing module in the image generation model to flatten the noise image sample tokens in the same training batch, that is, flatten these noise image sample tokens with multiple different image object categories and / or different noise intensities to construct a sample token pool. The sample token pool is a global token pool at the batch processing level. The sample token pool includes these multiple noise image sample tokens after being flattened. Among them, the expression of the flattening process can be expressed as:

[0128] X ∈ R B×S×D → X pool ∈ R BS×D

[0129] Among them, X is the noise image sample token before flattening, B is the size of the training batch, S is the token length of a sample, D is the hidden dimension, and X pool is the sample token pool.

[0130] In practical applications, such as Figure 2As described above, through the token processing module in the image generation model, the token of each noise image sample in the same training batch is flattened to obtain a sample token pool at the training batch level. Since the token of the processed noise image sample has multiple different image object categories and / or different noise intensities, the sample token pool includes multiple noise image sample tokens with multiple different image object categories and / or different noise intensities.

[0131] In step S107, the image generation model is called to perform denoising processing on the sample token pool and the conditional sample feature information to obtain predicted image information; during the denoising processing, based on a preset expert capacity, the target sample tokens required to be processed by each expert sub-network in the mixture-of-experts network in the image generation model are determined from the sample token pool.

[0132] In the embodiment of the present disclosure, the sample token pool and the conditional sample feature information are used as model inputs and input into the image generation model for denoising processing to obtain predicted image information. The predicted image information may refer to information for training the image generation model, such as predicted image noise, etc.

[0133] Optionally, the image generation model is implemented based on a diffusion model and a Transformer network, that is, the Transformer network is used to replace the U-Net network as the denoising network in the diffusion model. In addition, on the basis of introducing the Transformer network, a mixture-of-experts network (MoE) is also introduced into the image generation model, that is, at least part of the feed-forward network (FNN) in the Transformer network is replaced with a mixture-of-experts network, and other modules of the Transformer network remain unchanged. Exemplarily, the mixture-of-experts network may include N expert sub-networks (such as E1, E2,..., EN). Optionally, each expert sub-network (abbreviated as an expert) may be implemented as a feed-forward network with the same architecture.

[0134] During the denoising processing in the training stage, the electronic device can obtain a preset expert capacity, which is used to define the number of tokens that each expert sub-network can process. Then, according to the preset expert capacity, the target sample tokens required to be processed by each expert sub-network in the mixture-of-experts network are determined from the sample token pool, so as to use each expert sub-network in the mixture-of-experts network to separately process the corresponding target sample tokens in the sample token pool.

[0135] In an alternative embodiment, the mixture-of-experts network includes a routing sub-network; during the denoising processing, the following operations are respectively performed on the mixture-of-experts network at each denoising time step:

[0136] Using the routing sub-network, determine the sample affinity matrix between each sample token in the sample token pool and each expert sub-network in the mixture-of-experts network;

[0137] Based on the sample affinity matrix and the preset expert capacity, determine the target sample tokens that each of the expert sub-networks needs to process from the sample token pool;

[0138] Process the corresponding target sample tokens through each of the expert sub-networks, and perform a fusion process on the processing results of each expert, so as to use the expert fusion processing result as the output of the mixture-of-experts network.

[0139] In the embodiments of the present disclosure, the sample affinity matrix is also referred to as the similarity matrix, which can be used to measure the distance or similarity between two points in a space. The above-mentioned determination of the sample affinity matrix M Dy The calculation expression of can be expressed as:

[0140] M Dy = X pool Wr, Xpool ∈ R BS×D , Wr ∈ R D×N

[0141] where, X pool is the input sample token pool, Wr is the routing sub-network, B is the size of the training batch, S is the label length of a sample, D is the hidden dimension, and N is the total number of expert sub-networks.

[0142] After calculating the sample affinity matrix, based on the sample affinity matrix and the preset expert capacity, determine the target sample tokens that each of the expert sub-networks needs to process from the sample token pool, so as to process the corresponding target sample tokens through each of the expert sub-networks, and perform a weighted sum process on the processing results of each expert sub-network, so as to use the expert fusion processing result as the output of the mixture-of-experts network.

[0143] In an alternative embodiment, the determining the target sample tokens that each of the expert sub-networks needs to process from the sample token pool based on the sample affinity matrix and the preset expert capacity includes: determining the token quantity threshold that each of the expert sub-networks needs to process based on the preset expert capacity; sorting in descending order according to the values of the sample affinity matrix corresponding to each sample token, and determining the filtered sample tokens with the sorting serial number less than or equal to the token quantity threshold from the sample token pool as the target sample tokens that each of the expert sub-networks needs to process.

[0144] In step S109, based on the difference between the predicted image information and the target image information, train the image generation model to obtain the trained image generation model.

[0145] In an embodiment of the present disclosure, the electronic device guides the training of the image generation model according to the difference between the predicted image information and the target image information until the training end condition is met, and obtains the trained image generation model. Here, the target image information corresponds to the predicted image information. Only as an example, if the predicted image information is predicted image noise, the target image information is target image noise. The trained image generation model is at least used to perform tasks such as generating images based on text or generating images based on categories. Of course, in some embodiments, the trained image generation model can also be applied to other tasks in the field of visual generation, such as image restoration, image style adjustment, image object replacement, image scene replacement, video generation, etc.

[0146] In the related art, if the mixture-of-experts model is directly applied to the diffusion model to perform visual generation tasks, in practical applications, the performance of the diffusion model will be damaged to a certain extent, thus limiting the effectiveness of the mixture-of-experts model in the diffusion model. The inventors found that the main reason for this performance degradation is the neglect of the accessibility of the global token distribution information, and the global token distribution is all the characteristics of the diffusion process. Specifically, the gating selection mechanism of the mixture-of-experts model mainly includes a token selection mechanism (Token Choice, TC) and an expert selection mechanism (Expert Choice, EC), as Figure 3 shown. The former selects the expert subnetwork in the mixture-of-experts model based on the input single token, while the latter is that each expert subnetwork selects tokens from samples of the same object category. Obviously, these two gating selection mechanisms limit the interaction between tokens during the expert selection process or limit the token interaction between samples of different object categories. In addition, since the mixture-of-experts model replaces the feed-forward network layer, these two gating mechanisms in the diffusion model are essentially limited to processing tokens in a single sample, where the gating mechanism only operates on tokens sharing the same condition and noise level, and cannot access the tokens of other samples with different noise levels and category conditions, thus hindering the model's ability to capture the inherent heterogeneity between samples in the diffusion process, that is, limiting the mixture-of-experts model's learning of the key contrast information between different samples, and it is difficult for the mixture-of-experts model to learn to process different task intensities affected by noise, thereby resulting in a certain degree of damage to the model performance.

[0147] To solve the above problem of impaired performance, the present disclosure flattens the noise image sample tokens in the same training batch by introducing a token processing module to obtain a sample token pool at the training batch level. Then, during the diffusion process, the target sample tokens to be processed by each expert sub-network in the mixture of experts network are determined from the sample token pool. Since the noise image sample tokens in the same training batch are sample images from different image object categories and / or different noise intensities, the sample token pool contains tokens of sample images from different image object categories and / or different noise intensities, realizing cross-sample interaction and enabling experts to access the global token distribution information across different noise levels and image object category conditions, helping the model capture the inherent heterogeneity among samples during the diffusion process, which can improve the usage effect of the mixture of experts model in the diffusion model and thus contribute to improving the performance of the trained model.

[0148] In addition, enhancing cross-sample token interaction during training, adopting an approximation of the complete token distribution between different noise levels and samples, and promoting more specialized expert learning through comprehensive access to global token information. Empirical analysis shows that the adopted global sample token pool can accelerate loss convergence and outperform dense models with equivalent activation parameters.

[0149] For the stability and simplicity of training, the expert capacity C is maintained during the training process Ei as a fixed expert capacity C, where i ∈ {1…N}. However, this processing method has two drawbacks: (1) If the expert capacity C remains unchanged during the inference stage, the layer performs the same calculations in all sampling steps. (2) When generating a single sample or a small number of samples, the size of the batch-level token pool will be significantly different from that in the training stage, that is, there is a misalignment between training and inference. To address these drawbacks, as Figure 4 shown, in an optional embodiment, after obtaining the sample token pool at the training batch level, the method further includes:

[0150] In step S401, through a capacity prediction network, determine the predicted expert capacity of each expert sub-network at each denoising time step; the predicted expert capacity is used to dynamically determine the number of tokens to be processed by each expert sub-network.

[0151] In the embodiments of the present disclosure, during the training stage, the electronic device inputs the sample token pool into the capacity prediction network to be trained for capacity prediction, and determines the predicted expert capacity of each expert sub-network at each denoising time step. The capacity prediction network is used to dynamically determine the token selection of each expert sub-network. The predicted expert capacity is used to dynamically determine the number of tokens to be processed by each expert sub-network.

[0152] During the training process, since each MoE layer will have a result of the final selection of a token and an expert, a capacity prediction network is introduced to predict this selection result, and the dynamic token selection during inference is guided based on the predicted result. By selecting, in each expert sub-network, the target sample tokens with the highest token-expert sample affinity matrix in the sample token pool and maintaining the same total activation capacity as the baseline model, the model captures tokens at different time steps and different image object categories while maintaining the same computational cost. By learning from the training-time tokens and the dynamic routing pattern, it is beneficial to effectively allocate computational resources between complex and simple situations.

[0153] Optionally, the capacity prediction network is a lightweight structure with a simple two-layer multi-layer perceptron (MLP) with the SiLU activation function. If CP(X pool ) ∈ R BS×D represents the output of the capacity predictor, then CP(X pool ) can be expressed as:

[0154] CP(X pool ) = W2σSiLU(W1X pool )

[0155] where W1 and W2 are weight matrices, σSiLU is the activation function, and X pool is the input sample token pool.

[0156] In the above embodiment, by introducing a capacity prediction network to implement a dynamic token selection mechanism, it can be dynamically adjusted according to the category condition information and the noise level during the diffusion process, improving the model training effect. In addition, considering the characteristics of different time steps of the diffusion model, computational resources are adaptively allocated according to the task complexity, which is determined by the noise level and token-specific requirements, while maintaining a consistent average computational cost. Specifically, more computational resources are dynamically allocated to challenging tasks and fewer resources are allocated to simpler tasks, which can adapt to the token complexity and the requirements of specific time steps.

[0157] In an alternative embodiment, determining the predicted expert capacity of each expert sub-network at each denoising time step through the capacity prediction network includes:

[0158] For each expert sub-network at each denoising time step, input each sample token in the sample token pool into the capacity prediction network to output the capacity prediction probability of each sample token;

[0159] For each expert sub-network at each denoising time step, determine the total number of sample tokens with a capacity prediction probability greater than a preset probability threshold, and use the determined total number of the first sample tokens as the predicted expert capacity of each expert sub-network at each denoising time step.

[0160] In an embodiment of the present disclosure, for each expert sub-network at each denoising time step, the electronic device inputs each sample token in the sample token pool into the capacity prediction network and outputs the capacity prediction probability of each sample token. Then, the total number of sample tokens with a capacity prediction probability greater than a preset probability threshold is counted to obtain the total number of the first sample tokens, and the total number of the first sample tokens is used as the predicted expert capacity of each expert sub-network at each denoising time step. Exemplarily, during the training process, the preset probability threshold can be set to a fixed value less than 1 (such as 0.9, 0.8, etc.).

[0161] In this case, continue as Figure 4 shown, the training of the image generation model based on the difference between the predicted image information and the target image information to obtain the trained image generation model includes:

[0162] In step S403, based on the difference between the predicted image information and the target image information, and the difference between the trained expert capacity and the predicted expert capacity, the image generation model is trained to obtain the trained image generation model; the mixture-of-experts network in the trained image generation model includes the trained capacity prediction network.

[0163] In an embodiment of the present disclosure, the electronic device guides the training of the image generation model according to the difference between the predicted image information and the target image information until the training end condition is met to obtain the trained image generation model. At the same time, the electronic device also guides the training of the capacity prediction network separately according to the difference between the trained expert capacity and the predicted expert capacity until the training end condition is met to obtain the trained capacity prediction network. The mixture-of-experts network in the trained image generation model includes the trained capacity prediction network.

[0164] In an alternative embodiment, the training of the image generation model based on the difference between the predicted image information and the target image information, and the difference between the trained expert capacity and the predicted expert capacity to obtain the trained image generation model includes:

[0165] Determine the first training loss information based on the difference between the predicted image information and the target image information;

[0166] Determine the second training loss information based on the difference between the trained expert capacity and the predicted expert capacity;

[0167] Taking minimizing the first training loss information and minimizing the second training loss information as objectives, the diffusion model and the capacity prediction network are independently trained respectively to obtain the trained image generation model.

[0168] In the embodiments of the present disclosure, the first training loss information refers to the diffusion loss, and the loss function it adopts can be determined according to any one of the aforementioned training paradigms of the denoising diffusion probability model method or the flow model method based on flow matching, and no specific limitation is made thereto. The second training loss information L CP can be represented by the BCELoss loss function. With the goal of minimizing the first training loss information and minimizing the second training loss information, the diffusion model and the capacity prediction network are independently optimized respectively to obtain the trained image generation model. During the training process, the stop gradient technique is used to train the capacity prediction network to ensure that the capacity prediction network only focuses on the input features of the current layer, and at the same time prevent the capacity prediction network from interfering with the training of the main diffusion transformers model. Therefore, the second training loss information L CP will not affect the actual first training loss information.

[0169] In an alternative embodiment, the mixture of experts network includes a dynamic threshold module; the method further includes:

[0170] During the training process, the exponential moving average technique is adopted by the dynamic threshold module to determine the prediction probability threshold of each expert sub-network at each denoising time step; during the inference process, based on the prediction probability threshold determined by the capacity prediction network and the dynamic threshold module, the predicted expert capacity of each expert sub-network at each denoising time step is determined.

[0171] In the embodiments of the present disclosure, in order to ensure that the predicted expert capacity during the inference process is approximately equal to the predicted expert capacity during the training process, the dynamic threshold technique is implemented by introducing a dynamic threshold module. During the training process, the exponential moving average (EMA) technique is used by the dynamic threshold module to determine the prediction probability threshold of each expert sub-network at each denoising time step. Specifically, the predicted expert capacity of the selected k-th token can be used as the exponential moving average, and based on the exponential moving average and the previously determined prediction probability threshold, the current prediction probability threshold is determined. In the inference result, the prediction probability threshold determined by the dynamic threshold module replaces the fixed preset probability threshold to constrain the capacity prediction network to determine the predicted expert capacity, improving the reliability and flexibility of the dynamic token selection.

[0172] Figure 5 is a flowchart of an image generation method shown according to an exemplary embodiment. Figure 6 is a schematic diagram of a partial processing process of an image generation model during the inference phase shown according to an exemplary embodiment. The image generation method can be applied to an electronic device, and the electronic device can be a terminal, a server, or a similar computing device. As Figure 5 andFigure 6 As shown, the method includes the following steps.

[0173] In step S501, multiple noise images and conditional information are obtained.

[0174] Optionally, the multiple noise images can be sample images with multiple different image object categories and / or different noise intensities. The conditional sample information refers to the guiding conditions for constraining image generation. Optionally, the conditional sample information can include, but is not limited to, at least one of image object category, text, label, etc.

[0175] In step S503, the noise image tokens corresponding to each of the noise images and the conditional feature information corresponding to the conditional information are obtained.

[0176] Optionally, the noise image tokens can be obtained through encoding processing by a first encoder. Optionally, the first encoder can include, but is not limited to, a variational autoencoder (VAE Encoder). The conditional feature information can be the encoded feature for characterizing the conditional information. The conditional information can be obtained through encoding processing by a second encoder, and the second encoder can be an encoder that matches the conditional content included in the conditional sample information, and specific limitations are not made here.

[0177] In step S505, the token processing module in the trained image generation model is called to flatten each of the noise image tokens to obtain an image token pool; the trained image generation model is trained based on any one of the above methods.

[0178] Optionally, the token processing module in the trained image generation model is used to flatten multiple noise image sample tokens to construct an image token pool. The image token pool can be a global token pool at the batch level. The image token pool includes these flattened multiple noise image tokens.

[0179] In step S507, the trained image generation model is called to perform denoising processing on the image token pool and the conditional feature information to generate a target image; during the denoising processing, based on a preset expert capacity, the target tokens required by each expert sub-network in the mixture of experts network in the trained image generation model are determined from the image token pool.

[0180] Optionally, the image token pool and conditional information are used as model inputs and input into the trained image generation model for denoising processing, and the denoised output tokens are decoded by a decoder to generate a target image. During the denoising process, a preset expert capacity can be obtained, which is used to define the number of tokens that each expert subnet can process. Then, according to the preset expert capacity, the target tokens to be processed by each expert subnet in the mixture-of-experts network are determined from the image token pool, so as to use each expert subnet in the mixture-of-experts network to process the corresponding target tokens in the sample token pool respectively.

[0181] It should be noted that the relevant content about the image generation model here can be found in the description of the previous embodiments, and some details about the inference stage can refer to the previous training stage, which will not be elaborated here.

[0182] In the above embodiments, during the training process, each expert subnet in the image generation model can access the global token distribution information across different noise levels and image object category conditions, which helps the model capture the inherent heterogeneity among samples during the diffusion process, and can improve the usage effect of the mixture-of-experts model in the diffusion model. Therefore, using the trained image generation model for image generation can improve the image generation quality.

[0183] In an alternative embodiment, the trained image generation model includes a capacity prediction network. After obtaining the image token pool, the method further includes:

[0184] Determine the predicted expert capacity of each expert subnet at each denoising time step through the capacity prediction network; the predicted expert capacity is used to dynamically determine the number of tokens that each expert subnet needs to process; during the denoising process, based on the predicted expert capacity, the target tokens to be processed by each expert subnet in the mixture-of-experts network of the image generation model are determined from the image token pool.

[0185] Optionally, as Figure 6 shown, the capacity prediction network is used to determine the number of tokens that each expert subnet needs to process. For example, the number of tokens that expert subnet - 1 needs to process is 3, the number of tokens that expert subnet - 2 needs to process is 2, and the number of tokens that expert subnet - 3 needs to process is 1. Obviously, through the capacity prediction network, the number of tokens that each expert subnet needs to process is neither a fixed value nor the same as each other.

[0186] In the above embodiments, by introducing a capacity prediction network to implement a dynamic token selection mechanism, it is possible to adaptively allocate computing resources according to task complexity, which is determined by the noise level and token-specific requirements, while maintaining a consistent average computing cost. Specifically, more computing resources are dynamically allocated to challenging tasks and fewer resources are allocated to simpler tasks, which can adapt to token complexity and requirements at specific time steps.

[0187] In an alternative embodiment, determining the predicted expert capacity of each expert sub-network at each denoising time step through the capacity prediction network includes:

[0188] For each expert sub-network at each denoising time step, input each image token in the image token pool into the capacity prediction network, and output the capacity prediction probability of each image token;

[0189] For each expert sub-network at each denoising time step, determine the total number of the first tokens of the sample tokens whose capacity prediction probability is greater than a preset probability threshold, or determine the total number of the second tokens whose capacity prediction probability is greater than the prediction probability threshold, as the predicted expert capacity of each expert sub-network at each denoising time step;

[0190] Wherein, the prediction probability threshold is determined by a dynamic threshold module in the trained image generation model.

[0191] Optionally, in the case where the trained image generation model does not include a dynamic threshold module, for each expert sub-network at each denoising time step, determine the total number of the first tokens of the sample tokens whose capacity prediction probability is greater than a preset probability threshold, as the predicted expert capacity of each expert sub-network at each denoising time step. And in the case where the trained image generation model includes a dynamic threshold module, for each expert sub-network at each denoising time step, determine the total number of the second tokens whose capacity prediction probability is greater than the prediction probability threshold, as the predicted expert capacity of each expert sub-network at each denoising time step. In this way, by replacing the fixed preset probability threshold with the prediction probability threshold determined by the dynamic threshold module to constrain the capacity prediction network to determine the predicted expert capacity, the reliability and flexibility of dynamic token selection are improved.

[0192] In an alternative embodiment, the mixture-of-experts network includes a routing sub-network; during the denoising process, the following operations are respectively performed on the mixture-of-experts network at each denoising time step:

[0193] Use the routing sub-network to determine the affinity matrix between each image token in the image token pool and each expert sub-network in the mixture-of-experts network;

[0194] Based on the affinity matrix and the preset expert capacity or the predicted expert capacity, determine the target tokens that each of the expert sub-networks needs to process from the image token pool;

[0195] Process the corresponding target tokens through each of the expert sub-networks, and perform a fusion process on the processing results of each expert, so as to use the expert fusion processing result as the output of the mixture-of-experts network.

[0196] In an alternative embodiment, the above determination of the target tokens that each of the expert sub-networks needs to process may include: based on the preset expert capacity, determine the token quantity threshold that each expert sub-network needs to process; sort in descending order according to the values of the affinity matrix corresponding to each image token, and determine the screened image tokens whose sorting serial numbers are less than or equal to the token quantity threshold from the image token pool as the target tokens that each expert sub-network needs to process.

[0197] In another alternative embodiment, the above determination of the target tokens that each of the expert sub-networks needs to process may include: obtain the predicted expert capacity determined by the capacity prediction network; based on the predicted expert capacity, determine the token quantity threshold that each expert sub-network needs to process; sort in descending order according to the values of the affinity matrix corresponding to each image token, and determine the screened image tokens whose sorting serial numbers are less than or equal to the token quantity threshold from the image token pool as the target tokens that each expert sub-network needs to process.

[0198] It should be noted that some specific details of the embodiments of the above image generation method can refer to the foregoing training embodiments, and will not be elaborated here.

[0199] It should be noted that the above image generation method only involves the image generation model for performing tasks such as generating images based on text or generating images based on categories. Of course, in some embodiments, the image generation model can also be extended to other tasks in the field of visual generation, such as image restoration, image style adjustment, image object replacement, image scene replacement, video generation, etc.

[0200] To illustrate the difference between the dynamic token selection mechanism provided by the present disclosure and the traditional two gating selection mechanisms, the following is combined with Figures 7 - 9 for illustration.

[0201] Figure 7 In (a) shows a schematic diagram of the selection result of the expert-token using the traditional Excpert Choice gating selection mechanism, as shown in Figure 7As shown in (a), the sum of the columns of the affinity matrix is 1, and each expert selects TopK (k=3) tokens, that is, 3 tokens are selected from each column, a total of 24 tokens are selected, and the processing amount of each denoising time step is fixed. This gating mechanism is prone to miss the defect of selecting tokens.

[0202] Figure 7 (b) shows a schematic diagram of the selection results of the expert-token using the traditional Toke Choice gating selection mechanism. Figure 7 As shown in (b), the sum of the rows of the affinity matrix is 1, and each token selects the TopK (k=4) experts, that is, 4 experts are selected in each row, a total of 24 tokens are selected, and the processing volume of each denoising time step is fixed. In this gating mechanism, the probability of each expert being selected will be unequal, which is prone to imbalance.

[0203] Figure 8 (a) and (b) show schematic diagrams of the selection results of expert-tokens under different noise intensities using the dynamic token selection mechanism disclosed in the present invention. Figure 8 As shown in (a), under high noise intensity, the sum of the columns of the affinity matrix is 1, and each expert selects TopDyK (k is not fixed) tokens, for a total of 9 tokens. Figure 8 As shown in (b), under low noise intensity, the sum of the columns of the affinity matrix is 1, and each expert selects TopDyK (k is not fixed) tokens, for a total of 37 tokens. Obviously, the dynamic token selection mechanism disclosed in the present invention also requires experts to select tokens, but the number of tokens selected is not fixed. The number of tokens selected is jointly determined by the capacity prediction network and the token itself, which is equivalent to the token also having the ability to select experts, thus realizing the two-way selection of tokens and experts.

[0204] Figure 9 The following figure shows the comparison of the computational amount of the above three selection mechanisms at different time steps. Each line connecting a token and an expert represents a selection (i.e., computation). In the denoising process of the diffusion model at different time steps, both Token Choice (TC) and Expert Choice (EC) have a fixed computational amount (e.g. Figure 9, both are two connection lines). The dynamic token selection mechanism of the present disclosure adaptively allocates different amounts of computation at different time steps, but maintains the same average activation amount throughout the entire complete inference process. Specifically, the total activation amounts of each selection mechanism are TC: 2 + 2, EC: 2 + 2, and the present disclosure: 3 + 1. Obviously, the average activation amount of all three is 2. Through such a design, the expert model can be made more specialized at different time steps, improving the training efficiency and obtaining better inference results. This is because the TC scheme and the EC scheme ignore the different characteristics of the diffusion model at different time steps. For example, in the early time steps, the diffusion model pays more attention to the overall contour composition, and in the later time steps, it pays more attention to the detailed texture. The present disclosure configures different amounts of computation at different time steps, improving the utilization rate and effect of computing resources. In addition, the experimental results show that under the same average activation parameters, the scheme of the present disclosure is significantly better than the TC scheme and the EC scheme, and can reach about 1.4 times the activation parameters.

[0205] Figure 10 is a block diagram of a training device for an image generation model shown according to an exemplary embodiment. Referring to Figure 10 , the device includes:

[0206] A training data acquisition module 1010, configured to execute acquiring training data, where the training data includes a plurality of noise sample images and corresponding conditional sample information; the plurality of noise sample images are sample images having various different image object categories and / or different noise intensities;

[0207] An input acquisition module 1030, configured to execute acquiring noise image sample tokens corresponding to each of the noise sample images, and conditional sample feature information corresponding to the conditional sample information;

[0208] A sample token processing module 1050, configured to execute calling a token processing module in the image generation model to flatten each of the noise image sample tokens in the same training batch to obtain a training batch-level sample token pool; the image generation model is implemented based on a diffusion model;

[0209] A sample denoising module 1070, configured to execute calling the image generation model to perform denoising processing on the sample token pool and the conditional sample feature information to obtain predicted image information; during the denoising processing, based on a preset expert capacity, determine target sample tokens required for each expert sub-network in the mixture-of-experts network in the image generation model from the sample token pool;

[0210] A training module 1090, configured to execute training the image generation model based on the difference between the predicted image information and the target image information to obtain a trained image generation model.

[0211] In an alternative embodiment, the apparatus further comprises:

[0212] a capacity prediction module configured to execute, through a capacity prediction network, to determine the predicted expert capacity of each expert sub-network at each denoising time step; the predicted expert capacity is used to dynamically determine the number of tokens that each expert sub-network needs to process respectively;

[0213] The training module is further configured to execute:

[0214] Based on the difference between the predicted image information and the target image information, and the difference between the trained expert capacity and the predicted expert capacity, train the image generation model to obtain a trained image generation model; the mixture-of-experts network in the trained image generation model includes the trained capacity prediction network.

[0215] In an alternative embodiment, the training module includes:

[0216] a first loss determination unit configured to execute to determine first training loss information based on the difference between the predicted image information and the target image information;

[0217] a second loss determination unit configured to execute to determine second training loss information based on the difference between the trained expert capacity and the predicted expert capacity;

[0218] a training unit configured to execute to independently train the diffusion model and the capacity prediction network respectively with the goal of minimizing the first training loss information and minimizing the second training loss information to obtain a trained image generation model.

[0219] In an alternative embodiment, the capacity prediction module includes:

[0220] a capacity prediction probability determination unit configured to execute, for each expert sub-network at each denoising time step, input each sample token in the sample token pool into the capacity prediction network and output the capacity prediction probability of each sample token;

[0221] a first capacity prediction unit configured to execute, for each expert sub-network at each denoising time step, determine the total number of sample tokens whose capacity prediction probability is greater than a preset probability threshold, and use the determined total number of first sample tokens as the predicted expert capacity of each expert sub-network at each denoising time step.

[0222] In an alternative embodiment, the mixture-of-experts network includes a dynamic threshold module; the apparatus further comprises:

[0223] A threshold determination module, configured to execute during the training process, by using the exponential moving average technique in the dynamic threshold module, determine the prediction probability thresholds of each expert sub-network at each denoising time step; during the inference process, based on the prediction probability thresholds determined by the capacity prediction network and the dynamic threshold module, determine the prediction expert capacities of each expert sub-network at each denoising time step.

[0224] In an alternative embodiment, the mixture-of-experts network includes a routing sub-network; during the denoising process, the following operations are respectively performed on the mixture-of-experts network at each denoising time step:

[0225] Use the routing sub-network to determine the sample affinity matrix between each sample token in the sample token pool and each expert sub-network in the mixture-of-experts network;

[0226] Based on the sample affinity matrix and the preset expert capacity or the predicted expert capacity, determine the target sample tokens required to be processed by each of the expert sub-networks from the sample token pool;

[0227] Process the corresponding target sample tokens through each of the expert sub-networks, and perform a fusion process on the processing results of each expert, so as to use the expert fusion processing result as the output of the mixture-of-experts network.

[0228] In an alternative embodiment, the determining the target sample tokens required to be processed by each of the expert sub-networks from the sample token pool based on the sample affinity matrix and the preset expert capacity or the predicted expert capacity includes:

[0229] Based on the preset expert capacity or the predicted expert capacity, determine the token quantity threshold required to be processed by each of the expert sub-networks;

[0230] Sort in descending order according to the values of the sample affinity matrix corresponding to each sample token, and determine the screened sample tokens with the sorting serial number less than or equal to the token quantity threshold from the sample token pool as the target sample tokens required to be processed by each of the expert sub-networks.

[0231] Figure 11 It is a block diagram of an image generation device shown according to an exemplary embodiment. Refer to Figure 11 , the device includes:

[0232] A first acquisition module 1110, configured to execute acquiring a plurality of noise images and conditional information;

[0233] A second acquisition module 1120, configured to execute acquiring the noise image tokens corresponding to each of the noise images, and the conditional feature information corresponding to the conditional information;

[0234] A token processing module 1130, configured to execute calling the token processing module in the trained image generation model to flatten each of the noise image tokens to obtain an image token pool; the trained image generation model is trained by any of the above methods;

[0235] An image generation module 1140, configured to execute calling the trained image generation model to perform denoising processing on the image token pool and the conditional feature information to generate a target image; during the denoising processing, based on a preset expert capacity, determine target tokens required to be processed by each expert sub-network in the mixed expert network in the trained image generation model from the image token pool.

[0236] In an optional implementation manner, the device further includes:

[0237] A capacity prediction module, configured to execute determining, through the capacity prediction network, the predicted expert capacity of each expert sub-network at each denoising time step; the predicted expert capacity is used to dynamically determine the number of tokens required to be processed by each expert sub-network; during the denoising processing, based on the predicted expert capacity, determine target tokens required to be processed by each expert sub-network in the mixed expert network in the image generation model from the image token pool.

[0238] In an optional implementation manner, the capacity prediction module includes:

[0239] A capacity prediction probability determination unit, configured to execute, for each expert sub-network at each denoising time step, input each image token in the image token pool into the capacity prediction network and output the capacity prediction probability of each image token;

[0240] A capacity prediction unit, configured to execute, for each expert sub-network at each denoising time step, determine the total number of the first tokens of the sample tokens whose capacity prediction probability is greater than a preset probability threshold, or determine the total number of the second tokens whose capacity prediction probability is greater than the prediction probability threshold as the predicted expert capacity of each expert sub-network at each denoising time step;

[0241] Wherein, the prediction probability threshold is determined by a dynamic threshold module in the trained image generation model.

[0242] In an optional implementation manner, the mixed expert network includes a routing sub-network; during the denoising processing, the following operations are respectively performed on the mixed expert network at each denoising time step:

[0243] Adopt the routing sub-network to determine the affinity matrix between each image token in the image token pool and each expert sub-network in the mixed expert network;

[0244] Based on the affinity matrix and the preset expert capacity or the predicted expert capacity, determine the target tokens that each of the expert sub-networks needs to process from the image token pool;

[0245] Process the corresponding target tokens through each of the expert sub-networks respectively, and perform a fusion process on the processing results of each expert, so as to use the expert fusion processing result as the output of the mixture-of-experts network.

[0246] It should be noted that regarding the device in the above embodiments, the specific manners and beneficial effects of each step have been described in detail in the embodiments of the foregoing method, and will not be elaborated herein.

[0247] Figure 12 is a block diagram of an electronic device shown according to an exemplary embodiment. Refer to Figure 12 , the electronic device includes a processor; a memory for storing instructions executable by the processor; wherein, when the processor is configured to execute the instructions stored on the memory, the steps of any of the image generation model training methods and / or image generation methods in the above embodiments are implemented.

[0248] This electronic device may be a terminal, a server, or a similar computing device. Taking this electronic device as a server as an example, Figure 12 is a block diagram of an electronic device for an image generation model training method and / or an image generation method shown according to an exemplary embodiment. The electronic device 1200 may vary greatly due to configuration or performance differences, and may include one or more central processing units (CPUs) 1210 (the processor 1210 may include, but is not limited to, a processing device such as a microprocessor MCU or a field programmable gate array FPGA), a memory 1230 for storing data, and one or more storage media 1220 (such as one or more mass storage devices) for storing application programs 1223 or data 1222. Among them, the memory 1230 and the storage media 1220 may be transient storage or persistent storage. The program stored in the storage media 1220 may include one or more modules, and each module may include a series of instruction operations on the electronic device. Further, the central processor 1210 may be configured to communicate with the storage media 1220 and execute a series of instruction operations in the storage media 1220 on the electronic device 1200.

[0249] The electronic device 1200 may also include one or more power supplies 1260, one or more wired or wireless network interfaces 1250, one or more input / output interfaces 1240, and / or one or more operating systems 1221, such as Windows Server, Mac OS X, Unix, Linux, FreeBSD, and so on.

[0250] The input / output interface 1240 can be used to receive or send data via a network. Specific examples of the above network may include a wireless network provided by a communication provider of the electronic device 1200. In one example, the input / output interface 1240 includes a network adapter (Network Interface Controller, NIC), which can be connected to other network devices through a base station so as to communicate with the Internet. In an exemplary embodiment, the input / output interface 1240 may be a radio frequency (RF) module, which is used to communicate with the Internet wirelessly.

[0251] Those of ordinary skill in the art can understand that Figure 12 the structure shown is only schematic and does not limit the structure of the above electronic device. For example, the electronic device 1200 may also include more or fewer components than Figure 12 shown, or have a different configuration from Figure 12 shown.

[0252] In an exemplary embodiment, there is also provided a computer-readable storage medium including instructions, such as a memory including instructions, and the above instructions can be executed by a processor of the electronic device 1200 to complete the above method. Optionally, the computer-readable storage medium may be a ROM, a random access memory (RAM), a CD-ROM, a magnetic tape, a floppy disk, and an optical data storage device, etc.

[0253] In an exemplary embodiment, there is also provided a computer storage medium, when the instructions in the computer storage medium are executed by a processor of an electronic device, enabling the electronic device to execute the steps of the method provided in any of the above embodiments.

[0254] In an exemplary embodiment, there is also provided a computer program product, which includes a computer program / instructions, and when the computer program / instructions are executed by a processor, the method provided in any of the above embodiments is implemented. Optionally, the computer program is stored in a computer-readable storage medium. The processor of the electronic device reads the computer program from the computer-readable storage medium, and the processor executes the computer program, so that the electronic device executes the method provided in any of the above embodiments.

[0255] Those of ordinary skill in the art can understand that all or part of the processes in the methods of the above embodiments can be completed by instructing relevant hardware through a computer program. The computer program can be stored in a non-volatile computer-readable storage medium. When the computer program is executed, it can include the processes of the embodiments of the above methods. Among them, any reference to a memory, storage, database, or other medium used in the various embodiments provided in the present application can include non-volatile and / or volatile memories. Non-volatile memories can include read-only memory (ROM), programmable ROM (PROM), electrically programmable ROM (EPROM), electrically erasable programmable ROM (EEPROM), or flash memory. Volatile memories can include random access memory (RAM) or external cache memory. By way of illustration and not limitation, RAM is available in many forms, such as static RAM (SRAM), dynamic RAM (DRAM), synchronous DRAM (SDRAM), double data rate SDRAM (DDR SDRAM), enhanced SDRAM (ESDRAM), synchronous link DRAM (SLDRAM), Rambus direct RAM (RDRAM), direct memory bus dynamic RAM (DRDRAM), and Rambus dynamic RAM (RDRAM), etc.

[0256] After considering the specification and practicing the invention disclosed herein, those skilled in the art will readily conceive of other embodiments of the present disclosure. The present application is intended to cover any variations, uses, or adaptations of the present disclosure, which follow the general principles of the present disclosure and include known common general knowledge or conventional technical means in the technical field not disclosed in the present disclosure. The specification and embodiments are only regarded as exemplary, and the true scope and spirit of the present disclosure are pointed out by the following claims.

[0257] It should be understood that the present disclosure is not limited to the exact structures already described and shown in the drawings, and various modifications and changes can be made without departing from its scope. The scope of the present disclosure is only limited by the appended claims.

Claims

1. A training method for an image generation model, characterized in that Including: Obtain training data, where the training data includes a plurality of noisy sample images and corresponding conditional sample information; The plurality of noisy sample images are sample images with various different image object categories and / or different noise intensities; Obtain the noisy image sample tokens corresponding to each of the noisy sample images, and the conditional sample feature information corresponding to the conditional sample information; Invoke the token processing module in the image generation model to flatten each of the noisy image sample tokens in the same training batch to obtain a sample token pool at the training batch level; The image generation model is implemented based on a diffusion model; Invoke the image generation model to perform denoising processing on the sample token pool and the conditional sample feature information to obtain predicted image information; During the denoising process, based on a preset expert capacity, determine the target sample tokens required to be processed by each expert sub-network in the mixture expert network in the image generation model from the sample token pool; Based on the difference between the predicted image information and the target image information, train the image generation model to obtain a trained image generation model.

2. The method according to claim 1, wherein After obtaining the sample token pool at the training batch level, the method further includes: Determine the predicted expert capacity of each expert sub-network at each denoising time step through a capacity prediction network; the predicted expert capacity is used to dynamically determine the number of tokens required to be processed by each expert sub-network; The training the image generation model based on the difference between the predicted image information and the target image information to obtain a trained image generation model includes: Based on the difference between the predicted image information and the target image information, and the difference between the trained expert capacity and the predicted expert capacity, train the image generation model to obtain a trained image generation model; the mixture expert network in the trained image generation model includes the trained capacity prediction network.

3. The method according to claim 2, wherein The training the image generation model based on the difference between the predicted image information and the target image information, and the difference between the trained expert capacity and the predicted expert capacity to obtain a trained image generation model includes: Based on the difference between the predicted image information and the target image information, determine the first training loss information; Based on the difference between the trained expert capacity and the predicted expert capacity, determine the second training loss information; With the goal of minimizing the first training loss information and minimizing the second training loss information, independently train the diffusion model and the capacity prediction network respectively to obtain a trained image generation model.

4. The method according to claim 2, characterized in that, The determining the predicted expert capacity of each expert sub-network at each denoising time step through a capacity prediction network includes: For each expert sub-network at each denoising time step, input each sample token in the sample token pool into the capacity prediction network to output the capacity prediction probability of each sample token; For each expert sub-network at each denoising time step, determine the total number of sample tokens whose capacity prediction probability is greater than a preset probability threshold, and use the determined total number of the first sample tokens as the predicted expert capacity of each expert sub-network at each denoising time step.

5. The method according to claim 2, characterized in that The mixture-of-experts network includes a dynamic threshold module; the method further includes: During the training process, use the exponential moving average technique through the dynamic threshold module to determine the prediction probability threshold of each expert sub-network at each denoising time step; during the inference process, based on the prediction probability threshold determined by the capacity prediction network and the dynamic threshold module, determine the predicted expert capacity of each expert sub-network at each denoising time step.

6. The method according to any one of claims 1-5, characterized in that The mixture-of-experts network includes a routing sub-network; during the denoising process, perform the following operations on the mixture-of-experts network at each denoising time step respectively: Use the routing sub-network to determine the sample affinity matrix between each sample token in the sample token pool and each expert sub-network in the mixture-of-experts network; Based on the sample affinity matrix and the preset expert capacity, determine the target sample tokens to be processed by each of the expert sub-networks from the sample token pool; Process the corresponding target sample tokens through each of the expert sub-networks respectively, and perform a fusion process on the processing results of each expert, so as to use the expert fusion processing result as the output of the mixture-of-experts network.

7. The method according to claim 6, wherein The determining the target sample tokens to be processed by each of the expert sub-networks from the sample token pool based on the sample affinity matrix and the preset expert capacity includes: Based on the preset expert capacity, determine the token quantity threshold required by each of the expert sub-networks; Sort in descending order according to the values of the sample affinity matrix corresponding to each sample token, and determine the filtered sample tokens whose sorting serial number is less than or equal to the token quantity threshold from the sample token pool as the target sample tokens to be processed by each of the expert sub-networks.

8. An image generation method, characterized in that, Include: Obtain a plurality of noisy images and conditional information; Obtain the noisy image tokens corresponding to each of the noisy images, and the conditional feature information corresponding to the conditional information; Call the token processing module in the trained image generation model to flatten each of the noisy image tokens to obtain an image token pool; the trained image generation model is trained based on the method according to any one of claims 1-7; Call the trained image generation model to perform denoising processing on the image token pool and the conditional feature information to generate a target image; During the denoising process, based on the preset expert capacity, determine the target tokens to be processed by each of the expert sub-networks in the mixture-of-experts network in the trained image generation model from the image token pool.

9. The method according to claim 8, wherein The trained image generation model includes a capacity prediction network, and after obtaining the image token pool, the method further includes: Through the capacity prediction network, determine the predicted expert capacity of each expert sub-network at each denoising time step; the predicted expert capacity is used to dynamically determine the number of tokens that each expert sub-network needs to process; during the denoising process, based on the predicted expert capacity, determine the target tokens that each expert sub-network in the mixture-of-experts network in the image generation model needs to process from the image token pool.

10. The method according to claim 8, characterized in that, The determining, through the capacity prediction network, the predicted expert capacity of each expert sub-network at each denoising time step includes: For each expert sub-network at each denoising time step, input each image token in the image token pool into the capacity prediction network, and output the capacity prediction probability of each image token; For each expert sub-network at each denoising time step, determine the total number of the first tokens of the sample tokens whose capacity prediction probability is greater than the preset probability threshold, or determine the total number of the second tokens whose capacity prediction probability is greater than the prediction probability threshold, as the predicted expert capacity of each expert sub-network at each denoising time step; Wherein, the prediction probability threshold is determined by a dynamic threshold module in the trained image generation model.

11. According to the method as claimed in any one of claims 8 - 10, characterized in that, The mixture-of-experts network includes a routing sub-network; during the denoising process, the following operations are respectively performed on the mixture-of-experts network at each denoising time step: Use the routing sub-network to determine the affinity matrix between each image token in the image token pool and each expert sub-network in the mixture-of-experts network; Based on the affinity matrix and the preset expert capacity or the predicted expert capacity, determine the target tokens that each expert sub-network needs to process from the image token pool; Process the corresponding target tokens through each expert sub-network respectively, and perform fusion processing on the processing results of each expert, so as to use the expert fusion processing result as the output of the mixture-of-experts network.

12. A training device for an image generation model, characterized in that, Including: A training data acquisition module, configured to acquire training data, where the training data includes a plurality of noisy sample images and corresponding conditional sample information; The plurality of noisy sample images are sample images with multiple different image object categories and / or different noise intensities; An input acquisition module, configured to acquire the noisy image sample tokens corresponding to each of the noisy sample images, and the conditional sample feature information corresponding to the conditional sample information; A sample token processing module, configured to call the token processing module in the image generation model to flatten each of the noisy image sample tokens in the same training batch to obtain a sample token pool at the training batch level; The image generation model is implemented based on a diffusion model; A sample denoising module, configured to call the image generation model to perform denoising processing on the sample token pool and the conditional sample feature information to obtain predicted image information; During the denoising process, based on the preset expert capacity, determine the target sample tokens that each expert sub-network in the mixture-of-experts network in the image generation model needs to process from the sample token pool; A training module, configured to perform training on the image generation model based on the difference between the predicted image information and the target image information, to obtain a trained image generation model.

13. A training device for an image generation model, characterized in that, Comprising: A first acquisition module, configured to acquire a plurality of noise images and conditional information; A second acquisition module, configured to acquire noise image tokens corresponding to each of the noise images, and conditional feature information corresponding to the conditional information; A token processing module, configured to call the token processing module in the trained image generation model to perform flattening processing on each of the noise image tokens to obtain an image token pool; the trained image generation model is trained based on the method according to any one of claims 1-7; An image generation module, configured to call the trained image generation model to perform denoising processing on the image token pool and the conditional feature information to generate a target image; During the denoising processing, based on a preset expert capacity, determine target tokens required to be processed by each expert sub-network in the mixture of experts network in the trained image generation model from the image token pool.

14. An electronic device, characterized in that, Comprising: A processor; A memory for storing executable instructions of the processor; Wherein, the processor is configured to execute the instructions to implement the training method of the image generation model according to any one of claims 1 to 6, and / or implement the image generation method according to any one of claims 8 to 11.

15. A computer-readable storage medium, characterized in that, When the instructions in the computer-readable storage medium are executed by the processor of the electronic device, the electronic device is enabled to execute the training method of the image generation model according to any one of claims 1 to 6, and / or implement the image generation method according to any one of claims 8 to 11.

Images