/infusion

Infusion: Preventing Customized Text-to-Image Diffusion from Overfitting

Primary LanguagePython

Infusion

Infusion: Preventing Customized Text-to-Image Diffusion from Overfitting [Project Page] | [Paper] Abstract: Text-to-image (T2I) customization aims to create images that embody specific visual concepts delineated in textual descriptions However, existing works still face a main challenge, concept overfitting. To tackle this challenge, we first analyze overfitting, categorizing it into concept-agnostic overfitting, which undermines non-customized concept knowledge, and concept-specific overfitting, which is confined to customize on limited modalities, i.e, backgrounds, layouts, styles. To evaluate the overfitting degree, we further introduce two metrics, i.e, Latent Fisher divergence and Wasserstein metric to measure the distribution changes of non-customized and customized concept respectively. Drawing from the analysis, we propose Infusion, a T2I customization method that enables the learning of target concepts to avoid being constrained by limited training modalities, while preserving non-customized knowledge. Remarkably, Infusion achieves this feat with remarkable efficiency, requiring a mere 11KB of trained parameters. Extensive experiments also demonstrate that our approach outperforms state-of-the-art methods in both single and multi-concept customized generation.

⏳ To do

  • Customized training
  • Single concepts inference
  • Multi concepts inference
  • overfitting metrics
  • Support SDXL-1.0

Environment

Create and activate the conda environment:

conda env create -f environment.yaml
conda activate infusion

Training

Download the SD V1.5 or SD V2.1 to ./ckpt/.

Then run the commands:

sh train.sh

To prepare your own training data, please ensure that they are placed in a folder /path/to/your/images/. You need to download pretrained weight of clipseg:

wget https://owncloud.gwdg.de/index.php/s/ioHbRzFx6th32hn/download -O ./clipseg/weights.zip
unzip -d ./clipseg/weights -j ./clipseg/weights.zip

Then run:

python ./data/soft_segment.py --image_dir /path/to/your/images/ --super_class your_own_super_class

Modify the initializer_words, data_root, flip_p in ./configs/infusion_custom.yaml or ./configs/infusion_custom_sd_v2.yaml.

Finally, run:

sh train.sh

You can find weights along with tensorboard in ./ckpt.

Generating

Personalized samples can be obtained by running the command

sh inference.sh

Citation

If you find this work useful for your research, please consider citing our paper:

@article{zeng2024infusion,
  title={Infusion: Preventing Customized Text-to-Image Diffusion from Overfitting},
  author={Zeng, Weili and Yan, Yichao and Zhu, Qi and Chen, Zhuo and Chu, Pengzhi and Zhao, Weiming and Yang, Xiaokang},
  journal={arXiv preprint arXiv:2404.14007},
  year={2024}
}

Acknowledgement

This repo benefits from Custom Diffusion, Perfusion Thanks for their wonderful works.