Update README with links, add citation file

.gitignore

@@ -18,8 +18,6 @@ _temp
 
 sd2_decoder.pth
 
-*.png
-
 *.csv
 *.pt
 

CITATION.cff

@@ -0,0 +1,24 @@
+cff-version: 1.2.0
+message: If you use this software, please cite both the article from preferred-citation and the software itself.
+type: software
+title: 'METR: Image Watermarking with Large Number of Unique Messages'
+version: 1.0.0
+authors:
+  - given-names: Alexander
+    family-names: Varlamov
+    affiliation: MIPT, VK Lab
+    email: [email protected]
+  - given-names: Daria
+    family-names: Diatlova
+    affiliation: deepvk
+    email: [email protected]
+  - given-names: Egor
+    family-names: Spirin
+    affiliation: VK Lab
+    email: [email protected]
+identifiers:
+  - type: url
+    value: 'https://arxiv.org/abs/2408.08340'
+    description: arXiv preprint
+repository-code: 'https://github.com/deepvk/metr'
+license: Apache-2.0

README.md

@@ -1,29 +1,39 @@
-# 🚜 METR: Message Enhanced Tree-Ring Watermarking
+# 🚜 METR: Image Watermarking with Large Number of Unique Messages
 
-This is the implementation of METR: attack-resistant watermarks with support of a large number of unique messages,
+<img src="./schema/metr.png" style="max-width: 100%">
+
+`METR` is an attack-resistant watermarks with support of numerous unique messages,
 but without image quality reduction.
 
