Add on-device AI image enhancement feature to enhance images by MayuriKhinvasara · Pull Request #194 · android/socialite

MayuriKhinvasara · 2026-06-09T10:29:42Z

Add on-device AI image enhancement feature to enhance images on the Chat screen.

Major changes:

Integrate the Google Play Services Media Effect Enhancement library to support on-device AI processing.
Implement ImageEnhancementScreen and EnhancementViewModel to handle image enhancement features, including tonemapping, deblurring, and upscaling.
Add EnhancementUtils.kt providing coroutine-based wrappers for enhancement session management, bitmap processing, and module installation.

Minor changes:

Add an "AI Enhance" action button to image bubbles within the chat UI.
Update the ChatMessage model to include a unique ID and add support for updating media URIs in the ChatRepository and MessageDao.
Define navigation logic for the new enhancement pane in Main.kt and SocialiteNavigation.kt.
Update the README to include documentation for the Media Effect Enhancement integration.

…hat screen Major changes: - Integrate the Google Play Services Media Effect Enhancement library to support on-device AI processing. - Implement `ImageEnhancementScreen` and `EnhancementViewModel` to handle image enhancement features, including tonemapping, deblurring, and upscaling. - Add `EnhancementUtils.kt` providing coroutine-based wrappers for enhancement session management, bitmap processing, and module installation. Minor changes: - Add an "AI Enhance" action button to image bubbles within the chat UI. - Update the `ChatMessage` model to include a unique ID and add support for updating media URIs in the `ChatRepository` and `MessageDao`. - Define navigation logic for the new enhancement pane in `Main.kt` and `SocialiteNavigation.kt`. - Update the README to include documentation for the Media Effect Enhancement integration.

gemini-code-assist

Code Review

This pull request adds on-device AI image enhancement features (tonemapping, deblurring, and upscaling) using Google Play Services, introducing a new enhancement screen, ViewModel, and database updates to persist enhanced image URIs. The review feedback identifies several critical issues: a data loss bug where saving an enhanced image deletes all other enhanced images in the chat history, thread safety issues with concurrent access to enhancementSession, and incorrect coroutine continuation handling during module installation. Additionally, the reviewer recommends removing silent background module downloads to conserve data, avoiding passing short-lived Context objects to the ViewModel to prevent memory leaks, and catching Throwable to safely handle OutOfMemoryError during bitmap decoding.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

gemini-code-assist · 2026-06-09T10:31:51Z

+            override fun onInstalled() {
+                Log.d("EnhancementUtils", "onInstalled")
+            }


Resume the coroutine continuation with true when the module is successfully installed. Currently, the continuation is not resumed here, which means if the installation takes time, the caller is never notified of the actual completion via this callback.

Suggested change

override fun onInstalled() {

Log.d("EnhancementUtils", "onInstalled")

}

override fun onInstalled() {

Log.d("EnhancementUtils", "onInstalled")

if (continuation.isActive) continuation.resume(true)

}

Done. Added continuation.resume(true) to ensure the caller is correctly notified when the module finishes installing.

gemini-code-assist · 2026-06-09T10:31:51Z

+            .addOnSuccessListener { result ->
+                if (continuation.isActive) continuation.resume(result)
+            }


Only resume the continuation immediately if the task result is true (indicating the module is already installed). If the result is false, the installation has just been initiated, and we must wait for the onInstalled() callback to resume the continuation. Resuming immediately on task success causes the ViewModel to prematurely mark the module as ready while it is still downloading in the background.

.addOnSuccessListener { result -> if (result && continuation.isActive) { continuation.resume(true) } }

Done. Updated the listener to only resume immediately if the result is true. Otherwise, we properly wait for the onInstalled() callback.

gemini-code-assist · 2026-06-09T10:31:51Z

+                // First check if it's already installed
+                val installed = client.isModuleInstalledAsync()
+                if (!installed) {
+                    // If not installed, attempt to install it before querying device support
+                    try {
+                        client.installModuleAsync { progress ->
+                            // Silent background install progress
+                        }
+                    } catch (e: Exception) {
+                        Log.e("EnhancementSupport", "Failed to silently install module", e)
+                    }
+                }


Remove the silent background installation from checkSupport. This function is called on the main chat screen to quickly determine if the "AI Enhance" button should be shown. Silently downloading and installing a large AI module in the background without user consent can consume significant cellular data and battery. Additionally, isDeviceSupportedAsync() can be queried directly without downloading the module first. The module installation is already handled with a proper progress UI when the user explicitly navigates to the ImageEnhancementScreen.

Done. Removed the silent background installation. The module is now only installed when the user explicitly navigates to the enhancement screen.

