Towards Aligned Layout Generation via Diffusion Model with Aesthetic Constraints

source code of the layout generation model, LACE.

1. Installation

1.1 Prepare environment

Install package for python 3.9 or later version:

conda create --name LACE python=3.9
conda activate LACE
python -m pip install -r requirements.txt

1.2 Checkpoints

Download the trained checkpoints for diffusion model and FID model at Hugging Face or through command line:

wget https://huggingface.co/datasets/puar-playground/LACE/resolve/main/model.tar.gz
wget https://huggingface.co/datasets/puar-playground/LACE/resolve/main/fid.tar.gz
tar -xvzf model.tar.gz
tar -xvzf fid.tar.gz

Model hyper-parameter:
for Publaynet: --dim_transformer 1024 --nhead 16 --nlayer 4 --feature_dim 2048
for Rico13 and Rico25: --dim_transformer 512 --nhead 16 --nlayer 4 --feature_dim 2048

1.3 Datasets

The datasets are also available at:

wget https://huggingface.co/datasets/puar-playground/LACE/resolve/main/datasets.tar.gz
tar -xvzf datasets.tar.gz

Alternatively, you can download from the source and prepare each dataset as following:

• PubLayNet: Download the labels.tar.gz and decompress to ./dataset/publaynet-max25/raw folder.
  
• Rico: Download the rico_dataset_v0.1_semantic_annotations.zip and decompress to ./dataset/rico25-max25/raw folder.
  

When the dataset is initialized for the first time, a new folder callled processed will be created at e.g., ./dataset/magazine-max25/processed containing the formatted dataset for future uses. Training split of smaller dataset: Rico and Magazine will be duplicated to reach a reasonable epoch size.

2. Testing

Run python script test.py to test. Please run python test.py -h to see detailed explaination.
For PubLayNet:

python test.py --dataset publaynet --experiment all --device cuda:0 --dim_transformer 1024 --nhead 16 --batch_size 2048 --beautify

For Rico:

python test.py --dataset rico25 --experiment all --device cuda:0 --dim_transformer 512 --nhead 16 --batch_size 2048 --beautify

3. Training

Run python script train_diffusion.py to train.
The script takes several command line arguments. Please run python train_diffusion.py -h to see detailed explaination.
Example command for training:

python train.py --device cuda:1 --dataset rico25 --no-load_pre --lr 1e-6 --n_save_epoch 10

Reference

@inproceedings{
    chen2024towards,
    title={Towards Aligned Layout Generation via Diffusion Model with Aesthetic Constraints},
    author={Jian Chen and Ruiyi Zhang and Yufan Zhou and Changyou Chen},
    booktitle={The Twelfth International Conference on Learning Representations},
    year={2024},
    url={https://openreview.net/forum?id=kJ0qp9Xdsh}
}