Plug-and-Play Diffusion Features for Text-Driven Image-to-Image Translation

In plain words

Text-to-image models such as Stable Diffusion can draw almost anything you describe, but you cannot easily tell them where things go. If you have a photo of a horse and want "a robot horse" in exactly the same pose and spot, typing the prompt gives you a robot horse in some random pose.

This paper's trick is to run the model twice side by side. One run rebuilds your original picture; the other draws the new prompt. While both run, the method copies a few internal signals from the first run into the second: one layer's feature map, which the authors found encodes "which part of the object is where" (head, legs, torso), and the self-attention patterns, which encode "which pixels belong together". The new image therefore inherits the layout and shape of the original, while the text decides what it looks like.

Nothing is trained or fine-tuned; it works with the off-the-shelf model. Two knobs control how many denoising steps receive the copied signals, trading faithfulness to the original layout against freedom to change its appearance. It works poorly when the input carries no meaning the model can latch onto, such as segmentation masks painted in arbitrary colors.

TL;DR

Plug-and-Play (PnP) does text-guided image-to-image translation with a frozen Stable Diffusion v1.4. The guidance image is DDIM-inverted (empty prompt) to noise x_T; the translated image is sampled from the same x_T with the target prompt, while the decoder ResBlock feature at layer 4 and the self-attention q/k in the decoder are replaced by those of the guidance trajectory for the early part of sampling. A negative prompt describing the source can further push appearance away from it. Relevant if you need zero-shot, structure- preserving edits; it costs an inversion plus a second U-Net trajectory and, in the official code, gigabytes of saved feature tensors per image.

Problem

Text-to-image diffusion models give users little control over the spatial layout of the result: the same prompt yields different poses and layouts from different initial noise §1 §4. Prior options either need training on extra conditioning such as segmentation masks (e.g. Make-A-Scene) §1, trade structure for text fidelity through the noise level (SDEdit) §2 §5.1, or, like Prompt-to-Prompt, control layout through cross-attention to words, which needs word-aligned source/target prompts and does not preserve parts that the prompt does not name §1 §2. The goal is: given a guidance image I^G and a target prompt P, produce I* that follows P while keeping the structure and semantic layout of I^G, with no training §4.

Method

Observations. PCA of decoder ResBlock features f_t^l over 20 humanoid images (real and generated) shows that the coarsest decoder layer mostly separates foreground and background, the intermediate layer 4 encodes semantic parts (legs, torso, head) consistently across images of very different appearance, and deeper layers become high-frequency and noise-like §4 fig. 3. This holds across timesteps fig. 4. PCA of self-attention matrices A_t^l shows they group pixels by layout and semantic region fig. 6. Encoder features show a mirrored trend but less localized semantics app. A.2 fig. 15.

Self-attention in a U-Net layer §3 eq. 2:

\hat{f}^l_t = A^l_t v^l_t, \quad A^l_t = \mathrm{Softmax}(q^l_t {k^l_t}^T)

(The code also applies the usual 1/sqrt(d) scale attention.py.)

Injection. Starting from the same initial noise x*_T = x^G_T, at each step the guidance features are taken from denoising x^G_t with the empty prompt (or with the generating prompt P_G for generated guidance images) and replace the corresponding features in the target pass §4 eq. 3 eq. 4:

z^*_{t-1} = \hat{\epsilon}_\theta(x_t, P, t;\ f^4_t, \{A^l_t\})

Injecting features alone at layer 4 does not keep structure; adding features at deeper layers keeps structure but leaks appearance (colors of clothes) §4 fig. 5a. Features at layer 4 plus self-attention at higher-resolution layers balances the two fig. 5b; self-attention alone loses the semantic correspondence and misaligns structure fig. 5c. By default self-attention is injected into all decoder layers §4 (the code uses decoder blocks 4 to 11, see Paper vs code).

Negative prompting. Classifier-free guidance §4 eq. 5:

\epsilon = w\,\epsilon_\theta(x_t, P, t) + (1-w)\,\epsilon_\theta(x_t, \varnothing, t)

The unconditional term is replaced by a blend with a negative prompt P_n (e.g. a description of the guidance image) eq. 6:

\tilde{\epsilon} = \alpha\,\epsilon_\theta(x_t, \varnothing, t) + (1-\alpha)\,\epsilon_\theta(x_t, P_n, t)

and ε = w ε_θ(x_t, P, t) + (1 − w) ε̃ §4. α starts at α_0 and decreases over sampling app. C. It mainly helps textureless "primitive" inputs such as silhouettes; for natural images its effect is minor §4 app. A.1 fig. 13.

Steps:

