diff --git a/res/layout/color_option_with_background.xml b/res/layout/color_option_with_background.xml
index dc6f3ed..36990b5 100644
--- a/res/layout/color_option_with_background.xml
+++ b/res/layout/color_option_with_background.xml
@@ -22,8 +22,8 @@
     android:orientation="vertical" >
     <FrameLayout
         android:id="@+id/icon_container"
-        android:layout_width="@dimen/option_tile_width"
-        android:layout_height="@dimen/option_tile_width"
+        android:layout_width="@dimen/option_item_size"
+        android:layout_height="@dimen/option_item_size"
         android:importantForAccessibility="yes">
 
         <ImageView
@@ -43,45 +43,45 @@
 
         <ImageView
             android:id="@+id/color_preview_0"
-            android:layout_width="wrap_content"
-            android:layout_height="wrap_content"
+            android:layout_width="@dimen/component_color_chip_small_radius_default2"
+            android:layout_height="@dimen/component_color_chip_small_radius_default2"
             android:layout_gravity="center"
-            android:layout_marginRight="@dimen/color_seed_chip_margin"
-            android:layout_marginBottom="@dimen/color_seed_chip_margin"
+            android:layout_marginRight="@dimen/color_seed_chip_margin2"
+            android:layout_marginBottom="@dimen/color_seed_chip_margin2"
             android:src="@drawable/color_chip_seed_filled0"
             android:importantForAccessibility="no"/>
 
         <ImageView
             android:id="@+id/color_preview_1"
-            android:layout_width="wrap_content"
-            android:layout_height="wrap_content"
+            android:layout_width="@dimen/component_color_chip_small_radius_default2"
+            android:layout_height="@dimen/component_color_chip_small_radius_default2"
             android:layout_gravity="center"
-            android:layout_marginLeft="@dimen/color_seed_chip_margin"
-            android:layout_marginBottom="@dimen/color_seed_chip_margin"
+            android:layout_marginLeft="@dimen/color_seed_chip_margin2"
+            android:layout_marginBottom="@dimen/color_seed_chip_margin2"
             android:src="@drawable/color_chip_seed_filled2"
             android:importantForAccessibility="no"/>
 
         <ImageView
             android:id="@+id/color_preview_2"
-            android:layout_width="wrap_content"
-            android:layout_height="wrap_content"
+            android:layout_width="@dimen/component_color_chip_small_radius_default2"
+            android:layout_height="@dimen/component_color_chip_small_radius_default2"
             android:layout_gravity="center"
-            android:layout_marginRight="@dimen/color_seed_chip_margin"
-            android:layout_marginTop="@dimen/color_seed_chip_margin"
+            android:layout_marginRight="@dimen/color_seed_chip_margin2"
+            android:layout_marginTop="@dimen/color_seed_chip_margin2"
             android:src="@drawable/color_chip_seed_filled1"
             android:importantForAccessibility="no"/>
 
         <ImageView
             android:id="@+id/color_preview_3"
-            android:layout_width="wrap_content"
-            android:layout_height="wrap_content"
+            android:layout_width="@dimen/component_color_chip_small_radius_default2"
+            android:layout_height="@dimen/component_color_chip_small_radius_default2"
             android:layout_gravity="center"
-            android:layout_marginLeft="@dimen/color_seed_chip_margin"
-            android:layout_marginTop="@dimen/color_seed_chip_margin"
+            android:layout_marginLeft="@dimen/color_seed_chip_margin2"
+            android:layout_marginTop="@dimen/color_seed_chip_margin2"
             android:src="@drawable/color_chip_seed_filled3"
-            android:importantForAccessibility="no" />
+            android:importantForAccessibility="no"/>
     </FrameLayout>
-
+    
     <TextView
         android:id="@+id/option_title"
         android:layout_width="@dimen/option_tile_width"
@@ -91,4 +91,3 @@
         android:visibility="gone"
         android:gravity="center" />
 </LinearLayout>
-
diff --git a/res/layout/fragment_color_picker.xml b/res/layout/fragment_color_picker.xml
index efb5a87..7cf94d6 100644
--- a/res/layout/fragment_color_picker.xml
+++ b/res/layout/fragment_color_picker.xml
@@ -35,7 +35,7 @@
         android:layout_height="0dp"
         android:layout_weight="1"
         android:layout_marginHorizontal="24dp"
-        android:paddingTop="36dp"
+        android:paddingTop="20dp"
         android:paddingBottom="40dp"
         android:orientation="horizontal"
         android:gravity="center_vertical">
@@ -70,10 +70,11 @@
 
             <androidx.recyclerview.widget.RecyclerView
                 android:id="@+id/color_type_tabs"
-                android:layout_width="match_parent"
+                android:layout_width="wrap_content"
                 android:layout_height="wrap_content"
                 android:clipToPadding="false"
-                android:paddingHorizontal="16dp" />
+                android:paddingHorizontal="16dp"
+                android:layout_gravity="center_horizontal"/>
 
             <!--
             This is just an invisible placeholder put in place so that the parent keeps its height
@@ -94,7 +95,20 @@
 
         <View
             android:layout_width="0dp"
-            android:layout_height="22dp" />
+            android:layout_height="6dp" />
+
+        <TextView
+            android:id="@+id/color_type_tab_subhead"
+            android:layout_width="match_parent"
+            android:layout_height="wrap_content"
+            android:textColor="@color/text_color_secondary"
+            android:textSize="12sp"
+            android:gravity="center"
+            android:paddingHorizontal="16dp"/>
+
+        <View
+            android:layout_width="0dp"
+            android:layout_height="8dp" />
 
         <FrameLayout
             android:layout_width="match_parent"
diff --git a/res/values/dimens.xml b/res/values/dimens.xml
index e98a47e..0d5a9a8 100644
--- a/res/values/dimens.xml
+++ b/res/values/dimens.xml
@@ -130,20 +130,25 @@
     <dimen name="beta_tag_background_width">46dp</dimen>
     <dimen name="beta_tag_background_height">24dp</dimen>
 
-    <!-- For the color option section -->
+    <!-- For the color picker section -->
+    <!-- Dimension for the original color picker section -->
     <dimen name="color_options_container_top_margin">24dp</dimen>
-
+    <dimen name="color_page_indicator_margin_top">16dp</dimen>
+    <!-- Dimension for the original color chip resource -->
+    <dimen name="component_color_chip_small_radius_default">29dp</dimen>
+    <dimen name="component_color_chip_small_diameter_default">58dp</dimen>
+    <dimen name="color_seed_chip_margin">14dp</dimen>
+    <!-- Dimension for the revamped UI color chip -->
     <dimen name="component_color_overflow_small_radius_default">20dp</dimen>
     <dimen name="component_color_overflow_small_diameter_default">40dp</dimen>
     <dimen name="component_color_selected_small_radius_default">22dp</dimen>
     <dimen name="component_color_selected_small_diameter_default">44dp</dimen>
 
-    <!-- For the color page. -->
-    <dimen name="color_page_indicator_margin_top">16dp</dimen>
-
-    <dimen name="component_color_chip_small_radius_default">29dp</dimen>
-    <dimen name="component_color_chip_small_diameter_default">58dp</dimen>
-    <dimen name="color_seed_chip_margin">14dp</dimen>
+    <!-- For the color picker page. -->
+    <!-- Dimension for the revamped UI color chip -->
+    <dimen name="component_color_chip_small_radius_default2">24dp</dimen>
+    <dimen name="component_color_chip_small_diameter_default2">48dp</dimen>
+    <dimen name="color_seed_chip_margin2">11dp</dimen>
 
     <!-- Keyguard quick affordances -->
     <!-- Size for the container for the icon of a quick affordance for the lock screen in the picker experience. -->
diff --git a/res/values/strings.xml b/res/values/strings.xml
index 18c0103..01113c4 100755
--- a/res/values/strings.xml
+++ b/res/values/strings.xml
@@ -311,11 +311,17 @@
 
     <!-- The title for the tab with colors from wallpaper. [CHAR_LIMIT=20]-->
     <string name="wallpaper_color_tab">Wallpaper colors</string>
+    <!-- The subheader for the tab with colors from wallpaper. [CHAR_LIMIT=55]-->
+    <string name="wallpaper_color_subheader">Icons, text, and more match colors in your wallpaper</string>
     <!-- The description on an item that shows a color obtained from the wallpaper
         (used mainly for accessibility). [CHAR_LIMIT=NONE] -->
     <string name="wallpaper_color_title">Wallpaper color</string>
     <!-- The title for the tab with a default set of color options. [CHAR_LIMIT=20]-->
     <string name="preset_color_tab">Basic colors</string>
