Virtual hairstyle changing method based on diffusion model

Abstract

The invention discloses a virtual hairstyle changing method based on a diffusion model, and relates to the technical field of virtual hairstyle changing, a two-stage generation method is adopted, in the first stage, a rough hairstyle picture is generated, in the second stage, the rough hairstyle picture is refined through a redrawing mechanism, so that the generated hairstyle picture has higher detail fidelity, and the hairstyle changing efficiency is improved. In the second stage, a hair style redrawing diffusion model is adopted, so that the hair style can be better presented on complex textures and details, and the problems that in a traditional method, a generated hair style picture is prone to blurring or detail loss, and in a traditional GAN-based method, the hair style with complex textures is often difficult to process, artifacts are prone to being generated or details are ignored are solved; based on the scheme of a diffusion model, in combination with a cross attention mechanism, hairstyle details are better captured, and more complex hairstyle styles can be processed by injecting hairstyle reference information in the generation process and ensuring transmission of the hairstyle details, and relatively high stability is kept.

Description

Technical Field

[0001] The present invention relates to the technical field of virtual hair-changing technology, and in particular to a virtual hair-changing method based on a diffusion model. Background Art

[0002] With the development of science and technology, virtual reality technology has gradually matured, bringing more convenient virtual reality applications. Virtual try-on has always been a popular virtual reality application, and virtual hairstyle technology is a more eye-catching and challenging application in this field. Users can choose different hairstyles, colors and textures according to their preferences, and make personalized combinations and attempts. The application of this technology has greatly enriched the user's hairdressing experience, while also avoiding disappointment and economic losses caused by unsatisfactory haircut results. The focus of this task is to perfectly transfer the shape, color and texture of the target hairstyle to the user's head, while ensuring that the user's identity information and background do not change.

[0003] In recent years, most virtual hairstyle changing technologies are based on GAN methods. However, the problem that comes with it is that it cannot handle hairstyles with complex textures, and it is easy to ignore details or produce artifacts. Based on the above reasons, the present invention proposes a virtual hairstyle changing technology based on a diffusion model. Summary of the invention

[0004] In view of the above existing problems, the present invention is proposed.

[0005] Therefore, the present invention provides a virtual hairstyle changing method based on a diffusion model to solve the problems of being unable to handle hairstyles with complex textures, easily ignoring details, or generating artifacts.

[0006] In order to solve the above technical problems, the present invention provides the following technical solutions:

[0007] In one aspect, the present invention provides a virtual hair-changing system based on a diffusion model, comprising:

[0008] The first-stage image processing module processes the input source image and hairstyle reference image, and performs face detection, alignment, and image cropping operations;

[0009] The baldness generator is used to generate bald pictures of the people in the source pictures, that is, pictures without hair;

[0010] Hairstyle reference network, which extracts hairstyle details from hairstyle reference images. This network extracts hairstyle information through latent space encoding and cross-attention mechanism, and injects hairstyle information details into the hairstyle generation diffusion model;

[0011] The hairstyle generation diffusion model generates rough hairstyle images based on the bald pictures generated by the baldness generator and the hairstyle details provided by the hairstyle reference network to capture the overall outline of the hairstyle;

[0012] The second-stage image processing module processes the rough hairstyle image, including face detection, alignment and cropping operations, and also generates a hairstyle binary mask;

[0013] The hairstyle redrawing diffusion model redraws the rough hairstyle change picture generated in the first stage in a refined manner, combines the hairstyle features trained by hairstyle LoRA, redraws the hairstyle area, makes the hairstyle blend harmoniously with other parts, and outputs a refined hairstyle picture.

[0014] On the other hand, the present invention provides a virtual hair-changing method based on a diffusion model, comprising:

[0015] Step S1, during training, select pictures of the same person with different hairstyles, such as the SketchHairSalon dataset;

[0016] Step S2, randomly selecting one picture as a source picture and one picture as a hairstyle reference picture from pictures of the same person with different hairstyles, and then processing them into a desired format through an image processing module;

[0017] Step S3: source image I src Input into the baldness generator to get the baldness picture I bald ;

[0018] Step S4: The hairstyle reference image I ref Input into the hairstyle reference network to extract hairstyle related details;

[0019] Step S5: bald picture I bald The extracted hairstyle details are input into the hairstyle generation diffusion model through the cross-attention mechanism, and finally a rough hairstyle change picture is obtained;

[0020] Step S6, prepare 510 pictures of different angles (back side is not required) for each hairstyle as the LoRA hairstyle training set, input the pictures in the hairstyle training set into the image processing module, and obtain the hairstyle pictures and hairstyle description words;

[0021] Step S7, inputting the hairstyle image and hairstyle description words obtained in step S6 into the diffusion model to perform LoRA training of targeted hairstyles;

[0022] Step S8, generating a detailed hair style picture through the VAE decoder;

[0023] Step S1 to step S5 are the first-stage hairstyle generation steps, and step S6 to step S8 are the second-stage hairstyle generation steps;

[0024] Since the hairstyle images generated in one stage are relatively rough, a two-stage virtual hairstyle image generation method is adopted;

