diff --git a/res/layout/clock_style_option2.xml b/res/layout/clock_style_option2.xml
new file mode 100644
index 0000000..8bb60d1
--- /dev/null
+++ b/res/layout/clock_style_option2.xml
@@ -0,0 +1,46 @@
+<?xml version="1.0" encoding="utf-8"?><!--
+     Copyright (C) 2024 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 -->
+<FrameLayout
+    xmlns:android="http://schemas.android.com/apk/res/android"
+    android:layout_width="80dp"
+    android:layout_height="80dp"
+    android:clipToPadding="false"
+    android:clipChildren="false">
+
+    <com.android.wallpaper.picker.option.ui.view.OptionItemBackground
+        android:id="@id/background"
+        android:layout_width="match_parent"
+        android:layout_height="match_parent"
+        android:importantForAccessibility="no" />
+
+    <ImageView
+        android:id="@+id/foreground"
+        android:layout_width="match_parent"
+        android:layout_height="match_parent"
+        android:layout_margin="@dimen/floating_sheet_clock_style_thumbnail_margin"
+        android:src="@drawable/ic_clock_24px" />
+
+    <ImageView
+        android:id="@+id/edit_icon"
+        android:layout_width="@dimen/floating_sheet_clock_edit_icon_size"
+        android:layout_height="@dimen/floating_sheet_clock_edit_icon_size"
+        android:layout_marginTop="-16dp"
+        android:layout_marginStart="50dp"
+        android:src="@drawable/edit_icon"
+        android:importantForAccessibility="no" />
+</FrameLayout>
+
diff --git a/res/layout/floating_sheet_clock.xml b/res/layout/floating_sheet_clock.xml
index ec2e6c1..cdaf3de 100644
--- a/res/layout/floating_sheet_clock.xml
+++ b/res/layout/floating_sheet_clock.xml
@@ -46,6 +46,7 @@
                 app:layout_constraintEnd_toEndOf="parent"
                 app:layout_constraintTop_toTopOf="parent"
                 app:layout_constraintBottom_toTopOf="@+id/clock_style_clock_size_title"
+                android:layout_marginBottom="8dp"
                 android:clipToPadding="false"
                 android:clipChildren="false">
 
@@ -64,10 +65,15 @@
                     android:layout_height="@dimen/floating_sheet_clock_style_option_size"
                     android:visibility="invisible" />
 
+                <!--
+                TODO (b/377528523): We intentionally disable over scroll mode since it will clip the
+                  edit icon when in over scroll.
+                -->
                 <androidx.recyclerview.widget.RecyclerView
                     android:id="@+id/clock_style_list"
                     android:layout_width="match_parent"
                     android:layout_height="wrap_content"
+                    android:overScrollMode="never"
                     android:clipChildren="false"
                     android:clipToPadding="false"/>
             </FrameLayout>
@@ -167,7 +173,6 @@
             android:layout_width="match_parent"
             android:layout_height="wrap_content"
             android:paddingHorizontal="@dimen/floating_sheet_content_horizontal_padding"
-            android:paddingVertical="@dimen/floating_sheet_content_vertical_padding"
             android:orientation="vertical"
             android:clipToPadding="false"
             android:clipChildren="false">
diff --git a/src/com/android/wallpaper/customization/ui/binder/ClockFloatingSheetBinder.kt b/src/com/android/wallpaper/customization/ui/binder/ClockFloatingSheetBinder.kt
index 2d53c0e..5f982e2 100644
--- a/src/com/android/wallpaper/customization/ui/binder/ClockFloatingSheetBinder.kt
+++ b/src/com/android/wallpaper/customization/ui/binder/ClockFloatingSheetBinder.kt
@@ -22,11 +22,13 @@
 import android.view.LayoutInflater
 import android.view.View
 import android.view.ViewGroup
+import android.view.ViewTreeObserver.OnGlobalLayoutListener
+import android.widget.FrameLayout
 import android.widget.ImageView
+import android.widget.LinearLayout
 import android.widget.SeekBar
 import android.widget.Switch
 import android.widget.TextView
