HierEdit: Region-Aware Hierarchical Diffusion for Efficient High-Resolution Editing
Yuyao Zhang, Alexander Huang-Menders, Yu-Wing Tai
CVPR 2026
[Project Page] [arXiv] [CVPR OpenAccess]
HierEdit is a region-aware hierarchical diffusion framework for efficient high-resolution image editing. It first edits a low-resolution proxy to obtain a reference and localize modified regions, then refines only the edited regions in the original high-resolution image with a Local-Window MMDiT while reusing unedited regions as conditioning. This design avoids dense full-canvas attention and enables high-fidelity editing up to 4K resolution without specialized high-resolution training data. Notice that this is a flex-attention version of HierEdit for better deployment, so the running time may be different from the original version.
- HierEdit is accepted to CVPR 2026.
Create an environment and install the Python dependencies:
conda create -n hieredit python=3.10
conda activate hieredit
pip install -r train/requirements.txtThis code builds on FLUX.1-dev through Diffusers. Please make sure you have access to the base model and follow the corresponding model license.
Download the trained HierEdit LoRA adapters and place them in a local directory
(default ./checkpoints/hieredit):
checkpoints/
hieredit/
edit.safetensors
ref.safetensors
Download link: Google Drive
Pass the adapter directory to HierEditHighResInpainter(lora_dir=...).
High-resolution inpainting from a single (image, bbox, prompt). A low-res inpaint produces LR guidance, then HierEdit regenerates only the bbox at native resolution and composites it back with a hidden, feathered seam.
Notebook: examples/example.ipynb — set the
image path, bbox, and prompt.
Or from Python:
from PIL import Image
from hieredit.highres_inpaint import HierEditHighResInpainter
inpainter = HierEditHighResInpainter("./checkpoints/hieredit", device="cuda")
res = inpainter.inpaint(
Image.open("path/to/base.jpg").convert("RGB"),
bbox=(700, 1320, 1380, 1960), # (x1, y1, x2, y2) in image pixels
prompt="a small waterfall",
)
res["final"].save("out.png") # seamless result (also res["raw"], res["lr"])bbox: edited region in image pixels.target_size(default = native resolution) caps HierEdit resolution; 4k is heavy.- The seam-hiding blend is applied internally.
Core inference files:
hieredit/highres_inpaint.pyhieredit/pipeline/flux_block_attn.py
The default task is reference-guided inpainting (edit_ref):
bash train/script/train_guided_inpainting_edit_ref.shSet your Weights & Biases key in the script (or run wandb login) before training. Configs and entry points:
train/config/token_integration_with_reference_edit_ref.yaml— default edit_ref config (public text-to-image-2M shards)train/script/train_guided_inpainting_edit_ref.sh
train/script/train_guided_inpainting_guide.sh (→ hieredit/train_flux/train_block_attn_guided.py, config token_integration_with_reference_guidance.yaml) trains from a locally prepared dataset of (condition, result, bbox) triples loaded via datasets.load_from_disk. Point it at your data with dataset.data_dir in the config or the HIEREDIT_GUIDED_DATASET environment variable.
examples/ # Inference notebook (example.ipynb)
hieredit/highres_inpaint.py # High-res inpainting API (2-stage: LR inpaint -> HierEdit)
hieredit/pipeline/ # HierEdit/FLUX generation pipeline
hieredit/train_flux/ # Training code
train/config/ # Training YAML configs
train/script/ # Training launch scripts
If you find this work useful, please cite:
@inproceedings{zhang2026hieredit,
title={HierEdit: Region-Aware Hierarchical Diffusion for Efficient High-Resolution Editing},
author={Zhang, Yuyao and Huang-Menders, Alexander and Tai, Yu-Wing},
booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
pages={43546--43557},
year={2026}
}We thank the authors and maintainers of FLUX, Diffusers, PEFT, and the open-source sparse attention tools used by this project.