Refactor Imagen sample with improved UI and state management

bentrengrove · bentrengrove · commit a3de5a759251 · 2025-06-25T09:22:38.000+10:00
* Refactor the Imagen sample to use a sealed UI state (`Initial`, `Loading`, `ImageGenerated`, `Error`)
* Break down `ImagenScreen` into smaller, reusable composables with previews
* Add error handling for the image generation API call
* Improve keyboard handling by using `adjustResize` and IME insets
* Allow image generation via the keyboard's 'Send' action
* Update dependencies to latest versions
diff --git a/ai-catalog/app/src/main/AndroidManifest.xml b/ai-catalog/app/src/main/AndroidManifest.xml
@@ -17,6 +17,7 @@
         <activity
             android:name=".MainActivity"
             android:exported="true"
+            android:windowSoftInputMode="adjustResize"
             android:theme="@style/Theme.AISampleCatalog">
             <intent-filter>
                 <action android:name="android.intent.action.MAIN" />
diff --git a/ai-catalog/gradle/libs.versions.toml b/ai-catalog/gradle/libs.versions.toml
@@ -1,5 +1,5 @@
 [versions]
-agp = "8.8.0"
+agp = "8.8.2"
 coilCompose = "3.1.0"
 firebaseBom = "33.14.0"
 mlkitGenAi = "1.0.0-beta1"
@@ -11,13 +11,13 @@ espressoCore = "3.6.1"
 kotlinxCoroutinesGuava = "1.10.2"
 kotlinxSerializationJson = "1.6.2"
 lifecycleRuntimeKtx = "2.8.7"
-activityCompose = "1.10.0"
-composeBom = "2025.01.00"
-navigationCompose = "2.8.5"
-navigationRuntimeKtx = "2.8.5"
+activityCompose = "1.10.1"
+composeBom = "2025.06.01"
+navigationCompose = "2.9.0"
+navigationRuntimeKtx = "2.9.0"
 appcompat = "1.7.0"
 googleGmsGoogleServices = "4.4.2"
-hilt = "2.51.1"
+hilt = "2.56.2"
 hiltNavigationCompose = "1.2.0"
 ksp = "2.1.0-1.0.29"
 runtimeLivedata = "1.7.6"
@@ -26,6 +26,8 @@ media3 = "1.6.1"
 firebaseCommonKtx = "21.0.0"
 uiToolingPreviewAndroid = "1.8.1"
 spotless = "7.0.4"
+uiToolingPreview = "1.8.3"
+uiTooling = "1.8.3"
 
 [libraries]
 androidx-core-ktx = { group = "androidx.core", name = "core-ktx", version.ref = "coreKtx" }
@@ -64,6 +66,8 @@ androidx-material3-android = { group = "androidx.compose.material3", name = "mat
 androidx-media3-exoplayer = { module = "androidx.media3:media3-exoplayer", version.ref = "media3" }
 androidx-media3-ui = { module = "androidx.media3:media3-ui", version.ref = "media3" }
 androidx-ui-tooling-preview-android = { group = "androidx.compose.ui", name = "ui-tooling-preview-android", version.ref = "uiToolingPreviewAndroid" }
+ui-tooling-preview = { group = "androidx.compose.ui", name = "ui-tooling-preview", version.ref = "uiToolingPreview" }
+ui-tooling = { group = "androidx.compose.ui", name = "ui-tooling", version.ref = "uiTooling" }
 
 [plugins]
 android-application = { id = "com.android.application", version.ref = "agp" }
diff --git a/ai-catalog/samples/imagen/build.gradle.kts b/ai-catalog/samples/imagen/build.gradle.kts
@@ -44,16 +44,19 @@ android {
             )
         }
     }
+
     compileOptions {
         sourceCompatibility = JavaVersion.VERSION_17
         targetCompatibility = JavaVersion.VERSION_17
     }
-    composeOptions {
-        kotlinCompilerExtensionVersion = "1.5.15"
-    }
+
     kotlinOptions {
         jvmTarget = "17"
     }
