beenull/ente
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Initialize ORT only once

Browse source
This commit is contained in:
vishnukvmd 2024-01-04 11:39:06 +05:30
parent 5413646a28
commit 83200fbb24
3 changed files with 10 additions and 10 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

12
lib/services/semantic_search/frameworks/onnx/onnx.dart
View file

@ -32,10 +32,15 @@ class ONNX extends MLFramework {
return kModelBucketEndpoint + kTextModel;
}
@override
Future<void> init() async {
await _computer.compute(initOrtEnv);
await super.init();
}
@override
Future<void> loadImageModel(String path) async {
final startTime = DateTime.now();
await _computer.compute(_clipImage.init);
_imageEncoderAddress = await _computer.compute(
_clipImage.loadModel,
param: {
@ -51,7 +56,6 @@ class ONNX extends MLFramework {
@override
Future<void> loadTextModel(String path) async {
final startTime = DateTime.now();
await _computer.compute(_clipText.init);
// Doing this from main isolate since `rootBundle` cannot be accessed outside it
await _clipText.initTokenizer();
_textEncoderAddress = await _computer.compute(
@ -120,3 +124,7 @@ class ONNX extends MLFramework {
_logger.info('Released');
}
}
void initOrtEnv() async {
OrtEnv.instance.init();
}

4
lib/services/semantic_search/frameworks/onnx/onnx_image_encoder.dart
View file

@ -9,10 +9,6 @@ import "package:onnxruntime/onnxruntime.dart";
class OnnxImageEncoder {
final _logger = Logger("OnnxImageEncoder");
Future<void> init() async {
OrtEnv.instance.init();
}
Future<int> loadModel(Map args) async {
final sessionOptions = OrtSessionOptions()
..setInterOpNumThreads(1)

4
lib/services/semantic_search/frameworks/onnx/onnx_text_encoder.dart
View file

@ -12,10 +12,6 @@ class OnnxTextEncoder {
final _logger = Logger("OnnxTextEncoder");
final OnnxTextTokenizer _tokenizer = OnnxTextTokenizer();
Future<void> init() async {
OrtEnv.instance.init();
}
// Do not run in an isolate since rootBundle can only be accessed in the main isolate
Future<void> initTokenizer() async {
final vocab = await rootBundle.loadString(kVocabFilePath);