Demo do Stable Diffusion com WebGPU

 (islamov.ai)
1 pontos por GN⁺ 2023-07-19 | 1 comentários | Compartilhar no WhatsApp
  • Demo de geração de imagens com Stable Diffusion executada diretamente no navegador, funcionando com base em WebGPU e WebAssembly, sem necessidade de servidor separado
  • Para usar, é necessário ativar as flags "Experimental WebAssembly" e "Experimental WebAssembly JavaScript Promise Integration (JSPI)" no Chrome mais recente
  • Gera imagens a partir de Prompt, Negative Prompt e número de inference steps, e os arquivos do modelo são armazenados em cache, então não precisam ser baixados novamente a cada vez
  • Resultado de patches em onnxruntime, emscripten, binaryen para permitir alocação de memória acima de 4GB e de um port de StableDiffusionPipeline em Python para JS
  • Atualmente é lento devido à falta de suporte a multithreading e à ausência de implementação de operações WebGPU, mas indica potencial de melhoria futura com implementação de kernels em JS e suporte a memory64

Requisitos para execução

  • É necessário ativar as flags Experimental WebAssembly e Experimental WebAssembly JavaScript Promise Integration (JSPI) no navegador Chrome mais recente
  • Os Model files são armazenados em cache, então não é preciso baixá-los novamente em visitas futuras

Entradas e modo de funcionamento

  • Os campos de entrada são Prompt, Negative Prompt e número de inference steps
    • Como usa o PNDM Scheduler, o step real é processado como i+1 a partir do valor inserido i
  • Cada step leva cerca de 1 minuto, e a execução do VAE decoder para gerar a imagem leva cerca de 10 segundos adicionais
  • Com o DevTools aberto, a velocidade total fica aproximadamente 2 vezes mais lenta
  • O UNET roda apenas na CPU, sendo 10% mais rápido que na GPU; na GPU ele não produz resultados corretos e faz a aba do navegador travar
  • O número mínimo de steps para obter resultados aceitáveis é 20, embora 3 steps também sirvam para fins de demonstração

FAQ

  • Erro protobuf parsing failed

    • No DevTools, vá em Application → Storage e execute "Clear site data"

  • Erro sbox_fatal_memory_exceeded

    • Indica falta de RAM necessária para executar o SD; tente recarregar a aba ou o navegador

  • Forma de implementação

    • StableDiffusionPipeline em Python foi portada para JS
    • Foram aplicados patches em onnxruntime e em emscripten+binaryen (toolchain compiladora de WebAssembly) para suportar alocação e uso de memória acima de 4GB
    • Quando os pull requests relacionados forem incorporados às releases, qualquer pessoa poderá compilar e executar no navegador código que use mais de 4GB de memória

  • Motivos da lentidão

    • Ainda não há suporte a multithreading, então apenas um núcleo de CPU é usado
    • Não é possível criar memória de 64 bits com SharedArrayBuffer por meio do construtor WebAssembly.Memory
    • Foi proposta uma mudança na especificação da flag "memory64", e após adoção haverá suporte por meio de patch no motor V8

  • Execução na GPU

    • Funciona na GPU, mas o WebGPU do onnxruntime ainda está em estágio inicial e muitas operações não foram implementadas
    • Os dados continuam sendo transferidos continuamente entre CPU e JS
    • Quando kernels em JS forem implementados para a maioria das operações, a velocidade poderá melhorar bastante

  • É possível rodar localmente?

    • Sim, e o código da página está disponível no repositório stable-diffusion-webgpu-minimal

  • É possível executar LLMs grandes com transformers.js?

    • É possível usar o pacote onnxruntime com patch (@aislamov/onnxruntime-web64), mas não há garantia de funcionamento em todos os casos
    • Esse build é limitado a 8GB de memória, então é possível carregar pesos de aproximadamente 4GB

  • Planos de pull request para o repositório onnxruntime

    • Está previsto, sendo este o segundo trabalho após a adição anterior de aceleração por GPU ao binding de node.js

Leituras relacionadas

