# Score-Based Image-to-Image Brownian Bridge
This is the official implementation for the paper Score-Based Image-to-Image Brownian Bridge.

![](res/structure.png)

* For training and evaluating examples, see [Example Scripts Usage](#example-scripts-usage).
* For pre-trained models including VQGAN and DeepLabV3, see [Pre-trained Models](#pre-trained-models).

## Example Scripts Usage
This section describes how to use the example scripts to train and evaluate Score-Based Image-to-Image Brownian Bridge.

### Dependencies for example scripts

Ubuntu 22.04+ is recommended for running on CUDA devices. All the required packages for example scripts can be installed with the following command:
```bash
pip install -r requirements.txt
```

### Install package
To run examples, install the package first before get into the `examples` folder. The following command installs the package in editable mode.

```bash
pip install -e .
cd examples
```

### Training Script
Use `train.py` to train a Score-Based Image-to-Image Brownian Bridge model. The script supports training in image space and latent space. The following examples show how to train the model in image space and latent space using edge2shoes dataset.

```bash
# train in image space
python train.py \
    edge2shoes \
    <data_dir> \
    <output_model_path> \
    -exp <experiment_name>

# train in latent space
python train.py \
    edge2shoes \
    <data_dir> \
    <output_model_path> \
    -exp <experiment_name> \
    -vq <vqgan_model_path>
```

If directory `experiments/<experiment_name>` is already existed, the script will be terminated. Use `--replace_experiment` to replace existing experiment and avoid the termination.

Specify `--ckpt_path <checkpoint_path>` to resume from a specific checkpoint.

Use `--show_verbose` to display the training progress bar for each epoch. Set `--device` as `cuda:<gpu_id>` to use specific GPU (GPU version of PyTorch is required). Use `-use_multi_gpus` without `--device` argument to use multiple GPUs.

### Evaluation Script
Use `eval.py` and `eval_miou.py` to evaluate a Score-Based Image-to-Image Brownian Bridge model. The script supports evaluation of checkpoints and pre-trained PyTorch model. The following examples show how to evaluate a checkpoint or a pre-trained PyTorch model with edge2shoes dataset using fast sampling method.

```bash
# evaluate checkpoint
python eval.py \
    edge2shoes \
    <data_dir> \
    <checkpoint_path> \
    --fast_sampling

# evaluate pre-trained PyTorch model
python eval.py \
    edge2shoes \
    <data_dir> \
    <model_path> \
    --fast_sampling \
    --t 1000 \\
    -vq <vqgan_model_path>
```

If directory `experiments/test.exp` is already existed, the script will be terminated. Use `--replace_experiment` to replace existing experiment and avoid the termination. 

Use `-r` to repeat the whole dataset for multiple times so that the number of sampled images matches the requirement of FID and LPIPS metrics.

Use `eval_miou.py` to evaluate the model with mIoU metric for cityscapes. The following examples show how to evaluate a checkpoint or a pre-trained PyTorch model using fast sampling method.

```bash
# evaluate checkpoint
python eval_miou.py \
    <deeplabv3_model_path> \
    <data_dir> \
    <checkpoint_path> \
    --fast_sampling

# evaluate pre-trained PyTorch model
python eval_miou.py \
    <deeplabv3_model_path> \
    <data_dir> \
    <model_path> \
    --fast_sampling \
    --t 1000 \\
    -vq <vqgan_model_path>
```

Again, use `--show_verbose` to display the training progress bar for each epoch. Set `--device` as `cuda:<gpu_id>` to use specific GPU. Use `-use_multi_gpus` without `--device` argument to use multiple GPUs.

### Generation Script
Use `generate.py` to generate images from a Score-Based Image-to-Image Brownian Bridge model. The script supports generation from checkpoints and pre-trained PyTorch model. The following examples show how to generate images from a checkpoint or a pre-trained PyTorch model with edge2shoes dataset using fast sampling method.

```bash
python generate.py \
    edge2shoes \
    <data_dir> \
    <checkpoint_path> \
    --fast_sampling
```

## Pre-trained Models
We used pre-trained VQGAN separated from the [official LDM OpenImage checkpoint](https://github.com/CompVis/latent-diffusion). We first exported the state dict of the VQGAN model (autoencoder) in the official checkpoints, then use `convert_vqgan.py` script to convert the state dict to our `vqgan.VQGAN` PyTorch model for easy loading. The following command shows how to convert the state dict to a PyTorch model.

```bash
python convert_vqgan.py \
    <state_dict_path> \
    <output_model_path>
```

We used the pre-trained deeplabv3+ to evaluate the mIoU on Cityscapes from [here](https://github.com/VainF/DeepLabV3Plus-Pytorch). Again, we convert the checkpoints in state dict to `deeplabv3.DeepLabV3` PyTorch model using `convert_deeplabv3.py` script for easy loading. The following command shows how to convert the state dict to a PyTorch model.

```bash
python convert_deeplabv3.py \
    <state_dict_path> \
    <output_model_path>
```
