DOCSEARCHbyAlgolia
Submit Docs

OnnxStream Documentation

repository·master·Indexed 24 days ago

https://github.com/vitoplantamura/onnxstream

A memory-efficient inference library designed to run large models like Stable Diffusion, SDXL, and LLMs on resource-constrained hardware such as the Raspberry Pi Zero. It features a WeightsProvider abstraction to decouple the inference engine from weight loading, supporting techniques like weight streaming, tiled VAE decoding, attention slicing, and quantization. OnnxStream also provides a WebAssembly (WASM) implementation for running models like OpenAI's Whisper directly in the browser.

Tokens
9.8K
Snippets
10
Records
41
Agent score
85%

What's inside OnnxStream

  1. What is OnnxStream and how does it work?

    master

    OnnxStream is a lightweight inference library designed to minimize memory consumption, enabling large machine learning models (like Stable Diffusion) to run on hardware with extremely limited RAM, such as a Raspberry Pi Zero 2 (512MB RAM).

    Unlike standard frameworks that prioritize latency or throughput at the cost of high RAM usage, OnnxStream achieves low memory footprints by decoupling the inference engine from the model weight loading mechanism. This is achieved through a WeightsProvider abstraction.

    By implementing a custom WeightsProvider, you can control how model parameters are loaded, cached, or prefetched. This allows for techniques like streaming weights directly from an HTTP server without ever writing them to disk.

  2. Use WeightsProvider to manage model weights

    master

    The core architecture of OnnxStream relies on the WeightsProvider class. To customize how weights are handled during inference, you can derive a new class from WeightsProvider.

    OnnxStream provides three default implementations:

    • DiskNoCache: Loads weights from disk without caching.
    • DiskPrefetch: Loads weights from disk with prefetching capabilities.
    • Ram: Loads weights directly into RAM.

    Developers can implement custom providers to support specialized loading strategies, such as streaming data from a remote server to avoid disk I/O.

  3. Use attention slicing to reduce memory consumption

    master

    Attention slicing is a technique used to avoid materializing the full Q @ K^T matrix during scaled dot-product attention in the UNET model. This is particularly useful for running large models on memory-constrained hardware like a Raspberry Pi Zero 2.

    By splitting the Q matrix into chunks (defined by model.m_attention_fused_ops_parts), you can lower the memory consumption of the UNET model significantly (e.g., from 1.1GB to 300MB in FP32 precision).

    To enable this, set: model.m_fuse_ops_in_attention = true;

  4. Build the Stable Diffusion example

    master

    To build the Stable Diffusion example, follow these platform-specific preparation steps and build commands.

    Prerequisites

    • Linux (+Termux): Install build-essential git cmake python3.
    • Windows: Use the Visual Studio Tools > x64 Native Tools Command Prompt.
    • Mac: Install cmake via Homebrew: brew install cmake.
    • FreeBSD: Requires manual modifications to XNNPACK CMake files (specifically CMAKE_SYSTEM_NAME, CMAKE_SYSTEM_PROCESSOR, and cpuinfo dependency) as XNNPACK does not support FreeBSD out-of-the-box.

    Build Steps

    git clone https://github.com/vitoplantamura/OnnxStream.git
cd OnnxStream/src
mkdir build
cd build
cmake ..
cmake --build . --config Release

    Performance Tuning

    The MAX_SPEED option is enabled by default. It can increase performance by ~10% on Windows and >50% on Raspberry Pi, but it consumes significantly more build-time memory and may cause execution issues (e.g., on Termux). If you encounter problems, rebuild with:

    cmake .. -DMAX_SPEED=OFF

  5. Build OnnxStream for LLM support

    master

    To run Large Language Models (LLMs) like TinyLlama or Mistral 7B, you must build OnnxStream from source.

    For Windows users: You can download a pre-built EXE from the Releases page. The EXE will automatically download model parameters from Hugging Face on its first run.

    For other platforms (Linux, macOS, Termux): Build the application using cmake with the following flags:

    • -DOS_LLM=ON: Enables LLM support.
    • -DOS_CUDA=ON (Optional): Enables GPU acceleration via cuBLAS for Nvidia cards.

  6. How to export a model to ONNX for OnnxStream

    master

    To use a new model with OnnxStream, follow these steps:

    1. Export to ONNX: Export the original PyTorch model to the ONNX format.
    2. Simplify (Optional): Run ONNX Simplifier on the exported file to optimize it.
    3. Convert to Text: Run onnx2txt to convert the ONNX file into the specific format required by OnnxStream.

    Note: The current LLM implementation supports FP16 and FP32 precision. 8-bit quantization for LLMs is planned for a future release.

  7. Convert a custom Stable Diffusion 1.5 model

    master

    To use a custom model, you must convert it to a format compatible with OnnxStream. It is highly recommended to use the Hugging Face implementation rather than AUTO1111 to avoid unsupported Einsum operations.

    Step 1: Export from Hugging Face (Recommended)

    Use the diffusers library to export the UNET to ONNX. This method requires significant swap space (approx. 100GB).

    from diffusers import StableDiffusionPipeline
