Stable Diffusion XL 1.0 ejecutándose con 298 MB de RAM

• OnnxStream es una biblioteca de inferencia ultracompacta que busca ejecutar Stable Diffusion 1.5 y SDXL 1.0 Base en una Raspberry Pi Zero 2 con 512 MB de RAM, sin agregar swap ni guardar resultados intermedios en disco
• Separa el motor de inferencia y WeightsProvider, admite DiskNoCache, DiskPrefetch, carga/caché/prefetching personalizados, y también puede recibir pesos directamente por HTTP para procesarlos sin escribir en disco
• SDXL 1.0 Base usa archivos ONNX exportados desde la implementación de Hugging Face Diffusers 0.19.3, y generar una imagen de 10 steps en una Raspberry Pi Zero 2 toma alrededor de 11 horas
• El decodificador VAE de SDXL usa 4.4 GB de RAM en FP32, y como FP16 o UINT8 por sí solos no alcanzan para resolverlo, se reduce la memoria hasta 298 MB mediante decodificación por tiles superpuestos y blending
• En CPU usa hasta 55 veces menos memoria que OnnxRuntime, aunque puede ser entre 0.5 y 2 veces más lento, y tiene limitaciones como batch size 1 y foco en algunos operadores ONNX

---

Objetivos y diseño de OnnxStream

• OnnxStream comenzó con el objetivo de ejecutar Stable Diffusion 1.5 en una Raspberry Pi Zero 2
  • La Raspberry Pi Zero 2 es una microcomputadora con 512 MB de RAM
  • La RAM/VRAM mínima recomendada para Stable Diffusion suele ser de 8 GB
  • Las condiciones objetivo eran no agregar espacio de swap ni descargar resultados intermedios al disco
• Los frameworks de machine learning comunes suelen priorizar la latencia de inferencia o el throughput, por lo que su uso de RAM tiende a crecer
• OnnxStream es una biblioteca de inferencia pequeña y fácil de modificar, enfocada en minimizar el uso de memoria
• Su estructura central es la separación entre el motor de inferencia y WeightsProvider, el componente que provee los pesos del modelo
  • Las clases derivadas pueden implementar cómo cargar, cachear y hacer prefetch de los parámetros del modelo
  • Un WeightsProvider personalizado puede descargar datos directamente desde un servidor HTTP y procesarlos sin lecturas ni escrituras en disco
  • Las implementaciones básicas son DiskNoCache y DiskPrefetch
• En CPU usa hasta 55 veces menos memoria que OnnxRuntime, con una velocidad de entre 0.5 y 2 veces más lenta

Ejecución de Stable Diffusion 1.5

• El ejemplo de Stable Diffusion 1.5 usa OnnxStream para generar imágenes con distintas precisiones del decodificador VAE
• El único componente que no cabía en la RAM de la Raspberry Pi Zero 2 era el decodificador VAE
  • La causa eran las residual connections internas del modelo, tensores grandes y convoluciones grandes
  • Solo con precisión simple o media precisión no entraba en la RAM de la Raspberry Pi Zero 2
  • La solución fue la cuantización estática de 8 bits
• La imagen con el decodificador VAE en precisión W8A8 se generó en la Raspberry Pi Zero 2, y al usar la opción MAX_SPEED tardó alrededor de 1.5 horas
• La imagen W16A16 de comparación se generó en una PC usando el mismo latent

Soporte para Stable Diffusion XL 1.0 Base

• El ejemplo de Stable Diffusion de OnnxStream admite SDXL 1.0 Base y no incluye Refiner
• Los archivos ONNX se exportaron desde la implementación de SDXL 1.0 de la biblioteca Diffusers de Hugging Face, usando la versión 0.19.3
• SDXL 1.0 tiene un costo computacional mucho mayor que SD 1.5
  • La mayor diferencia es que puede generar imágenes de 1024x1024 en lugar de 512x512
  • En una PC de 12 núcleos y 32 GB de RAM, generar una imagen de 10 steps con Hugging Face Diffusers tardó 26 minutos
  • La VRAM mínima recomendada para SDXL suele ser de 12 GB
