diff --git a/src/com/android/wallpaper/customization/ui/binder/ShapeAndGridFloatingSheetBinder.kt b/src/com/android/wallpaper/customization/ui/binder/ShapeAndGridFloatingSheetBinder.kt
index f03bc9b..7217f61 100644
--- a/src/com/android/wallpaper/customization/ui/binder/ShapeAndGridFloatingSheetBinder.kt
+++ b/src/com/android/wallpaper/customization/ui/binder/ShapeAndGridFloatingSheetBinder.kt
@@ -51,7 +51,17 @@
 
         lifecycleOwner.lifecycleScope.launch {
             lifecycleOwner.repeatOnLifecycle(Lifecycle.State.STARTED) {
-                launch { viewModel.optionItems.collect { options -> adapter.setItems(options) } }
+                launch {
+                    viewModel.optionItems.collect { options ->
+                        adapter.setItems(options) {
+                            val indexToFocus =
+                                options.indexOfFirst { it.isSelected.value }.coerceAtLeast(0)
+                            (gridOptionList.layoutManager as LinearLayoutManager).scrollToPosition(
+                                indexToFocus
+                            )
+                        }
+                    }
+                }
             }
         }
     }
diff --git a/src/com/android/wallpaper/customization/ui/binder/ThemePickerToolbarBinder.kt b/src/com/android/wallpaper/customization/ui/binder/ThemePickerToolbarBinder.kt
index 6bddf4a..088a741 100644
--- a/src/com/android/wallpaper/customization/ui/binder/ThemePickerToolbarBinder.kt
+++ b/src/com/android/wallpaper/customization/ui/binder/ThemePickerToolbarBinder.kt
@@ -19,6 +19,7 @@
 import android.widget.Button
 import android.widget.FrameLayout
 import android.widget.Toolbar
+import androidx.core.view.isVisible
 import androidx.lifecycle.Lifecycle
 import androidx.lifecycle.LifecycleOwner
 import androidx.lifecycle.lifecycleScope
@@ -54,12 +55,16 @@
         lifecycleOwner.lifecycleScope.launch {
             lifecycleOwner.repeatOnLifecycle(Lifecycle.State.STARTED) {
                 launch {
-                    viewModel.keyguardQuickAffordancePickerViewModel2.onApply.collect { onApply ->
+                    viewModel.onApplyButtonClicked.collect { onApplyButtonClicked ->
                         applyButton.setOnClickListener {
-                            onApply?.invoke()?.let { viewModel.deselectOption() }
+                            onApplyButtonClicked?.invoke()?.let { viewModel.deselectOption() }
                         }
                     }
                 }
+
+                launch { viewModel.isOnApplyVisible.collect { applyButton.isVisible = it } }
+
+                launch { viewModel.isOnApplyEnabled.collect { applyButton.isEnabled = it } }
             }
         }
     }
diff --git a/src/com/android/wallpaper/customization/ui/viewmodel/KeyguardQuickAffordancePickerViewModel2.kt b/src/com/android/wallpaper/customization/ui/viewmodel/KeyguardQuickAffordancePickerViewModel2.kt
index 86cdd8a..658e435 100644
--- a/src/com/android/wallpaper/customization/ui/viewmodel/KeyguardQuickAffordancePickerViewModel2.kt
+++ b/src/com/android/wallpaper/customization/ui/viewmodel/KeyguardQuickAffordancePickerViewModel2.kt
@@ -83,12 +83,12 @@
                 started = SharingStarted.WhileSubscribed(),
                 initialValue = "",
             )
-    private val _selectedQuickAffordances = MutableStateFlow<Map<String, String>>(emptyMap())
-    val selectedQuickAffordances: Flow<Map<String, String>> =
-        _selectedQuickAffordances.asStateFlow()
+    private val _previewingQuickAffordances = MutableStateFlow<Map<String, String>>(emptyMap())
+    val previewingQuickAffordances: Flow<Map<String, String>> =
+        _previewingQuickAffordances.asStateFlow()
 
     fun resetPreview() {
-        _selectedQuickAffordances.tryEmit(emptyMap())
+        _previewingQuickAffordances.tryEmit(emptyMap())
         _selectedSlotId.tryEmit(SLOT_ID_BOTTOM_START)
     }
 
@@ -98,7 +98,7 @@
                 quickAffordanceInteractor.slots,
                 quickAffordanceInteractor.affordances,
                 quickAffordanceInteractor.selections,
