SwiftBrush v2: Faça seu Modelo de Difusão de Uma Etapa Ser Melhor Que Seu Professor

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

Resumo

Neste artigo, nosso objetivo é aprimorar o desempenho do SwiftBrush, um modelo proeminente de difusão de texto para imagem em um passo, para ser competitivo com seu equivalente de difusão estável em vários passos. Inicialmente, exploramos a compensação qualidade-diversidade entre o SwiftBrush e o SD Turbo: o primeiro se destaca na diversidade de imagens, enquanto o último se destaca na qualidade da imagem. Essa observação motiva nossas modificações propostas na metodologia de treinamento, incluindo uma melhor inicialização de pesos e um treinamento eficiente LoRA. Além disso, nossa introdução de uma nova perda CLIP limitada aprimora o alinhamento imagem-texto e resulta em uma melhor qualidade de imagem. Notavelmente, ao combinar os pesos de modelos treinados com LoRA eficiente e treinamento completo, alcançamos um novo modelo de difusão em um passo de ponta, atingindo um FID de 8,14 e superando todos os modelos baseados em GAN e modelos de difusão estável em vários passos. O código de avaliação está disponível em: 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.