Add audio input type

Author: kirklandsign

extension/android/jni/jni_layer_llama.cpp

@@ -273,6 +273,9 @@ class ExecuTorchLlmJni : public facebook::jni::HybridClass<ExecuTorchLlmJni> {
       jint width,
       jint height,
       jint channels) {
+    if (image == nullptr) {
+      return Error::InvalidArgument;
+    }
     std::vector<llm::Image> images;
     auto image_size = image->size();
     if (image_size != 0) {
@@ -290,6 +293,29 @@ class ExecuTorchLlmJni : public facebook::jni::HybridClass<ExecuTorchLlmJni> {
     return 0;
   }
 
+  jint append_audio_input(
+      facebook::jni::alias_ref<jintArray> audio,
+      jint batch_size,
+      jint n_channels,
+      jint n_samples) {
+    if (audio == nullptr) {
+      return Error::InvalidArgument;
+    }
+    auto audio_size = audio->size();
+    std::vector<uint8_t> audio_data(audio_size);
+    if (audio_size != 0) {
+      std::vector<jint> audio_data_jint(audio_size);
+      audio->getRegion(0, audio_size, audio_data_jint.data());
+      for (int i = 0; i < audio_size; i++) {
+        audio_data[i] = audio_data_jint[i];
+      }
+      llm::RawAudio audio_input{audio_data, batch_size, n_channels, n_samples};
+      prefill_inputs_.emplace_back(
+          llm::MultimodalInput{std::move(audio_input)});
+    }
+    return 0;
+  }
+
   void stop() {
     if (model_type_category_ == MODEL_TYPE_CATEGORY_MULTIMODAL) {
       multi_modal_runner_->stop();
@@ -321,6 +347,8 @@ class ExecuTorchLlmJni : public facebook::jni::HybridClass<ExecuTorchLlmJni> {
             "appendImagesInput", ExecuTorchLlmJni::append_images_input),
         makeNativeMethod(
             "appendTextInput", ExecuTorchLlmJni::append_text_input),
+        makeNativeMethod(
+            "appendAudioInput", ExecuTorchLlmJni::append_audio_input),
         makeNativeMethod("resetContext", ExecuTorchLlmJni::reset_context),
     });
   }