# Data Provenance in IARS - Infinity

##  Installation

###  Environment

Set up an environment as follows by installing the requirements.txt file.

All experiments have been conducted on infinity_2b_reg with the vae infinity_vae_d32. These model weights (and other) can be downloaded from huggingface at <a href='https://huggingface.co/FoundationVision/infinity'><img src='https://img.shields.io/badge/%F0%9F%A4%97%20weights-FoundationVision/Infinity-yellow'></a> 

If required, further instructions can be obtained from the [official Infinity repository](https://github.com/FoundationVision/Infinity/tree/main).

## Files & Execution

Our main scripts can be found in tools/{codebook.py, finetune_vae.py, latenttracer.py, run_infinity.py}. 

### Data Generation

Generating data and preapring for it for finetuning can be done by executing tools/run_infinity.py:

```bash
PYTHONPATH=. python3 tools/run_infinity.py \
    --pn 1M \
    --model_path weights/infinity_2b_reg.pth \
    --cfg_insertion_layer 0 \
    --vae_type 32 \
    --vae_path weights/infinity_vae_d32_reg.pth \
    --add_lvl_embeding_only_first_block 1 \
    --use_bit_label 1 \
    --model_type infinity_2b \
    --rope2d_each_sa_layer 1 \
    --rope2d_normalized_by_hw 2 \
    --use_scale_schedule_embedding 0 \
    --sampling_per_bits 1 \
    --text_encoder_ckpt weights/flan-t5-xl \
    --text_channels 2048 \
    --apply_spatial_patchify 0 \
    --h_div_w_template 1.000 \
    --use_flex_attn 0 \
    --cache_dir /dev/shm \
    --checkpoint_type torch \
    --seed 42 \
    --bf16 1 \
    --prompt_txt_file path/to/file \
    --save_folder path_to_folder
```

### Finetuning

For finetuning, execute tools/finetune_vae.py:
```bash
PYTHONPATH=. python3 tools/finetune_vae.py \
    --pn 1M \
    --model_path weights/infinity_2b_reg.pth \
    --cfg_insertion_layer 0 \
    --vae_type 32 \
    --vae_path weights/infinity_vae_d32_reg.pth \
    --add_lvl_embeding_only_first_block 1 \
    --use_bit_label 1 \
    --model_type infinity_2b \
    --rope2d_each_sa_layer 1 \
    --rope2d_normalized_by_hw 2 \
    --use_scale_schedule_embedding 0 \
    --sampling_per_bits 1 \
    --text_encoder_ckpt weights/flan-t5-xl \
    --text_channels 2048 \
    --apply_spatial_patchify 0 \
    --h_div_w_template 1.000 \
    --use_flex_attn 0 \
    --cache_dir /dev/shm \
    --checkpoint_type torch \
    --seed 666 \
    --bf16 1 \
    --finetune-lr 1e-5 \
    --epochs 10 \
    --batch-size 2 \
    --seed 42 \
    --output-dir  path_to_folder \
    --preprocessed-dir  path_to_folder \
```

### Data Provenance - Losses Calculation

For calculating our losses and baselines, execute tools/codebook.py:
```bash
PYTHONPATH=. python3 tools/codebook.py \
    --pn 1M \
    --model_path weights/infinity_2b_reg.pth \
    --cfg_insertion_layer 0 \
    --vae_type 32 \
    --vae_path weights/infinity_vae_d32_reg.pth \
    --add_lvl_embeding_only_first_block 1 \
    --use_bit_label 1 \
    --model_type infinity_2b \
    --rope2d_each_sa_layer 1 \
    --rope2d_normalized_by_hw 2 \
    --use_scale_schedule_embedding 0 \
    --sampling_per_bits 1 \
    --text_encoder_ckpt weights/flan-t5-xl \
    --text_channels 2048 \
    --apply_spatial_patchify 0 \
    --h_div_w_template 1.000 \
    --use_flex_attn 0 \
    --cache_dir /dev/shm \
    --checkpoint_type torch \
    --seed 42 \
    --bf16 1 \
    --dataset_config ./dataset_config.json \
    --ft_path path/to/finetuned_encoder.pt # or empty string for original checkpoint
```

### Latenttracer Baseline 

To calculate results with latentteacer, execute tools/latenttracer.py:
```bash
PYTHONPATH=. python3 tools/latenttracer.py \
    --pn 1M \
    --model_path weights/infinity_2b_reg.pth \
    --cfg_insertion_layer 0 \
    --vae_type 32 \
    --vae_path weights/infinity_vae_d32_reg.pth \
    --add_lvl_embeding_only_first_block 1 \
    --use_bit_label 1 \
    --model_type infinity_2b \
    --rope2d_each_sa_layer 1 \
    --rope2d_normalized_by_hw 2 \
    --use_scale_schedule_embedding 0 \
    --sampling_per_bits 1 \
    --text_encoder_ckpt weights/flan-t5-xl \
    --text_channels 2048 \
    --apply_spatial_patchify 0 \
    --h_div_w_template 1.000 \
    --use_flex_attn 0 \
    --cache_dir /dev/shm \
    --checkpoint_type torch \
    --seed 42 \
    --bf16 1 \
    --num_samples 1000 \
    --batch_size 2 \
    --dataset_config ./dataset_config.json \
    --save_folder path/to/folder
```

### Setup Datasets
    
Datasets are supposed to be passed via a separate dataset_config.json file with the following format:
```json
{
    "Name": "Path to dataset"
}
```
