FARI: Robust One-Step Inversion for Watermarking in Diffusion Models

Official implementation for the ICLR 2026 paper FARI: Robust One-Step Inversion for Watermarking in Diffusion Models.

FARI is a fast and robust asymmetric inversion method designed for inversion-based watermarking in diffusion models. It compresses the inversion trajectory into a one-step approximation and fine-tunes lightweight LoRA adapters for robust watermark recovery under image distortions.

• One-step inversion for efficient watermark verification.
• LoRA-only fine-tuning on top of the pretrained denoiser.
• Robust evaluation for both Gaussian Shading and Tree-Ring watermarking.
• Pretrained FARI weights included at results/fari_default/fari_weights.pth.

.
├── train.py                  # Train FARI with adversarial distortions
├── val_gs.py                 # Evaluate on Gaussian Shading
├── val_tr.py                 # Evaluate on Tree-Ring
├── requirements.txt          # Python dependencies
├── results/fari_default/     # Released checkpoint
└── src/
    ├── fari.py               # FARI LoRA wrappers and one-step inversion
    ├── inversion/inv_pipe.py # Diffusion pipeline with inversion utilities
    ├── utils.py              # Prompt loading and image distortions
    ├── watermark.py          # Gaussian Shading implementation
    ├── io_utils.py           # Tree-Ring utility code
    └── optim_utils.py        # Tree-Ring evaluation utilities

Python 3.11 is recommended.

python -m venv .venv
source .venv/bin/activate
pip install -r requirements.txt

This codebase is intended for CUDA-enabled GPUs. The experiments in the paper use Stable Diffusion 2.1 Base and Hugging Face diffusers.

train.py accepts three prompt sources through --train_dataset_id:

• a Hugging Face dataset name, such as Gustavosta/Stable-Diffusion-Prompts
• a local JSON file containing an annotations field with caption
• a local CSV file containing a column named Our GT caption

For COCO-based experiments, place the metadata file at:

data/coco/meta_data.json

or pass a custom path with --train_dataset_id.

The COCO prompts and ground-truth images used in the paper are available from the Tree-Ring Watermark project:

• Google Drive download
• Tree-Ring Watermark repository

A pretrained checkpoint trained with the default settings is included:

results/fari_default/fari_weights.pth

You can use --name fari_default directly for evaluation.

Example training command:

python train.py \
  --name fari_default \
  --train_dataset_id data/coco/meta_data.json \
  --steps 1000 \
  --batch_size 4 \
  --lr 1e-4

Training outputs are written to results/<name>/:

• TRAIN.log
• training_settings.json
• fari_weights.pth

python val_gs.py \
  --name fari_default \
  --val_dataset_id Gustavosta/Stable-Diffusion-Prompts \
  --val_size 1000

This writes GS.log, gs_settings.json, and val_gs.json to results/<name>/.

python val_tr.py \
  --name fari_default \
  --dataset Gustavosta/Stable-Diffusion-Prompts \
  --start 0 \
  --end 1000

This writes TR.log, tr_settings.json, and val_tr.json to results/<name>/.

This repository reuses or adapts components from prior watermarking projects:

• Gaussian Shading implementation in src/watermark.py is based on the official Gaussian-Shading repository.
• Tree-Ring related utilities in src/io_utils.py and src/optim_utils.py are based on the official tree-ring-watermark repository.
• LoRA injection relies on cloneofsimo/lora.

Please also follow the original licenses of those upstream projects when reusing derived components.

If you are interested in a complementary line of work, check out our related project T2SMark: Balancing Robustness and Diversity in Noise-as-Watermark for Diffusion Models. While FARI focuses on making inversion-based watermark recovery fast and robust through one-step inversion, T2SMark focuses on the watermark design itself, addressing the robustness-diversity trade-off in Noise-as-Watermark methods via Tail-Truncated Sampling and a two-stage keying scheme. The released T2SMark code also covers both Stable Diffusion v2.1 and Stable Diffusion v3.5 Medium.

If you find this repository useful, please cite:

@inproceedings{
  yang2026fari,
  title={FARI: Robust One-Step Inversion for Watermarking in Diffusion Models},
  author={Jindong Yang and Han Fang and Weiming Zhang and Nenghai Yu and Kejiang Chen},
  booktitle={The Thirteenth International Conference on Learning Representations},
  year={2026},
  url={https://openreview.net/forum?id=YiGdNowqj6}
}

This repository is released under the MIT License. See LICENSE for details.