Image2StyleGAN++: How to Edit the Embedded Images?

In plain words

A GAN can draw very realistic faces and objects, but it cannot draw a given photo. To edit a real photo with one, you first have to find the input that makes the GAN produce that photo. Earlier work found such an input fairly well, but the copy was not exact: fine detail was lost, so the result was not good enough for editing.

This paper improves that search in three ways. It adds a second set of adjustable values, which act like a fine grain of noise, so the reconstruction can recover small detail. It lets the search cover only part of the picture, so a region can be left blank and filled in, or nudged towards a rough scribble. It also lets you copy the GAN's internal intermediate results from one image into another, which gives local edits that blend into the rest of the photo.

The results are sharper copies of photos and edits that stay local. Typical edits include filling in missing parts, changing a region's style, and swapping features between two faces.

TL;DR

• Find the latent code that reproduces a photo in a pretrained StyleGAN, first in its W+ space and then in per-layer noise maps §3 §4.2.
• Noise is optimised second, with W+ held fixed. Joint optimisation gives artifacts §4.2 fig. 2.
• Masks restrict the embedding to a region; activation-tensor copying then mixes two embeddings §4 §5.
• Author-reported PSNR for reconstruction is 39-45 dB, against about 20 dB for the previous best §1 §7.1 fig. 4. Inpainting beats partial and gated convolution in SSIM, MSE and PSNR at 256, 512 and 1024 pixels tab. 1.

Problem

Embedding a real image into a GAN latent space is needed for editing, but embeddings made by optimising only the W+ code reconstruct images with limited fidelity §1. The reported PSNR of the previous best embedding is about 20 dB §1. Local edits need an embedding that can leave a region unconstrained, or that matches only approximate colours, and a plain global code cannot do this §1 §3.

The method assumes a pretrained StyleGAN generator whose weights stay fixed §5. It uses StyleGAN's per-layer noise inputs as the extra variables §4 §6. The paper's examples are mostly faces, with some LSUN cars and bedrooms §9.4.

Method

The method has three parts. Section numbers below follow the LaTeX source.

1. Embedding (§4). The target image is reconstructed by optimising two groups of variables: a W+ code w, which carries semantic content, and noise maps n, which carry high-frequency detail §4. The loss has four terms (eq. 1): a style term, two masked pixel MSE terms and a perceptual term §4.1. The perceptual term uses VGG-16 features §4.1.

2. Optimisation order (§4.2). Optimising w and n jointly gives high-frequency artifacts, because the perceptual loss does not suit noise maps §4.2 [fig. 2]. Optimising n alone with pixel MSE gives a near-perfect reconstruction that is not useful for editing §4.2. The recommended order is to optimise w with n fixed, then n with w fixed, and to run each step once §4.2 fig. 3.

3. Building blocks (§6). The paper names four operations and gives their loss weights and optimisers. They are called from the algorithms in §7.

• Masked W+ optimisation W_l (eq. 2): optimises w, with n fixed.
• Masked noise optimisation Mk_n (eq. 3): optimises n, with w fixed.
• Masked style transfer M_st (eq. 4): optimises w against a style image.
• Masked activation operation I_att (eq. 5): mixes two intermediate activation tensors under two masks §5.

4. Activation tensors (§5). For layer l, the paper defines A_l as the activation of the image being embedded §5. It considers three operations: spatial copying, which replaces single spatial positions; averaging; and channel-wise copying. Spatial copying worked best in the authors' tests §5.

5. Applications (§7). Each application chains the operations above:

• Reconstruction: W+ optimisation, then noise optimisation §7.1 alg. 2.
• Image crossover: two masked W+ optimisations with blurred masks, then noise §7.2 alg. 3.
• Inpainting: masked W+ over selected layers, then noise on a larger blurred mask §7.3 alg. 4.
• Local edits from scribbles: W+ restricted to the first few layers, then noise §7.4 alg. 5.
• Local style transfer: masked W+ with a style term, then noise §7.5 alg. 6.
• Attribute-level transfer: activation mixing at layer 4 §7.6 eq. 5.

Generic embedding (alg. 1 §4). Start from initial values. Repeat a gradient step on the loss until a convergence test passes. The paper does not define the test.

