7. Optimizing AI Performance & Deployment with Play for On-device AI

Written by Zahidur Rahman Faisal

Hello reader! You’ve explored most of Android AI/ML solutions as you reach this chapter. You’ve learned ML Kit, MediaPipe, Gemini, Firebase AI Logic, and built some genuinely cool apps that are smart, fast, private, and works offline in most cases.

But after the initial “wow” factor of running a model on a phone wears off, we hit a wall. It’s what I call the “last mile” problem of on-device AI.

The proof-of-concept works great on your high-end test devices. But then the questions start rolling in:

“This new generative model is 500MB. We can’t just stick that in the APK, can we?

“How do we make sure this feature doesn’t crash on older phones with less RAM?”

“Our data science team has a new, better version of the model. How do we ship it to users without forcing a full app update?”

Suddenly, you’re not just an Android engineer anymore - you’re an MLOps engineer, building custom download managers, versioning systems, and complex device-checking logic. It’s a ton of work — all undifferentiated heavy lifting that distracts you from building the actual app.

For a while, that was just the cost of doing business on the cutting edge. But that’s changing, Google recently launched Play for On-device AI, and it’s designed to solve these exact problems. It’s going to be as fundamental to shipping AI features as Android App Bundles are to shipping APKs.

Remember building the On-device LLM app at Chapter 5? You had to download and copy the TinyLlama model, which was about 1.25 GB! Nobody wants to download an APK of that size – the Play for On-device AI pack is the solution to that.

What Is Play for On-device AI Pack?

Think about the problems that App Bundles solved. Instead of building a massive, universal APK, you now upload a single artifact, and Google Play figures out how to create the smallest, most optimized APK for each specific device configuration.

Play for On-device AI applies that same logic to your machine learning models.

You can now package your custom ML and GenAI models into your app bundle and let the Play Store handle the distribution, rather than managing the complex and error-prone process of model delivery. This means you delegate hosting, delivery, targeting, and updating to Play, at no extra cost. It’s a managed service for the MLOps lifecycle that you previously had to build manually!

The best part is that you control when the model is delivered to the user. You can configure model downloads as:

• Install-time – delivered with the app during installation.

• Fast-follow – downloaded immediately after the app is installed.

• On-demand – downloaded only when your app explicitly requests it.

Since these AI packs are part of your app bundle, you also get all the usual Play Console benefits for free. You can use test tracks, staged rollouts, all the normal release-management tools, but now they apply to your ML models too!

Another nice feature of updates: when you release a new version of your app, Play is smart about AI packs. If a particular pack hasn’t changed, users won’t need to download it again. Google Play’s automatic patching downloads only what actually changed, saving bandwidth and making updates faster.

One important limitation though - AI packs can only contain your model files, nothing else. You can’t put Java, Kotlin, or native libraries in there. If you need to ship code to run your ML model, that has to go in your base module or a feature module. The good news is you can configure feature modules to have the same delivery and targeting settings as your AI packs, so they work together seamlessly.

Solving the Deployment Puzzle: AI Packs and Delivery Modes

The core concept is the AI pack. It’s a special container within your app bundle that holds your models. The magic is in how these packs get delivered to the user. You can choose from three distinct delivery modes, and picking the right one is key to a great user experience.

1. Install-time Delivery

   This is the most straightforward option. The AI pack is delivered and installed alongside the main app, just like a regular APK.  

   • When to use: For critical, must-have AI features. If a core part of your app depends on a model, make sure it’s there the first time the user opens the app.

   • Trade-off: Increases the initial download size from the Play Store. For large models, this can deter users.

2. Fast-Follow Delivery

   This is the sweet spot for many features. The app itself is installed quickly, and the AI pack begins downloading automatically in the background right after.  

   • When to use: Ideal for important but non-critical AI features. The user can explore your app immediately, and the AI feature “lights up” a few moments later.

   • Trade-off: You must handle the state where the model isn’t ready yet. Show a loading UI or disable the feature until download completes. Your code should check download status on each app launch.

3. On-Demand Delivery

   This gives you the most control. The AI pack is only downloaded when your app explicitly requests it at runtime.

   • When to use: Best for large models powering niche features — for instance, a “pro” photo filter or a specialized generative-AI tool used by few users.

   • Trade-off: Requires more UI work — you must trigger the download, display progress, handle errors, and enable the feature when ready.

