diff --git a/res/layout/fragment_grid.xml b/res/layout/fragment_grid.xml
new file mode 100644
index 0000000..4f0aaef
--- /dev/null
+++ b/res/layout/fragment_grid.xml
@@ -0,0 +1,93 @@
+<?xml version="1.0" encoding="utf-8"?><!--
+  ~ 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.
+  ~
+  -->
+
+<LinearLayout xmlns:android="http://schemas.android.com/apk/res/android"
+    android:layout_width="match_parent"
+    android:layout_height="match_parent"
+    android:orientation="vertical"
+    android:clipChildren="false">
+
+    <FrameLayout
+        android:id="@+id/section_header_container"
+        android:layout_width="match_parent"
+        android:layout_height="wrap_content">
+
+        <include layout="@layout/section_header" />
+
+    </FrameLayout>
+
+    <com.android.wallpaper.picker.DisplayAspectRatioFrameLayout
+        android:layout_width="match_parent"
+        android:layout_height="0dp"
+        android:layout_weight="1"
+        android:paddingTop="20dp"
+        android:paddingBottom="40dp">
+
+        <include
+            android:id="@+id/preview"
+            layout="@layout/wallpaper_preview_card"
+            android:layout_width="0dp"
+            android:layout_height="match_parent"
+            android:layout_gravity="center"/>
+
+    </com.android.wallpaper.picker.DisplayAspectRatioFrameLayout>
+
+    <LinearLayout
+        android:layout_width="match_parent"
+        android:layout_height="wrap_content"
+        android:orientation="vertical"
+        android:layout_marginHorizontal="24dp"
+        android:layout_marginBottom="28dp"
+        android:background="@drawable/picker_fragment_background"
+        android:paddingBottom="62dp"
+        android:clipChildren="false">
+
+        <FrameLayout
+            android:layout_width="match_parent"
+            android:layout_height="wrap_content"
+            android:layout_marginTop="22dp"
+            android:clipChildren="false">
+
+            <androidx.recyclerview.widget.RecyclerView
+                android:id="@id/options"
+                android:layout_width="wrap_content"
+                android:layout_height="wrap_content"
+                android:layout_gravity="center_horizontal"
+                android:clipToPadding="false"
+                android:paddingHorizontal="16dp"
+                android:clipChildren="false" />
+
+            <!--
+            This is just an invisible placeholder put in place so that the parent keeps its height
+            stable as the RecyclerView updates from 0 items to N items. Keeping it stable allows the
+            layout logic to keep the size of the preview container stable as well, which bodes well
+            for setting up the SurfaceView for remote rendering without changing its size after the
+            content is loaded into the RecyclerView.
+
+            It's critical for any TextViews inside the included layout to have text.
+            -->
+            <include
+                layout="@layout/grid_option_2"
+                android:layout_width="wrap_content"
+                android:layout_height="wrap_content"
+                android:visibility="invisible" />
+
+        </FrameLayout>
+
+    </LinearLayout>
+
+</LinearLayout>
diff --git a/res/layout/grid_option_2.xml b/res/layout/grid_option_2.xml
new file mode 100644
index 0000000..a8b453a
--- /dev/null
+++ b/res/layout/grid_option_2.xml
@@ -0,0 +1,70 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+  ~ 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.
+  ~
+  -->
+
+<LinearLayout
+    xmlns:android="http://schemas.android.com/apk/res/android"
+    xmlns:tools="http://schemas.android.com/tools"
+    android:layout_width="@dimen/option_item_size"
+    android:layout_height="wrap_content"
+    android:orientation="vertical"
+    android:gravity="center_horizontal"
+    android:clipChildren="false">
+
+    <FrameLayout
+        android:layout_width="@dimen/option_item_size"
+        android:layout_height="@dimen/option_item_size"
+        android:clipChildren="false">
+
+        <ImageView
+            android:id="@id/selection_border"
+            android:layout_width="match_parent"
+            android:layout_height="match_parent"
+            android:background="@drawable/option_item_border"
+            android:alpha="0"
+            android:importantForAccessibility="no" />
+
+        <ImageView
+            android:id="@id/background"
+            android:layout_width="match_parent"
+            android:layout_height="match_parent"
+            android:background="@drawable/option_item_background"
+            android:importantForAccessibility="no" />
+
+        <ImageView
+            android:id="@id/foreground"
+            android:layout_width="58dp"
+            android:layout_height="58dp"
+            android:layout_gravity="center" />
+
+    </FrameLayout>
+
+    <View
+        android:layout_width="0dp"
+        android:layout_height="8dp" />
+
+    <TextView
+        android:id="@id/text"
+        android:layout_width="wrap_content"
+        android:layout_height="wrap_content"
+        android:textColor="@color/text_color_primary"
+        android:singleLine="true"
+        android:ellipsize="end"
+        android:text="Placeholder for stable size calculation, please do not remove."
+        tools:ignore="HardcodedText" />
+
+</LinearLayout>
diff --git a/res/layout/keyguard_quick_affordance.xml b/res/layout/keyguard_quick_affordance.xml
index 03b4e14..1e5c339 100644
--- a/res/layout/keyguard_quick_affordance.xml
+++ b/res/layout/keyguard_quick_affordance.xml
@@ -19,15 +19,15 @@
 <LinearLayout
     xmlns:android="http://schemas.android.com/apk/res/android"
     xmlns:tools="http://schemas.android.com/tools"
-    android:layout_width="@dimen/keyguard_quick_affordance_picker_item_width"
+    android:layout_width="@dimen/option_item_size"
     android:layout_height="wrap_content"
     android:orientation="vertical"
     android:gravity="center_horizontal"
     android:clipChildren="false">
 
     <FrameLayout
-        android:layout_width="@dimen/keyguard_quick_affordance_icon_container_size"
-        android:layout_height="@dimen/keyguard_quick_affordance_icon_container_size"
+        android:layout_width="@dimen/option_item_size"
+        android:layout_height="@dimen/option_item_size"
         android:clipChildren="false">
 
         <ImageView
diff --git a/res/values/dimens.xml b/res/values/dimens.xml
index b856c35..f79d9b1 100644
--- a/res/values/dimens.xml
+++ b/res/values/dimens.xml
@@ -150,8 +150,6 @@
     <dimen name="keyguard_quick_affordance_icon_container_size">74dp</dimen>
     <!-- Size for the icon of a quick affordance for the lock screen in the picker experience. -->
     <dimen name="keyguard_quick_affordance_icon_size">24dp</dimen>
-    <!-- Width of a single selectable item in the lock screen quick affordance picker. -->
-    <dimen name="keyguard_quick_affordance_picker_item_width">74dp</dimen>
 
     <dimen name="clock_carousel_item_width">190dp</dimen>
     <dimen name="clock_carousel_item_height">380dp</dimen>
diff --git a/src/com/android/customization/model/grid/GridOptionsManager.java b/src/com/android/customization/model/grid/GridOptionsManager.java
index 7f15d83..bff7933 100644
--- a/src/com/android/customization/model/grid/GridOptionsManager.java
+++ b/src/com/android/customization/model/grid/GridOptionsManager.java
@@ -21,7 +21,9 @@
 import android.os.Looper;
 import android.util.Log;
 