• OnnxStream puede ejecutar SDXL 1.0 con menos de 300 MB de RAM, por lo que también funciona en una Raspberry Pi Zero 2
  • Sin agregar swap
  • Sin escrituras en disco durante la inferencia
  • Generar una imagen de 10 steps en una Raspberry Pi Zero 2 toma alrededor de 11 horas

Optimizaciones de memoria específicas para SDXL

• En SDXL 1.0 se aplica el mismo conjunto de optimizaciones que en SD 1.5, aunque con algunas diferencias
• El modelo UNET usa cuantización dinámica UINT8 para ejecutarse en la Raspberry Pi Zero 2 con menos de 300 MB de RAM
  • La cuantización se limita a un subconjunto específico de grandes tensores intermedios
• El decodificador VAE de SDXL 1.0 es más difícil de manejar que el de SD 1.5
  • El decodificador VAE de SDXL 1.0 es 4 veces más grande que el decodificador VAE de SD 1.5
  • Al ejecutarlo en OnnxStream con precisión FP32, usa 4.4 GB de RAM
  • En SD 1.5, el decodificador VAE podía cuantizarse estáticamente a UINT8 para reducir el uso de RAM hasta 260 MB
• El decodificador VAE de SDXL 1.0 presenta overflow en aritmética FP16, y el rango numérico de las activations es grande, por lo que es difícil obtener imágenes de buena calidad con cuantización UINT8
• Existen soluciones FP16 como sdxl-vae-fp16-fix, pero incluso reduciendo la memoria a la mitad serían 2.2 GB, todavía demasiado para la Raspberry Pi Zero 2
• La solución final es la decodificación por tiles, inspirada en la implementación del decodificador VAE de Diffusers de Hugging Face
  • El tensor resultante de la diffusion tiene shape (1,4,128,128)
  • Se divide en 5x5, 25 tensores en total, superpuestos, con shape (1,4,32,32)
  • Cada tile tiene 25% de superposición con el tile de la izquierda y el de arriba
  • El resultado de decodificar cada tile es un tensor (1,3,256,256), que se combina mediante blending en la imagen final
  • Si se desactiva el blending, se ven los bordes de los tiles; si se activa, los bordes no se ven en el resultado final
• Con este método, el uso de RAM del decodificador VAE de SDXL baja de 4.4 GB a 298 MB

Funciones soportadas y dependencias

• OnnxStream agrupa de forma compacta las funciones necesarias para inferencia con ahorro de memoria
  • Separación entre motor de inferencia y WeightsProvider
  • DiskNoCache, DiskPrefetch, WeightsProvider personalizado
  • attention slicing
  • Cuantización dinámica unsigned asymmetric percentile de 8 bits
  • Cuantización estática W8A8 unsigned asymmetric percentile
  • Calibración de modelos cuantizados
  • Soporte para FP16
  • Implementación de 25 operadores ONNX de uso frecuente
• Las operaciones se ejecutan en secuencia, pero cada operator funciona con multithreading
• La implementación está estructurada como un único archivo de implementación y un archivo de header, y envuelve las llamadas a XNNPACK mediante la clase XnnPack
• Algunas primitives de aceleración dependen de XNNPACK
  • MatMul
  • Convolution
  • Add/Sub/Mul/Div element-wise
  • Sigmoid
  • Softmax

Comparación de rendimiento y limitaciones

• Stable Diffusion está compuesto por tres modelos
  • text encoder: 672 operations, 123 millones de parameters
  • UNET: 2050 operations, 854 millones de parameters
  • VAE decoder: 276 operations, 49 millones de parameters
• Para generar una imagen de 10 steps con batch size 1 se requieren las siguientes ejecuciones
  • text encoder 2 veces
  • UNET 20 veces
  • VAE decoder 1 vez
• Con UNET FP16, la diferencia de memoria y tiempo entre OnnxStream y OnnxRuntime es grande
  • OnnxStream: alrededor de 0.133 GB, 18.2~19.8 segundos
  • OnnxRuntime: 5.085~7.353 GB, 7.28~12.8 segundos
  • OnnxStream usa hasta 55 veces menos memoria, pero es entre 0.5 y 2 veces más lento