By offering these three modes, Play lets you balance feature availability, app size, and user experience intelligently.

An Overview of the Play for On-device AI Workflow

Using Play for On-device AI to package and deploy AI Packs is pretty straightforward. Let me break down the process for you:

1. You’ll need to package your trained models into AI packs and include them in your Android App Bundle. During this step, specify the delivery mode - install-time, fast-follow, or on-demand delivery for each pack.

2. If you choose fast-follow or on-demand delivery, you’ll need to integrate the Play AI Delivery Library into your app. This library handles the actual downloading of your AI packs when needed. If you went with install-time delivery, you can skip this step since the models will already be there when your app launches.

3. If you want more precise targeting, configure which devices receive which models. This step is optional, yet it’s quite handy! Imagine delivering a lightweight model to budget phones and a more sophisticated model to flagship devices with 6GB+ RAM. You can even target specific device models or exclude devices that can’t handle the model.

4. Finally, test and release your app bundle through the Play Console, just like any other update. You can use internal testing, staged rollouts, and every release-management tool you already know.

Later in this chapter, you’ll learn how to configure your app for on-demand delivery of an ML model, extending the “On-device LLM” project from Chapter 5.

Download the starter project and open it in Android Studio to begin.

Note that, the starter project contains some convenience methods in the MainViewModel, which won’t compile until you add the Play AI Delivery Library to the project. This is intentional so the chapter can focus on enabling Play-based-on-device capabilities rather than writing boilerplate code. Once you add the Play AI Delivery Library as a dependency (in Step 2), the project will compile and run successfully.

Note: The TinyLlama LLM model used in this sample project comes from the Hugging Face community and can be downloaded here.

Step 1: Configuring Your On-Demand AI Pack

• Create the AI Pack module: In the root directory of the Android project, create a new directory for the AI Pack. The directory name will be used as the AI Pack’s identifier. It must start with a letter and can only contain letters, numbers, and underscores (e.g., myModelPack or my_model_pack_1).

  For this project, name it onDemandAiPack. The project structure should look like this:

• Create the module’s build.gradle: Inside the onDemandAiPack directory, create a build.gradle file. In this file, apply the com.android.ai-pack plugin, then define the AI Pack’s name and specify a delivery mode (install-time, fast-follow, or on-demand).

// In onDemandAiPack/build.gradle
plugins {
  id 'com.android.ai-pack'
}

aiPack {
  packName = "onDemandAiPack"
  dynamicDelivery {
    deliveryType = "on-demand"
  }
}

• Register the pack in the app module: In the main application module’s build.gradle file (typically app/build.gradle), declare the new AI Pack inside the android block.

// In app/build.gradle
android {
  //...
  assetPacks = [":onDemandAiPack"]
}

• Include the module in settings.gradle file: Ensure the Gradle build system recognizes the new module by adding an include statement in the project’s root settings.gradle file.

// In settings.gradle
include ':app'
include ':onDemandAiPack'

• Structure the AI Pack and place the model: Within the AI Pack module directory, onDemandAiPack, create the standard source-set structure: src/main/assets/. Place your AI model file (for example, TinyLlama-1.1B-Chat-v1.0_multi-prefill-seq_q8_ekv1280.task) inside the assets directory.

  This path is preserved during the build and deployment process, and the final accessible path on the device will look like: [pack_path]/assets/[model_file_name]

Step 2: Integrating and Using Play AI Delivery

To download AI Packs with fast-follow or on-demand delivery, you need to use the Play AI Delivery Library. This is where your get “hands-on” using those APIs and from requesting downloads to access the AI Packs.

The process includes a few steps:

• Including the library dependency.

• Checking any existing or downloaded AI Pack status.

• Adding capability to download the AI Packs, or cancel any ongoing downloads.

• Monitoring status updates of the downloads and displaying them on the UI.

• Finally, accessing or using the AI Packs once it’s downloaded.

Adding the Dependency

Get started by adding this dependency In the app-level build.gradle file:

// In app/build.gradle
dependencies {
  ...
  // TUTORIAL DEPENDENCIES
  var aiDeliveryVersion = "0.1.1-alpha01"
  implementation "com.google.android.play:ai-delivery:$aiDeliveryVersion"
}

Sync the project with Gradle files from the Android Studio. Open MainViewModel.kt and import all the necessary libraries.

Build and run, the project should compile successfully. At this point, you’ll see a loading state: “Checking for AI Model…”.

