Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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),
Expand All @@ -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);

Expand Down Expand Up @@ -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) ?: ""
Expand Down
50 changes: 27 additions & 23 deletions app/src/main/kotlin/com/google/ai/sample/MenuScreen.kt
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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()) }

Expand Down Expand Up @@ -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.",
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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...")
Expand Down Expand Up @@ -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()) {
Expand All @@ -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") }
Expand All @@ -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") }
Expand All @@ -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") }
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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
}
Expand Down Expand Up @@ -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
}
Expand All @@ -144,7 +147,7 @@ object ModelDownloadManager {

isPaused = false
downloadJob = CoroutineScope(Dispatchers.IO).launch {
downloadWithResume(context, url)
downloadWithResume(context, model, url)
}
}

Expand All @@ -153,26 +156,26 @@ 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
}

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.")
Expand All @@ -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
}
Expand Down Expand Up @@ -348,4 +351,3 @@ object ModelDownloadManager {
}
}
}

Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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()) {
Expand Down
Loading
Loading