-import androidx.core.view.doOnLayout
 import androidx.core.view.isVisible
 import androidx.lifecycle.Lifecycle
 import androidx.lifecycle.LifecycleOwner
@@ -53,12 +55,15 @@
 import com.android.wallpaper.picker.customization.ui.view.adapter.FloatingToolbarTabAdapter
 import com.android.wallpaper.picker.customization.ui.viewmodel.ColorUpdateViewModel
 import com.android.wallpaper.picker.option.ui.adapter.OptionItemAdapter
+import com.android.wallpaper.picker.option.ui.adapter.OptionItemAdapter2
 import java.lang.ref.WeakReference
 import kotlin.math.abs
+import kotlinx.coroutines.DisposableHandle
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.asStateFlow
 import kotlinx.coroutines.flow.combine
+import kotlinx.coroutines.flow.filterNotNull
 import kotlinx.coroutines.launch
 
 object ClockFloatingSheetBinder {
@@ -68,8 +73,8 @@
 
     private val _clockFloatingSheetHeights: MutableStateFlow<ClockFloatingSheetHeightsViewModel?> =
         MutableStateFlow(null)
-    private val clockFloatingSheetHeights: Flow<ClockFloatingSheetHeightsViewModel?> =
-        _clockFloatingSheetHeights.asStateFlow()
+    private val clockFloatingSheetHeights: Flow<ClockFloatingSheetHeightsViewModel> =
+        _clockFloatingSheetHeights.asStateFlow().filterNotNull()
 
     fun bind(
         view: View,
@@ -90,7 +95,7 @@
                 .also { tabs.setAdapter(it) }
 
         val floatingSheetContainer =
-            view.requireViewById<ViewGroup>(R.id.clock_floating_sheet_content_container)
+            view.requireViewById<FrameLayout>(R.id.clock_floating_sheet_content_container)
 
         // Clock style
         val clockStyleContent = view.requireViewById<View>(R.id.clock_floating_sheet_style_content)
@@ -105,10 +110,10 @@
             view.requireViewById<ViewGroup>(R.id.clock_floating_sheet_font_content)
         val clockFontToolbar = view.requireViewById<ViewGroup>(R.id.clock_font_toolbar)
         clockFontToolbar.requireViewById<View>(R.id.clock_font_revert).setOnClickListener {
-            viewModel.revertFontAxes()
+            viewModel.cancelFontAxes()
         }
         clockFontToolbar.requireViewById<View>(R.id.clock_font_apply).setOnClickListener {
-            viewModel.applyFontAxes()
+            viewModel.confirmFontAxes()
         }
 
         // Clock color
@@ -137,16 +142,39 @@
         // Clock size switch
         val clockSizeSwitch = view.requireViewById<Switch>(R.id.clock_style_clock_size_switch)
 
-        view.doOnLayout {
-            if (_clockFloatingSheetHeights.value == null) {
-                _clockFloatingSheetHeights.value =
-                    ClockFloatingSheetHeightsViewModel(
-                        clockStyleContentHeight = clockStyleContent.height,
-                        clockColorContentHeight = clockColorContent.height,
-                        clockFontContentHeight = clockFontContent.height,
-                    )
+        clockStyleContent.viewTreeObserver.addOnGlobalLayoutListener(
+            object : OnGlobalLayoutListener {
+                override fun onGlobalLayout() {
+                    if (clockStyleContent.height != 0) {
+                        _clockFloatingSheetHeights.value =
+                            _clockFloatingSheetHeights.value?.copy(
+                                clockStyleContentHeight = clockStyleContent.height
+                            )
+                                ?: ClockFloatingSheetHeightsViewModel(
+                                    clockStyleContentHeight = clockStyleContent.height
+                                )
+                        clockStyleContent.viewTreeObserver.removeOnGlobalLayoutListener(this)
+                    }
+                }
             }
-        }
+        )
+
+        clockColorContent.viewTreeObserver.addOnGlobalLayoutListener(
+            object : OnGlobalLayoutListener {
+                override fun onGlobalLayout() {
+                    if (clockColorContent.height != 0) {
+                        _clockFloatingSheetHeights.value =
+                            _clockFloatingSheetHeights.value?.copy(
+                                clockColorContentHeight = clockColorContent.height
+                            )
+                                ?: ClockFloatingSheetHeightsViewModel(
+                                    clockColorContentHeight = clockColorContent.height
+                                )
+                        clockColorContent.viewTreeObserver.removeOnGlobalLayoutListener(this)
+                    }
+                }
+            }
+        )
 
         lifecycleOwner.lifecycleScope.launch {
             lifecycleOwner.repeatOnLifecycle(Lifecycle.State.STARTED) {
@@ -155,27 +183,40 @@
                 launch {
                     combine(clockFloatingSheetHeights, viewModel.selectedTab, ::Pair).collect {
                         (heights, selectedTab) ->
-                        heights ?: return@collect
+                        val (clockStyleContentHeight, clockColorContentHeight) = heights
+                        clockStyleContentHeight ?: return@collect
+                        clockColorContentHeight ?: return@collect
 
-                        val targetHeight =
-                            when (selectedTab) {
-                                Tab.STYLE -> heights.clockStyleContentHeight
-                                Tab.COLOR -> heights.clockColorContentHeight
-                                Tab.FONT -> heights.clockFontContentHeight
-                            } +
-                                view.resources.getDimensionPixelSize(
-                                    R.dimen.floating_sheet_content_vertical_padding
-                                ) * 2
+                        if (selectedTab == Tab.STYLE || selectedTab == Tab.COLOR) {
+                            val targetHeight =
+                                when (selectedTab) {
+                                    Tab.STYLE -> clockStyleContentHeight
+                                    Tab.COLOR -> clockColorContentHeight
+                                    else -> 0
+                                } +
+                                    view.resources.getDimensionPixelSize(
+                                        R.dimen.floating_sheet_content_vertical_padding
+                                    ) * 2
 
-                        val animationFloatingSheet =
                             ValueAnimator.ofInt(floatingSheetContainer.height, targetHeight)
-                        animationFloatingSheet.addUpdateListener { valueAnimator ->
-                            val value = valueAnimator.animatedValue as Int
+                                .apply {
+                                    addUpdateListener { valueAnimator ->
+                                        val value = valueAnimator.animatedValue as Int
+                                        floatingSheetContainer.layoutParams =
+                                            floatingSheetContainer.layoutParams.apply {
+                                                height = value
+                                            }
+                                    }
+                                    duration = ANIMATION_DURATION
+                                }
+                                .start()
+                        } else if (selectedTab == Tab.FONT) {
                             floatingSheetContainer.layoutParams =
-                                floatingSheetContainer.layoutParams.apply { height = value }
+                                LinearLayout.LayoutParams(
+                                    LinearLayout.LayoutParams.MATCH_PARENT,
+                                    LinearLayout.LayoutParams.WRAP_CONTENT,
+                                )
                         }
-                        animationFloatingSheet.setDuration(ANIMATION_DURATION)
-                        animationFloatingSheet.start()
 
                         clockStyleContent.isVisible = selectedTab == Tab.STYLE
                         clockColorContent.isVisible = selectedTab == Tab.COLOR
@@ -192,12 +233,6 @@
                     combine(viewModel.previewingClock, viewModel.previewingFontAxisMap, ::Pair)
                         .collect { pair ->
                             val (clock, fontAxisMap) = pair
-                            if (clock == null) {
-                                boundClockId = null
-                                boundEditorViews = mapOf()
-                                clockFontContent.removeAllViews()
-                                return@collect
-                            }
 
                             if (boundClockId != clock.clockId) {
                                 boundEditorViews =
@@ -281,6 +316,7 @@
         viewModel: ClockPickerViewModel,
     ): Map<String, Pair<View, ClockFontAxis>> {
         parent.removeAllViews()
+
         val inflater = LayoutInflater.from(parent.context)
         val axisMap = mutableMapOf<String, Pair<View, ClockFontAxis>>()
         var nextSwitch: View? = null
@@ -342,27 +378,35 @@
                 )
             }
         }
+
         return axisMap
     }
 
     private fun createClockStyleOptionItemAdapter(
         lifecycleOwner: LifecycleOwner
-    ): OptionItemAdapter<ClockStyleModel> =
-        OptionItemAdapter(
-            layoutResourceId = R.layout.clock_style_option,
+    ): OptionItemAdapter2<ClockStyleModel> =
+        OptionItemAdapter2(
+            layoutResourceId = R.layout.clock_style_option2,
             lifecycleOwner = lifecycleOwner,
-            bindIcon = { view: View, style: ClockStyleModel ->
-                (view.findViewById(R.id.clock_icon) as ImageView).setImageDrawable(style.thumbnail)
-                (view.findViewById(R.id.edit_icon) as ImageView).setVisibility(
-                    // TODO(b/364673969): Route isSelected here and hide when unselected
-                    if (style.isEditable) View.VISIBLE else View.GONE
-                )
+            bindPayload = { view: View, styleModel: ClockStyleModel ->
+                view
+                    .findViewById<ImageView>(R.id.foreground)
+                    ?.setImageDrawable(styleModel.thumbnail)
+                val job =
+                    lifecycleOwner.lifecycleScope.launch {
+                        lifecycleOwner.repeatOnLifecycle(Lifecycle.State.STARTED) {
+                            styleModel.showEditButton.collect {
+                                view.findViewById<ImageView>(R.id.edit_icon)?.isVisible = it
+                            }
+                        }
+                    }
+                return@OptionItemAdapter2 DisposableHandle { job.cancel() }
             },
         )
 
     private fun RecyclerView.initStyleList(
         context: Context,
-        adapter: OptionItemAdapter<ClockStyleModel>,
+        adapter: OptionItemAdapter2<ClockStyleModel>,
     ) {
         this.adapter = adapter
         layoutManager = LinearLayoutManager(context, LinearLayoutManager.HORIZONTAL, false)
diff --git a/src/com/android/wallpaper/customization/ui/binder/ColorsFloatingSheetBinder.kt b/src/com/android/wallpaper/customization/ui/binder/ColorsFloatingSheetBinder.kt
index 58a1fd9..7ddcb01 100644
--- a/src/com/android/wallpaper/customization/ui/binder/ColorsFloatingSheetBinder.kt
+++ b/src/com/android/wallpaper/customization/ui/binder/ColorsFloatingSheetBinder.kt
@@ -120,6 +120,9 @@
                     )
                 val night = uiMode and UI_MODE_NIGHT_MASK == UI_MODE_NIGHT_YES
                 ColorOptionIconBinder2.bind(colorOptionIconView, colorIcon, night)
+                // Return null since it does not need the lifecycleOwner to launch any job for later
+                // disposal when rebind.
+                return@OptionItemAdapter2 null
             },
         )
 
diff --git a/src/com/android/wallpaper/customization/ui/binder/ShortcutFloatingSheetBinder.kt b/src/com/android/wallpaper/customization/ui/binder/ShortcutFloatingSheetBinder.kt
index 4c72412..838ef87 100644
--- a/src/com/android/wallpaper/customization/ui/binder/ShortcutFloatingSheetBinder.kt
+++ b/src/com/android/wallpaper/customization/ui/binder/ShortcutFloatingSheetBinder.kt
@@ -156,6 +156,9 @@
                 val imageView =
                     itemView.requireViewById<ImageView>(com.android.wallpaper.R.id.foreground)
                 IconViewBinder.bind(imageView, gridIcon)
+                // Return null since it does not need the lifecycleOwner to launch any job for later
+                // disposal when rebind.
+                return@OptionItemAdapter2 null
             },
         )
 
diff --git a/src/com/android/wallpaper/customization/ui/binder/ThemePickerToolbarBinder.kt b/src/com/android/wallpaper/customization/ui/binder/ThemePickerToolbarBinder.kt
index 18b9584..357f131 100644
--- a/src/com/android/wallpaper/customization/ui/binder/ThemePickerToolbarBinder.kt
+++ b/src/com/android/wallpaper/customization/ui/binder/ThemePickerToolbarBinder.kt
@@ -16,20 +16,27 @@
 
 package com.android.wallpaper.customization.ui.binder
 
+import android.animation.ValueAnimator
+import android.view.ViewTreeObserver.OnGlobalLayoutListener
 import android.widget.Button
 import android.widget.FrameLayout
 import android.widget.Toolbar
-import androidx.core.view.isVisible
+import androidx.core.view.isInvisible
 import androidx.lifecycle.Lifecycle
 import androidx.lifecycle.LifecycleOwner
 import androidx.lifecycle.lifecycleScope
 import androidx.lifecycle.repeatOnLifecycle
 import com.android.wallpaper.customization.ui.viewmodel.ThemePickerCustomizationOptionsViewModel
+import com.android.wallpaper.customization.ui.viewmodel.ToolbarHeightsViewModel
 import com.android.wallpaper.picker.customization.ui.binder.DefaultToolbarBinder
 import com.android.wallpaper.picker.customization.ui.binder.ToolbarBinder
 import com.android.wallpaper.picker.customization.ui.viewmodel.CustomizationOptionsViewModel
 import javax.inject.Inject
 import javax.inject.Singleton
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.asStateFlow
+import kotlinx.coroutines.flow.combine
+import kotlinx.coroutines.flow.filterNotNull
 import kotlinx.coroutines.launch
 
 @Singleton
@@ -37,6 +44,9 @@
 @Inject
 constructor(private val defaultToolbarBinder: DefaultToolbarBinder) : ToolbarBinder {
 
+    private val _toolbarHeights: MutableStateFlow<ToolbarHeightsViewModel?> = MutableStateFlow(null)
+    private val toolbarHeights = _toolbarHeights.asStateFlow().filterNotNull()
+
     override fun bind(
         navButton: FrameLayout,
         toolbar: Toolbar,
@@ -60,6 +70,45 @@
             )
         }
 
+        navButton.viewTreeObserver.addOnGlobalLayoutListener(
+            object : OnGlobalLayoutListener {
+                override fun onGlobalLayout() {
+                    if (navButton.height != 0) {
+                        _toolbarHeights.value =
+                            _toolbarHeights.value?.copy(navButtonHeight = navButton.height)
+                                ?: ToolbarHeightsViewModel(navButtonHeight = navButton.height)
+                    }
+                    navButton.viewTreeObserver.removeOnGlobalLayoutListener(this)
+                }
+            }
+        )
+
+        toolbar.viewTreeObserver.addOnGlobalLayoutListener(
+            object : OnGlobalLayoutListener {
+                override fun onGlobalLayout() {
+                    if (toolbar.height != 0) {
+                        _toolbarHeights.value =
+                            _toolbarHeights.value?.copy(toolbarHeight = toolbar.height)
+                                ?: ToolbarHeightsViewModel(toolbarHeight = toolbar.height)
+                    }
+                    navButton.viewTreeObserver.removeOnGlobalLayoutListener(this)
+                }
+            }
+        )
+
+        applyButton.viewTreeObserver.addOnGlobalLayoutListener(
+            object : OnGlobalLayoutListener {
+                override fun onGlobalLayout() {
+                    if (applyButton.height != 0) {
+                        _toolbarHeights.value =
+                            _toolbarHeights.value?.copy(applyButtonHeight = applyButton.height)
+                                ?: ToolbarHeightsViewModel(applyButtonHeight = applyButton.height)
+                    }
+                    applyButton.viewTreeObserver.removeOnGlobalLayoutListener(this)
+                }
+            }
+        )
+
         lifecycleOwner.lifecycleScope.launch {
             lifecycleOwner.repeatOnLifecycle(Lifecycle.State.STARTED) {
                 launch {
@@ -68,10 +117,60 @@
                     }
                 }
 
-                launch { viewModel.isOnApplyVisible.collect { applyButton.isVisible = it } }
+                launch { viewModel.isOnApplyVisible.collect { applyButton.isInvisible = !it } }
 
                 launch { viewModel.isOnApplyEnabled.collect { applyButton.isEnabled = it } }
+
+                launch {
+                    combine(toolbarHeights, viewModel.isToolbarCollapsed, ::Pair).collect {
+                        (toolbarHeights, isToolbarCollapsed) ->
+                        val (navButtonHeight, toolbarHeight, applyButtonHeight) = toolbarHeights
+                        navButtonHeight ?: return@collect
+                        toolbarHeight ?: return@collect
+                        applyButtonHeight ?: return@collect
+
+                        val navButtonToHeight = if (isToolbarCollapsed) 0 else navButtonHeight
+                        val toolbarToHeight = if (isToolbarCollapsed) 0 else toolbarHeight
+                        val applyButtonToHeight = if (isToolbarCollapsed) 0 else applyButtonHeight
+                        ValueAnimator.ofInt(navButton.height, navButtonToHeight)
+                            .apply {
+                                addUpdateListener { valueAnimator ->
+                                    val value = valueAnimator.animatedValue as Int
+                                    navButton.layoutParams =
+                                        navButton.layoutParams.apply { height = value }
+                                }
+                                duration = ANIMATION_DURATION
+                            }
+                            .start()
+
+                        ValueAnimator.ofInt(toolbar.height, toolbarToHeight)
+                            .apply {
+                                addUpdateListener { valueAnimator ->
+                                    val value = valueAnimator.animatedValue as Int
+                                    toolbar.layoutParams =
+                                        toolbar.layoutParams.apply { height = value }
+                                }
+                                duration = ANIMATION_DURATION
+                            }
+                            .start()
+
+                        ValueAnimator.ofInt(applyButton.height, applyButtonToHeight)
+                            .apply {
+                                addUpdateListener { valueAnimator ->
+                                    val value = valueAnimator.animatedValue as Int
+                                    applyButton.layoutParams =
+                                        applyButton.layoutParams.apply { height = value }
+                                }
+                                duration = ANIMATION_DURATION
+                            }
+                            .start()
+                    }
+                }
             }
         }
     }
+
+    companion object {
+        private const val ANIMATION_DURATION = 200L
+    }
 }
