Estás leyendo la publicación: Compresión de imagen basada en difusión estable: hacia la IA
Publicado originalmente en Hacia la IA, la empresa líder mundial en noticias y medios de IA y tecnología. Si está creando un producto o servicio relacionado con la IA, lo invitamos a considerar convertirse en patrocinador de la IA. En Hacia la IA, ayudamos a escalar las empresas emergentes de IA y tecnología. Permítanos ayudarlo a dar rienda suelta a su tecnología a las masas.
Compresión de imagen basada en difusión estable
Stable Diffusion lo convierte en un códec de compresión de imágenes con pérdida muy potente.
Introducción
Stable Diffusion está inspirando actualmente a la comunidad de aprendizaje automático de código abierto, así como a artistas molestos de todo el mundo. Tenía curiosidad por ver para qué más podría usarse este impactante lanzamiento de tecnología además de hacer que los artistas y diseñadores profesionales reflexionen sobre su seguridad laboral.
Mientras experimentaba con el modelo, descubrí que es un códec de compresión de imágenes con pérdida extremadamente poderoso. Antes de describir mi método y compartir algo de código, aquí hay algunos resultados comparados con JPG y WebP con factores de compresión altos, todo con una resolución de 512×512 píxeles:
Estos ejemplos hacen bastante evidente que comprimir estas imágenes con Stable Diffusion da como resultado una calidad de imagen muy superior con un tamaño de archivo más pequeño en comparación con JPG y WebP. Esta calidad viene con algunas advertencias importantes que deben tenerse en cuenta, como explicaré en la sección de evaluación, pero a primera vista, esta es una opción muy prometedora para la compresión agresiva de imágenes con pérdida.
Espacio latente de difusión estable
Stable Diffusion utiliza tres redes neuronales artificiales entrenadas en tándem:
El Codificador automático variacional (VAE) codifica y decodifica imágenes del espacio de imágenes en alguna representación de espacio latente. La representación del espacio latente es una representación de menor resolución (64 x 64), mayor precisión (4 × 32 bits) de cualquier imagen de origen (512 x 512 a 3 × 8 o 4 × 8 bits).
Cómo el VAE codifica la imagen en este espacio latente, aprende por sí mismo durante el proceso de entrenamiento y, por lo tanto, las representaciones del espacio latente de diferentes versiones del modelo probablemente se verán diferentes a medida que el modelo se entrena más, pero la representación de Stable Diffusion v1. 4 se ve así (cuando se reasigna e interpreta como una imagen en color de 4 canales):
Las características principales de la imagen aún son visibles al volver a escalar e interpretar las latentes como valores de color (con canal alfa), pero el VAE también codifica las características de mayor resolución en estos valores de píxeles.
Un viaje de ida y vuelta de codificación/decodificación a través de este VAE se ve así:
Tenga en cuenta que este viaje de ida y vuelta no es sin pérdidas. Por ejemplo, el nombre de Anna en su collar es un poco menos legible después de la decodificación. El VAE del modelo de difusión estable 1.4 generalmente no es muy bueno para representar texto pequeño (así como caras, algo que espero que mejore la versión 1.5 del modelo entrenado).
El algoritmo principal de Stable Diffusion, que genera nuevas imágenes a partir de breves descripciones de texto, opera sobre esta representación espacial latente de imágenes. Comienza con ruido aleatorio en la representación del espacio latente y luego iterativamente elimina el ruido de esta imagen del espacio latente usando el entrenamiento Red en U, que en términos simples genera predicciones de lo que cree que “ve” en ese ruido, similar a cómo a veces vemos formas y caras cuando miramos las nubes. Cuando se usa Stable Diffusion para generar imágenes, este paso iterativo de eliminación de ruido está guiado por el tercer modelo ML, el codificador de texto, que brinda a U-Net información sobre lo que debe tratar de ver en el ruido. Para el códec de imagen experimental presentado aquí, no se necesita el codificador de texto. El código de Google Colab compartido a continuación todavía lo usa, pero solo para crear una codificación única de una cadena vacía, que se usa para decirle a U-Net que elimine el ruido sin guía durante la reconstrucción de la imagen.
Método de compresión
Para usar Stable Diffusion como un códec de compresión de imágenes, investigué cómo la representación latente generada por el VAE podría comprimirse de manera eficiente. Reducir la muestra de las latentes o aplicar métodos de compresión de imágenes con pérdida existentes a las latentes demostró degradar masivamente las imágenes reconstruidas en mis experimentos. Sin embargo, descubrí que la decodificación de VAE parece ser muy robusta para cuantización de lo latente.
La cuantificación de las latentes desde el punto flotante hasta los enteros sin signo de 8 bits al escalarlos, sujetarlos y luego reasignarlos da como resultado un error de reconstrucción visible muy pequeño:
Para cuantificar las latentes generadas por el VAE, primero las escalé en 1 / 0.18215, un número que quizás ya haya encontrado en el código fuente de Stable Diffusion. Parece que dividir las latentes por este número las asigna bastante bien a la [-1, 1] rango, aunque todavía se producirá algo de sujeción.
Al cuantificar las latentes a 8 bits, el tamaño de los datos de la representación de la imagen ahora es 64*64*4*8 bits = 16 kB (por debajo de los 512*512*3*8 bits = 768 kB de la realidad básica sin comprimir) .
Cuantificar las latentes a menos de 8 bits no arrojó buenos resultados en mis experimentos, pero lo que funcionó sorprendentemente bien fue cuantificar aún más por paletizado y difuminado a ellos. Creé una representación paletizada utilizando una paleta latente de 256 vectores de 4*8 bits y Difuminado de Floyd-Steinberg. El uso de una paleta con 256 entradas permite representar cada vector latente utilizando un solo índice de 8 bits, lo que eleva el tamaño de los datos a 64*64*8+256*4*8 bits = 5 kB.
Sin embargo, la representación paletizada ahora da como resultado algunos artefactos visibles al decodificarlos directamente con el VAE:
El tramado de las latentes paletizadas ha introducido ruido, lo que distorsiona el resultado decodificado. Pero dado que Stable Diffusion se basa en eliminar el ruido de latentes, podemos usar U-Net para eliminar el ruido introducido por el tramado. Después de solo 4 iteraciones, el resultado de la reconstrucción es visualmente muy parecido a la versión no cuantificada:
Si bien el resultado es muy bueno teniendo en cuenta la reducción extrema en el tamaño de los datos (un factor de compresión de 155x frente a la verdad del suelo sin comprimir), también se puede ver que se introdujeron artefactos, por ejemplo, un tono brillante en el símbolo del corazón que no estaba presente antes de la compresión. Sin embargo, es interesante cómo los artefactos introducidos por este esquema de compresión afectan el contenido de la imagen más que la calidad de la imagen, y es importante tener en cuenta que las imágenes comprimidas de esa manera pueden contener este tipo de artefactos de compresión.
Finalmente, comprimo sin pérdidas la paleta y los índices usando zlib, lo que da como resultado un poco menos de 5kB para la mayoría de las muestras en las que probé. Examiné la codificación de longitud de ejecución, pero el tramado deja solo tiradas muy cortas del mismo índice para la mayoría de las imágenes (incluso para aquellas imágenes que funcionarían muy bien para la codificación de longitud de ejecución cuando se comprimen en el espacio de la imagen, como gráficos sin mucho gradientes). Estoy bastante seguro de que todavía hay más potencial de optimización aquí para explorar en el futuro.
Evaluación
Para evaluar este códec de compresión experimental, no utilicé ninguna de las imágenes de prueba estándar o las imágenes que se encuentran en línea para asegurarme de que no lo estoy probando en ningún dato que podría haberse utilizado en el conjunto de entrenamiento del modelo Stable Diffusion. (porque tales imágenes podrían obtener una ventaja de compresión injusta, ya que parte de sus datos ya podrían estar codificados en el modelo entrenado). Además, para que la comparación sea lo más justa posible, utilicé la configuración de calidad de codificador más alta para los compresores JPG y WebP de la biblioteca de imágenes de Python y, además, apliqué una compresión sin pérdidas de los datos JPG comprimidos utilizando el mozjpeg biblioteca. Luego usé la fuerza de compresión 1 menos de lo que daría como resultado un tamaño de datos más pequeño que el resultado de Difusión estable, o la fuerza máxima de lo contrario (muchas imágenes no se vuelven más pequeñas que el resultado SD incluso con la máxima pérdida de compresión JPG o WebP).
Es importante tener en cuenta que, si bien los resultados de Stable Diffusion se ven subjetivamente mucho mejor que las imágenes comprimidas JPG y WebP, no son significativamente mejores (pero tampoco peores) en términos de métricas de medición estándar como PSNR o SSIM. Es solo que el tipo de artefactos introducidos son mucho menos notables, ya que afectan el contenido de la imagen más que la calidad de la imagen, lo que, sin embargo, también es un peligro de este método: uno no debe dejarse engañar por la calidad de las características reconstruidas. el contenido puede verse afectado por artefactos de compresión, incluso si se ve muy claro. por ejemplo, mirando un detalle en la imagen de prueba de San Francisco:
Como puede ver, mientras que Stable Diffusion como códec es mucho mejor para preservar los aspectos cualitativos de la imagen hasta el grano de la cámara (algo con lo que luchan la mayoría de los algoritmos de compresión tradicionales), el contenido aún se ve afectado por los artefactos de compresión y, por lo tanto, características finas como ya que la forma de los edificios puede cambiar. Si bien ciertamente no es posible reconocer más Ground Truth en la imagen comprimida JPG que en la imagen comprimida Stable Diffusion, la alta calidad visual del resultado SD puede ser engañosa, ya que los artefactos de compresión en JPG y WebP se identifican mucho más fácilmente. como tal.
Limitaciones
Algunas características actualmente no están bien conservadas por Stable Diffusion VAE, en particular, el texto pequeño y las caras parecen ser problemáticas (lo cual es consistente con las limitaciones y sesgos del modelo V1.4 SD enumerados en su modelo de tarjeta). Aquí hay tres ejemplos que ilustran esto:
Como puede ver en la columna del medio, el modelo Stable Diffusion 1.4 VAE actual no conserva muy bien el texto pequeño y las caras en su espacio latente; como tal, este esquema de compresión hereda esas limitaciones. Pero dado que la versión 1.5 de Stable Diffusion, disponible hasta ahora solo en Dreamstudio de stable.ai, parece funcionar mucho mejor con las caras, espero actualizar este artículo una vez que el modelo 1.5 esté disponible públicamente.
¿En qué se diferencia este enfoque de la restauración de imágenes basada en ML?
El aprendizaje profundo se usa con mucho éxito para restaurar imágenes y videos degradados o para escalarlos a una resolución más alta. ¿En qué se diferencia este enfoque de la restauración de imágenes que han sido degradadas por la compresión de imágenes?
Es importante entender que Stable Diffusion no se usa aquí para restaurar una imagen que ha sido degradada por la compresión en el espacio de la imagen, sino que aplica una compresión con pérdida a la comprensión interna de Stable Diffusion de esa imagen y luego intenta “reparar” el daño que compresión con pérdida causada a la representación interna (que no es lo mismo que reparar una imagen degradada).
Esto, por ejemplo, conserva el grano de la cámara (cualitativamente), así como cualquier otra degradación cualitativa, mientras que una restauración de IA de un jpg muy comprimido no podría restaurar ese grano de la cámara porque se ha perdido cualquier información al respecto de los datos.
Para dar una analogía de esta diferencia: Digamos que tenemos un artista altamente calificado que posee una memoria fotográfica. Les mostramos una imagen y luego les pedimos que la recreen y pueden crear una copia casi perfecta solo de su memoria. La memoria fotográfica de este artista es el VAE de Stable Diffusion.
En el caso de la restauración, le mostraríamos a este artista una imagen que ha sido muy degradada por la compresión de la imagen y luego le pediríamos que recreara esta imagen de memoria, pero de una forma en la que crea que la imagen podría haberse visto antes de degradarse.
Sin embargo, en el caso implementado aquí, les mostramos la imagen original sin comprimir y les pedimos que la memoricen lo mejor que puedan. Luego les realizamos una cirugía cerebral y reducimos los datos en su memoria aplicándoles una compresión con pérdida que elimina matices que parecen sin importancia y reemplaza variaciones muy similares de conceptos y aspectos de la imagen memorizada por la misma variación.
Después de la cirugía les pedimos que creen una reconstrucción perfecta de la imagen a partir de su memoria. Todavía recordarán todos los aspectos importantes de la imagen, desde el contenido hasta las propiedades cualitativas del grano de la cámara, por ejemplo, y la ubicación y el aspecto general de cada edificio que vieron, aunque la ubicación exacta de cada punto del grano de la cámara no es Ya no es lo mismo y algunos edificios que recordarán ahora con defectos extraños que realmente no tienen sentido.
Finalmente, les pedimos que dibujen la imagen nuevamente, pero esta vez, si recuerdan algunos aspectos de una manera realmente extraña que no tiene mucho sentido para ellos, deben usar su experiencia para dibujar estas cosas de una manera que tenga sentido. (Tenga en cuenta que no es lo mismo que: una foto borrosa o rayada es mala, pero ese defecto tiene mucho sentido).
Entonces, en ambos casos, se le pide al artista que dibuje cosas para que se vean como cree que deberían verse en función de su experiencia, pero dado que la compresión aquí se ha aplicado a su representación de memoria de la imagen, que almacena conceptos más que píxeles, solo el contenido informativo de los conceptos se ha reducido, mientras que en el caso de restaurar una imagen degradada por la compresión de la imagen, tanto la calidad visual como el contenido conceptual de la imagen se han reducido y deben ser inventados por el artista.
Esta analogía también deja en claro que este esquema de compresión está limitado por qué tan bien funciona la memoria fotográfica del artista. En el caso de Stable Diffusion v1.4 no son muy buenos recordando caras y además sufren de dislexia.
Código
Si desea probar este códec experimental usted mismo, puede hacerlo con el siguiente Google Colab:
https://colab.research.google.com/drive/1Ci1VYHuFJK5eOX9TB0Mq4NsqkeDrMaaH?usp=sharing
Conclusión
Stable Diffusion parece ser muy prometedor como base para un esquema de compresión de imágenes con pérdida.
Mis experimentos aquí fueron bastante superficiales, y espero que haya mucho más potencial para comprimir la representación del espacio latente más allá de la técnica descrita aquí, pero incluso este esquema relativamente simple de cuantización y tramado ya produce resultados muy impresionantes.
Si bien un VAE específicamente diseñado y capacitado para la compresión de imágenes probablemente podría desempeñarse mejor en esta tarea, la gran ventaja de usar Stable Diffusion VAE es que las 6 o 7 cifras que ya se invirtieron en la capacitación de Stable Diffusion se pueden aprovechar directamente sin tener que invertir grandes cantidades de dinero en la formación de un nuevo modelo para la tarea.
Apéndice: Resultados adicionales
La compresión de imagen basada en difusión estable se publicó originalmente en Hacia la IA en Medium, donde las personas continúan la conversación resaltando y respondiendo a esta historia.
Publicado a través de Hacia la IA
