Re-enable image pickers when config is changed, and other UI bug fixes.

Android/src/app/build.gradle.kts

@@ -30,7 +30,7 @@ android {
     minSdk = 26
     targetSdk = 35
     versionCode = 1
-    versionName = "0.9.1"
+    versionName = "0.9.2"
 
     // Needed for HuggingFace auth workflows.
     manifestPlaceholders["appAuthRedirectScheme"] = "com.google.aiedge.gallery.oauth"

Android/src/app/src/main/java/com/google/aiedge/gallery/GalleryApp.kt

@@ -23,6 +23,8 @@ import androidx.compose.foundation.layout.Column
 import androidx.compose.foundation.layout.Row
 import androidx.compose.foundation.layout.padding
 import androidx.compose.foundation.layout.size
+import androidx.compose.foundation.text.BasicText
+import androidx.compose.foundation.text.TextAutoSize
 import androidx.compose.material.icons.Icons
 import androidx.compose.material.icons.automirrored.rounded.ArrowBack
 import androidx.compose.material.icons.rounded.Refresh
@@ -44,6 +46,7 @@ import androidx.compose.ui.res.painterResource
 import androidx.compose.ui.res.stringResource
 import androidx.compose.ui.text.font.FontWeight
 import androidx.compose.ui.unit.dp
+import androidx.compose.ui.unit.sp
 import androidx.navigation.NavHostController
 import androidx.navigation.compose.rememberNavController
 import com.google.aiedge.gallery.data.AppBarAction
@@ -70,6 +73,7 @@ fun GalleryTopAppBar(
   scrollBehavior: TopAppBarScrollBehavior? = null,
   subtitle: String = "",
 ) {
+  val titleColor = MaterialTheme.colorScheme.primary
   CenterAlignedTopAppBar(
     title = {
       Column(horizontalAlignment = Alignment.CenterHorizontally) {
@@ -85,9 +89,16 @@ fun GalleryTopAppBar(
               tint = Color.Unspecified,
             )
           }
-          Text(
+          BasicText(
-            title,
+            text = title,
-            style = MaterialTheme.typography.titleLarge.copy(fontWeight = FontWeight.SemiBold)
+            maxLines = 1,
+            color = { titleColor },
+            style = MaterialTheme.typography.titleLarge.copy(fontWeight = FontWeight.SemiBold),
+            autoSize = TextAutoSize.StepBased(
+              minFontSize = 14.sp,
+              maxFontSize = 22.sp,
+              stepSize = 1.sp
+            )
           )
         }
         if (subtitle.isNotEmpty()) {

Android/src/app/src/main/java/com/google/aiedge/gallery/ui/common/chat/ChatPanel.kt

@@ -16,9 +16,10 @@
 
 package com.google.aiedge.gallery.ui.common.chat
 
-import androidx.compose.animation.AnimatedContent
+import androidx.compose.animation.AnimatedVisibility
 import androidx.compose.animation.ExperimentalSharedTransitionApi
-import androidx.compose.animation.SharedTransitionLayout
+import androidx.compose.animation.fadeIn
+import androidx.compose.animation.fadeOut
 import androidx.compose.foundation.Image
 import androidx.compose.foundation.background
 import androidx.compose.foundation.clickable
@@ -143,6 +144,18 @@ fun ChatPanel(
   val scope = rememberCoroutineScope()
   val haptic = LocalHapticFeedback.current
   var selectedImageMessage by remember { mutableStateOf<ChatMessageImage?>(null) }
+  val hasImageMessageToLastConfigChange = remember(messages) {
+    var foundImageMessage = false
+    for (message in messages.reversed()) {
+      if (message is ChatMessageConfigValuesChange) {
+        break
+      }
+      if (message is ChatMessageImage) {
+        foundImageMessage = true
+      }
+    }
+    foundImageMessage
+  }
 
   var curMessage by remember { mutableStateOf("") } // Correct state
   val focusManager = LocalFocusManager.current
@@ -222,378 +235,348 @@ fun ChatPanel(
     showErrorDialog = modelInitializationStatus?.status == ModelInitializationStatusType.ERROR
   }
 
-  SharedTransitionLayout(modifier = Modifier.fillMaxSize()) {
+  Column(
-    AnimatedContent(targetState = selectedImageMessage) { targetSelectedImageMessage ->
+    modifier = modifier.imePadding()
-      Column(
+  ) {
-        modifier = modifier.imePadding()
+    Box(contentAlignment = Alignment.BottomCenter, modifier = Modifier.weight(1f)) {
+      LazyColumn(
+        modifier = Modifier
+          .fillMaxSize()
+          .nestedScroll(nestedScrollConnection),
+        state = listState, verticalArrangement = Arrangement.Top,
       ) {
-        Box(contentAlignment = Alignment.BottomCenter, modifier = Modifier.weight(1f)) {
+        items(messages) { message ->
-          LazyColumn(
+          val imageHistoryCurIndex = remember { mutableIntStateOf(0) }
+          var hAlign: Alignment.Horizontal = Alignment.End
+          var backgroundColor: Color = MaterialTheme.customColors.userBubbleBgColor
+          var hardCornerAtLeftOrRight = false
+          var extraPaddingStart = 48.dp
+          var extraPaddingEnd = 0.dp
+          if (message.side == ChatSide.AGENT) {
+            hAlign = Alignment.Start
+            backgroundColor = MaterialTheme.customColors.agentBubbleBgColor
+            hardCornerAtLeftOrRight = true
+            extraPaddingStart = 0.dp
+            extraPaddingEnd = 48.dp
+          } else if (message.side == ChatSide.SYSTEM) {
+            extraPaddingStart = 24.dp
+            extraPaddingEnd = 24.dp
+            if (message.type == ChatMessageType.PROMPT_TEMPLATES) {
+              extraPaddingStart = 12.dp
+              extraPaddingEnd = 12.dp
+            }
+          }
+          if (message.type == ChatMessageType.IMAGE) {
+            backgroundColor = Color.Transparent
+          }
+          val bubbleBorderRadius = dimensionResource(R.dimen.chat_bubble_corner_radius)
+
+          Column(
             modifier = Modifier
-              .fillMaxSize()
+              .fillMaxWidth()
-              .nestedScroll(nestedScrollConnection),
+              .padding(
-            state = listState, verticalArrangement = Arrangement.Top,
+                start = 12.dp + extraPaddingStart,
+                end = 12.dp + extraPaddingEnd,
+                top = 6.dp,
+                bottom = 6.dp,
+              ),
+            horizontalAlignment = hAlign,
           ) {
-            items(messages) { message ->
+            // Sender row.
-              val imageHistoryCurIndex = remember { mutableIntStateOf(0) }
+            MessageSender(
-              var hAlign: Alignment.Horizontal = Alignment.End
+              message = message,
-              var backgroundColor: Color = MaterialTheme.customColors.userBubbleBgColor
+              agentNameRes = task.agentNameRes,
-              var hardCornerAtLeftOrRight = false
+              imageHistoryCurIndex = imageHistoryCurIndex.intValue
-              var extraPaddingStart = 48.dp
+            )
-              var extraPaddingEnd = 0.dp
+
-              if (message.side == ChatSide.AGENT) {
+            // Message body.
-                hAlign = Alignment.Start
+            when (message) {
-                backgroundColor = MaterialTheme.customColors.agentBubbleBgColor
+              // Loading.
-                hardCornerAtLeftOrRight = true
+              is ChatMessageLoading -> MessageBodyLoading()
-                extraPaddingStart = 0.dp
+
-                extraPaddingEnd = 48.dp
+              // Info.
-              } else if (message.side == ChatSide.SYSTEM) {
+              is ChatMessageInfo -> MessageBodyInfo(message = message)
-                extraPaddingStart = 24.dp
+
-                extraPaddingEnd = 24.dp
+              // Warning
-                if (message.type == ChatMessageType.PROMPT_TEMPLATES) {
+              is ChatMessageWarning -> MessageBodyWarning(message = message)
-                  extraPaddingStart = 12.dp
+
-                  extraPaddingEnd = 12.dp
+              // Config values change.
+              is ChatMessageConfigValuesChange -> MessageBodyConfigUpdate(message = message)
+
+              // Prompt templates.
+              is ChatMessagePromptTemplates -> MessageBodyPromptTemplates(
+                message = message,
+                task = task,
+                onPromptClicked = { template ->
+                  onSendMessage(
+                    selectedModel,
+                    listOf(ChatMessageText(content = template.prompt, side = ChatSide.USER))
+                  )
+                })
+
+              // Non-system messages.
+              else -> {
+                // The bubble shape around the message body.
+                var messageBubbleModifier = Modifier
+                  .clip(
+                    MessageBubbleShape(
+                      radius = bubbleBorderRadius, hardCornerAtLeftOrRight = hardCornerAtLeftOrRight
+                    )
+                  )
+                  .background(backgroundColor)
+                if (message is ChatMessageText) {
+                  messageBubbleModifier = messageBubbleModifier.pointerInput(Unit) {
+                    detectTapGestures(
+                      onLongPress = {
+                        haptic.performHapticFeedback(HapticFeedbackType.LongPress)
+                        longPressedMessage.value = message
+                        showMessageLongPressedSheet = true
+                      },
+                    )
+                  }
                 }
-              }
+                Box(
-              if (message.type == ChatMessageType.IMAGE) {
+                  modifier = messageBubbleModifier,
-                backgroundColor = Color.Transparent
+                ) {
-              }
+                  when (message) {
-              val bubbleBorderRadius = dimensionResource(R.dimen.chat_bubble_corner_radius)
+                    // Text
+                    is ChatMessageText -> MessageBodyText(message = message)
 
-              Column(
+                    // Image
-                modifier = Modifier
+                    is ChatMessageImage -> {
-                  .fillMaxWidth()
+                      MessageBodyImage(message = message, modifier = Modifier.clickable {
-                  .padding(
+                        selectedImageMessage = message
-                    start = 12.dp + extraPaddingStart,
+                      })
-                    end = 12.dp + extraPaddingEnd,
-                    top = 6.dp,
-                    bottom = 6.dp,
-                  ),
-                horizontalAlignment = hAlign,
-              ) {
-                // Sender row.
-                MessageSender(
-                  message = message,
-                  agentNameRes = task.agentNameRes,
-                  imageHistoryCurIndex = imageHistoryCurIndex.intValue
-                )
-
-                // Message body.
-                when (message) {
-                  // Loading.
-                  is ChatMessageLoading -> MessageBodyLoading()
-
-                  // Info.
-                  is ChatMessageInfo -> MessageBodyInfo(message = message)
-
-                  // Warning
-                  is ChatMessageWarning -> MessageBodyWarning(message = message)
-
-                  // Config values change.
-                  is ChatMessageConfigValuesChange -> MessageBodyConfigUpdate(message = message)
-
-                  // Prompt templates.
-                  is ChatMessagePromptTemplates -> MessageBodyPromptTemplates(message = message,
-                    task = task,
-                    onPromptClicked = { template ->
-                      onSendMessage(
-                        selectedModel,
-                        listOf(ChatMessageText(content = template.prompt, side = ChatSide.USER))
-                      )
-                    })
-
-                  // Non-system messages.
-                  else -> {
-                    // The bubble shape around the message body.
-                    var messageBubbleModifier = Modifier
-                      .clip(
-                        MessageBubbleShape(
-                          radius = bubbleBorderRadius,
-                          hardCornerAtLeftOrRight = hardCornerAtLeftOrRight
-                        )
-                      )
-                      .background(backgroundColor)
-                    if (message is ChatMessageText) {
-                      messageBubbleModifier = messageBubbleModifier.pointerInput(Unit) {
-                        detectTapGestures(
-                          onLongPress = {
-                            haptic.performHapticFeedback(HapticFeedbackType.LongPress)
-                            longPressedMessage.value = message
-                            showMessageLongPressedSheet = true
-                          },
-                        )
-                      }
                     }
-                    Box(
-                      modifier = messageBubbleModifier,
-                    ) {
-                      when (message) {
-                        // Text
-                        is ChatMessageText -> MessageBodyText(message = message)
 
-                        // Image
+                    // Image with history (for image gen)
-                        is ChatMessageImage -> {
+                    is ChatMessageImageWithHistory -> MessageBodyImageWithHistory(
-                          if (targetSelectedImageMessage != message) {
+                      message = message, imageHistoryCurIndex = imageHistoryCurIndex
-                            MessageBodyImage(
+                    )
-                              message = message,
+
-                              modifier = Modifier
+                    // Classification result
-                                .clickable {
+                    is ChatMessageClassification -> MessageBodyClassification(
-                                  selectedImageMessage = message
+                      message = message, modifier = Modifier.width(
-                                }
+                        message.maxBarWidth ?: CLASSIFICATION_BAR_MAX_WIDTH
-                                .sharedElement(
+                      )
-                                  sharedContentState = rememberSharedContentState(key = "selected_image"),
+                    )
-                                  animatedVisibilityScope = this@AnimatedContent,
+
-                                  clipInOverlayDuringTransition = OverlayClip(
+                    // Benchmark result.
-                                    MessageBubbleShape(
+                    is ChatMessageBenchmarkResult -> MessageBodyBenchmark(message = message)
-                                      radius = bubbleBorderRadius
+
-                                    )
+                    // Benchmark LLM result.
-                                  )
+                    is ChatMessageBenchmarkLlmResult -> MessageBodyBenchmarkLlm(
-                                ),
+                      message = message, modifier = Modifier.wrapContentWidth()
+                    )
+
+                    else -> {}
+                  }
+                }
+
+                if (message.side == ChatSide.AGENT) {
+                  Row(
+                    verticalAlignment = Alignment.CenterVertically,
+                    horizontalArrangement = Arrangement.spacedBy(8.dp),
+                  ) {
+                    LatencyText(message = message)
+                    // A button to show stats for the LLM message.
+                    if (task.type.id.startsWith("llm_") && message is ChatMessageText
+                      // This means we only want to show the action button when the message is done
+                      // generating, at which point the latency will be set.
+                      && message.latencyMs >= 0
+                    ) {
+                      val showingStats =
+                        viewModel.isShowingStats(model = selectedModel, message = message)
+                      MessageActionButton(
+                        label = if (showingStats) "Hide stats" else "Show stats",
+                        icon = Icons.Outlined.Timer,
+                        onClick = {
+                          // Toggle showing stats.
+                          viewModel.toggleShowingStats(selectedModel, message)
+
+                          // Add the stats message after the LLM message.
+                          if (viewModel.isShowingStats(
+                              model = selectedModel, message = message
+                            )
+                          ) {
+                            val llmBenchmarkResult = message.llmBenchmarkResult
+                            if (llmBenchmarkResult != null) {
+                              viewModel.insertMessageAfter(
+                                model = selectedModel,
+                                anchorMessage = message,
+                                messageToAdd = llmBenchmarkResult,
+                              )
+                            }
+                          }
+                          // Remove the stats message.
+                          else {
+                            val curMessageIndex = viewModel.getMessageIndex(
+                              model = selectedModel, message = message
+                            )
+                            viewModel.removeMessageAt(
+                              model = selectedModel, index = curMessageIndex + 1
                             )
                           }
-                        }
+                        },
-
+                        enabled = !uiState.inProgress
-                        // Image with history (for image gen)
+                      )
-                        is ChatMessageImageWithHistory -> MessageBodyImageWithHistory(
+                    }
-                          message = message, imageHistoryCurIndex = imageHistoryCurIndex
+                  }
-                        )
+                } else if (message.side == ChatSide.USER) {
-
+                  Row(
-                        // Classification result
+                    verticalAlignment = Alignment.CenterVertically,
-                        is ChatMessageClassification -> MessageBodyClassification(
+                    horizontalArrangement = Arrangement.spacedBy(4.dp)
-                          message = message, modifier = Modifier.width(
+                  ) {
-                            message.maxBarWidth ?: CLASSIFICATION_BAR_MAX_WIDTH
+                    // Run again button.
-                          )
+                    if (selectedModel.showRunAgainButton) {
-                        )
+                      MessageActionButton(
-
+                        label = stringResource(R.string.run_again),
-                        // Benchmark result.
+                        icon = Icons.Rounded.Refresh,
-                        is ChatMessageBenchmarkResult -> MessageBodyBenchmark(message = message)
+                        onClick = {
-
+                          onRunAgainClicked(selectedModel, message)
-                        // Benchmark LLM result.
+                        },
-                        is ChatMessageBenchmarkLlmResult -> MessageBodyBenchmarkLlm(
+                        enabled = !uiState.inProgress
-                          message = message, modifier = Modifier.wrapContentWidth()
+                      )
-                        )
-
-                        else -> {}
-                      }
                     }
-                    if (message.side == ChatSide.AGENT) {
-                      Row(
-                        verticalAlignment = Alignment.CenterVertically,
-                        horizontalArrangement = Arrangement.spacedBy(8.dp),
-                      ) {
-                        LatencyText(message = message)
-                        // A button to show stats for the LLM message.
-                        if (task.type.id.startsWith("llm_") && message is ChatMessageText
-                          // This means we only want to show the action button when the message is done
-                          // generating, at which point the latency will be set.
-                          && message.latencyMs >= 0
-                        ) {
-                          val showingStats =
-                            viewModel.isShowingStats(model = selectedModel, message = message)
-                          MessageActionButton(
-                            label = if (showingStats) "Hide stats" else "Show stats",
-                            icon = Icons.Outlined.Timer,
-                            onClick = {
-                              // Toggle showing stats.
-                              viewModel.toggleShowingStats(selectedModel, message)
 
-                              // Add the stats message after the LLM message.
+                    // Benchmark button
-                              if (viewModel.isShowingStats(
+                    if (selectedModel.showBenchmarkButton) {
-                                  model = selectedModel, message = message
+                      MessageActionButton(
-                                )
+                        label = stringResource(R.string.benchmark),
-                              ) {
+                        icon = Icons.Outlined.Timer,
-                                val llmBenchmarkResult = message.llmBenchmarkResult
+                        onClick = {
-                                if (llmBenchmarkResult != null) {
+                          showBenchmarkConfigsDialog = true
-                                  viewModel.insertMessageAfter(
+                          benchmarkMessage.value = message
-                                    model = selectedModel,
+                        },
-                                    anchorMessage = message,
+                        enabled = !uiState.inProgress
-                                    messageToAdd = llmBenchmarkResult,
+                      )
-                                  )
-                                }
-                              }
-                              // Remove the stats message.
-                              else {
-                                val curMessageIndex =
-                                  viewModel.getMessageIndex(
-                                    model = selectedModel,
-                                    message = message
-                                  )
-                                viewModel.removeMessageAt(
-                                  model = selectedModel, index = curMessageIndex + 1
-                                )
-                              }
-                            },
-                            enabled = !uiState.inProgress
-                          )
-                        }
-                      }
-                    } else if (message.side == ChatSide.USER) {
-                      Row(
-                        verticalAlignment = Alignment.CenterVertically,
-                        horizontalArrangement = Arrangement.spacedBy(4.dp)
-                      ) {
-                        // Run again button.
-                        if (selectedModel.showRunAgainButton) {
-                          MessageActionButton(
-                            label = stringResource(R.string.run_again),
-                            icon = Icons.Rounded.Refresh,
-                            onClick = {
-                              onRunAgainClicked(selectedModel, message)
-                            },
-                            enabled = !uiState.inProgress
-                          )
-                        }
-
-                        // Benchmark button
-                        if (selectedModel.showBenchmarkButton) {
-                          MessageActionButton(
-                            label = stringResource(R.string.benchmark),
-                            icon = Icons.Outlined.Timer,
-                            onClick = {
-                              showBenchmarkConfigsDialog = true
-                              benchmarkMessage.value = message
-                            },
-                            enabled = !uiState.inProgress
-                          )
-                        }
-                      }
                     }
                   }
                 }
               }
             }
           }
-
-          SnackbarHost(hostState = snackbarHostState, modifier = Modifier.padding(vertical = 4.dp))
-
-          // Show an info message for ask image task to get users started.
-          if (task.type == TaskType.LLM_ASK_IMAGE && messages.isEmpty()) {
-            Column(
-              modifier = Modifier
-                .padding(horizontal = 16.dp)
-                .fillMaxSize(),
-              horizontalAlignment = Alignment.CenterHorizontally,
-              verticalArrangement = Arrangement.Center
-            ) {
-              MessageBodyInfo(
-                ChatMessageInfo(content = "To get started, click + below to add an image and type a prompt to ask a question about it."),
-                smallFontSize = false
-              )
-            }
-          }
-        }
-
-        // Chat input
-        when (chatInputType) {
-          ChatInputType.TEXT -> {
-//        val isLlmTask = task.type == TaskType.LLM_CHAT
-//        val notLlmStartScreen = !(messages.size == 1 && messages[0] is ChatMessagePromptTemplates)
-            val hasImageMessage = messages.any { it is ChatMessageImage }
-            MessageInputText(
-              modelManagerViewModel = modelManagerViewModel,
-              curMessage = curMessage,
-              inProgress = uiState.inProgress,
-              isResettingSession = uiState.isResettingSession,
-              modelPreparing = uiState.preparing,
-              hasImageMessage = hasImageMessage,
-              modelInitializing = modelInitializationStatus?.status == ModelInitializationStatusType.INITIALIZING,
-              textFieldPlaceHolderRes = task.textInputPlaceHolderRes,
-              onValueChanged = { curMessage = it },
-              onSendMessage = {
-                onSendMessage(selectedModel, it)
-                curMessage = ""
-              },
-              onOpenPromptTemplatesClicked = {
-                onSendMessage(
-                  selectedModel, listOf(
-                    ChatMessagePromptTemplates(
-                      templates = selectedModel.llmPromptTemplates, showMakeYourOwn = false
-                    )
-                  )
-                )
-              },
-              onStopButtonClicked = onStopButtonClicked,
-//          showPromptTemplatesInMenu = isLlmTask && notLlmStartScreen,
-              showPromptTemplatesInMenu = false,
-              showImagePickerInMenu = selectedModel.llmSupportImage,
-              showStopButtonWhenInProgress = showStopButtonInInputWhenInProgress,
-            )
-          }
-
-          ChatInputType.IMAGE -> MessageInputImage(
-            disableButtons = uiState.inProgress,
-            streamingMessage = streamingMessage,
-            onImageSelected = { bitmap ->
-              onSendMessage(
-                selectedModel, listOf(
-                  ChatMessageImage(
-                    bitmap = bitmap, imageBitMap = bitmap.asImageBitmap(), side = ChatSide.USER
-                  )
-                )
-              )
-            },
-            onStreamImage = { bitmap ->
-              onStreamImageMessage(
-                selectedModel, ChatMessageImage(
-                  bitmap = bitmap, imageBitMap = bitmap.asImageBitmap(), side = ChatSide.USER
-                )
-              )
-            },
-            onStreamEnd = onStreamEnd,
-          )
         }
       }
 
-      // A full-screen image viewer.
+      SnackbarHost(hostState = snackbarHostState, modifier = Modifier.padding(vertical = 4.dp))
-      if (targetSelectedImageMessage != null) {
-        ZoomableBox(
-          modifier = Modifier
-            .fillMaxSize()
-            .background(Color.Black.copy(alpha = 0.9f))
-            .sharedElement(
-              rememberSharedContentState(key = "bounds"),
-              animatedVisibilityScope = this,
-            )
-            .skipToLookaheadSize(),
-        ) {
-          // Image.
-          Image(
-            bitmap = targetSelectedImageMessage.imageBitMap,
-            contentDescription = "",
-            modifier = modifier
-              .fillMaxSize()
-              .graphicsLayer(
-                scaleX = scale,
-                scaleY = scale,
-                translationX = offsetX,
-                translationY = offsetY
-              )
-              .sharedElement(
-                sharedContentState = rememberSharedContentState(key = "selected_image"),
-                animatedVisibilityScope = this@AnimatedContent,
-              ),
-            contentScale = ContentScale.Fit,
-          )
 
-          // Close button.
+      // Show an info message for ask image task to get users started.
-          IconButton(
+      if (task.type == TaskType.LLM_ASK_IMAGE && messages.isEmpty()) {
-            onClick = {
+        Column(
-              selectedImageMessage = null
+          modifier = Modifier
-            },
+            .padding(horizontal = 16.dp)
-            colors = IconButtonDefaults.iconButtonColors(
+            .fillMaxSize(),
-              containerColor = MaterialTheme.colorScheme.surfaceVariant,
+          horizontalAlignment = Alignment.CenterHorizontally,
-            ),
+          verticalArrangement = Arrangement.Center
-            modifier = Modifier.offset(x = (-8).dp, y = 8.dp)
+        ) {
-          ) {
+          MessageBodyInfo(
-            Icon(
+            ChatMessageInfo(content = "To get started, click + below to add an image and type a prompt to ask a question about it."),
-              Icons.Rounded.Close,
+            smallFontSize = false
-              contentDescription = "",
+          )
-              tint = MaterialTheme.colorScheme.primary
-            )
-          }
         }
       }
     }
+
+    // Chat input
+    when (chatInputType) {
+      ChatInputType.TEXT -> {
+//        val isLlmTask = task.type == TaskType.LLM_CHAT
+//        val notLlmStartScreen = !(messages.size == 1 && messages[0] is ChatMessagePromptTemplates)
+        MessageInputText(
+          modelManagerViewModel = modelManagerViewModel,
+          curMessage = curMessage,
+          inProgress = uiState.inProgress,
+          isResettingSession = uiState.isResettingSession,
+          modelPreparing = uiState.preparing,
+          hasImageMessage = hasImageMessageToLastConfigChange,
+          modelInitializing = modelInitializationStatus?.status == ModelInitializationStatusType.INITIALIZING,
+          textFieldPlaceHolderRes = task.textInputPlaceHolderRes,
+          onValueChanged = { curMessage = it },
+          onSendMessage = {
+            onSendMessage(selectedModel, it)
+            curMessage = ""
+          },
+          onOpenPromptTemplatesClicked = {
+            onSendMessage(
+              selectedModel, listOf(
+                ChatMessagePromptTemplates(
+                  templates = selectedModel.llmPromptTemplates, showMakeYourOwn = false
+                )
+              )
+            )
+          },
+          onStopButtonClicked = onStopButtonClicked,
+//          showPromptTemplatesInMenu = isLlmTask && notLlmStartScreen,
+          showPromptTemplatesInMenu = false,
+          showImagePickerInMenu = selectedModel.llmSupportImage,
+          showStopButtonWhenInProgress = showStopButtonInInputWhenInProgress,
+        )
+      }
+
+      ChatInputType.IMAGE -> MessageInputImage(
+        disableButtons = uiState.inProgress,
+        streamingMessage = streamingMessage,
+        onImageSelected = { bitmap ->
+          onSendMessage(
+            selectedModel, listOf(
+              ChatMessageImage(
+                bitmap = bitmap, imageBitMap = bitmap.asImageBitmap(), side = ChatSide.USER
+              )
+            )
+          )
+        },
+        onStreamImage = { bitmap ->
+          onStreamImageMessage(
+            selectedModel, ChatMessageImage(
+              bitmap = bitmap, imageBitMap = bitmap.asImageBitmap(), side = ChatSide.USER
+            )
+          )
+        },
+        onStreamEnd = onStreamEnd,
+      )
+    }
+  }
+
+  // A full-screen image viewer.
+  val curSelectedImageMessage = selectedImageMessage
+  AnimatedVisibility(
+    visible = curSelectedImageMessage != null,
+    enter = fadeIn(),
+    exit = fadeOut(),
+  ) {
+    if (curSelectedImageMessage == null) return@AnimatedVisibility
+
+    ZoomableBox(
+      modifier = Modifier
+        .fillMaxSize()
+        .background(Color.Black.copy(alpha = 0.9f))
+    ) {
+      // Image.
+      Image(
+        bitmap = curSelectedImageMessage.imageBitMap,
+        contentDescription = "",
+        modifier = modifier
+          .fillMaxSize()
+          .graphicsLayer(
+            scaleX = scale, scaleY = scale, translationX = offsetX, translationY = offsetY
+          ),
+        contentScale = ContentScale.Fit,
+      )
+
+      // Close button.
+      IconButton(
+        onClick = {
+          selectedImageMessage = null
+        }, colors = IconButtonDefaults.iconButtonColors(
+          containerColor = MaterialTheme.colorScheme.surfaceVariant,
+        ), modifier = Modifier.offset(x = (-8).dp, y = 8.dp)
+      ) {
+        Icon(
+          Icons.Rounded.Close, contentDescription = "", tint = MaterialTheme.colorScheme.primary
+        )
+      }
+    }
   }
 
   // Error dialog.
@@ -635,7 +618,7 @@ fun ChatPanel(
     }
   }
 
-  // Benchmark config dialog.
+// Benchmark config dialog.
   if (showBenchmarkConfigsDialog) {
     BenchmarkConfigDialog(onDismissed = { showBenchmarkConfigsDialog = false },
       messageToBenchmark = benchmarkMessage.value,
@@ -644,7 +627,7 @@ fun ChatPanel(
       })
   }
 
-  // Sheet to show when a message is long-pressed.
+// Sheet to show when a message is long-pressed.
   if (showMessageLongPressedSheet) {
     val message = longPressedMessage.value
     if (message != null && message is ChatMessageText) {
@@ -700,22 +683,20 @@ fun ZoomableBox(
   var offsetX by remember { mutableFloatStateOf(0f) }
   var offsetY by remember { mutableFloatStateOf(0f) }
   var size by remember { mutableStateOf(IntSize.Zero) }
-  Box(
+  Box(modifier = modifier
-    modifier = modifier
+    .clip(RectangleShape)
-      .clip(RectangleShape)
+    .onSizeChanged { size = it }
-      .onSizeChanged { size = it }
+    .pointerInput(Unit) {
-      .pointerInput(Unit) {
+      detectTransformGestures { _, pan, zoom, _ ->
-        detectTransformGestures { _, pan, zoom, _ ->
+        scale = maxOf(minScale, minOf(scale * zoom, maxScale))
-          scale = maxOf(minScale, minOf(scale * zoom, maxScale))
+        val maxX = (size.width * (scale - 1)) / 2
-          val maxX = (size.width * (scale - 1)) / 2
+        val minX = -maxX
-          val minX = -maxX
+        offsetX = maxOf(minX, minOf(maxX, offsetX + pan.x))
-          offsetX = maxOf(minX, minOf(maxX, offsetX + pan.x))
+        val maxY = (size.height * (scale - 1)) / 2
-          val maxY = (size.height * (scale - 1)) / 2
+        val minY = -maxY
-          val minY = -maxY
+        offsetY = maxOf(minY, minOf(maxY, offsetY + pan.y))
-          offsetY = maxOf(minY, minOf(maxY, offsetY + pan.y))
+      }
-        }
+    }, contentAlignment = Alignment.TopEnd
-      },
-    contentAlignment = Alignment.TopEnd
   ) {
     val scope = ZoomableBoxScopeImpl(scale, offsetX, offsetY)
     scope.content()
@@ -729,9 +710,7 @@ interface ZoomableBoxScope {
 }
 
 private data class ZoomableBoxScopeImpl(
-  override val scale: Float,
+  override val scale: Float, override val offsetX: Float, override val offsetY: Float
-  override val offsetX: Float,
-  override val offsetY: Float
 ) : ZoomableBoxScope
 
 @Preview(showBackground = true)

Android/src/app/src/main/java/com/google/aiedge/gallery/ui/home/HomeScreen.kt

@@ -16,7 +16,6 @@
 
 package com.google.aiedge.gallery.ui.home
 
-import android.app.Activity
 import android.content.Context
 import android.content.Intent
 import android.net.Uri
@@ -354,7 +353,7 @@ private fun TaskList(
   val linkColor = MaterialTheme.customColors.linkColor
 
   val introText = buildAnnotatedString {
-    append("Welcome to AI Edge Gallery! Explore a world of \namazing on-device models from ")
+    append("Welcome to Google AI Edge Gallery! Explore a world of \namazing on-device models from ")
     withLink(
       link = LinkAnnotation.Url(
         url = "https://huggingface.co/litert-community", // Replace with the actual URL

Android/src/app/src/main/java/com/google/aiedge/gallery/ui/llmchat/LlmChatModelHelper.kt

@@ -42,6 +42,7 @@ object LlmChatModelHelper {
   fun initialize(
     context: Context, model: Model, onDone: (String) -> Unit
   ) {
+    // Prepare options.
     val maxTokens =
       model.getIntConfigValue(key = ConfigKey.MAX_TOKENS, defaultValue = DEFAULT_MAX_TOKEN)
     val topK = model.getIntConfigValue(key = ConfigKey.TOPK, defaultValue = DEFAULT_TOPK)
@@ -62,7 +63,7 @@ object LlmChatModelHelper {
         .setMaxNumImages(if (model.llmSupportImage) 1 else 0)
         .build()
 
-    // Create an instance of the LLM Inference task
+    // Create an instance of the LLM Inference task and session.
     try {
       val llmInference = LlmInference.createFromOptions(context, options)
 
@@ -145,6 +146,9 @@ object LlmChatModelHelper {
     }
 
     // Start async inference.
+    //
+    // For a model that supports image modality, we need to add the text query chunk before adding
+    // image.
     val session = instance.session
     session.addQueryChunk(input)
     if (image != null) {

Android/src/app/src/main/java/com/google/aiedge/gallery/ui/theme/Theme.kt

@@ -205,7 +205,8 @@ fun GalleryTheme(
   content: @Composable () -> Unit
 ) {
   val themeOverride = ThemeSettings.themeOverride
-  val darkTheme: Boolean = isSystemInDarkTheme() || themeOverride.value == THEME_DARK
+  val darkTheme: Boolean =
+    (isSystemInDarkTheme() || themeOverride.value == THEME_DARK) && themeOverride.value != THEME_LIGHT
 
   StatusBarColorController(useDarkTheme = darkTheme)
 

Android/src/app/src/main/res/values/strings.xml

@@ -15,7 +15,7 @@
 -->
 
 <resources>
-    <string name="app_name">AI Edge Gallery</string>
+    <string name="app_name">Google AI Edge Gallery</string>
     <string name="model_manager">Model Manager</string>
     <string name="downloaded_size">%1$s downloaded</string>
     <string name="cancel">Cancel</string>