Stable Diffusion XL 1.0 rodando com 298MB de RAM

 (github.com/vitoplantamura)
4 pontos por GN⁺ 2023-10-04 | 1 comentários | Compartilhar no WhatsApp
  • OnnxStream é uma biblioteca de inferência ultracompacta criada para rodar Stable Diffusion 1.5 e SDXL 1.0 Base em um Raspberry Pi Zero 2 com 512MB de RAM, sem adicionar swap nem gravar resultados intermediários em disco
  • Ela separa o motor de inferência de WeightsProvider, oferecendo suporte a DiskNoCache, DiskPrefetch e carregamento, cache e prefetch personalizados; também pode buscar pesos diretamente via HTTP e processá-los sem gravar em disco
  • O SDXL 1.0 Base usa arquivos ONNX exportados da implementação do Hugging Face Diffusers 0.19.3, e leva cerca de 11 horas para gerar uma imagem de 10 steps no Raspberry Pi Zero 2
  • O decodificador VAE do SDXL consome 4,4GB de RAM em FP32 e não é fácil resolver isso só com FP16 ou UINT8, então o uso de memória caiu para 298MB com decodificação em blocos sobrepostos e blending
  • Em CPU, pode usar até 55 vezes menos memória que o OnnxRuntime, em troca de ser entre 0,5 e 2 vezes mais lento, com limitações como batch size 1 e foco em alguns operadores ONNX

Objetivos e design do OnnxStream

  • OnnxStream começou com o objetivo de rodar Stable Diffusion 1.5 em um Raspberry Pi Zero 2
    • O Raspberry Pi Zero 2 é um microcomputador com 512MB de RAM
    • O mínimo recomendado de RAM/VRAM para Stable Diffusion normalmente é 8GB
    • As condições do objetivo eram: sem espaço de swap adicional e sem descarregar resultados intermediários em disco
  • Frameworks comuns de machine learning tendem a priorizar latência de inferência ou throughput, o que facilita aumento no uso de RAM
  • O OnnxStream é uma biblioteca de inferência pequena e fácil de modificar, com foco em minimizar o uso de memória
  • A estrutura central separa o motor de inferência do componente de fornecimento de pesos do modelo, o WeightsProvider
    • Classes derivadas podem implementar formas de carregar, armazenar em cache e fazer prefetch dos parâmetros do modelo
    • Um WeightsProvider customizado pode baixar dados diretamente de um servidor HTTP e processá-los sem leitura ou gravação em disco
    • As implementações padrão são DiskNoCache e DiskPrefetch
  • Em CPU, usa até 55 vezes menos memória que o OnnxRuntime, embora possa ser de 0,5 a 2 vezes mais lento

Rodando Stable Diffusion 1.5

  • O exemplo de Stable Diffusion 1.5 gera imagens com o OnnxStream usando diferentes precisões do decodificador VAE
  • O único componente que não cabia na RAM do Raspberry Pi Zero 2 era o decodificador VAE
    • A causa eram conexões residuais, tensores grandes e convoluções grandes dentro do modelo
    • Nem precisão simples nem meia precisão cabiam na RAM do Raspberry Pi Zero 2
    • A solução foi quantização estática de 8 bits
  • A imagem com decodificador VAE em precisão W8A8 foi gerada no Raspberry Pi Zero 2 e levou cerca de 1,5 hora com a opção MAX_SPEED
  • A imagem W16A16 para comparação foi gerada em um PC usando o mesmo latent

