diff --git a/res/drawable/edit_icon.xml b/res/drawable/edit_icon.xml
deleted file mode 100644
index 9690d17..0000000
--- a/res/drawable/edit_icon.xml
+++ /dev/null
@@ -1,20 +0,0 @@
-<vector xmlns:android="http://schemas.android.com/apk/res/android"
-    android:width="48dp"
-    android:height="48dp"
-    android:viewportWidth="48"
-    android:viewportHeight="48">
-  <group>
-    <clip-path
-        android:pathData="M8,24C8,15.163 15.163,8 24,8C32.837,8 40,15.163 40,24C40,32.837 32.837,40 24,40C15.163,40 8,32.837 8,24Z"/>
-    <path
-        android:pathData="M8,24C8,15.163 15.163,8 24,8C32.837,8 40,15.163 40,24C40,32.837 32.837,40 24,40C15.163,40 8,32.837 8,24Z"
-        android:fillColor="@color/system_on_primary_fixed_variant"/>
-    <group>
-      <clip-path
-          android:pathData="M14,14h20v20h-20z"/>
-      <path
-          android:pathData="M17,31V27.813L27.375,17.438C27.528,17.285 27.694,17.174 27.875,17.104C28.056,17.035 28.243,17 28.438,17C28.632,17 28.819,17.035 29,17.104C29.181,17.174 29.347,17.285 29.5,17.438L30.563,18.5C30.715,18.653 30.826,18.819 30.896,19C30.965,19.181 31,19.368 31,19.563C31,19.757 30.965,19.944 30.896,20.125C30.826,20.306 30.715,20.472 30.563,20.625L20.188,31H17ZM28.438,20.625L29.5,19.563L28.438,18.5L27.375,19.563L28.438,20.625Z"
-          android:fillColor="#ffffff"/>
-    </group>
-  </group>
-</vector>
diff --git a/res/drawable/edit_icon_background.xml b/res/drawable/edit_icon_background.xml
new file mode 100644
index 0000000..920e57b
--- /dev/null
+++ b/res/drawable/edit_icon_background.xml
@@ -0,0 +1,11 @@
+<vector xmlns:android="http://schemas.android.com/apk/res/android"
+    android:width="48dp"
+    android:height="48dp"
+    android:viewportWidth="48"
+    android:viewportHeight="48">
+    <clip-path
+        android:pathData="M8,24C8,15.163 15.163,8 24,8C32.837,8 40,15.163 40,24C40,32.837 32.837,40 24,40C15.163,40 8,32.837 8,24Z"/>
+    <path
+        android:pathData="M8,24C8,15.163 15.163,8 24,8C32.837,8 40,15.163 40,24C40,32.837 32.837,40 24,40C15.163,40 8,32.837 8,24Z"
+        android:fillColor="@color/system_on_primary_fixed_variant"/>
+</vector>
diff --git a/res/drawable/edit_icon_foreground.xml b/res/drawable/edit_icon_foreground.xml
new file mode 100644
index 0000000..f98ab93
--- /dev/null
+++ b/res/drawable/edit_icon_foreground.xml
@@ -0,0 +1,11 @@
+<vector xmlns:android="http://schemas.android.com/apk/res/android"
+    android:width="48dp"
+    android:height="48dp"
+    android:viewportWidth="48"
+    android:viewportHeight="48">
+    <clip-path
+        android:pathData="M14,14h20v20h-20z"/>
+    <path
+        android:pathData="M17,31V27.813L27.375,17.438C27.528,17.285 27.694,17.174 27.875,17.104C28.056,17.035 28.243,17 28.438,17C28.632,17 28.819,17.035 29,17.104C29.181,17.174 29.347,17.285 29.5,17.438L30.563,18.5C30.715,18.653 30.826,18.819 30.896,19C30.965,19.181 31,19.368 31,19.563C31,19.757 30.965,19.944 30.896,20.125C30.826,20.306 30.715,20.472 30.563,20.625L20.188,31H17ZM28.438,20.625L29.5,19.563L28.438,18.5L27.375,19.563L28.438,20.625Z"
+        android:fillColor="#ffffff"/>
+</vector>
diff --git a/res/layout/clock_style_option.xml b/res/layout/clock_style_option.xml
index e251c4a..fa20706 100644
--- a/res/layout/clock_style_option.xml
+++ b/res/layout/clock_style_option.xml
@@ -42,15 +42,29 @@
         app:layout_constraintBottom_toBottomOf="parent"
         android:src="@drawable/ic_clock_24px" />
 
