From ba40d689a10da25429ac766b878028f41ea916e6 Mon Sep 17 00:00:00 2001 From: Han Yin Date: Fri, 18 Apr 2025 12:11:34 -0700 Subject: [PATCH] UI: implement BenchmarkScreen's individual back handling --- .../revamp/ui/screens/BenchmarkScreen.kt | 46 ++++++++++++- .../revamp/viewmodel/BenchmarkViewModel.kt | 64 ++++++++++++++++++- 2 files changed, 105 insertions(+), 5 deletions(-) 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 ba6a2f4ade..00106fd39a 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 @@ -1,6 +1,7 @@ package com.example.llama.revamp.ui.screens import android.llama.cpp.InferenceEngine.State +import androidx.activity.compose.BackHandler import androidx.compose.foundation.background import androidx.compose.foundation.layout.Box import androidx.compose.foundation.layout.Column @@ -23,24 +24,32 @@ import androidx.compose.runtime.getValue import androidx.compose.ui.Alignment import androidx.compose.ui.Modifier import androidx.compose.ui.unit.dp -import androidx.hilt.navigation.compose.hiltViewModel import com.example.llama.revamp.ui.components.ModelCard +import com.example.llama.revamp.ui.components.UnloadDialogState +import com.example.llama.revamp.ui.components.UnloadModelConfirmationDialog import com.example.llama.revamp.ui.theme.MonospacedTextStyle import com.example.llama.revamp.viewmodel.BenchmarkViewModel @Composable fun BenchmarkScreen( - onBackPressed: () -> Unit, - viewModel: BenchmarkViewModel = hiltViewModel() + onNavigateBack: () -> Unit, + viewModel: BenchmarkViewModel ) { val engineState by viewModel.engineState.collectAsState() val benchmarkResults by viewModel.benchmarkResults.collectAsState() val selectedModel by viewModel.selectedModel.collectAsState() + val unloadDialogState by viewModel.unloadDialogState.collectAsState() + // Run benchmark when entering the screen LaunchedEffect(selectedModel) { viewModel.runBenchmark() } + // Handle back button press + BackHandler { + viewModel.onBackPressed() + } + Column( modifier = Modifier .fillMaxSize() @@ -115,4 +124,35 @@ fun BenchmarkScreen( } } } + + // Unload confirmation dialog + when (val state = unloadDialogState) { + is UnloadDialogState.Confirming -> { + UnloadModelConfirmationDialog( + onConfirm = { + viewModel.onUnloadConfirmed(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 */ } + } } 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 078992705c..32ef00e244 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 @@ -5,8 +5,11 @@ 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 @@ -14,10 +17,18 @@ import javax.inject.Inject class BenchmarkViewModel @Inject constructor( private val benchmarkService: BenchmarkService ) : ViewModel() { - + /** + * Core states + */ val engineState: StateFlow = benchmarkService.engineState - val benchmarkResults: StateFlow = benchmarkService.benchmarkResults 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 @@ -26,4 +37,53 @@ 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 + } + } }