1. Encode I^G with the SD autoencoder and DDIM-invert it with the empty prompt to x^G_T (1000 steps) alg. 1 app. C run_features_extraction.py.
2. Run the guidance trajectory from x^G_T with deterministic DDIM and record f^4_t and the decoder self-attention q/k at the 50 timesteps used for translation app. C run_features_extraction.py.
3. Sample I* from x*_T = x^G_T with the target prompt for 50 DDIM steps, overriding f^4 during the first τ_f steps and self-attention during the first τ_A steps, with CFG and optional negative prompting alg. 1 app. C run_pnp.py.

# Plug-and-Play Diffusion Features  [alg. 1]
# Input:  guidance image I_G, target prompt P, thresholds tau_f, tau_A,
#         negative prompt P_neg, CFG scale w, per-step alpha_schedule
# Output: translated image I_star

def pnp_translate(I_G, P, P_neg, w, alpha_schedule, tau_f, tau_A, steps=50):
    # Invert the guidance image with the empty prompt  [§4] [app. C]
    x_T_guid = ddim_invert(encode(I_G), prompt="")
    # Start the translation from the same noise  [§4] [app. B]
    x_guid = x_T_guid
    x_star = x_T_guid
    for i, t in enumerate(ddim_timesteps_descending(steps)):
        # Guidance pass; keep layer-4 ResBlock output and decoder q, k  [alg. 1]
        eps_guid, f4, attn_qk = unet(x_guid, t, prompt="", record=True)
        x_guid = ddim_step(x_guid, eps_guid, t)
        # Inject only during the early part of sampling  [alg. 1] [app. C]
        f_inj = f4 if i < tau_f else None
        a_inj = attn_qk if i < tau_A else None
        # Target pass with overridden features, CFG + negative prompt  [eq. 6]
        eps_star = guided_eps(
            x_star, t, P, P_neg, alpha_schedule[i], w, f_inj, a_inj
        )
        x_star = ddim_step(x_star, eps_star, t)
    return decode(x_star)

