diff --git a/app/src/main/kotlin/com/google/ai/sample/GenerativeAiViewModelFactory.kt b/app/src/main/kotlin/com/google/ai/sample/GenerativeAiViewModelFactory.kt index e8191d3..42215a9 100644 --- a/app/src/main/kotlin/com/google/ai/sample/GenerativeAiViewModelFactory.kt +++ b/app/src/main/kotlin/com/google/ai/sample/GenerativeAiViewModelFactory.kt @@ -27,7 +27,9 @@ enum class ModelOption( val apiProvider: ApiProvider = ApiProvider.GOOGLE, val downloadUrl: String? = null, val size: String? = null, - val supportsScreenshot: Boolean = true + val supportsScreenshot: Boolean = true, + val isOfflineModel: Boolean = false, + val offlineModelFilename: String? = null ) { PUTER_GLM5("GLM-5 (Puter)", "z-ai/glm-5", ApiProvider.PUTER, supportsScreenshot = false), MISTRAL_LARGE_3("Mistral Large 3", "mistral-large-latest", ApiProvider.MISTRAL), @@ -49,7 +51,17 @@ enum class ModelOption( "gemma-3n-e4b-it", ApiProvider.GOOGLE, "https://huggingface.co/na5h13/gemma-3n-E4B-it-litert-lm/resolve/main/gemma-3n-E4B-it-int4.litertlm?download=true", - "4.92 GB" + "4.92 GB", + isOfflineModel = true, + offlineModelFilename = "gemma-3n-e4b-it-int4.litertlm" + ), + GEMMA_4_E4B_IT( + "Gemma 4 E4B it (offline)", + "gemma-4-e4b-it", + ApiProvider.GOOGLE, + "https://huggingface.co/litert-community/gemma-4-E4B-it-litert-lm/resolve/main/gemma-4-E4B-it.litertlm?download=true", + isOfflineModel = true, + offlineModelFilename = "gemma-4-E4B-it.litertlm" ), HUMAN_EXPERT("Human Expert", "human-expert", ApiProvider.HUMAN_EXPERT); @@ -77,7 +89,7 @@ val GenerativeViewModelFactory = object : ViewModelProvider.Factory { // Get the API key from MainActivity val mainActivity = MainActivity.getInstance() - val apiKey = if (currentModel == ModelOption.GEMMA_3N_E4B_IT || currentModel == ModelOption.HUMAN_EXPERT) { + val apiKey = if (currentModel.isOfflineModel || currentModel == ModelOption.HUMAN_EXPERT) { "offline-no-key-needed" // Dummy key for offline/human expert models } else { mainActivity?.getCurrentApiKey(currentModel.apiProvider) ?: "" diff --git a/app/src/main/kotlin/com/google/ai/sample/MenuScreen.kt b/app/src/main/kotlin/com/google/ai/sample/MenuScreen.kt index 6114b6d..10dbce5 100644 --- a/app/src/main/kotlin/com/google/ai/sample/MenuScreen.kt +++ b/app/src/main/kotlin/com/google/ai/sample/MenuScreen.kt @@ -210,10 +210,10 @@ fun MenuScreen( }, onClick = { expanded = false - val wasOfflineModel = selectedModel == ModelOption.GEMMA_3N_E4B_IT + val wasOfflineModel = selectedModel.isOfflineModel - if (modelOption == ModelOption.GEMMA_3N_E4B_IT) { - val isDownloaded = ModelDownloadManager.isModelDownloaded(context) + if (modelOption.isOfflineModel) { + val isDownloaded = ModelDownloadManager.isModelDownloaded(context, modelOption) if (!isDownloaded) { downloadDialogModel = modelOption showDownloadDialog = true @@ -257,7 +257,7 @@ fun MenuScreen( } // CPU/GPU Selection - only visible when offline model is selected - if (selectedModel == ModelOption.GEMMA_3N_E4B_IT) { + if (selectedModel.isOfflineModel) { item { val currentBackend = remember { mutableStateOf(GenerativeAiViewModelFactory.getBackend()) } @@ -444,7 +444,7 @@ fun MenuScreen( modifier = Modifier.fillMaxWidth().sliderFriendly() ) - if (selectedModel == ModelOption.GEMMA_3N_E4B_IT) { + if (selectedModel.isOfflineModel) { Spacer(modifier = Modifier.height(4.dp)) Text( text = "Note: LlmInference (offline model) may not support all generation parameters.", @@ -487,15 +487,13 @@ fun MenuScreen( val mainActivity = context as? MainActivity val activeModel = GenerativeAiViewModelFactory.getCurrentModel() // Check API Key for online models - if (activeModel.apiProvider != ApiProvider.GOOGLE || !activeModel.modelName.contains("litert")) { // Simple check, refine if needed. Actually offline model has specific Enum - if (activeModel != ModelOption.GEMMA_3N_E4B_IT && activeModel != ModelOption.HUMAN_EXPERT) { - val apiKey = mainActivity?.getCurrentApiKey(activeModel.apiProvider) - if (apiKey.isNullOrEmpty()) { - // Show API Key Dialog - onApiKeyButtonClicked(activeModel.apiProvider) // Or a specific callback to show dialog - return@TextButton - } - } + if (!activeModel.isOfflineModel && activeModel != ModelOption.HUMAN_EXPERT) { + val apiKey = mainActivity?.getCurrentApiKey(activeModel.apiProvider) + if (apiKey.isNullOrEmpty()) { + // Show API Key Dialog + onApiKeyButtonClicked(activeModel.apiProvider) // Or a specific callback to show dialog + return@TextButton + } } if (mainActivity != null) { // Ensure mainActivity is not null @@ -689,12 +687,12 @@ fun MenuScreen( } // Don't dismiss while downloading/paused }, - title = { Text("Download Model (4.92 GB)") }, + title = { Text("Download Model (${downloadDialogModel?.size ?: "unknown size"})") }, text = { Column { when (val state = dlState) { is ModelDownloadManager.DownloadState.Idle -> { - Text("Should the Gemma 3n E4B be downloaded?\n\n$formattedGbAvailable GB of storage available.") + Text("Should ${downloadDialogModel?.displayName ?: "this model"} be downloaded?\n\n$formattedGbAvailable GB of storage available.") } is ModelDownloadManager.DownloadState.Downloading -> { Text("Downloading...") @@ -741,8 +739,10 @@ fun MenuScreen( is ModelDownloadManager.DownloadState.Idle -> { TextButton( onClick = { - downloadDialogModel?.downloadUrl?.let { url -> - ModelDownloadManager.downloadModel(context, url) + downloadDialogModel?.let { model -> + model.downloadUrl?.let { url -> + ModelDownloadManager.downloadModel(context, model, url) + } // Task 2: Request notification permission when download starts val mainActivity = context as? MainActivity if (mainActivity != null && !mainActivity.isNotificationPermissionGranted()) { @@ -758,8 +758,10 @@ fun MenuScreen( is ModelDownloadManager.DownloadState.Paused -> { TextButton( onClick = { - downloadDialogModel?.downloadUrl?.let { url -> - ModelDownloadManager.resumeDownload(context, url) + downloadDialogModel?.let { model -> + model.downloadUrl?.let { url -> + ModelDownloadManager.resumeDownload(context, model, url) + } } } ) { Text("Resume") } @@ -777,8 +779,10 @@ fun MenuScreen( is ModelDownloadManager.DownloadState.Error -> { TextButton( onClick = { - downloadDialogModel?.downloadUrl?.let { url -> - ModelDownloadManager.downloadModel(context, url) + downloadDialogModel?.let { model -> + model.downloadUrl?.let { url -> + ModelDownloadManager.downloadModel(context, model, url) + } } } ) { Text("Retry") } @@ -794,7 +798,7 @@ fun MenuScreen( is ModelDownloadManager.DownloadState.Paused -> { TextButton( onClick = { - ModelDownloadManager.cancelDownload(context) + downloadDialogModel?.let { ModelDownloadManager.cancelDownload(context, it) } showDownloadDialog = false } ) { Text("Cancel Download") } diff --git a/app/src/main/kotlin/com/google/ai/sample/ScreenOperatorAccessibilityService.kt b/app/src/main/kotlin/com/google/ai/sample/ScreenOperatorAccessibilityService.kt index 7342c53..6c2616b 100644 --- a/app/src/main/kotlin/com/google/ai/sample/ScreenOperatorAccessibilityService.kt +++ b/app/src/main/kotlin/com/google/ai/sample/ScreenOperatorAccessibilityService.kt @@ -233,9 +233,9 @@ class ScreenOperatorAccessibilityService : AccessibilityService() { true // Asynchronous } is Command.TakeScreenshot -> { - val modelName = GenerativeAiViewModelFactory.getCurrentModel().modelName - if (modelName == "gemma-3n-e4b-it") { - Log.d(TAG, "Command.TakeScreenshot: Model is gemma-3n-e4b-it, capturing screen info only.") + val currentModel = GenerativeAiViewModelFactory.getCurrentModel() + if (currentModel.isOfflineModel) { + Log.d(TAG, "Command.TakeScreenshot: Model is offline, capturing screen info only.") this.showToast("Capturing screen info...", false) val screenInfo = captureScreenInformation() val mainActivity = MainActivity.getInstance() diff --git a/app/src/main/kotlin/com/google/ai/sample/feature/multimodal/ModelDownloadManager.kt b/app/src/main/kotlin/com/google/ai/sample/feature/multimodal/ModelDownloadManager.kt index fae1457..6bc58cd 100644 --- a/app/src/main/kotlin/com/google/ai/sample/feature/multimodal/ModelDownloadManager.kt +++ b/app/src/main/kotlin/com/google/ai/sample/feature/multimodal/ModelDownloadManager.kt @@ -7,6 +7,8 @@ import android.os.Build import android.util.Log import android.widget.Toast import androidx.core.app.NotificationCompat +import com.google.ai.sample.GenerativeAiViewModelFactory +import com.google.ai.sample.ModelOption import kotlin.coroutines.coroutineContext import kotlinx.coroutines.* import kotlinx.coroutines.flow.MutableStateFlow @@ -19,13 +21,12 @@ import java.net.HttpURLConnection import java.net.URL /** - * Custom download manager for the Gemma 3n model. + * Custom download manager for offline LiteRT models. * Uses HttpURLConnection with Range-Request support for resume capability. * Point 18: Includes Android notification for download progress. */ object ModelDownloadManager { private const val TAG = "ModelDownloadManager" - const val MODEL_FILENAME = "gemma-3n-e4b-it-int4.litertlm" private const val TEMP_SUFFIX = ".downloading" private const val BUFFER_SIZE = 8192 private const val MAX_RETRIES = 3 @@ -57,25 +58,27 @@ object ModelDownloadManager { private var downloadJob: Job? = null private var isPaused = false - fun isModelDownloaded(context: Context): Boolean { - val file = getModelFile(context) + fun isModelDownloaded(context: Context, model: ModelOption = GenerativeAiViewModelFactory.getCurrentModel()): Boolean { + val file = getModelFile(context, model) return file != null && file.exists() && file.length() > 0 } - fun getModelFile(context: Context): File? { + fun getModelFile(context: Context, model: ModelOption = GenerativeAiViewModelFactory.getCurrentModel()): File? { + val modelFilename = model.offlineModelFilename ?: return null val externalFilesDir = context.getExternalFilesDir(null) return if (externalFilesDir != null) { - File(externalFilesDir, MODEL_FILENAME) + File(externalFilesDir, modelFilename) } else { Log.e(TAG, "External files directory is not available.") null } } - private fun getTempFile(context: Context): File? { + private fun getTempFile(context: Context, model: ModelOption): File? { + val modelFilename = model.offlineModelFilename ?: return null val externalFilesDir = context.getExternalFilesDir(null) return if (externalFilesDir != null) { - File(externalFilesDir, MODEL_FILENAME + TEMP_SUFFIX) + File(externalFilesDir, modelFilename + TEMP_SUFFIX) } else { null } @@ -131,8 +134,8 @@ object ModelDownloadManager { notificationManager.cancel(DOWNLOAD_NOTIFICATION_ID) } - fun downloadModel(context: Context, url: String) { - if (isModelDownloaded(context)) { + fun downloadModel(context: Context, model: ModelOption, url: String) { + if (isModelDownloaded(context, model)) { Toast.makeText(context, "Model already downloaded.", Toast.LENGTH_SHORT).show() return } @@ -144,7 +147,7 @@ object ModelDownloadManager { isPaused = false downloadJob = CoroutineScope(Dispatchers.IO).launch { - downloadWithResume(context, url) + downloadWithResume(context, model, url) } } @@ -153,7 +156,7 @@ object ModelDownloadManager { isPaused = true } - fun resumeDownload(context: Context, url: String) { + fun resumeDownload(context: Context, model: ModelOption, url: String) { if (downloadJob?.isActive == true) { Log.d(TAG, "Download is still active, not resuming.") return @@ -161,18 +164,18 @@ object ModelDownloadManager { isPaused = false downloadJob = CoroutineScope(Dispatchers.IO).launch { - downloadWithResume(context, url) + downloadWithResume(context, model, url) } } - fun cancelDownload(context: Context) { + fun cancelDownload(context: Context, model: ModelOption) { Log.d(TAG, "Cancelling download...") isPaused = false downloadJob?.cancel() downloadJob = null // Delete temp file - val tempFile = getTempFile(context) + val tempFile = getTempFile(context, model) if (tempFile != null && tempFile.exists()) { tempFile.delete() Log.d(TAG, "Temp file deleted.") @@ -185,12 +188,12 @@ object ModelDownloadManager { } } - private suspend fun downloadWithResume(context: Context, url: String) { - val tempFile = getTempFile(context) ?: run { + private suspend fun downloadWithResume(context: Context, model: ModelOption, url: String) { + val tempFile = getTempFile(context, model) ?: run { _downloadState.value = DownloadState.Error("Storage not available.") return } - val finalFile = getModelFile(context) ?: run { + val finalFile = getModelFile(context, model) ?: run { _downloadState.value = DownloadState.Error("Storage not available.") return } @@ -348,4 +351,3 @@ object ModelDownloadManager { } } } - diff --git a/app/src/main/kotlin/com/google/ai/sample/feature/multimodal/PhotoReasoningScreen.kt b/app/src/main/kotlin/com/google/ai/sample/feature/multimodal/PhotoReasoningScreen.kt index db0f122..3b60b3b 100644 --- a/app/src/main/kotlin/com/google/ai/sample/feature/multimodal/PhotoReasoningScreen.kt +++ b/app/src/main/kotlin/com/google/ai/sample/feature/multimodal/PhotoReasoningScreen.kt @@ -370,7 +370,7 @@ fun PhotoReasoningScreen( ) } - val isGemma = modelName == "gemma-3n-e4b-it" + val isGemma = com.google.ai.sample.GenerativeAiViewModelFactory.getCurrentModel().isOfflineModel val isLoading = uiState is PhotoReasoningUiState.Loading val showStopButton = isGenerationRunning || isLoading || isOfflineGpuModelLoaded || isGemma val stopButtonText = if (isGenerationRunning || isLoading) "Stop" else "Model Unload" @@ -406,9 +406,9 @@ fun PhotoReasoningScreen( return@IconButton } - // Check MediaProjection for all models except gemma-3n-e4b-it and human-expert + // Check MediaProjection for all models except offline and human-expert // Human Expert uses its own MediaProjection for WebRTC, not ScreenCaptureService - if (!isMediaProjectionPermissionGranted && modelName != "gemma-3n-e4b-it" && modelName != "human-expert") { + if (!isMediaProjectionPermissionGranted && !com.google.ai.sample.GenerativeAiViewModelFactory.getCurrentModel().isOfflineModel && modelName != "human-expert") { mainActivity?.requestMediaProjectionPermission { // This block will be executed after permission is granted if (userQuestion.isNotBlank()) { diff --git a/app/src/main/kotlin/com/google/ai/sample/feature/multimodal/PhotoReasoningViewModel.kt b/app/src/main/kotlin/com/google/ai/sample/feature/multimodal/PhotoReasoningViewModel.kt index bad8b1c..61ae96f 100644 --- a/app/src/main/kotlin/com/google/ai/sample/feature/multimodal/PhotoReasoningViewModel.kt +++ b/app/src/main/kotlin/com/google/ai/sample/feature/multimodal/PhotoReasoningViewModel.kt @@ -286,8 +286,8 @@ class PhotoReasoningViewModel( // Initialize model if it's the offline one and already downloaded val currentModel = com.google.ai.sample.GenerativeAiViewModelFactory.getCurrentModel() val context = appContext - if (currentModel == ModelOption.GEMMA_3N_E4B_IT) { - if (ModelDownloadManager.isModelDownloaded(context)) { + if (currentModel.isOfflineModel) { + if (ModelDownloadManager.isModelDownloaded(context, currentModel)) { // Point 7 & 16: Initialize model asynchronously to not block UI viewModelScope.launch(Dispatchers.IO) { withContext(Dispatchers.Main) { @@ -324,7 +324,8 @@ class PhotoReasoningViewModel( private fun initializeOfflineModel(context: Context): String? { try { if (llmInference == null) { - val modelFile = ModelDownloadManager.getModelFile(context) + val currentModel = com.google.ai.sample.GenerativeAiViewModelFactory.getCurrentModel() + val modelFile = ModelDownloadManager.getModelFile(context, currentModel) if (modelFile != null && modelFile.exists()) { // Load backend preference GenerativeAiViewModelFactory.loadBackendPreference(context) @@ -406,7 +407,7 @@ class PhotoReasoningViewModel( } private fun isOfflineGpuModelLoaded(): Boolean { - return com.google.ai.sample.GenerativeAiViewModelFactory.getCurrentModel() == ModelOption.GEMMA_3N_E4B_IT && + return com.google.ai.sample.GenerativeAiViewModelFactory.getCurrentModel().isOfflineModel && com.google.ai.sample.GenerativeAiViewModelFactory.getBackend() == InferenceBackend.GPU && llmInference != null } @@ -649,10 +650,10 @@ class PhotoReasoningViewModel( } // Check for offline model (Gemma) - if (currentModel == ModelOption.GEMMA_3N_E4B_IT) { + if (currentModel.isOfflineModel) { val context = appContext - if (!ModelDownloadManager.isModelDownloaded(context)) { + if (!ModelDownloadManager.isModelDownloaded(context, currentModel)) { _uiState.value = PhotoReasoningUiState.Error("Model not downloaded.") return } @@ -893,7 +894,7 @@ class PhotoReasoningViewModel( val apiKeyManager = ApiKeyManager.getInstance(context) val currentKey = apiKeyManager.getCurrentApiKey(currentModel.apiProvider) - if (currentKey != null && currentModel != ModelOption.GEMMA_3N_E4B_IT && currentModel != ModelOption.HUMAN_EXPERT) { + if (currentKey != null && !currentModel.isOfflineModel && currentModel != ModelOption.HUMAN_EXPERT) { val genSettings = com.google.ai.sample.util.GenerationSettingsPreferences.loadSettings(context, currentModel.modelName) val config = com.google.ai.client.generativeai.type.generationConfig { temperature = genSettings.temperature @@ -2251,7 +2252,7 @@ private fun processCommands(text: String) { screenInfo: String? = null ) { if (screenshotUri == Uri.EMPTY) { - // This case is for gemma-3n-e4b-it, where we don't have a screenshot. + // This case is for offline models, where we don't have a screenshot. // We just want to send the screen info. val genericAnalysisPrompt = createGenericScreenshotPrompt() reason(