• En el text encoder FP32, OnnxStream usa 0.147 GB y OnnxRuntime 0.641 GB
• En el VAE decoder FP32, OnnxStream usa 1.004 GB y OnnxRuntime 1.330~2.026 GB
• Los resultados de comparación tienen varias condiciones
  • La primera ejecución de OnnxRuntime es una inferencia de warm-up antes de reutilizar InferenceSession
  • OnnxStream tiene un diseño eager, por lo que no existe el concepto de warm-up, aunque las ejecuciones posteriores pueden beneficiarse de la caché de archivos de weights del SO
  • OnnxStream actualmente no admite entradas que no sean batch size 1
  • OnnxRuntime puede usar batch size 2 en la ejecución de UNET para acelerar significativamente toda la diffusion
  • Cambiar EnableCpuMemArena y ExecutionMode en SessionOptions de OnnxRuntime no produjo diferencias significativas en las pruebas
  • NCNN fue muy similar a OnnxRuntime en términos de uso de memoria y tiempo de inferencia
  • El entorno de prueba fue Windows Server 2019, 16 GB de RAM, CPU 8750H con AVX2, SSD 970 EVO Plus y VMWare con 8 virtual cores

Attention slicing y cuantización

• Al ejecutar UNET, attention slicing y la cuantización W8A8 del decodificador VAE fueron importantes para reducir la memoria hasta un nivel ejecutable en la Raspberry Pi Zero 2
• Attention slicing evita materializar toda la matriz Q @ K^T al calcular scaled dot-product attention en multi-head attention
• Cuando el número de attention heads del modelo UNET es 8, los shapes de los tensores son los siguientes
  • Q: (8,4096,40)
  • K^T: (8,40,4096)
  • Primer resultado de MatMul: (8,4096,4096)
  • En precisión FP32, se convierte en un tensor de 512 MB
• La solución consiste en dividir Q verticalmente y realizar la operación de attention para cada chunk
  • El shape de Q_sliced es (1,x,40)
  • x es 4096 dividido por onnxstream::Model::m_attention_fused_ops_parts
  • El valor predeterminado es 2 y se puede personalizar
• Este método reduce el uso total de memoria del modelo UNET FP32 de 1.1 GB a 300 MB
• FlashAttention podría ser una alternativa más eficiente, pero habría que escribir kernels personalizados para cada arquitectura compatible, como AVX o NEON, y en ese caso habría que evitar XNNPACK

Conversión de modelos y forma de ejecución

• OnnxStream ejecuta el modelo definido en path_to_model_folder/model.txt
  • Todas las operaciones del modelo están en model.txt en formato ASCII
  • Los archivos de weights deben existir en la misma carpeta como una serie de archivos .bin
• En el objeto Model se pueden configurar varios parámetros opcionales
  • Especificar otro weights provider
  • Leer/escribir archivos de activation clipping range para el modelo cuantizado
  • Modo de calibración del modelo
  • Uso de aritmética FP16
  • Uso de aritmética UINT8
  • Uso de cuantización dinámica UINT8
  • Activación de attention slicing
  • Configuración del número de particiones de attention
• Para usar archivos ONNX en OnnxStream, el notebook onnx2txt.ipynb exporta model.txt y los archivos de weights .bin
• Al exportar un nn.Module de PyTorch a ONNX, hay varias condiciones
  • Al llamar a torch.onnx.export, dynamic_axes debe dejarse vacío
  • OnnxStream no admite entradas con shape dinámico
  • Se recomienda firmemente ejecutar ONNX Simplifier antes de la conversión

Preparación de build y ejecución

• El ejemplo de Stable Diffusion puede compilarse en Linux, Mac, Windows y Termux
  • En Windows se usa x64 Native Tools Command Prompt de Visual Studio Tools
  • En Mac se requiere instalar cmake con brew install cmake
