5. Custom ML with MediaPipe

Written by Zahidur Rahman Faisal

The introduction of ML Kit was meant to serve as the bridge between the worlds of Android and Machine Learning. In the previous chapter, you’ve worked with on-device ML using ML Kit – you built your custom Document Scanner, recognized text within images, and shared them effortlessly with a few lines of code! It feels powerful, right? ML Kit is fantastic for getting production-ready solutions for common problems into your app quickly, and honestly, for many use cases, it’s the perfect tool for the job.  

But sometimes, you need more than what ML Kit currently offers.

Maybe you want to build a cool LLM-based chat that works offline. Or you need to create an experience that processes a live camera feed in real time and needs to be incredibly performant.

That’s the moment you graduate from ML Kit to MediaPipe.

MediaPipe: A Complete Toolset for Custom Machine Learning Solutions

ML Kit gives you a set of specialized tools, whereas MediaPipe gives you the entire workshop! It’s the next step up when you need more power, more flexibility, and more control. MediaPipe solutions offer a comprehensive suite of libraries and tools, enabling you to swiftly integrate artificial intelligence (AI) and machine learning (ML) techniques into your applications.

MediaPipe provides two main resources to empower your intelligent apps:

• MediaPipe Tasks: Cross-platform APIs and libraries that make it easy to deploy and integrate ML solutions into your applications. Learn more.

• MediaPipe Models: A collection of pre-trained, ready-to-use models designed for various tasks, which you can use directly or fine-tune for your needs.

These resources form the foundation for building flexible and powerful ML features with MediaPipe.

The tools below enable you to use these Tasks and Models for your custom ML solutions:

• MediaPipe Model Maker: This is your entry point into the world of custom models. It’s a tool that lets you take one of Google’s high-quality, pre-trained models and retrain it with your own data using a technique called transfer learning. You don’t need to be an ML expert; you just need a good dataset.

The output of Model Maker is a TensorFlow Lite .tflite file, which you’ll need to convert into a MediaPipe-specific .task file. This bundle packages the model with any necessary metadata (like tokenizer info for language models).

You’ll integrate this custom .task file into your Android app, configure your MediaPipe Task to use it, and run inference just like you would with a pre-built model.

Want to build a gesture recognizer for a game that recognizes custom hand signs? Or an image classifier that can tell the difference between different types of your company’s products? Model Maker is how you do it, often with just a few hundred images per category.

• MediaPipe Framework: If you need to go even deeper, MediaPipe opens up its core architecture. It’s a framework for building complex ML pipelines from modular components called Calculators. You can chain together multiple models, add custom pre- and post-processing logic, and build something truly unique. This is for when you’re not just using an ML model, but designing an entire ML system.

Why MediaPipe?

Let’s break down why you’d switch to MediaPipe instead of using ML Kit.

When “Good Enough” Isn’t Custom Enough

ML Kit is excellent for common tasks because it uses models trained on general data. But what if your app needs to be something more specific, more granular? This is MediaPipe’s killer feature: Customization

ML Kit lets you use a custom TensorFlow Lite model, but MediaPipe is designed from the ground up to make building, customizing, and deploying those models a core part of the workflow.

When Every Millisecond Counts: Real-Time Performance

ML Kit’s on-device models are optimized for mobile, but MediaPipe is in a league of its own when it comes to processing live and streaming media. Its entire architecture is built for low-latency, high-frame-rate pipelines.

This is crucial for features like:

• Real-time chat: Large Language Model (LLM) applications that need to traverse through the language patterns and come back with incremental results.

• Augmented Reality (AR) effects: Applying filters to a face in a live camera feed.

• Fitness and physical therapy apps: Analyzing a user’s posture in real-time to give feedback on their form.  

• Gesture-based controls: Using hand movements to interact with your app’s UI.

MediaPipe achieves this through end-to-end hardware acceleration, making excellent use of the device’s GPU to handle the heavy lifting of both ML inference and video processing. When you’re processing a continuous video stream, this focus on performance is the difference between a smooth, magical experience and a laggy, frustrating one.

When Your App Lives Beyond Android

This is a big reason. ML Kit is fantastic for native mobile development on Android and iOS. But what happens when your team wants to launch a web version of your app?  

MediaPipe is a cross-platform framework. You can build your ML pipeline once and deploy it everywhere: Android, iOS, web, desktop, and even IoT devices. The APIs are designed to be consistent across platforms, meaning you can reuse a lot of your logic and don’t have to start from scratch for each new platform you support.  

This build once, deploy anywhere philosophy is a massive advantage for teams that need to maintain a consistent user experience across different ecosystems.

When You Want to Live on the Cutting Edge

