Demo de Stable Diffusion con WebGPU

• Demo de generación de imágenes con Stable Diffusion que se ejecuta directamente en el navegador, sin servidor adicional, y funciona con WebGPU y WebAssembly
• Para usarlo, en la versión más reciente de Chrome hay que activar las banderas "Experimental WebAssembly" y "Experimental WebAssembly JavaScript Promise Integration (JSPI)"
• Genera imágenes recibiendo Prompt, Negative Prompt y la cantidad de pasos de inferencia; además, los archivos del modelo se almacenan en caché, así que no hace falta volver a descargarlos cada vez
• Resultado de parchear onnxruntime, emscripten y binaryen para permitir asignación de memoria de más de 4 GB y portar a JS el StableDiffusionPipeline de Python
• Actualmente es lento por la falta de soporte de multihilo y por operaciones de WebGPU aún no implementadas, pero plantea mejoras posibles con implementación de kernels en JS y soporte de memory64

---

Requisitos de ejecución

• En la versión más reciente del navegador Chrome es necesario activar las banderas Experimental WebAssembly y Experimental WebAssembly JavaScript Promise Integration (JSPI)
• Los model files se almacenan en caché, por lo que no es necesario volver a descargarlos al regresar

Entradas y funcionamiento

• Como entradas ofrece Prompt, Negative Prompt y la cantidad de pasos de inferencia
  • Como usa PNDM Scheduler, los pasos reales se procesan de i a i+1 respecto al valor ingresado
• Cada paso tarda alrededor de 1 minuto, y ejecutar el VAE decoder para generar la imagen toma unos 10 segundos adicionales
• Si DevTools está abierto, la velocidad total se vuelve aproximadamente 2 veces más lenta
• UNET solo se ejecuta en CPU; ahí es 10% más rápido que en GPU, y en GPU no produce resultados correctos y termina congelando la pestaña del navegador
• La cantidad mínima de pasos para obtener resultados aceptables es 20, aunque para fines de demostración 3 pasos también sirven

Preguntas frecuentes

• Error protobuf parsing failed

  • En DevTools, ve a Application → Storage y ejecuta "Clear site data"

• Error sbox_fatal_memory_exceeded

  • Indica falta de RAM necesaria para ejecutar SD; intenta recargar la pestaña o el navegador

• Forma de implementación

  • Se portó a JS el StableDiffusionPipeline de Python
  • Se parchearon onnxruntime y emscripten+binaryen (toolchain compilador de WebAssembly) para permitir asignar y usar memoria de más de 4 GB
  • Cuando los pull requests relacionados se reflejen en una versión publicada, cualquiera podrá compilar y ejecutar en el navegador código que use más de 4 GB de memoria

• Causa de la baja velocidad

  • Aún no hay soporte de multihilo, así que solo se usa un núcleo de CPU
  • No es posible crear memoria de 64 bits con SharedArrayBuffer mediante el constructor de WebAssembly.Memory
  • Se propuso un cambio de especificación para la bandera "memory64", y si se adopta, se planea dar soporte parcheando el motor V8

• Ejecución en GPU

  • Sí se ejecuta en GPU, pero WebGPU en onnxruntime aún está en etapa inicial y muchas operaciones no están implementadas
  • Los datos se siguen transfiriendo constantemente entre CPU y GPU a través de JS
  • Si se implementan kernels en JS para la mayoría de las operaciones, la velocidad podría mejorar de forma importante

• Posibilidad de ejecución local

  • Sí, y el código de esa página está disponible en el repositorio stable-diffusion-webgpu-minimal

• Posibilidad de ejecutar LLM grandes con transformers.js

  • Se puede usar el paquete parcheado de onnxruntime (@aislamov/onnxruntime-web64), aunque no se garantiza que funcione en todos los casos
  • Esa build está limitada a 8 GB de memoria, por lo que puede cargar pesos de hasta unos 4 GB

• Plan de pull request para el repositorio de onnxruntime

  • Está previsto; sería el segundo trabajo después de haber agregado antes aceleración por GPU a los bindings de node.js

[GN⁺]

Comentarios en Hacker News

• El equipo de MLC ya había hecho funcionar esto en marzo: https://github.com/mlc-ai/web-stable-diffusion
  Lo más impresionante es que después incluso añadieron soporte para varios modelos de lenguaje grandes: https://webllm.mlc.ai/

  • De verdad es impresionante y el rendimiento también se ve mucho mejor. Yo estaba siguiendo otro enfoque para ejecutar modelos ONNX arbitrarios sin modificaciones previas

• “Carga 3.5 GB y usa 8 GB de RAM”
  Es curioso en lo que se ha convertido el navegador. La web se comió al sistema operativo

  • No entiendo muy bien esa crítica. Estamos ejecutando Stable Diffusion en mi computadora a través del navegador, así que ¿cómo se supondría que se haga sin descargarlo y cargarlo en la RAM?
    Incluso si lo descargas y lo ejecutas directamente sin navegador, el tamaño de descarga y el uso de RAM serían casi los mismos
  • La web se comió la premisa original de Java
  • Por eso hay una empresa que vende laptops donde el sistema operativo es básicamente el navegador, y otra que hace lo mismo en las smart TV
  • A medida que aumenta el ancho de banda y madura el sandbox de la web, es interesante ver cómo evolucionan las apps hacia algo que simplemente usas, en vez de descargar, instalar y mantener. A algunos no les gustará, pero para el público general abre muchas puertas
  • Ahora me parece que WebGPU sirve para esto además de los casos típicos como páginas WebXR interactivas en tiempo real, estado multijugador por streaming y muchos draw calls.
    Hay algo que se simplifica al ofrecer este tipo de experiencia desde el navegador, y no parece haber una forma más fácil de ejecutar Stable Diffusion, así que ojalá este tipo de proyectos siga recibiendo apoyo

