ArtBank: Artistic Style Transfer with Pre-trained Diffusion Model and Implicit Style Prompt Bank (AAAI2024)

Artistic style transfer aims to repaint the content image with the learned artistic style. Existing artistic style transfer methods can be divided into two categories: small model-based approaches and pre-trained large-scale model-based approaches. Small model-based approaches can preserve the content strucuture, but fail to produce highly realistic stylized images and introduce artifacts and disharmonious patterns; Pre-trained large-scale model-based approaches can generate highly realistic stylized images but struggle with preserving the content structure. To address the above issues, we propose ArtBank, a novel artistic style transfer framework, to generate highly realistic stylized images while preserving the content structure of the content images. Specifically, to sufficiently dig out the knowledge embedded in pre-trained large-scale models, an Implicit Style Prompt Bank (ISPB), a set of trainable parameter matrices, is designed to learn and store knowledge from the collection of artworks and behave as a visual prompt to guide pre-trained large-scale models to generate highly realistic stylized images while preserving content structure. Besides, to accelerate training the above ISPB, we propose a novel Spatial-Statistical-based self-Attention Module (SSAM). The qualitative and quantitative experiments demonstrate the superiority of our proposed method over state-of-the-art artistic style transfer methods. For details see the paper

(back to top)

Getting Started

Prerequisites

For packages, see environment.yaml.

conda env create -f environment.yaml
conda activate ldm

(back to top)

Installation

Clone the repo

git clone https://github.com/Jamie-Cheung/ArtBank.git

(back to top)

Train

Train ArtBank:

python main.py --base configs/stable-diffusion/v1-finetune.yaml
            -t 
            --actual_resume ./models/sd/sd-v1-4.ckpt
            -n <run_name> 
            --gpus 0, 
            --data_root /path/to/directory/with/images

python main.py --base configs/stable-diffusion/v1-finetune.yaml -t --actual_resume ./models/sd/sd-v1-4.ckpt -n 106 --gpus 0, --data_root ./image

See configs/stable-diffusion/v1-finetune.yaml for more options

Download the collection of artworks (https://drive.google.com/drive/folders/1_2jykbjVCF6SqJisvIt5-4fAFzVAj-F0?usp=drive_link) and put style images into ./images.

Download the pretrained Stable Diffusion Model and save it at ./models/sd/sd-v1-4.ckpt.

Test

Download the pretrained model [Artbank] (https://drive.google.com/drive/folders/1_2jykbjVCF6SqJisvIt5-4fAFzVAj-F0?usp=drive_link) You only need to change the line 177-178 with the pretrained model, then you can:

python test.py

News

We propose a new framework, called LSAST(IJCAI2024), which can generate more highly realistic artistic stylized images and preserve the content structure better. You can follow it from LSAST

(back to top)

Citation

@inproceedings{zhang2024artbank,
title={ArtBank: Artistic Style Transfer with Pre-trained Diffusion Model and Implicit Style Prompt Bank},
author={Zhang, Zhanjie and Zhang, Quanwei and Xing, Wei and Li, Guangyuan and Zhao, Lei and Sun, Jiakai and Lan, Zehua and Luan, Junsheng and Huang, Yiling 
and Lin, Huaizhong},
booktitle={Proceedings of the AAAI Conference on Artificial Intelligence},
volume={38},
number={7},
pages={7396--7404},
year={2024}
}

Contact

Please feel free to open an issue or contact us personally if you have questions, need help, or need explanations. Write to one of the following email addresses, and maybe put one other in the cc:

[email protected]

Thank you for my co-author Quanwei Zhang

(back to top)