import torch

pipe = StableDiffusionPipeline.from_single_file("https://huggingface.co/YourUsername/YourModel/blob/main/Model.safetensors")

dummy_input = (torch.randn(1, 4, 64, 64), torch.randn(1), torch.randn(1, 77, 768))
input_names = ["sample", "timestep", "encoder_hidden_states"]
output_names = ["out_sample"]

torch.onnx.export(pipe.unet, dummy_input, "/path/to/save/unet_temp.onnx", verbose=False, input_names=input_names, output_names=output_names, opset_version=14, do_constant_folding=True, export_params=True)

    Step 2: Simplify the ONNX model

    Run the ONNX simplifier to clean up the graph:

    python -m onnx_simplifier model_fixed3.onnx model_simplified.onnx

    Note: If you encounter issues with large models, consider using onnxsim_large_model.

    Step 3: Run the model

    Move the final model from onnx2txt into the unet_fp16 folder of your standard SD 1.5 model directory, then run:

    ./sd --models-path ./Converted/ --prompt "space landscape" --steps 28 --rpi

  8. Run OpenAI's Whisper in the browser via OnnxStream (WASM)

    master

    You can run OpenAI's Whisper models directly in a web browser using OnnxStream's WebAssembly (WASM) implementation. This allows for client-side speech-to-text processing without a backend server.

  9. Run Stable Diffusion XL (SDXL) on low-memory devices

    master

    OnnxStream supports Stable Diffusion XL 1.0 (base) and SDXL Turbo 1.0. While SDXL typically requires ~12GB of VRAM, OnnxStream can run it using less than 300MB of RAM on devices like the Raspberry Pi Zero 2.

    SDXL Optimizations

    • UNET: Uses UINT8 dynamic quantization on a specific subset of large intermediate tensors to keep memory usage low.
    • VAE Decoder: To prevent the 4.4GB RAM requirement of the SDXL VAE decoder, OnnxStream uses tiled decoding. This process splits the diffusion result tensor into overlapping tiles (e.g., 5x5 grid of 32x32 tiles), decodes them separately, and blends them back together. This reduces VAE memory consumption from 4.4GB to approximately 298MB.

  10. How to run a model with OnnxStream

    master

    To run an inference in OnnxStream, you must provide a model.txt file that defines the model operations in ASCII format. All associated weights must be stored as .bin files in the same directory as the model.txt.

    1. Prepare the model: Use the provided onnx2txt.ipynb notebook to export an ONNX file into a model.txt and a series of .bin weight files.
    2. Initialize the Model: Create an onnxstream::Model object.
    3. Configure (Optional): Set parameters like arithmetic precision or attention slicing on the Model object.
    4. Load the model: Call model.read_file("path/to/model.txt").
    5. Prepare Input: Create a Tensor object, set its name and shape, and populate it using a tensor_vector<float>.
    6. Execute: Push the tensor to the model using model.push_tensor() and call model.run().
    7. Retrieve Results: Access the output via model.m_data[0].get_vector<float>().
    #include "onnxstream.h"

using namespace onnxstream;

int main()
{
    Model model;

    // Optional parameters
    // model.m_use_fp16_arithmetic = true;
    // model.m_fuse_ops_in_attention = true;

    model.read_file("path_to_model_folder/model.txt");

    tensor_vector<float> data;
    // ... fill the tensor_vector with the tensor data.

    Tensor t;
    t.m_name = "input";
    t.m_shape = { 1, 4, 64, 64 };
    t.set_vector(std::move(data));
    model.push_tensor(std::move(t));

    model.run();
    
    auto& result = model.m_data[0].get_vector<float>();
    
    // ... process the result

    return 0;
}

  11. Exporting ONNX models for OnnxStream

    master

    When exporting a PyTorch nn.Module to ONNX for use with OnnxStream, follow these requirements:

    1. No Dynamic Axes: When calling torch.onnx.export, leave dynamic_axes empty. OnnxStream does not support inputs with dynamic shapes.
    2. Simplify the Model: It is strongly recommended to run ONNX Simplifier on the exported ONNX file before converting it to the model.txt format.
    3. Conversion: Use the onnx2txt.ipynb notebook to convert the simplified ONNX file into the required model.txt and .bin weight files.

  12. Run Stable Diffusion XL Turbo 1.0

    master

    SDXL Turbo 1.0 is supported and optimized for high-speed generation. It generates 512x512 images (instead of 1024x1024) and requires significantly fewer steps.

    Key characteristics:

    • High quality can be achieved in as little as 1 step.
    • Uses the same text encoder and VAE decoder as SDXL 1.0, meaning it also benefits from tiled decoding to maintain a memory footprint under 300MB.