-    <ImageView
+    <FrameLayout
         android:id="@+id/edit_icon"
         android:layout_width="@dimen/floating_sheet_clock_edit_icon_size"
         android:layout_height="@dimen/floating_sheet_clock_edit_icon_size"
         android:layout_marginTop="@dimen/floating_sheet_clock_edit_icon_margin"
         android:layout_marginEnd="@dimen/floating_sheet_clock_edit_icon_margin"
-        android:src="@drawable/edit_icon"
         app:layout_constraintEnd_toEndOf="@+id/background"
-        app:layout_constraintTop_toTopOf="@+id/background"
-        android:importantForAccessibility="no" />
+        app:layout_constraintTop_toTopOf="@+id/background" >
+        <ImageView
+            android:id="@+id/edit_icon_background"
+            android:layout_width="match_parent"
+            android:layout_height="match_parent"
+            android:src="@drawable/edit_icon_background"
+            android:importantForAccessibility="no" />
+
+        <ImageView
+            android:id="@+id/edit_icon_foreground"
+            android:layout_width="match_parent"
+            android:layout_height="match_parent"
+            android:src="@drawable/edit_icon_foreground"
+            android:importantForAccessibility="no" />
+    </FrameLayout>
+
+
 </androidx.constraintlayout.widget.ConstraintLayout>
 
diff --git a/res/layout/floating_sheet_clock_color_content.xml b/res/layout/floating_sheet_clock_color_content.xml
index fd218c6..acd274a 100644
--- a/res/layout/floating_sheet_clock_color_content.xml
+++ b/res/layout/floating_sheet_clock_color_content.xml
@@ -23,12 +23,29 @@
     android:clipToPadding="false"
     android:clipChildren="false">
 
+    <TextView
+        android:id="@+id/clock_color_title"
+        style="@style/CustomizationOptionEntryTitleTextStyle"
+        android:layout_width="match_parent"
+        android:layout_height="wrap_content"
+        android:text="@string/clock_color_title"
+        android:paddingHorizontal="@dimen/floating_sheet_content_horizontal_padding" />
+
+    <TextView
+        android:id="@+id/clock_color_description"
+        style="@style/CustomizationOptionEntrySubtitleTextStyle"
+        android:layout_width="match_parent"
+        android:layout_height="wrap_content"
+        android:paddingHorizontal="@dimen/floating_sheet_content_horizontal_padding"
+        android:text="@string/clock_color_description"
+        android:layout_marginBottom="16dp" />
+
     <FrameLayout
         android:layout_width="match_parent"
         android:layout_height="wrap_content"
         android:clipToPadding="false"
         android:clipChildren="false"
-        android:layout_marginBottom="@dimen/floating_sheet_clock_color_option_list_bottom_margin">
+        android:layout_marginBottom="20dp">
 
         <!--
         This is an invisible placeholder put in place so that the parent keeps its height
@@ -40,15 +57,14 @@
         It's critical for any TextViews inside the included layout to have text.
         -->
         <include
-            layout="@layout/color_option"
-            android:layout_width="@dimen/option_item_size"
-            android:layout_height="@dimen/option_item_size"
+            layout="@layout/color_option2"
             android:visibility="invisible" />
 
         <androidx.recyclerview.widget.RecyclerView
             android:id="@+id/clock_color_list"
             android:layout_width="match_parent"
             android:layout_height="wrap_content"