# Objective with four terms  [eq. 1] [§4.1]
def total_loss(generated, x, y, masks, weights):
    # Style term: VGG-16 conv3_3 features against style image y  [eq. 1]
    loss = weights.style * style_loss(masks.style, generated, y)
    # Masked pixel MSE against target x  [eq. 1]
    # N is not defined in the paper; see Pitfalls
    loss += weights.mse1 / N * squared_norm(masks.mse * (generated - x))
    # Masked pixel MSE against y outside the mask  [eq. 1]
    loss += weights.mse2 / N * squared_norm((1 - masks.mse) * (generated - y))
    # Perceptual term on VGG-16 layers conv1_1 to conv3_3  [eq. 1] [§4.1]
    loss += weights.percept * perceptual_loss(masks.percept, generated, x)
    return loss

# Sequential W+ then noise embedding, as in reconstruction  [alg. 2] [§4.2]
# Input: target image; pretrained StyleGAN generator G
# Output: W+ code w and noise maps n

def embed_image(target):
    # Initial W+ code: mean face code or random U[-1, 1]  [§7.1]
    w_init = initial_w_code()
    # Initial noise maps, sampled from N(0, I)  [§7.1]
    n_init = standard_normal_noise()
    # Stage 1: optimise w only, with n held at its start  [§7.1] [§4.2]
    w_out = optimise_w(target, w_init, n_init, steps=5000)
    # Stage 2: optimise n only, with w held at w_out  [§7.1] [§4.2]
    n_out = optimise_noise(target, w_out, n_init, steps=3000)
    return w_out, n_out

# Masked activation mixing at layer l  [eq. 5] [§5]
def mix_activations(acts_a, acts_b, mask_a, mask_b):
    # Resize both masks to the layer's H x W by nearest neighbour  [eq. 5]
    small_a = resize_nearest(mask_a, acts_a.shape)
    small_b = resize_nearest(mask_b, acts_b.shape)
    # Keep A where mask_a is 1, and add B where mask_b is 0  [eq. 5]
    # Shapes: (H, W, C)
    return small_a * acts_a + (1 - small_b) * acts_b

Equations, as printed. The symbols follow the paper: x and y are the target images, M_s, M_m and M_p are the spatial masks, G is the generator, and λ's are the loss weights §4.1.

L = \lambda_s L_{style}(M_s, G(w,n), y) + \frac{\lambda_{mse_1}}{N}\|M_m \odot (G(w,n) - x)\|_2^2 + \frac{\lambda_{mse_2}}{N}\|(1 - M_m) \odot (G(w,n) - y)\|_2^2 + \lambda_p L_{percept}(M_p, G(w,n), x) eq. 1

W_l(M_p, M_m, w_m, w_{ini}, n_{ini}, x) = \argmin_{w_m} \lambda_p L_{percept}(M_p, G(w,n), x) + \frac{\lambda_{mse_1}}{N}\|M_m \odot (G(w,n) - x)\|_2^2 eq. 2

Mk_n(M, w_{ini}, n_{ini}, x, y) = \argmin_n \frac{\lambda_{mse_2}}{N}\|M_m \odot (G(w,n) - x)\|_2^2 + \frac{\lambda_{mse_1}}{N}\|(1 - M_m) \odot (G(w,n) - y)\|_2^2 eq. 3

M_{st}(M_s, w_{ini}, n_{ini}, y) = \argmin_w \lambda_s L_{style}(M_s, G(w,n), y) eq. 4

I_{att}(M_1, M_2, w, n_{ini}, l) = G(w, n, M_1 \odot A_l^{I_1} + (1 - M_2) \odot B_l^{I_2}) eq. 5

Algorithms 2 to 6 in §7 combine these operations, as the paper writes them:

• Crossover alg. 3: W_l with M_blur on I_1 plus W_l with 1 - M_blur on I_2, then Mk_n with M_blur.
• Inpainting alg. 4: W_l with 1 - M on the defective image, then Mk_n with 1 - M_blur+ against G(w_out).
• Scribbles alg. 5: W_l on the scribble image, plus a term λ||w* - w_out||, then Mk_n with M_blur against G(w_out).
• Style transfer alg. 6: W_l with M_blur on I_1, plus M_st with 1 - M_blur on I_2, then Mk_n with M_blur against G(w_out).

Hyperparameters

Results

Author-reported. Inpainting is compared with partial convolution (Liu et al.) and gated convolution (Yu et al.) at 1024, 512 and 256 px §9.1. Edit realism is scored with a face-attribute classifier §9.3. The paper reports no significance tests and no error bars.

