From 93649246fe44218e7ef4c4a0df189fa81013c5ca Mon Sep 17 00:00:00 2001 From: Hansong Zhang <107070759+kirklandsign@users.noreply.github.com> Date: Wed, 15 Oct 2025 12:12:53 -0700 Subject: [PATCH] Android audio input API (#15166) Expose all llm runner API to java Mark some API as experimental, not deprecated (cherry picked from commit 1076686c5d41ee434a59b5e69232000fc3e676f0) --- .../executorch/extension/llm/LlmModule.java | 78 ++++++++++++++++-- extension/android/jni/jni_layer_llama.cpp | 80 +++++++++++++++++++ 2 files changed, 153 insertions(+), 5 deletions(-) diff --git a/extension/android/executorch_android/src/main/java/org/pytorch/executorch/extension/llm/LlmModule.java b/extension/android/executorch_android/src/main/java/org/pytorch/executorch/extension/llm/LlmModule.java index 289df5defd9..40e38afb8b9 100644 --- a/extension/android/executorch_android/src/main/java/org/pytorch/executorch/extension/llm/LlmModule.java +++ b/extension/android/executorch_android/src/main/java/org/pytorch/executorch/extension/llm/LlmModule.java @@ -167,7 +167,7 @@ public int generate( } /** - * Prefill an LLaVA Module with the given images input. + * Prefill a multimodal Module with the given images input. * * @param image Input image as a byte array * @param width Input image width @@ -177,7 +177,7 @@ public int generate( * exposed to user. * @throws RuntimeException if the prefill failed */ - @Deprecated + @Experimental public long prefillImages(int[] image, int width, int height, int channels) { int nativeResult = appendImagesInput(image, width, height, channels); if (nativeResult != 0) { @@ -189,14 +189,82 @@ public long prefillImages(int[] image, int width, int height, int channels) { private native int appendImagesInput(int[] image, int width, int height, int channels); /** - * Prefill an LLaVA Module with the given text input. + * Prefill a multimodal Module with the given images input. * - * @param prompt The text prompt to LLaVA. + * @param image Input normalized image as a float array + * @param width Input image width + * @param height Input image height + * @param channels Input image number of channels + * @return 0, as the updated starting position in KV cache of the input in the LLM is no longer + * exposed to user. + * @throws RuntimeException if the prefill failed + */ + @Experimental + public long prefillImages(float[] image, int width, int height, int channels) { + int nativeResult = appendNormalizedImagesInput(image, width, height, channels); + if (nativeResult != 0) { + throw new RuntimeException("Prefill failed with error code: " + nativeResult); + } + return 0; + } + + private native int appendNormalizedImagesInput( + float[] image, int width, int height, int channels); + + /** + * Prefill a multimodal Module with the given audio input. + * + * @param audio Input preprocessed audio as a byte array + * @param batch_size Input batch size + * @param n_bins Input number of bins + * @param n_frames Input number of frames + * @return 0, as the updated starting position in KV cache of the input in the LLM is no longer + * exposed to user. + * @throws RuntimeException if the prefill failed + */ + @Experimental + public long prefillAudio(byte[] audio, int batch_size, int n_bins, int n_frames) { + int nativeResult = appendAudioInput(audio, batch_size, n_bins, n_frames); + if (nativeResult != 0) { + throw new RuntimeException("Prefill failed with error code: " + nativeResult); + } + return 0; + } + + private native int appendAudioInput(byte[] audio, int batch_size, int n_bins, int n_frames); + + /** + * Prefill a multimodal Module with the given raw audio input. + * + * @param audio Input raw audio as a byte array + * @param batch_size Input batch size + * @param n_channels Input number of channels + * @param n_samples Input number of samples + * @return 0, as the updated starting position in KV cache of the input in the LLM is no longer + * exposed to user. + * @throws RuntimeException if the prefill failed + */ + @Experimental + public long prefillRawAudio(byte[] audio, int batch_size, int n_channels, int n_samples) { + int nativeResult = appendRawAudioInput(audio, batch_size, n_channels, n_samples); + if (nativeResult != 0) { + throw new RuntimeException("Prefill failed with error code: " + nativeResult); + } + return 0; + } + + private native int appendRawAudioInput( + byte[] audio, int batch_size, int n_channels, int n_samples); + + /** + * Prefill a multimodal Module with the given text input. + * + * @param prompt The text prompt to prefill. * @return 0, as the updated starting position in KV cache of the input in the LLM is no longer * exposed to user. * @throws RuntimeException if the prefill failed */ - @Deprecated + @Experimental public long prefillPrompt(String prompt) { int nativeResult = appendTextInput(prompt); if (nativeResult != 0) { diff --git a/extension/android/jni/jni_layer_llama.cpp b/extension/android/jni/jni_layer_llama.cpp index cabf30c42e4..ccb0d55dc41 100644 --- a/extension/android/jni/jni_layer_llama.cpp +++ b/extension/android/jni/jni_layer_llama.cpp @@ -276,6 +276,79 @@ class ExecuTorchLlmJni : public facebook::jni::HybridClass { return 0; } + // Returns status_code + jint append_normalized_images_input( + facebook::jni::alias_ref image, + jint width, + jint height, + jint channels) { + std::vector images; + if (image == nullptr) { + return static_cast(Error::EndOfMethod); + } + auto image_size = image->size(); + if (image_size != 0) { + std::vector image_data_jfloat(image_size); + std::vector image_data(image_size); + image->getRegion(0, image_size, image_data_jfloat.data()); + for (int i = 0; i < image_size; i++) { + image_data[i] = image_data_jfloat[i]; + } + llm::Image image_runner{std::move(image_data), width, height, channels}; + prefill_inputs_.emplace_back( + llm::MultimodalInput{std::move(image_runner)}); + } + + return 0; + } + + // Returns status_code + jint append_audio_input( + facebook::jni::alias_ref data, + jint batch_size, + jint n_bins, + jint n_frames) { + if (data == nullptr) { + return static_cast(Error::EndOfMethod); + } + auto data_size = data->size(); + if (data_size != 0) { + std::vector data_jbyte(data_size); + std::vector data_u8(data_size); + data->getRegion(0, data_size, data_jbyte.data()); + for (int i = 0; i < data_size; i++) { + data_u8[i] = data_jbyte[i]; + } + llm::Audio audio{std::move(data_u8), batch_size, n_bins, n_frames}; + prefill_inputs_.emplace_back(llm::MultimodalInput{std::move(audio)}); + } + return 0; + } + + // Returns status_code + jint append_raw_audio_input( + facebook::jni::alias_ref data, + jint batch_size, + jint n_channels, + jint n_samples) { + if (data == nullptr) { + return static_cast(Error::EndOfMethod); + } + auto data_size = data->size(); + if (data_size != 0) { + std::vector data_jbyte(data_size); + std::vector data_u8(data_size); + data->getRegion(0, data_size, data_jbyte.data()); + for (int i = 0; i < data_size; i++) { + data_u8[i] = data_jbyte[i]; + } + llm::RawAudio audio{ + std::move(data_u8), batch_size, n_channels, n_samples}; + prefill_inputs_.emplace_back(llm::MultimodalInput{std::move(audio)}); + } + return 0; + } + void stop() { if (model_type_category_ == MODEL_TYPE_CATEGORY_MULTIMODAL) { multi_modal_runner_->stop(); @@ -310,6 +383,13 @@ class ExecuTorchLlmJni : public facebook::jni::HybridClass { makeNativeMethod("load", ExecuTorchLlmJni::load), makeNativeMethod( "appendImagesInput", ExecuTorchLlmJni::append_images_input), + makeNativeMethod( + "appendNormalizedImagesInput", + ExecuTorchLlmJni::append_normalized_images_input), + makeNativeMethod( + "appendAudioInput", ExecuTorchLlmJni::append_audio_input), + makeNativeMethod( + "appendRawAudioInput", ExecuTorchLlmJni::append_raw_audio_input), makeNativeMethod( "appendTextInput", ExecuTorchLlmJni::append_text_input), makeNativeMethod("resetContext", ExecuTorchLlmJni::reset_context),