-                selectedQuickAffordances,
+                previewingQuickAffordances,
                 selectedSlotId,
             ) { slots, affordances, selections, selectedQuickAffordances, selectedSlotId ->
                 slots.associate { slot ->
@@ -194,7 +194,7 @@
             val isNoneSelected =
                 combine(
                         selectedSlotId,
-                        selectedQuickAffordances,
+                        previewingQuickAffordances,
                         selectedAffordanceIds,
                     ) { selectedSlotId, selectedQuickAffordances, selectedAffordanceIds ->
                         selectedQuickAffordances[selectedSlotId]?.let {
@@ -214,10 +214,10 @@
                             if (!isSelected) {
                                 {
                                     val newMap =
-                                        _selectedQuickAffordances.value.toMutableMap().apply {
+                                        _previewingQuickAffordances.value.toMutableMap().apply {
                                             put(selectedSlotId, KEYGUARD_QUICK_AFFORDANCE_ID_NONE)
                                         }
-                                    _selectedQuickAffordances.tryEmit(newMap)
+                                    _previewingQuickAffordances.tryEmit(newMap)
                                 }
                             } else {
                                 null
@@ -230,7 +230,7 @@
                     val isSelectedFlow: StateFlow<Boolean> =
                         combine(
                                 selectedSlotId,
-                                selectedQuickAffordances,
+                                previewingQuickAffordances,
                                 selectedAffordanceIds,
                             ) { selectedSlotId, selectedQuickAffordances, selectedAffordanceIds ->
                                 selectedQuickAffordances[selectedSlotId]?.let {
@@ -255,10 +255,10 @@
                                     if (!isSelected) {
                                         {
                                             val newMap =
-                                                _selectedQuickAffordances.value
+                                                _previewingQuickAffordances.value
                                                     .toMutableMap()
                                                     .apply { put(selectedSlotId, affordance.id) }
-                                            _selectedQuickAffordances.tryEmit(newMap)
+                                            _previewingQuickAffordances.tryEmit(newMap)
                                         }
                                     } else {
                                         null
@@ -287,7 +287,7 @@
         }
 
     val onApply: Flow<(() -> Unit)?> =
-        selectedQuickAffordances.map {
+        previewingQuickAffordances.map {
             if (it.isEmpty()) {
                 null
             } else {
diff --git a/src/com/android/wallpaper/customization/ui/viewmodel/ShapeAndGridPickerViewModel.kt b/src/com/android/wallpaper/customization/ui/viewmodel/ShapeAndGridPickerViewModel.kt
index c3ff291..2ce4747 100644
--- a/src/com/android/wallpaper/customization/ui/viewmodel/ShapeAndGridPickerViewModel.kt
+++ b/src/com/android/wallpaper/customization/ui/viewmodel/ShapeAndGridPickerViewModel.kt
@@ -58,6 +58,10 @@
             previewingGridOptionKey ?: currentlySetGridOption.key.value
         }
 
+    fun resetPreview() {
+        _previewingGridOptionKey.tryEmit(null)
+    }
+
     val optionItems: Flow<List<OptionItemViewModel<GridIconViewModel>>> =
         interactor.gridOptions.filterNotNull().map { gridOptions ->
             gridOptions.map { toOptionItemViewModel(it) }
@@ -81,8 +85,6 @@
             }
         }
 
-    val isOnApplyEnabled: Flow<Boolean> = onApply.map { it != null }
-
     private fun toOptionItemViewModel(
         option: GridOptionModel
     ): OptionItemViewModel<GridIconViewModel> {
diff --git a/src/com/android/wallpaper/customization/ui/viewmodel/ThemePickerCustomizationOptionsViewModel.kt b/src/com/android/wallpaper/customization/ui/viewmodel/ThemePickerCustomizationOptionsViewModel.kt
index 218818a..0adff1b 100644
--- a/src/com/android/wallpaper/customization/ui/viewmodel/ThemePickerCustomizationOptionsViewModel.kt
+++ b/src/com/android/wallpaper/customization/ui/viewmodel/ThemePickerCustomizationOptionsViewModel.kt
@@ -25,8 +25,13 @@
 import dagger.assisted.AssistedInject
 import dagger.hilt.android.scopes.ViewModelScoped
 import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.SharingStarted
+import kotlinx.coroutines.flow.flatMapLatest
+import kotlinx.coroutines.flow.flow
 import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.flow.stateIn
 
 class ThemePickerCustomizationOptionsViewModel
 @AssistedInject
@@ -53,6 +58,8 @@
 
     override fun deselectOption(): Boolean {
         keyguardQuickAffordancePickerViewModel2.resetPreview()
+        shapeAndGridPickerViewModel.resetPreview()
+
         return defaultCustomizationOptionsViewModel.deselectOption()
     }
 
@@ -110,6 +117,24 @@
             }
         }
 
+    @OptIn(ExperimentalCoroutinesApi::class)
+    val onApplyButtonClicked =
+        selectedOption
+            .flatMapLatest {
+                when (it) {
+                    ThemePickerCustomizationOptionUtil.ThemePickerLockCustomizationOption
+                        .SHORTCUTS -> keyguardQuickAffordancePickerViewModel2.onApply
+                    ThemePickerCustomizationOptionUtil.ThemePickerHomeCustomizationOption
+                        .APP_SHAPE_AND_GRID -> shapeAndGridPickerViewModel.onApply
+                    else -> flow { emit(null) }
+                }
+            }
+            .stateIn(viewModelScope, SharingStarted.Eagerly, null)
+
+    val isOnApplyEnabled: Flow<Boolean> = onApplyButtonClicked.map { it != null }
+
+    val isOnApplyVisible: Flow<Boolean> = selectedOption.map { it != null }
+
     @ViewModelScoped
     @AssistedFactory
     interface Factory : CustomizationOptionsViewModelFactory {
diff --git a/src/com/android/wallpaper/picker/common/preview/ui/binder/ThemePickerWorkspaceCallbackBinder.kt b/src/com/android/wallpaper/picker/common/preview/ui/binder/ThemePickerWorkspaceCallbackBinder.kt
index fac7230..eec7d5a 100644
--- a/src/com/android/wallpaper/picker/common/preview/ui/binder/ThemePickerWorkspaceCallbackBinder.kt
+++ b/src/com/android/wallpaper/picker/common/preview/ui/binder/ThemePickerWorkspaceCallbackBinder.kt
@@ -104,7 +104,7 @@
 
                         launch {
                             viewModel.keyguardQuickAffordancePickerViewModel2
-                                .selectedQuickAffordances
+                                .previewingQuickAffordances
                                 .collect {
                                     it[SLOT_ID_BOTTOM_START]?.let {
                                         workspaceCallback.sendMessage(
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 19d74f6..b6f249e 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
@@ -117,19 +117,19 @@
     @Test
     fun selectedQuickAffordancesMapUpdates_whenClickingOnQuickAffordanceOptionsAndCallingResetPreview() =
         testScope.runTest {
-            val selectedQuickAffordances = collectLastValue(underTest.selectedQuickAffordances)
+            val previewingQuickAffordances = collectLastValue(underTest.previewingQuickAffordances)
 
             val tabs = collectLastValue(underTest.tabs)
             val quickAffordances = collectLastValue(underTest.quickAffordances)
 
             // Default selectedQuickAffordances is an empty map
-            assertThat(selectedQuickAffordances()).isEqualTo(emptyMap<String, String>())
+            assertThat(previewingQuickAffordances()).isEqualTo(emptyMap<String, String>())
 
             // Click on quick affordance 1 when selected slot ID is bottom_start
             val onClickAffordance1 =
                 collectLastValue(quickAffordances()?.get(1)?.onClicked ?: emptyFlow())
             onClickAffordance1()?.invoke()
-            assertThat(selectedQuickAffordances())
+            assertThat(previewingQuickAffordances())
                 .isEqualTo(
                     mapOf(
                         KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START to
@@ -143,7 +143,7 @@
             val onClickAffordance2 =
                 collectLastValue(quickAffordances()?.get(2)?.onClicked ?: emptyFlow())
             onClickAffordance2()?.invoke()
-            assertThat(selectedQuickAffordances())
+            assertThat(previewingQuickAffordances())
                 .isEqualTo(
                     mapOf(
                         KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START to
@@ -154,7 +154,7 @@
                 )
 
             underTest.resetPreview()
-            assertThat(selectedQuickAffordances()).isEqualTo(emptyMap<String, String>())
+            assertThat(previewingQuickAffordances()).isEqualTo(emptyMap<String, String>())
         }
 
     @Test