+    <!-- The revised title for the tab with a default set of color options. [CHAR_LIMIT=20]-->
+    <string name="preset_color_tab_2">Other colors</string>
+    <!-- The subheader for the tab with a default set of color options. [CHAR_LIMIT=55]-->
+    <string name="preset_color_subheader">Choose any color for your icons, clock, and more</string>
     <!-- The text for A11y announcement when color changes. -->
     <string name="color_changed">Color changed</string>
     <!-- Title of a section of color selection option that obtains colors automatically from the
@@ -326,7 +332,7 @@
     color options.
     [CHAR LIMIT=32].
     -->
-    <string name="color_picker_title">System Colors</string>
+    <string name="color_picker_title">System colors</string>
 
     <!--
     Name of the slot on the "start" side of the bottom of the lock screen, where lock screen
diff --git a/src/com/android/customization/model/grid/ui/binder/GridIconViewBinder.kt b/src/com/android/customization/model/grid/ui/binder/GridIconViewBinder.kt
new file mode 100644
index 0000000..fba89a7
--- /dev/null
+++ b/src/com/android/customization/model/grid/ui/binder/GridIconViewBinder.kt
@@ -0,0 +1,17 @@
+package com.android.customization.model.grid.ui.binder
+
+import android.widget.ImageView
+import com.android.customization.model.grid.ui.viewmodel.GridIconViewModel
+import com.android.customization.widget.GridTileDrawable
+
+object GridIconViewBinder {
+    fun bind(view: ImageView, viewModel: GridIconViewModel) {
+        view.setImageDrawable(
+            GridTileDrawable(
+                viewModel.columns,
+                viewModel.rows,
+                viewModel.path,
+            )
+        )
+    }
+}
diff --git a/src/com/android/customization/model/grid/ui/binder/GridScreenBinder.kt b/src/com/android/customization/model/grid/ui/binder/GridScreenBinder.kt
index 78e0376..78536ca 100644
--- a/src/com/android/customization/model/grid/ui/binder/GridScreenBinder.kt
+++ b/src/com/android/customization/model/grid/ui/binder/GridScreenBinder.kt
@@ -25,11 +25,10 @@
 import androidx.lifecycle.repeatOnLifecycle
 import androidx.recyclerview.widget.LinearLayoutManager
 import androidx.recyclerview.widget.RecyclerView
+import com.android.customization.model.grid.ui.viewmodel.GridIconViewModel
 import com.android.customization.model.grid.ui.viewmodel.GridScreenViewModel
 import com.android.customization.picker.common.ui.view.ItemSpacing
 import com.android.wallpaper.R
-import com.android.wallpaper.picker.common.icon.ui.viewbinder.IconViewBinder
-import com.android.wallpaper.picker.common.icon.ui.viewmodel.Icon
 import com.android.wallpaper.picker.option.ui.adapter.OptionItemAdapter
 import com.android.wallpaper.picker.option.ui.binder.OptionItemBinder
 import kotlinx.coroutines.CoroutineDispatcher
@@ -61,9 +60,9 @@
                         selectedColor = view.context.getColor(R.color.text_color_primary),
                         unselectedColor = view.context.getColor(R.color.text_color_secondary),
                     ),
-                bindIcon = { foregroundView: View, gridIcon: Icon ->
+                bindIcon = { foregroundView: View, gridIcon: GridIconViewModel ->
                     val imageView = foregroundView as? ImageView
-                    imageView?.let { IconViewBinder.bind(imageView, gridIcon) }
+                    imageView?.let { GridIconViewBinder.bind(imageView, gridIcon) }
                 }
             )
         optionView.adapter = adapter
diff --git a/src/com/android/customization/model/grid/ui/fragment/GridFragment2.kt b/src/com/android/customization/model/grid/ui/fragment/GridFragment2.kt
index 4440b77..d8cad82 100644
--- a/src/com/android/customization/model/grid/ui/fragment/GridFragment2.kt
+++ b/src/com/android/customization/model/grid/ui/fragment/GridFragment2.kt
@@ -27,8 +27,10 @@
 import com.android.customization.module.ThemePickerInjector
 import com.android.wallpaper.R
 import com.android.wallpaper.module.CurrentWallpaperInfoFactory
+import com.android.wallpaper.module.CustomizationSections
 import com.android.wallpaper.module.InjectorProvider
 import com.android.wallpaper.picker.AppbarFragment
+import com.android.wallpaper.picker.customization.domain.interactor.WallpaperInteractor
 import com.android.wallpaper.picker.customization.ui.binder.ScreenPreviewBinder
 import com.android.wallpaper.picker.customization.ui.viewmodel.ScreenPreviewViewModel
 import com.android.wallpaper.util.PreviewUtils
@@ -55,7 +57,12 @@
         val injector = InjectorProvider.getInjector() as ThemePickerInjector
 
         val wallpaperInfoFactory = injector.getCurrentWallpaperInfoFactory(requireContext())
-        var screenPreviewBinding = bindScreenPreview(view, wallpaperInfoFactory)
+        var screenPreviewBinding =
+            bindScreenPreview(
+                view,
+                wallpaperInfoFactory,
+                injector.getWallpaperInteractor(requireContext())
+            )
 
         val viewModelFactory = injector.getGridScreenViewModelFactory(requireContext())
         GridScreenBinder.bind(
@@ -69,7 +76,12 @@
             backgroundDispatcher = Dispatchers.IO,
             onOptionsChanged = {
                 screenPreviewBinding.destroy()
-                screenPreviewBinding = bindScreenPreview(view, wallpaperInfoFactory)
+                screenPreviewBinding =
+                    bindScreenPreview(
+                        view,
+                        wallpaperInfoFactory,
+                        injector.getWallpaperInteractor(requireContext())
+                    )
             }
         )
 
@@ -83,6 +95,7 @@
     private fun bindScreenPreview(
         view: View,
         wallpaperInfoFactory: CurrentWallpaperInfoFactory,
+        wallpaperInteractor: WallpaperInteractor,
     ): ScreenPreviewBinder.Binding {
         return ScreenPreviewBinder.bind(
             activity = requireActivity(),
@@ -108,9 +121,12 @@
                             )
                         }
                     },
+                    wallpaperInteractor = wallpaperInteractor,
                 ),
             lifecycleOwner = this,
             offsetToStart = false,
+            screen = CustomizationSections.Screen.HOME_SCREEN,
+            onPreviewDirty = { activity?.recreate() },
         )
     }
 }
diff --git a/src/com/android/customization/model/grid/ui/viewmodel/GridIconViewModel.kt b/src/com/android/customization/model/grid/ui/viewmodel/GridIconViewModel.kt
new file mode 100644
index 0000000..3942d7c
--- /dev/null
+++ b/src/com/android/customization/model/grid/ui/viewmodel/GridIconViewModel.kt
@@ -0,0 +1,24 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package com.android.customization.model.grid.ui.viewmodel
+
+data class GridIconViewModel(
+    val columns: Int,
+    val rows: Int,
+    val path: String,
+)
diff --git a/src/com/android/customization/model/grid/ui/viewmodel/GridScreenViewModel.kt b/src/com/android/customization/model/grid/ui/viewmodel/GridScreenViewModel.kt
index 69d938c..c11a594 100644
--- a/src/com/android/customization/model/grid/ui/viewmodel/GridScreenViewModel.kt
+++ b/src/com/android/customization/model/grid/ui/viewmodel/GridScreenViewModel.kt
@@ -26,12 +26,11 @@
 import com.android.customization.model.ResourceConstants
 import com.android.customization.model.grid.domain.interactor.GridInteractor
 import com.android.customization.model.grid.shared.model.GridOptionItemsModel
-import com.android.customization.widget.GridTileDrawable
-import com.android.wallpaper.picker.common.icon.ui.viewmodel.Icon
 import com.android.wallpaper.picker.common.text.ui.viewmodel.Text
 import com.android.wallpaper.picker.option.ui.viewmodel.OptionItemViewModel
 import kotlinx.coroutines.flow.Flow
-import kotlinx.coroutines.flow.flowOf
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.launch
 
@@ -43,12 +42,12 @@
     @SuppressLint("StaticFieldLeak") // We're not leaking this context as it is the app context.
     private val applicationContext = context.applicationContext
 