Now you need to complete the status check and update the UI with results.

Checking the Status of AI Packs

Before using assets or models from an AI pack, you should check whether the pack has already been downloaded.

Open MainViewModel.kt, the aiPackManager serves as the primary interface for managing the download and state of AI Packs delivered via fast-follow or on-demand modes.

Now, update the checkPackStatus() function as follows:

fun checkPackStatus(packName: String) {
  val packLocation = aiPackManager.getPackLocation(packName)
  if (packLocation != null) {
    _aiPackStatus.value = AiPackStatus.Installed(packLocation.toString())
    return
  }

  aiPackManager.getPackStates(listOf(packName))
    .addOnSuccessListener { states ->
      val state = states.packStates()[packName]
      _aiPackStatus.value = mapAiPackStatus(state = state)
    }.addOnFailureListener { e ->
      _aiPackStatus.value = mapAiPackStatus(state = null)
    }
}

What it does is straightforward:

• Takes a packName to check if it’s already installed. If so, it’ll return the asset’s location.

• Otherwise, aiPackManager will invoke getPackStates() function. You can pass a list of packages to that function to check states sequentially and check state of the given packName here. On a successful completion, you need to map the state to local states and update the UI accordingly.

• In case of failure, the status is unknown, and you need to reflect that on the UI. Setting _aiPackStatus.value = mapAiPackStatus(state = null) does that.

But why mapAiPackStatus() function is needed? To understand, you need to have a deeper look at that; You’ll explore that shortly.

Fetching AI Packs

The aiPackManager handles the heavy-lifting. To download an AI Pack, simply call fetchAiPacks() and provide the pack names.

Add fetchAiPacks() function in the MainViewModel to do so.

fun fetchAiPacks() {
  aiPackManager.fetch(AI_PACKS)
}

Here, AI_PACKS is a list containing the names of all AI Packs you want to download.

Cancelling Requests

In case you want to cancel any fetch request, cancelling is as simple as fetching AI Packs using Play AI Delivery Library. Adding this function to the MainViewModel will allow you to do that:

fun cancelRequests() {
  aiPackManager.cancel(AI_PACKS)
}

The aiPackManager.cancel() method can take a list of AI Packs and cancels any ongoing downloads for those.

Displaying Status Updates

Play AI Delivery Library returns AiPackState while fetching AI Packs or checking AI Pack status. The AiPackState is a public interface that looks like this:

@Retention(RetentionPolicy.CLASS)
public @interface AiPackStatus {
  int UNKNOWN = 0;
  int NOT_INSTALLED = 8;
  int PENDING = 1;
  int WAITING_FOR_WIFI = 7;
  int REQUIRES_USER_CONFIRMATION = 9;
  int DOWNLOADING = 2;
  int TRANSFERRING = 3;
  int COMPLETED = 4;
  int FAILED = 5;
  int CANCELED = 6;
}

To translate them into a simpler from for the user, you have an AiPackStatus data class within the aipack package in the sample project. Let’s look at how the mapping works in the mapAiPackStatus() function inside MainViewModel:

private fun mapAiPackStatus(state: AiPackState?): AiPackStatus {
  return when (state?.status()) {
    AssetPackStatus.NOT_INSTALLED -> {
      AiPackStatus.NotInstalled
    }

    AssetPackStatus.WAITING_FOR_WIFI,
    AssetPackStatus.REQUIRES_USER_CONFIRMATION -> {
      AiPackStatus.RequestConfirmation
    }

    AssetPackStatus.PENDING,
    AssetPackStatus.DOWNLOADING,
    AssetPackStatus.TRANSFERRING -> {
      AiPackStatus.Downloading(state.transferProgressPercentage())
    }

    AssetPackStatus.COMPLETED -> {
      val assetLocation = getAssetFromAiPack(
        packName = AI_PACK_NAME,
        assetName = Model.TINYLLAMA_1_1B_CHAT_V1_0.modelName
      )
      if (assetLocation == null) {
        AiPackStatus.NotInstalled
      } else {
        AiPackStatus.Installed(location = assetLocation)
      }
    }

    AssetPackStatus.CANCELED,
    AssetPackStatus.FAILED -> {
      AiPackStatus.Failed(errorCode = state.errorCode())
    }

    else -> {
      AiPackStatus.Unknown
    }
  }
}

A few things to note for the above function:

1. If the pack isn’t installed yet, it displays the AiPackStatus.NotInstalled state.

2. The next block handles cases where user interaction is required:

   • WAITING_FOR_WIFI means the download is paused because the user is not on Wi-Fi.

   • REQUIRES_USER_CONFIRMATION means Google Play requires the user to explicitly approve the download of the requested asset.

   In either case, you map it to the RequestConfirmation state, which tells the UI to show a dialog.

3. AssetPackStatus.PENDING, AssetPackStatus.DOWNLOADING and AssetPackStatus.TRANSFERRING – these three states all indicate that the download is in progress. You map them to a single Downloading state and pass the progress percentage to the UI layer.

4. AssetPackStatus.COMPLETED is the “success” case. It means the pack is downloaded on the device. The Installed state uses getAssetFromAiPack() helper function to obtain the actual file path of the model. If for some reason the path isn’t available (e.g., the asset wasn’t downloaded properly), it’s treated as NotInstalled state.

5. If the download was CANCELED by the user or FAILED for some reason, the AiPackStatus.Failed state is displayed along with the errorCode, so you can log it or show a more specific message.

6. AiPackStatus.Unknown is the fallback state. It helps prevent crashes is case of any unexpected statuses.

Listening to Status Updates

Whenever you fetch an AI Pack or check status, listening to these status updates is made easy by the Play AI Delivery Library – you need to use the AiPackStateUpdateListener interface from the library.

For your convenience, the MainViewModel implements AiPackStateUpdateListener. All you need to do is as follows:

Attach the listener to aiPackManager within the init function of MainViewModel.

init {
  aiPackManager.registerListener(this)
}

Override the onStateUpdate() function to map AI Pack states to your own AiPackStatus.

override fun onStateUpdate(state: AiPackState) {
  Log.d(TAG, "onStateUpdate: $state")
  _aiPackStatus.value = mapAiPackStatus(state)
}

Accessing AI Packs

Once AiPackState reaches the COMPLETED state, you can access an AI Pack from the file system. Calling aiPackManager.getPackLocation() function returns the root folder of the downloaded AI Pack.

AI Packs (such as models) are stored in the assets directory within the root directory. You can retrieve the absolute path of a specific asset file as follows:

private fun getAssetFromAiPack(packName: String, assetName: String): String? {
  val aiPackLocation = aiPackManager.getPackLocation(packName)
  val assetsFolderPath = aiPackLocation?.assetsPath()
  val assetFile = File(assetsFolderPath, assetName)
  Log.d(TAG, "Asset path: ${assetFile.absolutePath}")

  return if (assetFile.exists()) assetFile.absolutePath else null
}

Here, assetFile.absolutePath is the full usable path for the MediaPipe ML inference engine to load the LLM Model.

[Optional] Step 3: Slaying the Fragmentation Dragon by Device Targeting

Remember the “last mile” theory at the beginning of this chapter? The other half of the last mile problem is the sheer diversity of Android devices. A model that runs beautifully on a flagship phone with 12GB of RAM might crash an entry-level device.

Historically, the options were limited – you could either ship a “one-size-fits-all” model that’s underpowered for high-end devices, or you could write a ton of custom logic to try and guess the device’s capabilities.

Play for On-device AI solves this with model variants and targeting. You can create different versions of your model optimized for different hardware, and Play will automatically deliver the right one. 

You can create targeting rules based on device properties like:

• Device RAM.

• System on Chip (SoC).

• Specific device models.

• System features.

You can enable Device Targeting with the following steps:

1. Enable the device targeting feature in the Android Gradle Plugin by adding this to gradle.properties file:

android.experimental.enableDeviceTargetingConfigApi=true

2. Create a device targeting configuration XML file. Define your custom device groups in that file. For example, you can create a device group called gemini_on_device that targets all devices powered by the Tensor G5 chip:

<config:device-targeting-config xmlns:config="http://schemas.android.com/apk/config">

  <config:device-group name="gemini_on_device">
      <config:device-selector>
          <config:system-on-chip manufacturer="Google" model="TensorG5"/>
      </config:device-selector>
  </config:device-group>

</config:device-targeting-config>

3. Include your device targeting configuration file in your app bundle by adding this to your main module’s build.gradle file:

android {
  ...
  bundle {
    deviceTargetingConfig = file('device_targeting_config.xml')
    deviceGroup {
      enableSplit = true   // split bundle by #group
      defaultGroup = "other"  // group used for standalone APKs
    }
  }
  ...
}