+import androidx.annotation.Nullable;
 import androidx.annotation.VisibleForTesting;
+import androidx.lifecycle.LiveData;
 
 import com.android.customization.model.CustomizationManager;
 import com.android.customization.module.CustomizationInjector;
@@ -110,6 +112,13 @@
         });
     }
 
+    /**
+     * Returns an observable that receives a new value each time that the grid options are changed.
+     */
+    public LiveData<Object> getOptionChangeObservable(@Nullable Handler handler) {
+        return mProvider.getOptionChangeObservable(handler);
+    }
+
     /** Call through content provider API to render preview */
     public void renderPreview(Bundle bundle, String gridName,
             PreviewUtils.WorkspacePreviewCallback callback) {
diff --git a/src/com/android/customization/model/grid/GridSectionController.java b/src/com/android/customization/model/grid/GridSectionController.java
index 2f54a1b..3e5dba0 100644
--- a/src/com/android/customization/model/grid/GridSectionController.java
+++ b/src/com/android/customization/model/grid/GridSectionController.java
@@ -22,8 +22,12 @@
 import android.widget.TextView;
 
 import androidx.annotation.Nullable;
+import androidx.fragment.app.Fragment;
+import androidx.lifecycle.LifecycleOwner;
+import androidx.lifecycle.Observer;
 
 import com.android.customization.model.CustomizationManager.OptionsFetchedListener;
+import com.android.customization.model.grid.ui.fragment.GridFragment2;
 import com.android.customization.picker.grid.GridFragment;
 import com.android.customization.picker.grid.GridSectionView;
 import com.android.wallpaper.R;
@@ -38,11 +42,22 @@
 
     private final GridOptionsManager mGridOptionsManager;
     private final CustomizationSectionNavigationController mSectionNavigationController;
+    private final boolean mIsRevampedUiEnabled;
+    private final Observer<Object> mOptionChangeObserver;
+    private final LifecycleOwner mLifecycleOwner;
+    private TextView mSectionDescription;
+    private View mSectionTile;
 
-    public GridSectionController(GridOptionsManager gridOptionsManager,
-            CustomizationSectionNavigationController sectionNavigationController) {
+    public GridSectionController(
+            GridOptionsManager gridOptionsManager,
+            CustomizationSectionNavigationController sectionNavigationController,
+            LifecycleOwner lifecycleOwner,
+            boolean isRevampedUiEnabled) {
         mGridOptionsManager = gridOptionsManager;
         mSectionNavigationController = sectionNavigationController;
+        mIsRevampedUiEnabled = isRevampedUiEnabled;
+        mLifecycleOwner = lifecycleOwner;
+        mOptionChangeObserver = o -> updateUi(/* reload= */ true);
     }
 
     @Override
@@ -52,34 +67,68 @@
 
     @Override
     public GridSectionView createView(Context context) {
-        GridSectionView gridSectionView = (GridSectionView) LayoutInflater.from(context)
+        final GridSectionView gridSectionView = (GridSectionView) LayoutInflater.from(context)
                 .inflate(R.layout.grid_section_view, /* root= */ null);
-        TextView sectionDescription = gridSectionView.findViewById(R.id.grid_section_description);
-        View sectionTile = gridSectionView.findViewById(R.id.grid_section_tile);
+        mSectionDescription = gridSectionView.findViewById(R.id.grid_section_description);
+        mSectionTile = gridSectionView.findViewById(R.id.grid_section_tile);
 
         // Fetch grid options to show currently set grid.
-        mGridOptionsManager.fetchOptions(new OptionsFetchedListener<GridOption>() {
-            @Override
-            public void onOptionsLoaded(List<GridOption> options) {
-                sectionDescription.setText(getActiveOption(options).getTitle());
-            }
-
-            @Override
-            public void onError(@Nullable Throwable throwable) {
-                if (throwable != null) {
-                    Log.e(TAG, "Error loading grid options", throwable);
-                }
-                sectionDescription.setText(R.string.something_went_wrong);
-                sectionTile.setVisibility(View.GONE);
-            }
-        }, /* The result is getting when calling isAvailable(), so reload= */ false);
+        updateUi(/* The result is getting when calling isAvailable(), so reload= */ false);
+        if (mIsRevampedUiEnabled) {
+            mGridOptionsManager.getOptionChangeObservable(/* handler= */ null).observe(
+                    mLifecycleOwner,
+                    mOptionChangeObserver);
+        }
 
         gridSectionView.setOnClickListener(
-                v -> mSectionNavigationController.navigateTo(new GridFragment()));
+                v -> {
+                    final Fragment gridFragment;
+                    if (mIsRevampedUiEnabled) {
+                        gridFragment = new GridFragment2();
+                    } else {
+                        gridFragment = new GridFragment();
+                    }
+                    mSectionNavigationController.navigateTo(gridFragment);
+                });
 
         return gridSectionView;
     }
 
+    @Override
+    public void release() {
+        if (mIsRevampedUiEnabled) {
+            mGridOptionsManager.getOptionChangeObservable(/* handler= */ null).removeObserver(
+                    mOptionChangeObserver
+            );
+        }
+    }
+
+    @Override
+    public void onTransitionOut() {
+        CustomizationSectionController.super.onTransitionOut();
+    }
+
+    private void updateUi(final boolean reload) {
+        mGridOptionsManager.fetchOptions(
+                new OptionsFetchedListener<GridOption>() {
+                    @Override
+                    public void onOptionsLoaded(List<GridOption> options) {
+                        final String title = getActiveOption(options).getTitle();
+                        mSectionDescription.setText(title);
+                    }
+
+                    @Override
+                    public void onError(@Nullable Throwable throwable) {
+                        if (throwable != null) {
+                            Log.e(TAG, "Error loading grid options", throwable);
+                        }
+                        mSectionDescription.setText(R.string.something_went_wrong);
+                        mSectionTile.setVisibility(View.GONE);
+                    }
+                },
+                reload);
+    }
+
     private GridOption getActiveOption(List<GridOption> options) {
         return options.stream()
                 .filter(option -> option.isActive(mGridOptionsManager))
diff --git a/src/com/android/customization/model/grid/LauncherGridOptionsProvider.java b/src/com/android/customization/model/grid/LauncherGridOptionsProvider.java
index fd40363..5ae283a 100644
--- a/src/com/android/customization/model/grid/LauncherGridOptionsProvider.java
+++ b/src/com/android/customization/model/grid/LauncherGridOptionsProvider.java
@@ -19,12 +19,17 @@
 import android.content.ContentValues;
 import android.content.Context;
 import android.content.res.Resources;
+import android.database.ContentObserver;
 import android.database.Cursor;
+import android.net.Uri;
 import android.os.Bundle;
+import android.os.Handler;
 import android.view.SurfaceView;
 
 import androidx.annotation.Nullable;
 import androidx.annotation.WorkerThread;
+import androidx.lifecycle.LiveData;
+import androidx.lifecycle.MutableLiveData;
 
 import com.android.customization.model.ResourceConstants;
 import com.android.wallpaper.R;
@@ -53,6 +58,7 @@
     private final Context mContext;
     private final PreviewUtils mPreviewUtils;
     private List<GridOption> mOptions;
+    private OptionChangeLiveData mLiveData;
 
     public LauncherGridOptionsProvider(Context context, String authorityMetadataKey) {
         mPreviewUtils = new PreviewUtils(context, authorityMetadataKey);
@@ -117,4 +123,51 @@
         return mContext.getContentResolver().update(mPreviewUtils.getUri(DEFAULT_GRID), values,
                 null, null);
     }
+
+    /**
+     * Returns an observable that receives a new value each time that the grid options are changed.
+     */
+    public LiveData<Object> getOptionChangeObservable(
+            @Nullable Handler handler) {
+        if (mLiveData == null) {
+            mLiveData = new OptionChangeLiveData(
+                    mContext, mPreviewUtils.getUri(DEFAULT_GRID), handler);
+        }
+
+        return mLiveData;
+    }
+
+    private static class OptionChangeLiveData extends MutableLiveData<Object> {
+
+        private final ContentResolver mContentResolver;
+        private final Uri mUri;
+        private final ContentObserver mContentObserver;
+
+        OptionChangeLiveData(
+                Context context,
+                Uri uri,
+                @Nullable Handler handler) {
+            mContentResolver = context.getContentResolver();
+            mUri = uri;
+            mContentObserver = new ContentObserver(handler) {
+                @Override
+                public void onChange(boolean selfChange) {
+                    postValue(new Object());
+                }
+            };
+        }
+
+        @Override
+        protected void onActive() {
+            mContentResolver.registerContentObserver(
+                    mUri,
+                    /* notifyForDescendants= */ true,
+                    mContentObserver);
+        }
+
+        @Override
+        protected void onInactive() {
+            mContentResolver.unregisterContentObserver(mContentObserver);
+        }
+    }
 }
diff --git a/src/com/android/customization/model/grid/data/repository/GridRepository.kt b/src/com/android/customization/model/grid/data/repository/GridRepository.kt
new file mode 100644
index 0000000..7c84aec
--- /dev/null
+++ b/src/com/android/customization/model/grid/data/repository/GridRepository.kt
@@ -0,0 +1,122 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package com.android.customization.model.grid.data.repository
+
+import androidx.lifecycle.asFlow
+import com.android.customization.model.CustomizationManager
+import com.android.customization.model.grid.GridOption
+import com.android.customization.model.grid.GridOptionsManager
+import com.android.customization.model.grid.shared.model.GridOptionItemModel
+import com.android.customization.model.grid.shared.model.GridOptionItemsModel
+import kotlin.coroutines.resume
+import kotlinx.coroutines.CoroutineDispatcher
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.SharingStarted
+import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.flow.stateIn
+import kotlinx.coroutines.suspendCancellableCoroutine
+import kotlinx.coroutines.withContext
+
+interface GridRepository {
+    val optionChanges: Flow<Unit>
+    suspend fun getOptions(): GridOptionItemsModel
+}
+
+class GridRepositoryImpl(
+    private val applicationScope: CoroutineScope,
+    private val manager: GridOptionsManager,
+    private val backgroundDispatcher: CoroutineDispatcher,
+) : GridRepository {
+
+    override val optionChanges: Flow<Unit> =
+        manager.getOptionChangeObservable(/* handler= */ null).asFlow().map {}
+
+    private val selectedOption = MutableStateFlow<GridOption?>(null)
+
+    override suspend fun getOptions(): GridOptionItemsModel {
+        return withContext(backgroundDispatcher) {
+            suspendCancellableCoroutine { continuation ->
+                manager.fetchOptions(
+                    object : CustomizationManager.OptionsFetchedListener<GridOption> {
+                        override fun onOptionsLoaded(options: MutableList<GridOption>?) {
+                            val optionsOrEmpty = options ?: emptyList()
+                            selectedOption.value = optionsOrEmpty.find { it.isActive(manager) }
+                            continuation.resume(
+                                GridOptionItemsModel.Loaded(
+                                    optionsOrEmpty.map { option -> toModel(option) }
+                                )
+                            )
+                        }
+
+                        override fun onError(throwable: Throwable?) {
+                            continuation.resume(
+                                GridOptionItemsModel.Error(
+                                    throwable ?: Exception("Failed to load grid options!")
+                                ),
+                            )
+                        }
+                    },
+                    /* reload= */ true,
+                )
+            }
+        }
+    }
+
+    private fun toModel(option: GridOption): GridOptionItemModel {
+        return GridOptionItemModel(
+            name = option.title,
+            rows = option.rows,
+            cols = option.cols,
+            isSelected =
+                selectedOption
+                    .map { it.key() }
+                    .map { selectedOptionKey -> option.key() == selectedOptionKey }
+                    .stateIn(
+                        scope = applicationScope,
+                        started = SharingStarted.Eagerly,
+                        initialValue = false,
+                    ),
+            onSelected = { onSelected(option) },
+        )
+    }
+
+    private suspend fun onSelected(option: GridOption) {
+        withContext(backgroundDispatcher) {
+            suspendCancellableCoroutine { continuation ->
+                manager.apply(
+                    option,
+                    object : CustomizationManager.Callback {
+                        override fun onSuccess() {
+                            continuation.resume(true)
+                        }
+
+                        override fun onError(throwable: Throwable?) {
+                            continuation.resume(false)
+                        }
+                    },
+                )
+            }
+        }
+    }
+
+    private fun GridOption?.key(): String? {
+        return if (this != null) "${cols}x${rows}" else null
+    }
+}
diff --git a/src/com/android/customization/model/grid/domain/interactor/GridInteractor.kt b/src/com/android/customization/model/grid/domain/interactor/GridInteractor.kt
new file mode 100644
index 0000000..307507b
--- /dev/null
+++ b/src/com/android/customization/model/grid/domain/interactor/GridInteractor.kt
@@ -0,0 +1,87 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package com.android.customization.model.grid.domain.interactor
+
+import com.android.customization.model.grid.data.repository.GridRepository
+import com.android.customization.model.grid.shared.model.GridOptionItemModel
+import com.android.customization.model.grid.shared.model.GridOptionItemsModel
+import javax.inject.Provider
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.SharingStarted
+import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.flow.onStart
+import kotlinx.coroutines.flow.shareIn
+
+class GridInteractor(
+    private val applicationScope: CoroutineScope,
+    private val repository: GridRepository,
+    private val snapshotRestorer: Provider<GridSnapshotRestorer>,
+) {
+    val options: Flow<GridOptionItemsModel> =
+    // this upstream flow tells us each time the options are changed.
+    repository.optionChanges
+            // when we start, we pretend the options _just_ changed. This way, we load something as
+            // soon as possible into the flow so it's ready by the time the first observer starts to
+            // observe.
+            .onStart { emit(Unit) }
+            // each time the options changed, we load them.
+            .map { reload() }
+            // we place the loaded options in a SharedFlow so downstream observers all
+            // share the same flow and don't trigger a new one each time they want to start
+            // observing.
+            .shareIn(
+                scope = applicationScope,
+                started = SharingStarted.WhileSubscribed(),
+                replay = 1,
+            )
+
+    suspend fun setSelectedOption(model: GridOptionItemModel) {
+        model.onSelected.invoke()
+    }
+
+    suspend fun getSelectedOption(): GridOptionItemModel? {
+        return (repository.getOptions() as? GridOptionItemsModel.Loaded)?.options?.firstOrNull {
+            optionItem ->
+            optionItem.isSelected.value
+        }
+    }
+
+    private suspend fun reload(): GridOptionItemsModel {
+        val model = repository.getOptions()
+        return if (model is GridOptionItemsModel.Loaded) {
+            GridOptionItemsModel.Loaded(
+                options =
+                    model.options.map { option ->
+                        GridOptionItemModel(
+                            name = option.name,
+                            cols = option.cols,
+                            rows = option.rows,
+                            isSelected = option.isSelected,
+                            onSelected = {
+                                option.onSelected()
+                                snapshotRestorer.get().store(option)
+                            },
+                        )
+                    }
+            )
+        } else {
+            model
+        }
+    }
+}
diff --git a/src/com/android/customization/model/grid/domain/interactor/GridSnapshotRestorer.kt b/src/com/android/customization/model/grid/domain/interactor/GridSnapshotRestorer.kt
new file mode 100644
index 0000000..583cbc7
--- /dev/null
+++ b/src/com/android/customization/model/grid/domain/interactor/GridSnapshotRestorer.kt
@@ -0,0 +1,73 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package com.android.customization.model.grid.domain.interactor
+
+import android.util.Log
+import com.android.customization.model.grid.shared.model.GridOptionItemModel
+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 GridSnapshotRestorer(
+    private val interactor: GridInteractor,
+) : SnapshotRestorer {
+
+    private lateinit var store: SnapshotStore
+    private var originalOption: GridOptionItemModel? = null
+
+    override suspend fun setUpSnapshotRestorer(store: SnapshotStore): RestorableSnapshot {
+        this.store = store
+        val option = interactor.getSelectedOption()
+        originalOption = option
+        return snapshot(option)
+    }
+
+    override suspend fun restoreToSnapshot(snapshot: RestorableSnapshot) {
+        val optionNameFromSnapshot = snapshot.args[KEY_GRID_OPTION_NAME]
+        originalOption?.let { optionToRestore ->
+            if (optionToRestore.name != optionNameFromSnapshot) {
+                Log.wtf(
+                    TAG,
+                    """Original snapshot name was ${optionToRestore.name} but we're being told to
+                        | restore to $optionNameFromSnapshot. The current implementation doesn't
+                        | support undo, only a reset back to the original grid option.""".trimMargin(),
+                )
+            }
+
+            interactor.setSelectedOption(optionToRestore)
+        }
+    }
+
+    fun store(option: GridOptionItemModel) {
+        store.store(snapshot(option))
+    }
+
+    private fun snapshot(option: GridOptionItemModel?): RestorableSnapshot {
+        return RestorableSnapshot(
+            args =
+                buildMap {
+                    option?.name?.let { optionName -> put(KEY_GRID_OPTION_NAME, optionName) }
+                }
+        )
+    }
+
+    companion object {
+        private const val TAG = "GridSnapshotRestorer"
+        private const val KEY_GRID_OPTION_NAME = "grid_option"
+    }
+}
diff --git a/src/com/android/customization/model/grid/shared/model/GridOptionItemModel.kt b/src/com/android/customization/model/grid/shared/model/GridOptionItemModel.kt
new file mode 100644
index 0000000..2eabeab
--- /dev/null
+++ b/src/com/android/customization/model/grid/shared/model/GridOptionItemModel.kt
@@ -0,0 +1,28 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package com.android.customization.model.grid.shared.model
+
+import kotlinx.coroutines.flow.StateFlow
+
+data class GridOptionItemModel(
+    val name: String,
+    val cols: Int,
+    val rows: Int,
+    val isSelected: StateFlow<Boolean>,
+    val onSelected: suspend () -> Unit,
+)
diff --git a/src/com/android/customization/model/grid/shared/model/GridOptionItemsModel.kt b/src/com/android/customization/model/grid/shared/model/GridOptionItemsModel.kt
new file mode 100644
index 0000000..e969be8
--- /dev/null
+++ b/src/com/android/customization/model/grid/shared/model/GridOptionItemsModel.kt
@@ -0,0 +1,27 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package com.android.customization.model.grid.shared.model
+
+sealed class GridOptionItemsModel {
+    data class Loaded(
+        val options: List<GridOptionItemModel>,
+    ) : GridOptionItemsModel()
+    data class Error(
+        val throwable: Throwable?,
+    ) : GridOptionItemsModel()
+}
diff --git a/src/com/android/customization/model/grid/ui/binder/GridScreenBinder.kt b/src/com/android/customization/model/grid/ui/binder/GridScreenBinder.kt
new file mode 100644
index 0000000..d24467a
--- /dev/null
+++ b/src/com/android/customization/model/grid/ui/binder/GridScreenBinder.kt
@@ -0,0 +1,75 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package com.android.customization.model.grid.ui.binder
+
+import android.view.View
+import androidx.lifecycle.Lifecycle
+import androidx.lifecycle.LifecycleOwner
+import androidx.lifecycle.lifecycleScope
+import androidx.lifecycle.repeatOnLifecycle
+import androidx.recyclerview.widget.LinearLayoutManager
+import androidx.recyclerview.widget.RecyclerView
+import com.android.customization.model.grid.ui.viewmodel.GridScreenViewModel
+import com.android.customization.picker.common.ui.view.ItemSpacing
+import com.android.wallpaper.R
+import com.android.wallpaper.picker.option.ui.adapter.OptionItemAdapter
+import com.android.wallpaper.picker.option.ui.binder.OptionItemBinder
+import kotlinx.coroutines.CoroutineDispatcher
+import kotlinx.coroutines.launch
+
+object GridScreenBinder {
+    fun bind(
+        view: View,
+        viewModel: GridScreenViewModel,
+        lifecycleOwner: LifecycleOwner,
+        backgroundDispatcher: CoroutineDispatcher,
+        onOptionsChanged: () -> Unit,
+    ) {
+        val optionView: RecyclerView = view.requireViewById(R.id.options)
+        optionView.layoutManager =
+            LinearLayoutManager(
+                view.context,
+                RecyclerView.HORIZONTAL,
+                /* reverseLayout= */ false,
+            )
+        optionView.addItemDecoration(ItemSpacing(ItemSpacing.ITEM_SPACING_DP))
+        val adapter =
+            OptionItemAdapter(
+                layoutResourceId = R.layout.grid_option_2,
+                lifecycleOwner = lifecycleOwner,
+                backgroundDispatcher = backgroundDispatcher,
+                foregroundTintSpec =
+                    OptionItemBinder.TintSpec(
+                        selectedColor = view.context.getColor(R.color.text_color_primary),
+                        unselectedColor = view.context.getColor(R.color.text_color_secondary),
+                    )
+            )
+        optionView.adapter = adapter
+
+        lifecycleOwner.lifecycleScope.launch {
+            lifecycleOwner.repeatOnLifecycle(Lifecycle.State.STARTED) {
+                launch {
+                    viewModel.optionItems.collect { options ->
+                        adapter.setItems(options)
+                        onOptionsChanged()
+                    }
+                }
+            }
+        }
+    }
+}
diff --git a/src/com/android/customization/model/grid/ui/fragment/GridFragment2.kt b/src/com/android/customization/model/grid/ui/fragment/GridFragment2.kt
new file mode 100644
index 0000000..4440b77
--- /dev/null
+++ b/src/com/android/customization/model/grid/ui/fragment/GridFragment2.kt
@@ -0,0 +1,116 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package com.android.customization.model.grid.ui.fragment
+
+import android.os.Bundle
+import android.view.LayoutInflater
+import android.view.View
+import android.view.ViewGroup
+import androidx.lifecycle.ViewModelProvider
+import com.android.customization.model.grid.ui.binder.GridScreenBinder
+import com.android.customization.model.grid.ui.viewmodel.GridScreenViewModel
+import com.android.customization.module.ThemePickerInjector
+import com.android.wallpaper.R
+import com.android.wallpaper.module.CurrentWallpaperInfoFactory
+import com.android.wallpaper.module.InjectorProvider
+import com.android.wallpaper.picker.AppbarFragment
+import com.android.wallpaper.picker.customization.ui.binder.ScreenPreviewBinder
+import com.android.wallpaper.picker.customization.ui.viewmodel.ScreenPreviewViewModel
+import com.android.wallpaper.util.PreviewUtils
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.suspendCancellableCoroutine
+
+@OptIn(ExperimentalCoroutinesApi::class)
+class GridFragment2 : AppbarFragment() {
+
+    override fun onCreateView(
+        inflater: LayoutInflater,
+        container: ViewGroup?,
+        savedInstanceState: Bundle?
+    ): View {
+        val view =
+            inflater.inflate(
+                R.layout.fragment_grid,
+                container,
+                false,
+            )
+        setUpToolbar(view)
+
+        val injector = InjectorProvider.getInjector() as ThemePickerInjector
+
+        val wallpaperInfoFactory = injector.getCurrentWallpaperInfoFactory(requireContext())
+        var screenPreviewBinding = bindScreenPreview(view, wallpaperInfoFactory)
+
+        val viewModelFactory = injector.getGridScreenViewModelFactory(requireContext())
+        GridScreenBinder.bind(
+            view = view,
+            viewModel =
+                ViewModelProvider(
+                    this,
+                    viewModelFactory,
+                )[GridScreenViewModel::class.java],
+            lifecycleOwner = this,
+            backgroundDispatcher = Dispatchers.IO,
+            onOptionsChanged = {
+                screenPreviewBinding.destroy()
+                screenPreviewBinding = bindScreenPreview(view, wallpaperInfoFactory)
+            }
+        )
+
+        return view
+    }
+
+    override fun getDefaultTitle(): CharSequence {
+        return getString(R.string.grid_title)
+    }
+
+    private fun bindScreenPreview(
+        view: View,
+        wallpaperInfoFactory: CurrentWallpaperInfoFactory,
+    ): ScreenPreviewBinder.Binding {
+        return ScreenPreviewBinder.bind(
+            activity = requireActivity(),
+            previewView = view.requireViewById(R.id.preview),
+            viewModel =
+                ScreenPreviewViewModel(
+                    previewUtils =
+                        PreviewUtils(
+                            context = requireContext(),
+                            authorityMetadataKey =
+                                requireContext()
+                                    .getString(
+                                        R.string.grid_control_metadata_name,
+                                    ),
+                        ),
+                    wallpaperInfoProvider = {
+                        suspendCancellableCoroutine { continuation ->
+                            wallpaperInfoFactory.createCurrentWallpaperInfos(
+                                { homeWallpaper, lockWallpaper, _ ->
+                                    continuation.resume(homeWallpaper ?: lockWallpaper, null)
+                                },
+                                /* forceRefresh= */ true,
+                            )
+                        }
+                    },
+                ),
+            lifecycleOwner = this,
+            offsetToStart = false,
+        )
+    }
+}
diff --git a/src/com/android/customization/model/grid/ui/viewmodel/GridScreenViewModel.kt b/src/com/android/customization/model/grid/ui/viewmodel/GridScreenViewModel.kt
new file mode 100644
index 0000000..af6ed0f
--- /dev/null
+++ b/src/com/android/customization/model/grid/ui/viewmodel/GridScreenViewModel.kt
@@ -0,0 +1,110 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package com.android.customization.model.grid.ui.viewmodel
+
+import android.annotation.SuppressLint
+import android.content.Context
+import android.content.res.Resources
+import androidx.lifecycle.ViewModel
+import androidx.lifecycle.ViewModelProvider
+import androidx.lifecycle.viewModelScope
+import com.android.customization.model.ResourceConstants
+import com.android.customization.model.grid.domain.interactor.GridInteractor
+import com.android.customization.model.grid.shared.model.GridOptionItemsModel
+import com.android.customization.widget.GridTileDrawable
+import com.android.wallpaper.picker.common.icon.ui.viewmodel.Icon
+import com.android.wallpaper.picker.common.text.ui.viewmodel.Text
+import com.android.wallpaper.picker.option.ui.viewmodel.OptionItemViewModel
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.flowOf
+import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.launch
+
+class GridScreenViewModel(
+    context: Context,
+    private val interactor: GridInteractor,
+) : ViewModel() {
+
+    @SuppressLint("StaticFieldLeak") // We're not leaking this context as it is the app context.
+    private val applicationContext = context.applicationContext
+
+    val optionItems: Flow<List<OptionItemViewModel>> =
+        interactor.options.map { model -> toViewModel(model) }
+
+    private fun toViewModel(
+        model: GridOptionItemsModel,
+    ): List<OptionItemViewModel> {
+        val iconShapePath =
+            applicationContext.resources.getString(
+                Resources.getSystem()
+                    .getIdentifier(
+                        ResourceConstants.CONFIG_ICON_MASK,
+                        "string",
+                        ResourceConstants.ANDROID_PACKAGE,
+                    )
+            )
+
+        return when (model) {
+            is GridOptionItemsModel.Loaded ->
+                model.options.map { option ->
+                    val text = Text.Loaded(option.name)
+                    OptionItemViewModel(
+                        key = flowOf("${option.cols}x${option.rows}"),
+                        icon =
+                            Icon.Loaded(
+                                drawable =
+                                    GridTileDrawable(
+                                        option.cols,
+                                        option.rows,
+                                        iconShapePath,
+                                    ),
+                                contentDescription = text
+                            ),
+                        text = text,
+                        isSelected = option.isSelected,
+                        onClicked =
+                            option.isSelected.map { isSelected ->
+                                if (!isSelected) {
+                                    { viewModelScope.launch { option.onSelected() } }
+                                } else {
+                                    null
+                                }
+                            },
+                    )
+                }
+            is GridOptionItemsModel.Error -> emptyList()
+        }
+    }
+
+    class Factory(
+        context: Context,
+        private val interactor: GridInteractor,
+    ) : ViewModelProvider.Factory {
+
+        private val applicationContext = context.applicationContext
+
+        @Suppress("UNCHECKED_CAST")
+        override fun <T : ViewModel> create(modelClass: Class<T>): T {
+            return GridScreenViewModel(
+                context = applicationContext,
+                interactor = interactor,
+            )
+                as T
+        }
+    }
+}
diff --git a/src/com/android/customization/module/DefaultCustomizationSections.java b/src/com/android/customization/module/DefaultCustomizationSections.java
index 18f2d37..7653095 100644
--- a/src/com/android/customization/module/DefaultCustomizationSections.java
+++ b/src/com/android/customization/module/DefaultCustomizationSections.java
@@ -183,8 +183,12 @@
                                 mThemedIconSnapshotRestorer));
 
                 // App grid section.
