# StudioGAN + Recursive Token Mapper (RTM)

StyleGAN2-ADA + RTM training and evaluation pipeline used for the AFHQ-v1
(512x512) and CIFAR-10 (32x32) StyleGAN baselines in the paper.

The whole [PyTorch-StudioGAN](https://github.com/POSTECH-CVLab/PyTorch-StudioGAN)
codebase is vendored under `src/` (see `LICENSE-StudioGAN` and
`LICENSE-NVIDIA` for the upstream licenses); 
## Setup

Use the same Python environment as the rest of this repository (PyTorch
≥ 2.1, CUDA ≥ 12, plus StudioGAN's extra deps below):

```bash
pip install h5py kornia==0.6.5 prdc seaborn imageio scikit-learn
```

Datasets are loaded with StudioGAN's stock loader. For AFHQ-v1 set
`DATA_DIR` to a folder with `train/{cat,dog,wild}` and `valid/{cat,dog,wild}`
sub-directories. For CIFAR-10 set `DATA_DIR` to the standard CIFAR-10 folder.

## Train

```bash
DATA_DIR=/path/to/AFHQ      bash scripts/train_afhq.sh
DATA_DIR=/path/to/cifar-10  bash scripts/train_cifar10.sh
```

Both scripts read RTM hyper-parameters from the YAML config (see the
`STYLEGAN.use_rtm_mapper` block in each config). The headline numbers in the
paper use the values shipped in those YAMLs:

| Knob                    | AFHQ (512) | CIFAR-10 (32) |
| ----------------------- | ---------- | ------------- |
| `rtm_num_tokens`        | 4 (a 2x2 grid) | 4 (a 2x2 grid) |
| `rtm_H_cycles`          | 4          | 4             |
| `rtm_L_cycles`          | 1          | 1             |
| `rtm_H_layers`          | 2          | 2             |
| `rtm_L_layers`          | 2          | 2             |
| `rtm_hidden_size`       | 256        | 256           |
| `rtm_refinement_steps`  | 4          | 4             |
| `use_rtm_equalized`     | True       | True          |
| `rtm_lr_multiplier`     | 0.01       | 0.01          |


## Evaluate

Point `CKPT_DIR` at the StudioGAN-format `checkpoints/<exp>/` folder produced
by the training run:

```bash
CKPT_DIR=runs/afhq_rtm/checkpoints/<exp>     DATA_DIR=/path/to/AFHQ      bash scripts/eval_afhq.sh
CKPT_DIR=runs/cifar10_rtm/checkpoints/<exp>  DATA_DIR=/path/to/cifar-10  bash scripts/eval_cifar10.sh
```