Suporte a Stable Diffusion XL 1.0 Base

  • O exemplo de Stable Diffusion do OnnxStream oferece suporte a SDXL 1.0 Base, sem incluir o Refiner
  • Os arquivos ONNX foram exportados da implementação SDXL 1.0 da biblioteca Diffusers do Hugging Face, na versão 0.19.3
  • O SDXL 1.0 tem custo computacional muito maior que o SD 1.5
    • A maior diferença é que ele pode gerar imagens de 1024x1024 em vez de 512x512
    • Em um PC de 12 núcleos com 32GB de RAM, gerar uma imagem de 10 steps com Hugging Face Diffusers leva 26 minutos
    • O mínimo recomendado de VRAM para SDXL normalmente é 12GB
  • O OnnxStream consegue rodar o SDXL 1.0 com menos de 300MB de RAM, então também funciona no Raspberry Pi Zero 2
    • Sem swap adicional
    • Sem escrita em disco durante a inferência
    • A geração de uma imagem de 10 steps no Raspberry Pi Zero 2 leva cerca de 11 horas

Otimizações de memória específicas para SDXL

  • O SDXL 1.0 usa o mesmo conjunto de otimizações do SD 1.5, com algumas diferenças
  • O modelo UNET usa quantização dinâmica UINT8 para rodar com menos de 300MB de RAM no Raspberry Pi Zero 2
    • A quantização é limitada a um subconjunto específico de grandes tensores intermediários
  • O decodificador VAE do SDXL 1.0 é mais difícil de tratar que o do SD 1.5
    • O decodificador VAE do SDXL 1.0 é 4 vezes maior que o do SD 1.5
    • Em FP32 no OnnxStream, ele usa 4,4GB de RAM
    • No SD 1.5, era possível reduzir o uso de RAM para 260MB com quantização estática UINT8 do decodificador VAE
  • O decodificador VAE do SDXL 1.0 sofre overflow em aritmética FP16, e a faixa numérica das ativações é grande demais para obter imagens de boa qualidade com quantização UINT8
  • Existem soluções FP16 como sdxl-vae-fp16-fix, mas mesmo reduzindo a memória pela metade ainda fica em 2,2GB, o que continua grande demais para o Raspberry Pi Zero 2
  • A solução final foi a decodificação em blocos inspirada na implementação do decodificador VAE do Diffusers
    • O tensor resultante da diffusion tem shape (1,4,128,128)
    • Ele é dividido em 5x5 blocos sobrepostos, 25 no total, cada um com shape (1,4,32,32)
    • Cada bloco se sobrepõe 25% ao bloco da esquerda e ao bloco de cima
    • O resultado da decodificação de cada bloco é um tensor (1,3,256,256), depois mesclado na imagem final
    • Com o blending desativado, as bordas dos blocos ficam visíveis; com ele ativado, elas desaparecem no resultado final
  • Com essa abordagem, o uso de RAM do decodificador VAE do SDXL cai de 4,4GB para 298MB

Recursos suportados e dependências

  • O OnnxStream reúne em um pacote pequeno os recursos necessários para inferência com economia de memória
    • Separação entre motor de inferência e WeightsProvider
    • DiskNoCache, DiskPrefetch e WeightsProvider customizado
    • attention slicing
    • quantização dinâmica assimétrica unsigned de 8 bits por percentil
    • quantização estática W8A8 assimétrica unsigned por percentil
    • calibração de modelo quantizado
    • suporte a FP16
    • implementação de 25 operadores ONNX comuns
  • As operações são executadas em sequência, mas cada operador funciona com multithreading
  • A implementação segue a estrutura de um único arquivo de implementação e um arquivo de cabeçalho, e a classe XnnPack encapsula as chamadas ao XNNPACK
  • Alguns primitives de aceleração dependem de XNNPACK
    • MatMul
    • Convolution
    • element-wise Add/Sub/Mul/Div
    • Sigmoid
    • Softmax