As MediaPipe is a more flexible and open framework, it’s often the place where you’ll first see support for more advanced and experimental on-device tasks, especially in the realm of generative AI.

The MediaPipe LLM Inference API, for example, allows you to run powerful Large Language Models like Gemma and Phi-2 entirely on-device. This unlocks features like offline chatbots, text summarization, and multimodal understanding (processing text, images, and audio together) - all private, fast, and free of server costs.

While ML Kit is now getting its own on-device GenAI APIs powered by Gemini Nano, MediaPipe often provides a more direct and customizable path for developers who want to experiment with a wider variety of open models and build more complex generative features.

Building Your First On-device LLM App

Remember the “Cat Breeds” app you built in chapter 2? What if you could chat with a veterinary specialist and ask about cats? Cool right? Let’s build that with an on-device LLM using MediaPipe!

Adding the LLM Inference API

The LLM Inference API enables Android apps to run large language models (LLMs) entirely on-device. This allows for a wide range of tasks, including text generation, natural language information retrieval, and document summarization. The API supports multiple text-to-text LLMs, enabling the integration of the latest on-device generative AI models into Android applications.

The LLM Inference API offers several key features:

1. Text-to-text generation: Produce relevant text responses from a given input prompt, enabling features like chat, summarization, and question answering.

2. Model flexibility: Choose from multiple LLMs to best suit your application’s needs. You can also fine-tune models with your own data or apply custom weights for specialized tasks.

3. LoRA integration: Enhance and personalize LLM performance using LoRA adapters—either by training on your own datasets or leveraging open-source LoRA models. (Note: LoRA is not supported for models converted via the AI Edge Torch Generative API.)

The LLM Inference API uses the com.google.mediapipe:tasks-genai library. To start using the API, launch the starter project and open the app-level build.gradle file.

Then, add the following code at the end of the dependencies section:

dependencies {
  ... ... ...

  // TUTORIAL DEPENDENCIES
  var MediaPipeVersion = "0.10.27"
  implementation "com.google.mediapipe:tasks-genai:$MediaPipeVersion"
}

Build and run the app, and watch the LogCat console in the IDE. You’ll notice a “Model Not Found” exception.

This occurs because the app is trying to initialize a model that isn’t available yet. You need to ensure the model is downloaded and placed in the correct path so it can be initialized properly.

Adding the Model

Add a language model to your test device via your computer before initializing the LLM Inference API. Run the following command in your terminal to check which devices are connected to your machine:

$ adb devices

You will see a list of all connected and recognized devices similar to the following output:

List of devices attached
ZRF198804FEBBD  device
emulator-5554   device    

In the list above, ZRF198804FEBBD and emulator-5554 are example device_id values. You’ll need the device_id of your connected test device when running upcoming ADB commands.

You can use a pre-converted model from the HuggingFace LiteRT Community, or convert a model to a MediaPipe-compatible format with the AI Edge Torch Generative Converter.

For this project, you’ll use TinyLlama/TinyLlama_v1.1, a lightweight language model containing approximately 1 billion parameters. Download the model from here.

Note: The model requires approximately 1.15 GB of storage and may take some time to download.

Once the model is downloaded, you need to add the model to your device through adb commands:

$ adb -s <device_id> shell mkdir -p /data/local/tmp/llm/
$ adb -s <device_id> push <model_download_path> /data/local/tmp/llm/TinyLlama-1.1B-Chat-v1.0_multi-prefill-seq_q8_ekv1280.task

1. The first line creates a folder on the connected device associated with the device_id to store the model.
2. The second line pushes the model from <model_download_path>, located on your computer, to the given folder on your test device.

Note: The <model_download_path> is relative. If you’re on Mac OS, it may look like: /Users//Downloads/TinyLlama-1.1B-Chat-v1.0_multi-prefill-seq_q8_ekv1280.task

Once those commands are successfully executed, you’ll see a confirmation in the terminal similar to the following:

.../TinyLlama-1.1B-Chat-v1.0_multi-prefill-seq_q8_ekv1280.task: 1 file pushed, 0 skipped. 361.2 MB/s (1148331545 bytes in 3.032s)

Alternatively, you can use the Device File Explorer in Android Studio to create the folders manually and drag the model file onto your device.

Build and run now - the app will compile and run without exceptions, and you’re ready to start using the model!

Initializing the Inference

Go to the InferenceManager class in the com.kodeco.android.aam.llm package in the started project. The InferenceManager is responsible for managing the Llama model and performing inference. It handles loading the model, creating an inference session, and generating responses based on user prompts. To do so, InferenceManager relies on two key objects:

1. llmInference: An instance of the LlmInference class. This is the primary class for interacting with the LLM Inference API. It creates the LLM engine with the correct model path, maximum number of tokens, and preferred backend.