+            android:paddingHorizontal="@dimen/floating_sheet_content_horizontal_padding"
             android:clipChildren="false"
             android:clipToPadding="false" />
     </FrameLayout>
diff --git a/res/values/dimens.xml b/res/values/dimens.xml
index be1e071..756ae2d 100644
--- a/res/values/dimens.xml
+++ b/res/values/dimens.xml
@@ -200,7 +200,6 @@
     <dimen name="floating_sheet_clock_edit_icon_margin">-18dp</dimen>
     <dimen name="floating_sheet_clock_style_thumbnail_margin_bottom">12dp</dimen>
     <dimen name="floating_sheet_clock_style_clock_size_text_margin_end">16dp</dimen>
-    <dimen name="floating_sheet_clock_color_option_list_bottom_margin">12dp</dimen>
     <dimen name="floating_sheet_color_option_size">54dp</dimen>
     <dimen name="floating_sheet_color_option_stroke_width">3dp</dimen>
     <dimen name="customization_option_entry_shortcut_icon_size">20dp</dimen>
diff --git a/res/values/strings.xml b/res/values/strings.xml
index b2ff8ac..8fdc97e 100755
--- a/res/values/strings.xml
+++ b/res/values/strings.xml
@@ -51,6 +51,12 @@
     <!-- Title of a tab to change the clock color. [CHAR LIMIT=15] -->
     <string name="clock_color">Color</string>
 
+    <!-- Title of a bottom sheet for changing the clock color. [CHAR LIMIT=15] -->
+    <string name="clock_color_title">Clock color</string>
+
+    <!-- Description of a bottom sheet for changing the clock color. [CHAR LIMIT=55] -->
+    <string name="clock_color_description">Match the wallpaper theme or choose your own color</string>
+
     <!-- Title of a clock color red option that will appear in the description of the custom clock section. [CHAR LIMIT=15] -->
     <string name="clock_color_red">Red</string>
 
diff --git a/src/com/android/customization/picker/color/ui/binder/ColorOptionIconBinder2.kt b/src/com/android/customization/picker/color/ui/binder/ColorOptionIconBinder2.kt
index 2c197ad..4758619 100644
--- a/src/com/android/customization/picker/color/ui/binder/ColorOptionIconBinder2.kt
+++ b/src/com/android/customization/picker/color/ui/binder/ColorOptionIconBinder2.kt
@@ -17,14 +17,36 @@
 
 package com.android.customization.picker.color.ui.binder
 
+import androidx.lifecycle.LifecycleOwner
 import com.android.customization.picker.color.ui.view.ColorOptionIconView2
 import com.android.customization.picker.color.ui.viewmodel.ColorOptionIconViewModel
