From 2b80ed922ef179bc42ad59920dc0f2021ba52ea1 Mon Sep 17 00:00:00 2001 From: Google AI Edge Gallery Date: Mon, 8 Jun 2026 16:44:06 -0700 Subject: [PATCH] Fix crash when benchmarking on-device AICore models. This routes all system-managed models (both AICore and Private Inference) to the Kotlin-side benchmark loop, bypassing the native JNI file check. PiperOrigin-RevId: 928840148 --- .../com/google/ai/edge/gallery/data/Model.kt | 4 + .../ai/edge/gallery/runtime/ModelHelperExt.kt | 4 + .../aicore/PrivateInferenceModelHelper.kt | 223 ++++++++++++++++++ .../gallery/ui/benchmark/BenchmarkScreen.kt | 11 +- .../ui/benchmark/BenchmarkViewModel.kt | 203 ++++++++++++---- .../gallery/ui/common/chat/ChatViewModel.kt | 8 +- .../gallery/ui/llmchat/LlmChatViewModel.kt | 8 +- .../ui/modelmanager/ModelManagerViewModel.kt | 120 ++++++---- 8 files changed, 487 insertions(+), 94 deletions(-) create mode 100644 Android/src/app/src/main/java/com/google/ai/edge/gallery/runtime/aicore/PrivateInferenceModelHelper.kt 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..779c0320c 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 @@ -41,6 +41,7 @@ enum class RuntimeType { @SerializedName("unknown") UNKNOWN, @SerializedName("litert_lm") LITERT_LM, @SerializedName("aicore") AICORE, + @SerializedName("private_inference") PRIVATE_INFERENCE, } enum class AICoreModelReleaseStage { @@ -333,6 +334,9 @@ data class Model( */ var latestModelFile: ModelFile? = null, ) { + val isSystemManaged: Boolean + get() = runtimeType == RuntimeType.AICORE || runtimeType == RuntimeType.PRIVATE_INFERENCE + init { normalizedName = NORMALIZE_NAME_REGEX.replace(name, "_") } diff --git a/Android/src/app/src/main/java/com/google/ai/edge/gallery/runtime/ModelHelperExt.kt b/Android/src/app/src/main/java/com/google/ai/edge/gallery/runtime/ModelHelperExt.kt index 83e09e2e4..d13cab0c8 100644 --- a/Android/src/app/src/main/java/com/google/ai/edge/gallery/runtime/ModelHelperExt.kt +++ b/Android/src/app/src/main/java/com/google/ai/edge/gallery/runtime/ModelHelperExt.kt @@ -19,6 +19,7 @@ package com.google.ai.edge.gallery.runtime import com.google.ai.edge.gallery.data.Model import com.google.ai.edge.gallery.data.RuntimeType import com.google.ai.edge.gallery.runtime.aicore.AICoreModelHelper +import com.google.ai.edge.gallery.runtime.aicore.PrivateInferenceModelHelper import com.google.ai.edge.gallery.ui.llmchat.LlmChatModelHelper val Model.runtimeHelper: LlmModelHelper @@ -26,5 +27,8 @@ val Model.runtimeHelper: LlmModelHelper if (this.runtimeType == RuntimeType.AICORE) { return AICoreModelHelper } + if (this.runtimeType == RuntimeType.PRIVATE_INFERENCE) { + return PrivateInferenceModelHelper + } return LlmChatModelHelper } diff --git a/Android/src/app/src/main/java/com/google/ai/edge/gallery/runtime/aicore/PrivateInferenceModelHelper.kt b/Android/src/app/src/main/java/com/google/ai/edge/gallery/runtime/aicore/PrivateInferenceModelHelper.kt new file mode 100644 index 000000000..632ab95b5 --- /dev/null +++ b/Android/src/app/src/main/java/com/google/ai/edge/gallery/runtime/aicore/PrivateInferenceModelHelper.kt @@ -0,0 +1,223 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.ai.edge.gallery.runtime.aicore + +import android.content.Context +import android.graphics.Bitmap +import android.util.Log +import androidx.concurrent.futures.await +import com.google.ai.edge.gallery.data.Model +import com.google.ai.edge.gallery.runtime.CleanUpListener +import com.google.ai.edge.gallery.runtime.LlmModelHelper +import com.google.ai.edge.gallery.runtime.ResultListener +import com.google.ai.edge.litertlm.Contents +import com.google.ai.edge.litertlm.Message +import com.google.ai.edge.litertlm.Role +import com.google.ai.edge.litertlm.ToolProvider +import com.google.android.apps.aicore.client.api.AiCoreClient +import com.google.android.apps.aicore.client.api.AiCoreClientOptions +import com.google.android.apps.aicore.client.api.AiFeature +import com.google.android.apps.aicore.client.api.legion.AuthKeyType +import com.google.android.apps.aicore.client.api.legion.ConnectionStrategy +import com.google.android.apps.aicore.client.api.legion.LegionServiceOptions +import com.google.android.apps.aicore.client.api.llm.LlmMessage +import com.google.android.apps.aicore.client.api.llm.LlmRequest +import com.google.android.apps.aicore.client.api.llm.LlmService +import java.time.Duration +import kotlinx.coroutines.CancellationException +import kotlinx.coroutines.CoroutineScope +import kotlinx.coroutines.Job +import kotlinx.coroutines.launch + +private const val TAG = "PrivateInferenceModelHelper" +private const val API_KEY = "[REDACTED]" // Prototyping API key + +data class PrivateInferenceChatMessage(val isUser: Boolean, val text: String) + +data class PrivateInferenceModelInstance( + val llmService: LlmService, + val chatHistory: MutableList = mutableListOf(), + var inferenceJob: Job? = null, +) + +object PrivateInferenceModelHelper : LlmModelHelper { + + private val cleanUpListeners: MutableMap = mutableMapOf() + + override fun initialize( + context: Context, + model: Model, + taskId: String, + supportImage: Boolean, + supportAudio: Boolean, + onDone: (String) -> Unit, + systemInstruction: Contents?, + tools: List, + enableConversationConstrainedDecoding: Boolean, + coroutineScope: CoroutineScope?, + ) { + if (coroutineScope == null) { + Log.e(TAG, "CoroutineScope is required for PrivateInferenceModelHelper") + onDone("Initialization failed: CoroutineScope is null") + return + } + + coroutineScope.launch { + try { + Log.d(TAG, "Initializing AiCoreClient...") + val client = AiCoreClient.create(AiCoreClientOptions.builder(context).build()) + + // Use LEGION_LLM for prototyping + val featureId = AiFeature.Id.LEGION_LLM + Log.d(TAG, "Fetching feature $featureId...") + val feature = client.getFeature(featureId).await() + + if (feature == null) { + Log.e(TAG, "LEGION_LLM feature is unavailable on this device.") + onDone("LEGION_LLM feature is unavailable on this device.") + return@launch + } + + Log.d(TAG, "Creating LlmService with Legion options...") + val optionsBuilder = + LegionServiceOptions.builder(client) + .setFeature(feature) + .setAuthKey(API_KEY) + .setAuthKeyType(AuthKeyType.API_KEY) + .setModelId("models/dev-v3p1-s") + .setConnectionStrategy(ConnectionStrategy.createCloseOnIdle(Duration.ofSeconds(30))) + + val llmService = LlmService.create(optionsBuilder.build()) + + Log.d(TAG, "Preparing inference engine...") + llmService.prepareInferenceEngine().await() + + model.instance = PrivateInferenceModelInstance(llmService) + Log.d(TAG, "Initialization completed successfully") + onDone("Feature is available (Private Inference)") + } catch (e: Exception) { + Log.e(TAG, "Initialization failed", e) + onDone("Initialization failed: ${e.message}") + } + } + } + + override fun resetConversation( + model: Model, + supportImage: Boolean, + supportAudio: Boolean, + systemInstruction: Contents?, + tools: List, + enableConversationConstrainedDecoding: Boolean, + initialMessages: List, + ) { + Log.d(TAG, "Resetting conversation") + val instance = model.instance as? PrivateInferenceModelInstance ?: return + instance.chatHistory.clear() + for (msg in initialMessages) { + instance.chatHistory.add( + PrivateInferenceChatMessage( + isUser = (msg.role == Role.USER), + text = msg.contents.toString(), + ) + ) + } + } + + override fun cleanUp(model: Model, onDone: () -> Unit) { + Log.d(TAG, "Cleaning up resources") + val instance = model.instance as? PrivateInferenceModelInstance + if (instance != null) { + instance.inferenceJob?.cancel() + // LlmService doesn't have a close/release, it relies on connection strategy + } + val onCleanUp = cleanUpListeners.remove(model.name) + onCleanUp?.invoke() + model.instance = null + onDone() + } + + override fun stopResponse(model: Model) { + Log.d(TAG, "Stopping response generation") + val instance = model.instance as? PrivateInferenceModelInstance ?: return + instance.inferenceJob?.cancel() + } + + override fun runInference( + model: Model, + input: String, + resultListener: ResultListener, + cleanUpListener: CleanUpListener, + onError: (message: String) -> Unit, + images: List, + audioClips: List, + coroutineScope: CoroutineScope?, + extraContext: Map?, + ) { + val instance = model.instance as? PrivateInferenceModelInstance + if (instance == null) { + onError("Private Inference model instance is not initialized.") + return + } + if (coroutineScope == null) { + Log.e(TAG, "CoroutineScope is required for inference") + onError("Inference failed: CoroutineScope is null") + return + } + + if (!cleanUpListeners.containsKey(model.name)) { + cleanUpListeners[model.name] = cleanUpListener + } + + instance.inferenceJob?.cancel() + instance.inferenceJob = coroutineScope.launch { + try { + Log.d(TAG, "Preparing LlmRequest...") + val requestBuilder = LlmRequest.builder() + + val messages = mutableListOf() + // Reconstruct history + for (msg in instance.chatHistory) { + val role = if (msg.isUser) LlmMessage.Role.USER else LlmMessage.Role.LLM + messages.add(LlmMessage.create(role, msg.text)) + } + // Add current input + messages.add(LlmMessage.create(LlmMessage.Role.USER, input)) + + requestBuilder.setMessages(messages) + val llmRequest = requestBuilder.build() + + Log.d(TAG, "Running Private Inference...") + val llmResult = instance.llmService.runInference(llmRequest).await() + + val reply = llmResult.results.firstOrNull() + val text = reply?.text ?: "" + + Log.d(TAG, "Inference successful, updating history") + instance.chatHistory.add(PrivateInferenceChatMessage(isUser = true, text = input)) + instance.chatHistory.add(PrivateInferenceChatMessage(isUser = false, text = text)) + + resultListener(text, true, null) + } catch (e: CancellationException) { + Log.i(TAG, "Inference cancelled") + } catch (e: Exception) { + Log.e(TAG, "Inference failed", e) + onError("Error: ${e.message}") + } + } + } +} 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..37bd9776c 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 @@ -72,6 +72,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 +106,17 @@ fun BenchmarkScreen( val configs = remember(selectedModel) { mutableStateListOf().apply { + val acceleratorOptions = + if (selectedModel.runtimeType == RuntimeType.PRIVATE_INFERENCE) { + listOf("Server") + } else { + 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..6a1538b0f 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 @@ -27,6 +27,7 @@ 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 +35,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 +44,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" @@ -90,6 +93,33 @@ constructor( collapseAll() } + private suspend fun ensureModelInitialized(model: Model): Boolean { + if (model.instance != null) { + return true + } + return suspendCancellableCoroutine { continuation -> + model.runtimeHelper.initialize( + context = appContext, + model = model, + taskId = "", + supportImage = false, + supportAudio = false, + onDone = { error -> + if (model.instance != null) { + continuation.resume(true) + } else { + Log.e(TAG, "Failed to initialize model for benchmark: $error") + continuation.resume(false) + } + }, + systemInstruction = null, + tools = emptyList(), + enableConversationConstrainedDecoding = false, + coroutineScope = viewModelScope, + ) + } + } + @OptIn(ExperimentalApi::class) fun runBenchmark( model: Model, @@ -121,56 +151,141 @@ 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() + + if (model.isSystemManaged) { + val initStart = System.currentTimeMillis() + val initialized = ensureModelInitialized(model) + val initLatency = System.currentTimeMillis() - initStart + if (!initialized) { + Log.e(TAG, "Failed to initialize system-managed model for benchmark") + setRunning(running = false) + return@launch } - 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") + firstInitTime = initLatency.toDouble() + + for (i in 0 until runCount) { + Log.d(TAG, "Start running system-managed benchmark #$i...") + val prompt = "a ".repeat(prefillTokens) + + var outputText = "" + var done = false + var errorOccurred = false + var timeToFirstTokenVal = -1.0 + + try { + val latency = + suspendCancellableCoroutine { continuation -> + val startTime = System.currentTimeMillis() + model.runtimeHelper.runInference( + model = model, + input = prompt, + images = emptyList(), + audioClips = emptyList(), + resultListener = { partial, isDone, _ -> + if (partial.isNotEmpty() && timeToFirstTokenVal < 0) { + timeToFirstTokenVal = (System.currentTimeMillis() - startTime) / 1000.0 + } + outputText += partial + if (isDone) { + done = true + continuation.resume(System.currentTimeMillis() - startTime) + } + }, + cleanUpListener = {}, + onError = { errMsg -> + Log.e(TAG, "System-managed benchmark run failed: $errMsg") + errorOccurred = true + continuation.resume(-1L) + }, + coroutineScope = viewModelScope, + extraContext = null, + ) + } + + if (errorOccurred || latency < 0) { + prefillSpeeds.add(0.0) + decodeSpeeds.add(0.0) + timesToFirstToken.add(0.0) + if (i > 0) nonFirstInitTimes.add(0.0) + continue + } + + val outputTokens = outputText.split(Regex("\\s+")).filter { it.isNotEmpty() }.size * 1.3 + val decodeSpeedVal = outputTokens / (latency / 1000.0) + + prefillSpeeds.add(0.0) + decodeSpeeds.add(decodeSpeedVal) + timesToFirstToken.add( + if (timeToFirstTokenVal >= 0) timeToFirstTokenVal else latency / 1000.0 + ) + if (i > 0) { + nonFirstInitTimes.add(0.0) + } + } catch (e: Exception) { + Log.e(TAG, "Exception during system-managed benchmark run", e) + prefillSpeeds.add(0.0) + decodeSpeeds.add(0.0) + timesToFirstToken.add(0.0) + if (i > 0) nonFirstInitTimes.add(0.0) + } - val initTimeMs = benchmarkInfo.initTimeInSecond * 1000.0 - if (i == 0) { - firstInitTime = initTimeMs - } else { - nonFirstInitTimes.add(initTimeMs) + setRunProgress(completedRunCount = i + 1) } - prefillSpeeds.add(benchmarkInfo.lastPrefillTokensPerSecond) - decodeSpeeds.add(benchmarkInfo.lastDecodeTokensPerSecond) - timesToFirstToken.add(benchmarkInfo.timeToFirstTokenInSecond) + } else { + // 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, + ) + Log.d(TAG, "Done #$i") - // 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}") + 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}") + } } + val endMs = System.currentTimeMillis() // 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..2344521d6 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 @@ -26,6 +26,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 +260,12 @@ abstract class ChatViewModel(val userDataDataStore: DataStore? = null) addItemDescription: String, customData: Any? = null, ) { - val accelerator = model.getStringConfigValue(key = ConfigKeys.ACCELERATOR, defaultValue = "") + val accelerator = + if (model.runtimeType == RuntimeType.PRIVATE_INFERENCE) { + "Server" + } else { + 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..20e813f1e 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,12 @@ open class LlmChatViewModelBase( onError: (String) -> Unit, allowThinking: Boolean = false, ) { - val accelerator = model.getStringConfigValue(key = ConfigKeys.ACCELERATOR, defaultValue = "") + val accelerator = + if (model.runtimeType == RuntimeType.PRIVATE_INFERENCE) { + "Server" + } else { + 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..89f844272 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 @@ -296,41 +296,54 @@ constructor( // TODO: b/494029782 - Both litertlm and aicore download and storage should be unified into a // model repository. - if (model.runtimeType == RuntimeType.AICORE) { - AICoreModelHelper.downloadModel( - context = context, - coroutineScope = viewModelScope, - model = model, - onProgress = { downloaded: Long, total: Long -> - setDownloadStatus( - curModel = model, - status = - ModelDownloadStatus( - status = ModelDownloadStatusType.IN_PROGRESS, - receivedBytes = downloaded, - totalBytes = total, - ), - ) - }, - onDone = { - setDownloadStatus( - curModel = model, - status = - ModelDownloadStatus( - status = ModelDownloadStatusType.SUCCEEDED, - receivedBytes = model.sizeInBytes, - totalBytes = model.sizeInBytes, - ), - ) - }, - onError = { error: String -> - setDownloadStatus( - curModel = model, - status = - ModelDownloadStatus(status = ModelDownloadStatusType.FAILED, errorMessage = error), - ) - }, - ) + if (model.isSystemManaged) { + if (model.runtimeType == RuntimeType.AICORE) { + AICoreModelHelper.downloadModel( + context = context, + coroutineScope = viewModelScope, + model = model, + onProgress = { downloaded: Long, total: Long -> + setDownloadStatus( + curModel = model, + status = + ModelDownloadStatus( + status = ModelDownloadStatusType.IN_PROGRESS, + receivedBytes = downloaded, + totalBytes = total, + ), + ) + }, + onDone = { + setDownloadStatus( + curModel = model, + status = + ModelDownloadStatus( + status = ModelDownloadStatusType.SUCCEEDED, + receivedBytes = model.sizeInBytes, + totalBytes = model.sizeInBytes, + ), + ) + }, + onError = { error: String -> + setDownloadStatus( + curModel = model, + status = + ModelDownloadStatus(status = ModelDownloadStatusType.FAILED, errorMessage = error), + ) + }, + ) + } else { + // For PRIVATE_INFERENCE, we mark it as SUCCEEDED immediately because it's system-managed. + setDownloadStatus( + curModel = model, + status = + ModelDownloadStatus( + status = ModelDownloadStatusType.SUCCEEDED, + receivedBytes = 0, + totalBytes = 0, + ), + ) + } return } @@ -349,7 +362,7 @@ constructor( // TODO: b/494029782 - Both litertlm and aicore download and storage should be unified into a // model repository. // AICore models cannot be deleted from the download repository within the app. - if (model.runtimeType == RuntimeType.AICORE) { + if (model.isSystemManaged) { return } downloadRepository.cancelDownloadModel(model) @@ -835,16 +848,16 @@ constructor( // TODO: b/494029782 - Both litertlm and aicore download and storage should be unified into a // model repository. - private fun checkAICoreModelStatuses() { + private fun checkSystemManagedModelStatuses() { viewModelScope.launch(Dispatchers.Main) { - val aicoreModels = + val systemManagedModels = uiState.value.tasks .flatMap { it.models } - .filter { it.runtimeType == RuntimeType.AICORE } + .filter { it.isSystemManaged } .distinctBy { it.name } - // Proactively attempt AICore model download upon app startup. - for (model in aicoreModels) { + // Proactively attempt system-managed model download upon app startup. + for (model in systemManagedModels) { downloadModel(task = null, model = model) } } @@ -904,11 +917,14 @@ constructor( modelAllowlist = readModelAllowlistFromDisk(fileName = MODEL_ALLOWLIST_TEST_FILENAME) // Local test only. - if (TEST_MODEL_ALLOW_LIST.isNotEmpty()) { + val allowListToUse = + if (TEST_MODEL_ALLOW_LIST.isNotEmpty()) TEST_MODEL_ALLOW_LIST + else TEST_MODEL_ALLOW_LIST_BAK + 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 +976,10 @@ constructor( continue } - if (allowedModel.runtimeType == RuntimeType.AICORE && !isAICoreAvailable) { + val isSystemManaged = + allowedModel.runtimeType == RuntimeType.AICORE || + allowedModel.runtimeType == RuntimeType.PRIVATE_INFERENCE + if (isSystemManaged && !isAICoreAvailable) { continue } @@ -1024,8 +1043,8 @@ constructor( // Process pending downloads. processPendingDownloads() - // Wait for AICore models statuses and update download indicators - checkAICoreModelStatuses() + // Wait for system-managed models statuses and update download indicators + checkSystemManagedModelStatuses() } catch (e: Exception) { e.printStackTrace() } @@ -1318,6 +1337,15 @@ constructor( private fun getModelDownloadStatus(model: Model): ModelDownloadStatus { Log.d(TAG, "Checking model ${model.name} download status...") + if (model.runtimeType == RuntimeType.PRIVATE_INFERENCE) { + Log.d(TAG, "Model is Private Inference. Set status to SUCCEEDED") + return ModelDownloadStatus( + status = ModelDownloadStatusType.SUCCEEDED, + receivedBytes = 0, + totalBytes = 0, + ) + } + if (model.localFileRelativeDirPathOverride.isNotEmpty()) { Log.d(TAG, "Model has localFileRelativeDirPathOverride set. Set status to SUCCEEDED") return ModelDownloadStatus(