Skip to content

THU-BPM/MarkDiffusion

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

An Open-Source Toolkit for Generative Watermarking of Latent Diffusion Models

Home Paper Models Colab DOC PYPI CONDA-FORGE

Language Versions: English | 中文 | Français | Español

🔥 As a new released project, We welcome PRs! If you have implemented a LDM watermarking algorithm or are interested in contributing one, we'd love to include it in MarkDiffusion. Join our community and help make generative watermarking more accessible to everyone!

Contents

🔥 Updates

🛠 (2025.12.19) Add a complete test suite for all functionality with 658 test cases.

🛠 (2025.12.10) Add a continuous integration testing system using github actions.

🎯 (2025.10.10) Add Mask, Overlay, AdaptiveNoiseInjection image attack tools, thanks Zheyu Fu for his PR!

🎯 (2025.10.09) Add FrameRateAdapter, FrameInterpolationAttack video attack tools, thanks Luyang Si for his PR!

🎯 (2025.10.08) Add SSIM, BRISQUE, VIF, FSIM image quality analyzer, thanks Huan Wang for her PR!

(2025.10.07) Add SFW watermarking method, thanks Huan Wang for her PR!

(2025.10.07) Add VideoMark watermarking method, thanks Hanqian Li for his PR!

(2025.9.29) Add GaussMarker watermarking method, thanks Luyang Si for his PR!

🔓 Introduction to MarkDiffusion

👀 Overview

MarkDiffusion is an open-source Python toolkit for generative watermarking of latent diffusion models. As the use of diffusion-based generative models expands, ensuring the authenticity and origin of generated media becomes critical. MarkDiffusion simplifies the access, understanding, and assessment of watermarking technologies, making it accessible to both researchers and the broader community. Note: if you are interested in LLM watermarking (text watermark), please refer to the MarkLLM toolkit from our group.

The toolkit comprises three key components: a unified implementation framework for streamlined watermarking algorithm integrations and user-friendly interfaces; a mechanism visualization suite that intuitively showcases added and extracted watermark patterns to aid public understanding; and a comprehensive evaluation module offering standard implementations of 31 tools across three essential aspects—detectability, robustness, and output quality, plus 6 automated evaluation pipelines.

MarkDiffusion Overview

💍 Key Features

  • Unified Implementation Framework: MarkDiffusion provides a modular architecture supporting eleven state-of-the-art generative image/video watermarking algorithms of LDMs.

  • Comprehensive Algorithm Support: Currently implements 11 watermarking algorithms from two major categories: Pattern-based methods (Tree-Ring, Ring-ID, ROBIN, WIND, SFW) and Key-based methods (Gaussian-Shading, PRC, SEAL, VideoShield, GaussMarker, VideoMark).

  • Visualization Solutions: The toolkit includes custom visualization tools that enable clear and insightful views into how different watermarking algorithms operate under various scenarios. These visualizations help demystify the algorithms' mechanisms, making them more understandable for users.

  • Evaluation Module: With 31 evaluation tools covering detectability, robustness, and impact on output quality, MarkDiffusion provides comprehensive assessment capabilities. It features 6 automated evaluation pipelines: Watermark Detection Pipeline, Image Quality Analysis Pipeline, Video Quality Analysis Pipeline, and specialized robustness assessment tools.

✨ Implemented Algorithms

Algorithm Category Target Reference
Tree-Ring Pattern Image Tree-Ring Watermarks: Fingerprints for Diffusion Images that are Invisible and Robust
Ring-ID Pattern Image RingID: Rethinking Tree-Ring Watermarking for Enhanced Multi-Key Identification
ROBIN Pattern Image ROBIN: Robust and Invisible Watermarks for Diffusion Models with Adversarial Optimization
WIND Pattern Image Hidden in the Noise: Two-Stage Robust Watermarking for Images
SFW Pattern Image Semantic Watermarking Reinvented: Enhancing Robustness and Generation Quality with Fourier Integrity
Gaussian-Shading Key Image Gaussian Shading: Provable Performance-Lossless Image Watermarking for Diffusion Models
GaussMarker Key Image GaussMarker: Robust Dual-Domain Watermark for Diffusion Models
PRC Key Image An undetectable watermark for generative image models
SEAL Key Image SEAL: Semantic Aware Image Watermarking
VideoShield Key Video VideoShield: Regulating Diffusion-based Video Generation Models via Watermarking
VideoMark Key Video VideoMark: A Distortion-Free Robust Watermarking Framework for Video Diffusion Models

🎯 Evaluation Module

Evaluation Pipelines

MarkDiffusion supports eight pipelines, two for detection (WatermarkedMediaDetectionPipeline and UnWatermarkedMediaDetectionPipeline), and six for quality analysis. The table below details the quality analysis pipelines.

