UI: show RAM warning if model too large

2025-07-10 20:00:35 -07:00 · 2025-07-10 20:00:35 -07:00 · 712bc745df
parent a5a54375a2
commit 712bc745df
7 changed files with 176 additions and 36 deletions
--- a/examples/llama.android/app/src/main/java/com/example/llama/MainActivity.kt
+++ b/examples/llama.android/app/src/main/java/com/example/llama/MainActivity.kt
@ -193,7 +193,7 @@ fun AppContent(
            val showSortMenu by modelSelectionViewModel.showSortMenu.collectAsState()
            val activeFilters by modelSelectionViewModel.activeFilters.collectAsState()
            val showFilterMenu by modelSelectionViewModel.showFilterMenu.collectAsState()
-            val preselectedModel by modelSelectionViewModel.preselectedModel.collectAsState()
+            val preselection by modelSelectionViewModel.preselection.collectAsState()

            ScaffoldConfig(
                topBarConfig =
@ -230,12 +230,13 @@ fun AppContent(
                        toggleMenu = modelSelectionViewModel::toggleFilterMenu
                    ),
                    runAction = BottomBarConfig.ModelSelection.RunActionConfig(
-                        selectedModel = preselectedModel,
-                        onRun = { model ->
-                            modelSelectionViewModel.confirmSelectedModel(model)
+                        preselection = preselection,
+                        onClickRun = { preselection ->
+                            if (modelSelectionViewModel.selectModel(preselection)) {
                                navigationActions.navigateToModelLoading()
                                modelSelectionViewModel.toggleSearchState(false)
                            }
+                        }
                    )
                )
            )
@ -437,6 +438,12 @@ fun AppContent(
                        onManageModelsClicked = {
                            navigationActions.navigateToModelsManagement()
                        },
+                        onConfirmSelection = { modelInfo, ramWarning ->
+                            if (modelSelectionViewModel.confirmSelectedModel(modelInfo, ramWarning)) {
+                                navigationActions.navigateToModelLoading()
+                                modelSelectionViewModel.toggleSearchState(false)
+                            }
+                        },
                        viewModel = modelSelectionViewModel
                    )
                }
--- a/examples/llama.android/app/src/main/java/com/example/llama/ui/components/InfoView.kt
+++ b/examples/llama.android/app/src/main/java/com/example/llama/ui/components/InfoView.kt
@ -3,7 +3,6 @@ package com.example.llama.ui.components
 import androidx.compose.foundation.layout.Arrangement
 import androidx.compose.foundation.layout.Column
 import androidx.compose.foundation.layout.Spacer
-import androidx.compose.foundation.layout.fillMaxSize
 import androidx.compose.foundation.layout.height
 import androidx.compose.foundation.layout.padding
 import androidx.compose.foundation.layout.size
