vitoplantamura/OnnxStream

解決的問題

OnnxStream 是一個輕量級的推論庫，專為在 RAM 極度有限的硬體（例如 512MB 記憶體的 Raspberry Pi Zero 2）上執行大型 AI 模型（如 Stable Diffusion 和 LLM）而設計。標準機器學習框架通常優先考量速度與吞吐量，而非記憶體效率，往往需要數 GB 的 RAM，而低階裝置無法滿足此需求。

工作原理

該庫透過 WeightsProvider 將推論引擎與權重載入流程解耦，從而最小化記憶體消耗。這使得權重可從磁碟甚至 HTTP 伺服器流式載入，而無需全部載入至 RAM。它採用多種記憶體節省技術：

• 注意力切片（Attention Slicing）：透過將查詢張量分割成較小的區塊，防止在多頭注意力過程中產生巨大的中間張量。
• 量化（Quantization）：支援動態（8 位無符號）與靜態（W8A8）量化，以降低權重與激活的精確度與大小。
• 分塊解碼（Tiled Decoding）：針對大型 VAE 解碼器（如 SDXL），將輸入張量分割為重疊的區塊，分別解碼後再融合回原圖。
• XNNPACK 集成：使用 XNNPACK 加速 MatMul 和 Convolution 等基本運算。

適用對象

• 面向超低資源嵌入式裝置（如 Raspberry Pi Zero）的開發者。
• 希望在無專用 GPU 或高記憶體硬體上執行大型生成模型的使用者。
• 尋找可修改、極簡 C++ 推論引擎並具備 Python 與 C# 繫結的工程師。

主要亮點

• 極致記憶體效率：對於某些模型，記憶體消耗可比 OnnxRuntime 減少最多 55 倍。
• 廣泛模型支援：可執行 Stable Diffusion 1.5、SDXL 1.0、SDXL Turbo、TinyLlama、Mistral 7B、YOLOv8 以及 OpenAI 的 Whisper。
• 跨平台支援：支援 Linux、Mac、Windows、Termux、FreeBSD 與 WebAssembly（瀏覽器中執行）。
• 靈活的權重載入：權重可透過 RAM、磁碟（支援或不支援預取）或自訂串流提供者載入。