Comparação de desempenho e limitações

  • O Stable Diffusion é composto por três modelos
    • text encoder: 672 operations, 123 milhões de parameters
    • UNET: 2050 operations, 854 milhões de parameters
    • VAE decoder: 276 operations, 49 milhões de parameters
  • Para gerar uma imagem de 10 steps com batch size 1, são necessárias as seguintes execuções
    • text encoder 2 vezes
    • UNET 20 vezes
    • VAE decoder 1 vez
  • Com base no UNET FP16, a diferença de memória e tempo entre OnnxStream e OnnxRuntime é grande
    • OnnxStream: cerca de 0,133GB, 18,2 a 19,8 segundos
    • OnnxRuntime: 5,085 a 7,353GB, 7,28 a 12,8 segundos
    • O OnnxStream usa até 55 vezes menos memória, mas é de 0,5 a 2 vezes mais lento
  • No text encoder FP32, o OnnxStream usa 0,147GB e o OnnxRuntime usa 0,641GB
  • No VAE decoder FP32, o OnnxStream usa 1,004GB e o OnnxRuntime usa 1,330 a 2,026GB
  • Os resultados da comparação vêm com algumas condições
    • A primeira execução do OnnxRuntime foi uma inferência de warm-up antes da reutilização do InferenceSession
    • O OnnxStream tem design eager e não possui conceito de warm-up, mas execuções posteriores podem se beneficiar do cache de arquivos de pesos do sistema operacional
    • Atualmente, o OnnxStream não oferece suporte a entradas com batch size diferente de 1
    • O OnnxRuntime pode usar batch size 2 na execução do UNET e acelerar bastante a diffusion como um todo
    • Alterar EnableCpuMemArena e ExecutionMode nas SessionOptions do OnnxRuntime não trouxe diferença significativa nos testes
    • O NCNN foi muito parecido com o OnnxRuntime em uso de memória e tempo de inferência
    • O ambiente de teste foi Windows Server 2019, 16GB de RAM, CPU 8750H com AVX2, SSD 970 EVO Plus e 8 núcleos virtuais no VMWare

Attention slicing e quantização

  • Durante a execução do UNET, o attention slicing e a quantização W8A8 do decodificador VAE foram cruciais para reduzir a memória a um nível viável no Raspberry Pi Zero 2
  • O attention slicing evita materializar a matriz completa Q @ K^T ao calcular scaled dot-product attention em multi-head attention
  • Quando o modelo UNET tem 8 attention heads, os shapes dos tensores são os seguintes
    • Q: (8,4096,40)
    • K^T: (8,40,4096)
    • Resultado do primeiro MatMul: (8,4096,4096)
    • Em precisão FP32, isso vira um tensor de 512MB
  • A solução é dividir Q verticalmente e executar a atenção para cada chunk
    • O shape de Q_sliced é (1,x,40)
    • x é o valor de 4096 dividido por onnxstream::Model::m_attention_fused_ops_parts
    • O valor padrão é 2, mas pode ser customizado
  • Com isso, o uso total de memória do modelo UNET FP32 cai de 1,1GB para 300MB
  • O FlashAttention pode ser uma alternativa mais eficiente, mas exigiria kernels customizados para cada arquitetura suportada, como AVX e NEON, além de contornar o XNNPACK

Conversão de modelos e modo de execução

  • O OnnxStream executa modelos definidos em path_to_model_folder/model.txt
    • Todas as operações do modelo ficam em model.txt em formato ASCII
    • Os arquivos de pesos devem existir como uma série de arquivos .bin na mesma pasta
  • O objeto Model permite configurar vários parâmetros opcionais
    • Especificar outro weights provider
    • Ler e gravar arquivos de faixa de clipping de ativações do modelo quantizado
    • Modo de calibração do modelo
    • Uso de aritmética FP16
    • Uso de aritmética UINT8
    • Uso de quantização dinâmica UINT8
    • Ativação de attention slicing
    • Definição da quantidade de divisões da attention
  • Para usar arquivos ONNX no OnnxStream, o notebook onnx2txt.ipynb exporta model.txt e os arquivos de pesos .bin
  • Ao exportar um nn.Module do PyTorch para ONNX, há algumas condições
    • Ao chamar torch.onnx.export, dynamic_axes deve ficar vazio
    • O OnnxStream não oferece suporte a entradas com shape dinâmico
    • É fortemente recomendado executar o ONNX Simplifier antes da conversão