• Primero hay que compilar XNNPACK
  • Como los prototypes de funciones de XNNPACK pueden cambiar en cualquier momento, se incluye un procedimiento para hacer checkout de un commit de un momento específico
  • El punto de referencia es un commit de master anterior a 2023-06-27 00:00
• Luego se clona el repositorio de OnnxStream y se compila con cmake desde src/build
  • MAX_SPEED=ON
  • Especificar con XNNPACK_DIR la ruta del clone de XNNPACK
• La opción MAX_SPEED puede mejorar el rendimiento, pero usa más memoria durante la compilación
  • En Windows mejora el rendimiento alrededor de 10%
  • En Raspberry Pi mejora el rendimiento en más de 50%
  • Es posible que el ejecutable generado no funcione, y hubo problemas en pruebas con Termux
  • Si hay problemas, se recomienda cambiar primero a MAX_SPEED=OFF
• Los weights del ejemplo de Stable Diffusion 1.5 se pueden descargar desde Releases del repositorio y pesan alrededor de 2 GB
• Los weights de Stable Diffusion XL 1.0 Base se pueden descargar desde Hugging Face y pesan alrededor de 8 GB

Opciones del ejemplo de Stable Diffusion

• El ejecutable de ejemplo controla la selección de modelo, entrada/salida, prompt y método de decodificación mediante opciones de línea de comandos
  • --xl: ejecutar Stable Diffusion XL 1.0 en lugar de Stable Diffusion 1.5
  • --models-path: especificar la carpeta de modelos de Stable Diffusion
  • --ops-printf: imprimir la operation actual en stdout durante la inferencia
  • --output: especificar el archivo PNG de salida
  • --decode-latents: omitir la diffusion y decodificar el archivo de latents especificado
  • --prompt: especificar el positive prompt
  • --neg-prompt: especificar el negative prompt
  • --steps: especificar la cantidad de diffusion steps
  • --save-latents: guardar los latents en el archivo especificado después de la diffusion
• También se ofrecen opciones separadas para Raspberry Pi y el decodificador
  • --decoder-calibrate: calibrar el VAE decoder cuantizado solo en SD 1.5
  • --decoder-fp16: usar el VAE decoder FP16 solo en SD 1.5
  • --not-tiled: no usar tiled VAE decoder solo en SDXL 1.0
  • --rpi: configurar el modelo para ejecutarlo en Raspberry Pi
  • --rpi-lowmem: aplicar configuración de baja memoria para Raspberry Pi Zero 2 solo en SDXL 1.0

[GN⁺]

Opiniones de Hacker News

• Interesante. La frase clave es esta: “OnnxStream puede ser solo entre 0,5 y 2 veces más lento que OnnxRuntime, pero con hasta 55 veces menos uso de memoria”.
  El compromiso entre memoria de video/uso de memoria y tiempo de inferencia parece que podría ser útil no solo en casos con memoria limitada como Raspberry Pi, sino también en otras situaciones.
  Me pregunto si, en la práctica, si este método de descarga de pesos permite manejar tamaños de lote más grandes con la misma memoria, no se podría aumentar mucho el throughput aunque suba la latencia.

  • Quiero usar esto en LLM. Si la huella de memoria baja tanto, se podrían cargar más modelos simultáneamente en una sola GPU y, suponiendo que el reloj acompañe, quizá se compense de sobra la pérdida de velocidad de inferencia de cada modelo.
  • ¿“0,5~2 veces más lento” no es un error tipográfico? Si es 0,5 veces más lento, en realidad significa que es 2 veces más rápido.
    Probablemente querían escribir “1,5~2 veces más lento”.
  • Según mi lectura, quizá algo simplista, no redujeron el ancho de banda de memoria necesario, sino solo el tamaño del conjunto de trabajo.
    La inferencia normalmente queda limitada por el ancho de banda de memoria una vez que se supera el nivel de “este modelo cabe en este sistema”, así que no creo que esta técnica ayude mucho a aumentar el throughput con tamaños de lote más grandes. Es muy probable que una sola instancia ya esté saturando el controlador de memoria.
    Aunque tal vez sí ayude del lado del entrenamiento.