2. llmInferenceSession: Creates an LLM inference session. It’s responsible for configuring the session with provided options, running the session, and then closing the session once you’re done.

Creating the LLM Inference Engine

First, initialize llmInference by adding the below function:

private fun createEngine(context: Context) {
  val inferenceOptions = LlmInference.LlmInferenceOptions.builder()
    .setModelPath(modelPath(context))
    .setMaxTokens(MAX_TOKENS)
    .setPreferredBackend(LLM_MODEL.llmInferenceBackend)
    .build()

  try {
    llmInference = LlmInference.createFromOptions(context, inferenceOptions)
  } catch (e: Exception) {
    Log.e(TAG, "Failed to load model. Error: ${e.message}", e)
  }
}

This function first prepares a configuration set named inferenceOptions for the LLM inference engine, then creates the engine using that configuration.

In inferenceOptions, you provide:

• The path where the model is stored.
• The maximum number of tokens (input tokens + output tokens) the model handles.
• Preferred backend – that could be your device’s CPU or GPU, depending on the computational load of your task.

The configurations for the model you’re using in this project are defined in the Model.TINYLLAMA_1_1B_CHAT_V1_0 enum inside the com.kodeco.android.aam.llm package.

Creating Your LLM Inference Session

Next is initializing llmInferenceSession. Add the below createSession() function:

private fun createSession() {
  val sessionOptions = LlmInferenceSessionOptions.builder()
    .setTemperature(LLM_MODEL.temperature)
    .setTopK(LLM_MODEL.topK)
    .setTopP(LLM_MODEL.topP)
    .build()

  try {
    llmInferenceSession = LlmInferenceSession.createFromOptions(llmInference, sessionOptions)
  } catch (e: Exception) {
    Log.e(TAG, "Failed to create LlmInference session. Error: ${e.message}", e)
  }
}

This function creates and configures the LLM inference session with parameters like temperature, top-k, and top-p in the LLM_MODEL object, which control the randomness and creativity of generated responses.

The following configuration options are available when you set up a LLM inference session:

Now, update the init block as follows to initialize your first inference session:

init {
  if (!modelExists(context)) {
    throw IllegalArgumentException("Model not found at path: ${LLM_MODEL.path}")
  }
  createEngine(context)
  createSession()
}

Do not forget to add necessary imports for all those changes.

Build and run the app. You should be able to compile without any error and see the Chat screen, but that’s not functional yet!

Streaming Responses

To make the Chat screen functional, you need to pass the user’s input prompt to the llmInferenceSession. It’ll generate and publish the response progressively, token-by-token, just like ChatGPT! To achieve this, you’ll need to attach a ProgressListener to it.

Add generateResponseAsync() function to InferenceManager class as follows:

fun generateResponseAsync(
  prompt: String,
  progressListener: ProgressListener<String>
): ListenableFuture<String> {
  llmInferenceSession.addQueryChunk(prompt)
  return llmInferenceSession.generateResponseAsync(progressListener)
}

The above function generates a response from the LLM. It does two things:

1. Takes the user’s prompt as input and adds it to the llmInferenceSession as a query.

2. Calls generateResponseAsync() with a ProgressListener. The ProgressListener will be called with chunks of response as they are generated — perfect for streaming updates to the UI.

The generateResponseAsync() function has to be executed when user sends a message from the UI layer. That part is handled by ChatViewModel.

Open ChatViewModel.kt inside the chat package and update the sendMessage() function:

fun sendMessage(userMessage: String) {
  viewModelScope.launch(Dispatchers.IO) {
    _uiState.value.addMessage(userMessage, USER_PREFIX)
    _uiState.value.createLoadingMessage()
    _isInputEnabled.value = false

    // Generating Async Response 
    try {
      inferenceManager.generateResponseAsync(prompt = userMessage) { progressiveResult, done ->
        _uiState.value.appendMessage(text = progressiveResult)
        _tokensRemaining.update { max(0, it - 1) }
        _isInputEnabled.value = done
      }
    } catch (e: Exception) {
      _uiState.value.addMessage(e.localizedMessage ?: "Error", MODEL_PREFIX)
    }
  }
}

These changes under the “Generating Async Response” block are self-explanatory; they execute as follows:

1. Asks inferenceManager to generate a response asynchronously, passing the prompt from the user.

2. Updates the UI once a progressiveResult becomes available. The UI restricts further input until the process is done.

3. If there’s any error, it’ll publish the error to the UI layer.

Include any required imports for the above updates. You’re now ready to test the LLM - Build and run, and then ask about cats. Try limiting the word count to keep the context window short with queries like this:

“Tell me about different types of cats in 100 words”

Estimating Remaining Tokens