(The loop counter i above counts sampling steps from the noisy end; the paper's t > τ notation in Alg. 1 means the same thing, see Pitfalls.)

# Guided noise prediction with negative prompting  [eq. 5] [eq. 6]
# [code: ldm/models/diffusion/ddim.py#L216-L237]

def guided_eps(x, t, P, P_neg, alpha, w, f_inj, a_inj):
    # All U-Net calls receive the same injected features  [eq. 6]
    eps_text = unet(x, t, P, f_inj, a_inj)
    eps_uncond = unet(x, t, "", f_inj, a_inj)
    eps_neg = unet(x, t, P_neg, f_inj, a_inj)
    # Blend neutral and negative; alpha decays over the steps  [app. C]
    eps_tilde = alpha * eps_uncond + (1 - alpha) * eps_neg
    # Classifier-free guidance with eps_tilde as the anchor  [eq. 5]
    return eps_tilde + w * (eps_text - eps_tilde)

What is replaced, concretely (from code):

• Feature f^4: the output of the residual branch (out_layers) of the ResBlock in decoder output_blocks[4]; the skip connection is still added from the target pass openaimodel.py. With the SD v1 config (model_channels 320, channel_mult [1,2,4,4], num_res_blocks 2) v1-inference.yaml this is in the second-coarsest decoder level (inferred from the config).
• Self-attention: the queries and keys of the first attention (attn1) in each injected decoder block are replaced; values come from the target pass attention.py. Replacing q and k is equivalent to replacing A_t^l in eq. 2.
• Saved tensors hold an unconditional and a conditional half; each half is broadcast over the corresponding half of the target batch attention.py openaimodel.py.

Hyperparameters

Results

Author-reported. SD v1.4, metrics: CLIP text-image cosine similarity (text fidelity, higher is better), distance between DINO-ViT self-similarity matrices (structure, lower is better), and in the appendix LPIPS between I^G and I* (deviation from guidance, higher means more change) §5.1 tab. 1. Benchmarks built by the authors: Wild-TI2I, 148 text-image pairs, 53% real web images, the rest generated §5 app. D.2; ImageNet-R-TI2I, 30 images (3 per class from 10 ImageNet-R classes) × 5 templated prompts = 150 pairs, two of the five prompts swapping the class for a related one app. D.1; and Generated-ImageNet-R-TI2I, built from text prompts with the same classes and renditions for comparing with P2P without inversion §5.1.

Main comparison (plotted, no numbers given) §5.1 fig. 9: against SDEdit at three noise levels, P2P, DiffuseIT and VQGAN-CLIP, PnP reaches structure distance comparable to SDEdit at a very low noise level and CLIP score comparable to SDEdit at a very high noise level. P2P is compared only on the two ImageNet-R benchmarks because it needs word-aligned prompts; PnP has a markedly lower structure distance than P2P even when P2P injects cross-attention over all timesteps §5.1. The authors note VQGAN-CLIP and DiffuseIT optimize the CLIP and DINO self-similarity metrics directly §5.1. Text2LIVE, DiffusionCLIP and FlexIT are compared only qualitatively fig. 11.

Ablations tab. 1:

Reading §5.2 app. A: removing either the feature or the self-attention injection roughly doubles the structure distance; negative prompting raises LPIPS (more appearance change) at a small structure cost; also injecting encoder layer-7 features lowers CLIP on every set.

Paper vs code

• Self-attention layers. Paper: injected into all decoder layers §4. Code: decoder blocks 4–11 only run_pnp.py; block 3 also has self-attention under the SD v1 config (attention_resolutions [4,2,1]) v1-inference.yaml but is not injected.
• τ_A is not configurable. Paper uses τ_A = 25, and 25 for primitive images too app. C; code hard-codes ddim_steps // 2 run_pnp.py. Only τ_f is exposed pnp-real.yaml.
• Threshold off-by-one. Code injects when the step index i <= τ run_pnp.py run_pnp.py, i.e. τ+1 steps (41 steps for τ_f = 40, 26 for τ_A = 25).
• CFG scale for real images. Paper 15.0 app. C; the shipped real-image config uses 10.0 pnp-real.yaml.
• Step counts. Paper 1000 inversion and 1000 reconstruction steps app. C; code 999 run_features_extraction.py feature-extraction-real.yaml.
• Exponential α schedule. Paper writes α(t) = e^{−6t} app. C; code multiplies by α_0 (0.75 for generated images), so it starts at 0.75, not 1 ddim.py.
• Negative prompt for real images. The paper suggests P_n describes the guidance image §4; in code it defaults to the source prompt run_pnp.py, which is forced to "" for real images run_features_extraction.py, making negative prompting a no-op unless negative_prompt is set (the example config sets it) pnp-real.yaml.
• Two-stage execution. Alg. 1 runs guidance and target passes in lockstep; the code first runs the guidance trajectory and saves all features to disk, then loads them for translation run_features_extraction.py run_pnp.py. Mathematically equivalent.

Pitfalls

• Meaning of τ. Alg. 1 writes if t > τ_f then inject with t running T…1, yet the values are given as "out of the 50 sampling steps" alg. 1 app. C. The code resolves this: inject during the first τ sampling steps (counting from the noisy end) run_pnp.py. Larger τ means more structure preservation.
• Same initial noise is mandatory. The translation must start from the inverted x^G_T; with a different seed the injected decoder features mismatch the encoder skip features and the method fails app. B fig. 14.
• Inversion with the empty prompt. PnP inverts with an empty prompt so any CFG scale and prompt can be used at sampling; the authors attribute P2P's weakness on real images to inverting with a source text §5.1. Do not invert with the target or source prompt.
• Guidance trajectory is a 1000-step reconstruction sampled only at the 50 translation timesteps app. C; extracting features from a 50-step reconstruction instead is not what the paper does.
• f^4 is the residual-branch output, not the block output. The skip connection of the target pass is still added openaimodel.py. Injecting the full block output would also copy the target's skip path out.
• Injection hits every CFG branch. Conditional, unconditional and negative passes all receive the injected features ddim.py.
• Cost. Negative prompting in the code runs the U-Net on two batches of two per step (negative+uncond, uncond+cond), recomputing the unconditional branch ddim.py; that is twice the usual CFG cost. Real images also need ~1000 inversion and ~1000 reconstruction U-Net calls before translation app. C, and the code stores q, k for 8 decoder blocks at every step on disk run_features_extraction.py.
• Q/K only for the attention of the injected blocks. The values still come from the target pass attention.py; that is why appearance can change while affinities are copied.
• Limitations. Fails when guidance content has no semantic association with the text, e.g. segmentation masks with arbitrary colors §6 fig. 12. DDIM inversion of textureless images may encode dominant low-frequency appearance, which then leaks into the result §6.
• Code license. The repository has no LICENSE file, so do not copy its code; re-implement from the description.

Lineage

• Builds on: Stable Diffusion / latent diffusion (arXiv 2112.10752) as the frozen backbone; DDIM (arXiv 2010.02502) sampling and inversion; classifier- free guidance (arXiv 2207.12598); DINO-ViT self-similarity from Splicing ViT features (arXiv 2201.00424) as the structure metric §3 §5.1.
• Compared against: SDEdit (arXiv 2108.01073), Prompt-to-Prompt (arXiv 2208.01626), DiffuseIT, VQGAN-CLIP, Text2LIVE, DiffusionCLIP, FlexIT §5.1.