-                sectionControllers.add(new GridSectionController(
-                        GridOptionsManager.getInstance(activity), sectionNavigationController));
+                sectionControllers.add(
+                        new GridSectionController(
+                                GridOptionsManager.getInstance(activity),
+                                sectionNavigationController,
+                                lifecycleOwner,
+                                /* isRevampedUiEnabled= */ true));
                 break;
         }
 
@@ -234,8 +238,12 @@
                 mThemedIconSnapshotRestorer));
 
         // App grid section.
-        sectionControllers.add(new GridSectionController(
-                GridOptionsManager.getInstance(activity), sectionNavigationController));
+        sectionControllers.add(
+                new GridSectionController(
+                        GridOptionsManager.getInstance(activity),
+                        sectionNavigationController,
+                        lifecycleOwner,
+                        /* isRevampedUiEnabled= */ false));
 
         return sectionControllers;
     }
diff --git a/src/com/android/customization/module/ThemePickerInjector.kt b/src/com/android/customization/module/ThemePickerInjector.kt
index f364e4c..9a58d70 100644
--- a/src/com/android/customization/module/ThemePickerInjector.kt
+++ b/src/com/android/customization/module/ThemePickerInjector.kt
@@ -25,6 +25,11 @@
 import androidx.fragment.app.Fragment
 import androidx.fragment.app.FragmentActivity
 import androidx.lifecycle.ViewModelProvider