You must have noticed the “0 tokens remaining” message in the chat screen right after sending your first query to the LLM. Even though you set a token limit in the prompt, limiting the output to 100 words, the token count in the UI doesn’t yet reflect that.

This happens because you haven’t implemented token estimation for the current chat session — that’s what you’ll do in this section.

Open InferenceManager class again and add the below function to estimate remaining tokens for the current chat session:

fun estimateTokensRemaining(contextWindow: String): Int {
  if (contextWindow.isEmpty()) return -1

  val sizeOfAllMessages = llmInferenceSession.sizeInTokens(contextWindow)
  val remainingTokens = MAX_TOKENS - sizeOfAllMessages

  return max(0, remainingTokens)
}

In the above method, you’re passing the contextWindow from ViewModel. The total context window will always consist of previous messages (context), added up with the latest prompt sent from the user.

The llmInferenceSession.sizeInTokens() function estimates the number of tokens used in the contextWindow and returns an integer representing the total tokens consumed.

The remainingTokens value is calculated by subtracting that size from MAX_TOKENS. The UI shouldn’t display any value beyond 0, hence max(0, remainingTokens) is used when returning the value.

Calculating Context Window

So, how does the ViewModel calculate context window? You’ll find out as you implement it! Open ChatViewModel and add the below function:

fun calculateRemainingTokens(prompt: String) {
  val contextWindow = uiState.value.messages.joinToString { it.rawMessage } + prompt
  val remainingTokens = inferenceManager.estimateTokensRemaining(contextWindow = contextWindow)
  _tokensRemaining.value = remainingTokens
}

The first line that calculates contextWindow is crucial. uiState.value.messages returns a list of ChatMessage objects in the current session. ChatMessage is a data class representing the state of each chat bubble displayed on the UI. The rawMessage property of a ChatMessage object holds the actual text content (token data).

The joinToString { it.rawMessage } + prompt part combines all previous messages with the user’s current prompt to form the complete conversation context.

The rest is simple, you call inferenceManager.estimateTokensRemaining() function passing the context window – that gives the number of remaining tokens. Then you update the UI with the new value.

Now open ChatRoute in ChatScreen.kt file and update the onChangedMessage lambda as follows:

onChangedMessage = { message ->
  viewModel.calculateRemainingTokens(message)
}

Now, Build and run the app. Enter the same prompt in the chat — “Tell me about different types of cats in 100 words.” You’ll notice that the token count updates dynamically with each interaction.

Resetting the Session

Did you try tapping the Reset button? If so, you may have noticed it’s not functional yet. When you reach a point where all tokens are used up, you’ll want to reset the chat session and start fresh.

Since InferenceManager controls llmInferenceSession, add the below function in the class:

fun resetSession() {
  llmInferenceSession.close()
  createSession()
}

Pretty simple, yet powerful! This function resets the LLM inference session by closing the current session and creating a new one. It’s useful for clearing the chat history (and the token count) to start a new conversation.

Now, the UI must also reset when the session is restarted. The ChatViewModel handles the logic. Open ChatViewModel, add the function below:

fun resetChat() {
  uiState.value.clearMessages()
  calculateRemainingTokens(prompt = "")
  inferenceManager.resetSession()
}

Here’s how resetChat() works:

1. Clears all existing messages from uiState.

2. Resets the remaining token count to its initial value.

3. Calls inferenceManager.resetSession() to restart the LLM inference session.

To execute resetChat() when the user clicks the Reset button, open ChatRoute and update onClickReset lambda like this:

onClickReset = {
  viewModel.resetChat()
}

Build and run the app once more. Try chatting, then hit Reset when your token count reaches to 0 – your LLM Chat app is now fully functional!

Conclusion: ML Kit vs. MediaPipe

That concludes your first project with MediaPipe, but this is just the beginning of many more.

So, when should you use which? Here’s a simple rule of thumb:

Stick with ML Kit when:

• You need a fast, reliable solution for a common, well-defined problem (e.g., text recognition, barcode scanning).

• You have limited or no ML expertise and prefer a simple, “plug-and-play” API.

• Your focus is on Android and iOS, with no immediate plans to expand to web or desktop.

Level up to MediaPipe when:

• You want to recognize something unique and train a custom model using your own dataset.

• Your feature depends on high-performance, real-time processing from a camera or microphone stream.

• You’re building a cross-platform application and want to reuse ML logic seamlessly across web, iOS, and desktop.

• You aim to design complex, multi-stage ML pipelines or experiment with the latest on-device Generative AI models.

ML Kit is the perfect tool to get your feet wet with on-device ML. It’s powerful and incredibly easy to use. But when you’re ready to remove the training wheels and build something truly custom, performant, and platform-agnostic, MediaPipe is ready for you. It’s the natural next step in your journey as an on-device AI developer.