Preparação de build e execução

  • É possível compilar o exemplo de Stable Diffusion em Linux, Mac, Windows e Termux
    • No Windows, usa-se o x64 Native Tools Command Prompt do Visual Studio Tools
    • No Mac, é preciso instalar o cmake com brew install cmake
  • Primeiro é necessário compilar o XNNPACK
    • Como os prototypes de função do XNNPACK podem mudar a qualquer momento, o procedimento inclui fazer checkout de um commit compatível
    • O ponto de referência é um commit da branch master anterior a 2023-06-27 00:00
  • Depois disso, clone o repositório do OnnxStream e faça o build em src/build com cmake
    • MAX_SPEED=ON
    • Definindo XNNPACK_DIR para o caminho do clone do XNNPACK
  • A opção MAX_SPEED pode melhorar o desempenho, mas usa mais memória durante o build
    • No Windows, traz cerca de 10% de ganho de desempenho
    • No Raspberry Pi, traz mais de 50% de ganho de desempenho
    • O executável gerado pode não funcionar, e houve problemas nos testes com Termux
    • Se houver problema, a recomendação é primeiro mudar para MAX_SPEED=OFF
  • Os weights do exemplo Stable Diffusion 1.5 podem ser baixados na área de Releases do repositório e têm cerca de 2GB
  • Os weights do Stable Diffusion XL 1.0 Base podem ser baixados no Hugging Face e têm cerca de 8GB

Opções do exemplo de Stable Diffusion

  • O executável de exemplo controla seleção de modelo, entrada e saída, prompt e método de decodificação por opções de linha de comando
    • --xl: roda Stable Diffusion XL 1.0 em vez de Stable Diffusion 1.5
    • --models-path: define a pasta do modelo Stable Diffusion
    • --ops-printf: imprime no stdout a operação atual durante a inferência
    • --output: define o arquivo PNG de saída
    • --decode-latents: pula a diffusion e decodifica o arquivo de latents especificado
    • --prompt: define o prompt positivo
    • --neg-prompt: define o prompt negativo
    • --steps: define o número de diffusion steps
    • --save-latents: salva os latents no arquivo especificado após a diffusion
  • Também há opções específicas para Raspberry Pi e para o decodificador
    • --decoder-calibrate: calibra o VAE decoder quantizado apenas no SD 1.5
    • --decoder-fp16: usa VAE decoder FP16 apenas no SD 1.5
    • --not-tiled: desativa o tiled VAE decoder apenas no SDXL 1.0
    • --rpi: configura o modelo para rodar no Raspberry Pi
    • --rpi-lowmem: aplica a configuração de baixa memória para Raspberry Pi Zero 2 apenas no SDXL 1.0

Leituras relacionadas