• Si necesita una versión reciente de Chrome con las flags Experimental WebAssembly y Experimental WebAssembly JavaScript Promise Integration (JSPI) activadas, entonces supongo que me tocará esperar a que llegue a Firefox

  • Al menos Firefox ahora sí permite estilizar las barras de desplazamiento con CSS, así que eso está bien
  • Activé las chrome:flags solicitadas en Brave, pero aun así no funcionó. Nunca he descargado Chrome en mis Macs con M1 y no pienso empezar ahora
  • Lo probé activándolo en Canary, pero tampoco logré hacerlo funcionar
  • ¿Ni siquiera vas a intentar hacerlo en Chrome?
  • ¿Qué hace la segunda flag?

• Esto ya estaba implementado aquí también: https://websd.mlc.ai/#text-to-image-generation-demo
  MLC usa Apache TVM para generación de código WebGPU y auto-tuning, y el rendimiento es bastante bueno

• Integré Stable Diffusion de MLC-AI en mi sitio web como un generador de fondos personalizado. Los datos del modelo vienen de Hugging Face
  https://dustinbrett.com/

  • Ver algo así de fluido y rápido hace triste que todas las apps del teléfono pidan 200 MB solo para mostrar unas cuantas pantallas y hacer unas llamadas a APIs
  • Cargó bien en un iPhone 12 mini y el teléfono ni siquiera se calentó. De verdad impresionante
  • El sitio web está impresionante y todo funciona con mucha fluidez. Al final terminé jugando Doom mucho más tiempo del necesario

• Para trabajos limitados por CPU, valdría la pena considerar service workers para que el hilo principal no se quede congelado

• ¿Lo siguiente será un chatbot con WebGPU? ¿Un chat al estilo ChatGPT en el navegador usando mi GPU?

  • Ya se puede hacer. Tiene algunas limitaciones
    Los modelos de lenguaje grandes consumen mucha RAM/VRAM cuanto mejores son, así que son especialmente complicados en WebGPU
  • https://webllm.mlc.ai/
    Ya está implementado

• Por curiosidad, ¿para qué se usaría esto?
  Lo que entiendo es que genera imágenes en el navegador y no en un servidor. Lo único que se me ocurre es que no haría falta recargar la página para modificar o generar una imagen nueva.
  Entonces quizá eso significaría sitios web donde el diseño visual cambia en tiempo real, y si también pudiera cambiar de forma funcional de una manera significativa sería bastante genial. Aun así, no tengo claro qué tan útil sería Stable Diffusion para generar componentes de UI o elementos visuales de un sitio

  • Es útil porque generar imágenes en masa sale caro. Básicamente mueve el costo de cómputo al cliente
  • El uso es poder ejecutarlo fácilmente en tu propia computadora. Ni siquiera hace falta ser programador ni instalar una aplicación
  • Ejecutar el modelo de forma completamente local tiene una gran ventaja de privacidad
    los prompts sensibles no se filtran hacia alguien remoto

• ¿Hay alguna métrica que compare el rendimiento en WebGPU con el rendimiento nativo?

  • UNET tarda alrededor de 1 minuto 10 segundos en WebGPU, y como 1 minuto en CPU de un solo hilo. VAE tarda 2 minutos en CPU y unos 10 segundos en GPU.
    Parece que eso se debe a que la mayor parte del cómputo de GPU para VAE ya está implementado, pero no el de UNET. En este último caso, el navegador sigue lanzando datos entre la GPU y la CPU en cada paso

• Si fuera lo bastante rápido, podría usarse para renderizar imágenes localmente para uso personal. Un sitio web podría pasar solo el prompt y luego renderizar una imagen distinta para cada usuario.
  ¿Cómo funcionaría el copyright en ese punto? ¿El modelo en sí tendría copyright o se derrumbaría todo el sistema?

  • Esa posibilidad es interesante, pero por ahora estamos muy lejos de un mundo así. En otras partes de este hilo se menciona que esto usa activamente 8 GB de RAM.
    Y no creo que muchos diseñadores web acepten el riesgo de que el modelo interprete mal un prompt, genere resultados distorsionados con un número incorrecto de dedos, o produzca accidentalmente contenido sexual o violento en contextos no deseados.
    Hoy en día, con muchos modelos generativos de imágenes, normalmente eliges la mejor de diez o más imágenes, y las demás que descartas pueden ser bastante malas en realidad.
    Para ilustrar sitios web dinámicamente de forma cotidiana, la calidad y previsibilidad del modelo tendrían que ser mucho mayores de lo que son ahora.
    Eso no significa que quiera decir que nunca llegaremos ahí. Los modelos recientes ya están haciendo cosas que hace apenas unos años se consideraban inimaginables. Comparado con https://xkcd.com/1425/, para los lectores jóvenes quizá hasta sea difícil explicar el problema de fondo detrás de ese chiste.