diff --git a/src/com/android/customization/picker/color/data/repository/ColorPickerRepositoryImpl.kt b/src/com/android/customization/picker/color/data/repository/ColorPickerRepositoryImpl.kt
index f5b4ac5..f393880 100644
--- a/src/com/android/customization/picker/color/data/repository/ColorPickerRepositoryImpl.kt
+++ b/src/com/android/customization/picker/color/data/repository/ColorPickerRepositoryImpl.kt
@@ -24,6 +24,7 @@
 import com.android.customization.picker.color.shared.model.ColorOptionModel
 import com.android.customization.picker.color.shared.model.ColorType
 import com.android.systemui.monet.Style
+import com.android.wallpaper.config.BaseFlags
 import com.android.wallpaper.picker.customization.data.repository.WallpaperColorsRepository
 import com.android.wallpaper.picker.customization.shared.model.WallpaperColorsModel
 import javax.inject.Inject
@@ -46,6 +47,8 @@
     private val colorManager: ColorCustomizationManager,
 ) : ColorPickerRepository {
 
+    private val isNewPickerUi = BaseFlags.get().isNewPickerUi()
+
     private val homeWallpaperColors: StateFlow<WallpaperColorsModel?> =
         wallpaperColorsRepository.homeWallpaperColors
     private val lockWallpaperColors: StateFlow<WallpaperColorsModel?> =
@@ -56,8 +59,7 @@
     private val _isApplyingSystemColor = MutableStateFlow(false)
     override val isApplyingSystemColor = _isApplyingSystemColor.asStateFlow()
 
-    // TODO (b/299510645): update color options on selected option change after restart is disabled
-    override val colorOptions: Flow<Map<ColorType, List<ColorOptionModel>>> =
+    private val generatedColorOptions: Flow<Map<ColorType, List<ColorOptionImpl>>> =
         combine(homeWallpaperColors, lockWallpaperColors) { homeColors, lockColors ->
                 homeColors to lockColors
             }
@@ -71,7 +73,7 @@
                             Result.success(
                                 mapOf(
                                     ColorType.WALLPAPER_COLOR to listOf(),
-                                    ColorType.PRESET_COLOR to listOf()
+                                    ColorType.PRESET_COLOR to listOf(),
                                 )
                             )
                         )
@@ -81,28 +83,27 @@
                     val lockColorsLoaded = lockColors as WallpaperColorsModel.Loaded
                     colorManager.setWallpaperColors(
                         homeColorsLoaded.colors,
-                        lockColorsLoaded.colors
+                        lockColorsLoaded.colors,
                     )
                     colorManager.fetchOptions(
                         object : CustomizationManager.OptionsFetchedListener<ColorOption?> {
                             override fun onOptionsLoaded(options: MutableList<ColorOption?>?) {
-                                val wallpaperColorOptions: MutableList<ColorOptionModel> =
+                                val wallpaperColorOptions: MutableList<ColorOptionImpl> =
                                     mutableListOf()
-                                val presetColorOptions: MutableList<ColorOptionModel> =
+                                val presetColorOptions: MutableList<ColorOptionImpl> =
                                     mutableListOf()
                                 options?.forEach { option ->
                                     when ((option as ColorOptionImpl).type) {
                                         ColorType.WALLPAPER_COLOR ->
-                                            wallpaperColorOptions.add(option.toModel())
-                                        ColorType.PRESET_COLOR ->
-                                            presetColorOptions.add(option.toModel())
+                                            wallpaperColorOptions.add(option)
+                                        ColorType.PRESET_COLOR -> presetColorOptions.add(option)
                                     }
                                 }
                                 continuation.resumeWith(
                                     Result.success(
                                         mapOf(
                                             ColorType.WALLPAPER_COLOR to wallpaperColorOptions,
-                                            ColorType.PRESET_COLOR to presetColorOptions
+                                            ColorType.PRESET_COLOR to presetColorOptions,
                                         )
                                     )
                                 )
@@ -117,11 +118,88 @@
                                 )
                             }
                         },