diff --git a/src/com/android/wallpaper/customization/ui/viewmodel/ClockFloatingSheetHeightsViewModel.kt b/src/com/android/wallpaper/customization/ui/viewmodel/ClockFloatingSheetHeightsViewModel.kt
index 1f86533..a39ee7d 100644
--- a/src/com/android/wallpaper/customization/ui/viewmodel/ClockFloatingSheetHeightsViewModel.kt
+++ b/src/com/android/wallpaper/customization/ui/viewmodel/ClockFloatingSheetHeightsViewModel.kt
@@ -17,7 +17,6 @@
 package com.android.wallpaper.customization.ui.viewmodel
 
 data class ClockFloatingSheetHeightsViewModel(
-    val clockStyleContentHeight: Int,
-    val clockColorContentHeight: Int,
-    val clockFontContentHeight: Int,
+    val clockStyleContentHeight: Int? = null,
+    val clockColorContentHeight: Int? = null,
 )
diff --git a/src/com/android/wallpaper/customization/ui/viewmodel/ClockPickerViewModel.kt b/src/com/android/wallpaper/customization/ui/viewmodel/ClockPickerViewModel.kt
index 11124f9..da3e65c 100644
--- a/src/com/android/wallpaper/customization/ui/viewmodel/ClockPickerViewModel.kt
+++ b/src/com/android/wallpaper/customization/ui/viewmodel/ClockPickerViewModel.kt
@@ -53,6 +53,7 @@
 import kotlinx.coroutines.flow.asStateFlow
 import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.flow.distinctUntilChanged
