🌊 MaskFlow: Precise, Consistent and Seamless Regional Image Editing

Overview

🌊 Models: This repository is the official implementation for paper "MaskFlow: Precise, Consistent and Seamless Regional Image Editing", including:

Pipelines

Schedulers

DataModule

Trainers

Evaluators

Editor

🎨 Dataset: MaskEdit-10k is available on 🤗 Hugging Face.

💜 Local Editor: Editor is relsead for pratical user-specified and freeform masks designing, the editor can be deployed on the local and server.

🩵 Demo (Comming soon): MaskFlow is integrated into LightX2V for an accessible inference workflow.

⭐️ Please leave your star if these can help you to create attractive artworks ⭐️

Contents
Introduction
🪄 [New] Editor
Quick Start
Distribution Matching Distillation
Configuration reference
Visualization
License
Citation

🌊 Introduction

MaskFlow is a mask-aware framework for precise regional image editing. Given a source image, a spatial mask, and a text instruction, it edits the selected region while preserving the surrounding content. Its localized generation process and Soft-Poisson refinement improve regional control, background consistency, and boundary quality.

🪄 [New] Editor

We release the image editor for convenient usage, supporting:

Freeform Mask: User can draw masks on the source image with any shapes to identify the editable region. The masks can also be saved for future use!
Online Inference: The editor can be deployed on the server to share the convenience to more people.

Deployment

After environment is ready, just run uv run python -m editor, and this editor will deployed on http://127.0.0.1:7890 on the local by default.

🍪 Quick Start

1. Set up the environment

MaskFlow requires Python 3.12 or later. An NVIDIA GPU is recommended for inference.

git clone https://github.com/ReyChiaro/MaskFlow.git
cd MaskFlow

# Install uv if it is not already available.
python -m pip install uv

# Reproduce the locked Python 3.12 environment.
uv python install 3.12
uv sync

Alternative installation with venv and pip

python3.12 -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip
python -m pip install -e .

Install a PyTorch build compatible with your CUDA environment if the automatically resolved build does not match your system.

2. Prepare the inputs

Prepare a source RGB image and a spatially aligned mask. White pixels in the mask indicate the region to edit; black pixels indicate the region to preserve. If a prompt comes from MaskEdit-10k and contains [MASK_AREA], replace the placeholder with a natural referring phrase before inference.

3. Choose a checkpoint

The LoRA adapters are hosted in ReyChiaro/MaskFlow. Diffusers downloads and caches the selected files automatically.

File	Variant	Steps	Text CFG	Intended use
`maskflow-S.safetensors`	S	50	4.0	Standard checkpoint trained on the `scene` split
`maskflow-S-tcfg4-step8.safetensors`	S distilled	8	4.0	Accelerated scene editing
`maskflow-S-tcfg4-step16.safetensors`	S distilled	16	4.0	Accelerated scene editing
`maskflow-SEC.safetensors`	SEC	50	4.0	Standard checkpoint trained on all MaskEdit-10k splits
`maskflow-SEC-tcfg4-step8.safetensors`	SEC distilled	8	4.0	Accelerated general editing
`maskflow-SEC-tcfg4-step16.safetensors`	SEC distilled	16	4.0	Accelerated general editing

S denotes training on the scene split, while SEC denotes training on all scene and infographic splits. A distilled LoRA is a residual adapter: it must be used with its matching standard SFT LoRA (S with S, or SEC with SEC). The SFT adapter is loaded as maskflow, and the distilled adapter is loaded as dmd.

4. Run standard 50-step inference

uv run python inference.py \
  input.source=/absolute/path/to/source.png \
  input.mask=/absolute/path/to/mask.png \
  'input.prompt=Replace the masked object with a red ceramic vase.' \
  checkpoint.sft_path=ReyChiaro/MaskFlow \
  checkpoint.sft_weight_name=maskflow-S.safetensors \
  runtime.num_inference_steps=50 \
  runtime.text_cfg_scale=4.0 \
  output.path=outputs/result.png

