SwiftBrush v2: Mejora tu modelo de difusión de un paso mejor que su maestro

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

August 26, 2024

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

cs.AI

Resumen

En este documento, nuestro objetivo es mejorar el rendimiento de SwiftBrush, un destacado modelo de difusión de texto a imagen de un solo paso, para que sea competitivo con su contraparte de difusión estable de múltiples pasos. Inicialmente, exploramos el equilibrio calidad-diversidad entre SwiftBrush y SD Turbo: el primero destaca en diversidad de imágenes, mientras que el segundo sobresale en calidad de imagen. Esta observación motiva nuestras modificaciones propuestas en la metodología de entrenamiento, que incluyen una mejor inicialización de pesos y un entrenamiento eficiente de LoRA. Además, nuestra introducción de una nueva pérdida CLIP limitada mejora el alineamiento imagen-texto y da como resultado una mejor calidad de imagen. Notablemente, al combinar los pesos de modelos entrenados con LoRA eficiente y entrenamiento completo, logramos un nuevo modelo de difusión de un solo paso de última generación, alcanzando un FID de 8.14 y superando a todos los modelos basados en GAN y de difusión estable de múltiples pasos. El código de evaluación está disponible en: 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.