Adding Custom Visualizers

Relevant source files

The following files were used as context for generating this wiki page:

• src/components/OnnxGraphVisualizer.tsx
• src/WorkbenchContext.tsx
• src/utils/modelInfo.ts
• src/types.ts
• src/components/CenterCanvas.tsx

The ONNX-Station workbench utilizes a modular visualization system to provide insights into exported model graphs, specifically focusing on operator distribution and VRAM allocation. Custom visualizers can be integrated into the CenterCanvas to extend the analysis of compiled ONNX models.

This system relies on a centralized data provider, the WorkbenchContext, which supplies compiledModelInfo containing structured metadata about the model's architecture, operator counts, and estimated memory usage.

Visualization Architecture

The visualization framework is built on a provider-consumer pattern. The WorkbenchContext acts as the single source of truth, calculating model metadata through the generateCompiledModelInfo utility whenever a compilation (export) event completes.

The diagram above illustrates how configuration and export states flow from the context into the visualization components.
Sources: src/WorkbenchContext.tsx:189-204, src/components/OnnxGraphVisualizer.tsx:11-13

Key Data Structures

Custom visualizers must consume the CompiledModelInfo interface to remain compatible with the existing pipeline.

Field	Type	Description
operators	CompiledModelOperator[]	List of operators with counts, memory, and categories.
totalOps	number	The aggregate count of all nodes in the graph.
totalMemoryMb	number	Estimated total VRAM footprint in Megabytes.
isSimulated	boolean	Flag indicating if the data is estimated via the workbench logic.

Sources: src/types.ts:84-95

Integrating with the D3 Pipeline

The current visualizers use D3.js for rendering. To add a new visualizer, developers should follow the pattern established in OnnxGraphVisualizer.tsx, which uses React's useEffect and useRef hooks to bridge the React lifecycle with D3's imperative DOM manipulation.

Rendering Cycle

1. Ref Hooks: Use useRef<SVGSVGElement | null>(null) to capture the SVG container.
2. Responsiveness: Implement a ResizeObserver to track the container's width and update the local state (size), triggering a D3 re-render.
3. Data-Driven Updates: Use a useEffect hook that depends on compiledModelInfo and size to clear the previous SVG content and redraw the visualization.

The sequence diagram shows the internal logic used to prevent duplicate renders and ensure smooth animations.
Sources: src/components/OnnxGraphVisualizer.tsx:24-42, src/components/OnnxGraphVisualizer.tsx:57-80

Data Simulation Logic

Visualizers are powered by the generateCompiledModelInfo utility. This function simulates the ONNX graph structure based on the ExportConfig, such as modality, model path, and optimization levels (O1 through O4).

Operator Categories

Operators are classified into five categories to allow for consistent coloring across different visualizers:

• compute: Heavy processing units (e.g., MatMul, Conv).
• memory: Data movement or retrieval (e.g., Gather, QuantizeLinear).
• activation: Non-linear layers (e.g., Softmax, Relu).
• normalization: Normalization layers (e.g., LayerNormalization).
• shape: Structural manipulations (e.g., Reshape, Transpose).

Sources: src/utils/modelInfo.ts:63-74, src/components/OnnxGraphVisualizer.tsx:47-53

Optimization Effects

The simulation logic adjusts operator counts and memory estimates based on the selected optimize flag. For instance, O3 (Complete Attention Fusion) reduces MatMul counts while adding a MultiHeadAttention (Fused) entry.

    if (config.optimize === 'O3') {
      addOp('MultiHeadAttention (Fused)', numLayers, numLayers * 6.0 * dtypeScale * optMemMultiplier, 'compute');
      addOp('Softmax (Remainder)', Math.round(numLayers * 0.2), numLayers * 0.4 * dtypeScale, 'activation');
    } else {
      addOp('Softmax', numLayers, numLayers * 1.8 * dtypeScale, 'activation');
    }

Sources: src/utils/modelInfo.ts:107-112

Placement in the Workbench

New visualizers should be placed within the CenterCanvas.tsx component. The OnnxGraphVisualizer is currently positioned within the left column (col-span-7) alongside the CLI command preview and export specifications.

Summary

Adding a custom visualizer involves creating a component that consumes compiledModelInfo from the useWorkbench hook and implements a rendering strategy (typically D3.js). By adhering to the CompiledModelOperator interface and using the established operator categories, new visualizers can maintain visual and technical consistency with the ONNX-Station VRAM and operator distribution mapping tools.