diff --git a/res/layout/color_option_2.xml b/res/layout/color_option_2.xml
new file mode 100644
index 0000000..2ac0fe6
--- /dev/null
+++ b/res/layout/color_option_2.xml
@@ -0,0 +1,90 @@
+<?xml version="1.0" encoding="utf-8"?><!--
+     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.
+-->
+<!-- Content description is set programmatically on the parent FrameLayout -->
+<LinearLayout
+    xmlns:android="http://schemas.android.com/apk/res/android"
+    android:layout_width="wrap_content"
+    android:layout_height="wrap_content"
+    android:orientation="vertical"
+    android:clipChildren="false">
+    <FrameLayout
+        android:id="@+id/icon_container"
+        android:layout_width="@dimen/option_tile_width"
+        android:layout_height="@dimen/option_tile_width"
+        android:clipChildren="false">
+
+        <ImageView
+            android:id="@id/selection_border"
+            android:layout_width="match_parent"
+            android:layout_height="match_parent"
+            android:background="@drawable/option_item_border"
+            android:alpha="0"
+            android:importantForAccessibility="no" />
+
+        <ImageView
+            android:id="@id/background"
+            android:layout_width="match_parent"
+            android:layout_height="match_parent"
+            android:background="@drawable/option_item_background"
+            android:importantForAccessibility="no" />
+
+        <FrameLayout
+            android:id="@id/foreground"
+            android:layout_width="match_parent"
+            android:layout_height="match_parent">
+            <ImageView
+                android:id="@+id/color_preview_0"
+                android:layout_width="wrap_content"
+                android:layout_height="wrap_content"
+                android:layout_gravity="center"
+                android:layout_marginRight="@dimen/color_seed_chip_margin"
+                android:layout_marginBottom="@dimen/color_seed_chip_margin"
+                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_gravity="center"
+                android:layout_marginLeft="@dimen/color_seed_chip_margin"
+                android:layout_marginBottom="@dimen/color_seed_chip_margin"
+                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_gravity="center"
+                android:layout_marginRight="@dimen/color_seed_chip_margin"
+                android:layout_marginTop="@dimen/color_seed_chip_margin"
+                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_gravity="center"
+                android:layout_marginLeft="@dimen/color_seed_chip_margin"
+                android:layout_marginTop="@dimen/color_seed_chip_margin"
+                android:src="@drawable/color_chip_seed_filled3"
+                android:importantForAccessibility="no" />
+        </FrameLayout>
+    </FrameLayout>
+</LinearLayout>
+
diff --git a/res/layout/color_option_with_background.xml b/res/layout/color_option_with_background.xml
index 67079f7..dc6f3ed 100644
--- a/res/layout/color_option_with_background.xml
+++ b/res/layout/color_option_with_background.xml
@@ -14,6 +14,7 @@
      limitations under the License.
 -->
 <!-- Content description is set programmatically on the parent FrameLayout -->
+<!-- TODO (b/272109171): Remove after clock settings is refactored to use OptionItemAdapter -->
 <LinearLayout
     xmlns:android="http://schemas.android.com/apk/res/android"
     android:layout_width="wrap_content"
diff --git a/res/layout/fragment_color_picker.xml b/res/layout/fragment_color_picker.xml
index a9d2adc..efb5a87 100644
--- a/res/layout/fragment_color_picker.xml
+++ b/res/layout/fragment_color_picker.xml
@@ -60,7 +60,8 @@
         android:layout_marginHorizontal="24dp"
         android:layout_marginBottom="28dp"
         android:background="@drawable/picker_fragment_background"
-        android:paddingTop="22dp">
+        android:paddingTop="22dp"
+        android:clipChildren="false">
 
         <FrameLayout
             android:layout_width="match_parent"
@@ -97,14 +98,16 @@
 
         <FrameLayout
             android:layout_width="match_parent"
-            android:layout_height="wrap_content">
+            android:layout_height="wrap_content"
+            android:clipChildren="false">
 
             <androidx.recyclerview.widget.RecyclerView
                 android:id="@+id/color_options"
                 android:layout_width="match_parent"
                 android:layout_height="wrap_content"
                 android:clipToPadding="false"
-                android:paddingHorizontal="16dp" />
+                android:paddingHorizontal="16dp"
+                android:clipChildren="false" />
 
             <!--
             This is just an invisible placeholder put in place so that the parent keeps its height
@@ -116,7 +119,7 @@
             It's critical for any TextViews inside the included layout to have text.
             -->
             <include
-                layout="@layout/color_option_with_background"
+                layout="@layout/color_option_2"
                 android:layout_width="wrap_content"
                 android:layout_height="wrap_content"
                 android:visibility="invisible" />
diff --git a/src/com/android/customization/model/color/ColorOption.java b/src/com/android/customization/model/color/ColorOption.java
index 66a3a3c..216bb9b 100644
--- a/src/com/android/customization/model/color/ColorOption.java
+++ b/src/com/android/customization/model/color/ColorOption.java
@@ -110,9 +110,12 @@
         if (mStyle != other.getStyle()) {
             return false;
         }