[0025] In the first stage, the trained hairstyle generation diffusion model is used to generate a rough hairstyle change picture. In the second stage, the hairstyle redrawing diffusion model is used to redraw the hairstyle area of ​​the rough hairstyle change picture to obtain a more refined hairstyle change picture.

[0026] The rough and fine here are relative.

[0027] Furthermore, in the first-stage image processing module:

[0028] Source Image I src , Hairstyle Reference Picture I ref Input to the image processing module, detect the face, align the face, crop and resize to the required size;

[0029] The baldness generator in one stage consists of a VAE encoder ε, a baldness generation diffusion model ∈ b , Bald ControlNetτ b , baldness generative diffusion model and a VAE decoder;

[0030] The baldness generation diffusion model is the latent space diffusion model StableDiffusion, which uses the pre-trained weights of SDv1.5. StableDiffusion is a classic latent space diffusion model (LDM, LatentDiffusionModel). LDM uses variational autoencoder VAE. Given an input image, the image is first mapped to the latent space through the encoder, so that diffusion is carried out in the latent space. In the forward process of diffusion, Gaussian noise will be gradually added to the latent variable in T iterations. The reverse diffusion process is to gradually remove noise, and finally it is restored to the RGB space through the decoder to obtain the final image.

[0031] The ControlNet branch creates a copy of the trainable diffusion model UNet encoding block and intermediate block, and adds an additional zero convolution layer. The encoding block and intermediate block of ControlNet are added to the decoding block and intermediate block of the diffusion model. The diffusion model here is the bald generation diffusion model, and the output of each replica block is added to the jump connection of the original diffusion model UNet.

[0032] Randomly sample 4-channel latent space Gaussian noise Input into the diffusion model of baldness generation;

[0033] Input the source image into the VAE encoder to get the latent space code z src =ε(Isrc ), encode the latent space z src Input to the Baldness ControlNet, i.e., the trainable copy of the decoding block and the intermediate block of the Baldness Generative Diffusion Model, after the zero convolution layer, the encoding block and the intermediate block of the Baldness ControlNet are added to the decoding block and the intermediate block of the Baldness Generative Diffusion Model through the residual connection, and the reference information of the source image is input to the Baldness Generative Diffusion Model;

[0034] The output of the final baldness generation diffusion model is passed through the VAE decoder to output the baldness image I of the person in the source image in RGB space. b ;

