Tianrui Zhu1*, Shiyi Zhang1*, Jiawei Shao2, Yansong Tang1†
1 Tsinghua University, 2 Institute of Artificial Intelligence (TeleAI)
We propose KV-Edit, a training-free image editing approach that strictly preserves background consistency between the original and edited images. Our method achieves impressive performance on various editing tasks, including object addition, removal, and replacement.
- [2025.3.12] Thanks to @smthemex for integrating KV-Edit into ComfyUI!
- [2025.3.4] We update "attention scale" feature to reduce the discontinuity with the background.
- [2025.2.26] Our paper is featured in huggingface Papers!
- [2025.2.25] Code for image editing is released!
- [2025.2.25] Paper released!
- [2025.2.25] More results can be found in our project page!
- ☑️ Release the gradio demo
- ☑️ Release the huggingface space for image editing
- ☑️ Release the paper
We implemented KV Cache in our DiT-based generative model, which stores the key-value pairs of background tokens during the inversion process and concatenates them with foreground content during denoising. Since background tokens are preserved rather than regenerated, KV-Edit can strictly maintain background consistency while generating seamlessly integrated new content.
The environment of our code is the same as FLUX, you can refer to the official repo of FLUX, or running the following command to construct a simplified environment.
Clone the repo:
git clone https://github.com/Xilluill/KV-Edit
We recommend you first use conda to create virtual environment, then run:
conda create --name KV-Edit python=3.10
conda activate KV-Edit
pip install -r requirements.txt
We provide three demo scripts for different hardware configurations. For users with server access and sufficient CPU/GPU memory ( >40/24 GB), we recommend you use:
python gradio_kv_edit.py
For users with 2 GPUs(like 3090/4090) which can avoid offload models to accelerate, you can use:
python gradio_kv_edit_gpu.py --gpus
For users with limited GPU, we recommend you use:
python gradio_kv_edit.py --offload
For users with limited CPU memory such as PC, we recommend you use:
python gradio_kv_edit_inf.py --offload
Here's a sample workflow for our demo:
1️⃣ Upload your image that needs to be edited.
2️⃣ Fill in your source prompt and click the "Inverse" button to perform image inversion.
3️⃣ Use the brush tool to draw your mask area.
4️⃣ Fill in your target prompt, then adjust the hyperparameters.
5️⃣ Click the "Edit" button to generate your edited image!
- 🎨 When using the inversion-based version, you only need to perform the inversion once for each image. You can then repeat steps 3-5 for multiple editing attempts!
- 🎨 "re_init" means using image blending with noise instead of result from inversion to generate new contents.
- 🎨 When the "attn_mask" option is checked, you need to input the mask before performing the inversion.
- 🎨 When the mask is large, and using less skip steps or "re_init", the content of the mask area may not be continuous with background, you can try to increase "attn_scale".
If you find our work helpful, please star 🌟 this repo and cite 📑 our paper. Thanks for your support!
@article{zhu2025kv,
title={KV-Edit: Training-Free Image Editing for Precise Background Preservation},
author={Zhu, Tianrui and Zhang, Shiyi and Shao, Jiawei and Tang, Yansong},
journal={arXiv preprint arXiv:2502.17363},
year={2025}
}
Our code is modified based on FLUX and RF-Solver-Edit. Special thanks to Wenke Huang for his early inspiration and helpful guidance to this project!
This repository is currently under active development and restructuring. The codebase is being optimized for better stability and reproducibility. While we strive to maintain code quality, you may encounter temporary issues during this transition period. For any questions or technical discussions, feel free to open an issue or contact us via email at xilluill070513@gmail.com.