-        if (mIsDefault) {
-            return other.isDefault() || TextUtils.isEmpty(other.getSerializedPackages())
-                    || EMPTY_JSON.equals(other.getSerializedPackages());
+        String thisSerializedPackages = getSerializedPackages();
+        if (mIsDefault || TextUtils.isEmpty(thisSerializedPackages)
+                || EMPTY_JSON.equals(thisSerializedPackages)) {
+            String otherSerializedPackages = other.getSerializedPackages();
+            return other.isDefault() || TextUtils.isEmpty(otherSerializedPackages)
+                    || EMPTY_JSON.equals(otherSerializedPackages);
         }
         // Map#equals ensures keys and values are compared.
         return mPackagesByCategory.equals(other.mPackagesByCategory);
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 d24467a..78536ca 100644
--- a/src/com/android/customization/model/grid/ui/binder/GridScreenBinder.kt
+++ b/src/com/android/customization/model/grid/ui/binder/GridScreenBinder.kt
@@ -18,12 +18,14 @@
 package com.android.customization.model.grid.ui.binder
 
 import android.view.View
+import android.widget.ImageView
 import androidx.lifecycle.Lifecycle
 import androidx.lifecycle.LifecycleOwner
 import androidx.lifecycle.lifecycleScope
 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
@@ -57,7 +59,11 @@
                     OptionItemBinder.TintSpec(
                         selectedColor = view.context.getColor(R.color.text_color_primary),
                         unselectedColor = view.context.getColor(R.color.text_color_secondary),
-                    )
+                    ),
+                bindIcon = { foregroundView: View, gridIcon: GridIconViewModel ->
+                    val imageView = foregroundView as? ImageView
+                    imageView?.let { GridIconViewBinder.bind(imageView, gridIcon) }
+                }
             )
         optionView.adapter = adapter
 
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 af6ed0f..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>> =
+    val optionItems: Flow<List<OptionItemViewModel<GridIconViewModel>>> =
         interactor.options.map { model -> toViewModel(model) }
 
     private fun toViewModel(
         model: GridOptionItemsModel,
-    ): List<OptionItemViewModel> {
+    ): 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(
-                        key = flowOf("${option.cols}x${option.rows}"),
-                        icon =
-                            Icon.Loaded(
-                                drawable =
-                                    GridTileDrawable(
-                                        option.cols,
-                                        option.rows,
-                                        iconShapePath,
-                                    ),
-                                contentDescription = text
+                    OptionItemViewModel<GridIconViewModel>(
+                        key =
+                            MutableStateFlow("${option.cols}x${option.rows}") as StateFlow<String>,
+                        payload =
+                            GridIconViewModel(
+                                columns = option.cols,
+                                rows = option.rows,
+                                path = iconShapePath,
                             ),
                         text = text,
                         isSelected = option.isSelected,
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 7553dcb..d94acfc 100644
--- a/src/com/android/customization/picker/color/data/repository/ColorPickerRepositoryImpl.kt
+++ b/src/com/android/customization/picker/color/data/repository/ColorPickerRepositoryImpl.kt
@@ -120,8 +120,10 @@
         for (overlay in overlays) {
             colorOptionBuilder.addOverlayPackage(overlay.key, overlay.value)
         }
+        val colorOption = colorOptionBuilder.build()
         return ColorOptionModel(
-            colorOption = colorOptionBuilder.build(),
+            key = "${colorOption.style}::${colorOption.serializedPackages}",
+            colorOption = colorOption,
             isSelected = false,
         )
     }
@@ -132,6 +134,7 @@
 
     private fun ColorOption.toModel(): ColorOptionModel {
         return ColorOptionModel(
+            key = "${this.style}::${this.serializedPackages}",
             colorOption = this,
             isSelected = isActive(colorManager),
         )
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 660a74f..b7d3a57 100644
--- a/src/com/android/customization/picker/color/data/repository/FakeColorPickerRepository.kt
+++ b/src/com/android/customization/picker/color/data/repository/FakeColorPickerRepository.kt
@@ -18,6 +18,7 @@
 
 import android.content.Context
 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
@@ -61,6 +62,7 @@
                                     selectedColorOptionIndex == index
                             val colorOption =
                                 ColorOptionModel(
+                                    key = "${ColorType.WALLPAPER_COLOR}::$index",
                                     colorOption = buildWallpaperOption(index),
                                     isSelected = isSelected,
                                 )
@@ -78,6 +80,7 @@
                                     selectedColorOptionIndex == index
                             val colorOption =
                                 ColorOptionModel(
+                                    key = "${ColorType.BASIC_COLOR}::$index",
                                     colorOption = buildPresetOption(index),
                                     isSelected =
                                         selectedColorOptionType == ColorType.BASIC_COLOR &&
@@ -131,6 +134,7 @@
             wallpaperColorOptions.forEach { option ->
                 add(
                     ColorOptionModel(
+                        key = option.key,
                         colorOption = option.colorOption,
                         isSelected = option.testEquals(colorOptionModel),
                     )
@@ -142,6 +146,7 @@
             basicColorOptions.forEach { option ->
                 add(
                     ColorOptionModel(
+                        key = option.key,
                         colorOption = option.colorOption,
                         isSelected = option.testEquals(colorOptionModel),
                     )
@@ -169,10 +174,7 @@
             return false
         }
         return if (other is ColorOptionModel) {
-            val thisColorOptionIsWallpaperColor = this.colorOption is ColorSeedOption
-            val otherColorOptionIsWallpaperColor = other.colorOption is ColorSeedOption
-            (thisColorOptionIsWallpaperColor == otherColorOptionIsWallpaperColor) &&
-                (this.colorOption.index == other.colorOption.index)
+            TextUtils.equals(this.key, other.key)
         } else {
             false
         }
diff --git a/src/com/android/customization/picker/color/domain/interactor/ColorPickerInteractor.kt b/src/com/android/customization/picker/color/domain/interactor/ColorPickerInteractor.kt
index a932067..8c7a4b7 100644
--- a/src/com/android/customization/picker/color/domain/interactor/ColorPickerInteractor.kt
+++ b/src/com/android/customization/picker/color/domain/interactor/ColorPickerInteractor.kt
@@ -16,12 +16,10 @@
  */
 package com.android.customization.picker.color.domain.interactor
 
-import androidx.annotation.VisibleForTesting
 import com.android.customization.picker.color.data.repository.ColorPickerRepository
 import com.android.customization.picker.color.shared.model.ColorOptionModel
 import javax.inject.Provider
 import kotlinx.coroutines.flow.MutableStateFlow
-import kotlinx.coroutines.flow.combine
 
 /** Single entry-point for all application state and business logic related to system color. */
 class ColorPickerInteractor(
@@ -32,31 +30,10 @@
      * The newly selected color option for overwriting the current active option during an
      * optimistic update, the value is set to null when update fails
      */
-    @VisibleForTesting private val activeColorOption = MutableStateFlow<ColorOptionModel?>(null)
+    val activeColorOption = MutableStateFlow<ColorOptionModel?>(null)
 
     /** List of wallpaper and preset color options on the device, categorized by Color Type */
-    val colorOptions =
-        combine(repository.colorOptions, activeColorOption) { colorOptions, activeOption ->
-            colorOptions
-                .map { colorTypeEntry ->
-                    colorTypeEntry.key to
-                        colorTypeEntry.value.map { colorOptionModel ->
-                            val isSelected =
-                                if (activeOption != null) {
-                                    colorOptionModel.colorOption.isEquivalent(
-                                        activeOption.colorOption
-                                    )
-                                } else {
-                                    colorOptionModel.isSelected
-                                }
-                            ColorOptionModel(
-                                colorOption = colorOptionModel.colorOption,
-                                isSelected = isSelected
-                            )
-                        }
-                }
-                .toMap()
-        }
+    val colorOptions = repository.colorOptions
 
     suspend fun select(colorOptionModel: ColorOptionModel) {
         activeColorOption.value = colorOptionModel
diff --git a/src/com/android/customization/picker/color/shared/model/ColorOptionModel.kt b/src/com/android/customization/picker/color/shared/model/ColorOptionModel.kt
index 69ef62a..5fde08e 100644
--- a/src/com/android/customization/picker/color/shared/model/ColorOptionModel.kt
+++ b/src/com/android/customization/picker/color/shared/model/ColorOptionModel.kt
@@ -21,6 +21,8 @@
 
 /** Models application state for a color option in a picker experience. */
 data class ColorOptionModel(
+    val key: String,
+
     /** Colors for the color option. */
     val colorOption: ColorOption,
 
diff --git a/src/com/android/customization/picker/color/ui/adapter/ColorOptionAdapter.kt b/src/com/android/customization/picker/color/ui/adapter/ColorOptionAdapter.kt
index 0e53766..7aa390d 100644
--- a/src/com/android/customization/picker/color/ui/adapter/ColorOptionAdapter.kt
+++ b/src/com/android/customization/picker/color/ui/adapter/ColorOptionAdapter.kt
@@ -32,7 +32,7 @@
 /**
  * Adapts between color option items and views.
  *
- * TODO (b/262924623): Refactor color picker with animated option framework ag/21132368
+ * TODO (b/272109171): Remove after clock settings is refactored to use OptionItemAdapter
  */
 class ColorOptionAdapter : RecyclerView.Adapter<ColorOptionAdapter.ViewHolder>() {
 
diff --git a/src/com/android/customization/picker/color/ui/binder/ColorOptionIconBinder.kt b/src/com/android/customization/picker/color/ui/binder/ColorOptionIconBinder.kt
new file mode 100644
index 0000000..1478cc4
--- /dev/null
+++ b/src/com/android/customization/picker/color/ui/binder/ColorOptionIconBinder.kt
@@ -0,0 +1,41 @@
+/*
+ * 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.picker.color.ui.binder
+
+import android.graphics.BlendMode
+import android.graphics.BlendModeColorFilter
+import android.view.ViewGroup
+import android.widget.ImageView
+import com.android.customization.picker.color.ui.viewmodel.ColorOptionIconViewModel
+import com.android.wallpaper.R
+
+object ColorOptionIconBinder {
+    fun bind(
+        view: ViewGroup,
+        viewModel: ColorOptionIconViewModel,
+    ) {
+        val color0View: ImageView = view.requireViewById(R.id.color_preview_0)
+        val color1View: ImageView = view.requireViewById(R.id.color_preview_1)
+        val color2View: ImageView = view.requireViewById(R.id.color_preview_2)
+        val color3View: ImageView = view.requireViewById(R.id.color_preview_3)
+        color0View.drawable.colorFilter = BlendModeColorFilter(viewModel.color0, BlendMode.SRC)
+        color1View.drawable.colorFilter = BlendModeColorFilter(viewModel.color1, BlendMode.SRC)
+        color2View.drawable.colorFilter = BlendModeColorFilter(viewModel.color2, BlendMode.SRC)
+        color3View.drawable.colorFilter = BlendModeColorFilter(viewModel.color3, BlendMode.SRC)
+    }
+}
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 887fef0..053d5dd 100644
--- a/src/com/android/customization/picker/color/ui/binder/ColorPickerBinder.kt
+++ b/src/com/android/customization/picker/color/ui/binder/ColorPickerBinder.kt
@@ -19,6 +19,7 @@
 
 import android.graphics.Rect
 import android.view.View
+import android.view.ViewGroup
 import androidx.core.view.ViewCompat
 import androidx.lifecycle.Lifecycle
 import androidx.lifecycle.LifecycleOwner
@@ -26,10 +27,11 @@
 import androidx.lifecycle.repeatOnLifecycle
 import androidx.recyclerview.widget.LinearLayoutManager
 import androidx.recyclerview.widget.RecyclerView
-import com.android.customization.picker.color.ui.adapter.ColorOptionAdapter
 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.wallpaper.R
+import com.android.wallpaper.picker.option.ui.adapter.OptionItemAdapter
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.launch
 
@@ -54,7 +56,15 @@
         colorTypeTabView.layoutManager =
             LinearLayoutManager(view.context, RecyclerView.HORIZONTAL, false)
         colorTypeTabView.addItemDecoration(ItemSpacing())
-        val colorOptionAdapter = ColorOptionAdapter()
+        val colorOptionAdapter =
+            OptionItemAdapter(
+                layoutResourceId = R.layout.color_option_2,
+                lifecycleOwner = lifecycleOwner,
+                bindIcon = { foregroundView: View, colorIcon: ColorOptionIconViewModel ->
+                    val viewGroup = foregroundView as? ViewGroup
+                    viewGroup?.let { ColorOptionIconBinder.bind(viewGroup, colorIcon) }
+                }
+            )
         colorOptionContainerView.adapter = colorOptionAdapter
         colorOptionContainerView.layoutManager =
             LinearLayoutManager(view.context, RecyclerView.HORIZONTAL, false)
diff --git a/src/com/android/customization/picker/color/ui/binder/ColorSectionViewBinder.kt b/src/com/android/customization/picker/color/ui/binder/ColorSectionViewBinder.kt
index 0842870..05b0916 100644
--- a/src/com/android/customization/picker/color/ui/binder/ColorSectionViewBinder.kt
+++ b/src/com/android/customization/picker/color/ui/binder/ColorSectionViewBinder.kt
@@ -17,10 +17,9 @@
 
 package com.android.customization.picker.color.ui.binder
 
-import android.graphics.BlendMode
-import android.graphics.BlendModeColorFilter
 import android.view.LayoutInflater
 import android.view.View
+import android.view.ViewGroup
 import android.widget.ImageView
 import android.widget.LinearLayout
 import androidx.core.view.isVisible
@@ -28,9 +27,10 @@
 import androidx.lifecycle.LifecycleOwner
 import androidx.lifecycle.lifecycleScope
 import androidx.lifecycle.repeatOnLifecycle
-import com.android.customization.picker.color.ui.viewmodel.ColorOptionViewModel
+import com.android.customization.picker.color.ui.viewmodel.ColorOptionIconViewModel
 import com.android.customization.picker.color.ui.viewmodel.ColorPickerViewModel
 import com.android.wallpaper.R
+import com.android.wallpaper.picker.option.ui.viewmodel.OptionItemViewModel
 import kotlinx.coroutines.launch
 
 object ColorSectionViewBinder {
@@ -63,8 +63,9 @@
                         setOptions(
                             options = colorOptions,
                             view = optionContainer,
+                            lifecycleOwner = lifecycleOwner,
                             addOverflowOption = !isConnectedHorizontallyToOtherSections,
-                            overflowOnClick = navigationOnClick
+                            overflowOnClick = navigationOnClick,
                         )
                     }
                 }
@@ -73,10 +74,11 @@
     }
 
     fun setOptions(
-        options: List<ColorOptionViewModel>,
+        options: List<OptionItemViewModel<ColorOptionIconViewModel>>,
         view: LinearLayout,
+        lifecycleOwner: LifecycleOwner,
         addOverflowOption: Boolean = false,
-        overflowOnClick: (View) -> Unit = {}
+        overflowOnClick: (View) -> Unit = {},
     ) {
         view.removeAllViews()
         // Color option slot size is the minimum between the color option size and the view column
@@ -92,26 +94,29 @@
             val itemView =
                 LayoutInflater.from(view.context)
                     .inflate(R.layout.color_option_no_background, view, false)
-
-            val color0View: ImageView = itemView.requireViewById(R.id.color_preview_0)
-            val color1View: ImageView = itemView.requireViewById(R.id.color_preview_1)
-            val color2View: ImageView = itemView.requireViewById(R.id.color_preview_2)
-            val color3View: ImageView = itemView.requireViewById(R.id.color_preview_3)
-            color0View.drawable.colorFilter = BlendModeColorFilter(item.color0, BlendMode.SRC)
-            color1View.drawable.colorFilter = BlendModeColorFilter(item.color1, BlendMode.SRC)
-            color2View.drawable.colorFilter = BlendModeColorFilter(item.color2, BlendMode.SRC)
-            color3View.drawable.colorFilter = BlendModeColorFilter(item.color3, BlendMode.SRC)
-
+            item.payload?.let { ColorOptionIconBinder.bind(itemView as ViewGroup, item.payload) }
             val optionSelectedView = itemView.findViewById<ImageView>(R.id.option_selected)
-            optionSelectedView.isVisible = item.isSelected
 
-            itemView.setOnClickListener(
-                if (item.onClick != null) {
-                    View.OnClickListener { item.onClick.invoke() }
-                } else {
-                    null
+            lifecycleOwner.lifecycleScope.launch {
+                lifecycleOwner.repeatOnLifecycle(Lifecycle.State.STARTED) {
+                    launch {
+                        item.isSelected.collect { isSelected ->
+                            optionSelectedView.isVisible = isSelected
+                        }
+                    }
+                    launch {
+                        item.onClicked.collect { onClicked ->
+                            itemView.setOnClickListener(
+                                if (onClicked != null) {
+                                    View.OnClickListener { onClicked.invoke() }
+                                } else {
+                                    null
+                                }
+                            )
+                        }
+                    }
                 }
-            )
+            }
             view.addView(itemView)
         }
         // add overflow option
diff --git a/src/com/android/customization/picker/color/ui/viewmodel/ColorOptionIconViewModel.kt b/src/com/android/customization/picker/color/ui/viewmodel/ColorOptionIconViewModel.kt
new file mode 100644
index 0000000..d32538d
--- /dev/null
+++ b/src/com/android/customization/picker/color/ui/viewmodel/ColorOptionIconViewModel.kt
@@ -0,0 +1,27 @@
+/*
+ * 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.picker.color.ui.viewmodel
+
+import android.annotation.ColorInt
+
+data class ColorOptionIconViewModel(
+    @ColorInt val color0: Int,
+    @ColorInt val color1: Int,
+    @ColorInt val color2: Int,
+    @ColorInt val color3: Int,
+)
diff --git a/src/com/android/customization/picker/color/ui/viewmodel/ColorOptionViewModel.kt b/src/com/android/customization/picker/color/ui/viewmodel/ColorOptionViewModel.kt
index 784ec2e..7af2aa5 100644
--- a/src/com/android/customization/picker/color/ui/viewmodel/ColorOptionViewModel.kt
+++ b/src/com/android/customization/picker/color/ui/viewmodel/ColorOptionViewModel.kt
@@ -19,7 +19,11 @@
 
 import android.annotation.ColorInt
 
-/** Models UI state for a color options in a picker experience. */
+/**
+ * Models UI state for a color options in a picker experience.
+ *
+ * TODO (b/272109171): Remove after clock settings is refactored to use OptionItemAdapter
+ */
 data class ColorOptionViewModel(
     /** Colors for the color option. */
     @ColorInt val color0: Int,
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 5784855..58bc858 100644
--- a/src/com/android/customization/picker/color/ui/viewmodel/ColorPickerViewModel.kt
+++ b/src/com/android/customization/picker/color/ui/viewmodel/ColorPickerViewModel.kt
@@ -25,12 +25,16 @@
 import com.android.customization.picker.color.domain.interactor.ColorPickerInteractor
 import com.android.customization.picker.color.shared.model.ColorType
 import com.android.wallpaper.R
+import com.android.wallpaper.picker.common.text.ui.viewmodel.Text
+import com.android.wallpaper.picker.option.ui.viewmodel.OptionItemViewModel
 import kotlin.math.max
 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.map
+import kotlinx.coroutines.flow.stateIn
 import kotlinx.coroutines.launch
 
 /** Models UI state for a color picker experience. */
@@ -75,7 +79,8 @@
         }
 
     /** The list of all color options mapped by their color type */
-    private val allColorOptions: Flow<Map<ColorType, List<ColorOptionViewModel>>> =
+    private val allColorOptions:
+        Flow<Map<ColorType, List<OptionItemViewModel<ColorOptionIconViewModel>>>> =
         interactor.colorOptions.map { colorOptions ->
             colorOptions
                 .map { colorOptionEntry ->
@@ -87,26 +92,45 @@
                                         colorOptionModel.colorOption as ColorSeedOption
                                     val colors =
                                         colorSeedOption.previewInfo.resolveColors(context.resources)
-                                    ColorOptionViewModel(
-                                        color0 = colors[0],
-                                        color1 = colors[1],
-                                        color2 = colors[2],
-                                        color3 = colors[3],
-                                        contentDescription =
-                                            colorSeedOption
-                                                .getContentDescription(context)
-                                                .toString(),
-                                        isSelected = colorOptionModel.isSelected,
-                                        onClick =
-                                            if (colorOptionModel.isSelected) {
-                                                null
-                                            } else {
-                                                {
-                                                    viewModelScope.launch {
-                                                        interactor.select(colorOptionModel)
+                                    val isSelectedFlow: StateFlow<Boolean> =
+                                        interactor.activeColorOption
+                                            .map {
+                                                it?.colorOption?.isEquivalent(
+                                                    colorOptionModel.colorOption
+                                                )
+                                                    ?: colorOptionModel.isSelected
+                                            }
+                                            .stateIn(viewModelScope)
+                                    OptionItemViewModel<ColorOptionIconViewModel>(
+                                        key =
+                                            MutableStateFlow(colorOptionModel.key)
+                                                as StateFlow<String>,
+                                        payload =
+                                            ColorOptionIconViewModel(
+                                                colors[0],
+                                                colors[1],
+                                                colors[2],
+                                                colors[3]
+                                            ),
+                                        text =
+                                            Text.Loaded(
+                                                colorSeedOption
+                                                    .getContentDescription(context)
+                                                    .toString()
+                                            ),
+                                        isSelected = isSelectedFlow,
+                                        onClicked =
+                                            isSelectedFlow.map { isSelected ->
+                                                if (isSelected) {
+                                                    null
+                                                } else {
+                                                    {
+                                                        viewModelScope.launch {
+                                                            interactor.select(colorOptionModel)
+                                                        }
                                                     }
                                                 }
-                                            }
+                                            },
                                     )
                                 }
                             }
@@ -122,21 +146,42 @@
                                         colorBundle.previewInfo.resolveSecondaryColor(
                                             context.resources
                                         )
-                                    ColorOptionViewModel(
-                                        color0 = primaryColor,
-                                        color1 = secondaryColor,
-                                        color2 = primaryColor,
-                                        color3 = secondaryColor,
-                                        contentDescription =
-                                            colorBundle.getContentDescription(context).toString(),
-                                        isSelected = colorOptionModel.isSelected,
-                                        onClick =
-                                            if (colorOptionModel.isSelected) {
-                                                null
-                                            } else {
-                                                {
-                                                    viewModelScope.launch {
-                                                        interactor.select(colorOptionModel)
+                                    val isSelectedFlow: StateFlow<Boolean> =
+                                        interactor.activeColorOption
+                                            .map {
+                                                it?.colorOption?.isEquivalent(
+                                                    colorOptionModel.colorOption
+                                                )
+                                                    ?: colorOptionModel.isSelected
+                                            }
+                                            .stateIn(viewModelScope)
+                                    OptionItemViewModel<ColorOptionIconViewModel>(
+                                        key =
+                                            MutableStateFlow(colorOptionModel.key)
+                                                as StateFlow<String>,
+                                        payload =
+                                            ColorOptionIconViewModel(
+                                                primaryColor,
+                                                secondaryColor,
+                                                primaryColor,
+                                                secondaryColor
+                                            ),
+                                        text =
+                                            Text.Loaded(
+                                                colorBundle
+                                                    .getContentDescription(context)
+                                                    .toString()
+                                            ),
+                                        isSelected = isSelectedFlow,
+                                        onClicked =
+                                            isSelectedFlow.map { isSelected ->
+                                                if (isSelected) {
+                                                    null
+                                                } else {
+                                                    {
+                                                        viewModelScope.launch {
+                                                            interactor.select(colorOptionModel)
+                                                        }
                                                     }
                                                 }
                                             },
@@ -149,7 +194,7 @@
         }
 
     /** The list of all available color options for the selected Color Type. */
-    val colorOptions: Flow<List<ColorOptionViewModel>> =
+    val colorOptions: Flow<List<OptionItemViewModel<ColorOptionIconViewModel>>> =
         combine(allColorOptions, selectedColorTypeId) { allColorOptions, selectedColorTypeIdOrNull
             ->
             val selectedColorTypeId = selectedColorTypeIdOrNull ?: ColorType.WALLPAPER_COLOR
@@ -157,7 +202,7 @@
         }
 
     /** The list of color options for the color section */
-    val colorSectionOptions: Flow<List<ColorOptionViewModel>> =
+    val colorSectionOptions: Flow<List<OptionItemViewModel<ColorOptionIconViewModel>>> =
         allColorOptions.map { allColorOptions ->
             val wallpaperOptions = allColorOptions[ColorType.WALLPAPER_COLOR]
             val presetOptions = allColorOptions[ColorType.BASIC_COLOR]
diff --git a/src/com/android/customization/picker/quickaffordance/ui/binder/KeyguardQuickAffordancePickerBinder.kt b/src/com/android/customization/picker/quickaffordance/ui/binder/KeyguardQuickAffordancePickerBinder.kt
index af5cd13..4395f5e 100644
--- a/src/com/android/customization/picker/quickaffordance/ui/binder/KeyguardQuickAffordancePickerBinder.kt
+++ b/src/com/android/customization/picker/quickaffordance/ui/binder/KeyguardQuickAffordancePickerBinder.kt
@@ -20,6 +20,7 @@
 import android.app.Dialog
 import android.content.Context
 import android.view.View
+import android.widget.ImageView
 import androidx.lifecycle.Lifecycle
 import androidx.lifecycle.LifecycleOwner
 import androidx.lifecycle.lifecycleScope
@@ -32,6 +33,8 @@
 import com.android.wallpaper.R
 import com.android.wallpaper.picker.common.dialog.ui.viewbinder.DialogViewBinder
 import com.android.wallpaper.picker.common.dialog.ui.viewmodel.DialogViewModel
+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 kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.flow.combine
@@ -62,6 +65,10 @@
             OptionItemAdapter(
                 layoutResourceId = R.layout.keyguard_quick_affordance,
                 lifecycleOwner = lifecycleOwner,
+                bindIcon = { foregroundView: View, gridIcon: Icon ->
+                    val imageView = foregroundView as? ImageView
+                    imageView?.let { IconViewBinder.bind(imageView, gridIcon) }
+                }
             )
         affordancesView.adapter = affordancesAdapter
         affordancesView.layoutManager =
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 d88edfa..14b6acc 100644
--- a/src/com/android/customization/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModel.kt
+++ b/src/com/android/customization/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModel.kt
@@ -147,16 +147,18 @@
                         isSelected = isSelected,
                         selectedQuickAffordances =
                             selectedAffordances.map { affordanceModel ->
-                                OptionItemViewModel(
-                                    key = flowOf("${slot.id}::${affordanceModel.id}"),
-                                    icon =
+                                OptionItemViewModel<Icon>(
+                                    key =
+                                        MutableStateFlow("${slot.id}::${affordanceModel.id}")
+                                            as StateFlow<String>,
+                                    payload =
                                         Icon.Loaded(
                                             drawable =
                                                 getAffordanceIcon(affordanceModel.iconResourceId),
                                             contentDescription = null,
                                         ),
                                     text = Text.Loaded(affordanceModel.name),
-                                    isSelected = flowOf(true),
+                                    isSelected = MutableStateFlow(true) as StateFlow<Boolean>,
                                     onClicked = flowOf(null),
                                     onLongClicked = null,
                                     isEnabled = true,
@@ -194,9 +196,9 @@
             )
 
     /** The list of all available quick affordances for the selected slot. */
-    val quickAffordances: Flow<List<OptionItemViewModel>> =
+    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,11 +222,16 @@
             ) +
                 affordances.map { affordance ->
                     val affordanceIcon = getAffordanceIcon(affordance.iconResourceId)
-                    val isSelectedFlow: Flow<Boolean> =
-                        selectedAffordanceIds.map { it.contains(affordance.id) }
-                    OptionItemViewModel(
-                        key = selectedSlotId.map { slotId -> "$slotId::${affordance.id}" },
-                        icon = Icon.Loaded(drawable = affordanceIcon, contentDescription = null),
+                    val isSelectedFlow: StateFlow<Boolean> =
+                        selectedAffordanceIds
+                            .map { it.contains(affordance.id) }
+                            .stateIn(viewModelScope)
+                    OptionItemViewModel<Icon>(
+                        key =
+                            selectedSlotId
+                                .map { slotId -> "$slotId::${affordance.id}" }
+                                .stateIn(viewModelScope),
+                        payload = Icon.Loaded(drawable = affordanceIcon, contentDescription = null),
                         text = Text.Loaded(affordance.name),
                         isSelected = isSelectedFlow,
                         onClicked =
@@ -273,15 +280,15 @@
     val summary: Flow<KeyguardQuickAffordanceSummaryViewModel> =
         slots.map { slots ->
             val icon2 =
-                slots[KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END]
-                    ?.selectedQuickAffordances
-                    ?.firstOrNull()
-                    ?.icon
+                (slots[KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END]
+                        ?.selectedQuickAffordances
+                        ?.firstOrNull())
+                    ?.payload
             val icon1 =
-                slots[KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START]
-                    ?.selectedQuickAffordances
-                    ?.firstOrNull()
-                    ?.icon
+                (slots[KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START]
+                        ?.selectedQuickAffordances
+                        ?.firstOrNull())
+                    ?.payload
 
             KeyguardQuickAffordanceSummaryViewModel(
                 description = toDescriptionText(context, slots),
@@ -359,14 +366,14 @@
 
     /** 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 {
-        return OptionItemViewModel(
-            key = slotId.map { "$it::none" },
-            icon = Icon.Resource(res = R.drawable.link_off, contentDescription = null),
+    ): OptionItemViewModel<Icon> {
+        return OptionItemViewModel<Icon>(
+            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,
             onClicked = onSelected,
diff --git a/src/com/android/customization/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordanceSlotViewModel.kt b/src/com/android/customization/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordanceSlotViewModel.kt
index 6d8195a..4d11346 100644
--- a/src/com/android/customization/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordanceSlotViewModel.kt
+++ b/src/com/android/customization/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordanceSlotViewModel.kt
@@ -17,6 +17,7 @@
 
 package com.android.customization.picker.quickaffordance.ui.viewmodel
 
+import com.android.wallpaper.picker.common.icon.ui.viewmodel.Icon
 import com.android.wallpaper.picker.option.ui.viewmodel.OptionItemViewModel
 
 /** Models UI state for a single lock screen quick affordance slot in a picker experience. */
@@ -32,7 +33,7 @@
      *
      * Useful for preview.
      */
-    val selectedQuickAffordances: List<OptionItemViewModel>,
+    val selectedQuickAffordances: List<OptionItemViewModel<Icon>>,
 
     /**
      * The maximum number of quick affordances that can be selected for this slot.
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 7d0a527..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
@@ -94,13 +94,17 @@
             assertThat(getOnClick(optionItemsValue[1])).isNull()
         }
 
-    private fun TestScope.getSelectedIndex(optionItems: List<OptionItemViewModel>): 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): (() -> 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/ui/viewmodel/ColorPickerViewModelTest.kt b/tests/src/com/android/customization/model/picker/color/ui/viewmodel/ColorPickerViewModelTest.kt
index 7d87a55..533d1dc 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
@@ -23,9 +23,10 @@
 import com.android.customization.picker.color.domain.interactor.ColorPickerInteractor
 import com.android.customization.picker.color.domain.interactor.ColorPickerSnapshotRestorer
 import com.android.customization.picker.color.shared.model.ColorType
-import com.android.customization.picker.color.ui.viewmodel.ColorOptionViewModel
+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.wallpaper.picker.option.ui.viewmodel.OptionItemViewModel
 import com.android.wallpaper.testing.FakeSnapshotStore
 import com.android.wallpaper.testing.collectLastValue
 import com.google.common.truth.Truth.assertThat
@@ -97,13 +98,13 @@
                 selectedColorOptionIndex = 0
             )
 
-            colorSectionOptions()?.get(2)?.onClick?.invoke()
+            selectColorOption(colorSectionOptions, 2)
             assertColorOptionUiState(
                 colorOptions = colorSectionOptions(),
                 selectedColorOptionIndex = 2
             )
 
-            colorSectionOptions()?.get(4)?.onClick?.invoke()
+            selectColorOption(colorSectionOptions, 4)
             assertColorOptionUiState(
                 colorOptions = colorSectionOptions(),
                 selectedColorOptionIndex = 4
@@ -134,7 +135,7 @@
             )
 
             // Select a color option
-            colorOptions()?.get(2)?.onClick?.invoke()
+            selectColorOption(colorOptions, 2)
 
             // Check original option is no longer selected
             colorTypes()?.get(ColorType.WALLPAPER_COLOR)?.onClick?.invoke()
@@ -155,6 +156,20 @@
             )
         }
 
+    /** Simulates a user selecting the affordance at the given index, if that is clickable. */
+    private fun TestScope.selectColorOption(
+        colorOptions: () -> List<OptionItemViewModel<ColorOptionIconViewModel>>?,
+        index: Int,
+    ) {
+        val onClickedFlow = colorOptions()?.get(index)?.onClicked
+        val onClickedLastValueOrNull: (() -> (() -> Unit)?)? =
+            onClickedFlow?.let { collectLastValue(it) }
+        onClickedLastValueOrNull?.let { onClickedLastValue ->
+            val onClickedOrNull: (() -> Unit)? = onClickedLastValue()
+            onClickedOrNull?.let { onClicked -> onClicked() }
+        }
+    }
+
     /**
      * Asserts the entire picker UI state is what is expected. This includes the color type tabs and
      * the color options list.
@@ -165,9 +180,9 @@
      * @param selectedColorOptionIndex The index of the color option that's expected to be selected,
      *   -1 stands for no color option should be selected
      */
-    private fun assertPickerUiState(
+    private fun TestScope.assertPickerUiState(
         colorTypes: Map<ColorType, ColorTypeViewModel>?,
-        colorOptions: List<ColorOptionViewModel>?,
+        colorOptions: List<OptionItemViewModel<ColorOptionIconViewModel>>?,
         selectedColorTypeText: String,
         selectedColorOptionIndex: Int,
     ) {
@@ -191,8 +206,8 @@
      * @param selectedColorOptionIndex The index of the color option that's expected to be selected,
      *   -1 stands for no color option should be selected
      */
-    private fun assertColorOptionUiState(
-        colorOptions: List<ColorOptionViewModel>?,
+    private fun TestScope.assertColorOptionUiState(
+        colorOptions: List<OptionItemViewModel<ColorOptionIconViewModel>>?,
         selectedColorOptionIndex: Int,
     ) {
         var foundSelectedColorOption = false
@@ -200,12 +215,13 @@
         if (colorOptions != null) {
             for (i in colorOptions.indices) {
                 val colorOptionHasSelectedIndex = i == selectedColorOptionIndex
+                val isSelected: Boolean? = collectLastValue(colorOptions[i].isSelected).invoke()
                 assertWithMessage(
                         "Expected color option with index \"${i}\" to have" +
                             " isSelected=$colorOptionHasSelectedIndex but it was" +
-                            " ${colorOptions[i].isSelected}, num options: ${colorOptions.size}"
+                            " ${isSelected}, num options: ${colorOptions.size}"
                     )
-                    .that(colorOptions[i].isSelected)
+                    .that(isSelected)
                     .isEqualTo(colorOptionHasSelectedIndex)
                 foundSelectedColorOption = foundSelectedColorOption || colorOptionHasSelectedIndex
             }
diff --git a/tests/src/com/android/customization/model/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModelTest.kt b/tests/src/com/android/customization/model/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModelTest.kt
index 5c99585..103ae84 100644
--- a/tests/src/com/android/customization/model/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModelTest.kt
+++ b/tests/src/com/android/customization/model/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModelTest.kt
@@ -382,7 +382,7 @@
 
     /** Simulates a user selecting the affordance at the given index, if that is clickable. */
     private fun TestScope.selectAffordance(
-        affordances: () -> List<OptionItemViewModel>?,
+        affordances: () -> List<OptionItemViewModel<Icon>>?,
         index: Int,
     ) {
         val onClickedFlow = affordances()?.get(index)?.onClicked
@@ -405,7 +405,7 @@
      */
     private fun TestScope.assertPickerUiState(
         slots: Map<String, KeyguardQuickAffordanceSlotViewModel>?,
-        affordances: List<OptionItemViewModel>?,
+        affordances: List<OptionItemViewModel<Icon>>?,
         selectedSlotText: String,
         selectedAffordanceText: String,
     ) {