internal changes

Android/src/app/src/main/java/com/google/ai/edge/gallery/data/Model.kt

@@ -333,6 +333,7 @@ data class Model(
    */
   var latestModelFile: ModelFile? = null,
 ) {
+
   init {
     normalizedName = NORMALIZE_NAME_REGEX.replace(name, "_")
   }

Android/src/app/src/main/java/com/google/ai/edge/gallery/ui/benchmark/BenchmarkScreen.kt

@@ -16,6 +16,8 @@
 
 package com.google.ai.edge.gallery.ui.benchmark
 
+// BEGIN_INTERNAL
+// END_INTERNAL
 import android.os.Bundle
 import androidx.compose.animation.AnimatedVisibility
 import androidx.compose.animation.fadeIn
@@ -72,6 +74,7 @@ import com.google.ai.edge.gallery.data.ConfigKey
 import com.google.ai.edge.gallery.data.ConfigKeys
 import com.google.ai.edge.gallery.data.Model
 import com.google.ai.edge.gallery.data.NumberSliderConfig
+import com.google.ai.edge.gallery.data.RuntimeType
 import com.google.ai.edge.gallery.data.SegmentedButtonConfig
 import com.google.ai.edge.gallery.data.ValueType
 import com.google.ai.edge.gallery.data.convertValueToTargetType
@@ -105,11 +108,13 @@ fun BenchmarkScreen(
   val configs =
     remember(selectedModel) {
       mutableStateListOf<Config>().apply {
+        val acceleratorOptions =
+          selectedModel.accelerators.map { it.label }
         add(
           SegmentedButtonConfig(
             key = ConfigKeys.ACCELERATOR,
-            defaultValue = selectedModel.accelerators.getOrNull(0)?.label ?: Accelerator.CPU.label,
-            options = selectedModel.accelerators.map { it.label },
+            defaultValue = acceleratorOptions.getOrNull(0) ?: Accelerator.CPU.label,
+            options = acceleratorOptions,
             allowMultiple = false,
           )
         )

Android/src/app/src/main/java/com/google/ai/edge/gallery/ui/benchmark/BenchmarkViewModel.kt

@@ -22,18 +22,21 @@ import androidx.lifecycle.viewModelScope
 import com.google.ai.edge.gallery.BuildConfig
 import com.google.ai.edge.gallery.data.DataStoreRepository
 import com.google.ai.edge.gallery.data.Model
+import com.google.ai.edge.gallery.data.RuntimeType
 import com.google.ai.edge.gallery.proto.BenchmarkResult
 import com.google.ai.edge.gallery.proto.LlmBenchmarkBasicInfo
 import com.google.ai.edge.gallery.proto.LlmBenchmarkResult
 import com.google.ai.edge.gallery.proto.LlmBenchmarkStats
 import com.google.ai.edge.gallery.proto.ValueSeries
+import com.google.ai.edge.gallery.runtime.runtimeHelper
 import com.google.ai.edge.litertlm.Backend
 import com.google.ai.edge.litertlm.ExperimentalApi
 import com.google.ai.edge.litertlm.benchmark
 import dagger.hilt.android.lifecycle.HiltViewModel
 import dagger.hilt.android.qualifiers.ApplicationContext
 import java.io.File
 import javax.inject.Inject
+import kotlin.coroutines.resume
 import kotlin.math.ceil
 import kotlin.math.floor
 import kotlin.random.Random
@@ -42,6 +45,7 @@ import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.asStateFlow
 import kotlinx.coroutines.flow.update
 import kotlinx.coroutines.launch
+import kotlinx.coroutines.suspendCancellableCoroutine
 
 private const val TAG = "AGBenchmarkVM"
 
@@ -121,56 +125,60 @@ constructor(
       val timesToFirstToken = mutableListOf<Double>()
       var firstInitTime = 0.0
       val nonFirstInitTimes = mutableListOf<Double>()
-      // Create a temporary cache dir to run benchmark in.
-      val timestamp = System.currentTimeMillis()
-      var needCleanUpCacheDir = true
-      val benchmarkCacheDir = File(appContext.cacheDir, "benchmark_$timestamp")
-      var cacheDirPath = benchmarkCacheDir.absolutePath
-      if (!benchmarkCacheDir.mkdirs()) {
-        Log.e(TAG, "Failed to create benchmark cache directory: ${benchmarkCacheDir.absolutePath}")
-        cacheDirPath = appContext.cacheDir.absolutePath
-        needCleanUpCacheDir = false
-      }
-      Log.d(TAG, "Using benchmark cache dir: $cacheDirPath")
-      val backend: Backend =
-        when (accelerator.lowercase()) {
-          "gpu" -> Backend.GPU()
-          "npu",
-          "tpu" -> Backend.NPU(nativeLibraryDir = appContext.applicationInfo.nativeLibraryDir)
-          else -> Backend.CPU()
-        }
-      val modelPath = model.getPath(context = appContext)
-      for (i in 0 until runCount) {
-        Log.d(TAG, "Start running #$i...")
-        val benchmarkInfo =
-          benchmark(
-            modelPath = modelPath,
-            backend = backend,
-            prefillTokens = prefillTokens,
-            decodeTokens = decodeTokens,
-            cacheDir = cacheDirPath,
+
+        // Create a temporary cache dir to run benchmark in.
+        val timestamp = System.currentTimeMillis()
+        var needCleanUpCacheDir = true
+        val benchmarkCacheDir = File(appContext.cacheDir, "benchmark_$timestamp")
+        var cacheDirPath = benchmarkCacheDir.absolutePath
+        if (!benchmarkCacheDir.mkdirs()) {
+          Log.e(
+            TAG,
+            "Failed to create benchmark cache directory: ${benchmarkCacheDir.absolutePath}",
           )
-        Log.d(TAG, "Done #$i")
+          cacheDirPath = appContext.cacheDir.absolutePath
+          needCleanUpCacheDir = false
+        }
+        Log.d(TAG, "Using benchmark cache dir: $cacheDirPath")
+        val backend: Backend =
+          when (accelerator.lowercase()) {
+            "gpu" -> Backend.GPU()
+            "npu",
+            "tpu" -> Backend.NPU(nativeLibraryDir = appContext.applicationInfo.nativeLibraryDir)
+            else -> Backend.CPU()
+          }
+        val modelPath = model.getPath(context = appContext)
+        for (i in 0 until runCount) {
+          Log.d(TAG, "Start running #$i...")
+          val benchmarkInfo =
+            benchmark(
+              modelPath = modelPath,
+              backend = backend,
+              prefillTokens = prefillTokens,
+              decodeTokens = decodeTokens,
+              cacheDir = cacheDirPath,
+            )
+          Log.d(TAG, "Done #$i")
 
-        val initTimeMs = benchmarkInfo.initTimeInSecond * 1000.0
-        if (i == 0) {
-          firstInitTime = initTimeMs
-        } else {
-          nonFirstInitTimes.add(initTimeMs)
+          val initTimeMs = benchmarkInfo.initTimeInSecond * 1000.0
+          if (i == 0) {
+            firstInitTime = initTimeMs
+          } else {
+            nonFirstInitTimes.add(initTimeMs)
+          }
+          prefillSpeeds.add(benchmarkInfo.lastPrefillTokensPerSecond)
+          decodeSpeeds.add(benchmarkInfo.lastDecodeTokensPerSecond)
+          timesToFirstToken.add(benchmarkInfo.timeToFirstTokenInSecond)
+
+          // Mark finish for this run.
+          setRunProgress(completedRunCount = i + 1)
+        }
+        if (needCleanUpCacheDir) {
+          benchmarkCacheDir.deleteRecursively()
+          Log.d(TAG, "Cleaned up benchmark cache dir: ${benchmarkCacheDir.absolutePath}")
         }
-        prefillSpeeds.add(benchmarkInfo.lastPrefillTokensPerSecond)
-        decodeSpeeds.add(benchmarkInfo.lastDecodeTokensPerSecond)
-        timesToFirstToken.add(benchmarkInfo.timeToFirstTokenInSecond)
 
-        // Mark finish for this run.
-        setRunProgress(completedRunCount = i + 1)
-      }
       val endMs = System.currentTimeMillis()
-      if (needCleanUpCacheDir) {
-        benchmarkCacheDir.deleteRecursively()
-        Log.d(TAG, "Cleaned up benchmark cache dir: ${benchmarkCacheDir.absolutePath}")
-      }
-
       // Create and add benchmark result.
       val basicInfo =
         LlmBenchmarkBasicInfo.newBuilder()

Android/src/app/src/main/java/com/google/ai/edge/gallery/ui/common/chat/ChatViewModel.kt

@@ -16,6 +16,8 @@
 
 package com.google.ai.edge.gallery.ui.common.chat
 
+// BEGIN_INTERNAL
+// END_INTERNAL
 import android.content.Context
 import android.graphics.Bitmap
 import android.util.Log
@@ -26,6 +28,7 @@ import androidx.lifecycle.viewModelScope
 import com.google.ai.edge.gallery.common.processLlmResponse
 import com.google.ai.edge.gallery.data.ConfigKeys
 import com.google.ai.edge.gallery.data.Model
+import com.google.ai.edge.gallery.data.RuntimeType
 import com.google.ai.edge.gallery.proto.AudioMessageProto
 import com.google.ai.edge.gallery.proto.ChatMessageProto
 import com.google.ai.edge.gallery.proto.ChatSessionProto
@@ -259,7 +262,8 @@ abstract class ChatViewModel(val userDataDataStore: DataStore<UserData>? = null)
     addItemDescription: String,
     customData: Any? = null,
   ) {
-    val accelerator = model.getStringConfigValue(key = ConfigKeys.ACCELERATOR, defaultValue = "")
+    val accelerator =
+      model.getStringConfigValue(key = ConfigKeys.ACCELERATOR, defaultValue = "")
     val newMessagesByModel = _uiState.value.messagesByModel.toMutableMap()
     val newMessages = newMessagesByModel[model.name]?.toMutableList() ?: mutableListOf()
 

Android/src/app/src/main/java/com/google/ai/edge/gallery/ui/llmchat/LlmChatViewModel.kt

@@ -24,6 +24,7 @@ import androidx.lifecycle.viewModelScope
 import com.google.ai.edge.gallery.common.SystemPromptHelper
 import com.google.ai.edge.gallery.data.ConfigKeys
 import com.google.ai.edge.gallery.data.Model
+import com.google.ai.edge.gallery.data.RuntimeType
 import com.google.ai.edge.gallery.data.SystemPromptRepository
 import com.google.ai.edge.gallery.data.Task
 import com.google.ai.edge.gallery.proto.UserData
@@ -129,7 +130,8 @@ open class LlmChatViewModelBase(
     onError: (String) -> Unit,
     allowThinking: Boolean = false,
   ) {
-    val accelerator = model.getStringConfigValue(key = ConfigKeys.ACCELERATOR, defaultValue = "")
+    val accelerator =
+      model.getStringConfigValue(key = ConfigKeys.ACCELERATOR, defaultValue = "")
     viewModelScope.launch(Dispatchers.Default) {
       setInProgress(true)
       setPreparing(true)

Android/src/app/src/main/java/com/google/ai/edge/gallery/ui/modelmanager/ModelManagerViewModel.kt

@@ -840,10 +840,12 @@ constructor(
       val aicoreModels =
         uiState.value.tasks
           .flatMap { it.models }
-          .filter { it.runtimeType == RuntimeType.AICORE }
+          .filter {
+            it.runtimeType == RuntimeType.AICORE
+          }
           .distinctBy { it.name }
 
-      // Proactively attempt AICore model download upon app startup.
+      // Proactively attempt system-managed model download upon app startup.
       for (model in aicoreModels) {
         downloadModel(task = null, model = model)
       }
@@ -904,11 +906,12 @@ constructor(
         modelAllowlist = readModelAllowlistFromDisk(fileName = MODEL_ALLOWLIST_TEST_FILENAME)
 
         // Local test only.
-        if (TEST_MODEL_ALLOW_LIST.isNotEmpty()) {
+        val allowListToUse = TEST_MODEL_ALLOW_LIST
+        if (allowListToUse.isNotEmpty()) {
           Log.d(TAG, "Loading local model allowlist for testing.")
           val gson = Gson()
           try {
-            modelAllowlist = gson.fromJson(TEST_MODEL_ALLOW_LIST, ModelAllowlist::class.java)
+            modelAllowlist = gson.fromJson(allowListToUse, ModelAllowlist::class.java)
           } catch (e: JsonSyntaxException) {
             Log.e(TAG, "Failed to parse local test json", e)
           }
@@ -960,7 +963,9 @@ constructor(
             continue
           }
 
-          if (allowedModel.runtimeType == RuntimeType.AICORE && !isAICoreAvailable) {
+          val isSystemManaged =
+            allowedModel.runtimeType == RuntimeType.AICORE
+          if (isSystemManaged && !isAICoreAvailable) {
             continue
           }