[0035] In the process of training the bald generator, only the bald ControlNet is trained. The weights of the bald generation diffusion model, VAE encoder and VAE decoder are fixed, and the pre-trained weights of SDv1.5 are used. The bald ControlNet participates in the training, and the initial weights also use the pre-trained weights of SDv1.5. By minimizing the loss function, the Adam optimizer is used for optimization iteration. The loss function for training the bald generator is: in, represents the loss function, Represents the expected value, that is, given the source image I src The latent space encoding, from the standard Gaussian distribution The sampled noise ∈, and the expectation at time step t, ∈ b represents the noise estimator of the baldness generative diffusion model, responsible for predicting the noise distribution of the input at each time step t, τ b (ε(I src ) represents the bald ControlNet model τ b The bald features of the generated input source image are processed by the VAE encoder ε on the source image I src Perform latent space encoding and process it at time step t, ∈ represents Gaussian noise, represents the L2 norm.

[0036] Furthermore, in the hairstyle reference network in the first stage,

[0037] Hairstyle Reference Network τ h It is obtained based on the latent space diffusion model StableDiffusion training:

[0038] The hairstyle reference picture I ref Input into the pre-trained VAE encoder ε, and get the latent space encoding z = ε (I ref );

[0039] Input the latent space code z into the hairstyle reference network τh In the above example, we extract the hairstyle detail feature c h , is input into the hairstyle generation diffusion model through the cross attention mechanism. The calculation method of the cross attention mechanism is:

[0040] Among them, Z″ represents the output result of the cross attention mechanism, Q = ZW q ,K′=c h W′ k ,V′=c h W′ v are the query, key, and value matrices of image features, respectively, and W q , W′ k and W′ v It is a trainable linear mapping layer. The query matrix uses the same query matrix as the self-attention mechanism. d represents the feature dimension and T represents the transposition operation.

[0041] Furthermore, in the first stage of the hairstyle generation diffusion model,

[0042] Hairstyle generation diffusion model ∈ b It is based on the pre-trained StableDiffusion model training:

[0043] The bald picture I generated by the bald generator b Input into the pre-trained VAE encoder to get the latent space code z b =ε(I ref );

[0044] Randomly generate Gaussian noise in latent space Add the obtained latent space code z b , is input into the hairstyle generation diffusion model, and the feature Z is obtained through the feature layer, and the attention is calculated through the self-attention mechanism:

[0045] Where Q = ZW q ,K=ZW k ,V=ZW v are the query, key, and value matrices of image features, respectively, and W k and W v It is a trainable linear mapping layer. The query matrix Q uses the same query matrix as the hairstyle reference network hairstyle cross attention mechanism. After the self-attention mechanism, it passes through the cross attention mechanism Z′ and adds the obtained attention: Z new =Z′+Z″, get new attention Z new , accurately transfer the hair to the area in the source image;

[0046] After the subsequent feature layers, the final output latent space feature layer is input into the pre-trained VAE decoder to obtain a rough hair-changing image in the RGB space.

[0047] Furthermore, during the two-stage hairstyle generation training process:

[0048] A small LoRA model is trained for each hairstyle: W = W 0 +ΔW=W 0 +BA, where W represents the fine-tuned weight matrix, ΔW represents the weight update amount, and the pre-trained weight is W 0 ∈R d×k , another low-rank matrix B∈R d×r ,A∈R r×k , and the rank r<<min(d,k), d, k represent the feature dimension and output dimension respectively, which reduces the training burden and enables the model to quickly capture the unique characteristics of each hairstyle, thereby significantly improving the quality of generated images while maintaining training efficiency.

[0049] Furthermore, in the two-stage image processing module:

[0050] Use the large model to label the input hairstyle, that is, to obtain the descriptive words of the hairstyle;

[0051] The open source image segmentation model SAM is used to segment the hair and obtain the hairstyle image I h .

[0052] Furthermore, in the two-stage LoRA training process, the diffusion model ∈ θ Use pre-trained SDv1.5 as pre-trained weights:

[0053] The hairstyle description words obtained by the two-stage image processing module are input into the pre-trained CLIP text encoder to obtain the text encoding c t , which is input into the diffusion model through the text cross-attention mechanism;

[0054] The hairstyle image obtained by the two-stage image processing module is input into the VAE encoder to obtain the latent space code z h =ε(I h ), plus randomly generated latent space Gaussian noise Input into the diffusion model and apply the weight of the hairstyle LoRA to the diffusion model;

[0055] The final feature layer output of the diffusion model passes through the VAE decoder to obtain a hairstyle image in RGB space;

[0056] During the training process, the Adam optimizer is used to perform optimization iterations by minimizing the loss function. The weight of the diffusion model is fixed, and only the hairstyle LoRA is trained. The loss function for training the hairstyle LoRA is:

[0057] in, is the loss function, represents the expected value, ∈ θ+Δθ represents the diffusion model of the hairstyle LoRA integration, and estimates the noise after weight adjustment θ+Δθ, Δθ represents the weight increment fine-tuned in the LoRA model training, z t represents the hairstyle image encoding result in the latent space, c t represents the encoding result of the text description word, and ∈ represents Gaussian noise.

[0058] Furthermore, in the two-stage image processing module:

[0059] In the inference phase, unlike the training phase, the diffusion model uses the pre-trained StableDiffusion redrawing model;

[0060] The goal of the redrawing diffusion model is to redraw the image required by the user in the user-specified area of ​​a picture, and ensure that the drawn area can be harmoniously integrated with the original image. The redrawing diffusion model is used in virtual hairstyle change, aiming to remove the hairstyle of the source picture and redraw the target hairstyle trained by hairstyle LoRA on the model, so that the other parts remain in the original state, and the generated hairstyle is more realistic and natural;

[0061] In the inference phase, the image processing module first performs face detection and face alignment on the input rough hairstyle image, and then crops it to the required size;

[0062] Then use the open source image segmentation method SAM to segment the hairstyle from the source image and the rough hairstyle image, and superimpose them to obtain the hairstyle binary mask I′ m , and then expand the hairstyle binary mask to obtain the hairstyle binary mask I after expansion m ;

[0063] The hairstyle binary mask I m Overlay with the rough hairstyle image to get the blocked hairstyle image I bg ;

[0064] Get the descriptive words for this hairstyle that meet the criteria during LoRA training.

[0065] Furthermore, in the two-stage hairstyle redrawing diffusion model:

[0066] Input the hairstyle description word into the CLIP text encoder to get the text encoding c t ;

[0067] Model image I that will cover the hairstyle bg Input into the VAE encoder ε to obtain the 4-channel redrawn background hidden code z bg =ε(I bg ), generating random Gaussian noise in latent space The 4-channel redraw background implicit code z bg , 4-channel latent space noise ∈ and 1-channel hairstyle binary mask image I obtained by the two-stage image processing module m , spliced ​​along the channel to obtain 9-channel input;

[0068] The 9-channel input is input into the hairstyle redrawing diffusion model integrated by the hairstyle LoRA, and after cyclic denoising, the feature layer output by the hairstyle redrawing diffusion model, and finally through the VAE decoder, a refined hairstyle change picture is obtained.

[0069] The beneficial effects of the present invention are:

[0070] The present invention adopts a two-stage generation method. The first stage generates a rough hairstyle image, and the second stage refines it through a redrawing mechanism, so that the generated hairstyle image has higher detail fidelity. The second stage adopts a hairstyle redrawing diffusion model, so that the hairstyle can be better presented in complex textures and details, avoiding the blurring or detail loss that occurs when the hairstyle image is generated in the traditional method.

[0071] In view of the problem that traditional GAN-based methods often have difficulty in processing hairstyles with complex textures and are prone to artifacts or ignoring details, the present invention adopts a diffusion model-based solution combined with a cross-attention mechanism to better capture hairstyle details. By injecting hairstyle reference information during the generation process, the transmission of hairstyle details is guaranteed, and more complex hairstyle styles can be processed while maintaining high stability.

[0072] The present invention introduces the LoRA model in the second stage to perform hairstyle fine-tuning, which reduces the number of training parameters and improves training efficiency. The lightweight fine-tuning method can quickly adapt to the generation requirements of different hairstyles, and does not require the entire model to be completely retrained, saving computing resources and time, and is suitable for the needs of personalized hairstyle generation. BRIEF DESCRIPTION OF THE DRAWINGS

[0073] In order to more clearly illustrate the technical solutions of the embodiments of the present invention, the accompanying drawings required for use in the description of the embodiments will be briefly introduced below. Obviously, the accompanying drawings described below are only some embodiments of the present invention. For ordinary technicians in this field, other accompanying drawings can be obtained based on these accompanying drawings without paying creative work.

[0074] Figure 1It is a schematic flow chart of the virtual hair-changing method of the present invention;

[0075] Figure 2 It is a schematic diagram of the first stage of the virtual hair-changing method of the present invention;

[0076] Figure 3 It is a two-stage flow chart of the virtual hair-changing method of the present invention;

[0077] Figure 4 A schematic diagram of the processing flow of an image processing module of the virtual hair-changing method of the present invention;

[0078] Figure 5 It is a schematic diagram of the two-stage hairstyle redrawing diffusion model processing flow of the virtual hairstyle changing method of the present invention. DETAILED DESCRIPTION

[0079] In order to make the above-mentioned objects, features and advantages of the present invention more obvious and easy to understand, the specific implementation methods of the present invention are described in detail below in conjunction with the accompanying drawings.

[0080] In the following description, many specific details are set forth to facilitate a full understanding of the present invention, but the present invention may also be implemented in other ways different from those described herein, and those skilled in the art may make similar generalizations without violating the connotation of the present invention. Therefore, the present invention is not limited to the specific embodiments disclosed below.

[0081] Secondly, the term "one embodiment" or "embodiment" as used herein refers to a specific feature, structure, or characteristic that may be included in at least one implementation of the present invention. The term "in one embodiment" that appears in different places in this specification does not necessarily refer to the same embodiment, nor does it refer to a separate or selective embodiment that is mutually exclusive with other embodiments.

[0082] Example 1, reference Figure 1 , this embodiment provides a virtual hair-changing system based on a diffusion model, including,

[0083] The first-stage image processing module processes the input source image and hairstyle reference image, and performs face detection, alignment, and image cropping operations;

[0084] The baldness generator is used to generate bald pictures of the people in the source pictures, that is, pictures without hair;

[0085] Hairstyle reference network, which extracts hairstyle details from hairstyle reference images. This network extracts hairstyle information through latent space encoding and cross-attention mechanism, and injects hairstyle information details into the hairstyle generation diffusion model;

[0086] The hairstyle generation diffusion model generates rough hairstyle images based on the bald pictures generated by the baldness generator and the hairstyle details provided by the hairstyle reference network to capture the overall outline of the hairstyle;

[0087] The second-stage image processing module processes the rough hairstyle image, including face detection, alignment and cropping operations, and also generates a hairstyle binary mask;

[0088] The hairstyle redrawing diffusion model redraws the rough hairstyle change picture generated in the first stage in a refined manner, combines the hairstyle features trained by hairstyle LoRA, redraws the hairstyle area, makes the hairstyle blend harmoniously with other parts, and outputs a refined hairstyle picture.

[0089] Example 2, reference Figure 1-Figure 5 This embodiment provides a virtual hair replacement method based on a diffusion model, comprising the following steps:

[0090] Step S1, during training, select pictures of the same person with different hairstyles, such as the SketchHairSalon dataset;

[0091] Step S2, randomly selecting one picture as a source picture and one picture as a hairstyle reference picture from pictures of the same person with different hairstyles, and then processing them into a desired format through an image processing module;

[0092] Step S3: source image I src Input into the baldness generator to get the baldness picture I bald ;

[0093] Step S4: The hairstyle reference image I ref Input into the hairstyle reference network to extract hairstyle related details;

[0094] Step S5: bald picture I bald The extracted hairstyle details are input into the hairstyle generation diffusion model through the cross-attention mechanism, and finally a rough hairstyle change picture is obtained;

[0095] Step S6, prepare 510 pictures of different angles (back side is not required) for each hairstyle as the LoRA hairstyle training set, input the pictures in the hairstyle training set into the image processing module, and obtain the hairstyle pictures and hairstyle description words;

[0096] Step S7, inputting the hairstyle image and hairstyle description words obtained in step S6 into the diffusion model to perform LoRA training of targeted hairstyles;

[0097] Step S8, generating a detailed hair style picture through the VAE decoder;

[0098] Step S1 to step S5 are the first-stage hairstyle generation steps, and step S6 to step S8 are the second-stage hairstyle generation steps;

[0099] Since the hairstyle images generated in one stage are relatively rough, a two-stage virtual hairstyle image generation method is adopted;

[0100] In the first stage, the trained hairstyle generation diffusion model is used to generate a rough hairstyle change picture. In the second stage, the hairstyle redrawing diffusion model is used to redraw the hairstyle area of ​​the rough hairstyle change picture to obtain a final more refined hairstyle change picture.

[0101] The coarseness and fineness here are relative;

[0102] In the first-stage image processing module:

[0103] Source Image I src , Hairstyle Reference Picture I ref Input to the image processing module, detect the face, align the face, crop and resize to the required size;

[0104] The baldness generator in one stage consists of a VAE encoder ε, a baldness generation diffusion model ∈ b , Bald ControlNetτ b , baldness generative diffusion model and a VAE decoder;

[0105] The baldness generation diffusion model is the latent space diffusion model StableDiffusion, which uses the pre-trained weights of SDv1.5. StableDiffusion is a classic latent space diffusion model (LDM, LatentDiffusionModel). LDM uses variational autoencoder VAE. Given an input image, the image is first mapped to the latent space through the encoder, so that diffusion is carried out in the latent space. In the forward process of diffusion, Gaussian noise will be gradually added to the latent variable in T iterations. The reverse diffusion process is to gradually remove noise, and finally it is restored to the RGB space through the decoder to obtain the final image.

[0106] The ControlNet branch creates a copy of the trainable diffusion model UNet encoding block and intermediate block, and adds an additional zero convolution layer. The encoding block and intermediate block of ControlNet are added to the decoding block and intermediate block of the diffusion model. The diffusion model here is the bald generation diffusion model, and the output of each replica block is added to the jump connection of the original diffusion model UNet.

[0107] Randomly sample 4-channel latent space Gaussian noise Input into the diffusion model of baldness generation;

[0108] Input the source image into the VAE encoder to get the latent space code z src =ε(I src ), encode the latent space z src Input to the Baldness ControlNet, i.e., the trainable copy of the decoding block and the intermediate block of the Baldness Generative Diffusion Model, after the zero convolution layer, the encoding block and the intermediate block of the Baldness ControlNet are added to the decoding block and the intermediate block of the Baldness Generative Diffusion Model through the residual connection, and the reference information of the source image is input to the Baldness Generative Diffusion Model;

[0109] The output of the final baldness generation diffusion model is passed through the VAE decoder to output the baldness image I of the person in the source image in RGB space. b ;

[0110] In the process of training the bald generator, only the bald ControlNet is trained. The weights of the bald generation diffusion model, VAE encoder and VAE decoder are fixed, and the pre-trained weights of SDv1.5 are used. The bald ControlNet participates in the training, and the initial weights also use the pre-trained weights of SDv1.5. By minimizing the loss function, the Adam optimizer is used for optimization iteration. The loss function for training the bald generator is: in, represents the loss function, Represents the expected value, that is, given the source image I src The latent space encoding, from the standard Gaussian distribution The sampled noise ∈, and the expectation at time step t, ∈ b represents the noise estimator of the baldness generative diffusion model, responsible for predicting the noise distribution of the input at each time step t, τ b (ε(I src ) represents the bald ControlNet model τ b The bald features of the generated input source image are processed by the VAE encoder ε on the source image I src Perform latent space encoding and process it at time step t, ∈ represents Gaussian noise, represents the L2 norm;

[0111] In the hairstyle reference network in the first stage,

[0112] Hairstyle Reference Network τ h It is obtained based on the latent space diffusion model StableDiffusion training:

[0113] The hairstyle reference picture I ref Input into the pre-trained VAE encoder ε, and get the latent space encoding z = ε (I ref );

[0114] Input the latent space code z into the hairstyle reference network τ h In the above example, we extract the hairstyle detail feature c h , is input into the hairstyle generation diffusion model through the cross attention mechanism. The calculation method of the cross attention mechanism is:

[0115] The same query matrix as the self-attention mechanism, d represents the feature dimension and T represents the transpose operation;

[0116] In the one-stage hair style generation diffusion model,

[0117] Hairstyle generation diffusion model ∈ b It is based on the pre-trained StableDiffusion model training:

[0118] The bald picture I generated by the bald generator b Input into the pre-trained VAE encoder to get the latent space code z b =ε(I ref );

[0119] Randomly generate Gaussian noise in latent space Add the obtained latent space code z b , is input into the hairstyle generation diffusion model, and the feature Z is obtained through the feature layer, and the attention is calculated through the self-attention mechanism:

[0120]

[0121] ZW v are the query, key, and value matrices of image features, respectively, and W k and W v It is a trainable linear mapping layer. The query matrix Q uses the same query matrix as the hairstyle reference network hairstyle cross attention mechanism. After the self-attention mechanism, it passes through the cross attention mechanism Z′ and adds the obtained attention: Z new =Z′+Z″, get new attention Z new , accurately transfer the hair to the area in the source image;

[0122] After the subsequent feature layers, the final output latent space feature layer is input into the pre-trained VAE decoder to obtain a rough hair-changing picture in the RGB space;

[0123] During the two-stage hairstyle generation training process:

[0124] A small LoRA model is trained for each hairstyle: W = W 0 +ΔW=W 0+BA, where W represents the fine-tuned weight matrix, ΔW represents the weight update amount, and the pre-trained weight is W 0 ∈R d×k , another low-rank matrix B∈R d×r ,A∈R r×k , and the rank r<<min(d,k), d, k represent the feature dimension and output dimension respectively, which reduces the training burden and enables the model to quickly capture the unique characteristics of each hairstyle, thereby significantly improving the quality of generated images while maintaining training efficiency;

[0125] In the two-stage image processing module:

[0126] Use the large model to label the input hairstyle, that is, to obtain the descriptive words of the hairstyle;

[0127] The open source image segmentation model SAM is used to segment the hair and obtain the hairstyle image I h ;

[0128] During the two-stage LoRA training process, the diffusion model ∈ θ Use pre-trained SDv1.5 as pre-trained weights:

[0129] The hairstyle description words obtained by the two-stage image processing module are input into the pre-trained CLIP text encoder to obtain the text encoding c t , which is input into the diffusion model through the text cross-attention mechanism;

[0130] The hairstyle image obtained by the two-stage image processing module is input into the VAE encoder to obtain the latent space code z h =ε(I h ), plus randomly generated latent space Gaussian noise Input into the diffusion model and apply the weight of the hairstyle LoRA to the diffusion model;

[0131] The final feature layer output of the diffusion model passes through the VAE decoder to obtain a hairstyle image in RGB space;

[0132] During the training process, the Adam optimizer is used to perform optimization iterations by minimizing the loss function. The weight of the diffusion model is fixed, and only the hairstyle LoRA is trained. The loss function for training the hairstyle LoRA is:

[0133] in, is the loss function, represents the expected value, ∈ θ+Δθ represents the diffusion model of the hairstyle LoRA integration, and estimates the noise after weight adjustment θ+Δθ, Δθ represents the weight increment fine-tuned in the LoRA model training, z trepresents the hairstyle image encoding result in the latent space, c t represents the encoding result of the text description word, ∈ represents Gaussian noise;

[0134] In the two-stage image processing module:

[0135] In the inference phase, unlike the training phase, the diffusion model uses the pre-trained StableDiffusion redrawing model;

[0136] The goal of the redrawing diffusion model is to redraw the image required by the user in the user-specified area of ​​a picture, and ensure that the drawn area can be harmoniously integrated with the original image. The redrawing diffusion model is used in virtual hairstyle change, aiming to remove the hairstyle of the source picture and redraw the target hairstyle trained by hairstyle LoRA on the model, so that the other parts remain in the original state, and the generated hairstyle is more realistic and natural;

[0137] In the inference phase, the image processing module first performs face detection and face alignment on the input rough hairstyle image, and then crops it to the required size;

[0138] Then use the open source image segmentation method SAM to segment the hairstyle from the source image and the rough hairstyle image, and superimpose them to obtain the hairstyle binary mask I′ m , and then expand the hairstyle binary mask to obtain the hairstyle binary mask I after expansion m ;

[0139] The hairstyle binary mask I m Overlay with the rough hairstyle image to get the blocked hairstyle image I bg ;

[0140] Get the descriptive words for this hairstyle that meet the criteria during LoRA training;

[0141] In the two-stage hair repainting diffusion model:

[0142] Input the hairstyle description word into the CLIP text encoder to get the text encoding c t ;

[0143] Model image I that will cover the hairstyle bg Input into the VAE encoder ε to obtain the 4-channel redrawn background hidden code z bg =ε(I bg ), generating random Gaussian noise in latent space The 4-channel redraw background implicit code z bg , 4-channel latent space noise ∈ and 1-channel hairstyle binary mask image I obtained by the two-stage image processing module m , spliced ​​along the channel to obtain 9-channel input;

[0144] The 9-channel input is input into the hairstyle redrawing diffusion model integrated by the hairstyle LoRA, and after cyclic denoising, the feature layer output by the hairstyle redrawing diffusion model, and finally through the VAE decoder, a refined hairstyle change picture is obtained.

[0145] It should be noted that the above embodiments are only used to illustrate the technical solutions of the present invention rather than to limit it. Although the present invention has been described in detail with reference to the preferred embodiments, those skilled in the art should understand that the technical solutions of the present invention may be modified or replaced by equivalents without departing from the spirit and scope of the technical solutions of the present invention, which should all be included in the scope of the claims of the present invention.

Claims

1. A virtual hair-changing system based on a diffusion model, characterized in that: include, The first-stage image processing module processes the input source image and hairstyle reference image, and performs face detection, alignment, and image cropping operations; The baldness generator is used to generate bald pictures of the people in the source pictures, that is, pictures without hair; Hairstyle reference network, which extracts hairstyle details from hairstyle reference images, extracts hairstyle information through latent space encoding and cross-attention mechanism, and injects hairstyle information details into the hairstyle generation diffusion model; The hairstyle generation diffusion model generates rough hairstyle change pictures based on the bald pictures generated by the baldness generator and the hairstyle details provided by the hairstyle reference network; The two-stage image processing module processes the rough hairstyle image, including face detection, alignment and cropping operations, and also generates a hairstyle binary mask; The hairstyle redrawing diffusion model redraws the rough hairstyle change picture generated in the first stage in a refined manner, and redraws the hairstyle area based on the hairstyle features trained by Hairstyle LoRA.

2. A virtual hair-changing method based on a diffusion model, characterized in that: Based on the virtual hair-changing system based on diffusion model as described in claim 1, The following steps are included: Step S1, selecting pictures of the same person with different hairstyles during training; Step S2, randomly selecting one picture as a source picture and one picture as a hairstyle reference picture from pictures of the same person with different hairstyles, and then processing them into a desired format through an image processing module; Step S3: source image I src Input into the baldness generator to get the baldness picture I bald ; Step S4: The hairstyle reference image I ref Input into the hairstyle reference network to extract hairstyle related details; Step S5: bald picture I bald The extracted hairstyle details are input into the hairstyle generation diffusion model through the cross-attention mechanism, and finally a rough hairstyle change picture is obtained; Step S6, preparing 510 pictures of each hairstyle at different angles as the LoRA hairstyle training set, inputting the pictures in the hairstyle training set into the image processing module to obtain the hairstyle pictures and hairstyle description words; Step S7, inputting the hairstyle image and hairstyle description words obtained in step S6 into the diffusion model to perform LoRA training of targeted hairstyles; Step S8, generating a detailed hair style picture through the VAE decoder; Step S1 to step S5 are the first-stage hairstyle generation steps, and step S6 to step S8 are the second-stage hairstyle generation steps.

3. A virtual hair change method based on a diffusion model according to claim 2, characterized in that: In the first-stage image processing module: Source Image I src , Hairstyle Reference Picture I ref Input to the image processing module, detect the face, align the face, crop and resize to the required size; The baldness generator in one stage consists of a VAE encoder ε, a baldness generation diffusion model ∈ b , Bald ControlNetτ b , baldness generative diffusion model and a VAE decoder; The baldness generation diffusion model is the latent space diffusion model StableDiffusion, which uses the pre-trained weights of SDv1.

5. The LDM uses the variational autoencoder VAE. Given an input image, the encoder first maps the image to the latent space so that the diffusion is carried out in the latent space. In the forward process of diffusion, Gaussian noise will be gradually added to the latent variable in T iterations. The reverse diffusion process is to gradually remove noise, and finally the decoder is used to restore it to the RGB space to obtain the final image. The ControlNet branch creates a copy of the trainable diffusion model UNet encoding block and intermediate block, and adds an additional zero convolution layer. The encoding block and intermediate block of ControlNet are added to the decoding block and intermediate block of the diffusion model. The diffusion model here is the bald generation diffusion model, and the output of each replica block is added to the jump connection of the original diffusion model UNet. Randomly sample 4-channel latent space Gaussian noise Input into the diffusion model of baldness generation; Input the source image into the VAE encoder to get the latent space code z src =ε(I src ), encode the latent space z src Input to the Baldness ControlNet, i.e., the trainable copy of the decoding block and the intermediate block of the Baldness Generative Diffusion Model, after the zero convolution layer, the encoding block and the intermediate block of the Baldness ControlNet are added to the decoding block and the intermediate block of the Baldness Generative Diffusion Model through the residual connection, and the reference information of the source image is input to the Baldness Generative Diffusion Model; The output of the final baldness generation diffusion model is passed through the VAE decoder to output the baldness image I of the person in the source image in RGB space. b ; In the process of training the bald generator, only the bald ControlNet is trained. The weights of the bald generation diffusion model, VAE encoder and VAE decoder are fixed, and the pre-trained weights of SDv1.5 are used. The bald ControlNet participates in the training, and the initial weights also use the pre-trained weights of SDv1.

5. By minimizing the loss function, the Adam optimizer is used for optimization iteration. The loss function for training the bald generator is: in, represents the loss function, Represents the expected value, that is, given the source image I src The latent space encoding, from the standard Gaussian distribution The sampled noise ∈, and the expectation at time step t, ∈ b represents the noise estimator of the baldness generative diffusion model, responsible for predicting the noise distribution of the input at each time step t, τ b (ε(I src ) represents the bald ControlNet model τ b The bald features of the generated input source image are processed by the VAE encoder ε on the source image I src Perform latent space encoding and process it at time step t, ∈ represents Gaussian noise, represents the L2 norm.

4. The virtual hair-changing method based on diffusion model according to claim 3, characterized in that: In the hairstyle reference network in the first stage, Hairstyle Reference Network τ h It is obtained based on the latent space diffusion model StableDiffusion training: The hairstyle reference picture I ref Input into the pre-trained VAE encoder ε, and get the latent space encoding z = ε (I ref ); Input the latent space code z into the hairstyle reference network τ h In the above example, we extract the hairstyle detail feature c h , is input into the hairstyle generation diffusion model through the cross attention mechanism. The calculation method of the cross attention mechanism is: Among them, Z″ represents the output result of the cross attention mechanism, Q = ZW q ,K ′ =c h W k ′ ,V ′ =c h W v ′ are the query, key, and value matrices of image features, respectively, and W q , W k ′ and W v ′ It is a trainable linear mapping layer. The query matrix uses the same query matrix as the self-attention mechanism. d represents the feature dimension and T represents the transposition operation.

5. A virtual hair change method based on a diffusion model according to claim 4, characterized in that: In the one-stage hair style generation diffusion model, Hairstyle generation diffusion model∈ b It is based on the pre-trained StableDiffusion model training: The bald picture I generated by the bald generator b Input into the pre-trained VAE encoder to get the latent space code z b =ε(I ref ); Randomly generate Gaussian noise in latent space Add the obtained latent space code z b , is input into the hairstyle generation diffusion model, and the feature Z is obtained through the feature layer, and the attention is calculated through the self-attention mechanism: Where Q = ZW q ,K=ZW k ,V=ZW v are the query, key, and value matrices of image features, respectively, and W k and W v It is a trainable linear mapping layer. The query matrix Q uses the same query matrix as the hairstyle reference network hairstyle cross attention mechanism. After the self-attention mechanism, it passes through the cross attention mechanism Z ′ , add up the attention obtained: Z new =Z ′ +Z″, get new attention Z new After the subsequent feature layers, the final output latent space feature layer is input into the pre-trained VAE decoder to obtain a rough hair-changing image in the RGB space.

6. A virtual hair change method based on a diffusion model according to claim 5, characterized in that: During the two-stage hairstyle generation training process: A small LoRA model is trained for each hairstyle: W = W0 + ΔW = W0 + BA, where W represents the fine-tuned weight matrix, ΔW represents the weight update amount, and the pre-trained weight is W0∈R d×k , another low-rank matrix B∈R d×r ,A∈R r×k , and the rank r<<min(d,k), d, k represent the feature dimension and output dimension respectively.

7. A virtual hair change method based on a diffusion model according to claim 6, characterized in that: In the two-stage image processing module: Use the large model to label the input hairstyle, that is, to obtain the descriptive words of the hairstyle; The open source image segmentation model SAM is used to segment the hair and obtain the hairstyle image I h .

8. The virtual hair-changing method based on diffusion model according to claim 7, characterized in that: In the two-stage LoRA training process, the diffusion model ∈ θ Use pre-trained SDv1.5 as pre-trained weights: The hairstyle description words obtained by the two-stage image processing module are input into the pre-trained CLIP text encoder to obtain the text encoding c t , which is input into the diffusion model through the text cross-attention mechanism; The hairstyle image obtained by the two-stage image processing module is input into the VAE encoder to obtain the latent space code z h =ε(I h ), plus randomly generated latent space Gaussian noise Input into the diffusion model and apply the weight of the hairstyle LoRA to the diffusion model; The last feature layer output of the diffusion model passes through the VAE decoder to obtain the hairstyle image in RGB space; During the training process, the Adam optimizer is used to perform optimization iterations by minimizing the loss function. The weight of the diffusion model is fixed, and only the hairstyle LoRA is trained. The loss function for training the hairstyle LoRA is: in, is the loss function, represents the expected value, ∈ θ+Δθ represents the diffusion model of the hairstyle LoRA integration, and estimates the noise after weight adjustment θ+Δθ, Δθ represents the weight increment fine-tuned in the LoRA model training, z t represents the hairstyle image encoding result in the latent space, c t represents the encoding result of the text description word, and ∈ represents Gaussian noise.

9. A virtual hair change method based on a diffusion model according to claim 8, characterized in that: In the two-stage image processing module: In the inference phase, the image processing module first performs face detection and face alignment on the input rough hairstyle image, and then crops it to the required size; Then use the open source image segmentation method SAM to segment the hairstyle from the source image and the rough hairstyle image, and superimpose them to obtain the hairstyle binary mask I ′ m , and then expand the hairstyle binary mask to obtain the hairstyle binary mask I after expansion m ; The hairstyle binary mask I m Overlay with the rough hairstyle image to get the blocked hairstyle image I bg ; Get the descriptive words for this hairstyle that meet the criteria during LoRA training.

10. A virtual hair change method based on a diffusion model according to claim 9, characterized in that: In the two-stage hair repainting diffusion model: Input the hairstyle description word into the CLIP text encoder to get the text encoding c t ; Model image I that will cover the hairstyle bg Input into the VAE encoder ε to obtain the 4-channel redrawn background hidden code z bg =ε(I bg ), generating random Gaussian noise in latent space The 4-channel redraw background implicit code z bg , 4-channel latent space noise ∈ and 1-channel hairstyle binary mask image I obtained by the two-stage image processing module m , spliced ​​along the channel to obtain 9-channel input; The 9-channel input is fed into the hairstyle redrawing diffusion model integrated by the hairstyle LoRA, and after cyclic denoising, the feature layer output by the hairstyle redrawing diffusion model, and finally through the VAE decoder, a refined hairstyle change picture is obtained.

Images