reflecting build changes

CMakeLists.txt

@@ -16,32 +16,27 @@ set(CMAKE_BUILD_TYPE Release)
 
 set(AXIE_SRC_DIR ${CMAKE_SOURCE_DIR}/src)
 set(AXIE_INC_DIR ${CMAKE_SOURCE_DIR}/inc)
+set(EXT_INC_DIR ${CMAKE_SOURCE_DIR}/external/inc)
 set(AXIE_CLI_DIR ${CMAKE_SOURCE_DIR}/cli)
 
 if(${ANDROID})
   message(NOTICE  "Config build for Android ARM64..")
+  OPTION(QNN_DELEGATE "QNN delegate" OFF) # Disabled by default
   set(CMAKE_TOOLCHAIN_FILE ${ANDROID_NDK_ROOT}/build/cmake/android.toolchain.cmake)
   set(ANDROID_ABI arm64-v8a)
-  set(AXIE_LIB_DIR ${CMAKE_SOURCE_DIR}/libs/android)
+  set(EXT_LIB_DIR ${CMAKE_SOURCE_DIR}/external/libs/android)
   find_library(android_log log)
 else()
-  message(NOTICE  "Config build for x86_64..")
-  set(AXIE_LIB_DIR ${CMAKE_SOURCE_DIR}/libs/x86)
+  message(NOTICE  "Config build for linux..")
+  set(EXT_LIB_DIR ${CMAKE_SOURCE_DIR}/external/libs/linux)
 endif()
 
-add_subdirectory(
-  "${TENSORFLOW_SOURCE_DIR}/tensorflow/lite"
-  "${CMAKE_CURRENT_BINARY_DIR}/tensorflow-lite"
-  EXCLUDE_FROM_ALL
-)
-
 link_directories(
-  ${AXIE_LIB_DIR}
+  ${EXT_LIB_DIR}
 )
 
 # build top library: whisperax
 if(${ANDROID})
