UI: expose a single facade ModelUnloadDialogHandler; move UnloadModelState into ModelUnloadingViewModel.kt

examples/llama.android/app/src/main/java/com/example/llama/revamp/ui/components/ModelUnloadDialogHandler.kt

@@ -0,0 +1,167 @@
+package com.example.llama.revamp.ui.components
+
+import androidx.compose.foundation.layout.Arrangement
+import androidx.compose.foundation.layout.Column
+import androidx.compose.foundation.layout.Row
+import androidx.compose.foundation.layout.Spacer
+import androidx.compose.foundation.layout.fillMaxWidth
+import androidx.compose.foundation.layout.height
+import androidx.compose.foundation.layout.size
+import androidx.compose.foundation.layout.width
+import androidx.compose.material3.AlertDialog
+import androidx.compose.material3.CircularProgressIndicator
+import androidx.compose.material3.MaterialTheme
+import androidx.compose.material3.Text
+import androidx.compose.material3.TextButton
+import androidx.compose.runtime.Composable
+import androidx.compose.ui.Alignment
+import androidx.compose.ui.Modifier
+import androidx.compose.ui.unit.dp
+import com.example.llama.revamp.viewmodel.UnloadModelState
+
+/**
+ * Reusable component for handling model unloading dialogs
+ *
+ * @param [UnloadModelState]:
+ *  - Hidden: default state without showing any UI
+ *  - Confirming: show dismissible [UnloadModelDialog] and asks for user confirmation to unload current model
+ *  - Unloading: show non-dismissible [UnloadModelDialog] while unloading model
+ *  - Error: show [UnloadModelErrorDialog] to prompt error message to user
+ */
+@Composable
+fun ModelUnloadDialogHandler(
+    unloadModelState: UnloadModelState,
+    onUnloadConfirmed: (onNavigateBack: () -> Unit) -> Unit,
+    onUnloadDismissed: () -> Unit,
+    onNavigateBack: () -> Unit
+) {
+    when (unloadModelState) {
+        is UnloadModelState.Confirming -> {
+            UnloadModelDialog(
+                onConfirm = {
+                    onUnloadConfirmed(onNavigateBack)
+                },
+                onDismiss = onUnloadDismissed,
+                isUnloading = false
+            )
+        }
+        is UnloadModelState.Unloading -> {
+            UnloadModelDialog(
+                onConfirm = {
+                    onUnloadConfirmed(onNavigateBack)
+                },
+                onDismiss = onUnloadDismissed,
+                isUnloading = true
+            )
+        }
+        is UnloadModelState.Error -> {
+            UnloadModelErrorDialog(
+                errorMessage = unloadModelState.message,
+                onConfirm = {
+                    onUnloadDismissed()
+                    onNavigateBack()
+                },
+                onDismiss = onUnloadDismissed
+            )
+        }
+        is UnloadModelState.Hidden -> {
+            // Dialog not shown
+        }
+    }
+}
+
+@Composable
+private fun UnloadModelDialog(
+    onConfirm: () -> Unit,
+    onDismiss: () -> Unit,
+    isUnloading: Boolean = false
+) {
+    AlertDialog(
+        onDismissRequest = {
+            // Ignore dismiss requests while unloading the model
+            if (!isUnloading) onDismiss()
+        },
+        title = {
+            Text("Confirm Exit")
+        },
+        text = {
+            Column {
+                Text(
+                    "Going back will unload the current model. " +
+                        "Any unsaved conversation will be lost."
+                )
+
+                if (isUnloading) {
+                    Spacer(modifier = Modifier.height(16.dp))
+                    Row(
+                        modifier = Modifier.fillMaxWidth(),
+                        horizontalArrangement = Arrangement.Center,
+                        verticalAlignment = Alignment.CenterVertically
+                    ) {
+                        CircularProgressIndicator(
+                            modifier = Modifier.size(24.dp),
+                            strokeWidth = 2.dp
+                        )
+                        Spacer(modifier = Modifier.width(8.dp))
+                        Text("Unloading model...")
+                    }
+                }
+            }
+        },
+        confirmButton = {
+            TextButton(
+                onClick = onConfirm,
+                enabled = !isUnloading
+            ) {
+                Text("Yes, Exit")
+            }
+        },
+        dismissButton = {
+            TextButton(
+                onClick = onDismiss,
+                enabled = !isUnloading
+            ) {
+                Text("Cancel")
+            }
+        }
+    )
+}
+
+@Composable
+private fun UnloadModelErrorDialog(
+    errorMessage: String,
+    onConfirm: () -> Unit,
+    onDismiss: () -> Unit,
+) {
+    AlertDialog(
+        onDismissRequest = onDismiss,
+        title = {
+            Text(
+                text = "Error Unloading Model",
+                color = MaterialTheme.colorScheme.error
+            )
+        },
+        text = {
+            Column {
+                Text(
+                    text = errorMessage,
+                    style = MaterialTheme.typography.bodyMedium
+                )
+
+                Spacer(modifier = Modifier.height(8.dp))
+
+                Text(
+                    text = "You may need to restart the app if this problem persists.",
+                    style = MaterialTheme.typography.bodySmall,
+                    color = MaterialTheme.colorScheme.onSurfaceVariant
+                )
+            }
+        },
+        confirmButton = {
+            TextButton(onClick = onConfirm) { Text("Continue") }
+        },
+        dismissButton = {
+            TextButton(onClick = onDismiss) { Text("Stay on Screen") }
+        }
+    )
+}

