# Main Architectural Components of the Google AI Edge Gallery App: A Deep Dive into the Android Codebase

> Explore the Google AI Edge Gallery app's main architectural components. Dive into its layered, modular design leveraging Jetpack Compose, Hilt, and a custom task framework for seamless AI integration. Understand the Android cod...

- Repository: [google-ai-edge/gallery](https://github.com/google-ai-edge/gallery)
- Tags: architecture
- Published: 2026-04-06

---

**The Google AI Edge Gallery app uses a layered, modular architecture built on Jetpack Compose, Hilt dependency injection, and a plugin-style custom-task framework that enables new AI capabilities without modifying core code.**

The `google-ai-edge/gallery` repository demonstrates a production-ready approach to on-device AI inference on Android. Understanding the main architectural components of the Google AI Edge Gallery app reveals how modern Android architecture components—Jetpack Compose, DataStore, and WorkManager—combine with a flexible plugin system to support diverse AI tasks, from large language model chat to specialized inference pipelines, while maintaining strict separation of concerns.

## Application Bootstrap and Entry Points

The app initializes through two primary entry points that establish the runtime environment.

**[`GalleryApplication.kt`](https://github.com/google-ai-edge/gallery/blob/main/GalleryApplication.kt)** serves as the Hilt entry point, responsible for application-wide initialization. It loads the persisted theme from `DataStoreRepository`, initializes Firebase for analytics, and configures the dependency injection graph before any UI renders.

**[`MainActivity.kt`](https://github.com/google-ai-edge/gallery/blob/main/MainActivity.kt)** handles the Android-specific lifecycle. It displays the system splash screen, then calls `setContent` to mount the Jetpack Compose UI tree. This activity acts as the host for the entire navigation graph and theme system, immediately delegating to the Compose layer once bootstrap completes.

## Navigation and UI Layer

Navigation is declarative and type-safe, following modern Compose conventions.

**[`GalleryNavGraph.kt`](https://github.com/google-ai-edge/gallery/blob/main/GalleryNavGraph.kt)** defines all top-level destinations—including home, model list, model detail pages, benchmarking screens, and custom-task screens—within a single navigation graph. It also manages deep-link routing, allowing external intents to launch specific AI tasks directly. The graph creates a `NavController` that coordinates transitions between the model management interface and individual task implementations.

UI components observe state through `ModelManagerViewModel.uiState`, a `StateFlow` that emits updates when download progress, initialization status, or configuration changes occur. Composables consume this state via `collectAsState`, ensuring the interface reacts instantly to background operations.

## State Management and Business Logic

**[`ModelManagerViewModel.kt`](https://github.com/google-ai-edge/gallery/blob/main/ModelManagerViewModel.kt)** functions as the central nervous system of the app. It encapsulates all business logic for model discovery, download orchestration, initialization, and cleanup. The view model exposes methods such as `downloadModel()`, `initializeModel()`, and `cleanupModel()` that coordinate between the UI and background workers.

Key responsibilities include:
- Fetching and parsing the model allow-list JSON from remote GitHub-hosted files or local test assets
- Maintaining the in-memory graph mapping `Task` objects to their associated `Model` instances
- Managing authentication tokens and user-specific settings
- Triggering Firebase analytics events through [`Analytics.kt`](https://github.com/google-ai-edge/gallery/blob/main/Analytics.kt)

## Data Persistence and Dependency Injection

The data layer combines Proto DataStore for structured persistence with WorkManager for background downloads.

**[`DataStoreRepository.kt`](https://github.com/google-ai-edge/gallery/blob/main/DataStoreRepository.kt)** persists user preferences—including theme settings, imported models, and authentication tokens—using Proto DataStore with `SettingsSerializer`. This provides type-safe, asynchronous storage that survives process death.

**[`DownloadRepository.kt`](https://github.com/google-ai-edge/gallery/blob/main/DownloadRepository.kt)** wraps Android’s WorkManager to handle large model file downloads in the background, reporting progress back to the view model through Kotlin coroutines.

**[`AppModule.kt`](https://github.com/google-ai-edge/gallery/blob/main/AppModule.kt)** provides the Hilt dependency injection bindings, exposing singletons for the DataStore instance, repository implementations, and protocol buffer serializers. This modular approach allows easy swapping of implementations for testing or different build variants.

## The Custom-Task Plugin Framework

The architecture’s most distinctive feature is its **plugin-style custom-task framework**, located in the `customtasks/` package. This system allows developers to add new AI capabilities without touching the core navigation or view model code.

**[`CustomTask.kt`](https://github.com/google-ai-edge/gallery/blob/main/CustomTask.kt)** defines the interface contract that every plugin must implement:
- `initializeModelFn`: A lambda that creates and configures the concrete model helper instance
- `cleanUpModelFn`: A lambda releasing native resources when the user navigates away
- Composable UI functions rendering the task-specific interface

Each capability—such as [`TinyGardenTask.kt`](https://github.com/google-ai-edge/gallery/blob/main/TinyGardenTask.kt) for generative audio or LLM chat implementations—lives in its own module, registering automatically with Hilt’s multibinding system. The `ModelManagerViewModel` discovers these tasks at startup by scanning the `customtasks/` package and matching them against the allow-list JSON definitions.

## Model Runtime Abstraction

AI inference follows an abstract contract pattern that decouples the UI from specific runtime implementations.

**[`LlmModelHelper.kt`](https://github.com/google-ai-edge/gallery/blob/main/LlmModelHelper.kt)** defines the base interface for model operations, specifying methods for `initialize()`, `runInference()`, and `cleanUp()`. Concrete implementations—such as [`LlmChatModelHelper.kt`](https://github.com/google-ai-edge/gallery/blob/main/LlmChatModelHelper.kt) in the `runtime/` folder—provide the actual native integration with on-device inference engines.

The **Model** class ([`Model.kt`](https://github.com/google-ai-edge/gallery/blob/main/Model.kt)) represents a downloadable AI asset, tracking its download status, configuration parameters (max tokens, accelerators), and the active runtime instance stored in the `instance` field. When a user selects a model, the view model binds the appropriate helper to this field, enabling type-safe inference calls throughout the UI layer.

## Model Lifecycle and Allow-List Management

The app dynamically populates its catalog through an external allow-list system rather than hardcoded model definitions.

At startup, `ModelManagerViewModel` calls `loadModelAllowlist()` to fetch JSON files hosted under the `model_allowlists/` directory in the repository. The `processTasks()` method parses these definitions, mapping model files to their supported tasks and loading per-model configurations such as token limits and hardware acceleration preferences.

This architecture enables the app to support new models immediately upon JSON updates, without requiring a full app redeployment.

## Key Implementation Patterns

### Starting a Model Download

To trigger a download from any composable:

```kotlin
val viewModel: ModelManagerViewModel = hiltViewModel()
val task = viewModel.getTaskById("llm_chat")
val model = task?.models?.firstOrNull { it.name == "gemma_2b" }

Button(onClick = {
    model?.let { viewModel.downloadModel(task = task, model = it) }
}) {
    Text("Download ${model?.displayName}")
}

```

> Source: `ModelManagerViewModel.downloadModel` in [`Android/src/app/src/main/java/com/google/ai/edge/gallery/ui/modelmanager/ModelManagerViewModel.kt`](https://github.com/google-ai-edge/gallery/blob/main/Android/src/app/src/main/java/com/google/ai/edge/gallery/ui/modelmanager/ModelManagerViewModel.kt)

### Initializing a Model After Download

Once download status becomes `SUCCEEDED`, initialization proceeds through the custom task:

```kotlin
val model = viewModel.getSelectedModel()
val task = viewModel.getTaskById("llm_chat")
if (model != null && task != null) {
    viewModel.initializeModel(
        context = LocalContext.current,
        task = task,
        model = model,
        onDone = { Log.d("Gallery", "Model ready for inference") }
    )
}

```

> Source: `ModelManagerViewModel.initializeModel` in [`Android/src/app/src/main/java/com/google/ai/edge/gallery/ui/modelmanager/ModelManagerViewModel.kt`](https://github.com/google-ai-edge/gallery/blob/main/Android/src/app/src/main/java/com/google/ai/edge/gallery/ui/modelmanager/ModelManagerViewModel.kt)

### Running Inference

The helper abstraction enables uniform inference across different model types:

```kotlin
val helper = model.instance as? LlmModelHelper
helper?.runInference(
    model = model,
    input = "Explain quantum computing in two sentences.",
    resultListener = { partial, done, _ ->
        // Append partial to UI state, complete when done == true
    }
)

```

> Source: `LlmModelHelper.runInference` in [`Android/src/app/src/main/java/com/google/ai/edge/gallery/runtime/LlmModelHelper.kt`](https://github.com/google-ai-edge/gallery/blob/main/Android/src/app/src/main/java/com/google/ai/edge/gallery/runtime/LlmModelHelper.kt)

### Adding a New Custom Task

To extend the app with a new AI capability:

```kotlin
class MyNewTask @Inject constructor() : CustomTask {
    override val task = Task(id = "my_task", displayName = "My AI Feature")
    
    override val initializeModelFn = { ctx, scope, model, onDone ->
        model.instance = MyCustomHelper(ctx)
        onDone()
    }
    
    override val cleanUpModelFn = { ctx, scope, model, onDone ->
        (model.instance as? MyCustomHelper)?.close()
        onDone()
    }
    
    @Composable
    override fun MainScreen(data: CustomTaskData) {
        // Task-specific UI implementation
    }
}

```

> Source: [`CustomTask.kt`](https://github.com/google-ai-edge/gallery/blob/main/CustomTask.kt) in [`Android/src/app/src/main/java/com/google/ai/edge/gallery/customtasks/common/CustomTask.kt`](https://github.com/google-ai-edge/gallery/blob/main/Android/src/app/src/main/java/com/google/ai/edge/gallery/customtasks/common/CustomTask.kt)

## Summary

- **Layered architecture**: The app separates concerns across application bootstrap ([`GalleryApplication.kt`](https://github.com/google-ai-edge/gallery/blob/main/GalleryApplication.kt)), UI hosting ([`MainActivity.kt`](https://github.com/google-ai-edge/gallery/blob/main/MainActivity.kt)), navigation ([`GalleryNavGraph.kt`](https://github.com/google-ai-edge/gallery/blob/main/GalleryNavGraph.kt)), business logic ([`ModelManagerViewModel.kt`](https://github.com/google-ai-edge/gallery/blob/main/ModelManagerViewModel.kt)), and data persistence ([`DataStoreRepository.kt`](https://github.com/google-ai-edge/gallery/blob/main/DataStoreRepository.kt), [`DownloadRepository.kt`](https://github.com/google-ai-edge/gallery/blob/main/DownloadRepository.kt)).
- **Plugin extensibility**: The `CustomTask` interface in the `customtasks/` folder enables new AI capabilities without core code modifications, using Hilt multibinding for automatic discovery.
- **Runtime abstraction**: [`LlmModelHelper.kt`](https://github.com/google-ai-edge/gallery/blob/main/LlmModelHelper.kt) provides a uniform contract for model initialization and inference, with concrete implementations handling specific engine integrations.
- **Dynamic model management**: The allow-list JSON system (`model_allowlists/`) allows remote configuration of available models and their supported tasks, fetched at runtime by `ModelManagerViewModel`.
- **Modern Android stack**: The implementation leverages Jetpack Compose for UI, Hilt for dependency injection, DataStore for persistence, and WorkManager for background downloads.

## Frequently Asked Questions

### How does the Google AI Edge Gallery app handle model downloads without blocking the UI?

The app delegates downloads to [`DownloadRepository.kt`](https://github.com/google-ai-edge/gallery/blob/main/DownloadRepository.kt), which wraps Android’s WorkManager API to perform file transfers on background threads. The `ModelManagerViewModel` exposes download progress through a `StateFlow` that composables observe via `collectAsState()`, ensuring the UI remains responsive while displaying real-time progress indicators.

### What is the purpose of the CustomTask interface in the architecture?

The `CustomTask` interface defines a plugin contract that standardizes how new AI capabilities integrate with the app. It specifies initialization callbacks (`initializeModelFn`), cleanup callbacks (`cleanUpModelFn`), and UI composables, allowing developers to add features like Tiny Garden or Mobile Actions by implementing the interface and placing the file in the `customtasks/` folder, without modifying the core navigation or view model logic.

### Where does the app store user preferences and theme settings?

User preferences—including dark mode overrides, imported models, and authentication tokens—persist through [`DataStoreRepository.kt`](https://github.com/google-ai-edge/gallery/blob/main/DataStoreRepository.kt), which uses Proto DataStore with [`SettingsSerializer.kt`](https://github.com/google-ai-edge/gallery/blob/main/SettingsSerializer.kt). This provides type-safe, asynchronous storage that survives process termination, with changes immediately reflected in the UI through [`ThemeSettings.kt`](https://github.com/google-ai-edge/gallery/blob/main/ThemeSettings.kt).

### How are available AI models populated in the app at startup?

[`ModelManagerViewModel.kt`](https://github.com/google-ai-edge/gallery/blob/main/ModelManagerViewModel.kt) calls `loadModelAllowlist()` during initialization to fetch JSON files from the remote `model_allowlists/` directory in the repository. The `processTasks()` method parses these definitions to build the in-memory graph of available tasks and their associated models, enabling dynamic model catalog updates without app redeployment.