diff --git a/Android.bp b/Android.bp
index 6d9ff8f..f6c8558 100644
--- a/Android.bp
+++ b/Android.bp
@@ -77,6 +77,7 @@
         "SettingsLibSettingsTheme",
         "SystemUI-statsd",
         "styleprotoslite",
+        "androidx.lifecycle_lifecycle-livedata-ktx",
         "androidx.lifecycle_lifecycle-runtime-ktx",
         "androidx.lifecycle_lifecycle-viewmodel-ktx",
         "androidx.recyclerview_recyclerview",
diff --git a/res/drawable/ic_nav_color.xml b/res/drawable/ic_nav_color.xml
new file mode 100644
index 0000000..cfa64e2
--- /dev/null
+++ b/res/drawable/ic_nav_color.xml
@@ -0,0 +1,23 @@
+<!--
+     Copyright (C) 2019 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.
+-->
+<!-- Represents the color icon (a palette) -->
+<vector xmlns:android="http://schemas.android.com/apk/res/android"
+    android:width="24dp"
+    android:height="24dp"
+    android:viewportWidth="24"
+    android:viewportHeight="24">
+    <path android:fillColor="@android:color/white" android:pathData="M12,22Q9.95,22 8.125,21.212Q6.3,20.425 4.938,19.062Q3.575,17.7 2.788,15.875Q2,14.05 2,12Q2,9.925 2.812,8.1Q3.625,6.275 5.013,4.925Q6.4,3.575 8.25,2.787Q10.1,2 12.2,2Q14.2,2 15.975,2.688Q17.75,3.375 19.087,4.588Q20.425,5.8 21.212,7.463Q22,9.125 22,11.05Q22,13.925 20.25,15.462Q18.5,17 16,17H14.15Q13.925,17 13.838,17.125Q13.75,17.25 13.75,17.4Q13.75,17.7 14.125,18.262Q14.5,18.825 14.5,19.55Q14.5,20.8 13.812,21.4Q13.125,22 12,22ZM12,12Q12,12 12,12Q12,12 12,12Q12,12 12,12Q12,12 12,12Q12,12 12,12Q12,12 12,12Q12,12 12,12Q12,12 12,12Q12,12 12,12Q12,12 12,12Q12,12 12,12Q12,12 12,12Q12,12 12,12Q12,12 12,12ZM6.5,13Q7.15,13 7.575,12.575Q8,12.15 8,11.5Q8,10.85 7.575,10.425Q7.15,10 6.5,10Q5.85,10 5.425,10.425Q5,10.85 5,11.5Q5,12.15 5.425,12.575Q5.85,13 6.5,13ZM9.5,9Q10.15,9 10.575,8.575Q11,8.15 11,7.5Q11,6.85 10.575,6.425Q10.15,6 9.5,6Q8.85,6 8.425,6.425Q8,6.85 8,7.5Q8,8.15 8.425,8.575Q8.85,9 9.5,9ZM14.5,9Q15.15,9 15.575,8.575Q16,8.15 16,7.5Q16,6.85 15.575,6.425Q15.15,6 14.5,6Q13.85,6 13.425,6.425Q13,6.85 13,7.5Q13,8.15 13.425,8.575Q13.85,9 14.5,9ZM17.5,13Q18.15,13 18.575,12.575Q19,12.15 19,11.5Q19,10.85 18.575,10.425Q18.15,10 17.5,10Q16.85,10 16.425,10.425Q16,10.85 16,11.5Q16,12.15 16.425,12.575Q16.85,13 17.5,13ZM12,20Q12.225,20 12.363,19.875Q12.5,19.75 12.5,19.55Q12.5,19.2 12.125,18.725Q11.75,18.25 11.75,17.3Q11.75,16.25 12.475,15.625Q13.2,15 14.25,15H16Q17.65,15 18.825,14.037Q20,13.075 20,11.05Q20,8.025 17.688,6.012Q15.375,4 12.2,4Q8.8,4 6.4,6.325Q4,8.65 4,12Q4,15.325 6.338,17.663Q8.675,20 12,20Z"/>
+</vector>
diff --git a/res/layout/color_section_view2.xml b/res/layout/color_section_view2.xml
index 0a3fc7f..687bcef 100644
--- a/res/layout/color_section_view2.xml
+++ b/res/layout/color_section_view2.xml
@@ -14,18 +14,14 @@
      See the License for the specific language governing permissions and
      limitations under the License.
 -->
-<com.android.customization.picker.color.ui.view.ColorSectionView2
-    xmlns:android="http://schemas.android.com/apk/res/android"
-    android:id="@+id/color_section_option_container"
+<com.android.customization.picker.color.ui.view.ColorSectionView2 xmlns:android="http://schemas.android.com/apk/res/android"
+    xmlns:tools="http://schemas.android.com/tools"
     android:layout_width="match_parent"
     android:layout_height="wrap_content"
     android:layout_marginBottom="@dimen/section_bottom_padding"
     android:layout_marginHorizontal="@dimen/section_horizontal_padding"
-    android:orientation="horizontal"
-    android:background="@drawable/top_connected_section_background"
-    android:paddingVertical="24dp"
-    android:paddingHorizontal="24dp"
-    android:weightSum="@integer/color_section_num_columns">
+    android:orientation="vertical"
+    android:background="@drawable/top_connected_section_background">
 
     <!--
         This is just an invisible placeholder put in place so that the parent keeps its height
@@ -36,10 +32,36 @@
 
         It's critical for any TextViews inside the included layout to have text.
         -->
-    <include
-        android:layout_width="0dp"
+    <LinearLayout
+        android:id="@+id/color_section_option_container"
+        android:layout_width="match_parent"
         android:layout_height="wrap_content"
-        layout="@layout/color_option_overflow_no_background"
-        android:visibility="invisible"
-        android:layout_weight="1"/>
+        android:orientation="horizontal"
+        android:paddingVertical="24dp"
+        android:paddingHorizontal="24dp"
+        android:weightSum="@integer/color_section_num_columns">
+        <include
+            android:layout_width="0dp"
+            android:layout_height="wrap_content"
+            layout="@layout/color_option_overflow_no_background"
+            android:visibility="invisible"
+            android:layout_weight="1"/>
+    </LinearLayout>
+
+    <TextView
+        android:id="@+id/more_colors"
+        android:layout_width="wrap_content"
+        android:layout_height="wrap_content"
+        android:layout_gravity="center_horizontal"
+        android:layout_marginTop="10dp"
+        android:minHeight="48dp"
+        android:gravity="center"
+        android:drawablePadding="12dp"
+        android:drawableStart="@drawable/ic_nav_color"
+        android:drawableTint="@color/text_color_primary"
+        android:text="@string/more_colors"
+        android:textAppearance="@android:style/TextAppearance.DeviceDefault.Small"
+        android:textColor="@color/text_color_primary"
+        android:visibility="gone"
+        tools:ignore="UseCompatTextViewDrawableXml" />
 </com.android.customization.picker.color.ui.view.ColorSectionView2>
diff --git a/res/values/strings.xml b/res/values/strings.xml
index f24ea28..f4d525e 100755
--- a/res/values/strings.xml
+++ b/res/values/strings.xml
@@ -412,4 +412,12 @@
     [CHAR LIMIT=64].
     -->
     <string name="more_settings_section_description">Text on lock screen, Now Playing, and more</string>
+
+    <!--
+    Label for button that lets the user navigate to a full-screen experience of selecting
+    system colors.
+
+    [CHAR LIMIT=128].
+    -->
+    <string name="more_colors">More Colors</string>
 </resources>
diff --git a/src/com/android/customization/model/color/ColorSectionController2.java b/src/com/android/customization/model/color/ColorSectionController2.java
deleted file mode 100644
index 791a9a0..0000000
--- a/src/com/android/customization/model/color/ColorSectionController2.java
+++ /dev/null
@@ -1,243 +0,0 @@
-/*
- * 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.color;
-
-import static com.android.customization.model.color.ColorOptionsProvider.COLOR_SOURCE_HOME;
-import static com.android.customization.model.color.ColorOptionsProvider.COLOR_SOURCE_LOCK;
-import static com.android.customization.model.color.ColorOptionsProvider.COLOR_SOURCE_PRESET;
-
-import android.app.Activity;
-import android.app.WallpaperColors;
-import android.content.Context;
-import android.os.Handler;
-import android.os.SystemClock;
-import android.stats.style.StyleEnums;
-import android.text.TextUtils;
-import android.util.Log;
-import android.view.LayoutInflater;
-
-import androidx.annotation.Nullable;
-import androidx.lifecycle.LifecycleOwner;
-
-import com.android.customization.model.CustomizationManager;
-import com.android.customization.model.theme.OverlayManagerCompat;
-import com.android.customization.module.CustomizationInjector;
-import com.android.customization.module.ThemesUserEventLogger;
-import com.android.customization.picker.color.ui.fragment.ColorPickerFragment;
-import com.android.customization.picker.color.ui.view.ColorSectionView2;
-import com.android.wallpaper.R;
-import com.android.wallpaper.model.CustomizationSectionController;
-import com.android.wallpaper.model.WallpaperColorsViewModel;
-import com.android.wallpaper.module.InjectorProvider;
-
-import com.google.common.collect.Iterables;
-import com.google.common.collect.Lists;
-
-import java.util.ArrayList;
-import java.util.List;
-
-/**
- * Color section view's controller for the logic of color customization.
- *
- * TODO (b/262924584): Convert ColorSectionController2 into Kotlin & use new architecture
- */
-public class ColorSectionController2 implements CustomizationSectionController<ColorSectionView2> {
-
-    private static final String TAG = "ColorSectionController";
-    private static final long MIN_COLOR_APPLY_PERIOD = 500L;
-
-    private final ThemesUserEventLogger mEventLogger;
-    private final ColorCustomizationManager mColorManager;
-    private final WallpaperColorsViewModel mWallpaperColorsViewModel;
-    private final LifecycleOwner mLifecycleOwner;
-    private final CustomizationSectionNavigationController mSectionNavigationController;
-
-    private List<ColorOption> mWallpaperColorOptions = new ArrayList<>();
-    private List<ColorOption> mPresetColorOptions = new ArrayList<>();
-    private ColorOption mSelectedColor;
-    @Nullable private WallpaperColors mHomeWallpaperColors;
-    @Nullable private WallpaperColors mLockWallpaperColors;
-    // Uses a boolean value to indicate whether wallpaper color is ready because WallpaperColors
-    // maybe be null when it's ready.
-    private boolean mHomeWallpaperColorsReady;
-    private boolean mLockWallpaperColorsReady;
-    private long mLastColorApplyingTime = 0L;
-    private ColorSectionView2 mColorSectionView;
-
-    public ColorSectionController2(Activity activity, WallpaperColorsViewModel viewModel,
-            LifecycleOwner lifecycleOwner,
-            CustomizationSectionNavigationController sectionNavigationController) {
-        CustomizationInjector injector = (CustomizationInjector) InjectorProvider.getInjector();
-        mEventLogger = (ThemesUserEventLogger) injector.getUserEventLogger(activity);
-        mColorManager = ColorCustomizationManager.getInstance(activity,
-                new OverlayManagerCompat(activity));
-        mWallpaperColorsViewModel = viewModel;
-        mLifecycleOwner = lifecycleOwner;
-        mSectionNavigationController = sectionNavigationController;
-    }
-
-    @Override
-    public boolean isAvailable(@Nullable Context context) {
-        return context != null && ColorUtils.isMonetEnabled(context) && mColorManager.isAvailable();
-    }
-
-    @Override
-    public ColorSectionView2 createView(Context context) {
-        mColorSectionView = (ColorSectionView2) LayoutInflater.from(context).inflate(
-                R.layout.color_section_view2, /* root= */ null);
-
-        mWallpaperColorsViewModel.getHomeWallpaperColorsLiveData().observe(mLifecycleOwner,
-                homeColors -> {
-                    mHomeWallpaperColors = homeColors;
-                    mHomeWallpaperColorsReady = true;
-                    maybeLoadColors();
-                });
-        mWallpaperColorsViewModel.getLockWallpaperColorsLiveData().observe(mLifecycleOwner,
-                lockColors -> {
-                    mLockWallpaperColors = lockColors;
-                    mLockWallpaperColorsReady = true;
-                    maybeLoadColors();
-                });
-        return mColorSectionView;
-    }
-
-    private void maybeLoadColors() {
-        if (mHomeWallpaperColorsReady && mLockWallpaperColorsReady) {
-            mColorManager.setWallpaperColors(mHomeWallpaperColors, mLockWallpaperColors);
-            loadColorOptions(/* reload= */ false);
-        }
-    }
-
-    private void loadColorOptions(boolean reload) {
-        mColorManager.fetchOptions(new CustomizationManager.OptionsFetchedListener<ColorOption>() {
-            @Override
-            public void onOptionsLoaded(List<ColorOption> options) {
-                List<ColorOption> wallpaperColorOptions = new ArrayList<>();
-                List<ColorOption> presetColorOptions = new ArrayList<>();
-                for (ColorOption option : options) {
-                    if (option instanceof ColorSeedOption) {
-                        wallpaperColorOptions.add(option);
-                    } else if (option instanceof ColorBundle) {
-                        presetColorOptions.add(option);
-                    }
-                }
-                mWallpaperColorOptions = wallpaperColorOptions;
-                mPresetColorOptions = presetColorOptions;
-                mSelectedColor = findActiveColorOption(mWallpaperColorOptions,
-                        mPresetColorOptions);
-
-                mColorSectionView.post(() -> setUpColorSectionView(mWallpaperColorOptions,
-                        mPresetColorOptions));
-            }
-
-            @Override
-            public void onError(@Nullable Throwable throwable) {
-                if (throwable != null) {
-                    Log.e(TAG, "Error loading theme bundles", throwable);
-                }
-            }
-        }, reload);
-    }
-
-    private void setUpColorSectionView(List<ColorOption> wallpaperColorOptions,
-            List<ColorOption> presetColorOptions) {
-        int wallpaperOptionSize = wallpaperColorOptions.size();
-
-        List<ColorOption> subOptions = wallpaperColorOptions.subList(0,
-                Math.min(5, wallpaperOptionSize));
-        // add additional options based on preset colors if there are less than 5 wallpaper colors
-        List<ColorOption> additionalSubOptions = presetColorOptions.subList(0,
-                Math.min(Math.max(0, 5 - wallpaperOptionSize), presetColorOptions.size()));
-        subOptions.addAll(additionalSubOptions);
-
-        mColorSectionView.setOverflowOnClick(() -> {
-            mSectionNavigationController.navigateTo(new ColorPickerFragment());
-            return null;
-        });
-        mColorSectionView.setColorOptionOnClick(selectedOption -> {
-            if (mSelectedColor.equals(selectedOption)) {
-                return null;
-            }
-            mSelectedColor = (ColorOption) selectedOption;
-            // Post with delay for color option to run ripple.
-            new Handler().postDelayed(()-> applyColor(mSelectedColor), /* delayMillis= */ 100);
-            return null;
-        });
-        mColorSectionView.setItems(subOptions, mColorManager);
-    }
-
-    private ColorOption findActiveColorOption(List<ColorOption> wallpaperColorOptions,
-            List<ColorOption> presetColorOptions) {
-        ColorOption activeColorOption = null;
-        for (ColorOption colorOption : Lists.newArrayList(
-                Iterables.concat(wallpaperColorOptions, presetColorOptions))) {
-            if (colorOption.isActive(mColorManager)) {
-                activeColorOption = colorOption;
-                break;
-            }
-        }
-        // Use the first one option by default. This should not happen as above should have an
-        // active option found.
-        if (activeColorOption == null) {
-            activeColorOption = wallpaperColorOptions.isEmpty()
-                    ? presetColorOptions.get(0)
-                    : wallpaperColorOptions.get(0);
-        }
-        return activeColorOption;
-    }
-
-    private void applyColor(ColorOption colorOption) {
-        if (SystemClock.elapsedRealtime() - mLastColorApplyingTime < MIN_COLOR_APPLY_PERIOD) {
-            return;
-        }
-        mLastColorApplyingTime = SystemClock.elapsedRealtime();
-        mColorManager.apply(colorOption, new CustomizationManager.Callback() {
-            @Override
-            public void onSuccess() {
-                mColorSectionView.announceForAccessibility(
-                        mColorSectionView.getContext().getString(R.string.color_changed));
-                mEventLogger.logColorApplied(getColorAction(colorOption), colorOption);
-            }
-
-            @Override
-            public void onError(@Nullable Throwable throwable) {
-                Log.w(TAG, "Apply theme with error: " + throwable);
-            }
-        });
-    }
-
-    private int getColorAction(ColorOption colorOption) {
-        int action = StyleEnums.DEFAULT_ACTION;
-        boolean isForBoth = mLockWallpaperColors == null || mLockWallpaperColors.equals(
-                mHomeWallpaperColors);
-
-        if (TextUtils.equals(colorOption.getSource(), COLOR_SOURCE_PRESET)) {
-            action = StyleEnums.COLOR_PRESET_APPLIED;
-        } else if (isForBoth) {
-            action = StyleEnums.COLOR_WALLPAPER_HOME_LOCK_APPLIED;
-        } else {
-            switch (colorOption.getSource()) {
-                case COLOR_SOURCE_HOME:
-                    action = StyleEnums.COLOR_WALLPAPER_HOME_APPLIED;
-                    break;
-                case COLOR_SOURCE_LOCK:
-                    action = StyleEnums.COLOR_WALLPAPER_LOCK_APPLIED;
-                    break;
-            }
-        }
-        return action;
-    }
-}
diff --git a/src/com/android/customization/model/mode/DarkModeSectionController.java b/src/com/android/customization/model/mode/DarkModeSectionController.java
index f56b709..ebeaa56 100644
--- a/src/com/android/customization/model/mode/DarkModeSectionController.java
+++ b/src/com/android/customization/model/mode/DarkModeSectionController.java
@@ -59,12 +59,17 @@
 
     private Context mContext;
     private DarkModeSectionView mDarkModeSectionView;
+    private final DarkModeSnapshotRestorer mSnapshotRestorer;
 
-    public DarkModeSectionController(Context context, Lifecycle lifecycle) {
+    public DarkModeSectionController(
+            Context context,
+            Lifecycle lifecycle,
+            DarkModeSnapshotRestorer snapshotRestorer) {
         mContext = context;
         mLifecycle = lifecycle;
         mPowerManager = context.getSystemService(PowerManager.class);
         mLifecycle.addObserver(this);
+        mSnapshotRestorer = snapshotRestorer;
     }
 
     @OnLifecycleEvent(Lifecycle.Event.ON_START)
@@ -132,6 +137,7 @@
                     mDarkModeSectionView.announceForAccessibility(
                             context.getString(R.string.mode_changed));
                     uiModeManager.setNightModeActivated(viewActivated);
+                    mSnapshotRestorer.store(viewActivated);
                 },
                 /* delayMillis= */ shortDelay);
     }
diff --git a/src/com/android/customization/model/mode/DarkModeSnapshotRestorer.kt b/src/com/android/customization/model/mode/DarkModeSnapshotRestorer.kt
new file mode 100644
index 0000000..aa8d97c
--- /dev/null
+++ b/src/com/android/customization/model/mode/DarkModeSnapshotRestorer.kt
@@ -0,0 +1,101 @@
+/*
+ * 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.mode
+
+import android.app.UiModeManager
+import android.content.Context
+import android.content.res.Configuration
+import androidx.annotation.VisibleForTesting
+import com.android.wallpaper.picker.undo.domain.interactor.SnapshotRestorer
+import com.android.wallpaper.picker.undo.domain.interactor.SnapshotStore
+import com.android.wallpaper.picker.undo.shared.model.RestorableSnapshot
+import kotlinx.coroutines.CoroutineDispatcher
+import kotlinx.coroutines.withContext
+
+class DarkModeSnapshotRestorer : SnapshotRestorer {
+
+    private val backgroundDispatcher: CoroutineDispatcher
+    private val isActive: () -> Boolean
+    private val setActive: suspend (Boolean) -> Unit
+
+    private lateinit var store: SnapshotStore
+
+    constructor(
+        context: Context,
+        manager: UiModeManager,
+        backgroundDispatcher: CoroutineDispatcher,
+    ) : this(
+        backgroundDispatcher = backgroundDispatcher,
+        isActive = {
+            context.applicationContext.resources.configuration.uiMode and
+                Configuration.UI_MODE_NIGHT_YES != 0
+        },
+        setActive = { isActive -> manager.setNightModeActivated(isActive) },
+    )
+
+    @VisibleForTesting
+    constructor(
+        backgroundDispatcher: CoroutineDispatcher,
+        isActive: () -> Boolean,
+        setActive: suspend (Boolean) -> Unit,
+    ) {
+        this.backgroundDispatcher = backgroundDispatcher
+        this.isActive = isActive
+        this.setActive = setActive
+    }
+
+    override suspend fun setUpSnapshotRestorer(store: SnapshotStore): RestorableSnapshot {
+        this.store = store
+        return snapshot(
+            isActivated = isActive(),
+        )
+    }
+
+    override suspend fun restoreToSnapshot(snapshot: RestorableSnapshot) {
+        val isActivated = snapshot.args[KEY]?.toBoolean() == true
+        withContext(backgroundDispatcher) { setActive(isActivated) }
+    }
+
+    fun store(
+        isActivated: Boolean,
+    ) {
+        store.store(
+            snapshot(
+                isActivated = isActivated,
+            ),
+        )
+    }
+
+    private fun snapshot(
+        isActivated: Boolean,
+    ): RestorableSnapshot {
+        return RestorableSnapshot(
+            args =
+                buildMap {
+                    put(
+                        KEY,
+                        isActivated.toString(),
+                    )
+                }
+        )
+    }
+
+    companion object {
+        private const val KEY = "is_activated"
+    }
+}
diff --git a/src/com/android/customization/model/themedicon/ThemedIconSectionController.java b/src/com/android/customization/model/themedicon/ThemedIconSectionController.java
index a1623d1..5d551a6 100644
--- a/src/com/android/customization/model/themedicon/ThemedIconSectionController.java
+++ b/src/com/android/customization/model/themedicon/ThemedIconSectionController.java
@@ -20,11 +20,13 @@
 import android.view.LayoutInflater;
 
 import androidx.annotation.Nullable;
+import androidx.lifecycle.Observer;
 
+import com.android.customization.model.themedicon.domain.interactor.ThemedIconInteractor;
+import com.android.customization.model.themedicon.domain.interactor.ThemedIconSnapshotRestorer;
 import com.android.customization.picker.themedicon.ThemedIconSectionView;
 import com.android.wallpaper.R;
 import com.android.wallpaper.model.CustomizationSectionController;
-import com.android.wallpaper.model.WorkspaceViewModel;
 
 /** The {@link CustomizationSectionController} for themed icon section. */
 public class ThemedIconSectionController implements
@@ -33,16 +35,26 @@
     private static final String KEY_THEMED_ICON_ENABLED = "SAVED_THEMED_ICON_ENABLED";
 
     private final ThemedIconSwitchProvider mThemedIconOptionsProvider;
-    private final WorkspaceViewModel mWorkspaceViewModel;
+    private final ThemedIconInteractor mInteractor;
+    private final ThemedIconSnapshotRestorer mSnapshotRestorer;
+    private final Observer<Boolean> mIsActivatedChangeObserver;
 
     private ThemedIconSectionView mThemedIconSectionView;
     private boolean mSavedThemedIconEnabled = false;
 
-
-    public ThemedIconSectionController(ThemedIconSwitchProvider themedIconOptionsProvider,
-            WorkspaceViewModel workspaceViewModel, @Nullable Bundle savedInstanceState) {
+    public ThemedIconSectionController(
+            ThemedIconSwitchProvider themedIconOptionsProvider,
+            ThemedIconInteractor interactor,
+            @Nullable Bundle savedInstanceState,
+            ThemedIconSnapshotRestorer snapshotRestorer) {
         mThemedIconOptionsProvider = themedIconOptionsProvider;
-        mWorkspaceViewModel = workspaceViewModel;
+        mInteractor = interactor;
+        mSnapshotRestorer = snapshotRestorer;
+        mIsActivatedChangeObserver = isActivated -> {
+            if (mThemedIconSectionView.isAttachedToWindow()) {
+                mThemedIconSectionView.getSwitch().setChecked(isActivated);
+            }
+        };
 
         if (savedInstanceState != null) {
             mSavedThemedIconEnabled = savedInstanceState.getBoolean(
@@ -64,15 +76,22 @@
         mThemedIconSectionView.getSwitch().setChecked(mSavedThemedIconEnabled);
         mThemedIconOptionsProvider.fetchThemedIconEnabled(
                 enabled -> mThemedIconSectionView.getSwitch().setChecked(enabled));
+        mInteractor.isActivatedAsLiveData().observeForever(mIsActivatedChangeObserver);
         return mThemedIconSectionView;
     }
 
+    @Override
+    public void release() {
+        mInteractor.isActivatedAsLiveData().removeObserver(mIsActivatedChangeObserver);
+    }
+
     private void onViewActivated(Context context, boolean viewActivated) {
         if (context == null) {
             return;
         }
         mThemedIconOptionsProvider.setThemedIconEnabled(viewActivated);
-        mWorkspaceViewModel.getUpdateWorkspace().setValue(viewActivated);
+        mInteractor.setActivated(viewActivated);
+        mSnapshotRestorer.store(viewActivated);
     }
 
     @Override
diff --git a/src/com/android/customization/model/themedicon/ThemedIconSwitchProvider.java b/src/com/android/customization/model/themedicon/ThemedIconSwitchProvider.java
index 9acd319..5e2a60a 100644
--- a/src/com/android/customization/model/themedicon/ThemedIconSwitchProvider.java
+++ b/src/com/android/customization/model/themedicon/ThemedIconSwitchProvider.java
@@ -118,7 +118,7 @@
      *
      * <p>The value would also be stored in SharedPreferences.
      */
-    protected void setThemedIconEnabled(boolean enabled) {
+    public void setThemedIconEnabled(boolean enabled) {
         mExecutorService.submit(() -> {
             ContentValues values = new ContentValues();
             values.put(COL_ICON_THEMED_VALUE, enabled);
diff --git a/src/com/android/customization/model/themedicon/data/repository/ThemedIconRepository.kt b/src/com/android/customization/model/themedicon/data/repository/ThemedIconRepository.kt
new file mode 100644
index 0000000..9108811
--- /dev/null
+++ b/src/com/android/customization/model/themedicon/data/repository/ThemedIconRepository.kt
@@ -0,0 +1,30 @@
+/*
+ * 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.themedicon.data.repository
+
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.asStateFlow
+
+class ThemeIconRepository {
+    private val _isActivated = MutableStateFlow(false)
+    val isActivated = _isActivated.asStateFlow()
+
+    fun setActivated(isActivated: Boolean) {
+        _isActivated.value = isActivated
+    }
+}
diff --git a/src/com/android/customization/model/themedicon/domain/interactor/ThemedIconInteractor.kt b/src/com/android/customization/model/themedicon/domain/interactor/ThemedIconInteractor.kt
new file mode 100644
index 0000000..1cfe877
--- /dev/null
+++ b/src/com/android/customization/model/themedicon/domain/interactor/ThemedIconInteractor.kt
@@ -0,0 +1,38 @@
+/*
+ * 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.themedicon.domain.interactor
+
+import androidx.lifecycle.LiveData
+import androidx.lifecycle.asLiveData
+import com.android.customization.model.themedicon.data.repository.ThemeIconRepository
+
+class ThemedIconInteractor(
+    private val repository: ThemeIconRepository,
+) {
+    val isActivated = repository.isActivated
+
+    private var isActivatedAsLiveData: LiveData<Boolean>? = null
+
+    fun isActivatedAsLiveData(): LiveData<Boolean> {
+        return isActivatedAsLiveData ?: isActivated.asLiveData().also { isActivatedAsLiveData = it }
+    }
+
+    fun setActivated(isActivated: Boolean) {
+        repository.setActivated(isActivated)
+    }
+}
diff --git a/src/com/android/customization/model/themedicon/domain/interactor/ThemedIconSnapshotRestorer.kt b/src/com/android/customization/model/themedicon/domain/interactor/ThemedIconSnapshotRestorer.kt
new file mode 100644
index 0000000..639a92d
--- /dev/null
+++ b/src/com/android/customization/model/themedicon/domain/interactor/ThemedIconSnapshotRestorer.kt
@@ -0,0 +1,60 @@
+/*
+ * 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.themedicon.domain.interactor
+
+import com.android.wallpaper.picker.undo.domain.interactor.SnapshotRestorer
+import com.android.wallpaper.picker.undo.domain.interactor.SnapshotStore
+import com.android.wallpaper.picker.undo.shared.model.RestorableSnapshot
+
+class ThemedIconSnapshotRestorer(
+    private val isActivated: () -> Boolean,
+    private val setActivated: (isActivated: Boolean) -> Unit,
+    private val interactor: ThemedIconInteractor,
+) : SnapshotRestorer {
+
+    private lateinit var store: SnapshotStore
+
+    override suspend fun setUpSnapshotRestorer(store: SnapshotStore): RestorableSnapshot {
+        this.store = store
+        return snapshot()
+    }
+
+    override suspend fun restoreToSnapshot(snapshot: RestorableSnapshot) {
+        val isActivated = snapshot.args[KEY]?.toBoolean() == true
+        setActivated(isActivated)
+        interactor.setActivated(isActivated)
+    }
+
+    fun store(
+        isActivated: Boolean,
+    ) {
+        store.store(snapshot(isActivated = isActivated))
+    }
+
+    private fun snapshot(
+        isActivated: Boolean? = null,
+    ): RestorableSnapshot {
+        return RestorableSnapshot(
+            args = buildMap { put(KEY, (isActivated ?: isActivated()).toString()) }
+        )
+    }
+
+    companion object {
+        private const val KEY = "is_activated"
+    }
+}
diff --git a/src/com/android/customization/module/DefaultCustomizationSections.java b/src/com/android/customization/module/DefaultCustomizationSections.java
index 16d4bc7..487bca8 100644
--- a/src/com/android/customization/module/DefaultCustomizationSections.java
+++ b/src/com/android/customization/module/DefaultCustomizationSections.java
@@ -8,13 +8,17 @@
 import androidx.lifecycle.ViewModelProvider;
 
 import com.android.customization.model.color.ColorSectionController;
-import com.android.customization.model.color.ColorSectionController2;
 import com.android.customization.model.grid.GridOptionsManager;
 import com.android.customization.model.grid.GridSectionController;
 import com.android.customization.model.mode.DarkModeSectionController;
+import com.android.customization.model.mode.DarkModeSnapshotRestorer;
 import com.android.customization.model.themedicon.ThemedIconSectionController;
 import com.android.customization.model.themedicon.ThemedIconSwitchProvider;
+import com.android.customization.model.themedicon.domain.interactor.ThemedIconInteractor;
+import com.android.customization.model.themedicon.domain.interactor.ThemedIconSnapshotRestorer;
 import com.android.customization.picker.clock.data.repository.ClockRegistryProvider;
+import com.android.customization.picker.color.ui.section.ColorSectionController2;
+import com.android.customization.picker.color.ui.viewmodel.ColorPickerViewModel;
 import com.android.customization.picker.notifications.ui.section.NotificationSectionController;
 import com.android.customization.picker.notifications.ui.viewmodel.NotificationSectionViewModel;
 import com.android.customization.picker.preview.ui.section.PreviewWithClockCarouselSectionController;
@@ -31,7 +35,6 @@
 import com.android.wallpaper.model.WallpaperColorsViewModel;
 import com.android.wallpaper.model.WallpaperPreviewNavigator;
 import com.android.wallpaper.model.WallpaperSectionController;
-import com.android.wallpaper.model.WorkspaceViewModel;
 import com.android.wallpaper.module.CurrentWallpaperInfoFactory;
 import com.android.wallpaper.module.CustomizationSections;
 import com.android.wallpaper.picker.customization.ui.section.ConnectedSectionController;
@@ -46,6 +49,7 @@
 /** {@link CustomizationSections} for the customization picker. */
 public final class DefaultCustomizationSections implements CustomizationSections {
 
+    private final ColorPickerViewModel.Factory mColorPickerViewModelFactory;
     private final KeyguardQuickAffordancePickerInteractor mKeyguardQuickAffordancePickerInteractor;
     private final KeyguardQuickAffordancePickerViewModel.Factory
             mKeyguardQuickAffordancePickerViewModelFactory;
@@ -56,8 +60,12 @@
             mClockCarouselViewModelProvider;
     private final PreviewWithClockCarouselSectionController.ClockViewFactoryProvider
             mClockViewFactoryProvider;
+    private final DarkModeSnapshotRestorer mDarkModeSnapshotRestorer;
+    private final ThemedIconSnapshotRestorer mThemedIconSnapshotRestorer;
+    private final ThemedIconInteractor mThemedIconInteractor;
 
     public DefaultCustomizationSections(
+            ColorPickerViewModel.Factory colorPickerViewModelFactory,
             KeyguardQuickAffordancePickerInteractor keyguardQuickAffordancePickerInteractor,
             KeyguardQuickAffordancePickerViewModel.Factory
                     keyguardQuickAffordancePickerViewModelFactory,
@@ -65,7 +73,11 @@
             BaseFlags flags,
             ClockRegistryProvider clockRegistryProvider,
             ClockCarouselViewModelProvider clockCarouselViewModelProvider,
-            ClockViewFactoryProvider clockViewFactoryProvider) {
+            ClockViewFactoryProvider clockViewFactoryProvider,
+            DarkModeSnapshotRestorer darkModeSnapshotRestorer,
+            ThemedIconSnapshotRestorer themedIconSnapshotRestorer,
+            ThemedIconInteractor themedIconInteractor) {
+        mColorPickerViewModelFactory = colorPickerViewModelFactory;
         mKeyguardQuickAffordancePickerInteractor = keyguardQuickAffordancePickerInteractor;
         mKeyguardQuickAffordancePickerViewModelFactory =
                 keyguardQuickAffordancePickerViewModelFactory;
@@ -74,6 +86,9 @@
         mClockRegistryProvider = clockRegistryProvider;
         mClockCarouselViewModelProvider = clockCarouselViewModelProvider;
         mClockViewFactoryProvider = clockViewFactoryProvider;
+        mDarkModeSnapshotRestorer = darkModeSnapshotRestorer;
+        mThemedIconSnapshotRestorer = themedIconSnapshotRestorer;
+        mThemedIconInteractor = themedIconInteractor;
     }
 
     @Override
@@ -82,7 +97,6 @@
             FragmentActivity activity,
             LifecycleOwner lifecycleOwner,
             WallpaperColorsViewModel wallpaperColorsViewModel,
-            WorkspaceViewModel workspaceViewModel,
             PermissionRequester permissionRequester,
             WallpaperPreviewNavigator wallpaperPreviewNavigator,
             CustomizationSectionNavigationController sectionNavigationController,
@@ -117,10 +131,12 @@
                 new ConnectedSectionController(
                         // Theme color section.
                         new ColorSectionController2(
-                                activity,
-                                wallpaperColorsViewModel,
-                                lifecycleOwner,
-                                sectionNavigationController),
+                                sectionNavigationController,
+                                new ViewModelProvider(
+                                        activity,
+                                        mColorPickerViewModelFactory)
+                                        .get(ColorPickerViewModel.class),
+                                lifecycleOwner),
                         // Wallpaper quick switch section.
                         new WallpaperQuickSwitchSectionController(
                                 screen,
@@ -157,13 +173,18 @@
 
             case HOME_SCREEN:
                 // Dark/Light theme section.
-                sectionControllers.add(new DarkModeSectionController(activity,
-                        lifecycleOwner.getLifecycle()));
+                sectionControllers.add(new DarkModeSectionController(
+                        activity,
+                        lifecycleOwner.getLifecycle(),
+                        mDarkModeSnapshotRestorer));
 
                 // Themed app icon section.
-                sectionControllers.add(new ThemedIconSectionController(
-                        ThemedIconSwitchProvider.getInstance(activity), workspaceViewModel,
-                        savedInstanceState));
+                sectionControllers.add(
+                        new ThemedIconSectionController(
+                                ThemedIconSwitchProvider.getInstance(activity),
+                                mThemedIconInteractor,
+                                savedInstanceState,
+                                mThemedIconSnapshotRestorer));
 
                 // App grid section.
                 sectionControllers.add(new GridSectionController(
@@ -179,7 +200,6 @@
             FragmentActivity activity,
             LifecycleOwner lifecycleOwner,
             WallpaperColorsViewModel wallpaperColorsViewModel,
-            WorkspaceViewModel workspaceViewModel,
             PermissionRequester permissionRequester,
             WallpaperPreviewNavigator wallpaperPreviewNavigator,
             CustomizationSectionNavigationController sectionNavigationController,
@@ -188,23 +208,34 @@
         List<CustomizationSectionController<?>> sectionControllers = new ArrayList<>();
 
         // Wallpaper section.
-        sectionControllers.add(new WallpaperSectionController(
-                activity, lifecycleOwner, permissionRequester, wallpaperColorsViewModel,
-                workspaceViewModel, sectionNavigationController, wallpaperPreviewNavigator,
-                savedInstanceState, displayUtils));
+        sectionControllers.add(
+                new WallpaperSectionController(
+                activity,
+                lifecycleOwner,
+                permissionRequester,
+                wallpaperColorsViewModel,
+                mThemedIconInteractor.isActivatedAsLiveData(),
+                sectionNavigationController,
+                wallpaperPreviewNavigator,
+                savedInstanceState,
+                displayUtils));
 
         // Theme color section.
         sectionControllers.add(new ColorSectionController(
                 activity, wallpaperColorsViewModel, lifecycleOwner, savedInstanceState));
 
         // Dark/Light theme section.
-        sectionControllers.add(new DarkModeSectionController(activity,
-                lifecycleOwner.getLifecycle()));
+        sectionControllers.add(new DarkModeSectionController(
+                activity,
+                lifecycleOwner.getLifecycle(),
+                mDarkModeSnapshotRestorer));
 
         // Themed app icon section.
         sectionControllers.add(new ThemedIconSectionController(
-                ThemedIconSwitchProvider.getInstance(activity), workspaceViewModel,
-                savedInstanceState));
+                ThemedIconSwitchProvider.getInstance(activity),
+                mThemedIconInteractor,
+                savedInstanceState,
+                mThemedIconSnapshotRestorer));
 
         // App grid section.
         sectionControllers.add(new GridSectionController(
diff --git a/src/com/android/customization/module/ThemePickerInjector.kt b/src/com/android/customization/module/ThemePickerInjector.kt
index c2a6565..56ef2e6 100644
--- a/src/com/android/customization/module/ThemePickerInjector.kt
+++ b/src/com/android/customization/module/ThemePickerInjector.kt
@@ -15,6 +15,7 @@
  */
 package com.android.customization.module
 
+import android.app.UiModeManager
 import android.content.Context
 import android.content.Intent
 import android.net.Uri
@@ -22,9 +23,15 @@
 import androidx.activity.ComponentActivity
 import androidx.fragment.app.Fragment
 import androidx.fragment.app.FragmentActivity
+import androidx.lifecycle.ViewModelProvider
+import com.android.customization.model.mode.DarkModeSnapshotRestorer
 import com.android.customization.model.theme.OverlayManagerCompat
 import com.android.customization.model.theme.ThemeBundleProvider
 import com.android.customization.model.theme.ThemeManager
+import com.android.customization.model.themedicon.ThemedIconSwitchProvider
+import com.android.customization.model.themedicon.data.repository.ThemeIconRepository
+import com.android.customization.model.themedicon.domain.interactor.ThemedIconInteractor
+import com.android.customization.model.themedicon.domain.interactor.ThemedIconSnapshotRestorer
 import com.android.customization.picker.clock.data.repository.ClockPickerRepositoryImpl
 import com.android.customization.picker.clock.data.repository.ClockRegistryProvider
 import com.android.customization.picker.clock.domain.interactor.ClockPickerInteractor
@@ -86,10 +93,18 @@
     private var notificationSectionViewModelFactory: NotificationSectionViewModel.Factory? = null
     private var colorPickerInteractor: ColorPickerInteractor? = null
     private var colorPickerViewModelFactory: ColorPickerViewModel.Factory? = null
+    private var darkModeSnapshotRestorer: DarkModeSnapshotRestorer? = null
+    private var themedIconSnapshotRestorer: ThemedIconSnapshotRestorer? = null
+    private var themedIconInteractor: ThemedIconInteractor? = null
 
     override fun getCustomizationSections(activity: ComponentActivity): CustomizationSections {
         return customizationSections
             ?: DefaultCustomizationSections(
+                    getColorPickerViewModelFactory(
+                        context = activity,
+                        wallpaperColorsViewModel =
+                            ViewModelProvider(activity)[WallpaperColorsViewModel::class.java],
+                    ),
                     getKeyguardQuickAffordancePickerInteractor(activity),
                     getKeyguardQuickAffordancePickerViewModelFactory(activity),
                     NotificationSectionViewModel.Factory(
@@ -112,7 +127,10 @@
                                 registry = registry,
                             )
                         }
-                    }
+                    },
+                    getDarkModeSnapshotRestorer(activity),
+                    getThemedIconSnapshotRestorer(activity),
+                    getThemedIconInteractor(),
                 )
                 .also { customizationSections = it }
     }
@@ -173,6 +191,8 @@
                 getKeyguardQuickAffordanceSnapshotRestorer(context)
             this[KEY_WALLPAPER_SNAPSHOT_RESTORER] = getWallpaperSnapshotRestorer(context)
             this[KEY_NOTIFICATIONS_SNAPSHOT_RESTORER] = getNotificationsSnapshotRestorer(context)
+            this[KEY_DARK_MODE_SNAPSHOT_RESTORER] = getDarkModeSnapshotRestorer(context)
+            this[KEY_THEMED_ICON_SNAPSHOT_RESTORER] = getThemedIconSnapshotRestorer(context)
         }
     }
 
@@ -348,6 +368,41 @@
                 .also { colorPickerViewModelFactory = it }
     }
 
+    protected fun getDarkModeSnapshotRestorer(
+        context: Context,
+    ): DarkModeSnapshotRestorer {
+        return darkModeSnapshotRestorer
+            ?: DarkModeSnapshotRestorer(
+                    context = context,
+                    manager = context.getSystemService(Context.UI_MODE_SERVICE) as UiModeManager,
+                    backgroundDispatcher = Dispatchers.IO,
+                )
+                .also { darkModeSnapshotRestorer = it }
+    }
+
+    protected fun getThemedIconSnapshotRestorer(
+        context: Context,
+    ): ThemedIconSnapshotRestorer {
+        val optionProvider = ThemedIconSwitchProvider.getInstance(context)
+        return themedIconSnapshotRestorer
+            ?: ThemedIconSnapshotRestorer(
+                    isActivated = { optionProvider.isThemedIconEnabled },
+                    setActivated = { isActivated ->
+                        optionProvider.isThemedIconEnabled = isActivated
+                    },
+                    interactor = getThemedIconInteractor(),
+                )
+                .also { themedIconSnapshotRestorer = it }
+    }
+
+    protected fun getThemedIconInteractor(): ThemedIconInteractor {
+        return themedIconInteractor
+            ?: ThemedIconInteractor(
+                    repository = ThemeIconRepository(),
+                )
+                .also { themedIconInteractor = it }
+    }
+
     companion object {
         @JvmStatic
         private val KEY_QUICK_AFFORDANCE_SNAPSHOT_RESTORER =
@@ -356,11 +411,15 @@
         private val KEY_WALLPAPER_SNAPSHOT_RESTORER = KEY_QUICK_AFFORDANCE_SNAPSHOT_RESTORER + 1
         @JvmStatic
         private val KEY_NOTIFICATIONS_SNAPSHOT_RESTORER = KEY_WALLPAPER_SNAPSHOT_RESTORER + 1
+        @JvmStatic
+        private val KEY_DARK_MODE_SNAPSHOT_RESTORER = KEY_NOTIFICATIONS_SNAPSHOT_RESTORER + 1
+        @JvmStatic
+        private val KEY_THEMED_ICON_SNAPSHOT_RESTORER = KEY_DARK_MODE_SNAPSHOT_RESTORER + 1
 
         /**
          * When this injector is overridden, this is the minimal value that should be used by
          * restorers returns in [getSnapshotRestorers].
          */
-        @JvmStatic protected val MIN_SNAPSHOT_RESTORER_KEY = KEY_NOTIFICATIONS_SNAPSHOT_RESTORER + 1
+        @JvmStatic protected val MIN_SNAPSHOT_RESTORER_KEY = KEY_THEMED_ICON_SNAPSHOT_RESTORER + 1
     }
 }
diff --git a/src/com/android/customization/picker/clock/ui/fragment/ClockSettingsFragment.kt b/src/com/android/customization/picker/clock/ui/fragment/ClockSettingsFragment.kt
index e3e2785..d716102 100644
--- a/src/com/android/customization/picker/clock/ui/fragment/ClockSettingsFragment.kt
+++ b/src/com/android/customization/picker/clock/ui/fragment/ClockSettingsFragment.kt
@@ -24,7 +24,9 @@
 import androidx.lifecycle.get
 import androidx.lifecycle.lifecycleScope
 import com.android.customization.module.ThemePickerInjector
+import com.android.customization.picker.clock.ui.binder.ClockCarouselViewBinder
 import com.android.customization.picker.clock.ui.binder.ClockSettingsBinder
+import com.android.customization.picker.clock.ui.view.ClockCarouselView
 import com.android.customization.picker.clock.ui.viewmodel.ClockSettingsViewModel
 import com.android.systemui.shared.clocks.shared.model.ClockPreviewConstants
 import com.android.wallpaper.R
@@ -115,14 +117,23 @@
             )
             .show()
 
+        val carouselView: ClockCarouselView = view.requireViewById(R.id.clock_carousel_view)
         lifecycleScope.launch {
-            val clockRegistry =
+            val registry =
                 withContext(Dispatchers.IO) { injector.getClockRegistryProvider(context).get() }
+            val clockViewFactory = injector.getClockViewFactory(context, registry)
+            ClockCarouselViewBinder.bind(
+                    view = carouselView,
+                    viewModel = injector.getClockCarouselViewModel(context, registry),
+                    clockViewFactory = { clockId -> clockViewFactory.getView(clockId) },
+                    lifecycleOwner = this@ClockSettingsFragment,
+                )
+                .show()
             ClockSettingsBinder.bind(
                 view,
                 ClockSettingsViewModel(
                     context,
-                    injector.getClockPickerInteractor(context, clockRegistry)
+                    injector.getClockPickerInteractor(context, registry)
                 ),
                 this@ClockSettingsFragment,
             )
diff --git a/src/com/android/customization/picker/clock/ui/section/ClockSectionController.kt b/src/com/android/customization/picker/clock/ui/section/ClockSectionController.kt
index bd81719..e1982d7 100644
--- a/src/com/android/customization/picker/clock/ui/section/ClockSectionController.kt
+++ b/src/com/android/customization/picker/clock/ui/section/ClockSectionController.kt
@@ -22,7 +22,7 @@
 import com.android.customization.module.ThemePickerInjector
 import com.android.customization.picker.clock.data.repository.ClockRegistryProvider
 import com.android.customization.picker.clock.ui.binder.ClockSectionViewBinder
-import com.android.customization.picker.clock.ui.fragment.ClockCustomDemoFragment
+import com.android.customization.picker.clock.ui.fragment.ClockSettingsFragment
 import com.android.customization.picker.clock.ui.view.ClockSectionView
 import com.android.wallpaper.R
 import com.android.wallpaper.config.BaseFlags
@@ -60,7 +60,7 @@
                     (InjectorProvider.getInjector() as ThemePickerInjector)
                         .getClockSectionViewModel(context, registry),
                 lifecycleOwner = lifecycleOwner
-            ) { navigationController.navigateTo(ClockCustomDemoFragment()) }
+            ) { navigationController.navigateTo(ClockSettingsFragment()) }
         }
         return view
     }
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 e163c15..4ce5ed9 100644
--- a/src/com/android/customization/picker/color/data/repository/ColorPickerRepositoryImpl.kt
+++ b/src/com/android/customization/picker/color/data/repository/ColorPickerRepositoryImpl.kt
@@ -31,8 +31,11 @@
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.combine
+import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.suspendCancellableCoroutine
 
 // TODO (b/262924623): refactor to remove dependency on ColorCustomizationManager & ColorOption
+// TODO (b/268203200): Create test for ColorPickerRepositoryImpl
 class ColorPickerRepositoryImpl(
     context: Context,
     wallpaperColorsViewModel: WallpaperColorsViewModel,
@@ -48,31 +51,48 @@
     /** List of wallpaper and preset color options on the device, categorized by Color Type */
     override val colorOptions: Flow<Map<ColorType, List<ColorOptionModel>>> =
         combine(homeWallpaperColors, lockWallpaperColors) { homeColors, lockColors ->
-            colorManager.setWallpaperColors(homeColors, lockColors)
-            val wallpaperColorOptions: MutableList<ColorOptionModel> = mutableListOf()
-            val presetColorOptions: MutableList<ColorOptionModel> = mutableListOf()
-            colorManager.fetchOptions(
-                object : CustomizationManager.OptionsFetchedListener<ColorOption?> {
-                    override fun onOptionsLoaded(options: MutableList<ColorOption?>?) {
-                        options?.forEach { option ->
-                            when (option) {
-                                is ColorSeedOption -> wallpaperColorOptions.add(option.toModel())
-                                is ColorBundle -> presetColorOptions.add(option.toModel())
+                homeColors to lockColors
+            }
+            .map { (homeColors, lockColors) ->
+                suspendCancellableCoroutine { continuation ->
+                    colorManager.setWallpaperColors(homeColors, lockColors)
+                    colorManager.fetchOptions(
+                        object : CustomizationManager.OptionsFetchedListener<ColorOption?> {
+                            override fun onOptionsLoaded(options: MutableList<ColorOption?>?) {
+                                val wallpaperColorOptions: MutableList<ColorOptionModel> =
+                                    mutableListOf()
+                                val presetColorOptions: MutableList<ColorOptionModel> =
+                                    mutableListOf()
+                                options?.forEach { option ->
+                                    when (option) {
+                                        is ColorSeedOption ->
+                                            wallpaperColorOptions.add(option.toModel())
+                                        is ColorBundle -> presetColorOptions.add(option.toModel())
+                                    }
+                                }
+                                continuation.resumeWith(
+                                    Result.success(
+                                        mapOf(
+                                            ColorType.WALLPAPER_COLOR to wallpaperColorOptions,
+                                            ColorType.BASIC_COLOR to presetColorOptions
+                                        )
+                                    )
+                                )
                             }
-                        }
-                    }
 
-                    override fun onError(throwable: Throwable?) {
-                        Log.e("ColorPickerRepository", "Error loading theme bundles", throwable)
-                    }
-                },
-                /* reload= */ false
-            )
-            mapOf(
-                ColorType.WALLPAPER_COLOR to wallpaperColorOptions,
-                ColorType.BASIC_COLOR to presetColorOptions
-            )
-        }
+                            override fun onError(throwable: Throwable?) {
+                                Log.e(TAG, "Error loading theme bundles", throwable)
+                                continuation.resumeWith(
+                                    Result.failure(
+                                        throwable ?: Throwable("Error loading theme bundles")
+                                    )
+                                )
+                            }
+                        },
+                        /* reload= */ false
+                    )
+                }
+            }
 
     override fun select(colorOptionModel: ColorOptionModel) {
         val colorOption: ColorOption = colorOptionModel.colorOption
@@ -82,7 +102,7 @@
                 override fun onSuccess() = Unit
 
                 override fun onError(throwable: Throwable?) {
-                    Log.w("ColorPickerRepository", "Apply theme with error", throwable)
+                    Log.w(TAG, "Apply theme with error", throwable)
                 }
             }
         )
@@ -94,4 +114,8 @@
             isSelected = isActive(colorManager),
         )
     }
+
+    companion object {
+        private const val TAG = "ColorPickerRepositoryImpl"
+    }
 }
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 6d8b7dc..331d635 100644
--- a/src/com/android/customization/picker/color/data/repository/FakeColorPickerRepository.kt
+++ b/src/com/android/customization/picker/color/data/repository/FakeColorPickerRepository.kt
@@ -136,13 +136,33 @@
 
     override fun select(colorOptionModel: ColorOptionModel) {
         val colorOptions = _colorOptions.value
-        colorOptions[ColorType.WALLPAPER_COLOR]?.forEach {
-            it.isSelected = (it.testEquals(colorOptionModel))
+        val wallpaperColorOptions = colorOptions[ColorType.WALLPAPER_COLOR]!!
+        val newWallpaperColorOptions = buildList {
+            wallpaperColorOptions.forEach { option ->
+                add(
+                    ColorOptionModel(
+                        colorOption = option.colorOption,
+                        isSelected = option.testEquals(colorOptionModel),
+                    )
+                )
+            }
         }
-        colorOptions[ColorType.BASIC_COLOR]?.forEach {
-            it.isSelected = (it.testEquals(colorOptionModel))
+        val basicColorOptions = colorOptions[ColorType.BASIC_COLOR]!!
+        val newBasicColorOptions = buildList {
+            basicColorOptions.forEach { option ->
+                add(
+                    ColorOptionModel(
+                        colorOption = option.colorOption,
+                        isSelected = option.testEquals(colorOptionModel),
+                    )
+                )
+            }
         }
-        _colorOptions.value = colorOptions
+        _colorOptions.value =
+            mapOf(
+                ColorType.WALLPAPER_COLOR to newWallpaperColorOptions,
+                ColorType.BASIC_COLOR to newBasicColorOptions
+            )
     }
 
     private fun ColorOptionModel.testEquals(other: Any?): Boolean {
diff --git a/src/com/android/customization/picker/color/ui/binder/ColorSectionViewBinder.kt b/src/com/android/customization/picker/color/ui/binder/ColorSectionViewBinder.kt
new file mode 100644
index 0000000..0842870
--- /dev/null
+++ b/src/com/android/customization/picker/color/ui/binder/ColorSectionViewBinder.kt
@@ -0,0 +1,126 @@
+/*
+ * 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.LayoutInflater
+import android.view.View
+import android.widget.ImageView
+import android.widget.LinearLayout
+import androidx.core.view.isVisible
+import androidx.lifecycle.Lifecycle
+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.ColorPickerViewModel
+import com.android.wallpaper.R
+import kotlinx.coroutines.launch
+
+object ColorSectionViewBinder {
+
+    /**
+     * Binds view with view-model for color picker section. The view should include a linear layout
+     * with id [R.id.color_section_option_container]
+     */
+    @JvmStatic
+    fun bind(
+        view: View,
+        viewModel: ColorPickerViewModel,
+        lifecycleOwner: LifecycleOwner,
+        navigationOnClick: (View) -> Unit,
+        isConnectedHorizontallyToOtherSections: Boolean = false,
+    ) {
+        val optionContainer: LinearLayout =
+            view.requireViewById(R.id.color_section_option_container)
+        val moreColorsButton: View = view.requireViewById(R.id.more_colors)
+        if (isConnectedHorizontallyToOtherSections) {
+            moreColorsButton.isVisible = true
+            moreColorsButton.setOnClickListener(navigationOnClick)
+        } else {
+            moreColorsButton.isVisible = false
+        }
+        lifecycleOwner.lifecycleScope.launch {
+            lifecycleOwner.repeatOnLifecycle(Lifecycle.State.STARTED) {
+                launch {
+                    viewModel.colorSectionOptions.collect { colorOptions ->
+                        setOptions(
+                            options = colorOptions,
+                            view = optionContainer,
+                            addOverflowOption = !isConnectedHorizontallyToOtherSections,
+                            overflowOnClick = navigationOnClick
+                        )
+                    }
+                }
+            }
+        }
+    }
+
+    fun setOptions(
+        options: List<ColorOptionViewModel>,
+        view: LinearLayout,
+        addOverflowOption: Boolean = false,
+        overflowOnClick: (View) -> Unit = {}
+    ) {
+        view.removeAllViews()
+        // Color option slot size is the minimum between the color option size and the view column
+        // count. When having an overflow option, a slot is reserved for the overflow option.
+        val colorOptionSlotSize =
+            (if (addOverflowOption) {
+                    minOf(view.weightSum.toInt() - 1, options.size)
+                } else {
+                    minOf(view.weightSum.toInt(), options.size)
+                })
+                .let { if (it < 0) 0 else it }
+        options.subList(0, colorOptionSlotSize).forEach { item ->
+            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)
+
+            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
+                }
+            )
+            view.addView(itemView)
+        }
+        // add overflow option
+        if (addOverflowOption) {
+            val itemView =
+                LayoutInflater.from(view.context)
+                    .inflate(R.layout.color_option_overflow_no_background, view, false)
+            itemView.setOnClickListener(overflowOnClick)
+            view.addView(itemView)
+        }
+    }
+}
diff --git a/src/com/android/customization/picker/color/ui/fragment/ColorPickerFragment.kt b/src/com/android/customization/picker/color/ui/fragment/ColorPickerFragment.kt
index fad7def..0bc22f8 100644
--- a/src/com/android/customization/picker/color/ui/fragment/ColorPickerFragment.kt
+++ b/src/com/android/customization/picker/color/ui/fragment/ColorPickerFragment.kt
@@ -37,6 +37,12 @@
 
 @OptIn(ExperimentalCoroutinesApi::class)
 class ColorPickerFragment : AppbarFragment() {
+    companion object {
+        @JvmStatic
+        fun newInstance(): ColorPickerFragment {
+            return ColorPickerFragment()
+        }
+    }
 
     override fun onCreateView(
         inflater: LayoutInflater,
diff --git a/src/com/android/customization/picker/color/ui/section/ColorSectionController2.kt b/src/com/android/customization/picker/color/ui/section/ColorSectionController2.kt
new file mode 100644
index 0000000..f1c982b
--- /dev/null
+++ b/src/com/android/customization/picker/color/ui/section/ColorSectionController2.kt
@@ -0,0 +1,67 @@
+/*
+ * 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.section
+
+import android.content.Context
+import android.view.LayoutInflater
+import androidx.lifecycle.LifecycleOwner
+import com.android.customization.picker.color.ui.binder.ColorSectionViewBinder
+import com.android.customization.picker.color.ui.fragment.ColorPickerFragment
+import com.android.customization.picker.color.ui.view.ColorSectionView2
+import com.android.customization.picker.color.ui.viewmodel.ColorPickerViewModel
+import com.android.wallpaper.R
+import com.android.wallpaper.model.CustomizationSectionController
+import com.android.wallpaper.model.CustomizationSectionController.CustomizationSectionNavigationController as NavigationController
+
+class ColorSectionController2(
+    private val navigationController: NavigationController,
+    private val viewModel: ColorPickerViewModel,
+    private val lifecycleOwner: LifecycleOwner
+) : CustomizationSectionController<ColorSectionView2> {
+
+    override fun isAvailable(context: Context): Boolean {
+        return true
+    }
+
+    override fun createView(context: Context): ColorSectionView2 {
+        return createView(context, CustomizationSectionController.ViewCreationParams())
+    }
+
+    override fun createView(
+        context: Context,
+        params: CustomizationSectionController.ViewCreationParams
+    ): ColorSectionView2 {
+        @SuppressWarnings("It is fine to inflate with null parent for our need.")
+        val view =
+            LayoutInflater.from(context)
+                .inflate(
+                    R.layout.color_section_view2,
+                    null,
+                ) as ColorSectionView2
+        ColorSectionViewBinder.bind(
+            view = view,
+            viewModel = viewModel,
+            lifecycleOwner = lifecycleOwner,
+            navigationOnClick = {
+                navigationController.navigateTo(ColorPickerFragment.newInstance())
+            },
+            isConnectedHorizontallyToOtherSections = params.isConnectedHorizontallyToOtherSections,
+        )
+        return view
+    }
+}
diff --git a/src/com/android/customization/picker/color/ui/view/ColorSectionView2.kt b/src/com/android/customization/picker/color/ui/view/ColorSectionView2.kt
index 358514e..7a8f21a 100644
--- a/src/com/android/customization/picker/color/ui/view/ColorSectionView2.kt
+++ b/src/com/android/customization/picker/color/ui/view/ColorSectionView2.kt
@@ -17,77 +17,10 @@
 
 import android.content.Context
 import android.util.AttributeSet
-import android.view.LayoutInflater
-import android.widget.ImageView
-import android.widget.LinearLayout
-import com.android.customization.model.color.ColorCustomizationManager
-import com.android.customization.model.color.ColorOption
-import com.android.wallpaper.R
 import com.android.wallpaper.picker.SectionView
 
 /**
  * The class inherits from {@link SectionView} as the view representing the color section of the
  * customization picker. It displays a list of color options and an overflow option.
  */
-class ColorSectionView2(context: Context, attrs: AttributeSet?) : SectionView(context, attrs) {
-
-    private val items = mutableListOf<ColorOption>()
-    private var onClick: ((ColorOption) -> Unit)? = null
-    private var overflowOnClick: (() -> Unit)? = null
-
-    // TODO (b/262924623): make adjustments for large screen
-    fun setItems(items: List<ColorOption>, manager: ColorCustomizationManager) {
-        this.items.clear()
-        this.items.addAll(items)
-        val optionContainer = findViewById<LinearLayout>(R.id.color_section_option_container)
-        optionContainer.removeAllViews()
-        // Last color option is either the last index of the items list, or the second last index
-        // of column count. Save the last column of the option container for the overflow option
-        val lastOptionIndex = minOf(optionContainer.weightSum.toInt() - 2, items.size - 1)
-        if (items.isNotEmpty()) {
-            for (position in 0..lastOptionIndex) {
-                val item = items[position]
-                val itemView =
-                    LayoutInflater.from(context)
-                        .inflate(R.layout.color_option_no_background, optionContainer, false)
-                item.bindThumbnailTile(itemView.findViewById(R.id.option_tile))
-                if (item.isActive(manager)) {
-                    val optionSelectedView = itemView.findViewById<ImageView>(R.id.option_selected)
-                    optionSelectedView.visibility = VISIBLE
-                }
-                itemView.setOnClickListener { onClick?.invoke(item) }
-                optionContainer.addView(itemView)
-            }
-        }
-        // add overflow option
-        val itemView =
-            LayoutInflater.from(context)
-                .inflate(R.layout.color_option_overflow_no_background, optionContainer, false)
-        itemView.setOnClickListener { overflowOnClick?.invoke() }
-        optionContainer.addView(itemView)
-    }
-
-    /** Sets the on click callback for a color option. */
-    fun setColorOptionOnClick(onClick: (ColorOption) -> Unit) {
-        this.onClick = onClick
-        if (items.isNotEmpty()) {
-            val optionContainer = findViewById<LinearLayout>(R.id.color_section_option_container)
-            val lastOptionIndex = minOf(optionContainer.childCount - 2, items.size - 1)
-            for (position in 0..lastOptionIndex) {
-                val item = items[position]
-                val itemView = optionContainer.getChildAt(position)
-                itemView.setOnClickListener { onClick.invoke(item) }
-            }
-        }
-    }
-
-    /** Sets the on click callback for the overflow option. */
-    fun setOverflowOnClick(onClick: () -> Unit) {
-        this.overflowOnClick = onClick
-        val optionContainer = findViewById<LinearLayout>(R.id.color_section_option_container)
-        if (optionContainer.childCount > 0) {
-            val itemView = optionContainer.getChildAt(optionContainer.childCount - 1)
-            itemView.setOnClickListener { onClick.invoke() }
-        }
-    }
-}
+class ColorSectionView2(context: Context, attrs: AttributeSet?) : SectionView(context, attrs)
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 15445fa..7eb5488 100644
--- a/src/com/android/customization/picker/color/ui/viewmodel/ColorPickerViewModel.kt
+++ b/src/com/android/customization/picker/color/ui/viewmodel/ColorPickerViewModel.kt
@@ -22,12 +22,14 @@
 import com.android.customization.model.color.ColorBundle
 import com.android.customization.model.color.ColorSeedOption
 import com.android.customization.picker.color.domain.interactor.ColorPickerInteractor
-import com.android.customization.picker.color.shared.model.ColorOptionModel
 import com.android.customization.picker.color.shared.model.ColorType
 import com.android.wallpaper.R
+import kotlin.math.max
+import kotlin.math.min
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.combine
+import kotlinx.coroutines.flow.map
 
 /** Models UI state for a color picker experience. */
 class ColorPickerViewModel
@@ -70,62 +72,85 @@
                 .toMap()
         }
 
+    /** The list of all available wallpaper colors */
+    private val wallpaperColorOptions: Flow<List<ColorOptionViewModel>> =
+        interactor.colorOptions.map { colorOptions ->
+            colorOptions[ColorType.WALLPAPER_COLOR]!!.map { colorOptionModel ->
+                val colorSeedOption: ColorSeedOption =
+                    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 {
+                            { interactor.select(colorOptionModel) }
+                        }
+                )
+            }
+        }
+
+    /** The list of all available preset colors */
+    private val presetColorOptions: Flow<List<ColorOptionViewModel>> =
+        interactor.colorOptions.map { colorOptions ->
+            colorOptions[ColorType.BASIC_COLOR]!!.map { colorOptionModel ->
+                val colorBundle: ColorBundle = colorOptionModel.colorOption as ColorBundle
+                val primaryColor = colorBundle.previewInfo.resolvePrimaryColor(context.resources)
+                val secondaryColor =
+                    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 {
+                            { interactor.select(colorOptionModel) }
+                        },
+                )
+            }
+        }
+
     /** The list of all available color options for the selected Color Type. */
     val colorOptions: Flow<List<ColorOptionViewModel>> =
-        combine(interactor.colorOptions, selectedColorTypeId) {
-            colorOptions,
+        combine(wallpaperColorOptions, presetColorOptions, selectedColorTypeId) {
+            wallpaperOptions,
+            presetOptions,
             selectedColorTypeIdOrNull ->
-            val selectedColorType: ColorType =
-                selectedColorTypeIdOrNull ?: ColorType.WALLPAPER_COLOR
-            val selectedColorOptions: List<ColorOptionModel> = colorOptions[selectedColorType]!!
-            selectedColorOptions.map { colorOptionModel ->
-                when (selectedColorType) {
-                    ColorType.BASIC_COLOR -> {
-                        val colorBundle: ColorBundle = colorOptionModel.colorOption as ColorBundle
-                        val primaryColor =
-                            colorBundle.previewInfo.resolvePrimaryColor(context.resources)
-                        val secondaryColor =
-                            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 {
-                                    { interactor.select(colorOptionModel) }
-                                },
-                        )
-                    }
-                    ColorType.WALLPAPER_COLOR -> {
-                        val colorSeedOption: ColorSeedOption =
-                            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 {
-                                    { interactor.select(colorOptionModel) }
-                                },
-                        )
-                    }
-                }
+            when (selectedColorTypeIdOrNull ?: ColorType.WALLPAPER_COLOR) {
+                ColorType.WALLPAPER_COLOR -> wallpaperOptions
+                ColorType.BASIC_COLOR -> presetOptions
             }
         }
 
+    /** The list of color options for the color section */
+    val colorSectionOptions: Flow<List<ColorOptionViewModel>> =
+        combine(wallpaperColorOptions, presetColorOptions) { wallpaperOptions, presetOptions ->
+            val subOptions =
+                wallpaperOptions.subList(0, min(COLOR_SECTION_OPTION_SIZE, wallpaperOptions.size))
+            // Add additional options based on preset colors if size of wallpaper color options is
+            // less than COLOR_SECTION_OPTION_SIZE
+            val additionalSubOptions =
+                presetOptions.subList(
+                    0,
+                    min(
+                        max(0, COLOR_SECTION_OPTION_SIZE - wallpaperOptions.size),
+                        presetOptions.size,
+                    )
+                )
+            subOptions + additionalSubOptions
+        }
+
     class Factory(
         private val context: Context,
         private val interactor: ColorPickerInteractor,
@@ -139,4 +164,8 @@
                 as T
         }
     }
+
+    companion object {
+        private const val COLOR_SECTION_OPTION_SIZE = 5
+    }
 }
diff --git a/src/com/android/customization/picker/preview/ui/section/PreviewWithClockCarouselSectionController.kt b/src/com/android/customization/picker/preview/ui/section/PreviewWithClockCarouselSectionController.kt
index f1319c6..6522d84 100644
--- a/src/com/android/customization/picker/preview/ui/section/PreviewWithClockCarouselSectionController.kt
+++ b/src/com/android/customization/picker/preview/ui/section/PreviewWithClockCarouselSectionController.kt
@@ -20,6 +20,7 @@
 import android.app.Activity
 import android.content.Context
 import android.view.ViewStub