examples/llama.android/app/src/main/java/com/example/llama/revamp/ui/components/UnloadModelConfirmationDialog.kt

@@ -1,91 +0,0 @@
-package com.example.llama.revamp.ui.components
-
-import androidx.compose.foundation.layout.Arrangement
-import androidx.compose.foundation.layout.Column
-import androidx.compose.foundation.layout.Row
-import androidx.compose.foundation.layout.Spacer
-import androidx.compose.foundation.layout.fillMaxWidth
-import androidx.compose.foundation.layout.height
-import androidx.compose.foundation.layout.size
-import androidx.compose.foundation.layout.width
-import androidx.compose.material3.AlertDialog
-import androidx.compose.material3.CircularProgressIndicator
-import androidx.compose.material3.MaterialTheme
-import androidx.compose.material3.Text
-import androidx.compose.material3.TextButton
-import androidx.compose.runtime.Composable
-import androidx.compose.ui.Alignment
-import androidx.compose.ui.Modifier
-import androidx.compose.ui.unit.dp
-
-/**
- * UI state for [UnloadModelConfirmationDialog]
- */
-sealed class UnloadDialogState {
-    object Hidden : UnloadDialogState()
-    object Confirming : UnloadDialogState()
-    object Unloading : UnloadDialogState()
-    data class Error(val message: String) : UnloadDialogState()
-}
-
-/**
- * Confirmation dialog shown when the user attempts to navigate away from
- * a screen that would require unloading the current model.
- */
-@Composable
-fun UnloadModelConfirmationDialog(
-    onConfirm: () -> Unit,
-    onDismiss: () -> Unit,
-    isUnloading: Boolean = false
-) {
-    AlertDialog(
-        onDismissRequest = {
-            // Ignore dismiss requests while unloading the model
-            if (!isUnloading) onDismiss()
-        },
-        title = {
-            Text("Confirm Exit")
-        },
-        text = {
-            Column {
-                Text(
-                    "Going back will unload the current model. " +
-                        "This operation cannot be undone. " +
-                        "Any unsaved conversation will be lost."
-                )
-
-                if (isUnloading) {
-                    Spacer(modifier = Modifier.height(16.dp))
-                    Row(
-                        modifier = Modifier.fillMaxWidth(),
-                        horizontalArrangement = Arrangement.Center,
-                        verticalAlignment = Alignment.CenterVertically
-                    ) {
-                        CircularProgressIndicator(
-                            modifier = Modifier.size(24.dp),
-                            strokeWidth = 2.dp
-                        )
-                        Spacer(modifier = Modifier.width(8.dp))
-                        Text("Unloading model...")
-                    }
-                }
-            }
-        },
-        confirmButton = {
-            TextButton(
-                onClick = onConfirm,
-                enabled = !isUnloading
-            ) {
-                Text("Yes, Exit")
-            }
-        },
-        dismissButton = {
-            TextButton(
-                onClick = onDismiss,
-                enabled = !isUnloading
-            ) {
-                Text("Cancel")
-            }
-        }
-    )
-}

examples/llama.android/app/src/main/java/com/example/llama/revamp/ui/screens/BenchmarkScreen.kt

@@ -25,8 +25,7 @@ import androidx.compose.ui.Alignment
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.unit.dp
 import com.example.llama.revamp.ui.components.ModelCard
-import com.example.llama.revamp.ui.components.UnloadDialogState
+import com.example.llama.revamp.ui.components.ModelUnloadDialogHandler
-import com.example.llama.revamp.ui.components.UnloadModelConfirmationDialog
 import com.example.llama.revamp.ui.theme.MonospacedTextStyle
 import com.example.llama.revamp.viewmodel.BenchmarkViewModel
 
