diff --git a/js/react_native/android/src/androidTest/java/ai/onnxruntime/reactnative/OnnxruntimeModuleTest.java b/js/react_native/android/src/androidTest/java/ai/onnxruntime/reactnative/OnnxruntimeModuleTest.java index 448ece3963..ea5a7d2745 100644 --- a/js/react_native/android/src/androidTest/java/ai/onnxruntime/reactnative/OnnxruntimeModuleTest.java +++ b/js/react_native/android/src/androidTest/java/ai/onnxruntime/reactnative/OnnxruntimeModuleTest.java @@ -15,6 +15,7 @@ import com.facebook.react.bridge.JavaOnlyMap; import com.facebook.react.bridge.ReactApplicationContext; import com.facebook.react.bridge.ReadableArray; import com.facebook.react.bridge.ReadableMap; +import java.io.ByteArrayOutputStream; import java.io.InputStream; import java.nio.ByteBuffer; import java.nio.ByteOrder; @@ -46,18 +47,31 @@ public class OnnxruntimeModuleTest { when(Arguments.createArray()).thenAnswer(i -> new JavaOnlyArray()); OnnxruntimeModule ortModule = new OnnxruntimeModule(reactContext); + String sessionKey = ""; // test loadModel() { InputStream modelStream = reactContext.getResources().openRawResource(ai.onnxruntime.reactnative.test.R.raw.test_types_float); + ByteArrayOutputStream byteBuffer = new ByteArrayOutputStream(); + + int bufferSize = 1024; + byte[] buffer = new byte[bufferSize]; + + int len; + while ((len = modelStream.read(buffer)) != -1) { + byteBuffer.write(buffer, 0, len); + } + + byte[] modelBuffer = byteBuffer.toByteArray(); + JavaOnlyMap options = new JavaOnlyMap(); try { - ReadableMap resultMap = ortModule.loadModel("test", modelStream, options); + ReadableMap resultMap = ortModule.loadModel(modelBuffer, options); + sessionKey = resultMap.getString("key"); ReadableArray inputNames = resultMap.getArray("inputNames"); ReadableArray outputNames = resultMap.getArray("outputNames"); - Assert.assertEquals(resultMap.getString("key"), "test"); Assert.assertEquals(inputNames.size(), 1); Assert.assertEquals(inputNames.getString(0), "input"); Assert.assertEquals(outputNames.size(), 1); @@ -103,7 +117,7 @@ public class OnnxruntimeModuleTest { options.putBoolean("encodeTensorData", true); try { - ReadableMap resultMap = ortModule.run("test", inputDataMap, outputNames, options); + ReadableMap resultMap = ortModule.run(sessionKey, inputDataMap, outputNames, options); ReadableMap outputMap = resultMap.getMap("output"); for (int i = 0; i < 2; ++i) { diff --git a/js/react_native/android/src/main/java/ai/onnxruntime/reactnative/OnnxruntimeModule.java b/js/react_native/android/src/main/java/ai/onnxruntime/reactnative/OnnxruntimeModule.java index fe59cefbee..d818733d7b 100644 --- a/js/react_native/android/src/main/java/ai/onnxruntime/reactnative/OnnxruntimeModule.java +++ b/js/react_native/android/src/main/java/ai/onnxruntime/reactnative/OnnxruntimeModule.java @@ -13,6 +13,7 @@ import ai.onnxruntime.OrtSession.RunOptions; import ai.onnxruntime.OrtSession.SessionOptions; import android.net.Uri; import android.os.Build; +import android.util.Base64; import android.util.Log; import androidx.annotation.NonNull; import androidx.annotation.RequiresApi; @@ -31,6 +32,7 @@ import java.io.IOException; import java.io.InputStream; import java.io.InputStreamReader; import java.io.Reader; +import java.math.BigInteger; import java.util.HashMap; import java.util.HashSet; import java.util.Iterator; @@ -46,6 +48,13 @@ public class OnnxruntimeModule extends ReactContextBaseJavaModule { private static OrtEnvironment ortEnvironment = OrtEnvironment.getEnvironment(); private static Map sessionMap = new HashMap<>(); + private static BigInteger nextSessionId = new BigInteger("0"); + private static String getNextSessionKey() { + String key = nextSessionId.toString(); + nextSessionId = nextSessionId.add(BigInteger.valueOf(1)); + return key; + } + public OnnxruntimeModule(ReactApplicationContext context) { super(context); reactContext = context; @@ -60,11 +69,11 @@ public class OnnxruntimeModule extends ReactContextBaseJavaModule { /** * React native binding API to load a model using given uri. * - * @param uri a model file location. it's used as a key when multiple sessions are created, i.e. multiple models are - * loaded. + * @param uri a model file location * @param options onnxruntime session options * @param promise output returning back to react native js - * @note when run() is called, the same uri must be passed into the first parameter. + * @note the value provided to `promise` includes a key representing the session. + * when run() is called, the key must be passed into the first parameter. */ @ReactMethod public void loadModel(String uri, ReadableMap options, Promise promise) { @@ -76,10 +85,30 @@ public class OnnxruntimeModule extends ReactContextBaseJavaModule { } } + /** + * React native binding API to load a model using the BASE64 encoded model data. + * + * @param data the BASE64 encoded model data. + * @param options onnxruntime session options + * @param promise output returning back to react native js + * @note the value provided to `promise` includes a key representing the session. + * when run() is called, the key must be passed into the first parameter. + */ + @ReactMethod + public void loadModelFromBase64EncodedBuffer(String data, ReadableMap options, Promise promise) { + try { + byte[] modelData = Base64.decode(data, Base64.DEFAULT); + WritableMap resultMap = loadModel(modelData, options); + promise.resolve(resultMap); + } catch (Exception e) { + promise.reject("Can't load model from buffer: " + e.getMessage(), e); + } + } + /** * React native binding API to run a model using given uri. * - * @param key a model file location given at loadModel() + * @param key session key representing a session given at loadModel() * @param input an input tensor * @param output an output names to be returned * @param options onnxruntime run options @@ -103,39 +132,49 @@ public class OnnxruntimeModule extends ReactContextBaseJavaModule { * @return model loading information, such as key, input names, and output names */ public WritableMap loadModel(String uri, ReadableMap options) throws Exception { - InputStream modelStream = reactContext.getApplicationContext().getContentResolver().openInputStream(Uri.parse(uri)); - WritableMap resultMap = loadModel(uri, modelStream, options); - modelStream.close(); - return resultMap; + return loadModelImpl(uri, null, options); } /** - * Load a model from raw resource directory. + * Load a model from buffer. * - * @param uri uri parameter from react native loadModel() or dummy string for unit testing purpose - * @param modelStream a input stream to read a model + * @param modelData the model data buffer * @param options onnxruntime session options * @return model loading information, such as key, input names, and output names */ - public WritableMap loadModel(String uri, InputStream modelStream, ReadableMap options) throws Exception { - OrtSession ortSession = null; + public WritableMap loadModel(byte[] modelData, ReadableMap options) throws Exception { + return loadModelImpl("", modelData, options); + } - if (!sessionMap.containsKey(uri)) { - byte[] modelArray = null; + /** + * Load model implementation method for either from model path or model data buffer. + * + * @param uri uri parameter from react native loadModel() + * @param modelData model data buffer + * @param options onnxruntime session options + * @return model loading information map, such as key, input names, and output names + */ + private WritableMap loadModelImpl(String uri, byte[] modelData, ReadableMap options) throws Exception { + OrtSession ortSession; + SessionOptions sessionOptions = parseSessionOptions(options); + if (modelData != null && modelData.length > 0) { // load model via model data array + ortSession = ortEnvironment.createSession(modelData, sessionOptions); + } else { // load model via model path string uri + InputStream modelStream = + reactContext.getApplicationContext().getContentResolver().openInputStream(Uri.parse(uri)); Reader reader = new BufferedReader(new InputStreamReader(modelStream)); - modelArray = new byte[modelStream.available()]; + byte[] modelArray = new byte[modelStream.available()]; modelStream.read(modelArray); - - SessionOptions sessionOptions = parseSessionOptions(options); + modelStream.close(); ortSession = ortEnvironment.createSession(modelArray, sessionOptions); - sessionMap.put(uri, ortSession); - } else { - ortSession = sessionMap.get(uri); } + String key = getNextSessionKey(); + sessionMap.put(key, ortSession); + WritableMap resultMap = Arguments.createMap(); - resultMap.putString("key", uri); + resultMap.putString("key", key); WritableArray inputNames = Arguments.createArray(); for (String inputName : ortSession.getInputNames()) { inputNames.pushString(inputName); @@ -153,7 +192,7 @@ public class OnnxruntimeModule extends ReactContextBaseJavaModule { /** * Run a model using given uri. * - * @param key a model file location given at loadModel() + * @param key a session key representing the session given at loadModel() * @param input an input tensor * @param output an output names to be returned * @param options onnxruntime run options @@ -162,7 +201,7 @@ public class OnnxruntimeModule extends ReactContextBaseJavaModule { public WritableMap run(String key, ReadableMap input, ReadableArray output, ReadableMap options) throws Exception { OrtSession ortSession = sessionMap.get(key); if (ortSession == null) { - throw new Exception("Model is not loaded: " + key); + throw new Exception("Model is not loaded."); } RunOptions runOptions = parseRunOptions(options); diff --git a/js/react_native/lib/backend.ts b/js/react_native/lib/backend.ts index 4ebc364cd8..6fd7d7bea4 100644 --- a/js/react_native/lib/backend.ts +++ b/js/react_native/lib/backend.ts @@ -48,25 +48,35 @@ class OnnxruntimeSessionHandler implements SessionHandler { #inferenceSession: Binding.InferenceSession; #key: string; + #pathOrBuffer: string|Uint8Array; + inputNames: string[]; outputNames: string[]; - constructor(path: string) { + constructor(pathOrBuffer: string|Uint8Array) { this.#inferenceSession = binding; - this.#key = normalizePath(path); + this.#pathOrBuffer = pathOrBuffer; + this.#key = ''; + this.inputNames = []; this.outputNames = []; } async loadModel(options: InferenceSession.SessionOptions): Promise { try { + let results: Binding.ModelLoadInfoType; // load a model - const results: Binding.ModelLoadInfoType = await this.#inferenceSession.loadModel(this.#key, options); - // resolve promise if onnxruntime session is successfully created - if (results.key !== this.#key) { - throw new Error('Session key is invalid'); + if (typeof this.#pathOrBuffer === 'string') { + results = await this.#inferenceSession.loadModel(normalizePath(this.#pathOrBuffer), options); + } else { + if (!this.#inferenceSession.loadModelFromBase64EncodedBuffer) { + throw new Error('Native module method "loadModelFromBase64EncodedBuffer" is not defined'); + } + const modelInBase64String = Buffer.from(this.#pathOrBuffer).toString('base64'); + results = await this.#inferenceSession.loadModelFromBase64EncodedBuffer(modelInBase64String, options); } - + // resolve promise if onnxruntime session is successfully created + this.#key = results.key; this.inputNames = results.inputNames; this.outputNames = results.outputNames; } catch (e) { @@ -156,9 +166,6 @@ class OnnxruntimeBackend implements Backend { async createSessionHandler(pathOrBuffer: string|Uint8Array, options?: InferenceSession.SessionOptions): Promise { - if (typeof pathOrBuffer !== 'string') { - throw new Error('Uint8Array is not supported'); - } const handler = new OnnxruntimeSessionHandler(pathOrBuffer); await handler.loadModel(options || {}); return handler; diff --git a/js/react_native/lib/binding.ts b/js/react_native/lib/binding.ts index afadbab971..4a4227e504 100644 --- a/js/react_native/lib/binding.ts +++ b/js/react_native/lib/binding.ts @@ -64,6 +64,7 @@ export declare namespace Binding { interface InferenceSession { loadModel(modelPath: string, options: SessionOptions): Promise; + loadModelFromBase64EncodedBuffer?(buffer: string, options: SessionOptions): Promise; run(key: string, feeds: FeedsType, fetches: FetchesType, options: RunOptions): Promise; } }