gemini-code-assist · 2026-06-09T10:31:51Z

+    fun saveEnhancedImage(messageId: Long, onComplete: () -> Unit) {
+        val enhancedBitmap = _uiState.value.enhancedImage?.bitmap ?: return
+        viewModelScope.launch(Dispatchers.IO) {
+            try {
+                val context = getApplication<Application>()
+                val directory = File(context.filesDir, "media")
+                if (!directory.exists()) {
+                    directory.mkdirs()
+                }
+
+                // Prevent storage bloat by deleting previous enhanced files
+                directory.listFiles()?.forEach { file ->
+                    if (file.name.startsWith("enhanced_")) {
+                        file.delete()
+                    }
+                }


Instead of deleting all files starting with "enhanced_", pass the oldUri to saveEnhancedImage and only delete that specific file if it was indeed a previously enhanced file for this message. Currently, this deletes enhanced images for all other messages in the chat history, causing broken image links and data loss.

Suggested change

fun saveEnhancedImage(messageId: Long, onComplete: () -> Unit) {

val enhancedBitmap = _uiState.value.enhancedImage?.bitmap ?: return

viewModelScope.launch(Dispatchers.IO) {

try {

val context = getApplication<Application>()

val directory = File(context.filesDir, "media")

if (!directory.exists()) {

directory.mkdirs()

}

// Prevent storage bloat by deleting previous enhanced files

directory.listFiles()?.forEach { file ->

if (file.name.startsWith("enhanced_")) {

file.delete()

}

}

fun saveEnhancedImage(messageId: Long, oldUri: String, onComplete: () -> Unit) {

val enhancedBitmap = _uiState.value.enhancedImage?.bitmap ?: return

viewModelScope.launch(Dispatchers.IO) {

try {

val context = getApplication<Application>()

val directory = File(context.filesDir, "media")

if (!directory.exists()) {

directory.mkdirs()

}

// Prevent storage bloat by deleting only the previous enhanced file for this message

val oldFileUri = Uri.parse(oldUri)

if (oldFileUri.scheme == "file") {

val oldFile = File(oldFileUri.path ?: "")

if (oldFile.exists() && oldFile.name.startsWith("enhanced_")) {

oldFile.delete()

}

}

Done. Now passing oldUri to selectively delete only the previously enhanced file for this specific message, preventing accidental deletion of others.

gemini-code-assist · 2026-06-09T10:31:51Z

+                }
+                Button(
+                    onClick = {
+                        enhancementViewModel.saveEnhancedImage(messageId, onFinishEditing)


Pass the current uri to saveEnhancedImage so that only the previous enhanced file for this specific message is deleted, preventing accidental deletion of enhanced images from other messages.

Suggested change

enhancementViewModel.saveEnhancedImage(messageId, onFinishEditing)

enhancementViewModel.saveEnhancedImage(messageId, uri, onFinishEditing)

Done. Passed the current uri into saveEnhancedImage to support the targeted deletion.

gemini-code-assist · 2026-06-09T10:31:51Z

+            enhancementSession?.release()
+            enhancementSession = null


Confine the release of enhancementSession to the Main thread to ensure thread safety.

Suggested change

enhancementSession?.release()

enhancementSession = null

// Release any previous session, since we have a new image. Confined to Main thread.

withContext(Dispatchers.Main) {

enhancementSession?.release()

enhancementSession = null

}

Done. Wrapped the session release in withContext(Dispatchers.Main) to synchronize access. This guarantees that our coroutines and the Android lifecycle's onCleared() method take turns accessing the enhancementSession variable one at a time, completely eliminating thread-safety issues and leaks without freezing the UI.

gemini-code-assist · 2026-06-09T10:31:51Z

+        super.onCleared()
+    }
+
+    fun onImageSelected(uri: Uri, context: Context) {


Remove the context parameter from onImageSelected. Passing short-lived contexts (like Activity context) to ViewModels is an architectural anti-pattern that can cause memory leaks. Since EnhancementViewModel is an AndroidViewModel, it can safely access the application context via getApplication().

Suggested change

fun onImageSelected(uri: Uri, context: Context) {

fun onImageSelected(uri: Uri) {

Done. Removed the context parameter and switched to using getApplication() to prevent memory leaks.

gemini-code-assist · 2026-06-09T10:31:51Z

+            enhancementSession?.release()
+            enhancementSession = null
+
+            val originalBitmap = decodeBitmapFromUri(uri, context)


Use the application context from getApplication() to decode the bitmap.

Suggested change

val originalBitmap = decodeBitmapFromUri(uri, context)

val originalBitmap = decodeBitmapFromUri(uri, getApplication())

Done. Updated to use getApplication() when decoding the bitmap.

gemini-code-assist · 2026-06-09T10:31:51Z

+
+    LaunchedEffect(uri) {
+        enhancementViewModel.setEnhancementMode(EnhancementMode.BITMAP)
+        enhancementViewModel.onImageSelected(Uri.parse(uri), context)


Remove the context argument when calling onImageSelected.

Suggested change

enhancementViewModel.onImageSelected(Uri.parse(uri), context)

enhancementViewModel.onImageSelected(Uri.parse(uri))

Done. Dropped the context argument here since the ViewModel handles it internally now.

gemini-code-assist · 2026-06-09T10:31:52Z

+            context.contentResolver.openInputStream(uri)?.use { inputStream ->
+                BitmapFactory.decodeStream(inputStream)
+            }
+        } catch (e: Exception) {


Catch Throwable instead of Exception to safely handle OutOfMemoryError during bitmap decoding. OutOfMemoryError is an Error (not an Exception), so it will bypass the current catch (e: Exception) block and crash the app if the decoded image is too large.

Suggested change

} catch (e: Exception) {

} catch (e: Throwable) {

Done. Changed the catch block to Throwable to safely catch and handle OutOfMemoryError crashes.

…hat screen Major changes: - Integrate the Google Play Services Media Effect Enhancement library to support on-device AI processing. - Implement `ImageEnhancementScreen` and `EnhancementViewModel` to handle image enhancement features, including tonemapping, deblurring, and upscaling. - Add `EnhancementUtils.kt` providing coroutine-based wrappers for enhancement session management, bitmap processing, and module installation. Minor changes: - Add an "AI Enhance" action button to image bubbles within the chat UI. - Update the `ChatMessage` model to include a unique ID and add support for updating media URIs in the `ChatRepository` and `MessageDao`. - Define navigation logic for the new enhancement pane in `Main.kt` and `SocialiteNavigation.kt`. - Update the README to include documentation for the Media Effect Enhancement integration.

…ncement # Conflicts: # app/build.gradle.kts # app/src/main/java/com/google/android/samples/socialite/repository/ChatRepository.kt # app/src/main/java/com/google/android/samples/socialite/ui/Main.kt # app/src/main/java/com/google/android/samples/socialite/ui/chat/component/MessageBubble.kt # app/src/main/java/com/google/android/samples/socialite/ui/mediaenhancement/EnhancementUtils.kt # app/src/main/java/com/google/android/samples/socialite/ui/mediaenhancement/EnhancementViewModel.kt # app/src/main/java/com/google/android/samples/socialite/ui/mediaenhancement/ImageEnhancementScreen.kt

…ncement # Conflicts: # app/src/main/java/com/google/android/samples/socialite/ui/mediaenhancement/EnhancementUtils.kt

MayuriKhinvasara requested review from SigmanZero, donovanfm and madebymozart as code owners June 9, 2026 10:29

MayuriKhinvasara removed request for donovanfm and madebymozart June 9, 2026 10:30

gemini-code-assist Bot reviewed Jun 9, 2026

View reviewed changes

MayuriKhinvasara added 6 commits June 9, 2026 22:51

Merge remote-tracking branch 'origin/mediaenhancement' into mediaenha…

3dc5007

…ncement # Conflicts: # app/src/main/java/com/google/android/samples/socialite/ui/mediaenhancement/EnhancementUtils.kt

Merge remote-tracking branch 'origin/mediaenhancement' into mediaenha…

98b6787

…ncement # Conflicts: # app/src/main/java/com/google/android/samples/socialite/ui/mediaenhancement/EnhancementUtils.kt

Merge branch 'evals' into mediaenhancement

e8e9750

SigmanZero approved these changes Jun 16, 2026

View reviewed changes

	enhancementViewModel.saveEnhancedImage(messageId, onFinishEditing)
	enhancementViewModel.saveEnhancedImage(messageId, uri, onFinishEditing)

	fun onImageSelected(uri: Uri, context: Context) {
	fun onImageSelected(uri: Uri) {

	val originalBitmap = decodeBitmapFromUri(uri, context)
	val originalBitmap = decodeBitmapFromUri(uri, getApplication())

	enhancementViewModel.onImageSelected(Uri.parse(uri), context)
	enhancementViewModel.onImageSelected(Uri.parse(uri))

Conversation

MayuriKhinvasara commented Jun 9, 2026

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Jun 9, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jun 9, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jun 9, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jun 9, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jun 9, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jun 9, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jun 9, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jun 9, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jun 9, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jun 9, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants