UI: Introduce an abstract ViewModel to handle additional model unloading logics

This commit is contained in:
Han Yin 2025-04-18 14:15:49 -07:00
parent e1c77c6bbd
commit c5a3ac7eb1
5 changed files with 120 additions and 67 deletions

View File

@ -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)
)

View File

@ -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.

View File

@ -47,7 +47,7 @@ fun BenchmarkScreen(
// Handle back button press
BackHandler {
viewModel.onBackPressed()
viewModel.onBackPressed(onNavigateBack)
}
Column(

View File

@ -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<State> = benchmarkService.engineState
val selectedModel: StateFlow<ModelInfo?> = benchmarkService.currentSelectedModel
val benchmarkResults: StateFlow<String?> = benchmarkService.benchmarkResults
/**
* Model unloading dialog state
*/
private val _unloadDialogState = MutableStateFlow<UnloadDialogState>(UnloadDialogState.Hidden)
val unloadDialogState: StateFlow<UnloadDialogState> = _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
}
}
}

View File

@ -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<State> = 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>(UnloadDialogState.Hidden)
val unloadDialogState: StateFlow<UnloadDialogState> = _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
}
}