+import kotlinx.coroutines.flow.filterNotNull
 import kotlinx.coroutines.flow.flowOn
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.mapLatest
@@ -123,7 +124,7 @@
             overridingClock ?: selectedClock
         }
 
-    data class ClockStyleModel(val thumbnail: Drawable, val isEditable: Boolean)
+    data class ClockStyleModel(val thumbnail: Drawable, val showEditButton: StateFlow<Boolean>)
 
     @OptIn(ExperimentalCoroutinesApi::class)
     val clockStyleOptions: StateFlow<List<OptionItemViewModel<ClockStyleModel>>> =
@@ -131,40 +132,10 @@
             .mapLatest { allClocks ->
                 // Delay to avoid the case that the full list of clocks is not initiated.
                 delay(CLOCKS_EVENT_UPDATE_DELAY_MILLIS)
-                allClocks.map { clockModel ->
-                    val isSelectedFlow =
-                        previewingClock
-                            .map { it.clockId == clockModel.clockId }
-                            .stateIn(viewModelScope)
-                    val contentDescription =
-                        resources.getString(
-                            R.string.select_clock_action_description,
-                            clockModel.description,
-                        )
-                    OptionItemViewModel<ClockStyleModel>(
-                        key = MutableStateFlow(clockModel.clockId) as StateFlow<String>,
-                        payload =
-                            ClockStyleModel(
-                                clockModel.thumbnail,
-                                isEditable = !clockModel.fontAxes.isEmpty(),
-                            ),
-                        text = Text.Loaded(contentDescription),
-                        isTextUserVisible = false,
-                        isSelected = isSelectedFlow,
-                        onClicked =
-                            isSelectedFlow.map { isSelected ->
-                                if (isSelected) {
-                                    fun() {
-                                        _selectedTab.value = Tab.FONT
-                                    }
-                                } else {
-                                    fun() {
-                                        overridingClock.value = clockModel
-                                        overrideFontAxisMap.value = null
-                                    }
-                                }
-                            },
-                    )
+                val allClockMap = allClocks.groupBy { it.fontAxes.isNotEmpty() }
+                buildList {
+                    allClockMap[true]?.map { add(it.toOption(resources)) }
+                    allClockMap[false]?.map { add(it.toOption(resources)) }
                 }
             }
             // makes sure that the operations above this statement are executed on I/O dispatcher
@@ -173,27 +144,65 @@
             .flowOn(backgroundDispatcher.limitedParallelism(1))
             .stateIn(viewModelScope, SharingStarted.Eagerly, emptyList())
 
+    private suspend fun ClockMetadataModel.toOption(
+        resources: Resources
+    ): OptionItemViewModel<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>(
+            key = MutableStateFlow(clockId) as StateFlow<String>,
+            payload = ClockStyleModel(thumbnail = thumbnail, showEditButton = showEditButton),
+            text = Text.Loaded(contentDescription),
+            isTextUserVisible = false,
+            isSelected = isSelectedFlow,
+            onClicked =
+                isSelectedFlow.map { isSelected ->
+                    if (isSelected && isEditable) {
+                        fun() {
+                            _selectedTab.value = Tab.FONT
+                        }
+                    } else {
+                        fun() {
+                            overridingClock.value = this
+                            overrideFontAxisMap.value = null
+                        }
+                    }
+                },
+        )
+    }
+
     // Clock Font Axis Editor
     private val overrideFontAxisMap = MutableStateFlow<Map<String, Float>?>(null)
