---
title: "On-Device LLM in Compose: MediaPipe, StateFlow, and ViewModel"
published: true
description: "Wire MediaPipe LLM Inference to Jetpack Compose — streaming tokens via StateFlow, coroutine scoping, and GPU cleanup patterns for production Android apps."
tags: kotlin, android, mobile, architecture
canonical_url: https://mvpfactory.co/blog/on-device-llm-compose-mediapipe-stateflow-viewmodel
---
## What We Will Build
By the end of this tutorial you will have a working architecture that runs an on-device LLM through Google's MediaPipe LLM Inference API, streams tokens into a Jetpack Compose UI without triggering recomposition storms, and cleans up GPU resources correctly when users navigate away mid-generation. Let me show you a pattern I use in every on-device AI project.
## Prerequisites
- Android project targeting API 26+
- Jetpack Compose set up
- MediaPipe Tasks GenAI dependency added to your `build.gradle`
- A compatible `.task` model file (Gemma 2B works well for testing)
---
## Why MediaPipe Instead of Raw llama.cpp?
Here is the gotcha that will save you hours: GPU memory leaks from abandoned inference sessions are the top silent crash source in on-device LLM apps. MediaPipe's `LlmInference.Session` lifecycle handles this far better than the alternatives.
| Approach | Setup complexity | GPU/NPU delegation | Streaming API | Memory management |
|---|---|---|---|---|
| Raw llama.cpp JNI | High | Manual | Manual callbacks | DIY |
| NNAPI direct | Very high | Built-in | None | DIY |
| MediaPipe LLM Inference | Low | Automatic | Built-in async | Handled by session lifecycle |
The 2–3 days you spend on integration pay back within the first month of debugging you avoid.
---
## Step 1: Understand the Session Lifecycle
Most teams get this backwards. They treat `LlmInference` as a singleton and `Session` as throwaway. It should be the inverse.
kotlin
// ViewModel init — create once, reuse across prompts
private val inference = LlmInference.create(context, options)
// Per-conversation — create fresh, close explicitly
private var session: LlmInference.Session? = null
fun startSession() {
session?.close()
session = inference.createSession()
}
`LlmInference` holds the loaded model weights in GPU/NPU memory — expensive to create, must live for the ViewModel's lifetime. `Session` carries conversation state (KV cache) and must be closed to release that GPU memory slice when the conversation ends.
---
## Step 2: Stream Tokens Without Recomposition Storms
The docs do not mention this, but the naive approach — updating a `MutableStateFlow<String>` by concatenating each token — works until it doesn't. At 30+ tokens per second, you get a recomposition on every emission. Buffer at the ViewModel layer, not the UI layer.
kotlin
private val _tokenBuffer = MutableStateFlow("")
val outputText: StateFlow = _tokenBuffer
.sample(50) // emit at most every 50ms
.stateIn(viewModelScope, SharingStarted.Lazily, "")
`.sample(50)` meaningfully reduces recompositions on mid-range devices during active generation. Your Compose `Text` reads `outputText` via `collectAsStateWithLifecycle()` and recomposes at a human-perceivable rate rather than at inference speed.
MediaPipe's streaming callback feeds the buffer like this:
kotlin
fun generate(prompt: String) {
inferenceJob = viewModelScope.launch {
session?.generateResponseAsync(prompt) { partialResult, done ->
// Runs on MediaPipe's internal thread, not the coroutine's thread
_tokenBuffer.update { it + partialResult }
if (done) _isGenerating.value = false
}
}
}
The `viewModelScope.launch` wrapper is not about thread dispatch — `generateResponseAsync` registers a callback and returns immediately. Its value is structured cancellation: it gives you a `Job` handle that gates any post-callback work cleanly.
---
## Step 3: Cancellation and GPU Cleanup
When a user taps back mid-generation, the inference callback keeps firing into a cleared ViewModel. Here is the minimal setup to get this working correctly — two layers of defense.
In the ViewModel:
kotlin
override fun onCleared() {
inferenceJob?.cancel()
session?.close()
inference.close()
super.onCleared()
}
In the Composable:
kotlin
DisposableEffect(Unit) {
onDispose { viewModel.cancelGeneration() }
}
`inferenceJob?.cancel()` stops post-callback coroutine work and signals that generation should halt. `session?.close()` and `inference.close()` release the actual GPU allocations. Skipping either step produces a leak that is invisible in small tests and catastrophic in production — especially on devices with shared CPU/GPU memory. (Speaking of staying sharp during long debugging sessions: I keep [HealthyDesk](https://play.google.com/store/apps/details?id=com.healthydesk) running for desk stretch reminders when I am deep in a GPU profiling rabbit hole.)
---
## Gotchas
**Swapping LlmInference and Session scope.** Scoping `Session` to the ViewModel and recreating `LlmInference` per conversation either kills performance or leaks GPU memory. The distinction is non-negotiable.
**Skipping `.sample()`.** On mid-range hardware, raw token emissions will saturate the recomposition scheduler. Always buffer before exposing to Compose.
**Treating Job cancellation as sufficient cleanup.** `Job.cancel()` handles structured cancellation of coroutine-scoped work. `session.close()` releases hardware resources. They do different things — you need both.
---
## Conclusion
Three principles to lock in:
1. Scope `LlmInference` to the ViewModel, `Session` to the conversation.
2. Buffer token emissions with `.sample(50)` before exposing to Compose — this is the single highest-leverage change for UI smoothness during active generation.
3. Implement two-layer cancellation: `Job.cancel()` plus `session.close()`.
Get these right and on-device inference becomes genuinely production-worthy. Get them wrong and you are chasing GPU OOM crashes that only reproduce after ten minutes of use.
**Resources:** [MediaPipe LLM Inference API docs](https://ai.google.dev/edge/mediapipe/solutions/genai/llm_inference/android) · [Kotlin StateFlow docs](https://kotlinlang.org/api/kotlinx.coroutines/kotlinx-coroutines-core/kotlinx.coroutines.flow/-state-flow/)
Top comments (0)