Quality Analysis Pipeline Input Type Required Data Applicable Metrics
DirectImageQualityAnalysisPipeline Single image Generated watermarked/unwatermarked image Metrics for single image evaluation
ReferencedImageQualityAnalysisPipeline Image + reference content Generated watermarked/unwatermarked image + reference image/text Metrics requiring computation between single image and reference content (text/image)
GroupImageQualityAnalysisPipeline Image set (+ reference image set) Generated watermarked/unwatermarked image set (+reference image set) Metrics requiring computation on image sets
RepeatImageQualityAnalysisPipeline Image set Repeatedly generated watermarked/unwatermarked image set Metrics for evaluating repeatedly generated image sets
ComparedImageQualityAnalysisPipeline Two images for comparison Generated watermarked and unwatermarked images Metrics measuring differences between two images
DirectVideoQualityAnalysisPipeline Single video Generated video frame set Metrics for overall video evaluation

Evaluation Tools

Tool Name Evaluation Category Function Description Output Metrics
FundamentalSuccessRateCalculator Detectability Calculate classification metrics for fixed-threshold watermark detection Various classification metrics
DynamicThresholdSuccessRateCalculator Detectability Calculate classification metrics for dynamic-threshold watermark detection Various classification metrics
Image Attack Tools
Rotation Robustness (Image) Image rotation attack, testing watermark resistance to rotation transforms Rotated images/frames
CrSc (Crop & Scale) Robustness (Image) Cropping and scaling attack, evaluating watermark robustness to size changes Cropped/scaled images/frames
GaussianNoise Robustness (Image) Gaussian noise attack, testing watermark resistance to noise interference Noise-corrupted images/frames
GaussianBlurring Robustness (Image) Gaussian blur attack, evaluating watermark resistance to blur processing Blurred images/frames
JPEGCompression Robustness (Image) JPEG compression attack, testing watermark robustness to lossy compression Compressed images/frames
Brightness Robustness (Image) Brightness adjustment attack, evaluating watermark resistance to brightness changes Brightness-modified images/frames
Mask Robustness (Image) Image masking attack, testing watermark resistance to partial occlusion by random black rectangles Masked images/frames
Overlay Robustness (Image) Image overlay attack, testing watermark resistance to graffiti-style strokes and annotations Overlaid images/frames
AdaptiveNoiseInjection Robustness (Image) Adaptive noise injection attack, testing watermark resistance to content-aware noise (Gaussian/Salt-pepper/Poisson/Speckle) Noisy images/frames with adaptive noise
Video Attack Tools
MPEG4Compression Robustness (Video) MPEG-4 video compression attack, testing video watermark compression robustness Compressed video frames
FrameAverage Robustness (Video) Frame averaging attack, destroying watermarks through inter-frame averaging Averaged video frames
FrameSwap Robustness (Video) Frame swapping attack, testing robustness by changing frame sequences Swapped video frames
FrameRateAdapter Robustness (Video) Frame rate conversion attack that resamples frames while preserving duration Resampled frame sequence
FrameInterpolationAttack Robustness (Video) Frame interpolation attack inserting blended frames to alter temporal density Interpolated video frames
Image Quality Analyzers
InceptionScoreCalculator Quality (Image) Evaluate generated image quality and diversity IS score
FIDCalculator Quality (Image) Fréchet Inception Distance, measuring distribution difference between generated and real images FID value
LPIPSAnalyzer Quality (Image) Learned Perceptual Image Patch Similarity, evaluating perceptual quality LPIPS distance
CLIPScoreCalculator Quality (Image) CLIP-based text-image consistency evaluation CLIP similarity score
PSNRAnalyzer Quality (Image) Peak Signal-to-Noise Ratio, measuring image distortion PSNR value (dB)
NIQECalculator Quality (Image) Natural Image Quality Evaluator, reference-free quality assessment NIQE score
SSIMAnalyzer Quality (Image) Structural Similarity Index between two images SSIM value
BRISQUEAnalyzer Quality (Image) Blind/Referenceless Image Spatial Quality Evaluator, evaluating perceptual quality of an image without requiring a reference BRISQUE score
VIFAnalyzer Quality (Image) Visual Information Fidelity analyzer, comparing a distorted image with a reference image to quantify the amount of visual information preserved VIF value
FSIMAnalyzer Quality (Image) Feature Similarity Index analyzer, comparing structural similarity between two images based on phase congruency and gradient magnitude FSIM value
Video Quality Analyzers
SubjectConsistencyAnalyzer Quality (Video) Evaluate consistency of subject objects in video Subject consistency score
BackgroundConsistencyAnalyzer Quality (Video) Evaluate background coherence and stability in video Background consistency score
MotionSmoothnessAnalyzer Quality (Video) Evaluate smoothness of video motion Motion smoothness metric
DynamicDegreeAnalyzer Quality (Video) Measure dynamic level and change magnitude in video Dynamic degree value
ImagingQualityAnalyzer Quality (Video) Comprehensive evaluation of video imaging quality Imaging quality score

🧩 Quick Start

Google Colab Demo

If you're interested in trying out MarkDiffusion without installing anything, you can use Google Colab to see how it works.

Installation