The base model defaults to Qwen/Qwen-Image-Edit-2511, and the output directory is created automatically.

Distribution Matching Distillation

To improve efficiency for practical deployment, we apply Distribution Matching Distillation (DMD) and provide accelerated 8-step and 16-step variants. The distilled LoRA represents a residual on top of the corresponding standard MaskFlow checkpoint, so both the matching SFT and DMD weights are required during inference. Although the student is distilled with teacher text classifier-free guidance, enabling CFG during student inference generally gives better performance.

Both the standard SFT LoRA and its matching distilled LoRA are required. The following example uses the SEC 8-step pair:

uv run python inference.py \
  input.source=/absolute/path/to/source.png \
  input.mask=/absolute/path/to/mask.png \
  'input.prompt=Replace the masked object with a red ceramic vase.' \
  checkpoint.sft_path=ReyChiaro/MaskFlow \
  checkpoint.sft_weight_name=maskflow-SEC.safetensors \
  checkpoint.dmd_path=ReyChiaro/MaskFlow \
  checkpoint.dmd_weight_name=maskflow-SEC-tcfg4-step8.safetensors \
  runtime.num_inference_steps=8 \
  runtime.text_cfg_scale=4.0 \
  output.path=outputs/result.png

To use local files, pass the two .safetensors paths and leave the corresponding weight_name fields unset:

checkpoint.sft_path=/absolute/path/to/maskflow-SEC.safetensors \
checkpoint.dmd_path=/absolute/path/to/maskflow-SEC-tcfg4-step8.safetensors

Configuration reference

Inference uses Hydra, so any field in configs/inference.yaml or the selected pipeline configuration can be overridden with key=value.

Override	Default	Description
`checkpoint.sft_path`	`ReyChiaro/MaskFlow`	Local SFT LoRA path or Hugging Face repository ID
`checkpoint.sft_weight_name`	`maskflow-S.safetensors`	SFT filename when loading from a multi-weight Hub repository
`checkpoint.dmd_path`	`null`	Local distilled LoRA path or Hugging Face repository ID
`checkpoint.dmd_weight_name`	`null`	Distilled filename when loading from a multi-weight Hub repository
`runtime.device`	`cuda`	Torch device used for inference
`runtime.dtype`	`bfloat16`	Torch compute dtype
`runtime.seed`	`42`	Random seed
`runtime.num_inference_steps`	`50`	Denoising steps; must match the selected distilled checkpoint
`runtime.text_cfg_scale`	`4.0`	Text classifier-free guidance scale
`pipeline.enable_pixel_blend`	`true`	Blend the unmasked pixels from the source image into the final result
`output.path`	timestamped path	Output image path

The adapter names are fixed to maskflow for SFT and dmd for step distillation. If checkpoint.dmd_path is provided without checkpoint.sft_path, inference stops with an error instead of silently producing an incorrectly initialized result.

Visualization

MaskFlow supports a diverse range of mask-guided image editing tasks. The comparisons below show that, relative to other models, MaskFlow localizes edits more precisely while better preserving the surrounding content. It also produces smoother transitions between edited and preserved regions, resulting in higher visual fidelity.

MaskFlow is also well suited to applications such as infographic editing, where the target location can be difficult to specify through language alone. Spatial masks provide direct and intuitive control, making the method practical for real-world editing workflows.

License

MaskFlow code and adapter weights are released under the MIT License. Use of the Qwen base model and third-party datasets remains subject to their respective licenses and terms.

Citation

@misc{xu2026maskflowpreciseconsistentseamless,
  title={MaskFlow: Precise, Consistent and Seamless Regional Image Editing},
  author={Rui Xu and Yang Yong and Shunzi Yang and Ruihao Gong and Chengtao Lv},
  year={2026},
  eprint={2608.06929},
  archivePrefix={arXiv},
  primaryClass={cs.CV},
  url={https://arxiv.org/abs/2608.06929},
}

FilesExpand file tree

README.md

Latest commit

History