-
   if(${QNN_DELEGATE}) # QCOM QNN delegate
     message(NOTICE  " Delegate via QCOM QNN..")
     add_library(whisperax   SHARED
@@ -75,6 +70,8 @@ endif()
 
 target_include_directories(whisperax PRIVATE
   ${AXIE_INC_DIR}
+  ${EXT_INC_DIR}
+  ${TENSORFLOW_SOURCE_DIR}
 )
 
 if(${ANDROID})
@@ -83,11 +80,11 @@ if(${ANDROID})
     ADD_DEFINITIONS(-DQNN_DELEGATE)
     add_library(qnn_delegate SHARED IMPORTED)
     set_property(TARGET qnn_delegate PROPERTY IMPORTED_LOCATION 
-      ${AXIE_LIB_DIR}/libQnnTFLiteDelegate.so
+      ${EXT_LIB_DIR}/libQnnTFLiteDelegate.so
     )
     target_link_libraries(whisperax PRIVATE 
       qnn_delegate 
-      tensorflow-lite 
+      tensorflowlite 
       tensorflowlite_gpu_delegate
       ${android_log}
       SDL3
@@ -96,7 +93,7 @@ if(${ANDROID})
   else() 
     ADD_DEFINITIONS(-DGPU_DELEGATE)
     target_link_libraries(whisperax PRIVATE 
-      tensorflow-lite 
+      tensorflowlite 
       tensorflowlite_gpu_delegate
       ${android_log}
       SDL3
@@ -105,7 +102,7 @@ if(${ANDROID})
 
 else()
   target_link_libraries(whisperax PRIVATE 
-    tensorflow-lite 
+    tensorflowlite 
     SDL3
   )
 endif()
@@ -126,6 +123,7 @@ add_executable(whisperax_cli
 target_include_directories(whisperax_cli PRIVATE
   ${AXIE_INC_DIR}
   ${AXIE_CLI_DIR}
+  ${EXT_INC_DIR}
 )
 
 target_link_libraries(whisperax_cli PRIVATE 

scripts/Dockerfile

@@ -33,6 +33,7 @@ ENV ANDROID_HOME=$ANDROID_SDK_HOME
 
 ENV QNN_SDK_ROOT=/opt/qnn-sdk
 ENV QNN_RUNTIME_ROOT=/opt/qnn-runtime
+ENV AXIE_ROOT=/src/AXIE
 
 ARG ANDROID_NDK_ZIP=$ANDROID_NDK_VERSION-linux.zip
 ARG BAZEL_INSTALLER=bazel-6.5.0-installer-linux-x86_64.sh
@@ -42,13 +43,11 @@ ARG QNN_TFLITE_DELEGATE=qnn-tflite-delegate-2.27.0.aar
 ARG ANDROID_COMMAND_LINE_TOOLS=commandlinetools-linux-11076708_latest.zip
 
 # Copy build dependencies
-ADD .build/tensorflow $TENSORFLOW_SOURCE_DIR/
-ADD scripts/build_tensorflow.sh /tmp/
-ADD .build/$ANDROID_NDK_ZIP $ANDROID_NDK_PATH/
-ADD .build/$BAZEL_INSTALLER $BAZEL_DIR/$BAZEL_INSTALLER
-ADD .build/$QNN_RUNTIME $QNN_RUNTIME_ROOT/
-ADD .build/$QNN_TFLITE_DELEGATE $QNN_SDK_ROOT/
-ADD .build/$ANDROID_COMMAND_LINE_TOOLS $ANDROID_SDK_HOME/
+ADD .source/$ANDROID_NDK_ZIP $ANDROID_NDK_PATH/
+ADD .source/$BAZEL_INSTALLER $BAZEL_DIR/$BAZEL_INSTALLER
+ADD .source/$QNN_RUNTIME $QNN_RUNTIME_ROOT/
+ADD .source/$QNN_TFLITE_DELEGATE $QNN_SDK_ROOT/
+ADD .source/$ANDROID_COMMAND_LINE_TOOLS $ANDROID_SDK_HOME/
 
 # Unzip and install dependencies
 RUN cd $QNN_RUNTIME_ROOT && unzip $QNN_RUNTIME && rm $QNN_RUNTIME && \
@@ -59,10 +58,10 @@ RUN cd $QNN_RUNTIME_ROOT && unzip $QNN_RUNTIME && rm $QNN_RUNTIME && \
     $CMD_TOOLS/bin/sdkmanager --install "platforms;android-34" --sdk_root=$ANDROID_SDK_HOME && \
     $CMD_TOOLS/bin/sdkmanager --install "build-tools;34.0.0" --sdk_root=$ANDROID_SDK_HOME && \
     $CMD_TOOLS/bin/sdkmanager --install "cmake;$CMAKE_VERION" --sdk_root=$ANDROID_SDK_HOME && \
-    chmod +x $BAZEL_DIR/$BAZEL_INSTALLER && $BAZEL_DIR/$BAZEL_INSTALLER && rm $BAZEL_DIR/$BAZEL_INSTALLER
-RUN chmod +x /tmp/build_tensorflow.sh && /tmp/build_tensorflow.sh
+    ln -s $AXIE_ROOT/.source/tensorflow $TENSORFLOW_SOURCE_DIR && \
+    chmod +x $BAZEL_DIR/$BAZEL_INSTALLER && $BAZEL_DIR/$BAZEL_INSTALLER && rm -rf $BAZEL_DIR
 
 # Set up the PATH
-ENV PATH=$PATH:$ANDROID_NDK_HOME:$TENSORFLOW_SOURCE_DIR:$ANDROID_SDK_HOME/$CMD_TOOLS/bin:$ANDROID_SDK_HOME/platform-tools:$ANDROID_SDK_HOME/cmake/$CMAKE_VERION/bin
+ENV PATH=$PATH:$ANDROID_NDK_HOME:$ANDROID_SDK_HOME/$CMD_TOOLS/bin:$ANDROID_SDK_HOME/platform-tools:$ANDROID_SDK_HOME/cmake/$CMAKE_VERION/bin
 
-WORKDIR /src/AXIE
+WORKDIR $AXIE_ROOT

scripts/adb_push.sh

@@ -2,18 +2,35 @@
 # For licensing see accompanying LICENSE file.
 # Copyright © 2024 Argmax, Inc. All rights reserved.
 
+for dev in `adb devices | grep -v "List" | awk '{print $1}'`
+do 
+  DEVICE=$dev
+  break
+done
+if [ "$DEVICE" = "" ]; then
+    echo "No Android device is connected via adb"
+    exit 0
+fi
+
+echo "Pushing to: $DEVICE"
+
 CURRENT_DIR="$(dirname "$(realpath "$0")")"
 SOURCE_DIR="$CURRENT_DIR/.."
 
-AXIE_TFLITE_CLI="$SOURCE_DIR/build/axie_tflite"
-LOCAL_LIBS="$SOURCE_DIR/libs"
-LOCAL_MODELS="$SOURCE_DIR/models"
-LOCAL_INPUTS="$SOURCE_DIR/inputs"
+AXIE_TFLITE_CLI="$SOURCE_DIR/build/android/whisperax_cli"
+AXIE_TFLITE_LIB="$SOURCE_DIR/build/android/libwhisperax.so"
+LOCAL_LIBS="$SOURCE_DIR/external/libs/android"
+LOCAL_TINY_DIR="$SOURCE_DIR/openai_whisper-tiny"
+LOCAL_BASE_DIR="$SOURCE_DIR/openai_whisper-base"
+LOCAL_SMALL_DIR="$SOURCE_DIR/openai_whisper-small"
 
 REMOTE_BIN_DIR="/data/local/tmp/bin"
 REMOTE_LIB_DIR="/data/local/tmp/lib"
-REMOTE_MODELS_DIR="/sdcard/argmax/tflite/models"
-REMOTE_INPUTS_DIR="/sdcard/argmax/tflite/inputs"
+REMOTE_SDROOT_DIR="/sdcard/argmax/tflite"
+REMOTE_TINY_DIR="${REMOTE_SDROOT_DIR}/openai_whisper-tiny"
+REMOTE_BASE_DIR="${REMOTE_SDROOT_DIR}/openai_whisper-base"
+REMOTE_SMALL_DIR="${REMOTE_SDROOT_DIR}/openai_whisper-small"
+REMOTE_INPUTS_DIR="${REMOTE_SDROOT_DIR}/inputs"
 
 # Function to push files only if they do not exist
 push_if_not_exists() {
@@ -29,8 +46,8 @@ push_if_not_exists() {
         done
     else
         # If it's a file, check if it exists on the remote device
-        if adb shell "[ ! -e $remote_path ]"; then
-            adb push "$local_path" "$remote_path"
+        if adb -s $DEVICE shell "[ ! -e $remote_path ]"; then
+            adb -s $DEVICE push "$local_path" "$remote_path"
         else
             # uncomment to debug
             # echo "$remote_path already exists. Skipping push."
@@ -40,7 +57,23 @@ push_if_not_exists() {
 }
 
 # Push the files and folders to the Android device
-adb push "$AXIE_TFLITE_CLI" "$REMOTE_BIN_DIR/"
+if adb -s $DEVICE shell "[ ! -e $REMOTE_SDROOT_DIR ]"; then
+    adb -s $DEVICE shell mkdir "$REMOTE_SDROOT_DIR"
+fi
+if adb -s $DEVICE shell "[ ! -e $REMOTE_INPUTS_DIR ]"; then
+    adb -s $DEVICE shell mkdir "$REMOTE_INPUTS_DIR"
+fi
+if adb -s $DEVICE shell "[ ! -e $REMOTE_BIN_DIR ]"; then
+    adb -s $DEVICE shell mkdir "$REMOTE_BIN_DIR"
+fi
+if adb -s $DEVICE shell "[ ! -e $REMOTE_LIB_DIR ]"; then
+    adb -s $DEVICE shell mkdir "$REMOTE_LIB_DIR"
+fi
+
+adb -s $DEVICE push "$AXIE_TFLITE_CLI" "$REMOTE_BIN_DIR/."
+adb -s $DEVICE push "$AXIE_TFLITE_LIB" "$REMOTE_LIB_DIR/."
+
 push_if_not_exists "$LOCAL_LIBS" "$REMOTE_LIB_DIR"
-push_if_not_exists "$LOCAL_MODELS" "$REMOTE_MODELS_DIR"
-push_if_not_exists "$LOCAL_INPUTS" "$REMOTE_INPUTS_DIR"
+push_if_not_exists "$LOCAL_TINY_DIR" "$REMOTE_TINY_DIR"
+push_if_not_exists "$LOCAL_BASE_DIR" "$REMOTE_BASE_DIR"
+push_if_not_exists "$LOCAL_SMALL_DIR" "$REMOTE_SMALL_DIR"

scripts/build_tensorflow.sh

@@ -5,21 +5,56 @@
 # This build script runs when docker image is created.
 # The resulting `libtensorflowlite_gpu_delegate.so` is copied into /libs folder in the build.sh
 
+CURRENT_DIR="$(dirname "$(realpath "$0")")"
+SOURCE_DIR="$CURRENT_DIR/.."
+PLATFORM=$1
+if [ "$PLATFORM" = "" ]; then
+    PLATFORM="android"
+fi
+
 export PYTHON_BIN_PATH=/usr/bin/python3
 export PYTHON_LIB_PATH=/usr/lib/python3/dist-packages
 export TF_NEED_ROCM=0
 export TF_NEED_CUDA=0
 export TF_NEED_CLANG=1
 export CLANG_COMPILER_PATH=/usr/bin/clang
 export CC_OPT_FLAGS=-Wno-sign-compare
-export TF_SET_ANDROID_WORKSPACE=1
-export ANDROID_NDK_API_LEVEL=24
-export ANDROID_API_LEVEL=34
-export ANDROID_BUILD_TOOLS_VERSION=34.0.0
 
-cd $TENSORFLOW_SOURCE_DIR && ./configure
+if [ "$PLATFORM" = "android" ]; then
+    export TF_SET_ANDROID_WORKSPACE=1
+    export ANDROID_NDK_API_LEVEL=24
+    export ANDROID_API_LEVEL=34
+    export ANDROID_BUILD_TOOLS_VERSION=34.0.0
+
+    cd $TENSORFLOW_SOURCE_DIR && ./configure
+
+    if [ ! -f $SOURCE_DIR/external/libs/$PLATFORM/libtensorflowlite_gpu_delegate.so ]; then
+        echo "$SOURCE_DIR/external/libs/$PLATFORM ..."
+        echo "Building libtensorflowlite_gpu_delegate.so ..."
+        printenv
+        mkdir -p tensorflow/lite/delegates/gpu
+        bazel build -c opt --config android_arm64 --cxxopt=--std=c++17 tensorflow/lite/delegates/gpu:libtensorflowlite_gpu_delegate.so
+        find "$TENSORFLOW_SOURCE_DIR/" $TENSORFLOW_SOURCE_DIR/bazel-bin/ \
+            -name libtensorflowlite_gpu_delegate.so -exec cp {} $SOURCE_DIR/external/libs/android/ \;
+    fi
+
+    if [ ! -f $SOURCE_DIR/external/libs/$PLATFORM/libtensorflowlite.so ]; then
+        bazel build -c opt --config android_arm64 --cxxopt=--std=c++17 //tensorflow/lite:libtensorflowlite.so
+        find "$TENSORFLOW_SOURCE_DIR/" $TENSORFLOW_SOURCE_DIR/bazel-bin/ \
+            -name libtensorflowlite.so -exec cp {} $SOURCE_DIR/external/libs/$PLATFORM/ \;
+    fi
+else 
+    export TF_SET_ANDROID_WORKSPACE=0
+    if [ ! -f $SOURCE_DIR/external/libs/$PLATFORM/libtensorflowlite.so ]; then
+        cd $TENSORFLOW_SOURCE_DIR && ./configure
+
+        bazel build //tensorflow/lite:libtensorflowlite.so
+        find "$TENSORFLOW_SOURCE_DIR/" $TENSORFLOW_SOURCE_DIR/bazel-bin/ \
+            -name libtensorflowlite.so -exec cp {} $SOURCE_DIR/external/libs/$PLATFORM/ \;
+    fi
+fi
 
-echo "Building libtensorflowlite_gpu_delegate.so ..."
-printenv
-mkdir -p tensorflow/lite/delegates/gpu
-bazel build -c opt --config android_arm64 --cxxopt=--std=c++17 tensorflow/lite/delegates/gpu:libtensorflowlite_gpu_delegate.so
+if [ ! -d $SOURCE_DIR/external/inc/flatbuffers ]; then
+    cp -rf $TENSORFLOW_SOURCE_DIR/bazel-tensorflow/external/flatbuffers/include/flatbuffers \
+        $SOURCE_DIR/external/inc/.
+fi

scripts/dev_env.sh

@@ -3,22 +3,24 @@
 # Copyright © 2024 Argmax, Inc. All rights reserved.
 
 # This script builds and runs `android-ndk-qnn-tensorflow-image` docker image with all dependencies
-# As part of build process, the script downloads all dependencies into the .build/ folder.
+# As part of build process, the script downloads all dependencies into the .source/ folder.
 # You will need `aria2` installed (see https://formulae.brew.sh/formula/aria2)
 
 IMAGE_NAME="android-ndk-qnn-tensorflow-image"
 CONTAINER_NAME="axie_tflite"
 FORCE_REBUILD=false
 FORCE_REMOVE=false
+CI_MODE=false
 
 CURRENT_DIR="$(dirname "$(realpath "$0")")"
 SOURCE_DIR="$CURRENT_DIR/.."
 
-while getopts "rf" opt; do
+while getopts "rfc" opt; do
   case ${opt} in
     r ) FORCE_REBUILD=true ;;
     f ) FORCE_REMOVE=true ;;
-    \? ) echo "Usage: cmd [-r] [-f]"
+    c ) CI_MODE=true ;;
+    \? ) echo "Usage: cmd [-r] [-f] [-c]"
          exit 1 ;;
   esac
 done
@@ -37,7 +39,7 @@ if ! $(docker image inspect $IMAGE_NAME > /dev/null 2>&1) || $FORCE_REBUILD; the
   # Set Aria options to download using 8 connections
   ARIA_OPTIONS="-x 8 -s 8 --continue --file-allocation=none"
 
-  BUILD_DIR="$SOURCE_DIR/.build"
+  BUILD_DIR="$SOURCE_DIR/.source"
   echo "Checking and retrieving dependencies..."
   if command -v aria2c &> /dev/null; then
     aria2c $ARIA_OPTIONS -d $BUILD_DIR https://github.com/bazelbuild/bazel/releases/download/6.5.0/bazel-6.5.0-installer-linux-x86_64.sh
@@ -68,24 +70,32 @@ else
   echo "Docker image $IMAGE_NAME already exists."
 fi
 
+BASH_CMD="echo 'Environment Variables:' && printenv && /bin/bash"
+if $CI_MODE; then
+  # in CI mode, keep running /bin/bash but in the background
+  RUN_OPTS="-d -i"
+else
+  RUN_OPTS="-it"
+fi
+
 # Check if the container exists
 if [ "$(docker ps -aq -f name=$CONTAINER_NAME)" ]; then
   if $FORCE_REMOVE; then
     echo "Removing existing container: $CONTAINER_NAME"
     docker rm -f $CONTAINER_NAME
   else
     if [ ! "$(docker ps -q -f name=$CONTAINER_NAME)" ]; then
-    echo "Starting existing container: $CONTAINER_NAME"
-    docker start $CONTAINER_NAME
+      echo "Starting existing container: $CONTAINER_NAME"
+      docker start $CONTAINER_NAME
     fi
     echo "SSHing into existing container: $CONTAINER_NAME"
-    docker exec -it $CONTAINER_NAME /bin/bash -c "echo 'Environment Variables:' && printenv && exec /bin/bash"
+    docker exec $RUN_OPTS $CONTAINER_NAME /bin/bash -c "${BASH_CMD}"
     exit 0
   fi
 fi
 
 # Run a new container
 echo "Starting new container: $CONTAINER_NAME"
-docker run -it --name $CONTAINER_NAME \
+docker run --platform linux/amd64 $RUN_OPTS --name $CONTAINER_NAME \
   --mount type=bind,source=$SOURCE_DIR,target=/src/AXIE \
-  $IMAGE_NAME /bin/bash -c "echo 'Environment Variables:' && printenv && exec /bin/bash"
+  $IMAGE_NAME /bin/bash -c "${BASH_CMD}"