feat: BaseModel abstract class, primitive Style Transfer, functions for retrieving model metadata (#44)

## Description
This pull request introduces Model base class on top of which every
other model should be built, there's also a primitive implementation of
StyleTransferModel and functions for retrieving model metadata. Fix for
bindings allowing to return multiple outputs.

**__During code review don't focus on StyleTransfer Module as it will
change and it's for test purposes right now__**

Don't review:
- StyleTransfer.kt
- StyleTransferModel.kt
- StyleTransfer.h
- StyleTransfer.mm
- StyleTransferModel.mm
- StyleTransferModel.h
- BitmapUtils.kt
- Imageprocessor.h/Imageprocessor.m
- TensorUtils.kt

### Type of change
- [ ] Bug fix (non-breaking change which fixes an issue)
- [x] New feature (non-breaking change which adds functionality)
- [ ] Breaking change (fix or feature that would cause existing
functionality to not work as expected)
- [ ] Documentation update (improves or adds clarity to existing
documentation)

### Tested on
- [x] iOS
- [x] Android

### Testing instructions
<!-- Provide step-by-step instructions on how to test your changes.
Include setup details if necessary. -->

### Screenshots
<!-- Add screenshots here, if applicable -->

### Related issues
<!-- Link related issues here using #issue-number -->

### Checklist
- [x] I have performed a self-review of my code
- [ ] I have commented my code, particularly in hard-to-understand areas
- [ ] I have updated the documentation accordingly
- [x] My changes generate no new warnings

### Additional notes
<!-- Include any additional information, assumptions, or context that
reviewers might need to understand this PR. -->

---------

Co-authored-by: jakmro <[email protected]>
Co-authored-by: Jakub Chmura <[email protected]>
Co-authored-by: mkopcins <[email protected]>

Author: 4 people

android/src/main/java/com/swmansion/rnexecutorch/ETModule.kt

@@ -1,57 +1,35 @@
 package com.swmansion.rnexecutorch
 
+import com.facebook.react.bridge.Arguments
 import com.facebook.react.bridge.Promise
 import com.facebook.react.bridge.ReactApplicationContext
 import com.facebook.react.bridge.ReadableArray
 import com.swmansion.rnexecutorch.utils.ArrayUtils
+import com.swmansion.rnexecutorch.utils.ETError
 import com.swmansion.rnexecutorch.utils.Fetcher
-import com.swmansion.rnexecutorch.utils.ProgressResponseBody
-import com.swmansion.rnexecutorch.utils.ResourceType
 import com.swmansion.rnexecutorch.utils.TensorUtils
-import okhttp3.OkHttpClient
 import org.pytorch.executorch.Module
-import org.pytorch.executorch.Tensor
-import java.net.URL
 
 class ETModule(reactContext: ReactApplicationContext) : NativeETModuleSpec(reactContext) {
   private lateinit var module: Module
-  private val client = OkHttpClient()
 
   override fun getName(): String {
     return NAME
   }
 
-  private fun downloadModel(
-    url: String, resourceType: ResourceType, callback: (path: String?, error: Exception?) -> Unit
-  ) {
-    Fetcher.downloadResource(reactApplicationContext,
-      client,
-      url,
-      resourceType,
-      false,
-      { path, error -> callback(path, error) },
-      object : ProgressResponseBody.ProgressListener {
-        override fun onProgress(bytesRead: Long, contentLength: Long, done: Boolean) {
-        }
-      })
-  }
-
   override fun loadModule(modelPath: String, promise: Promise) {
-    try {
-      downloadModel(
-        modelPath, ResourceType.MODEL
-      ) { path, error ->
-        if (error != null) {
-          promise.reject(error.message!!, "-1")
-          return@downloadModel
-        }
-
-        module = Module.load(path)
-        promise.resolve(0)
+    Fetcher.downloadModel(
+      reactApplicationContext,
+      modelPath,
+    ) { path, error ->
+      if (error != null) {
+        promise.reject(error.message!!, ETError.InvalidModelPath.toString())
         return@downloadModel
       }
-    } catch (e: Exception) {
-      promise.reject(e.message!!, "-1")
+
+      module = Module.load(path)
+      promise.resolve(0)
+      return@downloadModel
     }
   }
 
@@ -75,19 +53,21 @@ class ETModule(reactContext: ReactApplicationContext) : NativeETModuleSpec(react
       val executorchInput =
         TensorUtils.getExecutorchInput(input, ArrayUtils.createLongArray(shape), inputType.toInt())
 
-      lateinit var result: Tensor
-      module.forward(executorchInput)[0].toTensor().also { result = it }
+      val result = module.forward(executorchInput)
+      val resultArray = Arguments.createArray()
+
+      for (evalue in result) {
+        resultArray.pushArray(ArrayUtils.createReadableArray(evalue.toTensor()))
+      }
 
-      promise.resolve(ArrayUtils.createReadableArray(result))
+      promise.resolve(resultArray)
       return
     } catch (e: IllegalArgumentException) {
       //The error is thrown when transformation to Tensor fails
-      promise.reject("Forward Failed Execution", "18")
+      promise.reject("Forward Failed Execution", ETError.InvalidArgument.code.toString())
       return
     } catch (e: Exception) {
-      //Executorch forward method throws an exception with a message: "Method forward failed with code XX"
-      val exceptionCode = e.message!!.substring(e.message!!.length - 2)
-      promise.reject("Forward Failed Execution", exceptionCode)
+      promise.reject("Forward Failed Execution", e.message!!)
       return
     }
   }

android/src/main/java/com/swmansion/rnexecutorch/RnExecutorchModule.kt

@@ -9,18 +9,15 @@ import com.swmansion.rnexecutorch.utils.ResourceType
 import com.swmansion.rnexecutorch.utils.llms.ChatRole
 import com.swmansion.rnexecutorch.utils.llms.ConversationManager
 import com.swmansion.rnexecutorch.utils.llms.END_OF_TEXT_TOKEN
-import okhttp3.OkHttpClient
 import org.pytorch.executorch.LlamaCallback
 import org.pytorch.executorch.LlamaModule
-import java.net.URL
 
 class RnExecutorchModule(reactContext: ReactApplicationContext) :
   NativeRnExecutorchSpec(reactContext), LlamaCallback {
 
   private var llamaModule: LlamaModule? = null
   private var tempLlamaResponse = StringBuilder()
   private lateinit var conversationManager: ConversationManager
-  private val client = OkHttpClient()
   private var isFetching = false
 
   override fun getName(): String {
@@ -51,7 +48,7 @@ class RnExecutorchModule(reactContext: ReactApplicationContext) :
     callback: (path: String?, error: Exception?) -> Unit,
   ) {
     Fetcher.downloadResource(
-      reactApplicationContext, client, url, resourceType, isLargeFile,
+      reactApplicationContext, url, resourceType, isLargeFile,
       { path, error -> callback(path, error) },
       object : ProgressResponseBody.ProgressListener {
         override fun onProgress(bytesRead: Long, contentLength: Long, done: Boolean) {

android/src/main/java/com/swmansion/rnexecutorch/RnExecutorchPackage.kt

@@ -17,6 +17,8 @@ class RnExecutorchPackage : TurboReactPackage() {
       RnExecutorchModule(reactContext)
     } else if (name == ETModule.NAME) {
       ETModule(reactContext)
+    } else if(name == StyleTransfer.NAME){
+      StyleTransfer(reactContext)
     } else {
       null
     }
@@ -40,6 +42,15 @@ class RnExecutorchPackage : TurboReactPackage() {
         false,  // isCxxModule
         true
       )
+
+      moduleInfos[StyleTransfer.NAME] = ReactModuleInfo(
+        StyleTransfer.NAME,
+        StyleTransfer.NAME,
+        false,  // canOverrideExistingModule
+        false,  // needsEagerInit
+        false,  // isCxxModule
+        true
+      )
       moduleInfos
     }
   }

android/src/main/java/com/swmansion/rnexecutorch/StyleTransfer.kt

@@ -0,0 +1,49 @@
+package com.swmansion.rnexecutorch
+
+import android.graphics.BitmapFactory
+import android.net.Uri
+import com.facebook.react.bridge.Promise
+import com.facebook.react.bridge.ReactApplicationContext
+import com.swmansion.rnexecutorch.models.StyleTransferModel
+import com.swmansion.rnexecutorch.utils.BitmapUtils
+import com.swmansion.rnexecutorch.utils.ETError
+
+class StyleTransfer(reactContext: ReactApplicationContext) :
+  NativeStyleTransferSpec(reactContext) {
+
+  private lateinit var styleTransferModel: StyleTransferModel
+
+  companion object {
+    const val NAME = "StyleTransfer"
+  }
+
+  override fun loadModule(modelSource: String, promise: Promise) {
+    try {
+      styleTransferModel = StyleTransferModel(reactApplicationContext)
+      styleTransferModel.loadModel(modelSource)
+      promise.resolve(0)
+    } catch (e: Exception) {
+      promise.reject(e.message!!, ETError.InvalidModelPath.toString())
+    }
+  }
+
+  override fun forward(input: String, promise: Promise) {
+    try {
+      val uri = Uri.parse(input)
+      val bitmapInputStream = reactApplicationContext.contentResolver.openInputStream(uri)
+      val rawBitmap = BitmapFactory.decodeStream(bitmapInputStream)
+      bitmapInputStream!!.close()
+
+      val output = styleTransferModel.runModel(rawBitmap)
+      val outputUri = BitmapUtils.saveToTempFile(output, "test")
+
+      promise.resolve(outputUri.toString())
+    }catch(e: Exception){
+      promise.reject(e.message!!, e.message)
+    }
+  }
+
+  override fun getName(): String {
+    return NAME
+  }
+}

android/src/main/java/com/swmansion/rnexecutorch/models/BaseModel.kt

@@ -0,0 +1,43 @@
+package com.swmansion.rnexecutorch.models
+
+import android.content.Context
+import com.swmansion.rnexecutorch.utils.ETError
+import com.swmansion.rnexecutorch.utils.Fetcher
+import org.pytorch.executorch.EValue
+import org.pytorch.executorch.Module
+
+
+abstract class BaseModel<Input, Output>(val context: Context) {
+  protected lateinit var module: Module
+
+  fun loadModel(modelSource: String) {
+    Fetcher.downloadModel(
+      context,
+      modelSource
+    ) { path, error ->
+      if (error != null) {
+        throw Error(error.message!!)
+      }
+
+      module = Module.load(path)
+    }
+  }
+
+  protected fun forward(input: EValue): Array<EValue> {
+    try {
+      val result = module.forward(input)
+      return result
+    } catch (e: IllegalArgumentException) {
+      //The error is thrown when transformation to Tensor fails
+      throw Error(ETError.InvalidArgument.code.toString())
+    } catch (e: Exception) {
+      throw Error(e.message!!)
+    }
+  }
+
+  abstract fun runModel(input: Input): Output
+
+  protected abstract fun preprocess(input: Input): Input
+
+  protected abstract fun postprocess(input: Output): Output
+}

android/src/main/java/com/swmansion/rnexecutorch/models/StyleTransferModel.kt

@@ -0,0 +1,51 @@
+package com.swmansion.rnexecutorch.models
+
+import android.graphics.Bitmap
+import android.util.Log
+import com.facebook.react.bridge.ReactApplicationContext
+import com.swmansion.rnexecutorch.utils.TensorUtils
+import org.pytorch.executorch.EValue
+
+class StyleTransferModel(reactApplicationContext: ReactApplicationContext) : BaseModel<Bitmap, Bitmap>(reactApplicationContext) {
+  override fun runModel(input: Bitmap): Bitmap {
+      val processedData = preprocess(input)
+      val inputTensor = TensorUtils.bitmapToFloat32Tensor(processedData)
+
+      Log.d("RnExecutorch", module.numberOfInputs.toString())
+      for (i in 0 until module.numberOfInputs) {
+        Log.d("RnExecutorch", module.getInputType(i).toString())
+        for(shape in module.getInputShape(i)){
+          Log.d("RnExecutorch", shape.toString())
+        }
+      }
+
+      Log.d("RnExecutorch", module.numberOfOutputs.toString())
+      for(i in 0 until module.numberOfOutputs){
+        Log.d("RnExecutorch", module.getOutputType(i).toString())
+        for(shape in module.getOutputShape(i)){
+          Log.d("RnExecutorch", shape.toString())
+        }
+      }
+
+      val outputTensor = forward(EValue.from(inputTensor))
+      val outputData = postprocess(TensorUtils.float32TensorToBitmap(outputTensor[0].toTensor()))
+
+      return outputData
+  }
+
+  override fun preprocess(input: Bitmap): Bitmap {
+    val inputBitmap = Bitmap.createScaledBitmap(
+      input,
+      640, 640, true
+    )
+    return inputBitmap
+  }
+
+  override fun postprocess(input: Bitmap): Bitmap {
+    val scaledUpBitmap = Bitmap.createScaledBitmap(
+      input,
+      1280, 1280, true
+    )
+    return scaledUpBitmap
+  }
+}

android/src/main/java/com/swmansion/rnexecutorch/utils/BitmapUtils.kt

@@ -0,0 +1,44 @@
+package com.swmansion.rnexecutorch.utils
+
+import android.graphics.Bitmap
+import android.graphics.Matrix
+import android.net.Uri
+import androidx.core.net.toUri
+import java.io.File
+import java.io.FileOutputStream
+import java.io.IOException
+
+class BitmapUtils {
+  companion object {
+    fun saveToTempFile(bitmap: Bitmap, fileName: String): Uri {
+      val tempFile = File.createTempFile(fileName, ".png")
+      var outputStream : FileOutputStream? = null
+      try {
+        outputStream = FileOutputStream(tempFile)
+        bitmap.compress(Bitmap.CompressFormat.PNG, 100, outputStream)
+      } catch (e: IOException) {
+        e.printStackTrace()
+      }
+      finally {
+        outputStream?.close()
+      }
+      return tempFile.toUri()
+    }
+
+    private fun rotateBitmap(bitmap: Bitmap, angle: Float): Bitmap {
+      val matrix = Matrix()
+      matrix.postRotate(angle)
+      return Bitmap.createBitmap(bitmap, 0, 0, bitmap.width, bitmap.height, matrix, true)
+    }
+
+    private fun flipBitmap(bitmap: Bitmap, horizontal: Boolean, vertical: Boolean): Bitmap {
+      val matrix = Matrix()
+      matrix.preScale(
+        if (horizontal) -1f else 1f,
+        if (vertical) -1f else 1f
+      )
+      return Bitmap.createBitmap(bitmap, 0, 0, bitmap.width, bitmap.height, matrix, true)
+    }
+  }
+
+}

android/src/main/java/com/swmansion/rnexecutorch/utils/ETError.kt

@@ -0,0 +1,29 @@
+package com.swmansion.rnexecutorch.utils
+
+enum class ETError(val code: Int) {
+    InvalidModelPath(0xff),
+    
+    // System errors
+    Ok(0x00),
+    Internal(0x01),
+    InvalidState(0x02),
+    EndOfMethod(0x03),
+    
+    // Logical errors
+    NotSupported(0x10),
+    NotImplemented(0x11),
+    InvalidArgument(0x12),
+    InvalidType(0x13),
+    OperatorMissing(0x14),
+    
+    // Resource errors
+    NotFound(0x20),
+    MemoryAllocationFailed(0x21),
+    AccessFailed(0x22),
+    InvalidProgram(0x23),
+    
+    // Delegate errors
+    DelegateInvalidCompatibility(0x30),
+    DelegateMemoryAllocationFailed(0x31),
+    DelegateInvalidHandle(0x32);
+}