StyleDiffusion: Controllable Disentangled Style Transfer via Diffusion Models

In plain words

Style transfer takes the look of one picture (brushstrokes, colours) and applies it to another picture's content. The hard part is first separating "what is shown" from "how it is painted". This paper does that in two steps. First, a diffusion model, the kind of model that learns to undo added noise, is used to gently wash the painterly style out of the style picture, leaving a plain version of its content. Second, the difference between the style picture and that plain version is taken as a description of the style. The difference is measured in the embedding space of CLIP, an image-text model.

The paper then fine-tunes a diffusion model for each style, so that it turns ordinary photos into that style. The training asks that the change made to a photo points the same way as the change from the plain content to the real style picture, and that the photo's content is not lost. Because the content and the style are kept apart on purpose, the user can change how much style is removed or applied. The results are judged by the authors to look better than earlier methods, but they are not the strongest on every measure.

TL;DR

StyleDiffusion defines style as whatever a style image has beyond its content. It removes style with DDIM-based diffusion, measures the leftover style with the CLIP image encoder, and fine-tunes a diffusion model per style with a direction-plus-L1 loss and a reconstruction loss. It needs about 50 ImageNet content images and roughly 23 minutes per style on an A100. The authors report the best SSIM and CLIP score among 11 methods, but not the best style loss. There is no official code.

Problem

Style transfer methods based on Gram matrices (Gatys-type losses) or on GANs tie content and style together, so the two cannot be separated or tuned on their own §1. The authors argue that the content of an image is easier to define than its style, so the content should be extracted explicitly and the style learned as its complement §1. Making that extraction controllable and consistent across images needs a model that can remove style without losing structure §1 §4.1.

The method assumes the style image is a single reference image, and that about 50 content images are enough to fine-tune for one style §1 §5.1. Speed is not a goal: the authors note that inference is slower than the feed-forward baselines §5.2.

Method

Overview. Given a style image I_s and a content image I_c, the method extracts the content of both with a diffusion-based style removal step. It then fine-tunes a diffusion-based style transfer model so that the stylized content of I_c matches the style of I_s in CLIP space §4 fig. 1.

Step 1, style removal §4.1 alg. 1 app. C.

1. Convert the input to grayscale with the ITU-R 601-2 luma transform. Colour is treated as part of the style §4.1.
2. Run a pretrained diffusion model (ADM, trained on ImageNet §5.1) with deterministic DDIM steps. Push the image forward to a fixed return step T_{remov} (the ODE form of DDIM inversion, eq. 9), then reverse it with the DDIM update eq. 8.
3. Repeat the forward-then-reverse pass K_r times. The paper sets K_r = 5 app. C.
4. The output is the content I^c with style dispelled. Larger T_{remov} removes more style and keeps the main structures §4.1 fig. 4.

For content images that are photographs, the paper says the removal step is optional, because they are already in the pretrained domain. Its footnote says that, unless stated otherwise, the content images are not passed through the diffusion process §4.1. This footnote and the method text disagree (see Pitfalls).

Step 2, style transfer module §4.2. The stylized result I_{cs} is produced by fine-tuning the reverse process of the same diffusion model. The content I_c^c is first encoded to a latent by the forward DDIM process with return step T_{trans}. Sampling the reverse process from that latent gives the stylized result app. D alg. 2. At inference, a stochastic DDPM forward process can be used to get diverse results §4.2 app. G.

Step 3, CLIP-based style disentanglement loss §4.3. Let E be the CLIP image encoder, used off the shelf as a fixed projector. The style vector of the style image is D_s = E(I_s) - E(I_s^c) eq. 11. The style vector of the output is D_{cs} = E(I_{cs}) - E(I_c^c) eq. 12. Two terms are used:

• an L1 term, \mathcal{L}_{SD}^{L1} = \|D_{cs} - D_s\| eq. 12, weight \lambda_{L1} = 10;
• a direction term, \mathcal{L}_{SD}^{dir} = 1 - \cos(D_{cs}, D_s) eq. 13, weight \lambda_{dir} = 1.

The total is \mathcal{L}_{SD} = \lambda_{L1}\mathcal{L}_{SD}^{L1} + \lambda_{dir}\mathcal{L}_{SD}^{dir} eq. 14 §4.3. The paper says the direction term is needed because L1 alone can leave the output outside the target style §4.3 fig. 2.

Step 4, style reconstruction prior §4.3. A second loss, \mathcal{L}_{SR} = \|I_{ss} - I_s\| eq. 15, asks the transfer model to rebuild the style image from the style content I_s^c. The paper optimises it before \mathcal{L}_{SD}, to initialise the model §4.3. Algorithm 2 interleaves the two (see Pitfalls). No weight for \mathcal{L}_{SR} is stated eq. 15.