+import com.android.customization.model.grid.GridOptionsManager
+import com.android.customization.model.grid.data.repository.GridRepositoryImpl
+import com.android.customization.model.grid.domain.interactor.GridInteractor
+import com.android.customization.model.grid.domain.interactor.GridSnapshotRestorer
+import com.android.customization.model.grid.ui.viewmodel.GridScreenViewModel
 import com.android.customization.model.mode.DarkModeSnapshotRestorer
 import com.android.customization.model.theme.OverlayManagerCompat
 import com.android.customization.model.theme.ThemeBundleProvider
@@ -69,9 +74,11 @@
 import com.android.wallpaper.picker.LivePreviewFragment
 import com.android.wallpaper.picker.PreviewFragment
 import com.android.wallpaper.picker.undo.domain.interactor.SnapshotRestorer
+import kotlinx.coroutines.DelicateCoroutinesApi
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.GlobalScope
 
+@OptIn(DelicateCoroutinesApi::class)
 open class ThemePickerInjector : WallpaperPicker2Injector(), CustomizationInjector {
     private var customizationSections: CustomizationSections? = null
     private var userEventLogger: UserEventLogger? = null
@@ -99,6 +106,9 @@
     private var themedIconSnapshotRestorer: ThemedIconSnapshotRestorer? = null
     private var themedIconInteractor: ThemedIconInteractor? = null
     private var clockSettingsViewModelFactory: ClockSettingsViewModel.Factory? = null
+    private var gridInteractor: GridInteractor? = null
+    private var gridSnapshotRestorer: GridSnapshotRestorer? = null
+    private var gridScreenViewModelFactory: GridScreenViewModel.Factory? = null
 
     override fun getCustomizationSections(activity: ComponentActivity): CustomizationSections {
         return customizationSections
@@ -196,6 +206,7 @@
             this[KEY_NOTIFICATIONS_SNAPSHOT_RESTORER] = getNotificationsSnapshotRestorer(context)
             this[KEY_DARK_MODE_SNAPSHOT_RESTORER] = getDarkModeSnapshotRestorer(context)
             this[KEY_THEMED_ICON_SNAPSHOT_RESTORER] = getThemedIconSnapshotRestorer(context)
+            this[KEY_APP_GRID_SNAPSHOT_RESTORER] = getGridSnapshotRestorer(context)
         }
     }
 
