SwiftBrush v2 : Améliorez votre modèle de diffusion en une étape pour le rendre meilleur que son enseignant.

SwiftBrush v2: Make Your One-step Diffusion Model Better Than Its Teacher

August 26, 2024

Auteurs: Trung Dao, Thuan Hoang Nguyen, Thanh Le, Duc Vu, Khoi Nguyen, Cuong Pham, Anh Tran

cs.AI

Résumé

Dans cet article, notre objectif est d'améliorer les performances de SwiftBrush, un modèle de diffusion de texte en image en une étape de premier plan, afin qu'il soit compétitif avec son homologue Stable Diffusion à plusieurs étapes. Nous explorons initialement le compromis qualité-diversité entre SwiftBrush et SD Turbo : le premier excelle dans la diversité des images, tandis que le second excelle dans la qualité des images. Cette observation motive nos modifications proposées dans la méthodologie d'entraînement, comprenant une meilleure initialisation des poids et un entraînement LoRA efficace. De plus, notre introduction d'une nouvelle perte CLIP clippée améliore l'alignement image-texte et conduit à une meilleure qualité d'image. Remarquablement, en combinant les poids des modèles entraînés avec LoRA efficace et un entraînement complet, nous obtenons un nouveau modèle de diffusion en une étape de pointe, atteignant un FID de 8,14 et dépassant tous les modèles basés sur GAN et les modèles Stable Diffusion à plusieurs étapes. Le code d'évaluation est disponible sur : https://github.com/vinairesearch/swiftbrushv2.

English

In this paper, we aim to enhance the performance of SwiftBrush, a prominent one-step text-to-image diffusion model, to be competitive with its multi-step Stable Diffusion counterpart. Initially, we explore the quality-diversity trade-off between SwiftBrush and SD Turbo: the former excels in image diversity, while the latter excels in image quality. This observation motivates our proposed modifications in the training methodology, including better weight initialization and efficient LoRA training. Moreover, our introduction of a novel clamped CLIP loss enhances image-text alignment and results in improved image quality. Remarkably, by combining the weights of models trained with efficient LoRA and full training, we achieve a new state-of-the-art one-step diffusion model, achieving an FID of 8.14 and surpassing all GAN-based and multi-step Stable Diffusion models. The evaluation code is available at: https://github.com/vinairesearch/swiftbrushv2.