Diseñando un Mejor VQGAN Asimétrico para StableDiffusion

Designing a Better Asymmetric VQGAN for StableDiffusion

June 7, 2023

Autores: Zixin Zhu, Xuelu Feng, Dongdong Chen, Jianmin Bao, Le Wang, Yinpeng Chen, Lu Yuan, Gang Hua

cs.AI

Resumen

StableDiffusion es un revolucionario generador de texto a imagen que está causando un gran impacto en el mundo de la generación y edición de imágenes. A diferencia de los métodos tradicionales que aprenden un modelo de difusión en el espacio de píxeles, StableDiffusion aprende un modelo de difusión en el espacio latente a través de un VQGAN, garantizando tanto eficiencia como calidad. No solo admite tareas de generación de imágenes, sino que también permite la edición de imágenes reales, como la restauración de imágenes y la edición local. Sin embargo, hemos observado que el VQGAN estándar utilizado en StableDiffusion provoca una pérdida significativa de información, generando artefactos de distorsión incluso en regiones de la imagen no editadas. Para abordar este problema, proponemos un nuevo VQGAN asimétrico con dos diseños simples. En primer lugar, además de la entrada del codificador, el decodificador contiene una rama condicional que incorpora información de conocimientos previos específicos de la tarea, como la región de la imagen no enmascarada en la restauración. En segundo lugar, el decodificador es mucho más robusto que el codificador, permitiendo una recuperación más detallada mientras solo aumenta ligeramente el costo total de inferencia. El costo de entrenamiento de nuestro VQGAN asimétrico es bajo, y solo necesitamos reentrenar un nuevo decodificador asimétrico manteniendo el codificador VQGAN estándar y StableDiffusion sin cambios. Nuestro VQGAN asimétrico puede ser ampliamente utilizado en métodos de restauración y edición local basados en StableDiffusion. Experimentos exhaustivos demuestran que puede mejorar significativamente el rendimiento en tareas de restauración y edición, manteniendo al mismo tiempo la capacidad original de generación de texto a imagen. El código está disponible en https://github.com/buxiangzhiren/Asymmetric_VQGAN.

English

StableDiffusion is a revolutionary text-to-image generator that is causing a stir in the world of image generation and editing. Unlike traditional methods that learn a diffusion model in pixel space, StableDiffusion learns a diffusion model in the latent space via a VQGAN, ensuring both efficiency and quality. It not only supports image generation tasks, but also enables image editing for real images, such as image inpainting and local editing. However, we have observed that the vanilla VQGAN used in StableDiffusion leads to significant information loss, causing distortion artifacts even in non-edited image regions. To this end, we propose a new asymmetric VQGAN with two simple designs. Firstly, in addition to the input from the encoder, the decoder contains a conditional branch that incorporates information from task-specific priors, such as the unmasked image region in inpainting. Secondly, the decoder is much heavier than the encoder, allowing for more detailed recovery while only slightly increasing the total inference cost. The training cost of our asymmetric VQGAN is cheap, and we only need to retrain a new asymmetric decoder while keeping the vanilla VQGAN encoder and StableDiffusion unchanged. Our asymmetric VQGAN can be widely used in StableDiffusion-based inpainting and local editing methods. Extensive experiments demonstrate that it can significantly improve the inpainting and editing performance, while maintaining the original text-to-image capability. The code is available at https://github.com/buxiangzhiren/Asymmetric_VQGAN.