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/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/picker/color/ui/viewmodel/ColorPickerViewModel.kt b/src/com/android/customization/picker/color/ui/viewmodel/ColorPickerViewModel.kt
index 803663d..58bc858 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. */
@@ -91,15 +92,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],
@@ -141,15 +146,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,
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()
     }
 }