onnx

README

Shared ONNX Runtime Package

This package provides a shared ONNX Runtime initialization and session management layer for ZimaOS-Blue.

Purpose

Centralizes ONNX Runtime initialization to avoid duplicate initialization across multiple modules (pruner, TTS, etc.).

Features

• Thread-safe initialization: InitializeRuntime() uses sync.Once to ensure ONNX Runtime is initialized only once
• Two session types:
  • Session: For models with fixed input/output shapes (pre-allocated tensors)
  • DynamicSession: For models with variable input shapes (dynamic tensors)

Usage

Fixed-shape models (like pruner)

import "github.com/IceWhaleTech/ZimaOS-Blue/server/internal/onnx"

// Create pre-allocated tensors
inputTensor, _ := ort.NewEmptyTensor[int64](ort.NewShape(1, 4096))
outputTensor, _ := ort.NewEmptyTensor[float32](ort.NewShape(1, 4096))

// Create session with pre-allocated tensors
session, err := onnx.NewSession(
    modelPath,
    []string{"input"},
    []string{"output"},
    []ort.ArbitraryTensor{inputTensor},
    []ort.ArbitraryTensor{outputTensor},
)
defer session.Close()

// Run inference (no arguments needed)
err = session.Run()

Variable-shape models (like TTS)

import "github.com/IceWhaleTech/ZimaOS-Blue/server/internal/onnx"

// Create session for dynamic tensors
session, err := onnx.NewDynamicSession(
    modelPath,
    []string{"phonemes", "speaker_id", "speed", "pitch"},
    []string{"audio"},
)
defer session.Close()

// Create input tensors dynamically
phonemesTensor, _ := ort.NewTensor(ort.NewShape(1, len(phonemes)), phonemes)
speakerTensor, _ := ort.NewTensor(ort.NewShape(1), []int64{speakerID})
outputTensor, _ := ort.NewEmptyTensor[float32](ort.NewShape(1, 0))

// Run inference with dynamic tensors
err = session.Run(
    []ort.Value{phonemesTensor, speakerTensor, ...},
    []ort.Value{outputTensor},
)

Modules Using This Package

• internal/pruner: Neural code pruning (fixed-shape)
• internal/tts: Supertonic TTS provider (variable-shape)

Implementation Notes

• Session wraps ort.AdvancedSession for pre-allocated tensor workflows
• DynamicSession wraps ort.DynamicAdvancedSession for dynamic tensor workflows
• Both provide a unified Close() method for resource cleanup
• ONNX Runtime initialization is idempotent and thread-safe

Documentation

Index

• func EnsureRuntime(dataDir string) (string, error)
• func ExtractRuntimeFromTgz(tgzPath, dataDir string) (string, error)
• func InitializeRuntime() error
• func RuntimeLibPath(dataDir string) string
• func RuntimeTgzFilename() string
• func RuntimeTgzMirrors() []string
• func RuntimeTgzURL() string
• func SetDataDir(dataDir string)
• func SetLibraryPath(path string)
• type DynamicSession
  • func NewDynamicSession(modelPath string, inputNames, outputNames []string) (*DynamicSession, error)
  • func (s *DynamicSession) Close() error
• type Session
  • func NewSession(modelPath string, inputNames, outputNames []string, ...) (*Session, error)
  • func (s *Session) Close() error

Functions

func EnsureRuntime

func EnsureRuntime(dataDir string) (string, error)

EnsureRuntime downloads the ONNX Runtime shared library if not present. Returns the path to the library file.

func ExtractRuntimeFromTgz

func ExtractRuntimeFromTgz(tgzPath, dataDir string) (string, error)

ExtractRuntimeFromTgz extracts the ONNX Runtime dylib from a downloaded tgz file into the onnxruntime subdirectory under dataDir. Returns the dylib path.

func InitializeRuntime

func InitializeRuntime() error

InitializeRuntime initializes the ONNX Runtime environment. Unlike sync.Once, this retries on failure so that downloading the shared library after a failed attempt can succeed.

func RuntimeLibPath

func RuntimeLibPath(dataDir string) string

RuntimeLibPath returns the path to the ONNX Runtime shared library in the given data directory, or empty string if not present.

func RuntimeTgzFilename

func RuntimeTgzFilename() string

RuntimeTgzFilename returns the platform-specific tgz filename.

func RuntimeTgzMirrors

func RuntimeTgzMirrors() []string

RuntimeTgzMirrors returns CDN mirror URLs for the ONNX Runtime tgz.

func RuntimeTgzURL

func RuntimeTgzURL() string

RuntimeTgzURL returns the primary GitHub URL for the ONNX Runtime tgz.

func SetDataDir

func SetDataDir(dataDir string)

SetDataDir provides a hint for where to find the ONNX Runtime library. Must be called before the first session creation.

func SetLibraryPath

func SetLibraryPath(path string)

SetLibraryPath sets the ONNX Runtime shared library path. Must be called before any session creation.

Types

type DynamicSession

type DynamicSession struct {
	*ort.DynamicAdvancedSession
}

DynamicSession wraps ort.DynamicAdvancedSession with common lifecycle management. Use this for models with variable input shapes (like TTS).

func NewDynamicSession

func NewDynamicSession(modelPath string, inputNames, outputNames []string) (*DynamicSession, error)

NewDynamicSession creates a new ONNX session for dynamic tensor inputs. Use this for models with variable input shapes (like TTS).

func (*DynamicSession) Close

func (s *DynamicSession) Close() error

Close destroys the session and releases resources.

type Session

type Session struct {
	*ort.AdvancedSession
}

Session wraps ort.AdvancedSession with common lifecycle management. Use this for models with fixed input/output shapes (like pruner).

func NewSession

func NewSession(modelPath string, inputNames, outputNames []string, inputs, outputs []ort.ArbitraryTensor) (*Session, error)

NewSession creates a new ONNX session with pre-allocated tensors. Use this for models with fixed input/output shapes (like pruner).

func (*Session) Close

func (s *Session) Close() error

Close destroys the session and releases resources.