From 67ecaf3b5f74601caa14bf75d91bbf63016694dc Mon Sep 17 00:00:00 2001 From: Google AI Edge Gallery Date: Mon, 8 Jun 2026 16:01:14 -0700 Subject: [PATCH] internal changes PiperOrigin-RevId: 928819870 --- .../com/google/ai/edge/gallery/data/Model.kt | 1 + .../gallery/ui/benchmark/BenchmarkScreen.kt | 9 +- .../ui/benchmark/BenchmarkViewModel.kt | 98 ++++++++++--------- .../gallery/ui/common/chat/ChatViewModel.kt | 6 +- .../gallery/ui/llmchat/LlmChatViewModel.kt | 4 +- .../ui/modelmanager/ModelManagerViewModel.kt | 15 ++- 6 files changed, 79 insertions(+), 54 deletions(-) diff --git a/Android/src/app/src/main/java/com/google/ai/edge/gallery/data/Model.kt b/Android/src/app/src/main/java/com/google/ai/edge/gallery/data/Model.kt index a0dbe7275..6fce20726 100644 --- a/Android/src/app/src/main/java/com/google/ai/edge/gallery/data/Model.kt +++ b/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, "_") } diff --git a/Android/src/app/src/main/java/com/google/ai/edge/gallery/ui/benchmark/BenchmarkScreen.kt b/Android/src/app/src/main/java/com/google/ai/edge/gallery/ui/benchmark/BenchmarkScreen.kt index 69dc40613..44fd476ab 100644 --- a/Android/src/app/src/main/java/com/google/ai/edge/gallery/ui/benchmark/BenchmarkScreen.kt +++ b/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().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, ) ) diff --git a/Android/src/app/src/main/java/com/google/ai/edge/gallery/ui/benchmark/BenchmarkViewModel.kt b/Android/src/app/src/main/java/com/google/ai/edge/gallery/ui/benchmark/BenchmarkViewModel.kt index b79913439..314bfdda7 100644 --- a/Android/src/app/src/main/java/com/google/ai/edge/gallery/ui/benchmark/BenchmarkViewModel.kt +++ b/Android/src/app/src/main/java/com/google/ai/edge/gallery/ui/benchmark/BenchmarkViewModel.kt @@ -22,11 +22,13 @@ 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 @@ -34,6 +36,7 @@ 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() var firstInitTime = 0.0 val nonFirstInitTimes = mutableListOf() - // 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() diff --git a/Android/src/app/src/main/java/com/google/ai/edge/gallery/ui/common/chat/ChatViewModel.kt b/Android/src/app/src/main/java/com/google/ai/edge/gallery/ui/common/chat/ChatViewModel.kt index 54aaf5689..9f2af2a4e 100644 --- a/Android/src/app/src/main/java/com/google/ai/edge/gallery/ui/common/chat/ChatViewModel.kt +++ b/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? = 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() diff --git a/Android/src/app/src/main/java/com/google/ai/edge/gallery/ui/llmchat/LlmChatViewModel.kt b/Android/src/app/src/main/java/com/google/ai/edge/gallery/ui/llmchat/LlmChatViewModel.kt index 304ba739b..e5f9eafe8 100644 --- a/Android/src/app/src/main/java/com/google/ai/edge/gallery/ui/llmchat/LlmChatViewModel.kt +++ b/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) diff --git a/Android/src/app/src/main/java/com/google/ai/edge/gallery/ui/modelmanager/ModelManagerViewModel.kt b/Android/src/app/src/main/java/com/google/ai/edge/gallery/ui/modelmanager/ModelManagerViewModel.kt index be79fe899..120979d45 100644 --- a/Android/src/app/src/main/java/com/google/ai/edge/gallery/ui/modelmanager/ModelManagerViewModel.kt +++ b/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 }