-Our contribution
-- METR: improved Tree-Ring algorithm to support large number of unique messages for any diffusion model
-- METR++: special extension for LDM that integrates both METR and StableSignature watermarks
+This includes:
+- `METR`: Message Enhanced Tree-Ring algorithm to support large number of unique messages for any diffusion model
+- `METR++`: special extension for LDM that integrates both `METR` and `StableSignature` watermarks
+
+Our paper can be found at [**arXiv**](https://www.arxiv.org/abs/2408.08340)
+
+<details>
+<summary>Abstract</summary>
+
+> Improvements in diffusion models have boosted the quality of image generation, which has led researchers, companies, and creators to focus on improving watermarking algorithms. This provision would make it possible to clearly identify the creators of generative art. The main challenges that modern watermarking algorithms face have to do with their ability to withstand attacks and encrypt many unique messages, such as user IDs. In this paper, we present METR: Message Enhanced Tree-Ring, which is an approach that aims to address these challenges. METR is built on the Tree-Ring watermarking algorithm, a technique that makes it possible to encode multiple distinct messages without compromising attack resilience or image quality. This ensures the suitability of this watermarking algorithm for any Diffusion Model. In order to surpass the limitations on the quantity of encoded messages, we propose METR++, an enhanced version of METR. This approach, while limited to the Latent Diffusion Model architecture, is designed to inject a virtually unlimited number of unique messages. We demonstrate its robustness to attacks and ability to encrypt many unique messages while preserving image quality, which makes METR and METR++ hold great potential for practical applications in real-world settings. 
+
+</details>
 
 ## Table of Content
 
 <!-- TOC -->
-* [🚜 METR: Message Enhanced Tree-Ring Watermarking](#-metr-message-enhanced-tree-ring-watermarking)
-  * [Table of Content](#table-of-content)
-  * [Setup:](#setup)
-  * [Running METR:](#running-metr)
-    * [Generate images with a METR watermark:](#generate-images-with-a-metr-watermark)
-    * [Performs attacks](#performs-attacks)
-    * [Evaluate FID](#evaluate-fid)
-  * [Running METR++](#running-metr-1)
-    * [Fine-tune VAE decoder to given ID:](#fine-tune-vae-decoder-to-given-id)
-    * [Generate images with METR++ watermark:](#generate-images-with-metr-watermark)
-    * [Evaluate METR++](#evaluate-metr)
-  * [Reproducing](#reproducing)
-  * [References:](#references)
-  * [Citation](#citation)
+- [🚜 METR: Image Watermarking with Large Number of Unique Messages](#-metr-image-watermarking-with-large-number-of-unique-messages)
+  - [Table of Content](#table-of-content)
+  - [Setup:](#setup)
+  - [Running `METR`](#running-metr)
+    - [Perform attacks](#perform-attacks)
+    - [Evaluate FID](#evaluate-fid)
+  - [Running `METR++`](#running-metr-1)
+    - [Fine-tune VAE decoder to given ID:](#fine-tune-vae-decoder-to-given-id)
+    - [Generate images with `METR++` watermark:](#generate-images-with-metr-watermark)
+    - [Evaluate METR++](#evaluate-metr)
+  - [Reproducing](#reproducing)
+  - [References:](#references)
+  - [Citation](#citation)
 <!-- TOC -->
 
 ## Setup:
@@ -53,11 +63,10 @@ wandb login $WANDB_KEY
 accelerate config default
 ```
 
-## Running METR:
+## Running `METR`
 
-### Generate images with a METR watermark:
-
-1. Generate with a random message
+<details>
+<summary>Generate image with a random message</summary>
 
 To save images locally include additional argument `--save_locally` and provide path with `--local_path /path/to/save`.
 
@@ -78,7 +87,10 @@ accelerate launch -m metr.run_metr \
   --local_path generated_images
 ```
 
-2. Generate with a fixed message and fixed prompt:
+</details>
+
+<details>
+<summary>Generate image with a fixed message and fixed prompt</summary>
 
 ```bash
 accelerate launch -m metr.run_metr \
@@ -96,9 +108,12 @@ accelerate launch -m metr.run_metr \
   --given_prompt "sci-fi bernese mountain dog"
 ```
 
-### Performs attacks
+</details>
+
+### Perform attacks
 
-Perform diffusion model attack on METR:
+<details>
+<summary>Perform diffusion model attack on `METR`</summary>
 
 ```bash
 accelerate launch -m metr.run_metr \
@@ -117,17 +132,23 @@ accelerate launch -m metr.run_metr \
   --diff_attack_steps 150
 ```
 
+</details>
+
 ### Evaluate FID
 
-1. Download dataset and extract:
+<details>
+<summary>1. Download dataset and extract:</summary>
 
 ```bash
 wget -O fid_outputs.zip https://huggingface.co/datasets/Alphonsce/MSCOCO_zip/resolve/main/fid_outputs.zip?download=true
 unzip fid_outputs.zip
 ```
 [Google drive version if HF is not working](https://drive.google.com/drive/u/1/folders/1v0xj-8Yx8vZ_4qGsC5EU5FJBvEWTFmE3)
 
-2. FID on ground-truth images (FID gt):
+</details>
+
+<details>
+<summary>2. FID on ground-truth images (FID gt):</summary>
 
 Argument `--image_folder` is where images are saved.
 
@@ -148,7 +169,10 @@ accelerate launch -m metr.run_metr_fid \
   --msg_scaler 100
 ```
 
-3. FID on generated images (FID gen):
+</details>
+
+<details>
+<summary>3. FID on generated images (FID gen)</summary>
 To evaluate FID on generated images, add argument `--target_clean_generated`:
 
 ```bash
@@ -169,9 +193,11 @@ accelerate launch -m metr.run_metr_fid \
   --target_clean_generated
 ```
 
-## Running METR++
+</details>
+
+## Running `METR++`
 
-We forked Stable Signature repository to adjust it to be comparable with METR. It can be found [here](https://github.com/Alphonsce/stable_signature/).
+We forked Stable Signature repository to adjust it to be comparable with `METR`. It can be found [here](https://github.com/Alphonsce/stable_signature/).
 
 Install weights for WM extractor for Stable Signature (taken from [official Stable-Signature repository](https://github.com/facebookresearch/stable_signature) )
 ```bash
@@ -195,10 +221,12 @@ wget -O ldm_decoder_checkpoint_000.pth https://huggingface.co/Alphonsce/St_Sig_v
 
 ### Fine-tune VAE decoder to given ID:
 
-In example down below we fine-tune VAE decoder on samples from MSCOCO dataset and evaluate on images previously generated with METR watermark in `generated_images` folder.
+In example down below we fine-tune VAE decoder on samples from MSCOCO dataset and evaluate on images previously generated with `METR` watermark in `generated_images` folder.
 Pretrained VAE decoder weights will be saved in `finetune_ldm_decoder/ldm_decoder_checkpoint_000.pth` by default.
 You can change the name of checkpoint with `--checkpoint_name` argument.
 
+<details>
+
 ```bash
 TRAIN_DIR=fid_outputs/coco/ground_truth
 VAL_DIR=generated_images/imgs_w
@@ -220,9 +248,13 @@ VAL_DIR=generated_images/imgs_w
   --key_str 111010110101000001010111010011010100010000100111
 ```
 
-### Generate images with METR++ watermark:
+</details>
+
+### Generate images with `METR++` watermark:
+
+To generate images with `METR++` watermark, just remove `--no_stable_sig` argument and provide a path to tuned VAE decoder: `--decoder_state_dict_path /path/to/decoder/weights`:
 
-To generate images with METR++ watermark, just remove `--no_stable_sig` argument and provide a path to tuned VAE decoder: `--decoder_state_dict_path /path/to/decoder/weights`:
+<details>
 
 ```bash
 VAE_DECODER_PATH=finetune_ldm_decoder/ldm_decoder_checkpoint_000.pth
@@ -245,11 +277,14 @@ accelerate launch -m metr.run_metr \
   --decoder_state_dict_path $VAE_DECODER_PATH
 ```
 
+</details>
+
 ### Evaluate METR++
 
-1. Evaluate Stable Signature part of METR++
+<details>
+<summary>1. Evaluate Stable Signature part of METR++</summary>
 
-Evaluation is performed on a folder of generated images, you need to pass folder into `--` with images generated with Stable-Signature watermark.
+Evaluation is performed on a folder of generated images, you need to pass folder into `--` with images generated with `Stable-Signature` watermark.
 
 ```bash
 EVAL_FOLDER=metr_pp_generated_images/imgs_w
@@ -266,9 +301,12 @@ accelerate launch -m metr.metr_pp_eval_stable_sig \
   --key_str 111010110101000001010111010011010100010000100111
 ```
 
-2. Evaluate FID for METR++:
+</details>
+
+<details>
+<summary>2. Evaluate FID for METR++</summary>
 
-To evaluate FID for images with METR++ watermark pass `--use_stable_sig` argument.
+To evaluate FID for images with `METR++` watermark pass `--use_stable_sig` argument.
 ```bash
 accelerate launch -m metr.run_metr_fid \
   --project_name fid_gen \
@@ -288,16 +326,28 @@ accelerate launch -m metr.run_metr_fid \
   --use_stable_sig
 ```
 
+</details>
+
 ## Reproducing
 
 To reproduce all experiments from our paper see [`reproducibility.md`](./reproducibility.md)
 
 ## References:
 
-1. Tree-Ring watermark ([Repository](https://github.com/YuxinWenRick/tree-ring-watermark), [Paper](https://arxiv.org/abs/2305.20030))
+1. Tree-Ring Watermark ([Repository](https://github.com/YuxinWenRick/tree-ring-watermark), [Paper](https://arxiv.org/abs/2305.20030))
 2. Stable Signature ([Repository](https://github.com/facebookresearch/stable_signature), [Paper](https://arxiv.org/abs/2303.15435))
-3. Generative Model watermark attacker ([Repository](https://github.com/XuandongZhao/WatermarkAttacker), [Paper](https://arxiv.org/abs/2306.01953))
+3. Generative Model Watermark Attacker ([Repository](https://github.com/XuandongZhao/WatermarkAttacker), [Paper](https://arxiv.org/abs/2306.01953))
 
 ## Citation
 
-TODO
+```
+@misc{deepvk2024metr,
+      title={METR: Image Watermarking with Large Number of Unique Messages}, 
+      author={Alexander Varlamov and Daria Diatlova and Egor Spirin},
+      year={2024},
+      eprint={2408.08340},
+      archivePrefix={arXiv},
+      primaryClass={cs.LG},
+      url={https://arxiv.org/abs/2408.08340}, 
+}
+```