SERUM: Simple, Efficient, Robust, and Unifying Marking for Diffusion-Based Image Generation

Overview

We propose SERUM: an intriguingly simple yet highly effective method for marking images generated by diffusion models (DMs). We only add a unique watermark noise to the initial diffusion generation noise and train a lightweight detector to identify watermarked images, simplifying and unifying the strengths of prior approaches. SERUM provides robustness against any image augmentations or watermark removal attacks and is extremely efficient, all while maintaining negligible impact on image quality. In contrast to prior approaches, which are often only resilient to limited perturbations and incur significant training, injection, and detection costs, our SERUM achieves remarkable performance, with the highest true positive rate (TPR) at a 1% false positive rate (FPR) in most scenarios, along with fast injection and detection and low detector training overhead. Its decoupled architecture also seamlessly supports multiple users by embedding individualized watermarks with little interference between the marks. Overall, our method provides a practical solution to mark outputs from DMs and to reliably distinguish generated from natural images.

Installation

# 1. Create environment
conda create -n serum python=3.13.5
conda activate serum

# 2. Install PyTorch (CUDA 11.8 example)
pip install --upgrade pip
pip install torch==2.7.1 torchvision==0.22.1 --index-url https://download.pytorch.org/whl/cu118

# CPU-only (not recommended)
pip install torch==2.7.1 torchvision==0.22.1

# 3. Install dependencies
pip install -r requirements.txt

# 4. Install a C compiler (optional, for torch.compile)
conda install -c conda-forge gcc gxx

Quick Start

# Activate environment
conda activate serum

# Generate dataset (clean + augmented)
python -m src.data.dataset --config config.yaml --mode both

# Train watermark detector
python -m src.training.train --config config.yaml

# Evaluate performance
python -m src.evaluation.eval --config config.yaml

All results are saved under results/{experiment_name}

---

Project Structure

src/
├── data/                   # Dataset handling and augmentation
├── models/                 # Diffusion, watermark and watermark detector models
├── training/               # Training procedures
├── evaluation/             # Evaluation, visualization, and metrics
└── utils/                  # Utility functions and helpers

configs/                    # YAML configuration files
notebooks/                  # Jupyter notebooks for experiments (demo.ipynb)
results/                    # Model checkpoints, generated images, evaluation outputs

Demo Notebook

A demo.ipynb notebook demonstrates the full training and evaluation pipeline.

Custom Augmentations

Modify Augment class in src/data/augmentation.py. Default pipeline includes:

• Random Rotation
• JPEG Compression
• Crop & Scale
• Random Drop
• Gaussian Blur
• Salt & Pepper Noise
• Gaussian Noise
• Brightness Adjustments

Detailed Command Reference

See USAGE.md for all dataset, training, evaluation, generation, and scoring flags and examples.