StyleDrop

This is a copy of the unofficial PyTorch implementation.
A project that originates from the paper StyleDrop: Text-to-Image Generation in Any Style.

Difference to the project this is forked from

This project was tested to run under a v100 and L4 GPU using CUDA Version: 13.0

• The requirements have been updated
• code was changed to work under this GPU environment
• bug-fix: the original project only ever used one image as style reference (pull request was created for this fix)
• late-cleanup-commit: the original repository contained lots and lots of files, not needed if you only want a working pipeline

Also: This project was meant to compare the results of StyleDrop to StyleShot.

I used images of the Stylebench Database that was kinda part of the Styleshot Project on Github, as well as game screenshots - for example from the games "The Midnight Walk" and "Until Then" to compare the models.

See the results below.

And please note that StyleShot is a model that is meant to be used in Inference as it goes without any additional adapter training. It also only takes one style reference image which it is trying to get close to. You can find my (only slightly edited) copy of the StyleShot project here. Also note that for the shown results, the text driven demo of StyleShot has been used.

Data & Weights Preparation in order to train StyleDrop

First, download VQGAN from this link (from MAGE, thanks!), and put the downloaded VQGAN in assets/vqgan_jax_strongaug.ckpt.

Then, download the pre-trained checkpoints from this link to assets/ckpts for evaluation or to continue training for more iterations.

finally, prepare empty_feature by runnig command python extract_empty_feature.py

And the final directory structure is as follows:

.
├── assets
│   ├── ckpts
│   │   ├── cc3m-285000.ckpt
│   │   │   ├── lr_scheduler.pth
│   │   │   ├── nnet_ema.pth
│   │   │   ├── nnet.pth
│   │   │   ├── optimizer.pth
│   │   │   └── step.pth
│   │   └── imagenet256-450000.ckpt
│   │       ├── lr_scheduler.pth
│   │       ├── nnet_ema.pth
│   │       ├── nnet.pth
│   │       ├── optimizer.pth
│   │       └── step.pth
│   ├── fid_stats
│   │   ├── fid_stats_cc3m_val.npz
│   │   └── fid_stats_imagenet256_guided_diffusion.npz
│   ├── pipeline.png
└── └── vqgan_jax_strongaug.ckpt

Dependencies to work on a v100 or L4 GPU with CUDA 13

# great torch versions for different older gpu's
pip install torch==2.2.0 torchvision==0.17.0 torchaudio==2.2.0 --index-url https://download.pytorch.org/whl/cu118

# compatible with torch versions above and actually needed by this project
pip install numpy==1.26.4 scipy omegaconf

# listed by original project
pip install accelerate==0.12.0 absl-py ml_collections einops wandb ftfy==6.1.1 transformers==4.23.1 loguru webdataset==0.2.5 gradio

Train

All style data in the paper are placed in the data directory

1. Modify data/one_style.json (It should be noted that one_style.json and style data must be in the same directory), The format is file_name:[object,style]

{"image_03_05.jpg":["A bear","in kid crayon drawing style"]}

2. Training script as follows (you will find the code in run_training.sh)

#!/bin/bash
unset EVAL_CKPT
unset ADAPTER
export OUTPUT_DIR="output_dir/for/this/experiment"
accelerate launch --num_processes 8 --mixed_precision fp16 train_t2i_custom_v2.py --config=configs/custom.py

Inference

The pretrained style_adapter weights can be downloaded from 🤗 Hugging Face.

#!/bin/bash
export EVAL_CKPT="assets/ckpts/cc3m-285000.ckpt" 
export ADAPTER="path/to/your/style_adapter"

export OUTPUT_DIR="output/for/this/experiment"

accelerate launch --num_processes 8 --mixed_precision fp16 train_t2i_custom_v2.py --config=configs/custom.py

My Results

Styledrop Graffiti

Styleshot Graffiti

Styledrop "The Midnight Walk"

Styleshot "The Midnight Walk"

Styledrop "Until Then"

Styleshot "Until Then"

Conclusion

The original StyleShot paper reports that human evaluators preferred its stylized results more frequently, while the CLIP scores of StyleShot and StyleDrop were shown to be relatively similar. The experiments presented in this repository, however, provide a different perspective within the limited scope of this comparison.

Across the evaluated examples, StyleDrop's adapter training appeared to learn the target design more consistently. As training progressed, the generated samples visibly converged toward the reference style, capturing not only the overall color palette but also more of the structural and stylistic characteristics of the target design.

StyleShot in text-driven mode also adapted toward the reference style, but in these experiments the changes were more strongly reflected in the color distribution than in the overall visual appearance. While the influence of the reference image was clearly noticeable, the resulting images remained less similar to the target design than those produced by StyleDrop.

These observations are based on a small number of experiments and should not be interpreted as a general benchmark of either method. They simply reflect the behavior observed under the specific settings and examples evaluated in this project.

Citation

@article{sohn2023styledrop,
  title={StyleDrop: Text-to-Image Generation in Any Style},
  author={Sohn, Kihyuk and Ruiz, Nataniel and Lee, Kimin and Chin, Daniel Castro and Blok, Irina and Chang, Huiwen and Barber, Jarred and Jiang, Lu and Entis, Glenn and Li, Yuanzhen and others},
  journal={arXiv preprint arXiv:2306.00983},
  year={2023}
}

@article{gao2024styleshot,
  title={StyleShot: A Snapshot on Any Style},
  author={Gao, Junyao and Liu, Yanchen and Sun, Yanan and Tang, Yinhao and Zeng, Yanhong and Chen, Kai and Zhao, Cairong},
  journal={arXiv preprint arXiv:2407.01414},
  year={2024}
}

Further Acknowlegments

• The inofficial StyleDrop implementation is based on MUSE-PyTorch