diff --git a/examples/llama.android/app/src/main/java/com/example/llama/ui/components/ModelCards.kt b/examples/llama.android/app/src/main/java/com/example/llama/ui/components/ModelCards.kt
index f81bf954bb..6f39b721c3 100644
--- a/examples/llama.android/app/src/main/java/com/example/llama/ui/components/ModelCards.kt
+++ b/examples/llama.android/app/src/main/java/com/example/llama/ui/components/ModelCards.kt
@@ -35,6 +35,7 @@ import androidx.compose.runtime.CompositionLocalProvider
 import androidx.compose.runtime.remember
 import androidx.compose.ui.Alignment
 import androidx.compose.ui.Modifier
+import androidx.compose.ui.graphics.Color
 import androidx.compose.ui.text.font.FontStyle
 import androidx.compose.ui.text.font.FontWeight
 import androidx.compose.ui.unit.Dp
@@ -56,9 +57,15 @@ import java.util.Locale
 @Composable
 fun ModelCardCoreExpandable(
     model: ModelInfo,
+    containerColor: Color = MaterialTheme.colorScheme.primaryContainer,
     isExpanded: Boolean = false,
     onExpanded: ((Boolean) -> Unit)? = null
-) = ModelCardCoreExpandable(model, isExpanded, onExpanded) {
+) = ModelCardCoreExpandable(
+    model = model,
+    containerColor = containerColor,
+    isExpanded = isExpanded,
+    onExpanded = onExpanded
+) {
     Spacer(modifier = Modifier.height(8.dp))
 
     // Row 2: Context length, size label
@@ -80,6 +87,7 @@ fun ModelCardCoreExpandable(
 @Composable
 fun ModelCardCoreExpandable(
     model: ModelInfo,
+    containerColor: Color = MaterialTheme.colorScheme.primaryContainer,
     isExpanded: Boolean = false,
     onExpanded: ((Boolean) -> Unit)? = null,
     expandableSection: @Composable () -> Unit
@@ -89,9 +97,7 @@ fun ModelCardCoreExpandable(
             .fillMaxWidth()
             .clickable { onExpanded?.invoke(!isExpanded) },
         colors = when (isExpanded) {
-            true -> CardDefaults.cardColors(
-                containerColor = MaterialTheme.colorScheme.primaryContainer
-            )
+            true -> CardDefaults.cardColors(containerColor = containerColor)
             false -> CardDefaults.cardColors()
         },
         elevation = CardDefaults.cardElevation(defaultElevation = 12.dp)
@@ -151,6 +157,7 @@ fun ModelCardCoreExpandable(
 @Composable
 fun ModelCardFullExpandable(
     model: ModelInfo,
+    containerColor: Color = MaterialTheme.colorScheme.primaryContainer,
     isSelected: Boolean? = null,
     onSelected: ((Boolean) -> Unit)? = null,
     isExpanded: Boolean = false,
@@ -162,9 +169,7 @@ fun ModelCardFullExpandable(
                 .fillMaxWidth()
                 .clickable { onExpanded?.invoke(!isExpanded) },
             colors = when (isSelected) {
-                true -> CardDefaults.cardColors(
-                    containerColor = MaterialTheme.colorScheme.primaryContainer
-                )
+                true -> CardDefaults.cardColors(containerColor = containerColor)
                 false -> CardDefaults.cardColors()
                 else -> CardDefaults.cardColors()
             },
diff --git a/examples/llama.android/app/src/main/java/com/example/llama/ui/scaffold/topbar/ModelsManagementTopBar.kt b/examples/llama.android/app/src/main/java/com/example/llama/ui/scaffold/topbar/ModelsManagementTopBar.kt
index 0ddc3837e0..3db20d48bb 100644
--- a/examples/llama.android/app/src/main/java/com/example/llama/ui/scaffold/topbar/ModelsManagementTopBar.kt
+++ b/examples/llama.android/app/src/main/java/com/example/llama/ui/scaffold/topbar/ModelsManagementTopBar.kt
@@ -79,7 +79,7 @@ private fun StorageIndicator(
                 contentDescription = "Storage",
                 tint = when {
                     storageMetrics.availableGB < 5.0f -> MaterialTheme.colorScheme.error
-                    storageMetrics.availableGB < 10.0f -> MaterialTheme.colorScheme.tertiary
+                    storageMetrics.availableGB < 10.0f -> MaterialTheme.colorScheme.error.copy(alpha = 0.7f)
                     else -> MaterialTheme.colorScheme.onSurface
                 }
             )
diff --git a/examples/llama.android/app/src/main/java/com/example/llama/ui/screens/BenchmarkScreen.kt b/examples/llama.android/app/src/main/java/com/example/llama/ui/screens/BenchmarkScreen.kt
index 5457d688e3..3fd5768486 100644
--- a/examples/llama.android/app/src/main/java/com/example/llama/ui/screens/BenchmarkScreen.kt
+++ b/examples/llama.android/app/src/main/java/com/example/llama/ui/screens/BenchmarkScreen.kt
@@ -192,7 +192,12 @@ private fun ModelCardWithLoadingMetrics(
     loadingMetrics: ModelLoadingMetrics,
     isExpanded: Boolean = false,
     onExpanded: ((Boolean) -> Unit)? = null,
-) = ModelCardCoreExpandable(model, isExpanded, onExpanded) {
+) = ModelCardCoreExpandable(
+    model = model,
+    containerColor = MaterialTheme.colorScheme.tertiaryContainer,
+    isExpanded = isExpanded,
+    onExpanded = onExpanded
+) {
     Spacer(modifier = Modifier.height(8.dp))
 
     // Row 2: Context length, size label
diff --git a/examples/llama.android/app/src/main/java/com/example/llama/ui/screens/ConversationScreen.kt b/examples/llama.android/app/src/main/java/com/example/llama/ui/screens/ConversationScreen.kt
index 8e60014f7c..d5044d9488 100644
--- a/examples/llama.android/app/src/main/java/com/example/llama/ui/screens/ConversationScreen.kt
+++ b/examples/llama.android/app/src/main/java/com/example/llama/ui/screens/ConversationScreen.kt
@@ -237,7 +237,12 @@ fun ModelCardWithSystemPrompt(
     systemPrompt: String?,
     isExpanded: Boolean = true,
     onExpanded: ((Boolean) -> Unit)? = null,
-) = ModelCardCoreExpandable(model, isExpanded, onExpanded) {
+) = ModelCardCoreExpandable(
+    model = model,
+    containerColor = MaterialTheme.colorScheme.tertiaryContainer,
+    isExpanded = isExpanded,
+    onExpanded = onExpanded
+) {
     Spacer(modifier = Modifier.height(8.dp))
 
     // Row 2: Context length, size label
diff --git a/examples/llama.android/app/src/main/java/com/example/llama/ui/screens/ModelLoadingScreen.kt b/examples/llama.android/app/src/main/java/com/example/llama/ui/screens/ModelLoadingScreen.kt
index 451d5ab531..c3be3b674f 100644
--- a/examples/llama.android/app/src/main/java/com/example/llama/ui/screens/ModelLoadingScreen.kt
+++ b/examples/llama.android/app/src/main/java/com/example/llama/ui/screens/ModelLoadingScreen.kt
@@ -137,6 +137,7 @@ fun ModelLoadingScreen(
         selectedModel?.let { model ->
             ModelCardCoreExpandable(
                 model = model,
+                containerColor = MaterialTheme.colorScheme.tertiaryContainer,
                 isExpanded = isModelCardExpanded,
                 onExpanded = { isModelCardExpanded = !isModelCardExpanded },
             )
diff --git a/examples/llama.android/app/src/main/java/com/example/llama/viewmodel/ConversationViewModel.kt b/examples/llama.android/app/src/main/java/com/example/llama/viewmodel/ConversationViewModel.kt
index 2d29b3ab94..833f035acb 100644
--- a/examples/llama.android/app/src/main/java/com/example/llama/viewmodel/ConversationViewModel.kt
+++ b/examples/llama.android/app/src/main/java/com/example/llama/viewmodel/ConversationViewModel.kt
@@ -66,7 +66,8 @@ class ConversationViewModel @Inject constructor(
         // Add placeholder for assistant response
         val assistantMessage = Message.Assistant.Ongoing(
             content = "",
-            timestamp = System.currentTimeMillis()
+            timestamp = System.currentTimeMillis(),
+            tokensCount = 0,
         )
         _messages.value += assistantMessage
 
@@ -139,7 +140,8 @@ class ConversationViewModel @Inject constructor(
                     // Ongoing message update
                     add(Message.Assistant.Ongoing(
                         content = update.text,
-                        timestamp = it.timestamp
+                        timestamp = it.timestamp,
+                        tokensCount = it.tokensCount + 1,
                     ))
                 }
                 _messages.value = toList()
@@ -187,6 +189,7 @@ sealed class Message {
         data class Ongoing(
             override val timestamp: Long,
             override val content: String,
+            val tokensCount: Int,
         ) : Assistant()
 
         data class Stopped(