+import androidx.core.view.isGone
 import androidx.lifecycle.LifecycleOwner
 import androidx.lifecycle.lifecycleScope
 import com.android.customization.picker.clock.data.repository.ClockRegistryProvider
@@ -67,6 +68,7 @@
         val carouselViewStub: ViewStub = view.requireViewById(R.id.clock_carousel_view_stub)
         carouselViewStub.layoutResource = R.layout.clock_carousel_view
         val carouselView: ClockCarouselView = carouselViewStub.inflate() as ClockCarouselView
+        carouselView.isGone = true
         lifecycleOwner.lifecycleScope.launch {
             val registry = withContext(Dispatchers.IO) { clockRegistryProvider.get() }
             val clockViewFactory = clockViewFactoryProvider.get(registry)
diff --git a/tests/src/com/android/customization/model/mode/DarkModeSnapshotRestorerTest.kt b/tests/src/com/android/customization/model/mode/DarkModeSnapshotRestorerTest.kt
new file mode 100644
index 0000000..38067b7
--- /dev/null
+++ b/tests/src/com/android/customization/model/mode/DarkModeSnapshotRestorerTest.kt
@@ -0,0 +1,107 @@
+/*
+ * 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.mode
+
+import androidx.test.filters.SmallTest
+import com.android.wallpaper.testing.FakeSnapshotStore
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.test.StandardTestDispatcher
+import kotlinx.coroutines.test.TestScope
+import kotlinx.coroutines.test.runTest
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.junit.runners.JUnit4
+
+@OptIn(ExperimentalCoroutinesApi::class)
+@SmallTest
+@RunWith(JUnit4::class)
+class DarkModeSnapshotRestorerTest {
+
+    private lateinit var underTest: DarkModeSnapshotRestorer
+    private lateinit var testScope: TestScope
+
+    private var isActive = false
+
+    @Before
+    fun setUp() {
+        val testDispatcher = StandardTestDispatcher()
+        testScope = TestScope(testDispatcher)
+        underTest =
+            DarkModeSnapshotRestorer(
+                backgroundDispatcher = testDispatcher,
+                isActive = { isActive },
+                setActive = { isActive = it },
+            )
+    }
+
+    @Test
+    fun `set up and restore - active`() =
+        testScope.runTest {
+            isActive = true
+
+            val store = FakeSnapshotStore()
+            store.store(underTest.setUpSnapshotRestorer(store = store))
+            val storedSnapshot = store.retrieve()
+
+            underTest.restoreToSnapshot(snapshot = storedSnapshot)
+            assertThat(isActive).isTrue()
+        }
+
+    @Test
+    fun `set up and restore - inactive`() =
+        testScope.runTest {
+            isActive = false
+
+            val store = FakeSnapshotStore()
+            store.store(underTest.setUpSnapshotRestorer(store = store))
+            val storedSnapshot = store.retrieve()
+
+            underTest.restoreToSnapshot(snapshot = storedSnapshot)
+            assertThat(isActive).isFalse()
+        }
+
+    @Test
+    fun `set up - deactivate - restore to active`() =
+        testScope.runTest {
+            isActive = true
+            val store = FakeSnapshotStore()
+            store.store(underTest.setUpSnapshotRestorer(store = store))
+            val initialSnapshot = store.retrieve()
+
+            underTest.store(isActivated = false)
+
+            underTest.restoreToSnapshot(snapshot = initialSnapshot)
+            assertThat(isActive).isTrue()
+        }
+
+    @Test
+    fun `set up - activate - restore to inactive`() =
+        testScope.runTest {
+            isActive = false
+            val store = FakeSnapshotStore()
+            store.store(underTest.setUpSnapshotRestorer(store = store))
+            val initialSnapshot = store.retrieve()
+
+            underTest.store(isActivated = true)
+
+            underTest.restoreToSnapshot(snapshot = initialSnapshot)
+            assertThat(isActive).isFalse()
+        }
+}
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 19d5dd1..6e5f776 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
@@ -59,6 +59,19 @@
     }
 
     @Test
+    fun `Select a color section color`() = runTest {
+        val colorSectionOptions = collectLastValue(underTest.colorSectionOptions)
+
+        assertColorOptionUiState(colorOptions = colorSectionOptions(), selectedColorOptionIndex = 0)
+
+        colorSectionOptions()?.get(2)?.onClick?.invoke()
+        assertColorOptionUiState(colorOptions = colorSectionOptions(), selectedColorOptionIndex = 2)
+
+        colorSectionOptions()?.get(4)?.onClick?.invoke()
+        assertColorOptionUiState(colorOptions = colorSectionOptions(), selectedColorOptionIndex = 4)
+    }
+
+    @Test
     fun `Select a preset color`() = runTest {
         val colorTypes = collectLastValue(underTest.colorTypes)
         val colorOptions = collectLastValue(underTest.colorOptions)
@@ -128,7 +141,20 @@
             colorTypeId = ColorType.BASIC_COLOR,
             isSelected = "Basic colors" == selectedColorTypeText,
         )
+        assertColorOptionUiState(colorOptions, selectedColorOptionIndex)
+    }
 
+    /**
+     * Asserts the picker section UI state is what is expected.
+     *
+     * @param colorOptions The observed color options
+     * @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>?,
+        selectedColorOptionIndex: Int,
+    ) {
         var foundSelectedColorOption = false
         assertThat(colorOptions).isNotNull()
         if (colorOptions != null) {
diff --git a/tests/src/com/android/customization/model/themedicon/domain/interactor/ThemedIconInteractorTest.kt b/tests/src/com/android/customization/model/themedicon/domain/interactor/ThemedIconInteractorTest.kt
new file mode 100644
index 0000000..e6e30c3
--- /dev/null
+++ b/tests/src/com/android/customization/model/themedicon/domain/interactor/ThemedIconInteractorTest.kt
@@ -0,0 +1,56 @@
+/*
+ * 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.themedicon.domain.interactor
+
+import androidx.test.filters.SmallTest
+import com.android.customization.model.themedicon.data.repository.ThemeIconRepository
+import com.android.wallpaper.testing.collectLastValue
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.test.runTest
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.junit.runners.JUnit4
+
+@OptIn(ExperimentalCoroutinesApi::class)
+@SmallTest
+@RunWith(JUnit4::class)
+class ThemedIconInteractorTest {
+
+    private lateinit var underTest: ThemedIconInteractor
+
+    @Before
+    fun setUp() {
+        underTest =
+            ThemedIconInteractor(
+                repository = ThemeIconRepository(),
+            )
+    }
+
+    @Test
+    fun `end-to-end`() = runTest {
+        val isActivated = collectLastValue(underTest.isActivated)
+
+        underTest.setActivated(isActivated = true)
+        assertThat(isActivated()).isTrue()
+
+        underTest.setActivated(isActivated = false)
+        assertThat(isActivated()).isFalse()
+    }
+}
diff --git a/tests/src/com/android/customization/model/themedicon/domain/interactor/ThemedIconSnapshotRestorerTest.kt b/tests/src/com/android/customization/model/themedicon/domain/interactor/ThemedIconSnapshotRestorerTest.kt
new file mode 100644
index 0000000..df1fd20
--- /dev/null
+++ b/tests/src/com/android/customization/model/themedicon/domain/interactor/ThemedIconSnapshotRestorerTest.kt
@@ -0,0 +1,102 @@
+/*
+ * 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.themedicon.domain.interactor
+
+import androidx.test.filters.SmallTest
+import com.android.customization.model.themedicon.data.repository.ThemeIconRepository
+import com.android.wallpaper.testing.FakeSnapshotStore
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.test.runTest
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.junit.runners.JUnit4
+
+@OptIn(ExperimentalCoroutinesApi::class)
+@SmallTest
+@RunWith(JUnit4::class)
+class ThemedIconSnapshotRestorerTest {
+
+    private lateinit var underTest: ThemedIconSnapshotRestorer
+    private var isActivated = false
+
+    @Before
+    fun setUp() {
+        isActivated = false
+        underTest =
+            ThemedIconSnapshotRestorer(
+                isActivated = { isActivated },
+                setActivated = { isActivated = it },
+                interactor =
+                    ThemedIconInteractor(
+                        repository = ThemeIconRepository(),
+                    )
+            )
+    }
+
+    @Test
+    fun `set up and restore - active`() = runTest {
+        isActivated = true
+
+        val store = FakeSnapshotStore()
+        store.store(underTest.setUpSnapshotRestorer(store = store))
+        val storedSnapshot = store.retrieve()
+
+        underTest.restoreToSnapshot(snapshot = storedSnapshot)
+        assertThat(isActivated).isTrue()
+    }
+
+    @Test
+    fun `set up and restore - inactive`() = runTest {
+        isActivated = false
+
+        val store = FakeSnapshotStore()
+        store.store(underTest.setUpSnapshotRestorer(store = store))
+        val storedSnapshot = store.retrieve()
+
+        underTest.restoreToSnapshot(snapshot = storedSnapshot)
+        assertThat(isActivated).isFalse()
+    }
+
+    @Test
+    fun `set up - deactivate - restore to active`() = runTest {
+        isActivated = true
+        val store = FakeSnapshotStore()
+        store.store(underTest.setUpSnapshotRestorer(store = store))
+        val initialSnapshot = store.retrieve()
+
+        underTest.store(isActivated = false)
+
+        underTest.restoreToSnapshot(snapshot = initialSnapshot)
+        assertThat(isActivated).isTrue()
+    }
+
+    @Test
+    fun `set up - activate - restore to inactive`() = runTest {
+        isActivated = false
+        val store = FakeSnapshotStore()
+        store.store(underTest.setUpSnapshotRestorer(store = store))
+        val initialSnapshot = store.retrieve()
+
+        underTest.store(isActivated = true)
+
+        underTest.restoreToSnapshot(snapshot = initialSnapshot)
+        assertThat(isActivated).isFalse()
+    }
+}