diff --git a/src/com/android/customization/picker/clock/ui/fragment/ClockSettingsFragment.kt b/src/com/android/customization/picker/clock/ui/fragment/ClockSettingsFragment.kt
index 068f44f..dc70633 100644
--- a/src/com/android/customization/picker/clock/ui/fragment/ClockSettingsFragment.kt
+++ b/src/com/android/customization/picker/clock/ui/fragment/ClockSettingsFragment.kt
@@ -88,14 +88,14 @@
                             injector
                                 .getCurrentWallpaperInfoFactory(context)
                                 .createCurrentWallpaperInfos(
-                                    { homeWallpaper, lockWallpaper, _ ->
-                                        continuation.resume(
-                                            lockWallpaper ?: homeWallpaper,
-                                            null,
-                                        )
-                                    },
+                                    context,
                                     forceReload,
-                                )
+                                ) { homeWallpaper, lockWallpaper, _ ->
+                                    continuation.resume(
+                                        lockWallpaper ?: homeWallpaper,
+                                        null,
+                                    )
+                                }
                         }
                     },
                     onWallpaperColorChanged = { colors ->
diff --git a/src/com/android/customization/picker/color/ui/fragment/ColorPickerFragment.kt b/src/com/android/customization/picker/color/ui/fragment/ColorPickerFragment.kt
index 911b821..941d2c6 100644
--- a/src/com/android/customization/picker/color/ui/fragment/ColorPickerFragment.kt
+++ b/src/com/android/customization/picker/color/ui/fragment/ColorPickerFragment.kt
@@ -114,23 +114,23 @@
                         wallpaperInfoProvider = { forceReload ->
                             suspendCancellableCoroutine { continuation ->
                                 wallpaperInfoFactory.createCurrentWallpaperInfos(
-                                    { homeWallpaper, lockWallpaper, _ ->
-                                        lifecycleScope.launch {
-                                            if (
-                                                wallpaperColorsRepository.lockWallpaperColors.value
-                                                    is WallpaperColorsModel.Loading
-                                            ) {
-                                                loadInitialColors(
-                                                    wallpaperManager,
-                                                    wallpaperColorsRepository,
-                                                    CustomizationSections.Screen.LOCK_SCREEN
-                                                )
-                                            }
-                                        }
-                                        continuation.resume(lockWallpaper ?: homeWallpaper, null)
-                                    },
+                                    context,
                                     forceReload,
-                                )
+                                ) { homeWallpaper, lockWallpaper, _ ->
+                                    lifecycleScope.launch {
+                                        if (
+                                            wallpaperColorsRepository.lockWallpaperColors.value
+                                                is WallpaperColorsModel.Loading
+                                        ) {
+                                            loadInitialColors(
+                                                wallpaperManager,
+                                                wallpaperColorsRepository,
+                                                CustomizationSections.Screen.LOCK_SCREEN
+                                            )
+                                        }
+                                    }
+                                    continuation.resume(lockWallpaper ?: homeWallpaper, null)
+                                }
                             }
                         },
                         onWallpaperColorChanged = { colors ->
@@ -165,23 +165,23 @@
                     wallpaperInfoProvider = { forceReload ->
                         suspendCancellableCoroutine { continuation ->
                             wallpaperInfoFactory.createCurrentWallpaperInfos(
-                                { homeWallpaper, lockWallpaper, _ ->
-                                    lifecycleScope.launch {
-                                        if (
-                                            wallpaperColorsRepository.homeWallpaperColors.value
-                                                is WallpaperColorsModel.Loading
-                                        ) {
-                                            loadInitialColors(
-                                                wallpaperManager,
-                                                wallpaperColorsRepository,
-                                                CustomizationSections.Screen.HOME_SCREEN
-                                            )
-                                        }
-                                    }
-                                    continuation.resume(homeWallpaper ?: lockWallpaper, null)
-                                },
+                                context,
                                 forceReload,
-                            )
+                            ) { homeWallpaper, lockWallpaper, _ ->
+                                lifecycleScope.launch {
+                                    if (
+                                        wallpaperColorsRepository.homeWallpaperColors.value
+                                            is WallpaperColorsModel.Loading
+                                    ) {
+                                        loadInitialColors(
+                                            wallpaperManager,
+                                            wallpaperColorsRepository,
+                                            CustomizationSections.Screen.HOME_SCREEN
+                                        )
+                                    }
+                                }
+                                continuation.resume(homeWallpaper ?: lockWallpaper, null)
+                            }
                         }
                     },
                     onWallpaperColorChanged = { colors ->
diff --git a/src/com/android/customization/picker/grid/ui/fragment/GridFragment.kt b/src/com/android/customization/picker/grid/ui/fragment/GridFragment.kt
index 52bc5e1..2a301b4 100644
--- a/src/com/android/customization/picker/grid/ui/fragment/GridFragment.kt
+++ b/src/com/android/customization/picker/grid/ui/fragment/GridFragment.kt
@@ -185,11 +185,11 @@
                     wallpaperInfoProvider = {
                         suspendCancellableCoroutine { continuation ->
                             wallpaperInfoFactory.createCurrentWallpaperInfos(
-                                { homeWallpaper, lockWallpaper, _ ->
-                                    continuation.resume(homeWallpaper ?: lockWallpaper, null)
-                                },
+                                context,
                                 /* forceRefresh= */ true,
-                            )
+                            ) { homeWallpaper, lockWallpaper, _ ->
+                                continuation.resume(homeWallpaper ?: lockWallpaper, null)
+                            }
                         }
                     },
                     wallpaperInteractor = wallpaperInteractor,