(Recommended) We released pypi package for MarkDiffusion. You can install it directly with pip:

conda create -n markdiffusion python=3.11
conda activate markdiffusion
pip install markdiffusion[optional]

(Alternative) For users who are restricted only to use conda environment, we also provide a conda-forge package, which can be installed with the following commands:

conda create -n markdiffusion python=3.11
conda activate markdiffusion
conda config --add channels conda-forge
conda config --set channel_priority strict
conda install markdiffusion

However, please note that some advanced features require additional packages that are not available on conda and cannot be included in the release. You will need to install those separately if necessary.

How to Use the Toolkit

After installation, there are two ways to use MarkDiffusion:

  1. Clone the repository to try the demos or use it for custom development. The MarkDiffusion_demo.ipynb notebook offers detailed demonstrations for various use cases—please review it for guidance. Here’s a quick example of generating and detecting watermarked image with the TR algorithm:

    import torch
    from watermark.auto_watermark import AutoWatermark
    from utils.diffusion_config import DiffusionConfig
    from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
    
    # Device setup
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    
    # Configure diffusion pipeline
    scheduler = DPMSolverMultistepScheduler.from_pretrained("model_path", subfolder="scheduler")
    pipe = StableDiffusionPipeline.from_pretrained("model_path", scheduler=scheduler).to(device)
    diffusion_config = DiffusionConfig(
        scheduler=scheduler,
        pipe=pipe,
        device=device,
        image_size=(512, 512),
        num_inference_steps=50,
        guidance_scale=7.5,
        gen_seed=42,
        inversion_type="ddim"
    )
    
    # Load watermark algorithm
    watermark = AutoWatermark.load('TR', 
                                algorithm_config='config/TR.json',
                                diffusion_config=diffusion_config)
    
    # Generate watermarked media
    prompt = "A beautiful sunset over the ocean"
    watermarked_image = watermark.generate_watermarked_media(prompt)
    watermarked_image.save("watermarked_image.png")
    
    # Detect watermark
    detection_result = watermark.detect_watermark_in_media(watermarked_image)
    print(f"Watermark detected: {detection_result}")
  2. Import markdiffusion library directly in your code without cloning the repository. The MarkDiffusion_pypi_demo.ipynb notebook provides comprehensive examples for using MarkDiffusion via the markdiffusion library——please review it for guidance. Here's a quick example:

    import torch
    from markdiffusion.watermark import AutoWatermark
    from markdiffusion.utils import DiffusionConfig
    from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
    
    # Device
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"Using device: {device}")
    
    # Model path
    MODEL_PATH = "huanzi05/stable-diffusion-2-1-base"
    
    # Initialize scheduler and pipeline
    scheduler = DPMSolverMultistepScheduler.from_pretrained(MODEL_PATH, subfolder="scheduler")
    pipe = StableDiffusionPipeline.from_pretrained(
        MODEL_PATH,
        scheduler=scheduler,
        torch_dtype=torch.float16 if device == "cuda" else torch.float32,
        safety_checker=None,
    ).to(device)
    
    # Create DiffusionConfig for image generation
    image_diffusion_config = DiffusionConfig(
        scheduler=scheduler,
        pipe=pipe,
        device=device,
        image_size=(512, 512),
        guidance_scale=7.5,
        num_inference_steps=50,
        gen_seed=42,
        inversion_type="ddim"
    )
    
    # Load Tree-Ring watermark algorithm
    tr_watermark = AutoWatermark.load('TR', diffusion_config=image_diffusion_config)
    print("TR watermark algorithm loaded successfully!")
    
    # Generate watermarked image
    prompt = "A beautiful landscape with mountains and a river at sunset"
    
    watermarked_image = tr_watermark.generate_watermarked_media(input_data=prompt)
    
    # Display the watermarked image
    watermarked_image.save("watermarked_image.png")
    print("Watermarked image generated!")
    
    # Detect watermark in the watermarked image
    detection_result = tr_watermark.detect_watermark_in_media(watermarked_image)
    print("Watermarked image detection result:")
    print(detection_result)

🛠 Test Modules

We provide a comprehensive set of test modules to ensure the quality of the code. The module includes 658 unit tests of 95% code coverage. Please refer to the test/ directory for more details. Here are the full coverage report and the result report directly exported via pytest.

Citation

@article{pan2025markdiffusion,
  title={MarkDiffusion: An Open-Source Toolkit for Generative Watermarking of Latent Diffusion Models},
  author={Pan, Leyi and Guan, Sheng and Fu, Zheyu and Si, Luyang and Wang, Zian and Hu, Xuming and King, Irwin and Yu, Philip S and Liu, Aiwei and Wen, Lijie},
  journal={arXiv preprint arXiv:2509.10569},
  year={2025}
}

About

MarkDiffusion: An Open-Source Toolkit for Generative Watermarking of Latent Diffusion Models

Topics

Resources

License

Code of conduct

Contributing

Security policy

Stars

Watchers

Forks

Packages

No packages published

Contributors 6