# Data Provenance in IARS

##  Installation

###  Environment

Set up a conda environment as follows:
```bash
conda create --name wmar python=3.12
conda activate wmar
```

Install xformers (which will include Torch 2.7.0 CUDA 12.6) and other dependencies, and override the triton version.
```bash
pip install -U xformers --index-url https://download.pytorch.org/whl/cu126
pip install -r requirements.txt
pip install triton==3.1.0 
```

We next describe how to load all autoregressive models, finetuned tokenizer deltas, and other requirements.
We assume that all checkpoints will be placed under `checkpoints/`.

### Autoregressive Models

Instructions to download each of the three models evaluated in the paper are given below.

- **Taming**. You need to manually download the transfomer and VQGAN weights following the instructions from the [official repo](https://github.com/CompVis/taming-transformers). In particular, download `cin_transformer` from https://app.koofr.net/links/90cbd5aa-ef70-4f5e-99bc-f12e5a89380e and `VQGAN ImageNet (f=16), 16384` from https://heibox.uni-heidelberg.de/d/a7530b09fed84f80a887/ and set up the following folder structure under e.g., `checkpoints/2021-04-03T19-39-50_cin_transformer`:
  ```
  checkpoints/
      net2net.ckpt 
      vqgan.ckpt
  configs/
      net2net.yaml
      vqgan.yaml
  ```
  This directory should be also set as `--modeldir` when executing the code (see below).
  To adapt the model configs to the paths in our codebase execute:
  ```bash
  sed -i 's/ taming\./ deps.taming./g' checkpoints/2021-04-03T19-39-50_cin_transformer/configs/vqgan.yaml
  sed -i 's/ taming\./ deps.taming./g' checkpoints/2021-04-03T19-39-50_cin_transformer/configs/net2net.yaml
  ```

- **RAR**. RAR-XL is downloaded automatically on the first run; set `--modelpath` to the directory where you want to save the tokenizer and model weights, e.g., `checkpoints/rar`.

## Usage

### Large-scale generation and evaluation
We describe how to start a larger generation run and the follow-up evaluation and plotting that follows our experimental setup from the paper and reproduces our main results.
We focus on the Taming model. 
Before starting make sure to follow the relevant parts of the setup above.

We generate 50.000 images for finetuning using  `generate.py`. 

```bash
python3 generate.py --seed 1 --model taming \
--decoder_ft_ckpt none \
--encoder_ft_ckpt none  \
--modelpath checkpoints/2021-04-03T19-39-50_cin_transformer/ \
--wam False \
--wm_method none --wm_seed_strategy linear --wm_delta 2 --wm_gamma 0.25 \
--wm_context_size 1 --wm_split_strategy stratifiedrand \
--include_diffpure False --include_neural_compress False \
--orig_only True \
--top_p 0.92 --temperature 1.0 --top_k 250 --batch_size 10 \
--conditioning none \
--num_samples_per_conditioning 5 \
--chunk_id 0 --num_chunks 1 \
--outdir taming/taming_gen_images_original_decoder
```

### Finetuning

Run `finetune.py` using arguments such as documented in `scripts/finetune_taming_encoder.sh`. For Taming, an example command that runs finetuning is:

```bash
srun python finetune.py \
--model taming \
--modelpath checkpoints/2021-04-03T19-39-50_cin_transformer/ \
--dataset codes-imagenet \
--datapath taming/taming_gen_images_original_decoder_50k/codes \
--dataset_size 50000 \
--loss hard-to-soft-with-ae-no-reg \
--nb_epochs 50 \
--lr 0.0005 \
--n_step_lr 2 \
--augs none \
--augs_schedule 5,5,20,20 \
--batch_size_per_gpu 8 \
--idempotence_loss_weight 1.0 \
--outdir taming/finetune/taming_gen_10k_bs8_lr5e-4_step2_reg_no_aug552020/ \

```

The finetuning script also downloads the LPIPS checkpoint to `checkpoints/lpips` automatically (needed for perceptual loss).
Final checkpoints will be saved under `outdir` and can be used in evaluation by setting `encoder_ft_ckpt` and `decoder_ft_ckpt` flags as above. 

### Robustness

To evaluate the robustness of our provenance, you can use `scripts/robustness_script.sh` and set the model path as needed. The --clean_dir is the directory to the non-belonging images and the --watermarked_dir the path to the belonging images. 

## Acknowledgements

This code is based on the Official Implementation of [Watermark Autoregressive Image Generation](https://github.com/facebookresearch/wmar)