diff --git a/src/com/android/customization/picker/preview/ui/section/PreviewWithThemeSectionController.kt b/src/com/android/customization/picker/preview/ui/section/PreviewWithThemeSectionController.kt
index 0edc216..b3e778b 100644
--- a/src/com/android/customization/picker/preview/ui/section/PreviewWithThemeSectionController.kt
+++ b/src/com/android/customization/picker/preview/ui/section/PreviewWithThemeSectionController.kt
@@ -35,6 +35,7 @@
 import com.android.wallpaper.picker.customization.ui.viewmodel.ScreenPreviewViewModel
 import com.android.wallpaper.util.DisplayUtils
 import com.android.wallpaper.util.PreviewUtils
+import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.suspendCancellableCoroutine
 
 /**
@@ -69,6 +70,7 @@
         isTwoPaneAndSmallWidth,
         customizationPickerViewModel,
     ) {
+    @OptIn(ExperimentalCoroutinesApi::class)
     override fun createScreenPreviewViewModel(context: Context): ScreenPreviewViewModel {
         return PreviewWithThemeViewModel(
             previewUtils =
@@ -92,21 +94,21 @@
             wallpaperInfoProvider = { forceReload ->
                 suspendCancellableCoroutine { continuation ->
                     wallpaperInfoFactory.createCurrentWallpaperInfos(
-                        { homeWallpaper, lockWallpaper, _ ->
-                            val wallpaper =
-                                if (isOnLockScreen) {
-                                    lockWallpaper ?: homeWallpaper
-                                } else {
-                                    homeWallpaper ?: lockWallpaper
-                                }
-                            loadInitialColors(
-                                context = context,
-                                screen = screen,
-                            )
-                            continuation.resume(wallpaper, null)
-                        },
+                        context,
                         forceReload,
-                    )
+                    ) { homeWallpaper, lockWallpaper, _ ->
+                        val wallpaper =
+                            if (isOnLockScreen) {
+                                lockWallpaper ?: homeWallpaper
+                            } else {
+                                homeWallpaper ?: lockWallpaper
+                            }
+                        loadInitialColors(
+                            context = context,
+                            screen = screen,
+                        )
+                        continuation.resume(wallpaper, null)
+                    }
                 }
             },
             onWallpaperColorChanged = { colors ->
diff --git a/src/com/android/customization/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModel.kt b/src/com/android/customization/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModel.kt
index f832cde..809d20e 100644
--- a/src/com/android/customization/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModel.kt
+++ b/src/com/android/customization/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModel.kt
@@ -92,11 +92,11 @@
             wallpaperInfoProvider = { forceReload ->
                 suspendCancellableCoroutine { continuation ->
                     wallpaperInfoFactory.createCurrentWallpaperInfos(
-                        { homeWallpaper, lockWallpaper, _ ->
-                            continuation.resume(lockWallpaper ?: homeWallpaper, null)
-                        },
+                        context,
                         forceReload,
-                    )
+                    ) { homeWallpaper, lockWallpaper, _ ->
+                        continuation.resume(lockWallpaper ?: homeWallpaper, null)
+                    }
                 }
             },
             wallpaperInteractor = wallpaperInteractor,