Stable Diffusion XL 1.0 exécuté avec 298 Mo de RAM
(github.com/vitoplantamura)- OnnxStream est une bibliothèque d’inférence ultralégère visant à exécuter Stable Diffusion 1.5 et SDXL 1.0 Base sur un Raspberry Pi Zero 2 doté de 512 Mo de RAM, sans ajout de swap ni stockage intermédiaire sur disque
- Elle sépare le moteur d’inférence et le
WeightsProvider, prend en chargeDiskNoCache,DiskPrefetch, ainsi que le chargement, la mise en cache et le préchargement personnalisés, et peut aussi récupérer directement les poids via HTTP sans écriture sur disque - SDXL 1.0 Base utilise des fichiers ONNX exportés depuis l’implémentation Hugging Face Diffusers 0.19.3 ; sur Raspberry Pi Zero 2, la génération d’une image en 10 étapes prend environ 11 heures
- Le décodeur VAE de SDXL consomme 4,4 Go de RAM en FP32 et il est difficile de résoudre ce problème avec seulement FP16 ou UINT8 ; le décodage par tuiles superposées avec blending réduit la mémoire à 298 Mo
- Par rapport à OnnxRuntime sur CPU, il utilise jusqu’à 55 fois moins de mémoire, au prix d’une exécution pouvant être 0,5 à 2 fois plus lente, avec des limites comme une taille de batch de 1 et une prise en charge centrée sur certains opérateurs ONNX
Objectif et conception d’OnnxStream
- OnnxStream a été lancé avec pour objectif de faire tourner Stable Diffusion 1.5 sur un Raspberry Pi Zero 2
- Le Raspberry Pi Zero 2 est un micro-ordinateur doté de 512 Mo de RAM
- La RAM/VRAM minimale recommandée pour Stable Diffusion est généralement de 8 Go
- Les conditions visées sont : pas d’ajout d’espace de swap et pas d’offloading de résultats intermédiaires sur disque
- Les frameworks de machine learning classiques privilégient souvent la latence d’inférence ou le débit, ce qui tend à augmenter l’utilisation de la RAM
- OnnxStream est une petite bibliothèque d’inférence facile à modifier, centrée sur la minimisation de l’utilisation mémoire
- Son architecture clé consiste à séparer le moteur d’inférence du composant fournissant les poids du modèle, le
WeightsProvider- Des classes dérivées peuvent implémenter différentes stratégies de chargement, de mise en cache et de préchargement des paramètres du modèle
- Un
WeightsProviderpersonnalisé peut télécharger directement les données depuis un serveur HTTP et les traiter sans lecture ni écriture disque - Les implémentations par défaut sont
DiskNoCacheetDiskPrefetch
- Sur CPU, il utilise jusqu’à 55 fois moins de mémoire qu’OnnxRuntime, avec une exécution environ 0,5 à 2 fois plus lente
Exécution de Stable Diffusion 1.5
- L’exemple Stable Diffusion 1.5 utilise OnnxStream pour générer des images avec différentes précisions du décodeur VAE
- Le seul composant qui ne tenait pas dans la RAM du Raspberry Pi Zero 2 était le décodeur VAE
- La cause vient des connexions résiduelles internes au modèle, de grands tenseurs et de grandes convolutions
- La simple précision ou la demi-précision ne suffisaient pas à tenir dans la RAM du Raspberry Pi Zero 2
- La solution a été la quantification statique 8 bits
- Une image avec décodeur VAE en précision W8A8 a été générée sur Raspberry Pi Zero 2 ; avec l’option
MAX_SPEED, cela prend environ 1,5 heure - L’image W16A16 de comparaison a été générée sur PC avec le même latent
Prise en charge de Stable Diffusion XL 1.0 Base
- L’exemple Stable Diffusion d’OnnxStream prend en charge SDXL 1.0 Base, sans inclure le Refiner
- Les fichiers ONNX ont été exportés depuis l’implémentation SDXL 1.0 de la bibliothèque Hugging Face Diffusers, en version 0.19.3
- SDXL 1.0 est nettement plus coûteux en calcul que SD 1.5
- La principale différence est la capacité à générer des images en 1024x1024 au lieu de 512x512
- Sur un PC 12 cœurs avec 32 Go de RAM, la génération d’une image en 10 étapes avec Hugging Face Diffusers prend 26 minutes
- La VRAM minimale recommandée pour SDXL est généralement de 12 Go
- OnnxStream peut exécuter SDXL 1.0 avec moins de 300 Mo de RAM, ce qui lui permet de fonctionner même sur Raspberry Pi Zero 2
- Pas d’ajout de swap
- Pas d’écriture disque pendant l’inférence
- Sur Raspberry Pi Zero 2, la génération d’une image en 10 étapes prend environ 11 heures
Optimisations mémoire propres à SDXL
- SDXL 1.0 applique le même ensemble d’optimisations que SD 1.5, avec quelques différences
- Le modèle UNET utilise une quantification dynamique UINT8 afin de s’exécuter sur Raspberry Pi Zero 2 avec moins de 300 Mo de RAM
- La quantification est limitée à un sous-ensemble précis des grands tenseurs intermédiaires
- Le décodeur VAE de SDXL 1.0 est plus difficile à traiter que celui de SD 1.5
- Le décodeur VAE de SDXL 1.0 est 4 fois plus grand que celui de SD 1.5
- Exécuté en précision FP32 dans OnnxStream, il utilise 4,4 Go de RAM
- Dans SD 1.5, la quantification statique UINT8 du décodeur VAE permettait de réduire l’utilisation RAM à 260 Mo
- Le décodeur VAE de SDXL 1.0 produit des overflows en arithmétique FP16, et la plage numérique des activations est trop large pour obtenir des images de bonne qualité avec une quantification UINT8
- Il existe des solutions FP16 comme sdxl-vae-fp16-fix, mais même en divisant la mémoire par deux, on reste à 2,2 Go, ce qui est toujours trop pour un Raspberry Pi Zero 2
- La solution finale est le décodage par tuiles, inspiré de l’implémentation du décodeur VAE dans Hugging Face Diffusers
- Le tenseur résultat de la diffusion a pour shape
(1,4,128,128) - Il est divisé en tenseurs superposés de shape
(1,4,32,32), en une grille 5x5, soit 25 au total - Chaque tuile chevauche les tuiles de gauche et du haut de 25 %
- Le résultat du décodage de chaque tuile est un tenseur
(1,3,256,256), blendé dans l’image finale - Si le blending est désactivé, les frontières entre tuiles sont visibles ; s’il est activé, elles ne se voient pas dans le résultat final
- Le tenseur résultat de la diffusion a pour shape
- Cette méthode réduit l’utilisation RAM du décodeur VAE de SDXL de 4,4 Go à 298 Mo
Fonctionnalités prises en charge et dépendances
- OnnxStream regroupe de façon compacte les fonctionnalités nécessaires à l’inférence économe en mémoire
- Séparation du moteur d’inférence et du
WeightsProvider DiskNoCache,DiskPrefetch,WeightsProviderpersonnalisé- attention slicing
- quantification dynamique 8 bits unsigned asymmetric percentile
- quantification statique W8A8 unsigned asymmetric percentile
- calibration du modèle quantifié
- prise en charge de FP16
- implémentation de 25 opérateurs ONNX couramment utilisés
- Séparation du moteur d’inférence et du
- Les opérations sont exécutées séquentiellement, mais chaque opérateur fonctionne en multithread
- L’implémentation est structurée autour d’un unique fichier d’implémentation et d’un fichier d’en-tête, avec une classe
XnnPackqui encapsule les appels à XNNPACK - Certaines primitives d’accélération dépendent de XNNPACK
- MatMul
- Convolution
- Add/Sub/Mul/Div élément par élément
- Sigmoid
- Softmax
Comparaison des performances et limites
- Stable Diffusion est composé de trois modèles
- text encoder : 672 opérations, 123 millions de paramètres
- UNET : 2050 opérations, 854 millions de paramètres
- VAE decoder : 276 opérations, 49 millions de paramètres
- Pour générer une image en 10 étapes avec une taille de batch de 1, les exécutions suivantes sont nécessaires
- text encoder 2 fois
- UNET 20 fois
- VAE decoder 1 fois
- Avec UNET en FP16, l’écart de mémoire et de temps entre OnnxStream et OnnxRuntime est important
- OnnxStream : environ 0,133 Go, 18,2 à 19,8 secondes
- OnnxRuntime : 5,085 à 7,353 Go, 7,28 à 12,8 secondes
- OnnxStream utilise jusqu’à 55 fois moins de mémoire, mais est 0,5 à 2 fois plus lent
- Pour le text encoder FP32, OnnxStream utilise 0,147 Go contre 0,641 Go pour OnnxRuntime
- Pour le VAE decoder FP32, OnnxStream utilise 1,004 Go contre 1,330 à 2,026 Go pour OnnxRuntime
- Les résultats de comparaison sont soumis à plusieurs conditions
- La première exécution d’OnnxRuntime est une inférence de warm-up avant réutilisation de
InferenceSession - OnnxStream a une conception eager et n’a donc pas de notion de warm-up, mais les exécutions suivantes peuvent bénéficier du cache des fichiers de poids par l’OS
- OnnxStream ne prend actuellement pas en charge les entrées dont la batch size n’est pas 1
- OnnxRuntime peut accélérer nettement l’ensemble de la diffusion en utilisant une batch size de 2 pour l’exécution de l’UNET
- La modification de
EnableCpuMemArenaetExecutionModedans lesSessionOptionsd’OnnxRuntime n’a pas produit de différence significative dans les tests - NCNN était très proche d’OnnxRuntime en termes d’utilisation mémoire et de temps d’inférence
- L’environnement de test était Windows Server 2019, 16 Go de RAM, CPU 8750H avec AVX2, SSD 970 EVO Plus, VMWare avec 8 cœurs virtuels
- La première exécution d’OnnxRuntime est une inférence de warm-up avant réutilisation de
Attention slicing et quantification
- Lors de l’exécution de l’UNET, l’attention slicing et la quantification W8A8 du décodeur VAE ont été essentiels pour réduire la mémoire à un niveau exécutable sur Raspberry Pi Zero 2
- L’attention slicing évite de matérialiser toute la matrice
Q @ K^Tlors du calcul de la scaled dot-product attention dans le multi-head attention - Lorsque le nombre de têtes d’attention du modèle UNET est de 8, les shapes de tenseurs sont les suivantes
Q:(8,4096,40)K^T:(8,40,4096)- Résultat du premier MatMul :
(8,4096,4096) - En précision FP32, cela devient un tenseur de 512 Mo
- La solution consiste à découper
Qverticalement et à effectuer l’opération d’attention sur chaque chunk- La shape de
Q_slicedest(1,x,40) xest la valeur 4096 divisée paronnxstream::Model::m_attention_fused_ops_parts- La valeur par défaut est 2 et elle est personnalisable
- La shape de
- Cette méthode réduit l’utilisation mémoire totale du modèle UNET FP32 de 1,1 Go à 300 Mo
- FlashAttention pourrait être une alternative plus efficace, mais il faudrait écrire des kernels personnalisés pour chaque architecture prise en charge, comme AVX ou NEON, ce qui contournerait XNNPACK
Conversion du modèle et mode d’exécution
- OnnxStream exécute le modèle défini dans
path_to_model_folder/model.txt- Toutes les opérations du modèle se trouvent dans
model.txtau format ASCII - Les fichiers de poids doivent exister dans le même dossier sous la forme d’une série de fichiers
.bin
- Toutes les opérations du modèle se trouvent dans
- Plusieurs paramètres optionnels peuvent être configurés dans l’objet
Model- Choix d’un autre weights provider
- Lecture/écriture d’un fichier de plage de clipping des activations pour le modèle quantifié
- Mode de calibration du modèle
- Utilisation de l’arithmétique FP16
- Utilisation de l’arithmétique UINT8
- Utilisation de la quantification dynamique UINT8
- Activation de l’attention slicing
- Définition du nombre de partitions d’attention
- Pour utiliser un fichier ONNX avec OnnxStream, le notebook
onnx2txt.ipynbpermet d’exportermodel.txtet les fichiers de poids.bin - Lors de l’export d’un
nn.ModulePyTorch vers ONNX, certaines conditions s’appliquentdynamic_axesdoit rester vide lors de l’appel àtorch.onnx.export- OnnxStream ne prend pas en charge les entrées à shape dynamique
- Il est fortement recommandé d’exécuter ONNX Simplifier avant la conversion
Préparation de la compilation et de l’exécution
- L’exemple Stable Diffusion peut être compilé sous Linux, Mac, Windows et Termux
- Sous Windows, il faut utiliser le
x64 Native Tools Command PromptdeVisual Studio Tools - Sur Mac, il faut installer cmake avec
brew install cmake
- Sous Windows, il faut utiliser le
- Il faut d’abord compiler XNNPACK
- Les prototypes des fonctions XNNPACK pouvant changer à tout moment, la procédure inclut un checkout sur un commit correspondant à un moment précis
- Le repère est un commit master antérieur au
2023-06-27 00:00
- Il faut ensuite cloner le dépôt OnnxStream et compiler avec cmake depuis
src/buildMAX_SPEED=ON- Définir
XNNPACK_DIRavec le chemin du clone XNNPACK
- L’option
MAX_SPEEDpeut améliorer les performances, mais augmente l’utilisation mémoire lors de la compilation- Environ 10 % de gain de performance sous Windows
- Plus de 50 % de gain de performance sur Raspberry Pi
- L’exécutable généré peut ne pas fonctionner ; des problèmes ont été observés lors des tests sous Termux
- En cas de problème, il est recommandé d’essayer d’abord avec
MAX_SPEED=OFF
- Les poids de l’exemple Stable Diffusion 1.5 sont disponibles dans les Releases du dépôt et pèsent environ 2 Go
- Les poids de Stable Diffusion XL 1.0 Base sont disponibles sur Hugging Face et pèsent environ 8 Go
Options de l’exemple Stable Diffusion
- L’exécutable d’exemple contrôle le choix du modèle, les entrées/sorties, les prompts et le mode de décodage via des options de ligne de commande
--xl: exécuter Stable Diffusion XL 1.0 au lieu de Stable Diffusion 1.5--models-path: définir le dossier des modèles Stable Diffusion--ops-printf: afficher l’opération courante sur stdout pendant l’inférence--output: définir le fichier PNG de sortie--decode-latents: ignorer la diffusion et décoder le fichier de latents indiqué--prompt: définir le prompt positif--neg-prompt: définir le prompt négatif--steps: définir le nombre d’étapes de diffusion--save-latents: enregistrer les latents dans le fichier indiqué après la diffusion
- Des options liées au Raspberry Pi et au décodeur sont aussi fournies séparément
--decoder-calibrate: calibrer le décodeur VAE quantifié, uniquement pour SD 1.5--decoder-fp16: utiliser le décodeur VAE FP16, uniquement pour SD 1.5--not-tiled: ne pas utiliser le décodeur VAE tiled, uniquement pour SDXL 1.0--rpi: configurer le modèle pour une exécution sur Raspberry Pi--rpi-lowmem: appliquer la configuration basse mémoire pour Raspberry Pi Zero 2, uniquement avec SDXL 1.0
1 commentaires
Avis sur Hacker News
Intéressant. La phrase clé est celle-ci : « OnnxStream peut être seulement 0,5 à 2 fois plus lent qu’OnnxRuntime tout en utilisant jusqu’à 55 fois moins de mémoire »
Le compromis entre mémoire vidéo / utilisation mémoire et temps d’inférence semble pouvoir être avantageux non seulement dans les cas où la mémoire est limitée, comme sur Raspberry Pi, mais aussi dans d’autres situations.
Je me demande si, en pratique, cette méthode de déchargement des poids permettrait de traiter des tailles de batch plus grandes dans la même quantité de mémoire : la latence augmenterait, mais le débit ne pourrait-il pas beaucoup augmenter ?
Ils voulaient probablement écrire « 1,5 à 2 fois plus lent ».
Une fois qu’on a dépassé le simple stade de « ce modèle tient-il dans ce système ? », l’inférence est généralement limitée par la bande passante mémoire ; je ne pense donc pas que cette technique aide beaucoup à augmenter le débit via des batchs plus grands. Une seule instance sature probablement déjà le contrôleur mémoire.
En revanche, cela pourrait peut-être être utile côté entraînement.
11 heures, ça me rappelle l’époque où je faisais du ray tracing sur Amiga 500. Le rendu « final » était clairement un truc qu’on lançait pour la nuit.
Mais même si les caustiques paraissent lumineuses, statistiquement c’est un phénomène assez rare ; pour les obtenir correctement, il faut retirer les limites du moteur de rendu et le laisser tourner toute la nuit.
Le résultat, c’est une image d’une scène banale par un artiste médiocre, mais avec de jolies caustiques. Je vais garder mon vrai métier.
C’est à ce moment-là que je me suis dit qu’il me fallait un autre hobby. Juste avant, quelqu’un de très doué avait sorti un outil permettant de prévisualiser la scène avec OpenGL. Ça n’aurait pas marché sur Amiga, mais sur ma machine ça tournait de justesse.
J’utilise Stable Diffusion avec invoke.ai sur MBP ; avez-vous des recommandations pour mieux régler les paramètres SD ? Même avec le même prompt et des réglages en apparence identiques, par exemple le même modèle avec Euler A, je n’arrive pas du tout à atteindre la qualité des images que je vois en ligne.
Pour le SD 1.5 de base, j’utilise Volta parce qu’il est rapide : https://github.com/VoltaML/voltaML-fast-stable-diffusion/com...
Pour obtenir une vraiment bonne qualité d’image avec SD 1.5, il faut exploiter sans retenue les modèles fine-tunés, LoRA, ControlNet et autres fonctionnalités d’amélioration. Par exemple, faire suivre une image de base pour la structure, ou définir des prompts pour des zones spécifiques de l’image. InvokeAI a en réalité beaucoup de fonctionnalités, et nombre de ces améliorations sont cachées dans son interface à nœuds, mais d’autres interfaces comme Volta les exposent plus directement.
En général, les modèles fine-tunés impressionnants sont beaucoup moins polyvalents que les poids de base, mais en pratique ce n’est pas vraiment un problème et les résultats peuvent être nettement meilleurs.
Il demande plus de temps et de mémoire qu’Invoke ou qu’une carte graphique Nvidia, mais ce n’est pas si mauvais. Une image de qualité standard en 512x768 px prend environ 1 à 2 s/it, et une image haute qualité en 1024x1536 px avec Hires Fix prend autour de 14 à 20 s/it.
Ce serait vraiment génial de l’intégrer dans un cadre photo numérique ou un tableau mural.
https://github.com/rvdveen/epaper-slow-generative-art/
Un résultat impressionnant, mais évidemment, générer une image prend énormément de temps. Le README indique 11 heures.
Par exemple, comparer le temps, le coût, tout le matériel nécessaire et l’électricité pour produire 100 images avec 100 Pi Zero 2 — même pas forcément des W — à quelque chose comme un PC milieu de gamme moyen.
Je pense que le PC gagnerait probablement encore.
Le Zero 2 donne l’impression d’avoir été utilisé comme défi plus que pour son côté pratique ; un Pi 4 ou 5 serait peut-être un meilleur point de comparaison.
Impressionnant.
En vérité, l’époque où même les ampoules et les grille-pain abriteront une intelligence remarquable approche.
Je suis ce domaine depuis plusieurs années, mais les dix dernières ont été étonnantes.
Si je dis « ont été », c’est parce que l’accélération des 6 à 18 derniers mois est encore d’un tout autre niveau.
Ce qui m’inquiète, ce n’est pas ce qu’on pourra faire dans deux ans, mais à quel point le rythme des progrès va encore s’accélérer. Et puis encore, et encore.
Montons une startup.
À ce stade, les tentatives de réguler les technologies de type Stable Diffusion ne sont-elles pas terminées ? Si l’on réduit les modèles et l’infrastructure d’inférence au point qu’ils puissent tourner même sur une PS2, il semble impossible d’empêcher cette technologie sans un État de surveillance totalitaire, et même là, ce serait à peine possible.
Il en va de même pour le « piratage » de médias ou les ransomwares.
Les États réglementent depuis longtemps des choses qui ne peuvent pas être appliquées de manière purement technique.
Je sais que c’est un exemple volontairement exagéré pour s’amuser, mais pour le faire tourner sur une PS2, il faudrait réduire encore bien davantage.
J’ai hâte de voir Stable Diffusion pour Windows 3.1.
C’est dingue. Que ça prenne 11 heures ou non, je n’aurais jamais imaginé que SD puisse tourner sur du matériel comme un Pi Zero.