You can package all these variants into the app bundle, define your own targeting rules, and Play handles the rest. Every user gets the best possible experience their hardware can support, without us writing a single line of device-specific if/else logic!

Device Targeting is optional and not used in the sample project to allow it run on most configurations, even on the Emulators. For more details on Device Targeting, you may check official Device Targeting Configuration documentation.

Step 4: Testing On-Device AI Packs Locally

Testing on-demand features (such as AI Packs with models) traditionally requires uploading builds to the Google Play Console - a process that introduces significant latency into the development cycle, and that can be the trickiest part of the workflow.

Alternatively, you can generate APKs with AI Packs using bundletool and sideload them onto your device.

In this section, you’ll learn the step-by-step command-line operations required to build the application and deploy it for local AI Pack testing using bundletool.

Note: Before you begin, ensure you have Java Development Kit (JDK) and JetBrains Runtime (JBR) installed and configured on your machine.

What is Bundletool?

Bundletool is the underlying tool that Android Studio, the Android Gradle plugin, and Google Play use to build an Android App Bundle. It’s also available as a command-line tool.

Bundletool converts your app bundle into the APKs that would be installed on devices, replicating how Google Play generates and serves APKs — including those with AI Packs. This makes it possible to locally test and simulate the deployment process of your app with AI Packs before releasing it to users.

To begin with, download the latest release of bundletool (.jar). Ensure you have a test device or Emulator connected.

If you’re on Mac or Linux, you can install using this homebrew command:

brew install bundletool

The Core Workflow: Using Bundletool for On-Device Deployment

In this section, you’ll generate signed App Bundles or APKs, which require your keystore information – just like preparing an app for deployment to the Google Play Store.

You can initiate this process through the Android Studio menu:
Build > Generate Signed Bundle or APK...

You may create a new keystore by selecting Create new… and providing required information on the next screen as follows:

Or, if you already have an existing keystore file, simply enter your credentials and tap Next.

Once your keystore is configured, open the terminal and run the following commands step by step.

Note: The default debug.keystore will be used when you run these commands.

1. The workflow begins with creating the standard, release-signed Android App Bundle. You can use this terminal command (or use the Build menu) in the Android Studio:

./gradlew :app:bundleRelease

The resulting .aab file, typically located at app/build/outputs/bundle/release/app-release.aab directory, serves as the primary input for the subsequent bundletool commands.

2. Generate the local-testing APK set (.apks) with the command below:

bundletool build-apks \
--bundle=app/build/outputs/bundle/release/app-release.aab \
--output=app/build/outputs/bundle/release/app-release.apks \
--local-testing \
--connected-device

This command processes the .aab file and generates a set of split APKs. It also injects the necessary local testing metadata (using the --local-testing flag), and packages everything into a single app-release.apks output file.

Upon successful execution, the **outputs** folder will look like this:

3. Deploy the app and sideload the AI Pack assets with the below command:

bundletool install-apks --apks=app/build/outputs/bundle/release/app-release.apks

This command determines the required base and configuration APKs for the connected device and installs them.

Uninstall any previous version of the app from the device before executing this command, as the local-testing workflow does not support in-place updates. When complete, the application icon appears on the device, ready for testing.

Build and run the app and tap the Download Model button, verify that the download process and status updates work as expected.

Conclusion

Shipping on-device AI has always been a battle on two fronts: building the feature and building the infrastructure to deploy and manage it. Play for On-device AI effectively eliminates that second front.  

The end-to-end process can be summarized as follows:

1. Configure: Set up the AI Pack as a distinct Gradle module with an on-demand delivery type and place the model file in its assets directory.

2. Build: Generate a signed Android App Bundle using ./gradlew :app:bundleRelease.

3. Generate APKs: Use bundletool build-apks with the --local-testing flag to create .apks file.

4. Deploy: Uninstall any previous app version, then use bundletool install-apks to install the app and sideload AI Pack assets.

5. Test: Launch the app and verify that the code requests the pack via AiPackManager and accesses the model file using the path provided by AiPackLocation.

By handling the hosting, delivery, updating, and targeting of your models, the Play for On-device AI Library lets you offload the MLOps complexity to the Play Store. This means smaller app size, more optimized app updates, a better user experience tailored to each device, and most importantly, more time for you to focus on what you do best – building amazing, intelligent apps!