Clone this repo:
git clone https://github.com/ZYM-PKU/UDiffText.git
cd UDiffText
Install required Python packages
conda create -n udiff python=3.11
conda activate udiff
pip install torch==2.1.1 torchvision==0.16.1 --index-url https://download.pytorch.org/whl/cu121
pip install -r requirements.txt
mkdir ./checkpoints
checkpoints
โโโ AEs // AutoEncoder
โโโ encoders
โโโ LabelEncoder // Character-level encoder
โโโ ViTSTR // STR encoder
โโโ predictors // STR model
โโโ pretrained // Pretrained SD
โโโ ***.ckpt // UDiffText checkpoint
ICDAR13
โโโ train // training set
โโโ annos // annotations
โโโ gt_x.txt
โโโ ...
โโโ images // images
โโโ img_x.jpg
โโโ ...
โโโ val // validation set
โโโ annos // annotations
โโโ gt_img_x.txt
โโโ ...
โโโ images // images
โโโ img_x.jpg
โโโ ...
TextSeg
โโโ train // training set
โโโ annotation // annotations
โโโ x_anno.json // annotation json file
โโโ x_mask.png // character-level mask
โโโ ...
โโโ image // images
โโโ x.jpg.jpg
โโโ ...
โโโ val // validation set
โโโ annotation // annotations
โโโ x_anno.json // annotation json file
โโโ x_mask.png // character-level mask
โโโ ...
โโโ image // images
โโโ x.jpg
โโโ ...
SynthText
โโโ 1 // part 1
โโโ ant+hill_1_0.jpg // image
โโโ ant+hill_1_1.jpg
โโโ ...
โโโ 2 // part 2
โโโ ...
โโโ gt.mat // annotation file
Set the parameters in ./configs/pretrain.yaml and run:
python pretrain.py
Download the pretrained model and put it in ./checkpoints/pretrained/. You can ignore the "Missing Key" or "Unexcepted Key" warning when loading the checkpoint.
Set the parameters in ./configs/train.yaml, especially the paths:
load_ckpt_path: ./checkpoints/pretrained/512-inpainting-ema.ckpt // Checkpoint of the pretrained SD
model_cfg_path: ./configs/train/textdesign_sd_2.yaml // UDiffText model config
dataset_cfg_path: ./configs/dataset/locr.yaml // Use the Laion-OCR dataset
and run:
python train.py
Download our available checkpoints and put them in the corresponding directories in ./checkpoints.
Set the parameters in ./configs/test.yaml, especially the paths:
load_ckpt_path: "./checkpoints/***.ckpt" // UDiffText checkpoint
model_cfg_path: "./configs/test/textdesign_sd_2.yaml" // UDiffText model config
dataset_cfg_path: "./configs/dataset/locr.yaml" // LAION-OCR dataset config
and run:
python test.py
In order to run an interactive demo on your own machine, execute the code:
python demo.py
or try our online demo at hugging face:
Dataset: We sincerely thank the open-source large image-text dataset LAION-OCR with character-level segmentations provided by TextDiffuser.
Code & Model: We build our project based on the code repo of Stable Diffusion XL and leverage the pretrained checkpoint of Stable Diffusion 2.0.
@misc{zhao2023udifftext,
title={UDiffText: A Unified Framework for High-quality Text Synthesis in Arbitrary Images via Character-aware Diffusion Models},
author={Yiming Zhao and Zhouhui Lian},
year={2023},
eprint={2312.04884},
archivePrefix={arXiv},
primaryClass={cs.CV}
}