-                        /* reload= */ false
+                        /* reload= */ false,
                     )
                 }
             }
 
+    override val colorOptions: Flow<Map<ColorType, List<ColorOptionModel>>> =
+        if (isNewPickerUi) {
+            // Convert to ColorOptionModel. When the selected color option changes, update each
+            // ColorOptionModel's isSelected by calling toModel again.
+            combine(generatedColorOptions, selectedColorOption) { generatedColorOptions, _ ->
+                generatedColorOptions
+                    .map { entry ->
+                        entry.key to entry.value.map { colorOption -> colorOption.toModel() }
+                    }
+                    .toMap()
+            }
+        } else {
+            combine(homeWallpaperColors, lockWallpaperColors) { homeColors, lockColors ->
+                    homeColors to lockColors
+                }
+                .map { (homeColors, lockColors) ->
+                    suspendCancellableCoroutine { continuation ->
+                        if (
+                            homeColors is WallpaperColorsModel.Loading ||
+                                lockColors is WallpaperColorsModel.Loading
+                        ) {
+                            continuation.resumeWith(
+                                Result.success(
+                                    mapOf(
+                                        ColorType.WALLPAPER_COLOR to listOf(),
+                                        ColorType.PRESET_COLOR to listOf(),
+                                    )
+                                )
+                            )
+                            return@suspendCancellableCoroutine
+                        }
+                        val homeColorsLoaded = homeColors as WallpaperColorsModel.Loaded
+                        val lockColorsLoaded = lockColors as WallpaperColorsModel.Loaded
+                        colorManager.setWallpaperColors(
+                            homeColorsLoaded.colors,
+                            lockColorsLoaded.colors,
+                        )
+                        colorManager.fetchOptions(
+                            object : CustomizationManager.OptionsFetchedListener<ColorOption?> {
+                                override fun onOptionsLoaded(options: MutableList<ColorOption?>?) {
+                                    val wallpaperColorOptions: MutableList<ColorOptionModel> =
+                                        mutableListOf()
+                                    val presetColorOptions: MutableList<ColorOptionModel> =
+                                        mutableListOf()
+                                    options?.forEach { option ->
+                                        when ((option as ColorOptionImpl).type) {
+                                            ColorType.WALLPAPER_COLOR ->
+                                                wallpaperColorOptions.add(option.toModel())
+                                            ColorType.PRESET_COLOR ->
+                                                presetColorOptions.add(option.toModel())
+                                        }
+                                    }
+                                    continuation.resumeWith(
+                                        Result.success(
+                                            mapOf(
+                                                ColorType.WALLPAPER_COLOR to wallpaperColorOptions,
+                                                ColorType.PRESET_COLOR to presetColorOptions,
+                                            )
+                                        )
+                                    )
+                                }
+
+                                override fun onError(throwable: Throwable?) {
+                                    Log.e(TAG, "Error loading theme bundles", throwable)
+                                    continuation.resumeWith(
+                                        Result.failure(
+                                            throwable ?: Throwable("Error loading theme bundles")
+                                        )
+                                    )
+                                }
+                            },
+                            /* reload= */ false,
+                        )
+                    }
+                }
+        }
+
     override suspend fun select(colorOptionModel: ColorOptionModel) {
         _isApplyingSystemColor.value = true
         suspendCancellableCoroutine { continuation ->
@@ -141,7 +219,7 @@
                             Result.failure(throwable ?: Throwable("Error loading theme bundles"))
                         )
                     }
-                }
+                },
             )
         }
     }
@@ -158,11 +236,7 @@
             colorOptionBuilder.addOverlayPackage(overlay.key, overlay.value)
         }
         val colorOption = colorOptionBuilder.build()
-        return ColorOptionModel(
-            key = "",
-            colorOption = colorOption,
-            isSelected = false,
-        )
+        return ColorOptionModel(key = "", colorOption = colorOption, isSelected = false)
     }
 
     override fun getCurrentColorSource(): String? {
@@ -173,6 +247,8 @@
         return ColorOptionModel(
             key = "${this.type}::${this.style}::${this.serializedPackages}",
             colorOption = this,
+            // Instead of using the selectedColorOption flow to determine isSelected, we check the
+            // source of truth, which is the settings, using ColorOption::isActive
             isSelected = isActive(colorManager),
         )
     }
diff --git a/src/com/android/customization/picker/color/shared/model/ColorOptionModel.kt b/src/com/android/customization/picker/color/shared/model/ColorOptionModel.kt
index 5fde08e..ba477bc 100644
--- a/src/com/android/customization/picker/color/shared/model/ColorOptionModel.kt
+++ b/src/com/android/customization/picker/color/shared/model/ColorOptionModel.kt
@@ -27,5 +27,5 @@
     val colorOption: ColorOption,
 
     /** Whether this color option is selected. */
-    var isSelected: Boolean,
+    val isSelected: Boolean,
 )
diff --git a/src/com/android/wallpaper/customization/ui/viewmodel/ColorPickerViewModel2.kt b/src/com/android/wallpaper/customization/ui/viewmodel/ColorPickerViewModel2.kt
index a039996..fc78ac7 100644
--- a/src/com/android/wallpaper/customization/ui/viewmodel/ColorPickerViewModel2.kt
+++ b/src/com/android/wallpaper/customization/ui/viewmodel/ColorPickerViewModel2.kt
@@ -20,6 +20,7 @@
 import com.android.customization.model.color.ColorOptionImpl
 import com.android.customization.module.logging.ThemesUserEventLogger
 import com.android.customization.picker.color.domain.interactor.ColorPickerInteractor
+import com.android.customization.picker.color.shared.model.ColorOptionModel
 import com.android.customization.picker.color.shared.model.ColorType
 import com.android.customization.picker.color.ui.viewmodel.ColorOptionIconViewModel
 import com.android.themepicker.R
@@ -36,6 +37,7 @@
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.asStateFlow
 import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.stateIn
@@ -51,14 +53,16 @@
     @Assisted private val viewModelScope: CoroutineScope,
 ) {
 
+    private val overridingColorOption = MutableStateFlow<ColorOptionModel?>(null)
+    val previewingColorOption = overridingColorOption.asStateFlow()
+
     private val selectedColorTypeTabId = MutableStateFlow<ColorType?>(null)
 
     /** View-models for each color tab. */
     val colorTypeTabs: Flow<List<FloatingToolbarTabViewModel>> =
-        combine(
-            interactor.colorOptions,
-            selectedColorTypeTabId,
-        ) { colorOptions, selectedColorTypeIdOrNull ->
+        combine(interactor.colorOptions, selectedColorTypeTabId) {
+            colorOptions,
+            selectedColorTypeIdOrNull ->
             colorOptions.keys.mapIndexed { index, colorType ->
                 val isSelected =
                     (selectedColorTypeIdOrNull == null && index == 0) ||
@@ -118,7 +122,7 @@
                             val darkThemeColors =
                                 colorOption.previewInfo.resolveColors(/* darkTheme= */ true)
                             val isSelectedFlow: StateFlow<Boolean> =
-                                interactor.selectingColorOption
+                                previewingColorOption
                                     .map {
                                         it?.colorOption?.isEquivalent(colorOptionModel.colorOption)
                                             ?: colorOptionModel.isSelected
@@ -150,15 +154,7 @@
                                         } else {
                                             {
                                                 viewModelScope.launch {
-                                                    interactor.select(colorOptionModel)
-                                                    logger.logThemeColorApplied(
-                                                        colorOptionModel.colorOption
-                                                            .sourceForLogging,
-                                                        colorOptionModel.colorOption
-                                                            .styleForLogging,
-                                                        colorOptionModel.colorOption
-                                                            .seedColorForLogging,
-                                                    )
+                                                    overridingColorOption.value = colorOptionModel
                                                 }
                                             }
                                         }
@@ -169,6 +165,28 @@
                 .toMap()
         }
 
+    val onApply: Flow<(suspend () -> Unit)?> =
+        previewingColorOption.map { previewingColorOption ->
+            previewingColorOption?.let {
+                if (it.isSelected) {
+                    null
+                } else {
+                    {
+                        interactor.select(it)
+                        logger.logThemeColorApplied(
+                            previewingColorOption.colorOption.sourceForLogging,
+                            previewingColorOption.colorOption.styleForLogging,
+                            previewingColorOption.colorOption.seedColorForLogging,
+                        )
+                    }
+                }
+            }
+        }
+
+    fun resetPreview() {
+        overridingColorOption.value = null
+    }
+
     /** The list of all available color options for the selected Color Type. */
     val colorOptions: Flow<List<OptionItemViewModel<ColorOptionIconViewModel>>> =
         combine(allColorOptions, selectedColorTypeTabId) {
diff --git a/src/com/android/wallpaper/customization/ui/viewmodel/ThemePickerCustomizationOptionsViewModel.kt b/src/com/android/wallpaper/customization/ui/viewmodel/ThemePickerCustomizationOptionsViewModel.kt
index 03831bd..d114ac8 100644
--- a/src/com/android/wallpaper/customization/ui/viewmodel/ThemePickerCustomizationOptionsViewModel.kt
+++ b/src/com/android/wallpaper/customization/ui/viewmodel/ThemePickerCustomizationOptionsViewModel.kt
@@ -61,6 +61,7 @@
         keyguardQuickAffordancePickerViewModel2.resetPreview()
         shapeAndGridPickerViewModel.resetPreview()
         clockPickerViewModel.resetPreview()
+        colorPickerViewModel2.resetPreview()
         return defaultCustomizationOptionsViewModel.deselectOption()
     }
 
@@ -129,6 +130,8 @@
                         .SHORTCUTS -> keyguardQuickAffordancePickerViewModel2.onApply
                     ThemePickerCustomizationOptionUtil.ThemePickerHomeCustomizationOption
                         .APP_SHAPE_AND_GRID -> shapeAndGridPickerViewModel.onApply
+                    ThemePickerCustomizationOptionUtil.ThemePickerHomeCustomizationOption.COLORS ->
+                        colorPickerViewModel2.onApply
                     else -> flow { emit(null) }
                 }
             }
diff --git a/tests/robotests/src/com/android/wallpaper/customization/ui/viewmodel/ColorPickerViewModel2Test.kt b/tests/robotests/src/com/android/wallpaper/customization/ui/viewmodel/ColorPickerViewModel2Test.kt
index d13d4b1..c153de6 100644
--- a/tests/robotests/src/com/android/wallpaper/customization/ui/viewmodel/ColorPickerViewModel2Test.kt
+++ b/tests/robotests/src/com/android/wallpaper/customization/ui/viewmodel/ColorPickerViewModel2Test.kt
@@ -97,19 +97,19 @@
     }
 
     @Test
-    fun `Log selected wallpaper color`() =
+    fun onApply_wallpaperColor_shouldLogColor() =
         testScope.runTest {
             repository.setOptions(
                 listOf(
                     repository.buildWallpaperOption(
                         ColorOptionsProvider.COLOR_SOURCE_LOCK,
                         Style.EXPRESSIVE,
-                        "121212"
+                        "121212",
                     )
                 ),
                 listOf(repository.buildPresetOption(Style.FRUIT_SALAD, "#ABCDEF")),
                 ColorType.PRESET_COLOR,
-                0
+                0,
             )
 
             val colorTypes = collectLastValue(underTest.colorTypeTabs)
@@ -117,8 +117,10 @@
 
             // Select "Wallpaper colors" tab
             colorTypes()?.get(0)?.onClick?.invoke()
-            // Select a color option
+            // Select a color option to preview
             selectColorOption(colorOptions, 0)
+            // Apply the selected color option
+            applySelectedColorOption()
 
             assertThat(logger.themeColorSource)
                 .isEqualTo(StyleEnums.COLOR_SOURCE_LOCK_SCREEN_WALLPAPER)
@@ -127,19 +129,19 @@
         }
 
     @Test
-    fun `Log selected preset color`() =
+    fun onApply_presetColor_shouldLogColor() =
         testScope.runTest {
             repository.setOptions(
                 listOf(
                     repository.buildWallpaperOption(
                         ColorOptionsProvider.COLOR_SOURCE_LOCK,
                         Style.EXPRESSIVE,
-                        "121212"
+                        "121212",
                     )
                 ),
                 listOf(repository.buildPresetOption(Style.FRUIT_SALAD, "#ABCDEF")),
                 ColorType.WALLPAPER_COLOR,
-                0
+                0,
             )
 
             val colorTypes = collectLastValue(underTest.colorTypeTabs)
@@ -147,8 +149,10 @@
 
             // Select "Wallpaper colors" tab
             colorTypes()?.get(1)?.onClick?.invoke()
-            // Select a color option
+            // Select a color option to preview
             selectColorOption(colorOptions, 0)
+            // Apply the selected color option
+            applySelectedColorOption()
 
             assertThat(logger.themeColorSource).isEqualTo(StyleEnums.COLOR_SOURCE_PRESET_COLOR)
             assertThat(logger.themeColorStyle).isEqualTo(Style.FRUIT_SALAD.toString().hashCode())
@@ -156,7 +160,7 @@
         }
 
     @Test
-    fun `Select a preset color`() =
+    fun selectColorOption() =
         testScope.runTest {
             val colorTypes = collectLastValue(underTest.colorTypeTabs)
             val colorOptions = collectLastValue(underTest.colorOptions)
@@ -166,7 +170,7 @@
                 colorTypes = colorTypes(),
                 colorOptions = colorOptions(),
                 selectedColorTypeText = "Wallpaper colors",
-                selectedColorOptionIndex = 0
+                selectedColorOptionIndex = 0,
             )
 
             // Select "Basic colors" tab
@@ -175,7 +179,7 @@
                 colorTypes = colorTypes(),
                 colorOptions = colorOptions(),
                 selectedColorTypeText = "Basic colors",
-                selectedColorOptionIndex = -1
+                selectedColorOptionIndex = -1,
             )
 
             // Select a color option
@@ -187,7 +191,7 @@
                 colorTypes = colorTypes(),
                 colorOptions = colorOptions(),
                 selectedColorTypeText = "Wallpaper colors",
-                selectedColorOptionIndex = -1
+                selectedColorOptionIndex = -1,
             )
 
             // Check new option is selected
@@ -196,7 +200,7 @@
                 colorTypes = colorTypes(),
                 colorOptions = colorOptions(),
                 selectedColorTypeText = "Basic colors",
-                selectedColorOptionIndex = 2
+                selectedColorOptionIndex = 2,
             )
         }
 
@@ -210,10 +214,16 @@
             onClickedFlow?.let { collectLastValue(it) }
         onClickedLastValueOrNull?.let { onClickedLastValue ->
             val onClickedOrNull: (() -> Unit)? = onClickedLastValue()
-            onClickedOrNull?.let { onClicked -> onClicked() }
+            onClickedOrNull?.invoke()
         }
     }
 
+    /** Simulates a user selecting the affordance at the given index, if that is clickable. */
+    private suspend fun TestScope.applySelectedColorOption() {
+        val onApply = collectLastValue(underTest.onApply)()
+        onApply?.invoke()
+    }
+
     /**
      * Asserts the entire picker UI state is what is expected. This includes the color type tabs and
      * the color options list.