@@ -38,7 +37,7 @@ fun BenchmarkScreen(
     val engineState by viewModel.engineState.collectAsState()
     val benchmarkResults by viewModel.benchmarkResults.collectAsState()
     val selectedModel by viewModel.selectedModel.collectAsState()
-    val unloadDialogState by viewModel.unloadDialogState.collectAsState()
+    val unloadDialogState by viewModel.unloadModelState.collectAsState()
 
     // Run benchmark when entering the screen
     LaunchedEffect(selectedModel) {
@@ -126,33 +125,10 @@ fun BenchmarkScreen(
     }
 
     // Unload confirmation dialog
-    when (val state = unloadDialogState) {
+    ModelUnloadDialogHandler(
-        is UnloadDialogState.Confirming -> {
+        unloadModelState = unloadDialogState,
-            UnloadModelConfirmationDialog(
+        onUnloadConfirmed = { viewModel.onUnloadConfirmed(onNavigateBack) },
-                onConfirm = {
+        onUnloadDismissed = { viewModel.onUnloadDismissed() },
-                    viewModel.onUnloadConfirmed(onNavigateBack)
+        onNavigateBack = onNavigateBack,
-                },
-                onDismiss = {
-                    viewModel.onUnloadDismissed()
-                },
-                isUnloading = false
     )
-        }
-        is UnloadDialogState.Unloading -> {
-            UnloadModelConfirmationDialog(
-                onConfirm = {
-                    viewModel.onUnloadConfirmed(onNavigateBack)
-                },
-                onDismiss = {
-                    viewModel.onUnloadDismissed()
-                },
-                isUnloading = true
-            )
-        }
-        is UnloadDialogState.Error -> {
-            // TODO-han.yin: TBD
-            android.util.Log.e("JOJO", "Unload error: ${state.message}")
-        }
-        else -> { /* Dialog not shown */ }
-    }
 }

examples/llama.android/app/src/main/java/com/example/llama/revamp/viewmodel/ModelUnloadingViewModel.kt

@@ -7,12 +7,22 @@ import android.llama.cpp.isUninterruptible
 import androidx.lifecycle.ViewModel
 import androidx.lifecycle.viewModelScope
 import com.example.llama.revamp.engine.InferenceService
-import com.example.llama.revamp.ui.components.UnloadDialogState
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.asStateFlow
 import kotlinx.coroutines.launch
 
+
+/**
+ * UI states to be consumed by [ModelUnloadDialogHandler], etc.
+ */
+sealed class UnloadModelState {
+    object Hidden : UnloadModelState()
+    object Confirming : UnloadModelState()
+    object Unloading : UnloadModelState()
+    data class Error(val message: String) : UnloadModelState()
+}
+
 /**
  * Base ViewModel class for screens that requires additional model unloading functionality
  */
@@ -39,11 +49,13 @@ abstract class ModelUnloadingViewModel(
     /**
      * [UnloadModelConfirmationDialog]'s UI states
      */
-    private val _unloadDialogState = MutableStateFlow<UnloadDialogState>(UnloadDialogState.Hidden)
+    private val _unloadModelState = MutableStateFlow<UnloadModelState>(UnloadModelState.Hidden)
-    val unloadDialogState: StateFlow<UnloadDialogState> = _unloadDialogState.asStateFlow()
+    val unloadModelState: StateFlow<UnloadModelState> = _unloadModelState.asStateFlow()
 
     /**
      * Handle back press from both back button and top bar
+     *
+     * Subclass can override this default implementation
      */
     open fun onBackPressed(onNavigateBack: () -> Unit) =
         if (isUninterruptible) {
@@ -53,7 +65,7 @@ abstract class ModelUnloadingViewModel(
             onNavigateBack.invoke()
         } else {
             // If model is loaded, show confirmation dialog
-            _unloadDialogState.value = UnloadDialogState.Confirming
+            _unloadModelState.value = UnloadModelState.Confirming
         }
 
     /**
@@ -62,7 +74,7 @@ abstract class ModelUnloadingViewModel(
     fun onUnloadConfirmed(onNavigateBack: () -> Unit) =
         viewModelScope.launch {
             // Set unloading state to show progress
-            _unloadDialogState.value = UnloadDialogState.Unloading
+            _unloadModelState.value = UnloadModelState.Unloading
 
             try {
                 // Perform screen-specific cleanup
@@ -72,11 +84,11 @@ abstract class ModelUnloadingViewModel(
                 inferenceService.unloadModel()
 
                 // Reset state and navigate back
-                _unloadDialogState.value = UnloadDialogState.Hidden
+                _unloadModelState.value = UnloadModelState.Hidden
                 onNavigateBack()
             } catch (e: Exception) {
                 // Handle error
-                _unloadDialogState.value = UnloadDialogState.Error(
+                _unloadModelState.value = UnloadModelState.Error(
                     e.message ?: "Unknown error while unloading the model"
                 )
             }
@@ -86,11 +98,11 @@ abstract class ModelUnloadingViewModel(
      * Handle dismissal of unload dialog
      */
     fun onUnloadDismissed() =
-        when (_unloadDialogState.value) {
+        when (_unloadModelState.value) {
-            is UnloadDialogState.Unloading -> {
+            is UnloadModelState.Unloading -> {
                 // Ignore dismissing requests during unloading
             }
-            else -> _unloadDialogState.value = UnloadDialogState.Hidden
+            else -> _unloadModelState.value = UnloadModelState.Hidden
         }
 
     /**
@@ -98,7 +110,5 @@ abstract class ModelUnloadingViewModel(
      *
      * To be implemented by subclasses if needed
      */
-    protected open suspend fun performCleanup() {
+    protected open suspend fun performCleanup() {}
-        // Default empty implementation
-    }
 }