1 comentários

 
GN⁺ 2023-10-04
Opiniões no Hacker News

  • Interessante. A frase central é esta: “O OnnxStream pode ser apenas 0,5 a 2 vezes mais lento que o OnnxRuntime, usando até 55 vezes menos memória”
    A troca entre memória de vídeo/uso de memória e tempo de inferência parece vantajosa não só em casos com memória limitada, como no Raspberry Pi, mas também em outras situações
    Fico curioso se, na prática, caso esse método de descarregar pesos permita processar batches maiores dentro da mesma memória, não seria possível aumentar bastante o throughput mesmo com maior latência

    • Eu gostaria de usar isso em LLMs. Se a ocupação de memória cair tanto assim, dá para colocar mais modelos ao mesmo tempo em uma única GPU e, assumindo que o clock acompanhe, talvez compensar bem a perda de velocidade de inferência de cada modelo individual
    • “0,5 a 2 vezes mais lento” não é um erro de digitação? Se for 0,5 vez mais lento, isso na verdade significa 2 vezes mais rápido
      Provavelmente queriam escrever “1,5 a 2 vezes mais lento”
    • Pelo que entendi, de forma um tanto simplificada, eles não reduziram a largura de banda de memória necessária, apenas o tamanho do conjunto de trabalho
      Em inferência, depois que você passa da questão de “este modelo cabe neste sistema?”, normalmente o gargalo é a largura de banda de memória, então não acho que essa técnica ajude muito a aumentar o throughput com batches maiores. É bem provável que uma única instância já esteja saturando o controlador de memória
      Mas talvez ajude no lado de treinamento

  • 11 horas me faz lembrar da época em que eu fazia ray tracing em um Amiga 500. Um render “final” era definitivamente algo para deixar rodando durante a noite

    • Hoje em dia ainda faço isso às vezes. Ray tracers bidirecionais modernos permitem usar técnicas interessantes, e eu queria ver cáusticas, como aquelas linhas brilhantes que aparecem em poças d’água
      Mas, embora as cáusticas pareçam brilhantes, estatisticamente são um fenômeno bem raro; para obtê-las direito, é preciso remover os limites do motor de renderização e simplesmente deixar rodando a noite toda
      O resultado é uma imagem de uma cena comum feita por um artista pouco habilidoso, mas com cáusticas bonitas. Melhor eu continuar no meu emprego principal
    • Rodar primeiro uma renderização de baixa qualidade era porque é melhor desperdiçar só uma hora para confirmar se está certo do que desperdiçar a noite toda com algo errado
      Foi mais ou menos quando achei que precisava de outro hobby. Logo antes disso, alguém muito talentoso lançou uma ferramenta que renderizava a cena primeiro em OpenGL para visualização. Não teria funcionado no Amiga, mas na minha máquina rodava por pouco
    • No meu 286 era igual. Eu deixava o povray configurado, ia dormir e conferia a imagem de manhã antes de ir para a escola
    • Um pouco depois, fiz algo parecido em um 386 com uma cópia suspeita do 3DSMAX
    • Lembro de rodar fractais de Mandelbrot no C64. Depurar o código era realmente difícil

  • Estou usando Stable Diffusion com invoke.ai em um MBP; alguém tem recomendações para ajustar melhor os parâmetros do SD? Mesmo usando o mesmo prompt e configurações aparentemente iguais, por exemplo o mesmo modelo com Euler A, não consigo de jeito nenhum chegar à qualidade das imagens que vejo na internet

    • Entre os que testei até agora, este foi o melhor, mas parece não ter suporte para Mac. É um fork cheio de recursos do Fooocus, feito pelo desenvolvedor original do ControlNet, e a qualidade que sai com prompts curtos é impressionante: https://github.com/MoonRide303/Fooocus-MRE
      Para SD 1.5 básico uso o Volta por ser rápido: https://github.com/VoltaML/voltaML-fast-stable-diffusion/com...
      Uma qualidade de imagem realmente boa no SD 1.5 exige usar sem economizar modelos com fine-tuning, LoRA, ControlNet e outros recursos de reforço. Por exemplo, fazer a imagem seguir uma imagem-base como estrutura, ou definir prompts por áreas específicas da imagem. O InvokeAI também tem muitos recursos de fato, e muitos desses reforços ficam escondidos na UI de nós, mas outras UIs, como o Volta, os expõem de forma mais direta
    • Você está usando pesos customizados? Imagino que sim, mas há uma grande diferença entre os pesos básicos RunwayML 1.5 e modelos ajustados por fine-tuning para objetivos específicos
      Em geral, modelos com fine-tuning impressionantes são bem menos versáteis que os pesos básicos, mas no uso real isso não costuma ser um grande problema e pode melhorar bastante os resultados
    • Tive a mesma experiência usando Invoke.ai ou MochiDiffusion em um MBP M1. O único com que consegui igualar a qualidade de outras imagens foi o Automatic1111(https://github.com/AUTOMATIC1111/stable-diffusion-webui)
      Ele exige mais tempo e memória do que o Invoke ou uma placa de vídeo Nvidia, mas não é tão ruim. Uma imagem padrão de qualidade em 512x768px fica em torno de 1 a 2 s/it, e uma imagem de alta qualidade em 1024x1536px usando Hires Fix fica por volta de 14 a 20 s/it
    • Está explícito que essas imagens saíram diretamente do gerador? Nos vídeos de processo que vi, eles começavam com “uma garota em pé em um campo verde” e depois passavam mais de uma hora fazendo inpainting para corrigir mãos, pose e afins
    • O Draw Things adicionou um modo de seed compatível com CUDA, então no Mac também dá para reproduzir imagens geradas em placas NVDIA

  • Acho que seria muito legal embutir isso em um porta-retratos digital ou quadro de parede

  • É uma conquista impressionante, mas é claro que leva um tempo enorme para gerar uma imagem. No README diz 11 horas

    • Sim. Eu não preciso nem vou usar esta implementação, mas as técnicas usadas aqui vão migrar para outras ferramentas, e isso vai ser ótimo
    • Seria interessante ver como isso se compara, em termos de custo e energia, com abordagens mais avançadas
      Por exemplo, comparar o tempo, o custo, todo o hardware necessário e a energia para gerar 100 imagens com 100 Pi Zero 2 — nem precisaria ser o W — versus algo como um PC intermediário médio
      Provavelmente o PC ainda venceria
      O Zero 2 parece ter sido usado mais como desafio do que por praticidade, então o Pi 4 ou 5 talvez fossem pontos de referência melhores
    • O ponto principal é que roda em um Raspberry Pi, e ainda por cima em um Zero 2

  • Impressionante
    De fato, aproxima-se a era em que até lâmpadas e torradeiras abrigarão uma inteligência excepcional
    Acompanho esta área há muitos anos, mas a última década foi surpreendente
    Digo “foi” porque a aceleração dos últimos 6 a 18 meses está em outro patamar completamente diferente
    Não me preocupa o que poderemos fazer daqui a 2 anos, mas sim o quanto a velocidade de avanço vai aumentar ainda mais. E então vai acelerar de novo, e de novo

    • Uma torradeira que recebe um prompt e desenha a imagem na torrada seria ótima. Também daria para aproveitar o calor da GPU para torrar de verdade
      Vamos criar uma startup

  • A esta altura, as tentativas de regular tecnologias do tipo Stable Diffusion não já acabaram? Se conseguirem reduzir o modelo e a infraestrutura de inferência a ponto de rodar até em um PS2, parece impossível impedir essa tecnologia sem um Estado de vigilância totalitário — e talvez nem assim

    • A guerra contra a computação de uso geral continua em andamento, mas ainda não chegou ao ponto de impedir as pessoas de possuírem dispositivos de computação de uso geral
    • Por essa lógica, a regulação contra furtos também não acabou? Se você tem um corpo, consegue abrir uma janela sem ferramentas, então seria impossível impedir furtos sem um Estado de vigilância totalitário
      O mesmo vale para “pirataria” de mídia ou ransomware
      Estados regulam há muito tempo coisas que não conseguem impor de forma puramente técnica
    • O requisito original desses modelos é 16 GB de RAM, e isso pode ser comprado por menos de 20 dólares. Eles rodam muito mais rápido em GPU, e essas GPUs também podem ser compradas por menos de 200 dólares. Milhões de pessoas comuns já têm os dois
    • O PS2 tinha apenas 32 MB de RAM. O PS3 também tinha só 256 MB
      Sei que é uma hipérbole divertida, mas para rodar em um PS2 teria que reduzir muito mais do que isso
    • Eu achava que a maioria das tentativas de regulação estava focada nos processos de treinamento, cada vez maiores, mais do que no uso generativo de modelos existentes. Há alguma regulação específica sobre o uso dos modelos em si?

  • Mal posso esperar pelo dia em que sair um Stable Diffusion para Windows 3.1

  • Insano. Levando 11 horas ou não, eu jamais teria imaginado que SD pudesse rodar em um hardware como o Pi Zero