@ -27,13 +26,14 @@ data class InfoAction(

@Composable
 fun InfoView(
+    modifier: Modifier = Modifier,
    title: String,
    icon: ImageVector,
    message: String? = null,
    action: InfoAction? = null
 ) {
    Column(
-        modifier = Modifier.fillMaxSize().padding(16.dp),
+        modifier = modifier.padding(16.dp),
        horizontalAlignment = Alignment.CenterHorizontally,
        verticalArrangement = Arrangement.Center
    ) {
@ -62,10 +62,9 @@ fun InfoView(
            )
        }

+        action?.let {
            Spacer(modifier = Modifier.height(24.dp))

-
-        action?.let {
            Button(onClick = action.onAction) {
                Icon(
                    imageVector = action.icon,
--- a/examples/llama.android/app/src/main/java/com/example/llama/ui/scaffold/bottombar/BottomBarConfig.kt
+++ b/examples/llama.android/app/src/main/java/com/example/llama/ui/scaffold/bottombar/BottomBarConfig.kt
@ -4,6 +4,7 @@ import androidx.compose.foundation.text.input.TextFieldState
 import com.example.llama.data.model.ModelFilter
 import com.example.llama.data.model.ModelInfo
 import com.example.llama.data.model.ModelSortOrder
+import com.example.llama.viewmodel.Preselection

 /**
 * [BottomAppBar] configurations
@ -42,8 +43,8 @@ sealed class BottomBarConfig {
        )

        data class RunActionConfig(
-            val selectedModel: ModelInfo?,
-            val onRun: (ModelInfo) -> Unit
+            val preselection: Preselection?,
+            val onClickRun: (Preselection) -> Unit,
        )
    }

--- a/examples/llama.android/app/src/main/java/com/example/llama/ui/scaffold/bottombar/ModelSelectionBottomBar.kt
+++ b/examples/llama.android/app/src/main/java/com/example/llama/ui/scaffold/bottombar/ModelSelectionBottomBar.kt
@ -173,12 +173,12 @@ fun ModelSelectionBottomBar(
        floatingActionButton = {
            // Only show FAB if a model is selected
            AnimatedVisibility(
-                visible = runAction.selectedModel != null,
+                visible = runAction.preselection != null,
                enter = scaleIn() + fadeIn(),
                exit = scaleOut() + fadeOut()
            ) {
                FloatingActionButton(
-                    onClick = { runAction.selectedModel?.let { runAction.onRun(it) } },
+                    onClick = { runAction.preselection?.let { runAction.onClickRun(it) } },
                    containerColor = MaterialTheme.colorScheme.primary
                ) {
                    Icon(
--- a/examples/llama.android/app/src/main/java/com/example/llama/ui/screens/ModelSelectionScreen.kt
+++ b/examples/llama.android/app/src/main/java/com/example/llama/ui/screens/ModelSelectionScreen.kt
@ -20,6 +20,8 @@ import androidx.compose.material.icons.filled.FolderOpen
 import androidx.compose.material.icons.filled.MoreVert
 import androidx.compose.material.icons.filled.Search
 import androidx.compose.material.icons.filled.SearchOff
+import androidx.compose.material.icons.filled.Warning
+import androidx.compose.material3.AlertDialog
 import androidx.compose.material3.Button
 import androidx.compose.material3.DockedSearchBar
 import androidx.compose.material3.ExperimentalMaterial3Api
@ -27,6 +29,7 @@ import androidx.compose.material3.Icon
 import androidx.compose.material3.MaterialTheme
 import androidx.compose.material3.SearchBarDefaults
 import androidx.compose.material3.Text
+import androidx.compose.material3.TextButton
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.LaunchedEffect
 import androidx.compose.runtime.collectAsState
@ -40,20 +43,24 @@ import androidx.compose.ui.focus.focusRequester
 import androidx.compose.ui.platform.LocalSoftwareKeyboardController
 import androidx.compose.ui.text.style.TextAlign
 import androidx.compose.ui.unit.dp
+import com.example.llama.data.model.ModelInfo
 import com.example.llama.ui.components.InfoAction
 import com.example.llama.ui.components.InfoView
 import com.example.llama.ui.components.ModelCardFullExpandable
+import com.example.llama.util.formatFileByteSize
 import com.example.llama.viewmodel.ModelSelectionViewModel
+import com.example.llama.viewmodel.Preselection.RamWarning

@OptIn(ExperimentalMaterial3Api::class)
@Composable
 fun ModelSelectionScreen(
    onManageModelsClicked: () -> Unit,
+    onConfirmSelection: (ModelInfo, RamWarning) -> Unit,
    viewModel: ModelSelectionViewModel,
 ) {
    // Data: models
    val filteredModels by viewModel.filteredModels.collectAsState()
-    val preselectedModel by viewModel.preselectedModel.collectAsState()
+    val preselection by viewModel.preselection.collectAsState()

    // Query states
    val textFieldState = viewModel.searchFieldState
@ -83,7 +90,7 @@ fun ModelSelectionScreen(
    }

    // Handle back button press
-    BackHandler(preselectedModel != null || isSearchActive) {
+    BackHandler(preselection != null || isSearchActive) {
        if (isSearchActive) {
            viewModel.toggleSearchState(false)
        } else {
@ -143,7 +150,7 @@ fun ModelSelectionScreen(
                        items(items = queryResults, key = { it.id }) { model ->
                            ModelCardFullExpandable(
                                model = model,
-                                isSelected = if (model == preselectedModel) true else null,
+                                isSelected = if (model == preselection?.modelInfo) true else null,
                                onSelected = { selected ->
                                    if (selected) {
                                        toggleSearchFocusAndIme(false)
@ -152,7 +159,7 @@ fun ModelSelectionScreen(
                                        toggleSearchFocusAndIme(true)
                                    }
                                },
-                                isExpanded = model == preselectedModel,
+                                isExpanded = model == preselection?.modelInfo,
                                onExpanded = { expanded ->
                                    viewModel.preselectModel(model, expanded)
                                    toggleSearchFocusAndIme(!expanded)
@ -176,11 +183,11 @@ fun ModelSelectionScreen(
                    items(items = filteredModels, key = { it.id }) { model ->
                        ModelCardFullExpandable(
                            model = model,
-                            isSelected = if (model == preselectedModel) true else null,
+                            isSelected = if (model == preselection?.modelInfo) true else null,
                            onSelected = { selected ->
                                if (!selected) viewModel.resetSelection()
                            },
-                            isExpanded = model == preselectedModel,
+                            isExpanded = model == preselection?.modelInfo,
                            onExpanded = { expanded ->
                                viewModel.preselectModel(model, expanded)
                            }
@ -189,6 +196,19 @@ fun ModelSelectionScreen(
                }
            }
        }
+
+        // Show insufficient RAM warning
+        preselection?.let {
+            it.ramWarning?.let { warning ->
+                if (warning.showing) {
+                    RamErrorDialog(
+                        warning,
+                        onDismiss = { viewModel.dismissRamWarning() },
+                        onConfirm = { onConfirmSelection(it.modelInfo, warning) }
+                    )
+                }
+            }
+        }
    }
 }

@ -203,6 +223,7 @@ private fun EmptyModelsView(
        else -> "No models match the selected filters"
    }
    InfoView(
+        modifier = Modifier.fillMaxSize(),
        title = "No Models Available",
        icon = Icons.Default.FolderOpen,
        message = message,
@ -253,3 +274,42 @@ private fun EmptySearchResultsView(
        }
    }
 }
+
+@Composable
+private fun RamErrorDialog(
+    ramError: RamWarning,
+    onDismiss: () -> Unit,
+    onConfirm: () -> Unit,
+) {
+    val requiredRam = formatFileByteSize(ramError.requiredRam)
+    val availableRam = formatFileByteSize(ramError.availableRam)
+
+    AlertDialog(
+        text = {
+            InfoView(
+                modifier = Modifier.fillMaxWidth(),
+                title = "Insufficient RAM",
+                icon = Icons.Default.Warning,
+                message = "You are trying to run a $requiredRam size model, " +
+                    "but currently there's only $availableRam memory available!",
+            )
+       },
+        containerColor = MaterialTheme.colorScheme.errorContainer,
+        titleContentColor = MaterialTheme.colorScheme.onErrorContainer,
+        textContentColor = MaterialTheme.colorScheme.onErrorContainer,
+        dismissButton = {
+            TextButton(onClick = onDismiss) {
+                Text("Cancel")
+            }
+        },
+        onDismissRequest = onDismiss,
+        confirmButton = {
+            TextButton(onClick = onConfirm) {
+                Text(
+                    text = "Proceed",
+                    color = MaterialTheme.colorScheme.error
+                )
+            }
+        }
+    )
+}
--- a/examples/llama.android/app/src/main/java/com/example/llama/ui/screens/ModelsManagementScreen.kt
+++ b/examples/llama.android/app/src/main/java/com/example/llama/ui/screens/ModelsManagementScreen.kt
@ -115,6 +115,7 @@ fun ModelsManagementScreen(
                else -> "No models match the selected filters"
            }
            InfoView(
+                modifier = Modifier.fillMaxSize(),
                title = "No Models Available",
                icon = Icons.Default.FolderOpen,
                message = message,
--- a/examples/llama.android/app/src/main/java/com/example/llama/viewmodel/ModelSelectionViewModel.kt
+++ b/examples/llama.android/app/src/main/java/com/example/llama/viewmodel/ModelSelectionViewModel.kt
@ -13,14 +13,18 @@ import com.example.llama.data.model.queryBy
 import com.example.llama.data.model.sortByOrder
 import com.example.llama.data.repo.ModelRepository
 import com.example.llama.engine.InferenceService
+import com.example.llama.monitoring.PerformanceMonitor
+import com.example.llama.viewmodel.Preselection.RamWarning
 import dagger.hilt.android.lifecycle.HiltViewModel
 import kotlinx.coroutines.FlowPreview
 import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.asStateFlow
 import kotlinx.coroutines.flow.collectLatest
 import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.flow.debounce
+import kotlinx.coroutines.flow.stateIn
 import kotlinx.coroutines.flow.update
 import kotlinx.coroutines.launch
 import javax.inject.Inject
@ -29,13 +33,14 @@ import javax.inject.Inject
@OptIn(FlowPreview::class)
@HiltViewModel
 class ModelSelectionViewModel @Inject constructor(
+    modelRepository: ModelRepository,
+    private val performanceMonitor: PerformanceMonitor,
    private val inferenceService: InferenceService,
-    modelRepository: ModelRepository
 ) : ViewModel() {

    // UI state: search mode
    private val _isSearchActive = MutableStateFlow(false)
-    val isSearchActive: StateFlow<Boolean> = _isSearchActive.asStateFlow()
+    val isSearchActive = _isSearchActive.asStateFlow()

    fun toggleSearchState(active: Boolean) {
        _isSearchActive.value = active
@ -50,14 +55,14 @@ class ModelSelectionViewModel @Inject constructor(

    // UI state: sort menu
    private val _sortOrder = MutableStateFlow(ModelSortOrder.LAST_USED)
-    val sortOrder: StateFlow<ModelSortOrder> = _sortOrder.asStateFlow()
+    val sortOrder = _sortOrder.asStateFlow()

    fun setSortOrder(order: ModelSortOrder) {
        _sortOrder.value = order
    }

    private val _showSortMenu = MutableStateFlow(false)
-    val showSortMenu: StateFlow<Boolean> = _showSortMenu.asStateFlow()
+    val showSortMenu = _showSortMenu.asStateFlow()

    fun toggleSortMenu(visible: Boolean) {
        _showSortMenu.value = visible
@ -84,7 +89,7 @@ class ModelSelectionViewModel @Inject constructor(
    }

    private val _showFilterMenu = MutableStateFlow(false)
-    val showFilterMenu: StateFlow<Boolean> = _showFilterMenu.asStateFlow()
+    val showFilterMenu = _showFilterMenu.asStateFlow()

    fun toggleFilterMenu(visible: Boolean) {
        _showFilterMenu.value = visible
@ -92,15 +97,34 @@ class ModelSelectionViewModel @Inject constructor(

    // Data: filtered & sorted models
    private val _filteredModels = MutableStateFlow<List<ModelInfo>>(emptyList())
-    val filteredModels: StateFlow<List<ModelInfo>> = _filteredModels.asStateFlow()
+    val filteredModels = _filteredModels.asStateFlow()

    // Data: queried models
    private val _queryResults = MutableStateFlow<List<ModelInfo>>(emptyList())
-    val queryResults: StateFlow<List<ModelInfo>> = _queryResults.asStateFlow()
+    val queryResults = _queryResults.asStateFlow()

    // Data: pre-selected model in expansion mode
    private val _preselectedModel = MutableStateFlow<ModelInfo?>(null)
-    val preselectedModel: StateFlow<ModelInfo?> = _preselectedModel.asStateFlow()
+    private val _showRamWarning = MutableStateFlow(false)
+    val preselection = combine(
+        _preselectedModel,
+        performanceMonitor.monitorMemoryUsage(),
+        _showRamWarning,
+    ) { model, memory, show ->
+        if (model == null) {
+            null
+        } else {
+            if (memory.availableMem >= model.sizeInBytes + RAM_LOAD_MODEL_BUFFER_BYTES) {
+                Preselection(model, null)
+            } else {
+                Preselection(model, RamWarning(model.sizeInBytes, memory.availableMem, show))
+            }
+        }
+    }.stateIn(
+        scope = viewModelScope,
+        started = SharingStarted.WhileSubscribed(SUBSCRIPTION_TIMEOUT_MS),
+        initialValue = null
+    )

    init {
        viewModelScope.launch {
@ -132,24 +156,58 @@ class ModelSelectionViewModel @Inject constructor(
    }

    /**
-     * Pre-select a model
+     * Pre-select a model to expand its details and show Run FAB
     */
-    fun preselectModel(modelInfo: ModelInfo, preselected: Boolean) =
-        _preselectedModel.update { current ->
-            if (preselected) modelInfo else null
+    fun preselectModel(modelInfo: ModelInfo, preselected: Boolean) {
+        _preselectedModel.value = if (preselected) modelInfo else null
+        _showRamWarning.value = false
    }

    /**
-     * Confirm currently selected model
+     * Select the currently pre-selected model
+     *
+     * @return True if RAM enough, otherwise False.
     */
-    fun confirmSelectedModel(modelInfo: ModelInfo) =
+    fun selectModel(preselection: Preselection) =
+        when (preselection.ramWarning?.showing) {
+            null -> {
+                inferenceService.setCurrentModel(preselection.modelInfo)
+                true
+            }
+            false -> {
+                _showRamWarning.value = true
+                false
+            }
+            else -> false
+        }
+
+    /**
+     * Acknowledge RAM warnings and confirm currently pre-selected model
+     *
+     * @return True if confirmed, otherwise False.
+     */
+    fun confirmSelectedModel(modelInfo: ModelInfo, ramWarning: RamWarning): Boolean =
+        if (ramWarning.showing) {
            inferenceService.setCurrentModel(modelInfo)
+            _showRamWarning.value = false
+            true
+        } else {
+            false
+        }
+
+    /**
+     * Dismiss the RAM warnings
+     */
+    fun dismissRamWarning() {
+        _showRamWarning.value = false
+    }

    /**
     * Reset selected model to none (before navigating away)
     */
    fun resetSelection() {
        _preselectedModel.value = null
+        _showRamWarning.value = false
    }

    /**
@ -164,6 +222,20 @@ class ModelSelectionViewModel @Inject constructor(
    companion object {
        private val TAG = ModelSelectionViewModel::class.java.simpleName

+        private const val SUBSCRIPTION_TIMEOUT_MS = 5000L
        private const val QUERY_DEBOUNCE_TIMEOUT_MS = 500L
+
+        private const val RAM_LOAD_MODEL_BUFFER_BYTES = 300 * 1024
    }
 }
+
+data class Preselection(
+    val modelInfo: ModelInfo,
+    val ramWarning: RamWarning?,
+) {
+    data class RamWarning(
+        val requiredRam: Long,
+        val availableRam: Long,
+        val showing: Boolean,
+    )
+}