• 11 horas me recuerda a cuando hacía ray tracing en una Amiga 500. El render “final” definitivamente era algo que se dejaba corriendo toda la noche.

  • Hoy en día todavía lo hago a veces. Un ray tracer bidireccional moderno puede usar técnicas interesantes, y quería ver cáusticas como esas líneas brillantes que aparecen en los charcos.
    Pero aunque las cáusticas se vean brillantes, estadísticamente son fenómenos bastante raros, así que para obtenerlas bien hay que quitar los límites del motor de render y dejarlo corriendo toda la noche.
    El resultado es una imagen de una escena mediocre hecha por un artista sin mucha habilidad, pero con cáusticas geniales. Mejor sigo con mi trabajo de siempre.
  • La razón para hacer primero un render de baja calidad era que es mejor perder solo una hora para confirmar que está bien que desperdiciar toda la noche con algo incorrecto.
    Por esa época pensé que necesitaba otro hobby. Justo antes, una persona muy talentosa había lanzado una herramienta que permitía previsualizar la escena renderizándola con OpenGL. En Amiga no habría funcionado, pero en mi máquina apenas corría.
  • En mi 286 era lo mismo. Dejaba configurado povray, me iba a dormir y por la mañana revisaba la imagen antes de ir a la escuela.
  • Aunque fue más tarde, hice algo parecido en una 386 con una copia dudosa de 3DSMAX.
  • Recuerdo correr fractales de Mandelbrot en una C64. Depurar el código era realmente difícil.

