How the Gallery Model Management System Handles Downloading and Caching ML Models

The Google AI Edge Gallery implements a three-layer architecture that leverages WorkManager for background downloads, supports resume capability via HTTP Range headers, and caches models in the app's external files directory with automatic ZIP extraction.

The google-ai-edge/gallery repository provides a complete Android reference implementation for on-device ML inference. Its model management system handles downloading and caching ML models through a robust pipeline that ensures large model files are acquired reliably, even across process restarts or interrupted network conditions.

Three-Layer Architecture for Model Acquisition

The Gallery app organizes model acquisition into three distinct layers, each handling specific responsibilities from metadata retrieval to byte-level storage.

Layer 1: Model Metadata and Allow-list Loading

The process begins in android/src/app/src/main/java/com/google/ai/edge/gallery/ui/modelmanager/ModelManagerViewModel.kt, where the system loads a JSON allow-list describing every downloadable model. The loadModelAllowlist() function fetches this catalog from GitHub's model_allowlists/ directory or falls back to a locally cached version. Each entry maps to a Model data class (defined in Model.kt) containing critical metadata including download URLs, file sizes, version identifiers, required accelerators, and whether the asset requires ZIP extraction.

Layer 2: Download Orchestration with WorkManager

When a user initiates a download, DefaultDownloadRepository.downloadModel() creates a OneTimeWorkRequest<DownloadWorker> tagged with the unique model name. The request carries all necessary parameters—URL, authentication tokens, ZIP flags, and total size—via KEY_MODEL_* constants. WorkManager enqueues this as an expedited job, ensuring the download receives system resources even when the app is in the background.

Layer 3: Worker-Level Download and Cache Management

The DownloadWorker class, extending CoroutineWorker, executes the actual byte streaming in android/src/app/src/main/java/com/google/ai/edge/gallery/worker/DownloadWorker.kt. It writes to the app's external files directory (context.getExternalFilesDir(null)), specifically to a path structured as <model-normalized-name>/<version>/. The worker publishes foreground notifications to keep the process alive and reports real-time progress through WorkInfo.progress LiveData, which the ViewModel observes to update the UI with bytes received and transfer rates.

Resume Capability and Partial Download Handling

The system supports resumable downloads through HTTP Range headers. When DownloadWorker.doWork() executes, it checks for existing files ending in .$TMP_FILE_EXT (typically .tmp). If a partial file exists, the worker calculates the byte offset and adds a Range header to the HTTP request, allowing the download to continue from the interruption point rather than restarting.

The getModelDownloadStatus() method in ModelManagerViewModel determines the initial state by inspecting the external files directory. It returns PARTIALLY_DOWNLOADED if a .tmp file exists, NOT_DOWNLOADED if no files are present, or a success status if the final model file is already cached. This enables the UI to display appropriate actions—resume, cancel, or install—based on the current filesystem state.

Cache Storage Strategy and ZIP Extraction

Downloaded models reside under /sdcard/Android/data/com.google.ai.edge.gallery/files/, organized by normalized model name and version. During active downloads, files carry the .tmp extension; upon completion, the worker renames the file to remove this extension. For models distributed as ZIP archives (isZip = true), the worker extracts contents to a dedicated unzipDir subdirectory and deletes the original archive to conserve storage.

Imported local models bypass this pipeline entirely, storing assets in IMPORTS_DIR with the model.imported flag set to true. All download start timestamps persist in SharedPreferences under download_start_time_ms, enabling analytics to calculate total download duration even if the app restarted during the transfer.

Implementing Model Downloads in Code

To trigger a download from a composable or ViewModel, obtain the ModelManagerViewModel instance and invoke the download method:

val viewModel: ModelManagerViewModel = hiltViewModel()
val model = selectedTask.models.first()
val task = selectedTask

// Starts the WorkManager job and updates UI state
viewModel.downloadModel(task = task, model = model)

Observe live download status reactively to update the user interface:

val uiState by viewModel.uiState.collectAsState()
val downloadStatus = uiState.modelDownloadStatus[model.name]

when (downloadStatus?.status) {
    ModelDownloadStatusType.IN_PROGRESS -> ShowProgressBar(
        downloadStatus.receivedBytes, 
        downloadStatus.totalBytes
    )
    ModelDownloadStatusType.SUCCEEDED -> ShowReadyState()
    ModelDownloadStatusType.FAILED -> ShowError(downloadStatus.errorMessage)
    else -> ShowDownloadButton()
}

Cancel an active download using the model reference:

viewModel.cancelDownloadModel(model)

For testing or custom implementations, interact directly with the repository:

val repo: DownloadRepository = DefaultDownloadRepository(context, lifecycleProvider)
repo.downloadModel(
    task = null,
    model = model,
    onStatusUpdated = { m, status -> 
        Log.d("Test", "Status: $status") 
    }
)

Summary

• The Gallery app uses a three-layer architecture: metadata loading via ModelManagerViewModel, WorkManager orchestration through DefaultDownloadRepository, and byte-level handling in DownloadWorker.
• Resume support is implemented via .tmp files and HTTP Range headers, allowing downloads to survive network interruptions and process restarts.
• Models cache to the external files directory under versioned subdirectories, with automatic ZIP extraction and cleanup to optimize storage.
• WorkManager ensures downloads proceed as expedited foreground jobs, with WorkInfo.progress providing real-time UI updates.
• The system distinguishes between remote downloads, partial transfers, and locally imported models through getModelDownloadStatus() filesystem checks.

Frequently Asked Questions

How does the Gallery app resume interrupted model downloads?

When DownloadWorker starts, it checks for an existing .tmp file in the model's version directory. If present, it calculates the current file size and sends an HTTP Range header to request only the remaining bytes. The server resumes streaming from that offset, appending to the partial file rather than overwriting it.

Where are downloaded ML models stored on Android?

According to the source code in DownloadWorker.kt, models stream into the app's external files directory returned by context.getExternalFilesDir(null), typically located at /sdcard/Android/data/com.google.ai.edge.gallery/files/. The final path follows the pattern <model-normalized-name>/<version>/<file>.

What is the purpose of the model allow-list JSON?

The allow-list JSON in model_allowlists/ serves as the single source of truth for available models. ModelManagerViewModel.loadModelAllowlist() parses this file to populate the UI with downloadable options, ensuring the app only attempts to fetch validated assets with correct URLs, checksums, and hardware requirements.

How does the system handle ZIP archives containing multiple model files?

When the model metadata specifies isZip = true, DownloadWorker downloads the archive to a .tmp file, then extracts its contents to a dedicated unzipDir subdirectory upon completion. The worker deletes the original ZIP file after extraction to conserve device storage, leaving only the unzipped assets in the cache.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}