+import com.android.wallpaper.picker.customization.ui.binder.ColorUpdateBinder
+import com.android.wallpaper.picker.customization.ui.viewmodel.ColorUpdateViewModel
 
 object ColorOptionIconBinder2 {
-    fun bind(view: ColorOptionIconView2, viewModel: ColorOptionIconViewModel, darkTheme: Boolean) {
+
+    interface Binding {
+        /** Destroys the color update binding, in spite of lifecycle state. */
+        fun destroy()
+    }
+
+    fun bind(
+        view: ColorOptionIconView2,
+        viewModel: ColorOptionIconViewModel,
+        darkTheme: Boolean,
+        colorUpdateViewModel: ColorUpdateViewModel,
+        shouldAnimateColor: () -> Boolean,
+        lifecycleOwner: LifecycleOwner,
+    ): Binding {
+        val binding =
+            ColorUpdateBinder.bind(
+                setColor = { color -> view.bindStrokeColor(color) },
+                color = colorUpdateViewModel.colorPrimary,
+                shouldAnimate = shouldAnimateColor,
+                lifecycleOwner = lifecycleOwner,
+            )
         if (darkTheme) {
             view.bindColor(
-                view.resources.getColor(android.R.color.system_primary_dark, view.context.theme),
                 viewModel.darkThemeColor0,
                 viewModel.darkThemeColor1,
                 viewModel.darkThemeColor2,
@@ -32,12 +54,16 @@
             )
         } else {
             view.bindColor(
-                view.resources.getColor(android.R.color.system_primary_light, view.context.theme),
                 viewModel.lightThemeColor0,
                 viewModel.lightThemeColor1,
                 viewModel.lightThemeColor2,
                 viewModel.lightThemeColor3,
             )
         }
+        return object : Binding {
+            override fun destroy() {
+                binding.destroy()
+            }
+        }
     }
 }
diff --git a/src/com/android/customization/picker/color/ui/view/ColorOptionIconView2.kt b/src/com/android/customization/picker/color/ui/view/ColorOptionIconView2.kt
index 3fc6324..ea20af3 100644
--- a/src/com/android/customization/picker/color/ui/view/ColorOptionIconView2.kt
+++ b/src/com/android/customization/picker/color/ui/view/ColorOptionIconView2.kt
@@ -55,13 +55,11 @@
      * @param color3 the color in the bottom right quadrant
      */
     fun bindColor(
-        @ColorInt strokeColor: Int,
         @ColorInt color0: Int,
         @ColorInt color1: Int,
         @ColorInt color2: Int,
         @ColorInt color3: Int,
     ) {
-        this.strokeColor = strokeColor
         this.color0 = color0
         this.color1 = color1
         this.color2 = color2
@@ -69,6 +67,11 @@
         invalidate()
     }
 
+    fun bindStrokeColor(@ColorInt strokeColor: Int) {
+        this.strokeColor = strokeColor
+        invalidate()
+    }
+
     override fun onSizeChanged(w: Int, h: Int, oldw: Int, oldh: Int) {
         this.w = w
         this.h = h
diff --git a/src/com/android/wallpaper/customization/ui/binder/ClockFloatingSheetBinder.kt b/src/com/android/wallpaper/customization/ui/binder/ClockFloatingSheetBinder.kt
index 943729b..a5dd89a 100644
--- a/src/com/android/wallpaper/customization/ui/binder/ClockFloatingSheetBinder.kt
+++ b/src/com/android/wallpaper/customization/ui/binder/ClockFloatingSheetBinder.kt
@@ -20,10 +20,12 @@
 import android.animation.AnimatorListenerAdapter
 import android.animation.ValueAnimator
 import android.content.Context
-import android.content.res.Configuration
+import android.content.res.Configuration.UI_MODE_NIGHT_MASK
+import android.content.res.Configuration.UI_MODE_NIGHT_YES
 import android.view.View
 import android.view.ViewGroup
 import android.view.ViewTreeObserver.OnGlobalLayoutListener
+import android.widget.FrameLayout
 import android.widget.ImageView
 import android.widget.SeekBar
 import android.widget.Switch
@@ -36,8 +38,8 @@
 import androidx.recyclerview.widget.LinearLayoutManager
 import androidx.recyclerview.widget.RecyclerView
 import com.android.customization.picker.clock.shared.ClockSize
-import com.android.customization.picker.color.ui.binder.ColorOptionIconBinder
-import com.android.customization.picker.color.ui.view.ColorOptionIconView
+import com.android.customization.picker.color.ui.binder.ColorOptionIconBinder2
+import com.android.customization.picker.color.ui.view.ColorOptionIconView2
 import com.android.customization.picker.color.ui.viewmodel.ColorOptionIconViewModel
 import com.android.customization.picker.common.ui.view.SingleRowListItemSpacing
 import com.android.systemui.plugins.clocks.AxisType
@@ -54,7 +56,6 @@
 import com.android.wallpaper.picker.customization.ui.view.FloatingToolbar
 import com.android.wallpaper.picker.customization.ui.view.adapter.FloatingToolbarTabAdapter
 import com.android.wallpaper.picker.customization.ui.viewmodel.ColorUpdateViewModel
-import com.android.wallpaper.picker.option.ui.adapter.OptionItemAdapter
 import com.android.wallpaper.picker.option.ui.adapter.OptionItemAdapter2
 import java.lang.ref.WeakReference
 import kotlinx.coroutines.DisposableHandle
@@ -144,10 +145,17 @@
         // Clock color
         val clockColorContent = view.requireViewById<View>(R.id.clock_floating_sheet_color_content)
         val clockColorAdapter =
-            createClockColorOptionItemAdapter(view.resources.configuration.uiMode, lifecycleOwner)
+            createClockColorOptionItemAdapter(
+                uiMode = view.resources.configuration.uiMode,
+                colorUpdateViewModel = colorUpdateViewModel,
+                shouldAnimateColor = isFloatingSheetActive,
+                lifecycleOwner = lifecycleOwner,
+            )
         val clockColorList =
             view.requireViewById<RecyclerView>(R.id.clock_color_list).apply {
-                initColorList(appContext, clockColorAdapter)
+                adapter = clockColorAdapter
+                layoutManager =
+                    LinearLayoutManager(appContext, LinearLayoutManager.HORIZONTAL, false)
             }
         val clockColorSlider: SeekBar = view.requireViewById(R.id.clock_color_slider)
         clockColorSlider.setOnSeekBarChangeListener(
@@ -438,11 +446,25 @@
                     lifecycleOwner.lifecycleScope.launch {
                         lifecycleOwner.repeatOnLifecycle(Lifecycle.State.STARTED) {
                             styleModel.showEditButton.collect {
-                                view.findViewById<ImageView>(R.id.edit_icon)?.isVisible = it
+                                view.findViewById<FrameLayout>(R.id.edit_icon)?.isVisible = it
                             }
                         }
                     }
-                return@OptionItemAdapter2 DisposableHandle { job.cancel() }
+                val binding =
+                    ColorUpdateBinder.bind(
+                        setColor = { color ->
+                            view.findViewById<ImageView>(R.id.edit_icon_background)?.drawable?.let {
+                                DrawableCompat.setTint(DrawableCompat.wrap(it), color)
+                            }
+                        },
+                        color = colorUpdateViewModel.colorOnPrimaryFixedVariant,
+                        shouldAnimate = shouldAnimateColor,
+                        lifecycleOwner = lifecycleOwner,
+                    )
+                return@OptionItemAdapter2 DisposableHandle {
+                    job.cancel()
+                    binding.destroy()
+                }
             },
             colorUpdateViewModel = WeakReference(colorUpdateViewModel),
             shouldAnimateColor = shouldAnimateColor,
@@ -468,39 +490,34 @@
 
     private fun createClockColorOptionItemAdapter(
         uiMode: Int,
+        colorUpdateViewModel: ColorUpdateViewModel,
+        shouldAnimateColor: () -> Boolean,
         lifecycleOwner: LifecycleOwner,
-    ): OptionItemAdapter<ColorOptionIconViewModel> =
-        OptionItemAdapter(
-            layoutResourceId = R.layout.color_option,
+    ): OptionItemAdapter2<ColorOptionIconViewModel> =
+        OptionItemAdapter2(
+            layoutResourceId = R.layout.color_option2,
             lifecycleOwner = lifecycleOwner,
-            bindIcon = { foregroundView: View, colorIcon: ColorOptionIconViewModel ->
-                val colorOptionIconView = foregroundView as? ColorOptionIconView
-                val night =
-                    uiMode and Configuration.UI_MODE_NIGHT_MASK == Configuration.UI_MODE_NIGHT_YES
-                colorOptionIconView?.let { ColorOptionIconBinder.bind(it, colorIcon, night) }
+            bindPayload = { itemView: View, colorIcon: ColorOptionIconViewModel ->
+                val colorOptionIconView =
+                    itemView.requireViewById<ColorOptionIconView2>(
+                        com.android.wallpaper.R.id.background
+                    )
+                val night = uiMode and UI_MODE_NIGHT_MASK == UI_MODE_NIGHT_YES
+                val binding =
+                    ColorOptionIconBinder2.bind(
+                        view = colorOptionIconView,
+                        viewModel = colorIcon,
+                        darkTheme = night,
+                        colorUpdateViewModel = colorUpdateViewModel,
+                        shouldAnimateColor = shouldAnimateColor,
+                        lifecycleOwner = lifecycleOwner,
+                    )
+                return@OptionItemAdapter2 DisposableHandle { binding.destroy() }
             },
+            colorUpdateViewModel = WeakReference(colorUpdateViewModel),
+            shouldAnimateColor = shouldAnimateColor,
         )
 
-    private fun RecyclerView.initColorList(
-        context: Context,
-        adapter: OptionItemAdapter<ColorOptionIconViewModel>,
-    ) {
-        apply {
-            this.adapter = adapter
-            layoutManager = LinearLayoutManager(context, LinearLayoutManager.HORIZONTAL, false)
-            addItemDecoration(
-                SingleRowListItemSpacing(
-                    context.resources.getDimensionPixelSize(
-                        R.dimen.floating_sheet_content_horizontal_padding
-                    ),
-                    context.resources.getDimensionPixelSize(
-                        R.dimen.floating_sheet_list_item_horizontal_space
-                    ),
-                )
-            )
-        }
-    }
-
     // Alpha is 1 when current height is from height, and 0 when current height is to height.
     private fun getAlpha(fromHeight: Int, toHeight: Int, currentHeight: Int): Float =
         (1 - (currentHeight - fromHeight).toFloat() / (toHeight - fromHeight).toFloat())
diff --git a/src/com/android/wallpaper/customization/ui/binder/ColorsFloatingSheetBinder.kt b/src/com/android/wallpaper/customization/ui/binder/ColorsFloatingSheetBinder.kt
index 4f5adab..a097828 100644
--- a/src/com/android/wallpaper/customization/ui/binder/ColorsFloatingSheetBinder.kt
+++ b/src/com/android/wallpaper/customization/ui/binder/ColorsFloatingSheetBinder.kt
@@ -41,6 +41,7 @@
 import com.android.wallpaper.picker.customization.ui.viewmodel.ColorUpdateViewModel
 import com.android.wallpaper.picker.option.ui.adapter.OptionItemAdapter2
 import java.lang.ref.WeakReference
+import kotlinx.coroutines.DisposableHandle
 import kotlinx.coroutines.launch
 
 object ColorsFloatingSheetBinder {
@@ -158,10 +159,16 @@
                         com.android.wallpaper.R.id.background
                     )
                 val night = uiMode and UI_MODE_NIGHT_MASK == UI_MODE_NIGHT_YES
-                ColorOptionIconBinder2.bind(colorOptionIconView, colorIcon, night)
-                // Return null since it does not need the lifecycleOwner to launch any job for later
-                // disposal when rebind.
-                return@OptionItemAdapter2 null
+                val binding =
+                    ColorOptionIconBinder2.bind(
+                        view = colorOptionIconView,
+                        viewModel = colorIcon,
+                        darkTheme = night,
+                        colorUpdateViewModel = colorUpdateViewModel,
+                        shouldAnimateColor = shouldAnimateColor,
+                        lifecycleOwner = lifecycleOwner,
+                    )
+                return@OptionItemAdapter2 DisposableHandle { binding.destroy() }
             },
             colorUpdateViewModel = WeakReference(colorUpdateViewModel),
             shouldAnimateColor = shouldAnimateColor,
diff --git a/src/com/android/wallpaper/customization/ui/binder/ThemePickerCustomizationOptionBinder.kt b/src/com/android/wallpaper/customization/ui/binder/ThemePickerCustomizationOptionBinder.kt
index 1faf971..4bbb789 100644
--- a/src/com/android/wallpaper/customization/ui/binder/ThemePickerCustomizationOptionBinder.kt
+++ b/src/com/android/wallpaper/customization/ui/binder/ThemePickerCustomizationOptionBinder.kt
@@ -95,6 +95,11 @@
             navigateToLockScreenNotificationsSettingsActivity,
         )
 
+        val optionsViewModel =
+            viewModel.customizationOptionsViewModel as ThemePickerCustomizationOptionsViewModel
+
+        val isOnMainScreen = { optionsViewModel.selectedOption.value == null }
+
         val optionClock: View =
             lockScreenCustomizationOptionEntries
                 .first { it.first == ThemePickerLockCustomizationOption.CLOCK }
@@ -161,8 +166,6 @@
         val optionThemedIconsSwitch =
             optionThemedIcons?.findViewById<Switch>(R.id.option_entry_themed_icons_switch)
 
-        val optionsViewModel =
-            viewModel.customizationOptionsViewModel as ThemePickerCustomizationOptionsViewModel
         lifecycleOwner.lifecycleScope.launch {
             lifecycleOwner.repeatOnLifecycle(Lifecycle.State.STARTED) {
                 launch {
@@ -252,14 +255,21 @@
                 }
 
                 launch {
+                    var binding: ColorOptionIconBinder2.Binding? = null
                     optionsViewModel.colorPickerViewModel2.selectedColorOption.collect { colorOption
                         ->
                         (colorOption as? ColorOptionImpl)?.let {
-                            ColorOptionIconBinder2.bind(
-                                view = optionColorsIcon,
-                                viewModel = ColorOptionIconViewModel.fromColorOption(colorOption),
-                                darkTheme = view.resources.configuration.isNightModeActive,
-                            )
+                            binding?.destroy()
+                            binding =
+                                ColorOptionIconBinder2.bind(
+                                    view = optionColorsIcon,
+                                    viewModel =
+                                        ColorOptionIconViewModel.fromColorOption(colorOption),
+                                    darkTheme = view.resources.configuration.isNightModeActive,
+                                    colorUpdateViewModel = colorUpdateViewModel,
+                                    shouldAnimateColor = isOnMainScreen,
+                                    lifecycleOwner = lifecycleOwner,
+                                )
                         }
                     }
                 }
diff --git a/src/com/android/wallpaper/customization/ui/viewmodel/ClockPickerViewModel.kt b/src/com/android/wallpaper/customization/ui/viewmodel/ClockPickerViewModel.kt
index e055d48..a5f5d22 100644
--- a/src/com/android/wallpaper/customization/ui/viewmodel/ClockPickerViewModel.kt
+++ b/src/com/android/wallpaper/customization/ui/viewmodel/ClockPickerViewModel.kt
@@ -34,7 +34,6 @@
 import com.android.wallpaper.picker.common.text.ui.viewmodel.Text
 import com.android.wallpaper.picker.customization.ui.viewmodel.FloatingToolbarTabViewModel
 import com.android.wallpaper.picker.di.modules.BackgroundDispatcher
-import com.android.wallpaper.picker.option.ui.viewmodel.OptionItemViewModel
 import com.android.wallpaper.picker.option.ui.viewmodel.OptionItemViewModel2
 import dagger.assisted.Assisted
 import dagger.assisted.AssistedFactory
@@ -290,7 +289,7 @@
             }
         }
 
-    val clockColorOptions: Flow<List<OptionItemViewModel<ColorOptionIconViewModel>>> =
+    val clockColorOptions: Flow<List<OptionItemViewModel2<ColorOptionIconViewModel>>> =
         colorPickerInteractor.selectedColorOption.map { selectedColorOption ->
             // Use mapLatest and delay(100) here to prevent too many selectedClockColor update
             // events from ClockRegistry upstream, caused by sliding the saturation level bar.
@@ -304,7 +303,7 @@
                             .map { colorMap.keys.indexOf(it) == index }
                             .stateIn(viewModelScope)
                     add(
-                        OptionItemViewModel<ColorOptionIconViewModel>(
+                        OptionItemViewModel2<ColorOptionIconViewModel>(
                             key = MutableStateFlow(colorModel.colorId) as StateFlow<String>,
                             payload =
                                 ColorOptionIconViewModel(
@@ -346,7 +345,7 @@
 
     private suspend fun ColorOption.toOptionItemViewModel(
         context: Context
-    ): OptionItemViewModel<ColorOptionIconViewModel> {
+    ): OptionItemViewModel2<ColorOptionIconViewModel> {
         val lightThemeColors =
             (this as ColorOptionImpl)
                 .previewInfo
@@ -362,7 +361,7 @@
         val isSelectedFlow =
             previewingClockColorId.map { it == DEFAULT_CLOCK_COLOR_ID }.stateIn(viewModelScope)
         val key = "${this.type}::${this.style}::${this.serializedPackages}"
-        return OptionItemViewModel<ColorOptionIconViewModel>(
+        return OptionItemViewModel2<ColorOptionIconViewModel>(
             key = MutableStateFlow(key) as StateFlow<String>,
             payload =
                 ColorOptionIconViewModel(
diff --git a/src/com/android/wallpaper/customization/ui/viewmodel/ShapeGridPickerViewModel.kt b/src/com/android/wallpaper/customization/ui/viewmodel/ShapeGridPickerViewModel.kt
index 85422dd..60b5c35 100644
--- a/src/com/android/wallpaper/customization/ui/viewmodel/ShapeGridPickerViewModel.kt
+++ b/src/com/android/wallpaper/customization/ui/viewmodel/ShapeGridPickerViewModel.kt
@@ -133,18 +133,23 @@
             .shareIn(scope = viewModelScope, started = SharingStarted.Lazily, replay = 1)
 
     val onApply: Flow<(suspend () -> Unit)?> =
-        combine(previewingGridKey, selectedGridOption, previewingShapeKey, selectedShapeKey) {
-            previewingGridOptionKey,
+        combine(overridingGridKey, selectedGridOption, overridingShapeKey, selectedShapeKey) {
+            overridingGridKey,
             selectedGridOption,
-            previewingShapeKey,
+            overridingShapeKey,
             selectedShapeKey ->
             if (
-                previewingGridOptionKey == selectedGridOption.key.value &&
-                    previewingShapeKey == selectedShapeKey
+                (overridingGridKey != null && overridingGridKey != selectedGridOption.key.value) ||
+                    (overridingShapeKey != null && overridingShapeKey != selectedShapeKey)
             ) {
-                null
+                {
+                    interactor.applySelectedOption(
+                        overridingShapeKey ?: selectedShapeKey,
+                        overridingGridKey ?: selectedGridOption.key.value,
+                    )
+                }
             } else {
-                { interactor.applySelectedOption(previewingShapeKey, previewingGridOptionKey) }
+                null
             }
         }
 
diff --git a/tests/robotests/Android.bp b/tests/robotests/Android.bp
index a4fd7f3..e775a4d 100644
--- a/tests/robotests/Android.bp
+++ b/tests/robotests/Android.bp
@@ -1,4 +1,3 @@
-
 package {
     default_team: "trendy_team_android_kernel",
     default_applicable_licenses: ["Android-Apache-2.0"],
@@ -43,8 +42,6 @@
         "androidx.test.runner",
     ],
 
-    upstream: true,
-
     instrumentation_for: "ThemePickerShell",
 
     strict_mode: false,