The paper names the non-masked interpolation as an upper bound on the perceptual length for each model §9.4. Masked results on FFHQ are below that bound. The LSUN results are higher, which the authors attribute to translation in the data §9.4.

Paper vs code

No official code.

Searched: the arXiv abstract page, whose comments field and code section name no repository (only a CVPR 2020 note and a video link); a project page at rameenabdal.github.io/Image2StyleGAN++, which returned HTTP 404; a GitHub repository search for "Image2StyleGAN++", which returned only third-party reimplementations, not used; and the first author's GitHub repositories, which include none for this paper. The last author's GitHub user page returned HTTP 404. No repository was cloned.

Pitfalls

• Joint optimisation fails. Optimising w and n together gives high-frequency artifacts, and the perceptual loss does not suit noise maps §4.2 fig. 2. Optimise w first, then n, and run each once. A second pass over w lowered PSNR from 39.5 dB to 20 dB in the authors' example §4.2 fig. 3.
• Weight names swap between eq. 1 and eq. 3. Eq. 1 pairs λ_mse1 with the masked term against x, while eq. 3 pairs λ_mse2 with that term. The inpainting setting λ_mse2 = 1e-4 §7.3 follows eq. 3. Use eq. 3's convention for the noise step. Both weights are 1e-5 in §6, so only the inpainting case matters.
• N is not defined. The MSE terms divide by N, which the paper never defines. It is probably the pixel count of the image, but that is an inference, not stated.
• Convergence is undefined. Algorithm 1 loops "until converged", and no test is given. The applications state iteration counts, so use those.
• Optimiser for each step is not fixed. §6 offers Adam (0.01) or gradient descent (0.8) for the W+ step. §7 says Adam unless specified. The reconstruction step is not specified, so Adam at 0.01 is the reading by default. The noise step uses Adam at learning rate 5, which the paper itself calls very high §6.
• Summed results in algorithms 3, 5 and 6. Algorithms 3 and 6 write the W+ output as a sum of two optimisation results, and algorithm 5 adds a regulariser. The text does not say whether the results are summed as latent vectors or the two losses are optimised together. The scribble regulariser λ is not stated.
• Eq. 5 is not an average. Section 9.5 uses masks (1, 0) and calls the result a feature average. With eq. 5 that gives A + B, not (A + B) / 2, and §5 defines averaging separately as λA + (1 - λ)B. Treat §9.5 as unclear.
• Argmin variable clashes with the mask symbol. Eq. 2 takes the argmin over w_m, but §6 defines w_m as a binary mask over W+ layers. Read the argmin as a variable over w, restricted by that mask.
• Initial code for inpainting is stated twice. §7.3 sets the initial code to the mean face code, and the discussion of different initialisations uses an offset sampled from U[-0.4, 0.4]. The paper does not say which was used for the figures.
• Blur and resampling are not specified. The Gaussian kernel size for blurred masks, the interpolation of the perceptual masks M_p, and the blur used to build M_blur+ are not stated.
• Section numbers. Citations follow the numbered sections in the LaTeX source. The "Additional Results" section comes after the references and is §9.
• Compute. Reconstruction runs 5000 W+ steps and 3000 noise steps §7.1, and the other applications run 200 to 1000 steps per stage §7.2 §7.3 §7.4 §7.5. The paper gives no wall-clock time, and the card does not estimate one.
• Comparisons use different output resolutions. Partial and gated convolution are compared at 1024, 512 and 256 px by bilinear interpolation of their outputs §9.1. The gated convolution model is a 256 px model, which affects its score at other resolutions §9.1.
• Evaluation is small and face-heavy. Inpainting uses 10 ground-truth images and 7 masks, and edit realism uses one classifier §9.1 §9.3. The paper's own limitations: lighting, age and skin colour are not learned well §9.1.
• Edits can look smooth. The authors note that the edited region is sometimes overly smooth §9.3.
• Perceptual length is not defined here. The paper cites StyleGAN's metric and does not restate its formula §9.4.

Lineage

• Builds on: Image2StyleGAN (Abdal et al., 2019), which embeds into W+ only §1 §2. Activation-tensor manipulation builds on GAN dissection (Bau et al.) and on Frühstück et al.'s work on blending §2. The perceptual and style losses cite Gatys et al. and other VGG-feature work §4.1. No card exists yet for Image2StyleGAN.
• Compared against: partial convolution (Liu et al., 2018) and gated convolution (Yu et al., 2018) for inpainting §9.1 tab. 1.
• Superseded by: not stated.