OSI: One-step Inversion Excels in Extracting Diffusion Watermarks
Yuwei Chen^1,2, Zhenliang He^1,2✉, Jia Tang^1,3, Meina Kan^1,2, Shiguang Shan^1,2
1State Key Laboratory of AI Safety, Institute of Computing Technology, CAS, China
²University of Chinese Academy of Sciences (CAS), China
³School of Information Science and Technology, ShanghaiTech University, China

We introduce OSI, a fast and accurate method for extracting the initial noise sign from diffusion-generated images, with exceptional performance in Gaussian Shading style watermark extraction.

Clone this repo:

git clone https://github.com/VIPL-GENUN/OSI.git
cd OSI

Create and activate a new conda environment:

conda create -n osi python=3.8.20
conda activate osi

Install dependencies:

pip install -r requirements.txt

We have provided the OSI model for Stable Diffusion v2.1 on HuggingFace. Weights for other versions, including SD3.5 and SDXL, will be updated soon. You can download it using the following command:

huggingface-cli download VIPL-GENUN/OSI

Prepare the pre-trained Stable Diffusion weights and the OSI model weights, then update paths in the command/script:

• --model_path: Stable Diffusion checkpoint path
• --unet_path / --encoder_path: OSI UNet / encoder checkpoints
• --dataset_path: Stable Diffusion Prompts or MS-COCO datasets path

Our inference process run_osi_sd21.py starts by loading the Stable Diffusion Prompts dataset, the pre-trained Stable Diffusion (v2.1), and our OSI model (both UNet and Encoder). In each step, the system takes a prompt from the dataset to generate a watermarked image. This image is then put through an attack to test robustness. Finally, the OSI model processes the attacked image to predict the signs of the initial latent noise, from which the hidden watermark message is extracted.

Run a single distortion example:

python run_osi_sd21.py \
  --num 1000 \
  --image_length 512 \
  --guidance_scale 7.5 \
  --num_inference_steps 50 \
  --channel_copy 1 \
  --hw_copy 8 \
  --fpr 0.000001 \
  --output_path ./output/ \
  --dataset_path <dataset_path> \
  --model_path <stable_diffusion_model_path> \
  --unet_path <osi_unet_ckpt> \
  --encoder_path <osi_encoder_ckpt> \
  --distortion_name Jpeg \
  --jpeg_ratio 25

Run all distortion presets in one go:

bash scripts/run_osi.sh

Note on Scripts: For more examples involving different adversarial attacks or clean (no-attack) evaluations, please refer to the configurations in scripts/run_osi.sh.

We borrow the code from Tree-Ring Watermark and Gaussian Shading. We appreciate the authors for sharing their code.

If you find our work helpful, please consider citing:

@article{chen2026osionestepinversionexcels,
      title={OSI: One-step Inversion Excels in Extracting Diffusion Watermarks}, 
      author={Yuwei Chen, Zhenliang He, Jia Tang, Meina Kan, Shiguang Shan},
      journal={arXiv preprint arXiv:2602.09494},
      year={2026}
}