Training procedure alg. 2 app. D.

# StyleDiffusion fine-tuning  [alg. 2] [app. D]
# Inputs: pretrained model; 50 content images; style content; style image
# Output: fine-tuned model for one style

def finetune(model, contents, style_content, style):
    # 1. Encode each content to a latent: DDIM forward to T_trans = 301,
    #    40 steps  [§5.1] [app. D]
    content_latents = [
        ddim_forward(c, model, t_end=301, steps=40) for c in contents
    ]
    # 2. Encode the style content the same way  [app. D]
    style_latent = ddim_forward(style_content, model, t_end=301, steps=40)
    tuned = copy(model)
    # 3. Repeat the whole training block K = 5 times  [app. D]
    for epoch in range(5):
        # 3a. Style reconstruction: K_s = 50 repeats on the style latent
        #     (6 reverse steps each)  [app. D] [eq. 15]
        for _ in range(50):
            rebuilt = ddim_reverse(style_latent, tuned, steps=6)
            step(tuned, style_recon_loss(rebuilt, style))
        # 3b. Style disentanglement: one update per content latent
        #     (6 reverse steps each)  [app. D] [eq. 14]
        for latent, content in zip(content_latents, contents):
            stylized = ddim_reverse(latent, tuned, steps=6)
            step(tuned, style_disent_loss(content, stylized,
                                          style_content, style))
    return tuned

CLIP style disentanglement loss. The function below restates eq. 11–eq. 14. The L1 norm and the cosine are over the CLIP embedding; the paper's \|\cdot\| is not specified further.

# CLIP-based style disentanglement loss  [eq. 11]-[eq. 14]  [§4.3]
# Inputs: content = I_c^c, stylized = I_cs, style_content = I_s^c, style = I_s

def style_disent_loss(content, stylized, style_content, style):
    # 1. Style vector of the style image: CLIP(style) - CLIP(style content)
    #    [eq. 11]
    style_vec = clip_image(style) - clip_image(style_content)
    # 2. Style vector of the output: CLIP(stylized) - CLIP(content)  [eq. 12]
    output_vec = clip_image(stylized) - clip_image(content)
    # 3. L1 term, weight 10  [eq. 12] [§4.3]
    l1_term = l1_norm(output_vec - style_vec)
    # 4. Direction term, weight 1: one minus the cosine similarity  [eq. 13]
    direction_term = 1 - cosine(output_vec, style_vec)
    return 10 * l1_term + 1 * direction_term  # [eq. 14]

Inference. Given any content image, the content is extracted (optionally) by Step 1, encoded with the forward DDIM process to T_{trans}, and reverse-sampled with the fine-tuned model §4.2 app. E. The default inference setting uses 40 forward steps and 6 reverse steps app. E.

Controls. The paper states two knobs. T_{remov} sets how much style is removed from the style image, and therefore how much is transferred §5.3 fig. 4 app. F. T_{trans} sets the return step of the transfer model at training or test time, which trades content preservation for style strength §5.3 fig. 5 app. F.

Hyperparameters

Timing and memory app. E, at 512 \times 512:

• Style removal: about 49.2 s total (2 × 4.921 s × 5) and about 11 GB.
• Latent precompute for 50 images: about 246 s, about 11 GB.
• Fine-tuning, 5 epochs: about 23 minutes in total, about 26 GB.
• Inference (40 forward, 6 reverse): 4.921 s + 0.691 s = 5.612 s, about 13 GB.

Results

Author-reported. Setup: 384 stylized images (32 content × 12 style), averaged; 512 \times 512 images; A100 GPU; each method fine-tuned or trained on the target style §5.2 tab. 1. SSIM measures content similarity, CLIP score measures style similarity, and Style Loss follows Gatys et al. and Huang et al. §5.2. The user study used 50 pairs per user and 20 users, 1000 votes per question §5.2. In the user columns, the figure is the share of votes that preferred the competitor to ours, so lower means ours was preferred.

Source: tab. 1. Training and testing times are measured on an A100, and testing is averaged over 512 \times 512 images.

Ablations §5.3 fig. 7 fig. 8 app. F, on one example set:

• L1 only: SSIM 0.660, CLIP 0.652. L1 + direction: 0.693, 0.705. Full model (L1 + direction + \mathcal{L}_{SR}): 0.672, 0.741. \mathcal{L}_{SR} alone: 0.793, 0.488. Gram loss: 0.429, 0.712. Gram + \mathcal{L}_{SR}: 0.367, 0.763 fig. 7.
• On the AE baseline: Gram loss 0.306 / 0.586 / style loss 0.263; the paper's losses 0.526 / 0.702 / 0.732; both 0.464 / 0.728 / 0.231 fig. 8.
• Diffusion vs AE for transfer: SSIM 0.672 vs 0.526; CLIP 0.741 vs 0.702 fig. 6.
• CLIP encoder ablation (SSIM / CLIP score): ImageNet VGG-19 0.541 / 0.602; CLIP ViT-B/16 0.672 / 0.741; ViT-B/32 0.591 / 0.705; RN50 0.529 / 0.677; RN101 0.582 / 0.719 app. F.