• Estoy usando Stable Diffusion con invoke.ai en una MBP. ¿Alguna recomendación para ajustar mejor los parámetros de SD? Incluso con el mismo prompt y configuraciones aparentemente iguales, por ejemplo el mismo modelo con Euler A, no logro ni acercarme a la calidad de las imágenes que veo en internet.

  • De lo que he probado hasta ahora, esto fue lo mejor, aunque parece que no tiene soporte para Mac. Es un fork con muchas funciones de Fooocus, hecho por el desarrollador original de ControlNet, y la calidad que logra con prompts cortos es sorprendente: https://github.com/MoonRide303/Fooocus-MRE
    Para SD 1.5 básico uso Volta porque es rápido: https://github.com/VoltaML/voltaML-fast-stable-diffusion/com...
    Una calidad de imagen realmente buena en SD 1.5 sale de usar sin escatimar modelos fine-tuned, LoRA, ControlNet y otras funciones de mejora. Por ejemplo, hacer que siga una imagen base como estructura, o definir prompts por regiones específicas de la imagen. InvokeAI en realidad también tiene muchas funciones, y muchas de estas mejoras están escondidas dentro de la UI de nodos, pero otras UIs como Volta las exponen de forma más directa.
  • ¿Estás usando pesos personalizados? Probablemente sí, pero hay una gran diferencia entre los pesos base de RunwayML 1.5 y modelos fine-tuned para objetivos específicos.
    En general, los modelos fine-tuned impresionantes son mucho menos versátiles que los pesos base, pero en la práctica eso no suele ser un gran problema y los resultados pueden mejorar bastante.
  • Tuve la misma experiencia usando Invoke.ai o MochiDiffusion en una MBP M1. Lo único con lo que pude igualar la calidad de otras imágenes fue Automatic1111(https://github.com/AUTOMATIC1111/stable-diffusion-webui).
    Requiere más tiempo y memoria que Invoke o una tarjeta gráfica Nvidia, pero no está nada mal. Una imagen de calidad estándar de 512x768 px tarda unos 1~2 s/it, y una imagen de alta calidad de 1024x1536 px con Hires Fix, unos 14~20 s/it.
  • ¿Está especificado que esas imágenes salen directamente del generador? Los videos de procesos que he visto empiezan con “una chica de pie en un campo verde” y luego pasan más de una hora haciendo inpainting para corregir manos, postura, etc.
  • Draw Things agregó un modo de seed compatible con CUDA, así que también en Mac se pueden igualar imágenes generadas en tarjetas NVDIA.

• Sería realmente genial meter esto en un marco digital o un cuadro de pared y hacerlo correr ahí.

  • Hace tiempo hice esto con una versión anterior que corría Stable Diffusion en una Raspberry Pi Zero 2 W: https://hackaday.com/2023/09/19/e-paper-news-feed-illustrate...
    https://github.com/rvdveen/epaper-slow-generative-art/
  • Ahora mismo estoy construyendo exactamente eso con una pantalla e-ink. Lamentablemente, no logro compilar en una Pi Zero 2W la parte del repositorio relacionada con XNNPACK.
  • Buena idea. Podría generar y actualizar una nueva imagen por sí mismo cada unas 10 horas, basándose en algo como un tema dado por el usuario.
  • Aunque no sería muy ecológico.

• Es un logro sorprendente, pero obviamente generar una imagen toma muchísimo tiempo. En el README dice 11 horas.

  • Así es. No es que yo necesite o vaya a usar esta implementación, pero las técnicas usadas aquí se trasladarán a otras herramientas, y eso será genial.
  • Sería interesante ver qué tan equivalente es en costo y energía comparado con métodos más avanzados.
    Por ejemplo, comparar el tiempo, el costo, todo el hardware necesario y la energía para generar 100 imágenes con 100 Pi Zero 2 —ni siquiera tienen que ser W— contra algo como una PC promedio de gama media.
    Probablemente la PC seguiría ganando.
    El Zero 2 se siente más como un desafío que como algo práctico, así que quizá una Pi 4 o 5 sería un mejor punto de referencia.
  • La clave es que corre en una Raspberry Pi, y encima una Zero 2.

• Impresionante.
  En verdad, se acerca la era en la que incluso las lámparas y las tostadoras albergarán una inteligencia notable.
  Llevo varios años siguiendo este campo, pero la última década ha sido asombrosa.
  Digo “ha sido” porque la aceleración de los últimos 6 a 18 meses está en otro nivel por completo.
  No me preocupa qué podremos hacer dentro de 2 años, sino cuánto más rápido se volverá el ritmo de avance. Y luego se acelerará otra vez, y otra vez.

  • Me encanta la idea de una tostadora que reciba un prompt y dibuje esa imagen en el pan tostado. Incluso se podría aprovechar el calor de la GPU para tostar de verdad.
    Hagamos una startup.

• Llegados a este punto, ¿no se acabaron los intentos de regular tecnologías tipo Stable Diffusion? Si se reducen el modelo y la infraestructura de inferencia hasta el punto de que puedan correr incluso en una PS2, parece imposible frenar esta tecnología sin un Estado de vigilancia totalitario, e incluso ahí apenas si sería posible.

  • La guerra contra la computación de propósito general sigue en curso, pero todavía no ha llegado al punto de impedir que la gente posea dispositivos de computación de propósito general.
  • Con esa lógica, ¿no se acabó también la regulación del robo? Con solo tener un cuerpo se puede abrir una ventana sin herramientas, así que sin un Estado de vigilancia totalitario no se podría impedir el robo.
    Lo mismo aplica a la “piratería” de medios o al ransomware.
    Los Estados llevan mucho tiempo regulando cosas que no pueden hacerse cumplir de forma puramente técnica.
  • Los requisitos originales de este modelo son 16 GB de RAM, y eso se puede comprar por menos de 20 dólares. Corre mucho más rápido en una GPU, y esas GPU también se pueden conseguir por menos de 200 dólares. Millones de personas comunes ya tienen ambas cosas.
  • La PS2 tenía solo 32 MB de RAM. La PS3 también tenía apenas 256 MB.
    Entiendo que es una exageración divertida, pero para correrlo en una PS2 habría que reducirlo muchísimo más que esto.
  • Pensaba que la mayoría de los intentos de regulación se enfocaban en las ejecuciones de entrenamiento, cada vez más grandes, más que en el uso generativo de modelos existentes. ¿Hay regulaciones específicas sobre el uso del modelo en sí?

• No puedo esperar al día en que salga Stable Diffusion para Windows 3.1.

• Una locura. Tarde 11 horas o no, jamás habría esperado que SD pudiera correr en hardware como una Pi Zero.