kv-edit-flux-modular

KV-Edit for FLUX — training-free editing with pixel-precise background preservation (Modular Diffusers custom block)

Text-based image editing on off-the-shelf FLUX.1-dev — no training, no extra weights — as a Modular Diffusers custom block. Implements KV-Edit (arXiv:2502.17363): the source image is rectified-flow inverted once under its source prompt while every attention block caches the K/V of the background tokens; the edit then denoises under the target prompt with those cached background K/V substituted in place. The unedited background is preserved, not re-synthesized — the failure mode this fixes vs. FlowEdit (structure-preserving, but the background still drifts).

Usage

import torch
from diffusers import ModularPipeline

pipe = ModularPipeline.from_pretrained("remyxai/kv-edit-flux-modular", trust_remote_code=True)
pipe.load_components(dtype=torch.bfloat16)
pipe.to("cuda")

edited = pipe(
    image="photo.png",
    mask="mask.png",             # white = region to edit, black = keep pixel-precise
    source_prompt="a cat sitting on a sofa",   # describes the source image
    prompt="a dog sitting on a sofa",          # the target edit
    T_steps=28, guidance_scale=3.5,
).images[0]
edited.save("edited.png")

Tip: mask is optional — with mask=None the whole image is editable and the K/V seam is disarmed (bit-exact stock FLUX attention, i.e. plain RF-inversion editing). A good source_prompt (describing what is actually in the image) is what makes the inversion — and therefore the background fidelity — faithful. Invert once, edit many times: the cached K/V are per image+mask.

How it works

1. Invert the source image with rectified-flow inversion under the source_prompt, running stock FLUX attention while a custom attention processor captures the normalized K/V of the background image tokens (the mask == keep region) at every block and step.
2. Denoise from the inverted latent under the target prompt: the processor substitutes the cached background K/V in place of the current ones — equivalent to the reference's "concat the background K/V memory with the foreground content" (softmax over keys is permutation-invariant).
3. Decode. Only the masked region is regenerated; background tokens attend to their own cached state, so the background is retained exactly rather than re-synthesized.

The seam is an attention-processor swap threaded through joint_attention_kwargs['kv_edit'] (concurrency-safe, no globals); the original processors are restored in a finally. Captured K/V are kept on CPU (O(#background tokens) per block per step) and sliced back per denoise step.

Key parameters

Attribution & AI assistance

Training-free reimplementation for Modular Diffusers of KV-Edit (Apache-2.0, Xilluill/KV-Edit). The Modular-Diffusers adaptation was authored with AI assistance (Claude) and validated by the Remyx AI team; method credit to the KV-Edit authors. Uses FLUX.1-dev under its non-commercial license.

Citation

@misc{zhu2025kvedit,
  title={KV-Edit: Training-Free Image Editing for Precise Background Preservation},
  author={Zhu, Tianrui and Zhang, Shiyi and Shao, Jiawei and Tang, Yansong},
  year={2025},
  eprint={2502.17363}, archivePrefix={arXiv}, primaryClass={cs.CV},
  url={https://arxiv.org/abs/2502.17363}
}