+    private val isFontAxisMapEdited = overrideFontAxisMap.map { it != null }
+    private val selectedClockFontAxes =
+        previewingClock
+            .map { clock -> clock.fontAxes.associate { it.key to it.currentValue } }
+            .stateIn(viewModelScope, SharingStarted.Eagerly, null)
     val previewingFontAxisMap =
-        combine(overrideFontAxisMap, previewingClock) { overrideAxes, previewingClock ->
-                overrideAxes ?: previewingClock.fontAxes.associate { it.key to it.currentValue }
+        combine(overrideFontAxisMap, selectedClockFontAxes.filterNotNull()) {
+                overrideAxes,
+                selectedFontAxes ->
+                selectedFontAxes.toMutableMap().let { mutableMap ->
+                    overrideAxes?.forEach { (key, value) -> mutableMap[key] = value }
+                    mutableMap.toMap()
+                }
             }
             .stateIn(viewModelScope, SharingStarted.Eagerly, emptyMap())
 
-    private val isFontAxisMapEdited = overrideFontAxisMap.map { it != null }
-
     fun updatePreviewFontAxis(key: String, value: Float) {
-        val axisMap = previewingFontAxisMap.value.toMutableMap()
+        val axisMap = (overrideFontAxisMap.value?.toMutableMap() ?: mutableMapOf())
         axisMap[key] = value
-        overrideFontAxisMap.value = axisMap
+        overrideFontAxisMap.value = axisMap.toMap()
     }
 
-    fun applyFontAxes() {
+    fun confirmFontAxes() {
         _selectedTab.value = Tab.STYLE
     }
 
-    fun revertFontAxes() {
+    fun cancelFontAxes() {
         overrideFontAxisMap.value = null
         _selectedTab.value = Tab.STYLE
     }
diff --git a/src/com/android/wallpaper/customization/ui/viewmodel/ThemePickerCustomizationOptionsViewModel.kt b/src/com/android/wallpaper/customization/ui/viewmodel/ThemePickerCustomizationOptionsViewModel.kt
index b2e2039..54ae132 100644
--- a/src/com/android/wallpaper/customization/ui/viewmodel/ThemePickerCustomizationOptionsViewModel.kt
+++ b/src/com/android/wallpaper/customization/ui/viewmodel/ThemePickerCustomizationOptionsViewModel.kt
@@ -30,6 +30,7 @@
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.combine
+import kotlinx.coroutines.flow.distinctUntilChanged
 import kotlinx.coroutines.flow.flatMapLatest
 import kotlinx.coroutines.flow.flow
 import kotlinx.coroutines.flow.map
@@ -171,6 +172,14 @@
 
     val isOnApplyVisible: Flow<Boolean> = selectedOption.map { it != null }
 
+    val isToolbarCollapsed: Flow<Boolean> =
+        combine(selectedOption, clockPickerViewModel.selectedTab) { selectedOption, selectedTab ->
+                selectedOption ==
+                    ThemePickerCustomizationOptionUtil.ThemePickerLockCustomizationOption.CLOCK &&
+                    selectedTab == ClockPickerViewModel.Tab.FONT
+            }
+            .distinctUntilChanged()
+
     @ViewModelScoped
     @AssistedFactory
     interface Factory : CustomizationOptionsViewModelFactory {
diff --git a/src/com/android/wallpaper/customization/ui/viewmodel/ToolbarHeightsViewModel.kt b/src/com/android/wallpaper/customization/ui/viewmodel/ToolbarHeightsViewModel.kt
new file mode 100644
index 0000000..0d859da
--- /dev/null
+++ b/src/com/android/wallpaper/customization/ui/viewmodel/ToolbarHeightsViewModel.kt
@@ -0,0 +1,23 @@
+/*
+ * Copyright (C) 2024 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.wallpaper.customization.ui.viewmodel
+
+data class ToolbarHeightsViewModel(
+    val navButtonHeight: Int? = null,
+    val toolbarHeight: Int? = null,
+    val applyButtonHeight: Int? = null,
+)
diff --git a/tests/robotests/src/com/android/wallpaper/customization/ui/viewmodel/ClockPickerViewModelTest.kt b/tests/robotests/src/com/android/wallpaper/customization/ui/viewmodel/ClockPickerViewModelTest.kt
index ab64f89..0a8aa53 100644
--- a/tests/robotests/src/com/android/wallpaper/customization/ui/viewmodel/ClockPickerViewModelTest.kt
+++ b/tests/robotests/src/com/android/wallpaper/customization/ui/viewmodel/ClockPickerViewModelTest.kt
@@ -243,7 +243,7 @@
         assertThat(selectedTab()).isEqualTo(Tab.FONT)
         assertThat(previewingFontAxes()).isEqualTo(mapOf("key" to 100f))
 
-        underTest.applyFontAxes()
+        underTest.confirmFontAxes()
 
         assertThat(selectedTab()).isEqualTo(Tab.STYLE)
         assertThat(previewingFontAxes()).isEqualTo(mapOf("key" to 100f))
@@ -269,7 +269,7 @@
         assertThat(selectedTab()).isEqualTo(Tab.FONT)
         assertThat(previewingFontAxes()).isEqualTo(mapOf("key" to 100f))
 
-        underTest.revertFontAxes()
+        underTest.cancelFontAxes()
 
         assertThat(selectedTab()).isEqualTo(Tab.STYLE)
         assertThat(previewingFontAxes()).isEqualTo(mapOf("key" to 50f))