diff --git a/src/com/android/wallpaper/customization/ui/viewmodel/ClockPickerViewModel.kt b/src/com/android/wallpaper/customization/ui/viewmodel/ClockPickerViewModel.kt
index da3e65c..c0f0ba7 100644
--- a/src/com/android/wallpaper/customization/ui/viewmodel/ClockPickerViewModel.kt
+++ b/src/com/android/wallpaper/customization/ui/viewmodel/ClockPickerViewModel.kt
@@ -36,6 +36,7 @@
 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
 import dagger.assisted.AssistedInject
@@ -127,7 +128,7 @@
     data class ClockStyleModel(val thumbnail: Drawable, val showEditButton: StateFlow<Boolean>)
 
     @OptIn(ExperimentalCoroutinesApi::class)
-    val clockStyleOptions: StateFlow<List<OptionItemViewModel<ClockStyleModel>>> =
+    val clockStyleOptions: StateFlow<List<OptionItemViewModel2<ClockStyleModel>>> =
         clockPickerInteractor.allClocks
             .mapLatest { allClocks ->
                 // Delay to avoid the case that the full list of clocks is not initiated.
@@ -146,13 +147,13 @@
 
     private suspend fun ClockMetadataModel.toOption(
         resources: Resources
-    ): OptionItemViewModel<ClockStyleModel> {
+    ): OptionItemViewModel2<ClockStyleModel> {
         val isSelectedFlow = previewingClock.map { it.clockId == clockId }.stateIn(viewModelScope)
         val isEditable = fontAxes.isNotEmpty()
         val showEditButton = isSelectedFlow.map { it && isEditable }.stateIn(viewModelScope)
         val contentDescription =
             resources.getString(R.string.select_clock_action_description, description)
-        return OptionItemViewModel<ClockStyleModel>(
+        return OptionItemViewModel2<ClockStyleModel>(
             key = MutableStateFlow(clockId) as StateFlow<String>,
             payload = ClockStyleModel(thumbnail = thumbnail, showEditButton = showEditButton),
             text = Text.Loaded(contentDescription),
diff --git a/src/com/android/wallpaper/customization/ui/viewmodel/ColorPickerViewModel2.kt b/src/com/android/wallpaper/customization/ui/viewmodel/ColorPickerViewModel2.kt
index 26e7867..e7efebd 100644
--- a/src/com/android/wallpaper/customization/ui/viewmodel/ColorPickerViewModel2.kt
+++ b/src/com/android/wallpaper/customization/ui/viewmodel/ColorPickerViewModel2.kt
@@ -27,7 +27,7 @@
 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.customization.ui.viewmodel.FloatingToolbarTabViewModel
-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
 import dagger.assisted.AssistedInject
@@ -109,7 +109,7 @@
 
     /** The list of all color options mapped by their color type */
     private val allColorOptions:
-        Flow<Map<ColorType, List<OptionItemViewModel<ColorOptionIconViewModel>>>> =
+        Flow<Map<ColorType, List<OptionItemViewModel2<ColorOptionIconViewModel>>>> =
         interactor.colorOptions.map { colorOptions ->
             colorOptions
                 .map { colorOptionEntry ->
@@ -128,7 +128,7 @@
                                             ?: colorOptionModel.isSelected
                                     }
                                     .stateIn(viewModelScope)
-                            OptionItemViewModel<ColorOptionIconViewModel>(
+                            OptionItemViewModel2<ColorOptionIconViewModel>(
                                 key = MutableStateFlow(colorOptionModel.key) as StateFlow<String>,
                                 payload =
                                     ColorOptionIconViewModel(
@@ -188,9 +188,9 @@
     }
 
     /** The list of all available color options for the selected Color Type. */
-    val colorOptions: Flow<List<OptionItemViewModel<ColorOptionIconViewModel>>> =
+    val colorOptions: Flow<List<OptionItemViewModel2<ColorOptionIconViewModel>>> =
         combine(allColorOptions, selectedColorTypeTabId) {
-            allColorOptions: Map<ColorType, List<OptionItemViewModel<ColorOptionIconViewModel>>>,
+            allColorOptions: Map<ColorType, List<OptionItemViewModel2<ColorOptionIconViewModel>>>,
             selectedColorTypeIdOrNull ->
             val selectedColorTypeId = selectedColorTypeIdOrNull ?: ColorType.WALLPAPER_COLOR
             allColorOptions[selectedColorTypeId]!!
diff --git a/src/com/android/wallpaper/customization/ui/viewmodel/KeyguardQuickAffordancePickerViewModel2.kt b/src/com/android/wallpaper/customization/ui/viewmodel/KeyguardQuickAffordancePickerViewModel2.kt
index fd94b78..fd04580 100644
--- a/src/com/android/wallpaper/customization/ui/viewmodel/KeyguardQuickAffordancePickerViewModel2.kt
+++ b/src/com/android/wallpaper/customization/ui/viewmodel/KeyguardQuickAffordancePickerViewModel2.kt
@@ -36,6 +36,7 @@
 import com.android.wallpaper.picker.common.text.ui.viewmodel.Text
 import com.android.wallpaper.picker.customization.ui.viewmodel.FloatingToolbarTabViewModel
 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
 import dagger.assisted.AssistedInject
@@ -174,7 +175,7 @@
             .shareIn(scope = viewModelScope, started = SharingStarted.WhileSubscribed(), replay = 1)
 
     /** The list of all available quick affordances for the selected slot. */
-    val quickAffordances: Flow<List<OptionItemViewModel<Icon>>> =
+    val quickAffordances: Flow<List<OptionItemViewModel2<Icon>>> =
         quickAffordanceInteractor.affordances.map { affordances ->
             val isNoneSelected =
                 combine(selectedSlotId, previewingQuickAffordances, selectedAffordanceIds) {
@@ -219,7 +220,7 @@
                                 } ?: selectedAffordanceIds.contains(affordance.id)
                             }
                             .stateIn(viewModelScope)
-                    OptionItemViewModel<Icon>(
+                    OptionItemViewModel2<Icon>(
                         key =
                             selectedSlotId
                                 .map { slotId -> "$slotId::${affordance.id}" }
@@ -374,8 +375,8 @@
         slotId: StateFlow<String>,
         isSelected: StateFlow<Boolean>,
         onSelected: Flow<(() -> Unit)?>,
-    ): OptionItemViewModel<Icon> {
-        return OptionItemViewModel<Icon>(
+    ): OptionItemViewModel2<Icon> {
+        return OptionItemViewModel2<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),
diff --git a/tests/robotests/src/com/android/wallpaper/customization/ui/viewmodel/ColorPickerViewModel2Test.kt b/tests/robotests/src/com/android/wallpaper/customization/ui/viewmodel/ColorPickerViewModel2Test.kt
index a4d0ba0..421d8a0 100644
--- a/tests/robotests/src/com/android/wallpaper/customization/ui/viewmodel/ColorPickerViewModel2Test.kt
+++ b/tests/robotests/src/com/android/wallpaper/customization/ui/viewmodel/ColorPickerViewModel2Test.kt
@@ -29,7 +29,7 @@
 import com.android.customization.picker.color.ui.viewmodel.ColorOptionIconViewModel
 import com.android.systemui.monet.Style
 import com.android.wallpaper.picker.customization.ui.viewmodel.FloatingToolbarTabViewModel
-import com.android.wallpaper.picker.option.ui.viewmodel.OptionItemViewModel
+import com.android.wallpaper.picker.option.ui.viewmodel.OptionItemViewModel2
 import com.android.wallpaper.testing.FakeSnapshotStore
 import com.android.wallpaper.testing.collectLastValue
 import com.google.common.truth.Truth.assertThat
@@ -205,7 +205,7 @@
 
     /** Simulates a user selecting the affordance at the given index, if that is clickable. */
     private fun TestScope.selectColorOption(
-        colorOptions: () -> List<OptionItemViewModel<ColorOptionIconViewModel>>?,
+        colorOptions: () -> List<OptionItemViewModel2<ColorOptionIconViewModel>>?,
         index: Int,
     ) {
         val onClickedFlow = colorOptions()?.get(index)?.onClicked
@@ -235,7 +235,7 @@
      */
     private fun TestScope.assertPickerUiState(
         colorTypes: List<FloatingToolbarTabViewModel>?,
-        colorOptions: List<OptionItemViewModel<ColorOptionIconViewModel>>?,
+        colorOptions: List<OptionItemViewModel2<ColorOptionIconViewModel>>?,
         selectedColorTypeText: String,
         selectedColorOptionIndex: Int,
     ) {
@@ -260,7 +260,7 @@
      *   -1 stands for no color option should be selected
      */
     private fun TestScope.assertColorOptionUiState(
-        colorOptions: List<OptionItemViewModel<ColorOptionIconViewModel>>?,
+        colorOptions: List<OptionItemViewModel2<ColorOptionIconViewModel>>?,
         selectedColorOptionIndex: Int,
     ) {
         var foundSelectedColorOption = false
diff --git a/tests/robotests/src/com/android/wallpaper/customization/ui/viewmodel/KeyguardQuickAffordancePickerViewModel2Test.kt b/tests/robotests/src/com/android/wallpaper/customization/ui/viewmodel/KeyguardQuickAffordancePickerViewModel2Test.kt
index b6f249e..a7efc45 100644
--- a/tests/robotests/src/com/android/wallpaper/customization/ui/viewmodel/KeyguardQuickAffordancePickerViewModel2Test.kt
+++ b/tests/robotests/src/com/android/wallpaper/customization/ui/viewmodel/KeyguardQuickAffordancePickerViewModel2Test.kt
@@ -30,7 +30,7 @@
 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.customization.ui.viewmodel.FloatingToolbarTabViewModel
-import com.android.wallpaper.picker.option.ui.viewmodel.OptionItemViewModel
+import com.android.wallpaper.picker.option.ui.viewmodel.OptionItemViewModel2
 import com.android.wallpaper.testing.collectLastValue
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.Dispatchers
@@ -149,7 +149,7 @@
                         KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START to
                             FakeCustomizationProviderClient.AFFORDANCE_1,
                         KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_END to
-                            FakeCustomizationProviderClient.AFFORDANCE_2
+                            FakeCustomizationProviderClient.AFFORDANCE_2,
                     )
                 )
 
@@ -203,7 +203,7 @@
                 icon =
                     Icon.Loaded(
                         FakeCustomizationProviderClient.ICON_1,
-                        Text.Loaded("Right shortcut")
+                        Text.Loaded("Right shortcut"),
                     ),
                 text = "Right shortcut",
                 isSelected = true,
@@ -399,7 +399,7 @@
 
     private fun assertQuickAffordance(
         testScope: TestScope,
-        quickAffordance: OptionItemViewModel<Icon>?,
+        quickAffordance: OptionItemViewModel2<Icon>?,
         key: String,
         icon: Icon,
         text: Text,