+
+    lint {
+        warningsAsErrors = true
+    }
 }
 
 dependencies {
@@ -68,6 +71,8 @@ dependencies {
     implementation(libs.hilt.android)
     implementation(libs.hilt.navigation.compose)
     implementation(libs.androidx.runtime.livedata)
+    implementation(libs.ui.tooling.preview)
+    debugImplementation(libs.ui.tooling)
     ksp(libs.hilt.compiler)
 
     testImplementation(libs.junit)
diff --git a/ai-catalog/samples/imagen/src/main/java/com/android/ai/samples/imagen/ImagenScreen.kt b/ai-catalog/samples/imagen/src/main/java/com/android/ai/samples/imagen/ImagenScreen.kt
@@ -15,25 +15,30 @@
  */
 package com.android.ai.samples.imagen
 
-import android.content.Context
 import android.content.Intent
-import android.net.Uri
 import androidx.compose.foundation.Image
+import androidx.compose.foundation.layout.Arrangement.Absolute.spacedBy
 import androidx.compose.foundation.layout.Column
-import androidx.compose.foundation.layout.Row
 import androidx.compose.foundation.layout.Spacer
+import androidx.compose.foundation.layout.WindowInsets
+import androidx.compose.foundation.layout.aspectRatio
 import androidx.compose.foundation.layout.fillMaxSize
 import androidx.compose.foundation.layout.fillMaxWidth
 import androidx.compose.foundation.layout.height
+import androidx.compose.foundation.layout.ime
 import androidx.compose.foundation.layout.padding
 import androidx.compose.foundation.layout.size
+import androidx.compose.foundation.layout.windowInsetsBottomHeight
 import androidx.compose.foundation.layout.wrapContentSize
 import androidx.compose.foundation.rememberScrollState
+import androidx.compose.foundation.text.KeyboardActions
+import androidx.compose.foundation.text.KeyboardOptions
 import androidx.compose.foundation.verticalScroll
 import androidx.compose.material.icons.Icons
 import androidx.compose.material.icons.filled.Code
 import androidx.compose.material.icons.filled.SmartToy
 import androidx.compose.material3.Button
+import androidx.compose.material3.ButtonDefaults
 import androidx.compose.material3.Card
 import androidx.compose.material3.ExperimentalMaterial3Api
 import androidx.compose.material3.Icon
@@ -46,30 +51,40 @@ import androidx.compose.material3.TopAppBarDefaults.topAppBarColors
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.collectAsState
 import androidx.compose.runtime.getValue
-import androidx.compose.runtime.livedata.observeAsState
 import androidx.compose.runtime.mutableStateOf
-import androidx.compose.runtime.remember
+import androidx.compose.runtime.saveable.rememberSaveable
 import androidx.compose.runtime.setValue
 import androidx.compose.ui.Alignment
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.graphics.asImageBitmap
 import androidx.compose.ui.layout.ContentScale
 import androidx.compose.ui.platform.LocalContext
 import androidx.compose.ui.res.stringResource
+import androidx.compose.ui.text.input.ImeAction
 import androidx.compose.ui.text.style.TextAlign
+import androidx.compose.ui.tooling.preview.Preview
 import androidx.compose.ui.unit.dp
-import androidx.compose.ui.unit.sp
+import androidx.core.net.toUri
 import androidx.hilt.navigation.compose.hiltViewModel
 
 @OptIn(ExperimentalMaterial3Api::class)
 @Composable
 fun ImagenScreen(viewModel: ImagenViewModel = hiltViewModel()) {
-    val context = LocalContext.current
-    val isGenerating by viewModel.isGenerating.observeAsState(false)
-    val generatedBitmap by viewModel.imageGenerated.collectAsState()
+    val uiState: ImagenUIState by viewModel.uiState.collectAsState()
 
-    val placeholder = stringResource(R.string.placeholder_prompt)
-    var editTextValue by remember { mutableStateOf(placeholder) }
+    ImagenScreen(
+        uiState = uiState,
+        onGenerateClick = viewModel::generateImage,
+    )
+}
+
+@Composable
+@OptIn(ExperimentalMaterial3Api::class)
+private fun ImagenScreen(
+    uiState: ImagenUIState,
+    onGenerateClick: (String) -> Unit
+) {
+    val isGenerating = uiState is ImagenUIState.Loading
 
     Scaffold(
         modifier = Modifier,
@@ -83,79 +98,173 @@ fun ImagenScreen(viewModel: ImagenViewModel = hiltViewModel()) {
                     Text(text = stringResource(R.string.title_image_generation_screen))
                 },
                 actions = {
-                    SeeCodeButton(context)
+                    SeeCodeButton()
                 },
             )
         },
     ) { innerPadding ->
         Column(
             Modifier
-                .padding(12.dp)
                 .verticalScroll(rememberScrollState())
+                .padding(16.dp)
                 .padding(innerPadding),
         ) {
-            Card(
-                modifier = Modifier.size(
-                    width = 400.dp,
-                    height = 400.dp,
-                ).align(Alignment.CenterHorizontally),
-            ) {
-                generatedBitmap?.let {
-                    Image(
-                        bitmap = it.asImageBitmap(),
-                        contentDescription = "Picture",
-                        contentScale = ContentScale.Fit,
-                        modifier = Modifier.fillMaxSize(),
-                    )
-                }
-                if (isGenerating) {
-                    Text(
-                        text = stringResource(R.string.generating_label),
-                        modifier = Modifier
-                            .fillMaxSize()
-                            .wrapContentSize(Alignment.Center),
-                        textAlign = TextAlign.Center,
-                    )
-                }
-            }
-            Spacer(modifier = Modifier.height(24.dp))
-            TextField(
-                value = editTextValue,
-                onValueChange = { editTextValue = it },
-                label = { Text(stringResource(R.string.prompt_label)) },
-                modifier = Modifier.fillMaxWidth(),
+            GeneratedContent(
+                uiState = uiState,
+                modifier = Modifier
+                    .fillMaxWidth()
+                    .aspectRatio(1f)
             )
-            Row {
-                Button(
-                    modifier = Modifier.padding(vertical = 8.dp),
-                    onClick = {
-                        viewModel.generateImage(editTextValue)
-                    },
-                    enabled = !isGenerating,
-                ) {
-                    Icon(Icons.Default.SmartToy, contentDescription = "Robot")
-                    Text(modifier = Modifier.padding(start = 8.dp), text = stringResource(R.string.generate_button))
-                }
+
+            Spacer(modifier = Modifier.height(16.dp))
+
+            GenerationInput(
+                onGenerateClick = onGenerateClick,
+                enabled = !isGenerating,
+                modifier = Modifier.fillMaxSize(),
+            )
+
+            // Ensure the screen scrolls when the keyboard appears
+            Spacer(Modifier.windowInsetsBottomHeight(WindowInsets.ime))
+        }
+    }
+}
+
+@Composable
+private fun GeneratedContent(uiState: ImagenUIState, modifier: Modifier = Modifier) {
+    Card(
+        modifier = modifier,
+    ) {
+        when (uiState) {
+            ImagenUIState.Initial -> {
+                // no-op
+            }
+
+            ImagenUIState.Loading -> {
+                Text(
+                    text = stringResource(R.string.generating_label),
+                    modifier = Modifier
+                        .fillMaxSize()
+                        .wrapContentSize(Alignment.Center),
+                    textAlign = TextAlign.Center,
+                )
+            }
+
+            is ImagenUIState.ImageGenerated -> {
+                Image(
+                    bitmap = uiState.bitmap.asImageBitmap(),
+                    contentDescription = uiState.contentDescription,
+                    contentScale = ContentScale.Fit,
+                    modifier = Modifier.fillMaxSize(),
+                )
+            }
+
+            is ImagenUIState.Error -> {
+                Text(
+                    text = uiState.message,
+                    modifier = Modifier
+                        .fillMaxSize()
+                        .wrapContentSize(Alignment.Center),
+                    textAlign = TextAlign.Center,
+                )
             }
         }
     }
 }
 
 @Composable
-fun SeeCodeButton(context: Context) {
+private fun GenerationInput(onGenerateClick: (String) -> Unit, enabled: Boolean, modifier: Modifier = Modifier) {
+    val placeholder = stringResource(R.string.placeholder_prompt)
+    var textFieldValue by rememberSaveable { mutableStateOf(placeholder) }
+
+    Column(
+        verticalArrangement = spacedBy(8.dp),
+        modifier = modifier,
+    ) {
+        TextField(
+            value = textFieldValue,
+            onValueChange = { textFieldValue = it },
+            label = { Text(stringResource(R.string.prompt_label)) },
+            modifier = Modifier.fillMaxWidth(),
+            enabled = enabled,
+            keyboardOptions = KeyboardOptions(imeAction = ImeAction.Send),
+            keyboardActions = KeyboardActions(
+                onSend = {
+                    onGenerateClick(textFieldValue)
+                },
+            ),
+        )
+        Button(
+            onClick = {
+                onGenerateClick(textFieldValue)
+            },
+            enabled = enabled,
+            contentPadding = ButtonDefaults.ButtonWithIconContentPadding,
+            modifier = modifier.fillMaxWidth(),
+        ) {
+            Icon(
+                Icons.Default.SmartToy,
+                contentDescription = null,
+                modifier = Modifier.size(ButtonDefaults.IconSize),
+            )
+            Spacer(Modifier.size(ButtonDefaults.IconSpacing))
+            Text(text = stringResource(R.string.generate_button))
+        }
+    }
+}
+
+@Composable
+private fun SeeCodeButton() {
+    val context = LocalContext.current
     val githubLink = "https://github.com/android/ai-samples/tree/main/ai-catalog/samples/imagen"
     Button(
         onClick = {
-            val intent = Intent(Intent.ACTION_VIEW, Uri.parse(githubLink))
+            val intent = Intent(Intent.ACTION_VIEW, githubLink.toUri())
             context.startActivity(intent)
         },
-        modifier = Modifier.padding(end = 8.dp),
+        contentPadding = ButtonDefaults.ButtonWithIconContentPadding,
     ) {
-        Icon(Icons.Filled.Code, contentDescription = "See code")
+        Icon(Icons.Filled.Code, contentDescription = null)
+        Spacer(Modifier.size(ButtonDefaults.IconSpacing))
         Text(
-            modifier = Modifier.padding(start = 8.dp),
-            fontSize = 12.sp,
             text = stringResource(R.string.see_code),
         )
     }
 }
+
+@Preview
+@Composable
+@OptIn(ExperimentalMaterial3Api::class)
+private fun ImagenScreenPreview() {
+    ImagenScreen(
+        uiState = ImagenUIState.Initial,
+        onGenerateClick = {},
+    )
+}
+
+@Preview
+@Composable
+private fun GeneratedContentPreview() {
+    GeneratedContent(
+        uiState = ImagenUIState.Initial,
+        modifier = Modifier.size(400.dp),
+    )
+}
+
+@Preview
+@Composable
+private fun GeneratedContentLoadingPreview() {
+    GeneratedContent(
+        uiState = ImagenUIState.Loading,
+        modifier = Modifier.size(400.dp),
+    )
+}
+
+@Preview
+@Composable
+private fun GenerationInputPreview() {
+    GenerationInput(
+        onGenerateClick = {},
+        enabled = true,
+    )
+}
diff --git a/ai-catalog/samples/imagen/src/main/java/com/android/ai/samples/imagen/ImagenViewModel.kt b/ai-catalog/samples/imagen/src/main/java/com/android/ai/samples/imagen/ImagenViewModel.kt
diff --git a/ai-catalog/samples/imagen/src/main/res/values/strings.xml b/ai-catalog/samples/imagen/src/main/res/values/strings.xml

Original file line number	Diff line number	Diff line change
`@@ -44,16 +44,19 @@ android {`
`44`	`44`	`)`
`45`	`45`	`}`
`46`	`46`	`}`
	`47`	`+`
`47`	`48`	`compileOptions {`
`48`	`49`	`sourceCompatibility = JavaVersion.VERSION_17`
`49`	`50`	`targetCompatibility = JavaVersion.VERSION_17`
`50`	`51`	`}`
`51`		`- composeOptions {`
`52`		`- kotlinCompilerExtensionVersion = "1.5.15"`
`53`		`- }`
	`52`	`+`
`54`	`53`	`kotlinOptions {`
`55`	`54`	`jvmTarget = "17"`
`56`	`55`	`}`
	`56`	`+`
	`57`	`+ lint {`
	`58`	`+ warningsAsErrors = true`
	`59`	`+ }`
`57`	`60`	`}`
`58`	`61`
`59`	`62`	`dependencies {`
`@@ -68,6 +71,8 @@ dependencies {`
`68`	`71`	`implementation(libs.hilt.android)`
`69`	`72`	`implementation(libs.hilt.navigation.compose)`
`70`	`73`	`implementation(libs.androidx.runtime.livedata)`
	`74`	`+ implementation(libs.ui.tooling.preview)`
	`75`	`+ debugImplementation(libs.ui.tooling)`
`71`	`76`	`ksp(libs.hilt.compiler)`
`72`	`77`
`73`	`78`	`testImplementation(libs.junit)`