Paper vs code

No official code.

Search: the paper and the wiki's own sources were checked for a link to the authors' code, and none was found. The paper links only the baseline repositories app. B. A GitHub search for "StyleDiffusion" and for the title's distinctive words returned forks and unrelated projects. The only repository that cites this paper is an unofficial reimplementation, which is not used here.

Pitfalls

• Per-style fine-tuning. Every new style needs its own fine-tuning run, about 23 minutes and about 26 GB on an A100 app. E. The authors list this as their first limitation §6. Latents of the 50 content images can be reused across styles app. D.
• Inference cost. One 512 \times 512 image takes 5.6 s and about 13 GB app. E, far slower than the feed-forward baselines in tab. 1. The paper calls it "not fast enough" §6.
• Content images, removal or not. The method text says the content I_c^c comes from the style removal module §4.2. Its footnote says that, unless otherwise specified, the diffusion removal step is not applied to content images, to keep their structure §4.1. Which one the reported results used is not explicit. Implement the footnote behaviour only if you match the reported results.
• Algorithm 1 uses the wrong model. Its reverse loop calls f_{\hat\theta} and \epsilon_{\hat\theta} alg. 1, but \hat\theta is the fine-tuned model, which does not exist yet at style-removal time. The removal step most likely uses the pretrained \epsilon_\theta. This is an inference from the text, which describes removal with the pretrained model §4.1.
• Reconstruction is ordered differently in text and algorithm. The text optimises \mathcal{L}_{SR} first, "before optimizing" \mathcal{L}_{SD} §4.3. Algorithm 2 runs K_s = 50 \mathcal{L}_{SR} steps then one pass of \mathcal{L}_{SD} per epoch, for K = 5 epochs alg. 2. The algorithm is the more specific source; the text is the paper's summary.
• Unstated weights and norms. The weight of \mathcal{L}_{SR} is not stated eq. 15. The paper's \|\cdot\| is called L1 (Manhattan) for \mathcal{L}_{SD}^{L1} §4.2, but whether it is a sum or a mean over the embedding is not stated. The same symbol is used for \mathcal{L}_{SR} on image pixels, with no norm type given eq. 15.
• Learning-rate schedule is ambiguous. "Increase it linearly by 1.2 per epoch" §5.1 can mean a multiplicative factor or an additive step. Neither reading is stated further. Starting at 4e-6 and rising is unusual, so check against any released code before relying on it.
• Diffusion model details. The ADM checkpoint is named only as "pre-trained on ImageNet" §5.1. Whether it is class-conditional, and which input resolution it was trained at, is not stated. Timings are reported at 512 \times 512 app. E, but the model's own resolution is not given.
• Asymmetric step counts. Forward uses 40 steps but reverse only 6 during transfer §5.1. The paper's reason is that 6 reverse steps rebuild clear structures. The approximation error from this choice is not measured.
• Metric bias. The paper says the baselines were trained to optimise Style Loss, which is why they score lower on SSIM and CLIP score than the method §5.2. For style loss, the method (0.837) is lower-scoring than every baseline except TPFR (0.989) tab. 1; Gatys scores 0.111. The paper calls its style loss "comparable" §5.2.
• Small user study. The user study used 20 people, 50 pairs each §5.2. No confidence intervals are reported.
• Removal strength is one-dimensional. Two knobs exist in the algorithm: T_{remov} and K_r alg. 1. The paper sweeps only T_{remov} fig. 4 and does not report K_r sensitivity app. C.
• Diversity at test time. The stochastic DDPM forward process gives diverse outputs §4.2 app. G. The paper gives no quantitative diversity metric; its evidence is the qualitative figure in the supplement app. G.

Lineage

• Builds on: DDPM (Ho et al., 2020) and DDIM (Song et al., 2020) for the sampling processes; ADM (Dhariwal and Nichol, 2021) as the pretrained model; fast sampling and ODE inversion from DiffusionCLIP (Kim et al., 2022) §4.1 §5.1. CLIP (Radford et al., 2021) supplies the projector E §4.3.
• Compared against: Gatys et al. (2016), CLIPstyler (text-driven CLIP-based style transfer, not a baseline in this paper), and the ten baselines in tab. 1.
• Style removal is compared with AE-based removal using the AEs of Li et al. (2017) app. F fig. 6.
• Superseded by: not known from the content of this wiki.