diff --git a/examples/llama.android/app/src/main/java/com/example/llama/revamp/MainActivity.kt b/examples/llama.android/app/src/main/java/com/example/llama/revamp/MainActivity.kt index 846fbea9ad..1379b9aa2f 100644 --- a/examples/llama.android/app/src/main/java/com/example/llama/revamp/MainActivity.kt +++ b/examples/llama.android/app/src/main/java/com/example/llama/revamp/MainActivity.kt @@ -147,9 +147,8 @@ fun AppContent( topBarConfig = TopBarConfig.Performance( title = "Benchmark", navigationIcon = NavigationIcon.Back { - android.util.Log.w("JOJO", "Benchmark navigation icon tapped") - benchmarkViewModel.onBackPressed() - }, + benchmarkViewModel.onBackPressed { navigationActions.navigateUp() } + }, memoryMetrics = memoryUsage, temperatureInfo = Pair(temperatureInfo, useFahrenheit) ) diff --git a/examples/llama.android/app/src/main/java/com/example/llama/revamp/ui/components/UnloadModelConfirmationDialog.kt b/examples/llama.android/app/src/main/java/com/example/llama/revamp/ui/components/UnloadModelConfirmationDialog.kt index 620b643d85..b25c488ddb 100644 --- a/examples/llama.android/app/src/main/java/com/example/llama/revamp/ui/components/UnloadModelConfirmationDialog.kt +++ b/examples/llama.android/app/src/main/java/com/example/llama/revamp/ui/components/UnloadModelConfirmationDialog.kt @@ -10,6 +10,7 @@ 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 @@ -17,6 +18,16 @@ 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. diff --git a/examples/llama.android/app/src/main/java/com/example/llama/revamp/ui/screens/BenchmarkScreen.kt b/examples/llama.android/app/src/main/java/com/example/llama/revamp/ui/screens/BenchmarkScreen.kt index 00106fd39a..a208bf77a4 100644 --- a/examples/llama.android/app/src/main/java/com/example/llama/revamp/ui/screens/BenchmarkScreen.kt +++ b/examples/llama.android/app/src/main/java/com/example/llama/revamp/ui/screens/BenchmarkScreen.kt @@ -47,7 +47,7 @@ fun BenchmarkScreen( // Handle back button press BackHandler { - viewModel.onBackPressed() + viewModel.onBackPressed(onNavigateBack) } Column( diff --git a/examples/llama.android/app/src/main/java/com/example/llama/revamp/viewmodel/BenchmarkViewModel.kt b/examples/llama.android/app/src/main/java/com/example/llama/revamp/viewmodel/BenchmarkViewModel.kt index 32ef00e244..b28981fc59 100644 --- a/examples/llama.android/app/src/main/java/com/example/llama/revamp/viewmodel/BenchmarkViewModel.kt +++ b/examples/llama.android/app/src/main/java/com/example/llama/revamp/viewmodel/BenchmarkViewModel.kt @@ -1,35 +1,23 @@ package com.example.llama.revamp.viewmodel -import android.llama.cpp.InferenceEngine.State -import androidx.lifecycle.ViewModel import androidx.lifecycle.viewModelScope import com.example.llama.revamp.data.model.ModelInfo import com.example.llama.revamp.engine.BenchmarkService -import com.example.llama.revamp.ui.components.UnloadDialogState import dagger.hilt.android.lifecycle.HiltViewModel -import kotlinx.coroutines.flow.MutableStateFlow import kotlinx.coroutines.flow.StateFlow -import kotlinx.coroutines.flow.asStateFlow import kotlinx.coroutines.launch import javax.inject.Inject @HiltViewModel class BenchmarkViewModel @Inject constructor( private val benchmarkService: BenchmarkService -) : ViewModel() { +) : ModelUnloadingViewModel(benchmarkService) { /** - * Core states + * UI states */ - val engineState: StateFlow = benchmarkService.engineState val selectedModel: StateFlow = benchmarkService.currentSelectedModel val benchmarkResults: StateFlow = benchmarkService.benchmarkResults - /** - * Model unloading dialog state - */ - private val _unloadDialogState = MutableStateFlow(UnloadDialogState.Hidden) - val unloadDialogState: StateFlow = _unloadDialogState.asStateFlow() - /** * Run benchmark with specified parameters */ @@ -37,53 +25,4 @@ class BenchmarkViewModel @Inject constructor( viewModelScope.launch { benchmarkService.benchmark(pp, tg, pl, nr) } - - /** - * Handle back press from both back button and top bar - */ - fun onBackPressed() { - when (engineState.value) { - State.Benchmarking -> { - // Ignore back navigation requests during active benchmarking - } - else -> _unloadDialogState.value = UnloadDialogState.Confirming - } - } - - /** - * Handle confirmation from unload dialog - */ - fun onUnloadConfirmed(onNavigateBack: () -> Unit) { - viewModelScope.launch { - // Set unloading state to show progress - _unloadDialogState.value = UnloadDialogState.Unloading - android.util.Log.d("JOJO", "onUnloadConfirmed $ state -> Unloading") - - try { - // Unload the model - benchmarkService.unloadModel() - android.util.Log.d("JOJO", "onUnloadConfirmed $ service unload model finished!") - - // Reset state and navigate back - _unloadDialogState.value = UnloadDialogState.Hidden - android.util.Log.d("JOJO", "onUnloadConfirmed $ state -> Hidden!") - onNavigateBack() - } catch (e: Exception) { - // Handle error if needed - _unloadDialogState.value = UnloadDialogState.Hidden - } - } - } - - /** - * Handle dismissal of unload dialog - */ - fun onUnloadDismissed() { - when (_unloadDialogState.value) { - is UnloadDialogState.Unloading -> { - // Ignore dismissing requests during active benchmarking - } - else -> _unloadDialogState.value = UnloadDialogState.Hidden - } - } } diff --git a/examples/llama.android/app/src/main/java/com/example/llama/revamp/viewmodel/ModelUnloadingViewModel.kt b/examples/llama.android/app/src/main/java/com/example/llama/revamp/viewmodel/ModelUnloadingViewModel.kt new file mode 100644 index 0000000000..61434895d7 --- /dev/null +++ b/examples/llama.android/app/src/main/java/com/example/llama/revamp/viewmodel/ModelUnloadingViewModel.kt @@ -0,0 +1,104 @@ +package com.example.llama.revamp.viewmodel + +import android.llama.cpp.InferenceEngine +import android.llama.cpp.InferenceEngine.State +import android.llama.cpp.isModelLoaded +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 + +/** + * Base ViewModel class for screens that requires additional model unloading functionality + */ +abstract class ModelUnloadingViewModel( + private val inferenceService: InferenceService +) : ViewModel() { + + /** + * [InferenceEngine]'s core state + */ + val engineState: StateFlow = inferenceService.engineState + + /** + * Determine if the screen is in an uninterruptible state + * + * Subclass can override this default implementation + */ + protected open val isUninterruptible: Boolean + get() = engineState.value.isUninterruptible + + protected open val isModelLoaded: Boolean + get() = engineState.value.isModelLoaded + + /** + * [UnloadModelConfirmationDialog]'s UI states + */ + private val _unloadDialogState = MutableStateFlow(UnloadDialogState.Hidden) + val unloadDialogState: StateFlow = _unloadDialogState.asStateFlow() + + /** + * Handle back press from both back button and top bar + */ + open fun onBackPressed(onNavigateBack: () -> Unit) = + if (isUninterruptible) { + // During uninterruptible operations, ignore back navigation requests + } else if (!isModelLoaded) { + // If model not loaded, no need to unload at all, directly perform back navigation + onNavigateBack.invoke() + } else { + // If model is loaded, show confirmation dialog + _unloadDialogState.value = UnloadDialogState.Confirming + } + + /** + * Handle confirmation from unload dialog + */ + fun onUnloadConfirmed(onNavigateBack: () -> Unit) = + viewModelScope.launch { + // Set unloading state to show progress + _unloadDialogState.value = UnloadDialogState.Unloading + + try { + // Perform screen-specific cleanup + performCleanup() + + // Unload the model + inferenceService.unloadModel() + + // Reset state and navigate back + _unloadDialogState.value = UnloadDialogState.Hidden + onNavigateBack() + } catch (e: Exception) { + // Handle error + _unloadDialogState.value = UnloadDialogState.Error( + e.message ?: "Unknown error while unloading the model" + ) + } + } + + /** + * Handle dismissal of unload dialog + */ + fun onUnloadDismissed() = + when (_unloadDialogState.value) { + is UnloadDialogState.Unloading -> { + // Ignore dismissing requests during unloading + } + else -> _unloadDialogState.value = UnloadDialogState.Hidden + } + + /** + * Perform any screen-specific cleanup before unloading the model + * + * To be implemented by subclasses if needed + */ + protected open suspend fun performCleanup() { + // Default empty implementation + } +}