Stable Diffusion XL 1.0 chạy trong 298MB RAM
(github.com/vitoplantamura)- OnnxStream là một thư viện suy luận siêu gọn nhằm chạy Stable Diffusion 1.5 và SDXL 1.0 Base trên Raspberry Pi Zero 2 chỉ với 512MB RAM, không cần thêm swap hay ghi tạm kết quả ra đĩa
- Thư viện tách biệt engine suy luận và
WeightsProvider, hỗ trợDiskNoCache,DiskPrefetch, cùng cơ chế tải·cache·prefetch tùy biến; cũng có thể nhận trực tiếp weights qua HTTP mà không cần ghi ra đĩa - SDXL 1.0 Base dùng các tệp ONNX được xuất từ bản triển khai Hugging Face Diffusers 0.19.3, và mất khoảng 11 giờ để tạo ảnh 10 bước trên Raspberry Pi Zero 2
- Bộ giải mã VAE của SDXL tiêu tốn 4.4GB RAM ở FP32 và khó giải quyết chỉ bằng FP16·UINT8, nên nhóm tác giả đã hạ mức dùng bộ nhớ xuống 298MB bằng giải mã theo ô chồng lấp và hòa trộn
- So với OnnxRuntime trên CPU, nó dùng ít bộ nhớ hơn tối đa 55 lần nhưng có thể chậm hơn 0.5~2 lần, đồng thời bị giới hạn ở batch size 1 và tập trung vào một số toán tử ONNX nhất định
Mục tiêu và thiết kế của OnnxStream
- OnnxStream bắt đầu từ mục tiêu chạy Stable Diffusion 1.5 trên Raspberry Pi Zero 2
- Raspberry Pi Zero 2 là một máy tính siêu nhỏ với 512MB RAM
- Mức RAM/VRAM tối thiểu được khuyến nghị cho Stable Diffusion thường là 8GB
- Điều kiện mục tiêu là không thêm swap và không offload kết quả trung gian ra đĩa
- Các framework machine learning thông thường thường ưu tiên độ trễ suy luận hoặc thông lượng, nên mức dùng RAM dễ tăng cao
- OnnxStream là một thư viện suy luận nhỏ gọn, dễ chỉnh sửa, tập trung vào việc giảm tối đa mức dùng bộ nhớ
- Cấu trúc cốt lõi là tách biệt engine suy luận và thành phần cung cấp trọng số mô hình là
WeightsProvider- Các lớp dẫn xuất có thể triển khai cách nạp, cache và prefetch tham số mô hình
WeightsProvidertùy biến có thể tải trực tiếp dữ liệu từ máy chủ HTTP và xử lý mà không cần đọc·ghi đĩa- Các triển khai mặc định là
DiskNoCachevàDiskPrefetch
- Trên CPU, nó dùng ít bộ nhớ hơn tối đa 55 lần so với OnnxRuntime, đổi lại chậm hơn khoảng 0.5~2 lần
Chạy Stable Diffusion 1.5
- Ví dụ Stable Diffusion 1.5 dùng OnnxStream để tạo ảnh với các mức độ chính xác VAE decoder khác nhau
- Thành phần duy nhất không thể vừa trong RAM của Raspberry Pi Zero 2 là VAE decoder
- Nguyên nhân là residual connection, tensor lớn và convolution lớn trong mô hình
- Chỉ dùng đơn chính xác hoặc bán chính xác vẫn không thể vừa với RAM của Raspberry Pi Zero 2
- Giải pháp là lượng tử hóa tĩnh 8-bit
- Ảnh từ VAE decoder ở độ chính xác W8A8 được tạo trên Raspberry Pi Zero 2, mất khoảng 1.5 giờ khi dùng tùy chọn
MAX_SPEED - Ảnh W16A16 để so sánh được tạo trên PC bằng cùng một latent
Hỗ trợ Stable Diffusion XL 1.0 Base
- Ví dụ Stable Diffusion của OnnxStream hỗ trợ SDXL 1.0 Base và không bao gồm Refiner
- Các tệp ONNX được xuất từ bản triển khai SDXL 1.0 trong thư viện Hugging Face Diffusers, với phiên bản dùng là 0.19.3
- SDXL 1.0 có chi phí tính toán lớn hơn nhiều so với SD 1.5
- Khác biệt lớn nhất là có thể tạo ảnh 1024x1024 thay vì 512x512
- Trên PC 12 lõi với 32GB RAM, Hugging Face Diffusers mất 26 phút để tạo ảnh 10 bước
- Mức VRAM tối thiểu được khuyến nghị cho SDXL thường là 12GB
- OnnxStream có thể chạy SDXL 1.0 với dưới 300MB RAM, nên cũng hoạt động được trên Raspberry Pi Zero 2
- Không thêm swap
- Không ghi ra đĩa trong lúc suy luận
- Mất khoảng 11 giờ để tạo ảnh 10 bước trên Raspberry Pi Zero 2
Tối ưu bộ nhớ dành riêng cho SDXL
- SDXL 1.0 áp dụng cùng nhóm tối ưu như SD 1.5 nhưng có một số khác biệt
- Mô hình UNET dùng lượng tử hóa động UINT8 để chạy trong dưới 300MB RAM trên Raspberry Pi Zero 2
- Phạm vi lượng tử hóa được giới hạn ở một tập con nhất định của các tensor trung gian lớn
- VAE decoder của SDXL 1.0 khó xử lý hơn SD 1.5
- VAE decoder của SDXL 1.0 lớn gấp 4 lần so với VAE decoder của SD 1.5
- Khi chạy ở độ chính xác FP32 trên OnnxStream, nó dùng tới 4.4GB RAM
- Ở SD 1.5, VAE decoder có thể được lượng tử hóa tĩnh UINT8 để giảm RAM xuống 260MB
- VAE decoder của SDXL 1.0 bị overflow ở số học FP16, và dải giá trị activation quá lớn nên khó tạo ảnh chất lượng tốt bằng lượng tử hóa UINT8
- Có các cách khắc phục FP16 như sdxl-vae-fp16-fix, nhưng dù giảm bộ nhớ còn một nửa thì 2.2GB vẫn quá lớn với Raspberry Pi Zero 2
- Giải pháp cuối cùng là giải mã theo ô lấy cảm hứng từ triển khai VAE decoder của Diffusers
- Tensor kết quả diffusion có shape
(1,4,128,128) - Tensor này được chia thành 5x5, tổng 25 ô chồng lấp nhau với shape
(1,4,32,32) - Mỗi ô chồng lấp 25% với ô bên trái và ô phía trên
- Kết quả giải mã của mỗi ô là tensor
(1,3,256,256)và được hòa trộn vào ảnh cuối cùng - Nếu tắt hòa trộn thì sẽ thấy ranh giới giữa các ô, còn bật lên thì ranh giới biến mất trong kết quả cuối
- Tensor kết quả diffusion có shape
- Với cách này, mức dùng RAM của VAE decoder SDXL giảm từ 4.4GB xuống 298MB
Tính năng hỗ trợ và phụ thuộc
- OnnxStream gom gọn các tính năng cần thiết cho suy luận tiết kiệm bộ nhớ
- Tách biệt engine suy luận và
WeightsProvider DiskNoCache,DiskPrefetch,WeightsProvidertùy biến- attention slicing
- lượng tử hóa động 8-bit unsigned asymmetric percentile
- lượng tử hóa tĩnh W8A8 unsigned asymmetric percentile
- hiệu chỉnh mô hình lượng tử hóa
- hỗ trợ FP16
- triển khai 25 toán tử ONNX thường dùng
- Tách biệt engine suy luận và
- Các phép toán chạy tuần tự, nhưng mỗi operator hoạt động đa luồng
- Cách triển khai theo cấu trúc một tệp triển khai và một tệp header, với lớp
XnnPackbao bọc lời gọi XNNPACK - Một số primitive tăng tốc phụ thuộc vào XNNPACK
- MatMul
- Convolution
- element-wise Add/Sub/Mul/Div
- Sigmoid
- Softmax
So sánh hiệu năng và các giới hạn
- Stable Diffusion gồm ba mô hình
- text encoder: 672 operations, 123 triệu parameters
- UNET: 2050 operations, 854 triệu parameters
- VAE decoder: 276 operations, 49 triệu parameters
- Để tạo ảnh 10 bước với batch size 1 cần chạy
- text encoder 2 lần
- UNET 20 lần
- VAE decoder 1 lần
- Với UNET FP16, chênh lệch về bộ nhớ và thời gian giữa OnnxStream và OnnxRuntime là đáng kể
- OnnxStream: khoảng 0.133GB, 18.2~19.8 giây
- OnnxRuntime: 5.085~7.353GB, 7.28~12.8 giây
- OnnxStream dùng ít bộ nhớ hơn tối đa 55 lần nhưng chậm hơn 0.5~2 lần
- Với text encoder FP32, OnnxStream dùng 0.147GB còn OnnxRuntime dùng 0.641GB
- Với VAE decoder FP32, OnnxStream dùng 1.004GB còn OnnxRuntime dùng 1.330~2.026GB
- Kết quả so sánh đi kèm một số điều kiện
- Lần chạy đầu tiên của OnnxRuntime là suy luận warm-up trước khi tái sử dụng
InferenceSession - OnnxStream có thiết kế eager nên không có khái niệm warm-up, nhưng các lần chạy sau có thể hưởng lợi từ cache tệp weights của hệ điều hành
- OnnxStream hiện không hỗ trợ đầu vào khác batch size 1
- OnnxRuntime có thể dùng batch size 2 khi chạy UNET để tăng tốc đáng kể toàn bộ quá trình diffusion
- Trong thử nghiệm, việc thay đổi
EnableCpuMemArenavàExecutionModetrongSessionOptionscủa OnnxRuntime không tạo ra khác biệt đáng kể - NCNN có mức dùng bộ nhớ và thời gian suy luận rất giống OnnxRuntime
- Môi trường thử nghiệm là Windows Server 2019, RAM 16GB, CPU 8750H AVX2, SSD 970 EVO Plus, 8 lõi ảo VMWare
- Lần chạy đầu tiên của OnnxRuntime là suy luận warm-up trước khi tái sử dụng
Attention slicing và lượng tử hóa
- Khi chạy UNET, attention slicing và lượng tử hóa W8A8 của VAE decoder là những yếu tố quan trọng để hạ bộ nhớ xuống mức có thể chạy được trên Raspberry Pi Zero 2
- attention slicing giúp tránh phải materialize toàn bộ ma trận
Q @ K^Tkhi tính scaled dot-product attention trong multi-head attention - Khi mô hình UNET có 8 attention head, shape tensor sẽ như sau
Q:(8,4096,40)K^T:(8,40,4096)- Kết quả MatMul đầu tiên:
(8,4096,4096) - Ở độ chính xác FP32, đây là tensor 512MB
- Cách giải quyết là chia
Qtheo chiều dọc và thực hiện attention trên từng chunk- Shape của
Q_slicedlà(1,x,40) xlà giá trị 4096 chia choonnxstream::Model::m_attention_fused_ops_parts- Giá trị mặc định là 2 và có thể tùy chỉnh
- Shape của
- Cách này giúp hạ tổng mức dùng bộ nhớ của mô hình UNET FP32 từ 1.1GB xuống 300MB
- FlashAttention có thể là một lựa chọn hiệu quả hơn, nhưng sẽ cần viết custom kernel cho từng kiến trúc hỗ trợ như AVX·NEON, và trong trường hợp đó phải обход XNNPACK
Chuyển đổi mô hình và cách chạy
- OnnxStream chạy mô hình được định nghĩa trong
path_to_model_folder/model.txt- Mọi phép toán của mô hình đều nằm trong
model.txtở định dạng ASCII - Các tệp weights phải tồn tại dưới dạng chuỗi tệp
.bintrong cùng thư mục
- Mọi phép toán của mô hình đều nằm trong
- Có thể thiết lập nhiều tham số tùy chọn trên đối tượng
Model- chỉ định weights provider khác
- đọc·ghi tệp dải clipping activation của mô hình lượng tử hóa
- chế độ hiệu chỉnh mô hình
- dùng số học FP16
- dùng số học UINT8
- dùng lượng tử hóa động UINT8
- bật attention slicing
- đặt số lượng phân chia attention
- Để dùng tệp ONNX trong OnnxStream, hãy dùng notebook
onnx2txt.ipynbđể xuấtmodel.txtvà các tệp weights.bin - Khi xuất
PyTorch nn.Modulesang ONNX có một số điều kiện- Khi gọi
torch.onnx.export, phải để trốngdynamic_axes - OnnxStream không hỗ trợ đầu vào dynamic shape
- Rất khuyến nghị chạy ONNX Simplifier trước khi chuyển đổi
- Khi gọi
Chuẩn bị build và chạy
- Có thể build ví dụ Stable Diffusion trên Linux, Mac, Windows và Termux
- Trên Windows dùng
x64 Native Tools Command PromptcủaVisual Studio Tools - Trên Mac cần cài cmake bằng
brew install cmake
- Trên Windows dùng
- Trước tiên cần build XNNPACK
- Prototype hàm của XNNPACK có thể thay đổi bất kỳ lúc nào, nên quy trình có bước checkout về commit phù hợp với một thời điểm cụ thể
- Mốc tham chiếu là commit trên nhánh master trước
2023-06-27 00:00
- Sau đó clone kho OnnxStream và build bằng cmake trong
src/buildMAX_SPEED=ON- chỉ định đường dẫn clone XNNPACK bằng
XNNPACK_DIR
- Tùy chọn
MAX_SPEEDcó thể tăng hiệu năng nhưng cũng dùng nhiều bộ nhớ hơn khi build- Trên Windows tăng hiệu năng khoảng 10%
- Trên Raspberry Pi tăng hiệu năng hơn 50%
- File thực thi được tạo ra có thể không chạy được, và đã có vấn đề trong thử nghiệm trên Termux
- Nếu gặp sự cố, nên thử đổi sang
MAX_SPEED=OFFtrước
- Weights cho ví dụ Stable Diffusion 1.5 có thể tải từ mục Releases của kho, dung lượng khoảng 2GB
- Weights cho Stable Diffusion XL 1.0 Base có thể tải từ Hugging Face, dung lượng khoảng 8GB
Tùy chọn ví dụ Stable Diffusion
- File thực thi ví dụ điều khiển việc chọn mô hình, input/output, prompt và cách giải mã bằng các tùy chọn dòng lệnh
--xl: chạy Stable Diffusion XL 1.0 thay cho Stable Diffusion 1.5--models-path: chỉ định thư mục mô hình Stable Diffusion--ops-printf: in operation hiện tại ra stdout trong lúc suy luận--output: chỉ định tệp PNG đầu ra--decode-latents: bỏ qua diffusion và giải mã tệp latents được chỉ định--prompt: chỉ định positive prompt--neg-prompt: chỉ định negative prompt--steps: chỉ định số bước diffusion--save-latents: lưu latents sau diffusion vào tệp chỉ định
- Ngoài ra còn có các tùy chọn riêng cho Raspberry Pi và decoder
--decoder-calibrate: hiệu chỉnh VAE decoder đã lượng tử hóa, chỉ áp dụng cho SD 1.5--decoder-fp16: dùng VAE decoder FP16, chỉ áp dụng cho SD 1.5--not-tiled: không dùng tiled VAE decoder, chỉ áp dụng cho SDXL 1.0--rpi: cấu hình mô hình để chạy trên Raspberry Pi--rpi-lowmem: áp dụng cấu hình bộ nhớ thấp cho Raspberry Pi Zero 2, chỉ áp dụng cho SDXL 1.0
1 bình luận
Ý kiến trên Hacker News
Thú vị đấy. Câu cốt lõi là: “OnnxStream có thể dùng ít bộ nhớ hơn OnnxRuntime tới 55 lần mà tốc độ chỉ chậm hơn 0,5~2 lần”
Sự đánh đổi giữa bộ nhớ video/mức dùng bộ nhớ và thời gian suy luận có vẻ có lợi không chỉ trong những trường hợp bị giới hạn bộ nhớ như Raspberry Pi, mà cả ở các tình huống khác
Nếu cách dỡ tải trọng số này thật sự có thể xử lý batch size lớn hơn trong cùng lượng bộ nhớ, tôi tự hỏi liệu dù độ trễ tăng lên, throughput có thể tăng đáng kể hay không
Có lẽ họ định viết “chậm hơn 1,5~2 lần”
Suy luận thường bị ràng buộc bởi băng thông bộ nhớ một khi vượt qua mức “mô hình này có nhét vừa vào hệ thống này không”, nên kỹ thuật này có lẽ không giúp nhiều cho việc tăng throughput bằng batch size lớn hơn. Chỉ một instance thôi cũng có khả năng đã làm bão hòa memory controller rồi
Tuy vậy, có thể nó sẽ hữu ích ở phía huấn luyện
11 tiếng làm tôi nhớ thời ngày xưa ray tracing trên Amiga 500. Bản render “cuối cùng” đúng là việc để chạy qua đêm
Nhưng dù caustics trông sáng, về mặt thống kê chúng là hiện tượng khá hiếm, nên muốn có được đúng thì phải gỡ giới hạn của render engine và cứ để chạy qua đêm
Kết quả là một ảnh cảnh bình thường của một tác giả thiếu kỹ năng, nhưng có caustics đẹp. Chắc tôi vẫn nên tiếp tục công việc chính
Tầm đó tôi nghĩ mình cần một sở thích khác. Ngay trước đó, một người rất giỏi đã tung ra công cụ cho phép xem trước cảnh bằng OpenGL. Có lẽ nó không chạy được trên Amiga, nhưng trên máy của tôi thì chạy được một cách chật vật
Tôi đang dùng Stable Diffusion bằng invoke.ai trên MBP, có gợi ý nào để chỉnh tham số SD tốt hơn không? Dù dùng cùng prompt và các thiết lập nhìn bề ngoài giống nhau, chẳng hạn cùng model kiểu Euler A, tôi vẫn không thể đạt được chất lượng ảnh như thấy trên internet
Với SD 1.5 cơ bản thì tôi dùng Volta vì nó nhanh: https://github.com/VoltaML/voltaML-fast-stable-diffusion/com...
Chất lượng ảnh SD 1.5 thật sự tốt chỉ có được khi dùng thoải mái các mô hình fine-tune, LoRA, ControlNet và các tính năng tăng cường khác. Ví dụ như bắt ảnh cơ sở làm cấu trúc để bám theo, hoặc chỉ định prompt theo từng vùng cụ thể của ảnh. InvokeAI thực ra cũng có nhiều tính năng, và nhiều tính năng tăng cường kiểu này ẩn trong node UI, nhưng các UI khác như Volta thì phơi bày trực tiếp hơn
Nhìn chung, các mô hình fine-tune ấn tượng kém đa dụng hơn nhiều so với trọng số cơ bản, nhưng trong sử dụng thực tế thì đó không phải vấn đề lớn, và kết quả có thể tốt hơn đáng kể
Nó cần nhiều thời gian và bộ nhớ hơn Invoke hoặc card đồ họa Nvidia, nhưng cũng không quá tệ. Ảnh chất lượng tiêu chuẩn 512x768px khoảng 1~2 s/it, ảnh chất lượng cao 1024x1536px dùng Hires Fix khoảng 14~20 s/it
Nếu nhúng cái này vào khung ảnh kỹ thuật số hoặc tranh treo tường để chạy thì chắc sẽ rất tuyệt
https://github.com/rvdveen/epaper-slow-generative-art/
Đây là một thành tựu đáng kinh ngạc, nhưng tất nhiên việc tạo ảnh sẽ mất cực kỳ lâu. README ghi là 11 giờ
Ví dụ như so sánh thời gian, chi phí, toàn bộ phần cứng cần thiết và điện năng để tạo 100 ảnh bằng 100 chiếc Pi Zero 2—thậm chí không nhất thiết phải là bản W—với một PC tầm trung trung bình chẳng hạn
Có lẽ PC vẫn sẽ thắng
Zero 2 có cảm giác được dùng như một thử thách hơn là vì tính thực dụng, nên Pi 4 hoặc 5 có thể là mốc so sánh tốt hơn
Ấn tượng thật
Quả thật, thời đại mà ngay cả bóng đèn và lò nướng bánh mì cũng chứa trí tuệ vượt trội đang đến gần
Tôi đã theo dõi lĩnh vực này nhiều năm, nhưng 10 năm qua thật đáng kinh ngạc
Lý do tôi nói “đã từng” là vì sự tăng tốc trong 6–18 tháng gần đây lại ở một đẳng cấp hoàn toàn khác
Tôi không lo 2 năm nữa ta sẽ làm được gì, mà lo tốc độ phát triển sẽ còn nhanh hơn bao nhiêu. Và rồi nó sẽ lại nhanh hơn nữa, rồi lại nhanh hơn nữa
Lập startup thôi
Đến mức này thì các nỗ lực quản lý những công nghệ kiểu Stable Diffusion coi như chấm hết rồi không phải sao? Nếu có thể thu nhỏ mô hình và hạ tầng suy luận xuống mức chạy được cả trên PS2, thì có vẻ như không thể ngăn công nghệ này nếu không có một nhà nước giám sát toàn trị, mà ngay cả trong một nhà nước như vậy cũng chưa chắc làm được
“Vi phạm bản quyền” truyền thông hay ransomware cũng tương tự
Nhà nước từ lâu đã quản lý cả những thứ không thể thực thi thuần túy bằng kỹ thuật
Tôi biết đây là một ví dụ phóng đại cho vui, nhưng để chạy trên PS2 thì còn phải thu nhỏ hơn thế này rất nhiều
Không thể chờ đến ngày có Stable Diffusion cho Windows 3.1
Điên thật. Dù mất 11 giờ hay không, tôi hoàn toàn không ngờ SD có thể chạy được trên phần cứng như Pi Zero