84 lines
3.0 KiB
JavaScript
84 lines
3.0 KiB
JavaScript
|
|
// This script is run within the webview itself
|
|
(function () {
|
|
const vscode = acquireVsCodeApi();
|
|
|
|
//the onnx proto model
|
|
let onnxProtoModel = null;
|
|
//the onnx proto graphs
|
|
let onnxProtoGraphs = null;
|
|
//the onnx model name
|
|
let onnxModelName = '';
|
|
//the tensor external data map. key: location, value: Uint8Array
|
|
let externalTensorMap = new Map();
|
|
//the tensor external locations
|
|
let externalTensorLocations = [];
|
|
|
|
// Handle messages from the extension
|
|
window.addEventListener('message', e => {
|
|
const { type, body } = e.data;
|
|
switch (type) {
|
|
case 'init':
|
|
{
|
|
//the onnx file content, Uint8Array
|
|
let onnxValue = body.modelData;
|
|
//the onnx file name
|
|
onnxModelName = body.modelName;
|
|
|
|
onnxProtoModel = window.proto.onnx.ModelProto.decode(onnxValue);
|
|
let { locations, graphs } = window.getExternalLocations(onnxProtoModel);
|
|
|
|
//refer to https://github.com/onnx/onnx/blob/main/docs/ExternalData.md
|
|
if (locations.length > 0) {
|
|
onnxProtoGraphs = graphs;
|
|
externalTensorMap.clear();
|
|
externalTensorLocations = locations;
|
|
//load tensor external data
|
|
vscode.postMessage({ type: 'external_tensor', external_locations: locations });
|
|
} else {
|
|
let model = new window.onnx.Model(onnxProtoModel, graphs, new Map());
|
|
|
|
let host = new window.host.BrowserHost(vscode);
|
|
let view = new window.View(host);
|
|
view.start().then((data) => {
|
|
view.open(model, onnxModelName);
|
|
}).catch((err) => {
|
|
console.log(err);
|
|
})
|
|
}
|
|
|
|
break;
|
|
}
|
|
|
|
case 'external_tensor': {
|
|
//the external tensor data, Uint8Array
|
|
let tensorData = body.tensorData;
|
|
let loc = body.location;
|
|
|
|
externalTensorMap.set(loc, tensorData);
|
|
if (externalTensorMap.size == externalTensorLocations.length) {
|
|
let model = new window.onnx.Model(onnxProtoModel, onnxProtoGraphs, externalTensorMap);
|
|
|
|
let host = new window.host.BrowserHost(vscode);
|
|
let view = new window.View(host);
|
|
view.start().then((data) => {
|
|
view.open(model, onnxModelName);
|
|
}).catch((err) => {
|
|
console.log(err);
|
|
})
|
|
}
|
|
|
|
break;
|
|
}
|
|
|
|
default:
|
|
{
|
|
break;
|
|
}
|
|
}
|
|
});
|
|
|
|
// Signal to VS Code that the webview is initialized.
|
|
vscode.postMessage({ type: 'ready' });
|
|
}());
|