1 comentários

 
GN⁺ 2023-07-19
Comentários no Hacker News

  • A equipe da MLC já tinha feito isso funcionar em março: https://github.com/mlc-ai/web-stable-diffusion
    Mais impressionante ainda é que depois eles adicionaram suporte a vários modelos de linguagem grandes: https://webllm.mlc.ai/

    • Realmente impressionante, e o desempenho também parece bem melhor. Eu estava seguindo outra abordagem para rodar modelos ONNX arbitrários sem modificações prévias

  • “Carrega 3,5 GB e usa 8 GB de RAM”
    É curioso que o navegador tenha chegado a esse ponto. A web engoliu o sistema operacional

    • Não entendo muito bem essa crítica. Você está rodando o Stable Diffusion no seu computador pelo navegador; como faria isso sem baixar e colocar na RAM?
      Mesmo baixando e executando diretamente sem navegador, o tamanho do download e o uso de RAM seriam quase os mesmos
    • A web engoliu a premissa original do Java
    • Por isso existem empresas que vendem notebooks em que o sistema operacional é basicamente o navegador, e empresas que fazem a mesma coisa em smart TVs
    • É interessante ver como, com mais largura de banda e um sandbox da web mais maduro, isso evolui para apps que você simplesmente usa, em vez de baixar, instalar e manter. Algumas pessoas vão odiar, mas isso abre muitas portas para o público em geral
    • Agora estou achando que o WebGPU é útil para isso também, e não só para casos como páginas WebXR com interação em tempo real, estado multiplayer por streaming e muitos draw calls.
      Há uma simplificação em oferecer esse tipo de experiência pelo navegador, e não parece haver jeito mais fácil que esse para rodar Stable Diffusion, então espero que projetos assim continuem recebendo apoio

  • Então precisa do Chrome mais recente com as flags Experimental WebAssembly e Experimental WebAssembly JavaScript Promise Integration (JSPI) ativadas; acho que vou esperar chegar no Firefox

    • Pelo menos o Firefox agora permite estilização da barra de rolagem em CSS, então isso é bom
    • Ativei as chrome:flags pedidas no Brave, mas ainda assim não funcionou. Nunca baixei o Chrome nos meus Macs M1 e não pretendo começar agora
    • Testei no Canary, mas mesmo assim não consegui fazer funcionar
    • Você quer dizer que nem vai tentar no Chrome?
    • O que a segunda flag faz?

  • Isso também já foi implementado aqui: https://websd.mlc.ai/#text-to-image-generation-demo
    A MLC usa o Apache TVM para geração de código WebGPU e autotuning, e o desempenho é bem decente

  • Integrei o Stable Diffusion da MLC-AI no meu site como um gerador de fundos personalizado. Os dados do modelo vêm do Hugging Face
    https://dustinbrett.com/

    • Ver algo tão fluido e rápido assim só me deixa triste com o fato de todos os apps de celular exigirem 200 MB só para mostrar algumas telas e fazer algumas chamadas de API
    • Carregou bem no iPhone 12 mini e o celular nem esquentou. Realmente impressionante
    • O site é impressionante e tudo funciona de forma muito fluida. Acabei jogando Doom por muito mais tempo do que precisava

  • Para tarefas limitadas pela CPU, talvez valha considerar usar service workers para evitar que a thread principal trave

  • O próximo passo é um chatbot em WebGPU? Chat no estilo ChatGPT no navegador usando minha GPU?

    • Isso já existe. Há algumas limitações
      Modelos de linguagem grandes consomem muita RAM/VRAM quanto melhores são, então no WebGPU isso é especialmente complicado
    • https://webllm.mlc.ai/
      Já foi implementado

  • Curiosidade: para que isso serve?
    O que eu entendo é que as imagens são geradas no navegador, e não no servidor. A única vantagem que me vem à cabeça é não precisar recarregar a página para alterar ou gerar uma nova imagem.
    Isso talvez signifique que dá para ter sites com design visual mudando em tempo real, e seria bem legal se isso também pudesse mudar de forma funcional e significativa. Só não sei o quanto o Stable Diffusion seria útil para gerar componentes de UI ou elementos visuais de um site

    • É útil porque gerar imagens em grande volume é caro. Você está transferindo o custo computacional para o cliente
    • O uso é poder rodar isso facilmente no próprio computador. Não é preciso ser programador nem instalar um aplicativo
    • Rodar o modelo de forma puramente local traz uma grande vantagem de privacidade
      Prompts sensíveis não vazam para alguém remoto

  • Existe alguma métrica comparando o desempenho no WebGPU com o desempenho nativo?

    • O UNET leva cerca de 1 minuto e 10 segundos no WebGPU, e cerca de 1 minuto em CPU single-thread. O VAE leva 2 minutos na CPU e cerca de 10 segundos na GPU.
      Parece que a maior parte da computação de GPU para o VAE já foi implementada, mas a do UNET ainda não. Nesse último caso, o navegador fica jogando dados entre GPU e CPU a cada etapa

  • Se for rápido o suficiente, isso poderia ser usado para renderizar imagens localmente para uso pessoal. Um site poderia enviar só o prompt e renderizar imagens diferentes para cada usuário.
    Como ficaria a questão dos direitos autorais nesse ponto? O próprio modelo estaria sujeito a copyright, ou todo o sistema deixaria de fazer sentido?

    • Essa possibilidade é interessante, mas no momento estamos muito longe desse mundo. Em outras partes desta discussão, foi dito que isso usa ativamente 8 GB de RAM.
      E acho que muitos web designers não aceitariam o risco de o modelo interpretar mal o prompt, produzir resultados distorcidos com número errado de dedos, ou acidentalmente gerar conteúdo sexual ou violento em contextos inadequados.
      Hoje, muitos modelos de geração de imagem normalmente exigem escolher a melhor entre dez ou mais imagens, e as descartadas podem ser realmente bem ruins.
      Para ilustrar sites de forma dinâmica no dia a dia, a qualidade e previsibilidade do modelo teriam de ser muito melhores do que são agora.
      Isso não quer dizer que eu ache impossível chegar lá. Os modelos recentes já fazem coisas que, poucos anos atrás, seriam consideradas inimagináveis. Comparando com https://xkcd.com/1425/, pode até ser difícil explicar para leitores mais jovens qual era o contexto da piada!