-    val optionItems: Flow<List<OptionItemViewModel<Icon>>> =
+    val optionItems: Flow<List<OptionItemViewModel<GridIconViewModel>>> =
         interactor.options.map { model -> toViewModel(model) }
 
     private fun toViewModel(
         model: GridOptionItemsModel,
-    ): List<OptionItemViewModel<Icon>> {
+    ): List<OptionItemViewModel<GridIconViewModel>> {
         val iconShapePath =
             applicationContext.resources.getString(
                 Resources.getSystem()
@@ -63,17 +62,14 @@
             is GridOptionItemsModel.Loaded ->
                 model.options.map { option ->
                     val text = Text.Loaded(option.name)
-                    OptionItemViewModel<Icon>(
-                        key = flowOf("${option.cols}x${option.rows}"),
+                    OptionItemViewModel<GridIconViewModel>(
+                        key =
+                            MutableStateFlow("${option.cols}x${option.rows}") as StateFlow<String>,
                         payload =
-                            Icon.Loaded(
-                                drawable =
-                                    GridTileDrawable(
-                                        option.cols,
-                                        option.rows,
-                                        iconShapePath,
-                                    ),
-                                contentDescription = text
+                            GridIconViewModel(
+                                columns = option.cols,
+                                rows = option.rows,
+                                path = iconShapePath,
                             ),
                         text = text,
                         isSelected = option.isSelected,
diff --git a/src/com/android/customization/module/DefaultCustomizationSections.java b/src/com/android/customization/module/DefaultCustomizationSections.java
index 482ece01..232e948 100644
--- a/src/com/android/customization/module/DefaultCustomizationSections.java
+++ b/src/com/android/customization/module/DefaultCustomizationSections.java
@@ -36,6 +36,7 @@
 import com.android.wallpaper.model.WallpaperSectionController;
 import com.android.wallpaper.module.CurrentWallpaperInfoFactory;
 import com.android.wallpaper.module.CustomizationSections;
+import com.android.wallpaper.picker.customization.domain.interactor.WallpaperInteractor;
 import com.android.wallpaper.picker.customization.ui.section.ConnectedSectionController;
 import com.android.wallpaper.picker.customization.ui.section.ScreenPreviewSectionController;
 import com.android.wallpaper.picker.customization.ui.section.WallpaperQuickSwitchSectionController;
@@ -97,7 +98,8 @@
             @Nullable Bundle savedInstanceState,
             CurrentWallpaperInfoFactory wallpaperInfoFactory,
             DisplayUtils displayUtils,
-            WallpaperQuickSwitchViewModel wallpaperQuickSwitchViewModel) {
+            WallpaperQuickSwitchViewModel wallpaperQuickSwitchViewModel,
+            WallpaperInteractor wallpaperInteractor) {
         List<CustomizationSectionController<?>> sectionControllers = new ArrayList<>();
 
         // Wallpaper section.
@@ -112,7 +114,8 @@
                         displayUtils,
                         mClockCarouselViewModel,
                         mClockViewFactory,
-                        sectionNavigationController)
+                        sectionNavigationController,
+                        wallpaperInteractor)
                         : new ScreenPreviewSectionController(
                                 activity,
                                 lifecycleOwner,
@@ -120,7 +123,8 @@
                                 wallpaperInfoFactory,
                                 wallpaperColorsViewModel,
                                 displayUtils,
-                                sectionNavigationController));
+                                sectionNavigationController,
+                                wallpaperInteractor));
 
         sectionControllers.add(
                 new ConnectedSectionController(
diff --git a/src/com/android/customization/module/ThemePickerInjector.kt b/src/com/android/customization/module/ThemePickerInjector.kt
index e06a1b5..eb20037 100644
--- a/src/com/android/customization/module/ThemePickerInjector.kt
+++ b/src/com/android/customization/module/ThemePickerInjector.kt
@@ -21,10 +21,13 @@
 import android.content.Intent
 import android.net.Uri
 import android.os.Bundle
+import android.text.TextUtils
 import androidx.activity.ComponentActivity
 import androidx.fragment.app.Fragment
 import androidx.fragment.app.FragmentActivity
 import androidx.lifecycle.ViewModelProvider
+import com.android.customization.model.color.ColorCustomizationManager
+import com.android.customization.model.color.ColorOptionsProvider
 import com.android.customization.model.grid.GridOptionsManager
 import com.android.customization.model.grid.data.repository.GridRepositoryImpl
 import com.android.customization.model.grid.domain.interactor.GridInteractor
@@ -72,6 +75,9 @@
 import com.android.wallpaper.picker.ImagePreviewFragment
 import com.android.wallpaper.picker.LivePreviewFragment
 import com.android.wallpaper.picker.PreviewFragment
+import com.android.wallpaper.picker.customization.data.content.WallpaperClientImpl
+import com.android.wallpaper.picker.customization.data.repository.WallpaperRepository
+import com.android.wallpaper.picker.customization.domain.interactor.WallpaperInteractor
 import com.android.wallpaper.picker.undo.domain.interactor.SnapshotRestorer
 import kotlinx.coroutines.DelicateCoroutinesApi
 import kotlinx.coroutines.Dispatchers
@@ -82,6 +88,7 @@
     private var customizationSections: CustomizationSections? = null
     private var userEventLogger: UserEventLogger? = null
     private var prefs: WallpaperPreferences? = null
+    private var wallpaperInteractor: WallpaperInteractor? = null
     private var keyguardQuickAffordancePickerInteractor: KeyguardQuickAffordancePickerInteractor? =
         null
     private var keyguardQuickAffordancePickerViewModelFactory:
@@ -102,6 +109,7 @@
     private var colorPickerInteractor: ColorPickerInteractor? = null
     private var colorPickerViewModelFactory: ColorPickerViewModel.Factory? = null
     private var colorPickerSnapshotRestorer: ColorPickerSnapshotRestorer? = null
+    private var colorCustomizationManager: ColorCustomizationManager? = null
     private var darkModeSnapshotRestorer: DarkModeSnapshotRestorer? = null
     private var themedIconSnapshotRestorer: ThemedIconSnapshotRestorer? = null
     private var themedIconInteractor: ThemedIconInteractor? = null
@@ -209,6 +217,25 @@
         return ThemeManager(provider, activity, overlayManagerCompat, logger)
     }
 
+    override fun getWallpaperInteractor(context: Context): WallpaperInteractor {
+        return wallpaperInteractor
+            ?: WallpaperInteractor(
+                    repository =
+                        WallpaperRepository(
+                            scope = GlobalScope,
+                            client = WallpaperClientImpl(context = context),
+                            backgroundDispatcher = Dispatchers.IO,
+                        ),
+                    shouldHandleReload = {
+                        TextUtils.equals(
+                            getColorCustomizationManager(context).currentColorSource,
+                            ColorOptionsProvider.COLOR_SOURCE_PRESET
+                        )
+                    }
+                )
+                .also { wallpaperInteractor = it }
+    }
+
     override fun getKeyguardQuickAffordancePickerInteractor(
         context: Context
     ): KeyguardQuickAffordancePickerInteractor {
@@ -352,7 +379,11 @@
     ): ColorPickerInteractor {
         return colorPickerInteractor
             ?: ColorPickerInteractor(
-                    repository = ColorPickerRepositoryImpl(context, wallpaperColorsViewModel),
+                    repository =
+                        ColorPickerRepositoryImpl(
+                            wallpaperColorsViewModel,
+                            getColorCustomizationManager(context)
+                        ),
                     snapshotRestorer = {
                         getColorPickerSnapshotRestorer(context, wallpaperColorsViewModel)
                     }
@@ -383,6 +414,13 @@
                 .also { colorPickerSnapshotRestorer = it }
     }
 
+    private fun getColorCustomizationManager(context: Context): ColorCustomizationManager {
+        return colorCustomizationManager
+            ?: ColorCustomizationManager.getInstance(context, OverlayManagerCompat(context)).also {
+                colorCustomizationManager = it
+            }
+    }
+
     fun getDarkModeSnapshotRestorer(
         context: Context,
     ): DarkModeSnapshotRestorer {
diff --git a/src/com/android/customization/picker/clock/ui/viewmodel/ClockSettingsViewModel.kt b/src/com/android/customization/picker/clock/ui/viewmodel/ClockSettingsViewModel.kt
index 23fbc7e..c3cd217 100644
--- a/src/com/android/customization/picker/clock/ui/viewmodel/ClockSettingsViewModel.kt
+++ b/src/com/android/customization/picker/clock/ui/viewmodel/ClockSettingsViewModel.kt
@@ -123,7 +123,7 @@
                                 context,
                                 selectedColorId,
                             )
-                            ?: (colorOptions[ColorType.BASIC_COLOR]
+                            ?: (colorOptions[ColorType.PRESET_COLOR]
                                     ?.find { it.isSelected }
                                     ?.colorOption as? ColorBundle)
                                 ?.toColorOptionViewModel(
diff --git a/src/com/android/customization/picker/color/data/repository/ColorPickerRepository.kt b/src/com/android/customization/picker/color/data/repository/ColorPickerRepository.kt
index 1a0f5a9..7cf9fd0 100644
--- a/src/com/android/customization/picker/color/data/repository/ColorPickerRepository.kt
+++ b/src/com/android/customization/picker/color/data/repository/ColorPickerRepository.kt
@@ -34,4 +34,7 @@
 
     /** Returns the current selected color option based on system settings */
     fun getCurrentColorOption(): ColorOptionModel
+
+    /** Returns the current selected color source based on system settings */
+    fun getCurrentColorSource(): String?
 }
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 dfc7a1b..512a500 100644
--- a/src/com/android/customization/picker/color/data/repository/ColorPickerRepositoryImpl.kt
+++ b/src/com/android/customization/picker/color/data/repository/ColorPickerRepositoryImpl.kt
@@ -17,14 +17,12 @@
 package com.android.customization.picker.color.data.repository
 
 import android.app.WallpaperColors
-import android.content.Context
 import android.util.Log
 import com.android.customization.model.CustomizationManager
 import com.android.customization.model.color.ColorBundle
 import com.android.customization.model.color.ColorCustomizationManager
 import com.android.customization.model.color.ColorOption
 import com.android.customization.model.color.ColorSeedOption
-import com.android.customization.model.theme.OverlayManagerCompat
 import com.android.customization.picker.color.shared.model.ColorOptionModel
 import com.android.customization.picker.color.shared.model.ColorType
 import com.android.systemui.monet.Style
@@ -38,16 +36,14 @@
 // TODO (b/262924623): refactor to remove dependency on ColorCustomizationManager & ColorOption
 // TODO (b/268203200): Create test for ColorPickerRepositoryImpl
 class ColorPickerRepositoryImpl(
-    context: Context,
     wallpaperColorsViewModel: WallpaperColorsViewModel,
+    private val colorManager: ColorCustomizationManager,
 ) : ColorPickerRepository {
 
     private val homeWallpaperColors: StateFlow<WallpaperColors?> =
         wallpaperColorsViewModel.homeWallpaperColors
     private val lockWallpaperColors: StateFlow<WallpaperColors?> =
         wallpaperColorsViewModel.lockWallpaperColors
-    private val colorManager: ColorCustomizationManager =
-        ColorCustomizationManager.getInstance(context, OverlayManagerCompat(context))
 
     override val colorOptions: Flow<Map<ColorType, List<ColorOptionModel>>> =
         combine(homeWallpaperColors, lockWallpaperColors) { homeColors, lockColors ->
@@ -74,7 +70,7 @@
                                     Result.success(
                                         mapOf(
                                             ColorType.WALLPAPER_COLOR to wallpaperColorOptions,
-                                            ColorType.BASIC_COLOR to presetColorOptions
+                                            ColorType.PRESET_COLOR to presetColorOptions
                                         )
                                     )
                                 )
@@ -132,6 +128,10 @@
         )
     }
 
+    override fun getCurrentColorSource(): String? {
+        return colorManager.currentColorSource
+    }
+
     private fun ColorOption.toModel(): ColorOptionModel {
         return ColorOptionModel(
             key = "${this.style}::${this.serializedPackages}",
diff --git a/src/com/android/customization/picker/color/data/repository/FakeColorPickerRepository.kt b/src/com/android/customization/picker/color/data/repository/FakeColorPickerRepository.kt
index f581c89..edbf6dc 100644
--- a/src/com/android/customization/picker/color/data/repository/FakeColorPickerRepository.kt
+++ b/src/com/android/customization/picker/color/data/repository/FakeColorPickerRepository.kt
@@ -20,6 +20,7 @@
 import android.graphics.Color
 import android.text.TextUtils
 import com.android.customization.model.color.ColorBundle
+import com.android.customization.model.color.ColorOptionsProvider
 import com.android.customization.model.color.ColorSeedOption
 import com.android.customization.picker.color.shared.model.ColorOptionModel
 import com.android.customization.picker.color.shared.model.ColorType
@@ -35,7 +36,7 @@
         MutableStateFlow(
             mapOf<ColorType, List<ColorOptionModel>>(
                 ColorType.WALLPAPER_COLOR to listOf(),
-                ColorType.BASIC_COLOR to listOf()
+                ColorType.PRESET_COLOR to listOf()
             )
         )
     override val colorOptions: StateFlow<Map<ColorType, List<ColorOptionModel>>> =
@@ -71,18 +72,18 @@
                             add(colorOption)
                         }
                     },
-                ColorType.BASIC_COLOR to
+                ColorType.PRESET_COLOR to
                     buildList {
                         repeat(times = numPresetOptions) { index ->
                             val isSelected =
-                                selectedColorOptionType == ColorType.BASIC_COLOR &&
+                                selectedColorOptionType == ColorType.PRESET_COLOR &&
                                     selectedColorOptionIndex == index
                             val colorOption =
                                 ColorOptionModel(
-                                    key = "${ColorType.BASIC_COLOR}::$index",
+                                    key = "${ColorType.PRESET_COLOR}::$index",
                                     colorOption = buildPresetOption(index),
                                     isSelected =
-                                        selectedColorOptionType == ColorType.BASIC_COLOR &&
+                                        selectedColorOptionType == ColorType.PRESET_COLOR &&
                                             selectedColorOptionIndex == index,
                                 )
                             if (isSelected) {
@@ -140,7 +141,7 @@
                 )
             }
         }
-        val basicColorOptions = colorOptions[ColorType.BASIC_COLOR]!!
+        val basicColorOptions = colorOptions[ColorType.PRESET_COLOR]!!
         val newBasicColorOptions = buildList {
             basicColorOptions.forEach { option ->
                 add(
@@ -155,12 +156,19 @@
         _colorOptions.value =
             mapOf(
                 ColorType.WALLPAPER_COLOR to newWallpaperColorOptions,
-                ColorType.BASIC_COLOR to newBasicColorOptions
+                ColorType.PRESET_COLOR to newBasicColorOptions
             )
     }
 
     override fun getCurrentColorOption(): ColorOptionModel = selectedColorOption
 
+    override fun getCurrentColorSource(): String? =
+        when (selectedColorOption.colorOption) {
+            is ColorSeedOption -> ColorOptionsProvider.COLOR_SOURCE_HOME
+            is ColorBundle -> ColorOptionsProvider.COLOR_SOURCE_PRESET
+            else -> null
+        }
+
     private fun ColorOptionModel.testEquals(other: Any?): Boolean {
         if (other == null) {
             return false
diff --git a/src/com/android/customization/picker/color/shared/model/ColorType.kt b/src/com/android/customization/picker/color/shared/model/ColorType.kt
index 631bb3d..c9a01d0 100644
--- a/src/com/android/customization/picker/color/shared/model/ColorType.kt
+++ b/src/com/android/customization/picker/color/shared/model/ColorType.kt
@@ -21,5 +21,5 @@
     WALLPAPER_COLOR,
 
     /** Preset colors */
-    BASIC_COLOR,
+    PRESET_COLOR,
 }
diff --git a/src/com/android/customization/picker/color/ui/adapter/ColorTypeTabAdapter.kt b/src/com/android/customization/picker/color/ui/adapter/ColorTypeTabAdapter.kt
index 8f1340a..bb9f082 100644
--- a/src/com/android/customization/picker/color/ui/adapter/ColorTypeTabAdapter.kt
+++ b/src/com/android/customization/picker/color/ui/adapter/ColorTypeTabAdapter.kt
@@ -22,15 +22,15 @@
 import android.view.ViewGroup
 import android.widget.TextView
 import androidx.recyclerview.widget.RecyclerView
-import com.android.customization.picker.color.ui.viewmodel.ColorTypeViewModel
+import com.android.customization.picker.color.ui.viewmodel.ColorTypeTabViewModel
 import com.android.wallpaper.R
 
 /** Adapts between color type items and views. */
 class ColorTypeTabAdapter : RecyclerView.Adapter<ColorTypeTabAdapter.ViewHolder>() {
 
-    private val items = mutableListOf<ColorTypeViewModel>()
+    private val items = mutableListOf<ColorTypeTabViewModel>()
 
-    fun setItems(items: List<ColorTypeViewModel>) {
+    fun setItems(items: List<ColorTypeTabViewModel>) {
         this.items.clear()
         this.items.addAll(items)
         notifyDataSetChanged()
diff --git a/src/com/android/customization/picker/color/ui/binder/ColorPickerBinder.kt b/src/com/android/customization/picker/color/ui/binder/ColorPickerBinder.kt
index 053d5dd..7623048 100644
--- a/src/com/android/customization/picker/color/ui/binder/ColorPickerBinder.kt
+++ b/src/com/android/customization/picker/color/ui/binder/ColorPickerBinder.kt
@@ -17,10 +17,9 @@
 
 package com.android.customization.picker.color.ui.binder
 
-import android.graphics.Rect
 import android.view.View
 import android.view.ViewGroup
-import androidx.core.view.ViewCompat
+import android.widget.TextView
 import androidx.lifecycle.Lifecycle
 import androidx.lifecycle.LifecycleOwner
 import androidx.lifecycle.lifecycleScope
@@ -30,8 +29,10 @@
 import com.android.customization.picker.color.ui.adapter.ColorTypeTabAdapter
 import com.android.customization.picker.color.ui.viewmodel.ColorOptionIconViewModel
 import com.android.customization.picker.color.ui.viewmodel.ColorPickerViewModel
+import com.android.customization.picker.common.ui.view.ItemSpacing
 import com.android.wallpaper.R
 import com.android.wallpaper.picker.option.ui.adapter.OptionItemAdapter
+import kotlinx.coroutines.flow.collect
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.launch
 
@@ -49,13 +50,14 @@
         lifecycleOwner: LifecycleOwner,
     ) {
         val colorTypeTabView: RecyclerView = view.requireViewById(R.id.color_type_tabs)
+        val colorTypeTabSubheaderView: TextView = view.requireViewById(R.id.color_type_tab_subhead)
         val colorOptionContainerView: RecyclerView = view.requireViewById(R.id.color_options)
 
         val colorTypeTabAdapter = ColorTypeTabAdapter()
         colorTypeTabView.adapter = colorTypeTabAdapter
         colorTypeTabView.layoutManager =
             LinearLayoutManager(view.context, RecyclerView.HORIZONTAL, false)
-        colorTypeTabView.addItemDecoration(ItemSpacing())
+        colorTypeTabView.addItemDecoration(ItemSpacing(ItemSpacing.TAB_ITEM_SPACING_DP))
         val colorOptionAdapter =
             OptionItemAdapter(
                 layoutResourceId = R.layout.color_option_2,
@@ -68,17 +70,23 @@
         colorOptionContainerView.adapter = colorOptionAdapter
         colorOptionContainerView.layoutManager =
             LinearLayoutManager(view.context, RecyclerView.HORIZONTAL, false)
-        colorOptionContainerView.addItemDecoration(ItemSpacing())
+        colorOptionContainerView.addItemDecoration(ItemSpacing(ItemSpacing.ITEM_SPACING_DP))
 
         lifecycleOwner.lifecycleScope.launch {
             lifecycleOwner.repeatOnLifecycle(Lifecycle.State.STARTED) {
                 launch {
-                    viewModel.colorTypes
+                    viewModel.colorTypeTabs
                         .map { colorTypeById -> colorTypeById.values }
                         .collect { colorTypes -> colorTypeTabAdapter.setItems(colorTypes.toList()) }
                 }
 
                 launch {
+                    viewModel.colorTypeTabSubheader.collect { subhead ->
+                        colorTypeTabSubheaderView.text = subhead
+                    }
+                }
+
+                launch {
                     viewModel.colorOptions.collect { colorOptions ->
                         colorOptionAdapter.setItems(colorOptions)
                     }
@@ -86,29 +94,4 @@
             }
         }
     }
-
-    // TODO (b/262924623): Remove function and use common ItemSpacing after ag/20929223 is merged
-    private class ItemSpacing : RecyclerView.ItemDecoration() {
-        override fun getItemOffsets(outRect: Rect, itemPosition: Int, parent: RecyclerView) {
-            val addSpacingToStart = itemPosition > 0
-            val addSpacingToEnd = itemPosition < (parent.adapter?.itemCount ?: 0) - 1
-            val isRtl = parent.layoutManager?.layoutDirection == ViewCompat.LAYOUT_DIRECTION_RTL
-            val density = parent.context.resources.displayMetrics.density
-            if (!isRtl) {
-                outRect.left = if (addSpacingToStart) ITEM_SPACING_DP.toPx(density) else 0
-                outRect.right = if (addSpacingToEnd) ITEM_SPACING_DP.toPx(density) else 0
-            } else {
-                outRect.left = if (addSpacingToEnd) ITEM_SPACING_DP.toPx(density) else 0
-                outRect.right = if (addSpacingToStart) ITEM_SPACING_DP.toPx(density) else 0
-            }
-        }
-
-        private fun Int.toPx(density: Float): Int {
-            return (this * density).toInt()
-        }
-
-        companion object {
-            private const val ITEM_SPACING_DP = 8
-        }
-    }
 }
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 fa7a344..ef68a4e 100644
--- a/src/com/android/customization/picker/color/ui/fragment/ColorPickerFragment.kt
+++ b/src/com/android/customization/picker/color/ui/fragment/ColorPickerFragment.kt
@@ -155,4 +155,8 @@
     override fun getDefaultTitle(): CharSequence {
         return requireContext().getString(R.string.color_picker_title)
     }
+
+    override fun getToolbarColorId(): Int {
+        return android.R.color.transparent
+    }
 }
diff --git a/src/com/android/customization/picker/color/ui/viewmodel/ColorPickerViewModel.kt b/src/com/android/customization/picker/color/ui/viewmodel/ColorPickerViewModel.kt
index 803663d..81a5810 100644
--- a/src/com/android/customization/picker/color/ui/viewmodel/ColorPickerViewModel.kt
+++ b/src/com/android/customization/picker/color/ui/viewmodel/ColorPickerViewModel.kt
@@ -31,9 +31,10 @@
 import kotlin.math.min
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.combine
-import kotlinx.coroutines.flow.flowOf
 import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.flow.stateIn
 import kotlinx.coroutines.launch
 
 /** Models UI state for a color picker experience. */
@@ -43,13 +44,13 @@
     private val interactor: ColorPickerInteractor,
 ) : ViewModel() {
 
-    private val selectedColorTypeId = MutableStateFlow<ColorType?>(null)
+    private val selectedColorTypeTabId = MutableStateFlow<ColorType?>(null)
 
-    /** View-models for each color type. */
-    val colorTypes: Flow<Map<ColorType, ColorTypeViewModel>> =
+    /** View-models for each color tab. */
+    val colorTypeTabs: Flow<Map<ColorType, ColorTypeTabViewModel>> =
         combine(
             interactor.colorOptions,
-            selectedColorTypeId,
+            selectedColorTypeTabId,
         ) { colorOptions, selectedColorTypeIdOrNull ->
             colorOptions.keys
                 .mapIndexed { index, colorType ->
@@ -57,26 +58,37 @@
                         (selectedColorTypeIdOrNull == null && index == 0) ||
                             selectedColorTypeIdOrNull == colorType
                     colorType to
-                        ColorTypeViewModel(
+                        ColorTypeTabViewModel(
                             name =
                                 when (colorType) {
                                     ColorType.WALLPAPER_COLOR ->
                                         context.resources.getString(R.string.wallpaper_color_tab)
-                                    ColorType.BASIC_COLOR ->
-                                        context.resources.getString(R.string.preset_color_tab)
+                                    ColorType.PRESET_COLOR ->
+                                        context.resources.getString(R.string.preset_color_tab_2)
                                 },
                             isSelected = isSelected,
                             onClick =
                                 if (isSelected) {
                                     null
                                 } else {
-                                    { this.selectedColorTypeId.value = colorType }
+                                    { this.selectedColorTypeTabId.value = colorType }
                                 },
                         )
                 }
                 .toMap()
         }
 
+    /** View-models for each color tab subheader */
+    val colorTypeTabSubheader: Flow<String> =
+        selectedColorTypeTabId.map { selectedColorTypeIdOrNull ->
+            when (selectedColorTypeIdOrNull ?: ColorType.WALLPAPER_COLOR) {
+                ColorType.WALLPAPER_COLOR ->
+                    context.resources.getString(R.string.wallpaper_color_subheader)
+                ColorType.PRESET_COLOR ->
+                    context.resources.getString(R.string.preset_color_subheader)
+            }
+        }
+
     /** The list of all color options mapped by their color type */
     private val allColorOptions:
         Flow<Map<ColorType, List<OptionItemViewModel<ColorOptionIconViewModel>>>> =
@@ -91,15 +103,19 @@
                                         colorOptionModel.colorOption as ColorSeedOption
                                     val colors =
                                         colorSeedOption.previewInfo.resolveColors(context.resources)
-                                    val isSelectedFlow: Flow<Boolean> =
-                                        interactor.activeColorOption.map {
-                                            it?.colorOption?.isEquivalent(
-                                                colorOptionModel.colorOption
-                                            )
-                                                ?: colorOptionModel.isSelected
-                                        }
+                                    val isSelectedFlow: StateFlow<Boolean> =
+                                        interactor.activeColorOption
+                                            .map {
+                                                it?.colorOption?.isEquivalent(
+                                                    colorOptionModel.colorOption
+                                                )
+                                                    ?: colorOptionModel.isSelected
+                                            }
+                                            .stateIn(viewModelScope)
                                     OptionItemViewModel<ColorOptionIconViewModel>(
-                                        key = flowOf(colorOptionModel.key),
+                                        key =
+                                            MutableStateFlow(colorOptionModel.key)
+                                                as StateFlow<String>,
                                         payload =
                                             ColorOptionIconViewModel(
                                                 colors[0],
@@ -129,7 +145,7 @@
                                     )
                                 }
                             }
-                            ColorType.BASIC_COLOR -> {
+                            ColorType.PRESET_COLOR -> {
                                 colorOptionEntry.value.map { colorOptionModel ->
                                     val colorBundle: ColorBundle =
                                         colorOptionModel.colorOption as ColorBundle
@@ -141,15 +157,19 @@
                                         colorBundle.previewInfo.resolveSecondaryColor(
                                             context.resources
                                         )
-                                    val isSelectedFlow: Flow<Boolean> =
-                                        interactor.activeColorOption.map {
-                                            it?.colorOption?.isEquivalent(
-                                                colorOptionModel.colorOption
-                                            )
-                                                ?: colorOptionModel.isSelected
-                                        }
+                                    val isSelectedFlow: StateFlow<Boolean> =
+                                        interactor.activeColorOption
+                                            .map {
+                                                it?.colorOption?.isEquivalent(
+                                                    colorOptionModel.colorOption
+                                                )
+                                                    ?: colorOptionModel.isSelected
+                                            }
+                                            .stateIn(viewModelScope)
                                     OptionItemViewModel<ColorOptionIconViewModel>(
-                                        key = flowOf(colorOptionModel.key),
+                                        key =
+                                            MutableStateFlow(colorOptionModel.key)
+                                                as StateFlow<String>,
                                         payload =
                                             ColorOptionIconViewModel(
                                                 primaryColor,
@@ -186,8 +206,9 @@
 
     /** The list of all available color options for the selected Color Type. */
     val colorOptions: Flow<List<OptionItemViewModel<ColorOptionIconViewModel>>> =
-        combine(allColorOptions, selectedColorTypeId) { allColorOptions, selectedColorTypeIdOrNull
-            ->
+        combine(allColorOptions, selectedColorTypeTabId) {
+            allColorOptions: Map<ColorType, List<OptionItemViewModel<ColorOptionIconViewModel>>>,
+            selectedColorTypeIdOrNull ->
             val selectedColorTypeId = selectedColorTypeIdOrNull ?: ColorType.WALLPAPER_COLOR
             allColorOptions[selectedColorTypeId]!!
         }
@@ -196,7 +217,7 @@
     val colorSectionOptions: Flow<List<OptionItemViewModel<ColorOptionIconViewModel>>> =
         allColorOptions.map { allColorOptions ->
             val wallpaperOptions = allColorOptions[ColorType.WALLPAPER_COLOR]
-            val presetOptions = allColorOptions[ColorType.BASIC_COLOR]
+            val presetOptions = allColorOptions[ColorType.PRESET_COLOR]
             val subOptions =
                 wallpaperOptions!!.subList(0, min(COLOR_SECTION_OPTION_SIZE, wallpaperOptions.size))
             // Add additional options based on preset colors if size of wallpaper color options is
diff --git a/src/com/android/customization/picker/color/ui/viewmodel/ColorTypeViewModel.kt b/src/com/android/customization/picker/color/ui/viewmodel/ColorTypeTabViewModel.kt
similarity index 96%
rename from src/com/android/customization/picker/color/ui/viewmodel/ColorTypeViewModel.kt
rename to src/com/android/customization/picker/color/ui/viewmodel/ColorTypeTabViewModel.kt
index 7343748..6a789cc 100644
--- a/src/com/android/customization/picker/color/ui/viewmodel/ColorTypeViewModel.kt
+++ b/src/com/android/customization/picker/color/ui/viewmodel/ColorTypeTabViewModel.kt
@@ -18,7 +18,7 @@
 package com.android.customization.picker.color.ui.viewmodel
 
 /** Models UI state for a single color type in a picker experience. */
-data class ColorTypeViewModel(
+data class ColorTypeTabViewModel(
     /** User-visible name for the color type. */
     val name: String,
 
diff --git a/src/com/android/customization/picker/preview/ui/section/PreviewWithClockCarouselSectionController.kt b/src/com/android/customization/picker/preview/ui/section/PreviewWithClockCarouselSectionController.kt
index 700439b..a2afc81 100644
--- a/src/com/android/customization/picker/preview/ui/section/PreviewWithClockCarouselSectionController.kt
+++ b/src/com/android/customization/picker/preview/ui/section/PreviewWithClockCarouselSectionController.kt
@@ -32,6 +32,7 @@
 import com.android.wallpaper.model.WallpaperColorsViewModel
 import com.android.wallpaper.module.CurrentWallpaperInfoFactory
 import com.android.wallpaper.module.CustomizationSections
+import com.android.wallpaper.picker.customization.domain.interactor.WallpaperInteractor
 import com.android.wallpaper.picker.customization.ui.section.ScreenPreviewSectionController
 import com.android.wallpaper.picker.customization.ui.section.ScreenPreviewView
 import com.android.wallpaper.util.DisplayUtils
@@ -48,6 +49,7 @@
     private val clockCarouselViewModel: ClockCarouselViewModel,
     private val clockViewFactory: ClockViewFactory,
     navigator: CustomizationSectionController.CustomizationSectionNavigationController,
+    wallpaperInteractor: WallpaperInteractor,
 ) :
     ScreenPreviewSectionController(
         activity,
@@ -57,6 +59,7 @@
         colorViewModel,
         displayUtils,
         navigator,
+        wallpaperInteractor,
     ) {
 
     private var clockCarouselBinding: ClockCarouselViewBinder.Binding? = null
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 cbc140e..14b6acc 100644
--- a/src/com/android/customization/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModel.kt
+++ b/src/com/android/customization/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModel.kt
@@ -148,7 +148,9 @@
                         selectedQuickAffordances =
                             selectedAffordances.map { affordanceModel ->
                                 OptionItemViewModel<Icon>(
-                                    key = flowOf("${slot.id}::${affordanceModel.id}"),
+                                    key =
+                                        MutableStateFlow("${slot.id}::${affordanceModel.id}")
+                                            as StateFlow<String>,
                                     payload =
                                         Icon.Loaded(
                                             drawable =
@@ -156,7 +158,7 @@
                                             contentDescription = null,
                                         ),
                                     text = Text.Loaded(affordanceModel.name),
-                                    isSelected = flowOf(true),
+                                    isSelected = MutableStateFlow(true) as StateFlow<Boolean>,
                                     onClicked = flowOf(null),
                                     onLongClicked = null,
                                     isEnabled = true,
@@ -196,7 +198,7 @@
     /** The list of all available quick affordances for the selected slot. */
     val quickAffordances: Flow<List<OptionItemViewModel<Icon>>> =
         quickAffordanceInteractor.affordances.map { affordances ->
-            val isNoneSelected = selectedAffordanceIds.map { it.isEmpty() }
+            val isNoneSelected = selectedAffordanceIds.map { it.isEmpty() }.stateIn(viewModelScope)
             listOf(
                 none(
                     slotId = selectedSlotId,
@@ -220,10 +222,15 @@
             ) +
                 affordances.map { affordance ->
                     val affordanceIcon = getAffordanceIcon(affordance.iconResourceId)
-                    val isSelectedFlow: Flow<Boolean> =
-                        selectedAffordanceIds.map { it.contains(affordance.id) }
+                    val isSelectedFlow: StateFlow<Boolean> =
+                        selectedAffordanceIds
+                            .map { it.contains(affordance.id) }
+                            .stateIn(viewModelScope)
                     OptionItemViewModel<Icon>(
-                        key = selectedSlotId.map { slotId -> "$slotId::${affordance.id}" },
+                        key =
+                            selectedSlotId
+                                .map { slotId -> "$slotId::${affordance.id}" }
+                                .stateIn(viewModelScope),
                         payload = Icon.Loaded(drawable = affordanceIcon, contentDescription = null),
                         text = Text.Loaded(affordance.name),
                         isSelected = isSelectedFlow,
@@ -359,13 +366,13 @@
 
     /** Returns a view-model for the special "None" option. */
     @SuppressLint("UseCompatLoadingForDrawables")
-    private fun none(
-        slotId: Flow<String>,
-        isSelected: Flow<Boolean>,
+    private suspend fun none(
+        slotId: StateFlow<String>,
+        isSelected: StateFlow<Boolean>,
         onSelected: Flow<(() -> Unit)?>,
     ): OptionItemViewModel<Icon> {
         return OptionItemViewModel<Icon>(
-            key = slotId.map { "$it::none" },
+            key = slotId.map { "$it::none" }.stateIn(viewModelScope),
             payload = Icon.Resource(res = R.drawable.link_off, contentDescription = null),
             text = Text.Resource(res = R.string.keyguard_affordance_none),
             isSelected = isSelected,
diff --git a/tests/src/com/android/customization/model/grid/ui/viewmodel/GridScreenViewModelTest.kt b/tests/src/com/android/customization/model/grid/ui/viewmodel/GridScreenViewModelTest.kt
index 301dbe8..58c5d99 100644
--- a/tests/src/com/android/customization/model/grid/ui/viewmodel/GridScreenViewModelTest.kt
+++ b/tests/src/com/android/customization/model/grid/ui/viewmodel/GridScreenViewModelTest.kt
@@ -22,7 +22,6 @@
 import com.android.customization.model.grid.data.repository.FakeGridRepository
 import com.android.customization.model.grid.domain.interactor.GridInteractor
 import com.android.customization.model.grid.domain.interactor.GridSnapshotRestorer
-import com.android.wallpaper.picker.common.icon.ui.viewmodel.Icon
 import com.android.wallpaper.picker.option.ui.viewmodel.OptionItemViewModel
 import com.android.wallpaper.testing.FakeSnapshotStore
 import com.android.wallpaper.testing.collectLastValue
@@ -95,13 +94,17 @@
             assertThat(getOnClick(optionItemsValue[1])).isNull()
         }
 
-    private fun TestScope.getSelectedIndex(optionItems: List<OptionItemViewModel<Icon>>): Int {
+    private fun TestScope.getSelectedIndex(
+        optionItems: List<OptionItemViewModel<GridIconViewModel>>
+    ): Int {
         return optionItems.indexOfFirst { optionItem ->
             collectLastValue(optionItem.isSelected).invoke() == true
         }
     }
 
-    private fun TestScope.getOnClick(optionItem: OptionItemViewModel<Icon>): (() -> Unit)? {
+    private fun TestScope.getOnClick(
+        optionItem: OptionItemViewModel<GridIconViewModel>
+    ): (() -> Unit)? {
         return collectLastValue(optionItem.onClicked).invoke()
     }
 }
diff --git a/tests/src/com/android/customization/model/picker/color/domain/interactor/ColorPickerInteractorTest.kt b/tests/src/com/android/customization/model/picker/color/domain/interactor/ColorPickerInteractorTest.kt
index 885d5a9..cbf1365 100644
--- a/tests/src/com/android/customization/model/picker/color/domain/interactor/ColorPickerInteractorTest.kt
+++ b/tests/src/com/android/customization/model/picker/color/domain/interactor/ColorPickerInteractorTest.kt
@@ -72,11 +72,11 @@
         val wallpaperColorOptionModelAfter = colorOptions()?.get(ColorType.WALLPAPER_COLOR)?.get(2)
         assertThat(wallpaperColorOptionModelAfter?.isSelected).isTrue()
 
-        val presetColorOptionModelBefore = colorOptions()?.get(ColorType.BASIC_COLOR)?.get(1)
+        val presetColorOptionModelBefore = colorOptions()?.get(ColorType.PRESET_COLOR)?.get(1)
         assertThat(presetColorOptionModelBefore?.isSelected).isFalse()
 
         presetColorOptionModelBefore?.let { underTest.select(colorOptionModel = it) }
-        val presetColorOptionModelAfter = colorOptions()?.get(ColorType.BASIC_COLOR)?.get(1)
+        val presetColorOptionModelAfter = colorOptions()?.get(ColorType.PRESET_COLOR)?.get(1)
         assertThat(presetColorOptionModelAfter?.isSelected).isTrue()
     }
 
diff --git a/tests/src/com/android/customization/model/picker/color/domain/interactor/ColorPickerSnapshotRestorerTest.kt b/tests/src/com/android/customization/model/picker/color/domain/interactor/ColorPickerSnapshotRestorerTest.kt
index 27b8550..71a8f23 100644
--- a/tests/src/com/android/customization/model/picker/color/domain/interactor/ColorPickerSnapshotRestorerTest.kt
+++ b/tests/src/com/android/customization/model/picker/color/domain/interactor/ColorPickerSnapshotRestorerTest.kt
@@ -70,9 +70,9 @@
         val initialSnapshot = underTest.setUpSnapshotRestorer(store = store)
         assertThat(initialSnapshot.args).isNotEmpty()
 
-        val colorOptionToSelect = colorOptions()?.get(ColorType.BASIC_COLOR)?.get(3)
+        val colorOptionToSelect = colorOptions()?.get(ColorType.PRESET_COLOR)?.get(3)
         colorOptionToSelect?.let { repository.select(it) }
-        assertState(colorOptions(), ColorType.BASIC_COLOR, 3)
+        assertState(colorOptions(), ColorType.PRESET_COLOR, 3)
 
         underTest.restoreToSnapshot(initialSnapshot)
         assertState(colorOptions(), ColorType.WALLPAPER_COLOR, 2)
@@ -86,11 +86,11 @@
         val initialSnapshot = underTest.setUpSnapshotRestorer(store = store)
         assertThat(initialSnapshot.args).isNotEmpty()
 
-        val colorOptionToSelect = colorOptions()?.get(ColorType.BASIC_COLOR)?.get(3)
+        val colorOptionToSelect = colorOptions()?.get(ColorType.PRESET_COLOR)?.get(3)
         colorOptionToSelect?.let { repository.select(it) }
-        assertState(colorOptions(), ColorType.BASIC_COLOR, 3)
+        assertState(colorOptions(), ColorType.PRESET_COLOR, 3)
 
-        val colorOptionToStore = colorOptions()?.get(ColorType.BASIC_COLOR)?.get(1)
+        val colorOptionToStore = colorOptions()?.get(ColorType.PRESET_COLOR)?.get(1)
         colorOptionToStore?.let { underTest.storeSnapshot(colorOptionToStore) }
 
         underTest.restoreToSnapshot(initialSnapshot)
diff --git a/tests/src/com/android/customization/model/picker/color/ui/viewmodel/ColorPickerViewModelTest.kt b/tests/src/com/android/customization/model/picker/color/ui/viewmodel/ColorPickerViewModelTest.kt
index 533d1dc..1d9457a 100644
--- a/tests/src/com/android/customization/model/picker/color/ui/viewmodel/ColorPickerViewModelTest.kt
+++ b/tests/src/com/android/customization/model/picker/color/ui/viewmodel/ColorPickerViewModelTest.kt
@@ -25,7 +25,7 @@
 import com.android.customization.picker.color.shared.model.ColorType
 import com.android.customization.picker.color.ui.viewmodel.ColorOptionIconViewModel
 import com.android.customization.picker.color.ui.viewmodel.ColorPickerViewModel
-import com.android.customization.picker.color.ui.viewmodel.ColorTypeViewModel
+import com.android.customization.picker.color.ui.viewmodel.ColorTypeTabViewModel
 import com.android.wallpaper.picker.option.ui.viewmodel.OptionItemViewModel
 import com.android.wallpaper.testing.FakeSnapshotStore
 import com.android.wallpaper.testing.collectLastValue
@@ -114,7 +114,7 @@
     @Test
     fun `Select a preset color`() =
         testScope.runTest {
-            val colorTypes = collectLastValue(underTest.colorTypes)
+            val colorTypes = collectLastValue(underTest.colorTypeTabs)
             val colorOptions = collectLastValue(underTest.colorOptions)
 
             // Initially, the wallpaper color tab should be selected
@@ -126,7 +126,7 @@
             )
 
             // Select "Basic colors" tab
-            colorTypes()?.get(ColorType.BASIC_COLOR)?.onClick?.invoke()
+            colorTypes()?.get(ColorType.PRESET_COLOR)?.onClick?.invoke()
             assertPickerUiState(
                 colorTypes = colorTypes(),
                 colorOptions = colorOptions(),
@@ -147,7 +147,7 @@
             )
 
             // Check new option is selected
-            colorTypes()?.get(ColorType.BASIC_COLOR)?.onClick?.invoke()
+            colorTypes()?.get(ColorType.PRESET_COLOR)?.onClick?.invoke()
             assertPickerUiState(
                 colorTypes = colorTypes(),
                 colorOptions = colorOptions(),
@@ -181,7 +181,7 @@
      *   -1 stands for no color option should be selected
      */
     private fun TestScope.assertPickerUiState(
-        colorTypes: Map<ColorType, ColorTypeViewModel>?,
+        colorTypes: Map<ColorType, ColorTypeTabViewModel>?,
         colorOptions: List<OptionItemViewModel<ColorOptionIconViewModel>>?,
         selectedColorTypeText: String,
         selectedColorOptionIndex: Int,
@@ -193,7 +193,7 @@
         )
         assertColorTypeTabUiState(
             colorTypes = colorTypes,
-            colorTypeId = ColorType.BASIC_COLOR,
+            colorTypeId = ColorType.PRESET_COLOR,
             isSelected = "Basic colors" == selectedColorTypeText,
         )
         assertColorOptionUiState(colorOptions, selectedColorOptionIndex)
@@ -251,7 +251,7 @@
      * @param isSelected Whether that color type should be selected
      */
     private fun assertColorTypeTabUiState(
-        colorTypes: Map<ColorType, ColorTypeViewModel>?,
+        colorTypes: Map<ColorType, ColorTypeTabViewModel>?,
         colorTypeId: ColorType,
         isSelected: Boolean,
     ) {
diff --git a/tests/src/com/android/customization/testing/TestCustomizationInjector.kt b/tests/src/com/android/customization/testing/TestCustomizationInjector.kt
index 5b0ecae..b49e654 100644
--- a/tests/src/com/android/customization/testing/TestCustomizationInjector.kt
+++ b/tests/src/com/android/customization/testing/TestCustomizationInjector.kt
@@ -2,7 +2,10 @@
 
 import android.app.Activity
 import android.content.Context
+import android.text.TextUtils
 import androidx.fragment.app.FragmentActivity
+import com.android.customization.model.color.ColorCustomizationManager
+import com.android.customization.model.color.ColorOptionsProvider
 import com.android.customization.model.theme.OverlayManagerCompat
 import com.android.customization.model.theme.ThemeBundleProvider
 import com.android.customization.model.theme.ThemeManager
@@ -16,6 +19,7 @@
 import com.android.customization.picker.clock.ui.viewmodel.ClockCarouselViewModel
 import com.android.customization.picker.clock.ui.viewmodel.ClockSectionViewModel
 import com.android.customization.picker.clock.ui.viewmodel.ClockSettingsViewModel
+import com.android.customization.picker.color.data.repository.ColorPickerRepository
 import com.android.customization.picker.color.data.repository.ColorPickerRepositoryImpl
 import com.android.customization.picker.color.domain.interactor.ColorPickerInteractor
 import com.android.customization.picker.color.domain.interactor.ColorPickerSnapshotRestorer
@@ -31,6 +35,9 @@
 import com.android.wallpaper.module.DrawableLayerResolver
 import com.android.wallpaper.module.PackageStatusNotifier
 import com.android.wallpaper.module.UserEventLogger
+import com.android.wallpaper.picker.customization.data.content.WallpaperClientImpl
+import com.android.wallpaper.picker.customization.data.repository.WallpaperRepository
+import com.android.wallpaper.picker.customization.domain.interactor.WallpaperInteractor
 import com.android.wallpaper.picker.undo.domain.interactor.SnapshotRestorer
 import com.android.wallpaper.testing.TestInjector
 import kotlinx.coroutines.Dispatchers
@@ -43,6 +50,7 @@
     private var packageStatusNotifier: PackageStatusNotifier? = null
     private var drawableLayerResolver: DrawableLayerResolver? = null
     private var userEventLogger: UserEventLogger? = null
+    private var wallpaperInteractor: WallpaperInteractor? = null
     private var keyguardQuickAffordancePickerInteractor: KeyguardQuickAffordancePickerInteractor? =
         null
     private var flags: BaseFlags? = null
@@ -53,9 +61,11 @@
     private var clockPickerInteractor: ClockPickerInteractor? = null
     private var clockSectionViewModel: ClockSectionViewModel? = null
     private var clockViewFactory: ClockViewFactory? = null
+    private var colorPickerRepository: ColorPickerRepository? = null
     private var colorPickerInteractor: ColorPickerInteractor? = null
     private var colorPickerViewModelFactory: ColorPickerViewModel.Factory? = null
     private var colorPickerSnapshotRestorer: ColorPickerSnapshotRestorer? = null
+    private var colorCustomizationManager: ColorCustomizationManager? = null
     private var clockCarouselViewModel: ClockCarouselViewModel? = null
     private var clockSettingsViewModelFactory: ClockSettingsViewModel.Factory? = null
 
@@ -92,6 +102,25 @@
         return userEventLogger ?: TestThemesUserEventLogger().also { userEventLogger = it }
     }
 
+    override fun getWallpaperInteractor(context: Context): WallpaperInteractor {
+        return wallpaperInteractor
+            ?: WallpaperInteractor(
+                    repository =
+                        WallpaperRepository(
+                            scope = GlobalScope,
+                            client = WallpaperClientImpl(context = context),
+                            backgroundDispatcher = Dispatchers.IO,
+                        ),
+                    shouldHandleReload = {
+                        TextUtils.equals(
+                            getColorCustomizationManager(context).currentColorSource,
+                            ColorOptionsProvider.COLOR_SOURCE_PRESET
+                        )
+                    }
+                )
+                .also { wallpaperInteractor = it }
+    }
+
     override fun getKeyguardQuickAffordancePickerInteractor(
         context: Context
     ): KeyguardQuickAffordancePickerInteractor {
@@ -169,13 +198,24 @@
             }
     }
 
+    private fun getColorPickerRepository(
+        context: Context,
+        wallpaperColorsViewModel: WallpaperColorsViewModel,
+    ): ColorPickerRepository {
+        return colorPickerRepository
+            ?: ColorPickerRepositoryImpl(
+                wallpaperColorsViewModel,
+                getColorCustomizationManager(context)
+            )
+    }
+
     override fun getColorPickerInteractor(
         context: Context,
         wallpaperColorsViewModel: WallpaperColorsViewModel,
     ): ColorPickerInteractor {
         return colorPickerInteractor
             ?: ColorPickerInteractor(
-                    repository = ColorPickerRepositoryImpl(context, wallpaperColorsViewModel),
+                    repository = getColorPickerRepository(context, wallpaperColorsViewModel),
                     snapshotRestorer = {
                         getColorPickerSnapshotRestorer(context, wallpaperColorsViewModel)
                     },
@@ -206,6 +246,13 @@
                 .also { colorPickerSnapshotRestorer = it }
     }
 
+    private fun getColorCustomizationManager(context: Context): ColorCustomizationManager {
+        return colorCustomizationManager
+            ?: ColorCustomizationManager.getInstance(context, OverlayManagerCompat(context)).also {
+                colorCustomizationManager = it
+            }
+    }
+
     override fun getClockCarouselViewModel(context: Context): ClockCarouselViewModel {
         return clockCarouselViewModel
             ?: ClockCarouselViewModel(getClockPickerInteractor(context)).also {