@@ -423,6 +434,44 @@
                 .also { clockSettingsViewModelFactory = it }
     }
 
+    fun getGridScreenViewModelFactory(
+        context: Context,
+    ): ViewModelProvider.Factory {
+        return gridScreenViewModelFactory
+            ?: GridScreenViewModel.Factory(
+                    context = context,
+                    interactor = getGridInteractor(context),
+                )
+                .also { gridScreenViewModelFactory = it }
+    }
+
+    private fun getGridInteractor(
+        context: Context,
+    ): GridInteractor {
+        return gridInteractor
+            ?: GridInteractor(
+                    applicationScope = GlobalScope,
+                    repository =
+                        GridRepositoryImpl(
+                            applicationScope = GlobalScope,
+                            manager = GridOptionsManager.getInstance(context),
+                            backgroundDispatcher = Dispatchers.IO,
+                        ),
+                    snapshotRestorer = { getGridSnapshotRestorer(context) },
+                )
+                .also { gridInteractor = it }
+    }
+
+    private fun getGridSnapshotRestorer(
+        context: Context,
+    ): GridSnapshotRestorer {
+        return gridSnapshotRestorer
+            ?: GridSnapshotRestorer(
+                    interactor = getGridInteractor(context),
+                )
+                .also { gridSnapshotRestorer = it }
+    }
+
     companion object {
         @JvmStatic
         private val KEY_QUICK_AFFORDANCE_SNAPSHOT_RESTORER =
@@ -435,11 +484,15 @@
         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
+        @JvmStatic
+        private val KEY_APP_GRID_SNAPSHOT_RESTORER = KEY_THEMED_ICON_SNAPSHOT_RESTORER + 1
 
         /**
          * When this injector is overridden, this is the minimal value that should be used by
          * restorers returns in [getSnapshotRestorers].
+         *
+         * It should always be greater than the biggest restorer key.
          */
-        @JvmStatic protected val MIN_SNAPSHOT_RESTORER_KEY = KEY_THEMED_ICON_SNAPSHOT_RESTORER + 1
+        @JvmStatic protected val MIN_SNAPSHOT_RESTORER_KEY = KEY_APP_GRID_SNAPSHOT_RESTORER + 1
     }
 }
diff --git a/tests/src/com/android/customization/model/grid/data/repository/FakeGridRepository.kt b/tests/src/com/android/customization/model/grid/data/repository/FakeGridRepository.kt
new file mode 100644
index 0000000..6291c21
--- /dev/null
+++ b/tests/src/com/android/customization/model/grid/data/repository/FakeGridRepository.kt
@@ -0,0 +1,87 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package com.android.customization.model.grid.data.repository
+
+import com.android.customization.model.grid.shared.model.GridOptionItemModel
+import com.android.customization.model.grid.shared.model.GridOptionItemsModel
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.channels.BufferOverflow
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.MutableSharedFlow
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.SharingStarted
+import kotlinx.coroutines.flow.asSharedFlow
+import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.flow.stateIn
+
+class FakeGridRepository(
+    private val scope: CoroutineScope,
+    initialOptionCount: Int,
+) : GridRepository {
+    private val _optionChanges =
+        MutableSharedFlow<Unit>(
+            replay = 1,
+            onBufferOverflow = BufferOverflow.DROP_OLDEST,
+        )
+    override val optionChanges: Flow<Unit> = _optionChanges.asSharedFlow()
+
+    private val selectedOptionIndex = MutableStateFlow(0)
+    private var options: GridOptionItemsModel = createOptions(count = initialOptionCount)
+
+    override suspend fun getOptions(): GridOptionItemsModel {
+        return options
+    }
+
+    fun setOptions(
+        count: Int,
+        selectedIndex: Int = 0,
+    ) {
+        options = createOptions(count, selectedIndex)
+        _optionChanges.tryEmit(Unit)
+    }
+
+    private fun createOptions(
+        count: Int,
+        selectedIndex: Int = 0,
+    ): GridOptionItemsModel {
+        selectedOptionIndex.value = selectedIndex
+        return GridOptionItemsModel.Loaded(
+            options =
+                buildList {
+                    repeat(times = count) { index ->
+                        add(
+                            GridOptionItemModel(
+                                name = "option_$index",
+                                cols = 4,
+                                rows = index * 2,
+                                isSelected =
+                                    selectedOptionIndex
+                                        .map { it == index }
+                                        .stateIn(
+                                            scope = scope,
+                                            started = SharingStarted.Eagerly,
+                                            initialValue = false,
+                                        ),
+                                onSelected = { selectedOptionIndex.value = index },
+                            )
+                        )
+                    }
+                }
+        )
+    }
+}
diff --git a/tests/src/com/android/customization/model/grid/domain/interactor/GridInteractorTest.kt b/tests/src/com/android/customization/model/grid/domain/interactor/GridInteractorTest.kt
new file mode 100644
index 0000000..20dd300
--- /dev/null
+++ b/tests/src/com/android/customization/model/grid/domain/interactor/GridInteractorTest.kt
@@ -0,0 +1,138 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package com.android.customization.model.grid.domain.interactor
+
+import androidx.test.filters.SmallTest
+import com.android.customization.model.grid.data.repository.FakeGridRepository
+import com.android.customization.model.grid.shared.model.GridOptionItemsModel
+import com.android.wallpaper.testing.FakeSnapshotStore
+import com.android.wallpaper.testing.collectLastValue
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.runBlocking
+import kotlinx.coroutines.test.TestScope
+import kotlinx.coroutines.test.runCurrent
+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 GridInteractorTest {
+
+    private lateinit var underTest: GridInteractor
+    private lateinit var testScope: TestScope
+    private lateinit var repository: FakeGridRepository
+    private lateinit var store: FakeSnapshotStore
+
+    @Before
+    fun setUp() {
+        testScope = TestScope()
+        repository =
+            FakeGridRepository(
+                scope = testScope.backgroundScope,
+                initialOptionCount = 3,
+            )
+        store = FakeSnapshotStore()
+        underTest =
+            GridInteractor(
+                applicationScope = testScope.backgroundScope,
+                repository = repository,
+                snapshotRestorer = {
+                    GridSnapshotRestorer(
+                            interactor = underTest,
+                        )
+                        .apply {
+                            runBlocking {
+                                setUpSnapshotRestorer(
+                                    store = store,
+                                )
+                            }
+                        }
+                },
+            )
+    }
+
+    @Test
+    fun selectingOptionThroughModel_updatesOptions() =
+        testScope.runTest {
+            val options = collectLastValue(underTest.options)
+            assertThat(options()).isInstanceOf(GridOptionItemsModel.Loaded::class.java)
+            (options() as? GridOptionItemsModel.Loaded)?.let { loaded ->
+                assertThat(loaded.options).hasSize(3)
+                assertThat(loaded.options[0].isSelected.value).isTrue()
+                assertThat(loaded.options[1].isSelected.value).isFalse()
+                assertThat(loaded.options[2].isSelected.value).isFalse()
+            }
+
+            val storedSnapshot = store.retrieve()
+            (options() as? GridOptionItemsModel.Loaded)?.let { loaded ->
+                loaded.options[1].onSelected()
+            }
+
+            assertThat(options()).isInstanceOf(GridOptionItemsModel.Loaded::class.java)
+            (options() as? GridOptionItemsModel.Loaded)?.let { loaded ->
+                assertThat(loaded.options).hasSize(3)
+                assertThat(loaded.options[0].isSelected.value).isFalse()
+                assertThat(loaded.options[1].isSelected.value).isTrue()
+                assertThat(loaded.options[2].isSelected.value).isFalse()
+            }
+            assertThat(store.retrieve()).isNotEqualTo(storedSnapshot)
+        }
+
+    @Test
+    fun selectingOptionThroughSetter_returnsSelectedOptionFromGetter() =
+        testScope.runTest {
+            val options = collectLastValue(underTest.options)
+            assertThat(options()).isInstanceOf(GridOptionItemsModel.Loaded::class.java)
+            (options() as? GridOptionItemsModel.Loaded)?.let { loaded ->
+                assertThat(loaded.options).hasSize(3)
+            }
+
+            val storedSnapshot = store.retrieve()
+            (options() as? GridOptionItemsModel.Loaded)?.let { loaded ->
+                underTest.setSelectedOption(loaded.options[1])
+                runCurrent()
+                assertThat(underTest.getSelectedOption()?.name).isEqualTo(loaded.options[1].name)
+                assertThat(store.retrieve()).isNotEqualTo(storedSnapshot)
+            }
+        }
+
+    @Test
+    fun externalUpdates_reloadInvoked() =
+        testScope.runTest {
+            val options = collectLastValue(underTest.options)
+            assertThat(options()).isInstanceOf(GridOptionItemsModel.Loaded::class.java)
+            (options() as? GridOptionItemsModel.Loaded)?.let { loaded ->
+                assertThat(loaded.options).hasSize(3)
+            }
+
+            val storedSnapshot = store.retrieve()
+            repository.setOptions(4)
+
+            assertThat(options()).isInstanceOf(GridOptionItemsModel.Loaded::class.java)
+            (options() as? GridOptionItemsModel.Loaded)?.let { loaded ->
+                assertThat(loaded.options).hasSize(4)
+            }
+            // External updates do not record a new snapshot with the undo system.
+            assertThat(store.retrieve()).isEqualTo(storedSnapshot)
+        }
+}
diff --git a/tests/src/com/android/customization/model/grid/domain/interactor/GridSnapshotRestorerTest.kt b/tests/src/com/android/customization/model/grid/domain/interactor/GridSnapshotRestorerTest.kt
new file mode 100644
index 0000000..c2712b1
--- /dev/null
+++ b/tests/src/com/android/customization/model/grid/domain/interactor/GridSnapshotRestorerTest.kt
@@ -0,0 +1,111 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package com.android.customization.model.grid.domain.interactor
+
+import androidx.test.filters.SmallTest
+import com.android.customization.model.grid.data.repository.FakeGridRepository
+import com.android.customization.model.grid.shared.model.GridOptionItemsModel
+import com.android.wallpaper.testing.FakeSnapshotStore
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.test.TestScope
+import kotlinx.coroutines.test.runCurrent
+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 GridSnapshotRestorerTest {
+
+    private lateinit var underTest: GridSnapshotRestorer
+    private lateinit var testScope: TestScope
+    private lateinit var repository: FakeGridRepository
+    private lateinit var store: FakeSnapshotStore
+
+    @Before
+    fun setUp() {
+        testScope = TestScope()
+        repository =
+            FakeGridRepository(
+                scope = testScope.backgroundScope,
+                initialOptionCount = 4,
+            )
+        underTest =
+            GridSnapshotRestorer(
+                interactor =
+                    GridInteractor(
+                        applicationScope = testScope.backgroundScope,
+                        repository = repository,
+                        snapshotRestorer = { underTest },
+                    )
+            )
+        store = FakeSnapshotStore()
+    }
+
+    @Test
+    fun restoreToSnapshot_noCallsToStore_restoresToInitialSnapshot() =
+        testScope.runTest {
+            runCurrent()
+            val initialSnapshot = underTest.setUpSnapshotRestorer(store = store)
+            assertThat(initialSnapshot.args).isNotEmpty()
+            repository.setOptions(
+                count = 4,
+                selectedIndex = 2,
+            )
+            runCurrent()
+            assertThat(getSelectedIndex()).isEqualTo(2)
+
+            underTest.restoreToSnapshot(initialSnapshot)
+            runCurrent()
+
+            assertThat(getSelectedIndex()).isEqualTo(0)
+        }
+
+    @Test
+    fun restoreToSnapshot_withCallToStore_restoresToInitialSnapshot() =
+        testScope.runTest {
+            runCurrent()
+            val initialSnapshot = underTest.setUpSnapshotRestorer(store = store)
+            assertThat(initialSnapshot.args).isNotEmpty()
+            repository.setOptions(
+                count = 4,
+                selectedIndex = 2,
+            )
+            runCurrent()
+            assertThat(getSelectedIndex()).isEqualTo(2)
+            underTest.store((repository.getOptions() as GridOptionItemsModel.Loaded).options[1])
+            runCurrent()
+
+            underTest.restoreToSnapshot(initialSnapshot)
+            runCurrent()
+
+            assertThat(getSelectedIndex()).isEqualTo(0)
+        }
+
+    private suspend fun getSelectedIndex(): Int {
+        return (repository.getOptions() as? GridOptionItemsModel.Loaded)?.options?.indexOfFirst {
+            optionItem ->
+            optionItem.isSelected.value
+        }
+            ?: -1
+    }
+}
diff --git a/tests/src/com/android/customization/model/grid/ui/viewmodel/GridScreenViewModelTest.kt b/tests/src/com/android/customization/model/grid/ui/viewmodel/GridScreenViewModelTest.kt
new file mode 100644
index 0000000..951638a
--- /dev/null
+++ b/tests/src/com/android/customization/model/grid/ui/viewmodel/GridScreenViewModelTest.kt
@@ -0,0 +1,104 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package com.android.customization.model.grid.ui.viewmodel
+
+import androidx.test.filters.SmallTest
+import androidx.test.platform.app.InstrumentationRegistry
+import com.android.customization.model.grid.data.repository.FakeGridRepository
+import com.android.customization.model.grid.domain.interactor.GridInteractor
+import com.android.customization.model.grid.domain.interactor.GridSnapshotRestorer
+import com.android.wallpaper.picker.option.ui.viewmodel.OptionItemViewModel
+import com.android.wallpaper.testing.FakeSnapshotStore
+import com.android.wallpaper.testing.collectLastValue
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.runBlocking
+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 GridScreenViewModelTest {
+
+    private lateinit var underTest: GridScreenViewModel
+    private lateinit var testScope: TestScope
+    private lateinit var interactor: GridInteractor
+    private lateinit var store: FakeSnapshotStore
+
+    @Before
+    fun setUp() {
+        testScope = TestScope()
+        store = FakeSnapshotStore()
+        interactor =
+            GridInteractor(
+                applicationScope = testScope.backgroundScope,
+                repository =
+                    FakeGridRepository(
+                        scope = testScope.backgroundScope,
+                        initialOptionCount = 4,
+                    ),
+                snapshotRestorer = {
+                    GridSnapshotRestorer(
+                            interactor = interactor,
+                        )
+                        .apply { runBlocking { setUpSnapshotRestorer(store) } }
+                }
+            )
+
+        underTest =
+            GridScreenViewModel(
+                context = InstrumentationRegistry.getInstrumentation().targetContext,
+                interactor = interactor,
+            )
+    }
+
+    @Test
+    fun clickOnItem_itGetsSelected() =
+        testScope.runTest {
+            val optionItemsValueProvider = collectLastValue(underTest.optionItems)
+            var optionItemsValue = checkNotNull(optionItemsValueProvider.invoke())
+            assertThat(optionItemsValue).hasSize(4)
+            assertThat(getSelectedIndex(optionItemsValue)).isEqualTo(0)
+            assertThat(getOnClick(optionItemsValue[0])).isNull()
+
+            val item1OnClickedValue = getOnClick(optionItemsValue[1])
+            assertThat(item1OnClickedValue).isNotNull()
+            item1OnClickedValue?.invoke()
+
+            optionItemsValue = checkNotNull(optionItemsValueProvider.invoke())
+            assertThat(optionItemsValue).hasSize(4)
+            assertThat(getSelectedIndex(optionItemsValue)).isEqualTo(1)
+            assertThat(getOnClick(optionItemsValue[0])).isNotNull()
+            assertThat(getOnClick(optionItemsValue[1])).isNull()
+        }
+
+    private fun TestScope.getSelectedIndex(optionItems: List<OptionItemViewModel>): Int {
+        return optionItems.indexOfFirst { optionItem ->
+            collectLastValue(optionItem.isSelected).invoke() == true
+        }
+    }
+
+    private fun TestScope.getOnClick(optionItem: OptionItemViewModel): (() -> Unit)? {
+        return collectLastValue(optionItem.onClicked).invoke()
+    }
+}