Snap for 6465939 from 3ca049371b56783d16b09b2aeca2e38c4684c79a to mainline-release

Change-Id: I4649a9fbc9a15e1ab38bf1912c1569d562156807
diff --git a/res/drawable/ic_check_circle_filled_24px.xml b/res/drawable/ic_check_circle_filled_24px.xml
index 01d2091..cacf1d8 100644
--- a/res/drawable/ic_check_circle_filled_24px.xml
+++ b/res/drawable/ic_check_circle_filled_24px.xml
@@ -19,7 +19,7 @@
     android:viewportWidth="24"
     android:viewportHeight="24">
   <path
-      android:fillColor="@color/accent_color"
+      android:fillColor="@color/selected_check_background_color"
       android:pathData="M12,2C6.5,2 2,6.5 2,12s4.5,10 10,10s10,-4.5 10,-10S17.5,2 12,2zM10,17l-4,-4l1.4,-1.4l2.6,2.6l6.6,-6.6L18,9L10,17z"/>
   <path
       android:pathData="m8.0085,14.9866 l-1.9939,-1.994 0.6892,-0.6889 0.6892,-0.6889 1.2925,1.2926c0.7109,0.711 1.3035,1.2926 1.3169,1.2926 0.0134,0 1.5034,-1.4789 3.3111,-3.2865l3.2866,-3.2865 0.689,0.689 0.689,0.689 -3.9878,3.9878 -3.9878,3.9878z"
diff --git a/res/layout-land/fragment_grid_picker.xml b/res/layout-land/fragment_grid_picker.xml
index 6c2d8c5..77330a8 100644
--- a/res/layout-land/fragment_grid_picker.xml
+++ b/res/layout-land/fragment_grid_picker.xml
@@ -39,18 +39,34 @@
             android:background="@color/secondary_color"
             app:card_style="screen_aspect_ratio"/>
 
-        <FrameLayout
+        <LinearLayout
             android:id="@+id/options_section"
             android:layout_width="0dp"
             android:layout_height="match_parent"
             android:layout_weight="1"
-            android:paddingVertical="10dp">
+            android:paddingVertical="10dp"
+            android:orientation="vertical">
 
             <androidx.recyclerview.widget.RecyclerView
                 android:id="@+id/options_container"
                 android:layout_width="match_parent"
-                android:layout_height="match_parent"/>
-        </FrameLayout>
+                android:layout_height="0dp"
+                android:layout_weight="1"/>
+
+            <RelativeLayout
+                android:layout_width="match_parent"
+                android:layout_height="wrap_content"
+                android:paddingHorizontal="10dp">
+                <Button
+                    android:id="@+id/apply_button"
+                    style="@style/ActionPrimaryButton"
+                    android:layout_width="wrap_content"
+                    android:layout_height="wrap_content"
+                    android:layout_alignParentEnd="true"
+                    android:text="@string/apply_btn"/>
+            </RelativeLayout>
+
+        </LinearLayout>
     </LinearLayout>
         <androidx.core.widget.ContentLoadingProgressBar
             android:id="@+id/loading_indicator"
diff --git a/res/layout-land/fragment_theme_picker.xml b/res/layout-land/fragment_theme_picker.xml
index dc30232..f2083f3 100644
--- a/res/layout-land/fragment_theme_picker.xml
+++ b/res/layout-land/fragment_theme_picker.xml
@@ -55,14 +55,30 @@
                     android:layout_height="0dp"
                     android:layout_weight="1"/>
 
-                <Button
-                    android:id="@+id/apply_button"
-                    style="@style/ActionPrimaryButton"
-                    android:layout_width="wrap_content"
+                <RelativeLayout
+                    android:layout_width="match_parent"
                     android:layout_height="wrap_content"
-                    android:layout_gravity="end"
-                    android:paddingHorizontal="10dp"
-                    android:text="@string/apply_theme_btn"/>
+                    android:background="@null"
+                    android:paddingHorizontal="10dp">
+                    <CheckBox
+                        android:id="@+id/use_my_wallpaper"
+                        android:layout_width="wrap_content"
+                        android:layout_height="wrap_content"
+                        android:layout_alignParentStart="true"
+                        android:layout_centerVertical="true"
+                        android:layout_toStartOf="@+id/apply_button"
+                        android:ellipsize="end"
+                        android:paddingLeft="4dp"
+                        android:minHeight="@dimen/min_taptarget_height"
+                        android:text="@string/keep_my_wallpaper"/>
+                    <Button
+                        android:id="@+id/apply_button"
+                        style="@style/ActionPrimaryButton"
+                        android:layout_width="wrap_content"
+                        android:layout_height="wrap_content"
+                        android:layout_alignParentEnd="true"
+                        android:text="@string/apply_theme_btn"/>
+                </RelativeLayout>
             </LinearLayout>
         </LinearLayout>
 
diff --git a/res/layout/fragment_grid_picker.xml b/res/layout/fragment_grid_picker.xml
index 5bf96e1..f664568 100644
--- a/res/layout/fragment_grid_picker.xml
+++ b/res/layout/fragment_grid_picker.xml
@@ -39,39 +39,49 @@
                 android:background="@color/preview_pager_background"
                 app:card_style="screen_aspect_ratio"
                 app:layout_constrainedHeight="true"
-                app:layout_constraintTop_toTopOf="parent"
-                app:layout_constraintBottom_toTopOf="@id/options_title"
-                app:layout_constraintStart_toStartOf="parent"
+                app:layout_constraintBottom_toTopOf="@id/options_container"
                 app:layout_constraintEnd_toEndOf="parent"
                 app:layout_constraintHeight_max="@dimen/preview_pager_max_height"
                 app:layout_constraintHeight_min="@dimen/preview_pager_min_height"
                 app:layout_constraintVertical_bias="0.0"
-                app:layout_constraintVertical_chainStyle="spread_inside"/>
-
-            <TextView
-                android:id="@+id/options_title"
-                android:layout_width="match_parent"
-                android:layout_height="wrap_content"
-                android:gravity="center"
-                android:paddingVertical="18dp"
-                android:lineHeight="24dp"
-                android:singleLine="true"
-                android:text="@string/grid_options_title"
-                android:textAppearance="@style/TitleTextAppearance"
-                app:layout_constraintTop_toBottomOf="@id/grid_preview_pager"
-                app:layout_constraintBottom_toTopOf="@id/options_container"
                 app:layout_constraintStart_toStartOf="parent"
-                app:layout_constraintEnd_toEndOf="parent"/>
+                app:layout_constraintTop_toTopOf="parent"
+                app:layout_constraintVertical_chainStyle="spread_inside"/>
 
             <androidx.recyclerview.widget.RecyclerView
                 android:id="@+id/options_container"
                 android:layout_width="match_parent"
                 android:layout_height="@dimen/options_container_height"
                 android:layout_gravity="center_horizontal"
-                app:layout_constraintTop_toBottomOf="@+id/options_title"
-                app:layout_constraintStart_toStartOf="parent"
+                android:layout_marginTop="10dp"
+                app:layout_constraintBottom_toTopOf="@id/placeholder"
                 app:layout_constraintEnd_toEndOf="parent"
+                app:layout_constraintStart_toStartOf="parent"
+                app:layout_constraintTop_toBottomOf="@+id/grid_preview_pager"
                 app:layout_constraintVertical_bias="1.0"/>
+
+            <Space
+                android:id="@+id/placeholder"
+                android:layout_width="match_parent"
+                android:layout_height="@dimen/min_taptarget_height"
+                app:layout_constraintBottom_toTopOf="@id/apply_button"
+                app:layout_constraintStart_toStartOf="parent"
+                app:layout_constraintTop_toBottomOf="@id/options_container"
+                app:layout_constraintVertical_bias="1.0"/>
+
+            <Button
+                android:id="@+id/apply_button"
+                style="@style/ActionPrimaryButton"
+                android:layout_width="wrap_content"
+                android:layout_height="wrap_content"
+                android:layout_gravity="end"
+                android:layout_marginEnd="10dp"
+                android:layout_marginVertical="10dp"
+                android:layout_weight="1"
+                android:text="@string/apply_theme_btn"
+                app:layout_constraintBottom_toBottomOf="parent"
+                app:layout_constraintEnd_toEndOf="parent"/>
+
         </androidx.constraintlayout.widget.ConstraintLayout>
 
         <androidx.core.widget.ContentLoadingProgressBar
diff --git a/res/layout/fragment_grid_scalable_picker.xml b/res/layout/fragment_grid_scalable_picker.xml
index 6a3a77a..de44923 100644
--- a/res/layout/fragment_grid_scalable_picker.xml
+++ b/res/layout/fragment_grid_scalable_picker.xml
@@ -44,29 +44,46 @@
         app:behavior_peekHeight="@dimen/content_container_minimum_height"
         app:layout_behavior="@string/bottom_sheet_behavior">
 
-        <LinearLayout
+        <androidx.constraintlayout.widget.ConstraintLayout
             android:id="@+id/content_section"
             android:layout_width="match_parent"
-            android:layout_height="match_parent"
-            android:orientation="vertical">
-
-            <TextView
-                android:id="@+id/options_title"
-                android:layout_width="match_parent"
-                android:layout_height="wrap_content"
-                android:gravity="center"
-                android:paddingVertical="18dp"
-                android:lineHeight="24dp"
-                android:singleLine="true"
-                android:text="@string/grid_options_title"
-                android:textAppearance="@style/TitleTextAppearance"/>
+            android:layout_height="match_parent">
 
             <androidx.recyclerview.widget.RecyclerView
                 android:id="@+id/options_container"
                 android:layout_width="match_parent"
                 android:layout_height="@dimen/options_container_height"
-                android:layout_gravity="center_horizontal"/>
-        </LinearLayout>
+                android:layout_gravity="center_horizontal"
+                android:layout_marginTop="10dp"
+                app:layout_constraintBottom_toTopOf="@id/placeholder"
+                app:layout_constraintEnd_toEndOf="parent"
+                app:layout_constraintStart_toStartOf="parent"
+                app:layout_constraintTop_toTopOf="parent"
+                app:layout_constraintVertical_bias="1.0" />
+
+            <Space
+                android:id="@+id/placeholder"
+                android:layout_width="match_parent"
+                android:layout_height="@dimen/min_taptarget_height"
+                app:layout_constraintBottom_toTopOf="@id/apply_button"
+                app:layout_constraintStart_toStartOf="parent"
+                app:layout_constraintTop_toBottomOf="@id/options_container"
+                app:layout_constraintVertical_bias="1.0" />
+
+            <Button
+                android:id="@+id/apply_button"
+                style="@style/ActionPrimaryButton"
+                android:layout_width="wrap_content"
+                android:layout_height="wrap_content"
+                android:layout_gravity="end"
+                android:layout_marginVertical="10dp"
+                android:layout_marginEnd="10dp"
+                android:layout_weight="1"
+                android:text="@string/apply_theme_btn"
+                app:layout_constraintBottom_toBottomOf="parent"
+                app:layout_constraintEnd_toEndOf="parent" />
+
+        </androidx.constraintlayout.widget.ConstraintLayout>
     </FrameLayout>
 
     <FrameLayout
diff --git a/res/layout/fragment_theme_picker.xml b/res/layout/fragment_theme_picker.xml
index d15d93f..1de5a3c 100644
--- a/res/layout/fragment_theme_picker.xml
+++ b/res/layout/fragment_theme_picker.xml
@@ -53,11 +53,29 @@
                 android:layout_gravity="bottom|center_horizontal"
                 android:layout_marginTop="10dp"
                 android:layout_weight="1"
+                app:layout_constraintBottom_toTopOf="@id/use_my_wallpaper"
                 app:layout_constraintEnd_toEndOf="parent"
                 app:layout_constraintStart_toStartOf="parent"
                 app:layout_constraintTop_toBottomOf="@+id/theme_preview_pager"
                 app:layout_constraintVertical_bias="1.0"/>
 
+            <CheckBox
+                android:id="@+id/use_my_wallpaper"
+                android:layout_width="wrap_content"
+                android:layout_height="@dimen/min_taptarget_height"
+                android:layout_marginStart="10dp"
+                android:ellipsize="end"
+                android:gravity="start|center_vertical"
+                android:paddingLeft="4dp"
+                android:text="@string/keep_my_wallpaper"
+                app:layout_constraintBottom_toTopOf="@id/apply_button"
+                app:layout_constraintEnd_toEndOf="parent"
+                app:layout_constraintHeight_min="@dimen/min_taptarget_height"
+                app:layout_constraintHorizontal_bias="0.0"
+                app:layout_constraintStart_toStartOf="parent"
+                app:layout_constraintTop_toBottomOf="@id/options_container"
+                app:layout_constraintVertical_bias="1.0"/>
+
             <Button
                 android:id="@+id/apply_button"
                 style="@style/ActionPrimaryButton"
diff --git a/res/layout/fragment_theme_scalable_picker.xml b/res/layout/fragment_theme_scalable_picker.xml
index ae147c2..e03df3e 100644
--- a/res/layout/fragment_theme_scalable_picker.xml
+++ b/res/layout/fragment_theme_scalable_picker.xml
@@ -55,11 +55,29 @@
                 android:layout_gravity="bottom|center_horizontal"
                 android:layout_marginTop="10dp"
                 android:layout_weight="1"
+                app:layout_constraintBottom_toTopOf="@id/use_my_wallpaper"
                 app:layout_constraintEnd_toEndOf="parent"
                 app:layout_constraintStart_toStartOf="parent"
                 app:layout_constraintTop_toTopOf="parent"
                 app:layout_constraintVertical_bias="1.0" />
 
+            <CheckBox
+                android:id="@+id/use_my_wallpaper"
+                android:layout_width="wrap_content"
+                android:layout_height="@dimen/min_taptarget_height"
+                android:layout_marginStart="10dp"
+                android:ellipsize="end"
+                android:gravity="start|center_vertical"
+                android:paddingLeft="4dp"
+                android:text="@string/keep_my_wallpaper"
+                app:layout_constraintBottom_toTopOf="@id/apply_button"
+                app:layout_constraintEnd_toEndOf="parent"
+                app:layout_constraintHeight_min="@dimen/min_taptarget_height"
+                app:layout_constraintHorizontal_bias="0.0"
+                app:layout_constraintStart_toStartOf="parent"
+                app:layout_constraintTop_toBottomOf="@id/options_container"
+                app:layout_constraintVertical_bias="1.0" />
+
             <Button
                 android:id="@+id/apply_button"
                 style="@style/ActionPrimaryButton"
diff --git a/res/layout/preview_card_wallpaper_content.xml b/res/layout/preview_card_wallpaper_content.xml
new file mode 100644
index 0000000..0b43c2f
--- /dev/null
+++ b/res/layout/preview_card_wallpaper_content.xml
@@ -0,0 +1,31 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+     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.
+-->
+<FrameLayout
+    xmlns:android="http://schemas.android.com/apk/res/android"
+    xmlns:tools="http://schemas.android.com/tools"
+    android:id="@+id/preview_static_image"
+    android:layout_width="wrap_content"
+    android:layout_height="wrap_content"
+    android:layout_gravity="bottom"
+    tools:showIn="@layout/theme_preview_card">
+    <TextView
+        style="@style/CardTitleTextAppearance"
+        android:id="@+id/wallpaper_description"
+        android:layout_width="wrap_content"
+        android:layout_height="wrap_content"
+        android:gravity="center_horizontal|bottom"/>
+</FrameLayout>
\ No newline at end of file
diff --git a/res/values/dimens.xml b/res/values/dimens.xml
index 3c78967..49daee4 100644
--- a/res/values/dimens.xml
+++ b/res/values/dimens.xml
@@ -21,15 +21,15 @@
     <dimen name="tip_dot_size">8dp</dimen>
     <dimen name="tip_dot_line_width">2dp</dimen>
 
-    <dimen name="check_size">16dp</dimen>
-    <dimen name="check_offset">4dp</dimen>
+    <dimen name="check_size">20dp</dimen>
+    <dimen name="check_offset">-4dp</dimen>
 
     <dimen name="component_preview_page_bottom_margin">16dp</dimen>
 
     <!-- Dimensions for the customization option tiles -->
-    <dimen name="options_container_height">120dp</dimen>
+    <dimen name="options_container_height">104dp</dimen>
     <dimen name="options_container_width">0dp</dimen>
-    <dimen name="option_tile_width">88dp</dimen>
+    <dimen name="option_tile_width">72dp</dimen>
     <dimen name="theme_option_icon_sample_height">18dp</dimen>
     <dimen name="theme_option_icon_sample_width">18dp</dimen>
     <dimen name="theme_option_shape_sample_height">16dp</dimen>
@@ -38,8 +38,8 @@
     <dimen name="theme_option_font_sample_width">52dp</dimen>
     <dimen name="theme_option_sample_padding">5dp</dimen>
     <!-- Note, using dp instead of sp as this is just the "+" symbol, not text -->
-    <dimen name="option_tile_padding_vertical">16dp</dimen>
-    <dimen name="option_tile_padding_horizontal">16dp</dimen>
+    <dimen name="option_tile_padding_vertical">12dp</dimen>
+    <dimen name="option_tile_padding_horizontal">10dp</dimen>
 
     <dimen name="option_bottom_margin">8dp</dimen>
     <!-- Note, using dp instead of sp as this text is more like a "snapshot" of the font -->
diff --git a/res/values/strings.xml b/res/values/strings.xml
index c0aef7a..c940efd 100755
--- a/res/values/strings.xml
+++ b/res/values/strings.xml
@@ -67,7 +67,7 @@
 
     <!-- Plus sign used to indicate that the user can add a custom theme -->
 
-    <!-- Name for the Android Theme (Style/Wallpaper/Grid/Clock...) that comes preset with the device [CHAR LIMIT=10]-->
+    <!-- Name for the Android Theme that comes preset with the device [CHAR LIMIT=10]-->
     <string name="default_theme_title">Default</string>
 
     <!-- Title of a page that shows the user the preview of a font selection [CHAR_LIMIT=20] -->
@@ -92,20 +92,14 @@
     <!-- Body text for previewing a font [CHAR LIMIT=160] -->
     <string name="font_card_body">Add your favorite fonts to every screen</string>
 
-    <!--Title for grid options [CHAR LIMIT=NONE] -->
-    <string name="grid_options_title">Choose a grid size</string>
-
     <!--Title for a grid option, describing the number of columns and rows, eg: 4x4 [CHAR LIMIT=10] -->
     <string name="grid_title_pattern"><xliff:g name="num_cols" example="1">%1$d</xliff:g>x<xliff:g name="num_rows" example="1">%2$d</xliff:g></string>
 
     <!-- Message shown when a theme has been applied successfully in the system [CHAR LIMIT=NONE] -->
-    <string name="applied_theme_msg">Style set successfully</string>
+    <string name="applied_theme_msg">Style applied</string>
 
     <!-- Message shown when a clock has been applied successfully in the system [CHAR LIMIT=NONE] -->
-    <string name="applied_clock_msg">Clock set successfully</string>
-
-    <!-- Message shown when a grid has been applied successfully in the system [CHAR LIMIT=NONE] -->
-    <string name="applied_grid_msg">Grid set successfully</string>
+    <string name="applied_clock_msg">Clock applied</string>
 
     <!-- Message shown when a theme couldn't be applied in the system because of an error
         [CHAR LIMIT=NONE] -->
diff --git a/robolectric_tests/src/com/android/customization/model/theme/ThemeManagerTest.java b/robolectric_tests/src/com/android/customization/model/theme/ThemeManagerTest.java
index cfb8a33..6f1f95e 100644
--- a/robolectric_tests/src/com/android/customization/model/theme/ThemeManagerTest.java
+++ b/robolectric_tests/src/com/android/customization/model/theme/ThemeManagerTest.java
@@ -40,11 +40,12 @@
 import androidx.annotation.Nullable;
 import androidx.fragment.app.FragmentActivity;
 
-import com.android.customization.model.CustomizationManager.Callback;
 import com.android.customization.model.CustomizationManager.OptionsFetchedListener;
+import com.android.customization.model.CustomizationManager.Callback;
 import com.android.customization.model.theme.custom.CustomTheme;
 import com.android.customization.module.ThemesUserEventLogger;
 import com.android.customization.testutils.OverlayManagerMocks;
+import com.android.wallpaper.module.WallpaperSetter;
 
 import org.json.JSONObject;
 import org.junit.After;
@@ -60,6 +61,7 @@
 public class ThemeManagerTest {
 
     @Mock OverlayManagerCompat mMockOm;
+    @Mock WallpaperSetter mMockWallpaperSetter;
     @Mock ThemesUserEventLogger mThemesUserEventLogger;
     @Mock ThemeBundleProvider mThemeBundleProvider;
     private OverlayManagerMocks mMockOmHelper;
@@ -73,7 +75,7 @@
         mActivity = spy(activity);
         mMockOmHelper = new OverlayManagerMocks();
         mMockOmHelper.setUpMock(mMockOm);
-        mThemeManager = new ThemeManager(mThemeBundleProvider, activity,
+        mThemeManager = new ThemeManager(mThemeBundleProvider, activity, mMockWallpaperSetter,
                 mMockOm, mThemesUserEventLogger);
     }
 
diff --git a/src/com/android/customization/model/grid/GridOptionsManager.java b/src/com/android/customization/model/grid/GridOptionsManager.java
index 84f5373..a334f5e 100644
--- a/src/com/android/customization/model/grid/GridOptionsManager.java
+++ b/src/com/android/customization/model/grid/GridOptionsManager.java
@@ -58,7 +58,7 @@
 
     @Override
     public void fetchOptions(OptionsFetchedListener<GridOption> callback, boolean reload) {
-        new FetchTask(mProvider, callback, reload).execute();
+        new FetchTask(mProvider, callback).execute();
     }
 
     /** See if using surface view to render grid options */
@@ -67,25 +67,23 @@
     }
 
     /** Call through content provider API to render preview */
-    public Bundle renderPreview(Bundle bundle, String gridName) {
-        return mProvider.renderPreview(gridName, bundle);
+    public void renderPreview(Bundle bundle, String gridName) {
+        mProvider.renderPreview(gridName, bundle);
     }
 
     private static class FetchTask extends AsyncTask<Void, Void, Pair<List<GridOption>, String>> {
         private final LauncherGridOptionsProvider mProvider;
         @Nullable private final OptionsFetchedListener<GridOption> mCallback;
-        private final boolean mReload;
 
         private FetchTask(@NonNull LauncherGridOptionsProvider provider,
-                @Nullable OptionsFetchedListener<GridOption> callback, boolean reload) {
+                @Nullable OptionsFetchedListener<GridOption> callback) {
             mCallback = callback;
             mProvider = provider;
-            mReload = reload;
         }
 
         @Override
         protected Pair<List<GridOption>, String> doInBackground(Void[] params) {
-            return mProvider.fetch(mReload);
+            return mProvider.fetch(false);
         }
 
         @Override
diff --git a/src/com/android/customization/model/grid/LauncherGridOptionsProvider.java b/src/com/android/customization/model/grid/LauncherGridOptionsProvider.java
index 6a47f95..7d1d5cd 100644
--- a/src/com/android/customization/model/grid/LauncherGridOptionsProvider.java
+++ b/src/com/android/customization/model/grid/LauncherGridOptionsProvider.java
@@ -29,6 +29,7 @@
 import androidx.annotation.WorkerThread;
 
 import com.android.customization.model.ResourceConstants;
+import com.android.systemui.shared.system.SurfaceViewRequestUtils;
 import com.android.wallpaper.R;
 import com.android.wallpaper.util.PreviewUtils;
 
@@ -52,9 +53,6 @@
     private static final String COL_PREVIEW_COUNT = "preview_count";
     private static final String COL_IS_DEFAULT = "is_default";
 
-    // Normal gird size name
-    private static final String GRID_NAME_NORMAL = "normal";
-
     private static final String METADATA_KEY_PREVIEW_VERSION = "preview_version";
 
     private final Context mContext;
@@ -102,9 +100,7 @@
                 int cols = c.getInt(c.getColumnIndex(COL_COLS));
                 int previewCount = c.getInt(c.getColumnIndex(COL_PREVIEW_COUNT));
                 boolean isSet = Boolean.valueOf(c.getString(c.getColumnIndex(COL_IS_DEFAULT)));
-                String title = GRID_NAME_NORMAL.equals(name)
-                        ? mContext.getString(R.string.default_theme_title)
-                        : mContext.getString(R.string.grid_title_pattern, cols, rows);
+                String title = mContext.getString(R.string.grid_title_pattern, cols, rows);
                 mOptions.add(new GridOption(title, name, isSet, rows, cols,
                         mPreviewUtils.getUri(PREVIEW), previewCount, iconPath));
             }
@@ -120,11 +116,11 @@
      * Request rendering of home screen preview via Launcher to Wallpaper using SurfaceView
      * @param name      the grid option name
      * @param bundle    surface view request bundle generated from
-     *    {@link com.android.wallpaper.util.SurfaceViewUtils#createSurfaceViewRequest(SurfaceView)}.
+     *                  {@link SurfaceViewRequestUtils#createSurfaceBundle(SurfaceView)}.
      */
-    Bundle renderPreview(String name, Bundle bundle) {
+    void renderPreview(String name, Bundle bundle) {
         bundle.putString("name", name);
-        return mPreviewUtils.renderPreview(bundle);
+        mPreviewUtils.renderPreview(bundle);
     }
 
     int applyGrid(String name) {
diff --git a/src/com/android/customization/model/theme/DefaultThemeProvider.java b/src/com/android/customization/model/theme/DefaultThemeProvider.java
index 7c99f9c..ce12d35 100644
--- a/src/com/android/customization/model/theme/DefaultThemeProvider.java
+++ b/src/com/android/customization/model/theme/DefaultThemeProvider.java
@@ -15,6 +15,8 @@
  */
 package com.android.customization.model.theme;
 
+import static android.content.res.Resources.ID_NULL;
+
 import static com.android.customization.model.ResourceConstants.ANDROID_PACKAGE;
 import static com.android.customization.model.ResourceConstants.ICONS_FOR_PREVIEW;
 import static com.android.customization.model.ResourceConstants.OVERLAY_CATEGORY_COLOR;
@@ -27,9 +29,14 @@
 import static com.android.customization.model.ResourceConstants.OVERLAY_CATEGORY_SHAPE;
 import static com.android.customization.model.ResourceConstants.SYSUI_PACKAGE;
 
+import android.content.ComponentName;
 import android.content.Context;
+import android.content.Intent;
+import android.content.pm.PackageManager;
 import android.content.pm.PackageManager.NameNotFoundException;
+import android.content.pm.ResolveInfo;
 import android.content.res.Resources.NotFoundException;
+import android.service.wallpaper.WallpaperService;
 import android.text.TextUtils;
 import android.util.Log;
 
@@ -42,6 +49,7 @@
 import com.android.customization.module.CustomizationPreferences;
 import com.android.wallpaper.R;
 import com.android.wallpaper.asset.ResourceAsset;
+import com.android.wallpaper.model.LiveWallpaperInfo;
 
 import com.bumptech.glide.request.RequestOptions;
 import com.google.android.apps.wallpaper.asset.ThemeBundleThumbAsset;
@@ -49,7 +57,9 @@
 import org.json.JSONArray;
 import org.json.JSONException;
 import org.json.JSONObject;
+import org.xmlpull.v1.XmlPullParserException;
 
+import java.io.IOException;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.Iterator;
@@ -73,6 +83,12 @@
     private static final String ICON_THEMEPICKER_PREFIX = "theme_overlay_icon_themepicker_";
     private static final String ICON_SETTINGS_PREFIX = "theme_overlay_icon_settings_";
     private static final String ICON_SYSUI_PREFIX = "theme_overlay_icon_sysui_";
+    private static final String WALLPAPER_PREFIX = "theme_wallpaper_";
+    private static final String WALLPAPER_TITLE_PREFIX = "theme_wallpaper_title_";
+    private static final String WALLPAPER_ATTRIBUTION_PREFIX = "theme_wallpaper_attribution_";
+    private static final String WALLPAPER_THUMB_PREFIX = "theme_wallpaper_thumbnail_";
+    private static final String WALLPAPER_ACTION_PREFIX = "theme_wallpaper_action_";
+    private static final String WALLPAPER_OPTIONS_PREFIX = "theme_wallpaper_options_";
 
     private static final String DEFAULT_THEME_NAME= "default";
     private static final String THEME_TITLE_FIELD = "_theme_title";
@@ -156,6 +172,8 @@
 
                 mOverlayProvider.addNoPreviewIconOverlay(builder, iconSettingsOverlayPackage);
 
+                addWallpaper(themeName, builder);
+
                 mThemes.add(builder.build(mContext));
             } catch (NameNotFoundException | NotFoundException e) {
                 Log.w(TAG, String.format("Couldn't load part of theme %s, will skip it", themeName),
@@ -166,6 +184,70 @@
         addCustomThemes();
     }
 
+    private void addWallpaper(String themeName, Builder builder) {
+        try {
+            String wallpaperResName = WALLPAPER_PREFIX + themeName;
+            int wallpaperResId = mStubApkResources.getIdentifier(wallpaperResName,
+                    "drawable", mStubPackageName);
+            // Check in case the theme has a separate thumbnail for the wallpaper
+            String wallpaperThumbnailResName = WALLPAPER_THUMB_PREFIX + themeName;
+            int wallpaperThumbnailResId = mStubApkResources.getIdentifier(wallpaperThumbnailResName,
+                    "drawable", mStubPackageName);
+            if (wallpaperResId != ID_NULL) {
+                builder.setWallpaperInfo(mStubPackageName, wallpaperResName,
+                        themeName, wallpaperResId,
+                        mStubApkResources.getIdentifier(WALLPAPER_TITLE_PREFIX + themeName,
+                                "string", mStubPackageName),
+                        mStubApkResources.getIdentifier(
+                                WALLPAPER_ATTRIBUTION_PREFIX + themeName, "string",
+                                mStubPackageName),
+                        mStubApkResources.getIdentifier(WALLPAPER_ACTION_PREFIX + themeName,
+                                "string", mStubPackageName))
+                        .setWallpaperAsset(wallpaperThumbnailResId != ID_NULL ?
+                                getThumbAsset(WALLPAPER_THUMB_PREFIX, themeName)
+                                : getDrawableResourceAsset(WALLPAPER_PREFIX, themeName));
+            } else {
+                // Try to see if it's a live wallpaper reference
+                wallpaperResId = mStubApkResources.getIdentifier(wallpaperResName,
+                        "string", mStubPackageName);
+                if (wallpaperResId != ID_NULL) {
+                    String wpComponent = mStubApkResources.getString(wallpaperResId);
+
+                    int wallpaperOptionsResId = mStubApkResources.getIdentifier(
+                            WALLPAPER_OPTIONS_PREFIX + themeName, "string", mStubPackageName);
+                    String wallpaperOptions = wallpaperOptionsResId != ID_NULL
+                            ? mStubApkResources.getString(wallpaperOptionsResId) : null;
+
+                    String[] componentParts = wpComponent.split("/");
+                    Intent liveWpIntent =  new Intent(WallpaperService.SERVICE_INTERFACE);
+                    liveWpIntent.setComponent(
+                            new ComponentName(componentParts[0], componentParts[1]));
+
+                    Context appContext = mContext.getApplicationContext();
+                    PackageManager pm = appContext.getPackageManager();
+                    ResolveInfo resolveInfo =
+                            pm.resolveService(liveWpIntent, PackageManager.GET_META_DATA);
+                    if (resolveInfo != null) {
+                        android.app.WallpaperInfo wallpaperInfo;
+                        try {
+                            wallpaperInfo = new android.app.WallpaperInfo(appContext, resolveInfo);
+                            LiveWallpaperInfo liveInfo = new LiveWallpaperInfo(wallpaperInfo);
+                            builder.setLiveWallpaperInfo(liveInfo).setWallpaperAsset(
+                                    wallpaperThumbnailResId != ID_NULL ?
+                                            getThumbAsset(WALLPAPER_THUMB_PREFIX, themeName)
+                                        : liveInfo.getThumbAsset(mContext))
+                                    .setWallpaperOptions(wallpaperOptions);
+                        } catch (XmlPullParserException | IOException e) {
+                            Log.w(TAG, "Skipping wallpaper " + resolveInfo.serviceInfo, e);
+                        }
+                    }
+                }
+            }
+        } catch (NotFoundException e) {
+            // Nothing to do here, if there's no wallpaper we'll just omit wallpaper
+        }
+    }
+
     /**
      * Default theme requires different treatment: if there are overlay packages specified in the
      * stub apk, we'll use those, otherwise we'll get the System default values. But we cannot skip
@@ -236,6 +318,8 @@
             mOverlayProvider.addSystemDefaultIcons(builder, SYSUI_PACKAGE, ICONS_FOR_PREVIEW);
         }
 
+        addWallpaper(DEFAULT_THEME_NAME, builder);
+
         mThemes.add(builder.build(mContext));
     }
 
diff --git a/src/com/android/customization/model/theme/ThemeBundle.java b/src/com/android/customization/model/theme/ThemeBundle.java
index 14b2dcd..e83e852 100644
--- a/src/com/android/customization/model/theme/ThemeBundle.java
+++ b/src/com/android/customization/model/theme/ThemeBundle.java
@@ -31,6 +31,7 @@
 import android.graphics.drawable.Drawable;
 import android.graphics.drawable.ShapeDrawable;
 import android.graphics.drawable.shapes.PathShape;
+import android.icu.text.SimpleDateFormat;
 import android.text.TextUtils;
 import android.util.Log;
 import android.view.View;
@@ -39,7 +40,9 @@
 
 import androidx.annotation.ColorInt;
 import androidx.annotation.Dimension;
+import androidx.annotation.DrawableRes;
 import androidx.annotation.Nullable;
+import androidx.annotation.StringRes;
 import androidx.core.graphics.PathParser;
 
 import com.android.customization.model.CustomizationManager;
@@ -48,6 +51,7 @@
 import com.android.wallpaper.R;
 import com.android.wallpaper.asset.Asset;
 import com.android.wallpaper.asset.BitmapCachingAsset;
+import com.android.wallpaper.model.LiveWallpaperInfo;
 import com.android.wallpaper.model.WallpaperInfo;
 
 import org.json.JSONException;
@@ -55,10 +59,12 @@
 
 import java.util.ArrayList;
 import java.util.Collections;
+import java.util.Date;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Iterator;
 import java.util.List;
+import java.util.Locale;
 import java.util.Map;
 import java.util.Set;
 
@@ -77,15 +83,20 @@
     private final PreviewInfo mPreviewInfo;
     private final boolean mIsDefault;
     protected final Map<String, String> mPackagesByCategory;
+    @Nullable private final WallpaperInfo mWallpaperInfo;
+    @Nullable private final String mWallpaperOptions;
     private WallpaperInfo mOverrideWallpaper;
     private Asset mOverrideWallpaperAsset;
     private CharSequence mContentDescription;
 
     protected ThemeBundle(String title, Map<String, String> overlayPackages,
-            boolean isDefault, PreviewInfo previewInfo) {
+            boolean isDefault, @Nullable WallpaperInfo wallpaperInfo,
+            @Nullable String wallpaperOptions, PreviewInfo previewInfo) {
         mTitle = title;
         mIsDefault = isDefault;
         mPreviewInfo = previewInfo;
+        mWallpaperInfo = wallpaperInfo;
+        mWallpaperOptions = wallpaperOptions;
         mPackagesByCategory = Collections.unmodifiableMap(overlayPackages);
     }
 
@@ -158,6 +169,16 @@
         mOverrideWallpaperAsset = null;
     }
 
+    public boolean shouldUseThemeWallpaper() {
+        return mOverrideWallpaper == null && mWallpaperInfo != null;
+    }
+
+    public Asset getWallpaperPreviewAsset(Context context) {
+        return mOverrideWallpaper != null ?
+                getOverrideWallpaperAsset(context) :
+                getPreviewInfo().wallpaperAsset;
+    }
+
     private Asset getOverrideWallpaperAsset(Context context) {
         if (mOverrideWallpaperAsset == null) {
             mOverrideWallpaperAsset = new BitmapCachingAsset(context,
@@ -166,6 +187,15 @@
         return mOverrideWallpaperAsset;
     }
 
+    public WallpaperInfo getWallpaperInfo() {
+        return mWallpaperInfo;
+    }
+
+    @Nullable
+    public String getWallpaperOptions() {
+        return mWallpaperOptions;
+    }
+
     boolean isDefault() {
         return mIsDefault;
     }
@@ -250,13 +280,14 @@
         @ColorInt public final int colorAccentDark;
         public final List<Drawable> icons;
         public final Drawable shapeDrawable;
+        @Nullable public final Asset wallpaperAsset;
         public final List<Drawable> shapeAppIcons;
         @Dimension public final int bottomSheeetCornerRadius;
 
         private PreviewInfo(Context context, Typeface bodyFontFamily, Typeface headlineFontFamily,
                 int colorAccentLight, int colorAccentDark, List<Drawable> icons,
                 Drawable shapeDrawable, @Dimension int cornerRadius,
-                List<Drawable> shapeAppIcons) {
+                @Nullable Asset wallpaperAsset, List<Drawable> shapeAppIcons) {
             this.bodyFontFamily = bodyFontFamily;
             this.headlineFontFamily = headlineFontFamily;
             this.colorAccentLight = colorAccentLight;
@@ -264,6 +295,8 @@
             this.icons = icons;
             this.shapeDrawable = shapeDrawable;
             this.bottomSheeetCornerRadius = cornerRadius;
+            this.wallpaperAsset = wallpaperAsset == null
+                    ? null : new BitmapCachingAsset(context, wallpaperAsset);
             this.shapeAppIcons = shapeAppIcons;
         }
 
@@ -290,11 +323,15 @@
         private Path mShapePath;
         private boolean mIsDefault;
         @Dimension private int mCornerRadius;
+        private Asset mWallpaperAsset;
+        private WallpaperInfo mWallpaperInfo;
+        private String mWallpaperOptions;
         protected Map<String, String> mPackages = new HashMap<>();
         private List<Drawable> mAppIcons = new ArrayList<>();
 
         public ThemeBundle build(Context context) {
-            return new ThemeBundle(mTitle, mPackages, mIsDefault, createPreviewInfo(context));
+            return new ThemeBundle(mTitle, mPackages, mIsDefault, mWallpaperInfo, mWallpaperOptions,
+                    createPreviewInfo(context));
         }
 
         public PreviewInfo createPreviewInfo(Context context) {
@@ -322,7 +359,8 @@
                 }
             }
             return new PreviewInfo(context, mBodyFontFamily, mHeadlineFontFamily, mColorAccentLight,
-                    mColorAccentDark, mIcons, shapeDrawable, mCornerRadius, shapeIcons);
+                    mColorAccentDark, mIcons, shapeDrawable, mCornerRadius,
+                    mWallpaperAsset, shapeIcons);
         }
 
         public Map<String, String> getPackages() {
@@ -378,6 +416,30 @@
             return this;
         }
 
+        public Builder setWallpaperInfo(String wallpaperPackageName, String wallpaperResName,
+                String themeId, @DrawableRes int wallpaperResId, @StringRes int titleResId,
+                @StringRes int attributionResId, @StringRes int actionUrlResId) {
+            mWallpaperInfo = new ThemeBundledWallpaperInfo(wallpaperPackageName, wallpaperResName,
+                    themeId, wallpaperResId, titleResId, attributionResId, actionUrlResId);
+            return this;
+        }
+
+        public Builder setLiveWallpaperInfo(LiveWallpaperInfo info) {
+            mWallpaperInfo = info;
+            return this;
+        }
+
+
+        public Builder setWallpaperAsset(Asset wallpaperAsset) {
+            mWallpaperAsset = wallpaperAsset;
+            return this;
+        }
+
+        public Builder setWallpaperOptions(String wallpaperOptions) {
+            mWallpaperOptions = wallpaperOptions;
+            return this;
+        }
+
         public Builder asDefault() {
             mIsDefault = true;
             return this;
diff --git a/src/com/android/customization/model/theme/ThemeManager.java b/src/com/android/customization/model/theme/ThemeManager.java
index 533fbd0..79b45fd 100644
--- a/src/com/android/customization/model/theme/ThemeManager.java
+++ b/src/com/android/customization/model/theme/ThemeManager.java
@@ -24,6 +24,7 @@
 import static com.android.customization.model.ResourceConstants.OVERLAY_CATEGORY_ICON_THEMEPICKER;
 import static com.android.customization.model.ResourceConstants.OVERLAY_CATEGORY_SHAPE;
 
+import android.graphics.Point;
 import android.provider.Settings;
 import android.text.TextUtils;
 
@@ -34,6 +35,13 @@
 import com.android.customization.model.ResourceConstants;
 import com.android.customization.model.theme.custom.CustomTheme;
 import com.android.customization.module.ThemesUserEventLogger;
+import com.android.wallpaper.R;
+import com.android.wallpaper.asset.Asset;
+import com.android.wallpaper.model.LiveWallpaperInfo;
+import com.android.wallpaper.model.WallpaperInfo;
+import com.android.wallpaper.module.WallpaperPersister.SetWallpaperCallback;
+import com.android.wallpaper.module.WallpaperSetter;
+import com.android.wallpaper.util.WallpaperCropUtils;
 
 import org.json.JSONObject;
 
@@ -53,22 +61,25 @@
         THEME_CATEGORIES.add(OVERLAY_CATEGORY_ICON_SYSUI);
         THEME_CATEGORIES.add(OVERLAY_CATEGORY_ICON_LAUNCHER);
         THEME_CATEGORIES.add(OVERLAY_CATEGORY_ICON_THEMEPICKER);
-    }
+    };
+
 
     private final ThemeBundleProvider mProvider;
     private final OverlayManagerCompat mOverlayManagerCompat;
 
+    private final WallpaperSetter mWallpaperSetter;
     protected final FragmentActivity mActivity;
     private final ThemesUserEventLogger mEventLogger;
 
     private Map<String, String> mCurrentOverlays;
 
     public ThemeManager(ThemeBundleProvider provider, FragmentActivity activity,
-            OverlayManagerCompat overlayManagerCompat,
+            WallpaperSetter wallpaperSetter, OverlayManagerCompat overlayManagerCompat,
             ThemesUserEventLogger logger) {
         mProvider = provider;
         mActivity = activity;
         mOverlayManagerCompat = overlayManagerCompat;
+        mWallpaperSetter = wallpaperSetter;
         mEventLogger = logger;
     }
 
@@ -79,7 +90,67 @@
 
     @Override
     public void apply(ThemeBundle theme, Callback callback) {
-        applyOverlays(theme, callback);
+        // Set wallpaper
+        if (theme.shouldUseThemeWallpaper()) {
+            mWallpaperSetter.requestDestination(mActivity, mActivity.getSupportFragmentManager(),
+                    R.string.set_theme_wallpaper_dialog_message,
+                    destination -> applyWallpaper(
+                            theme,
+                            destination,
+                            createSetWallpaperCallback(theme, callback)),
+                    theme.getWallpaperInfo() instanceof LiveWallpaperInfo);
+
+        } else {
+            applyOverlays(theme, callback);
+        }
+    }
+
+    private SetWallpaperCallback createSetWallpaperCallback(ThemeBundle theme, Callback callback) {
+        return new SetWallpaperCallback() {
+            @Override
+            public void onSuccess(WallpaperInfo wallpaperInfo) {
+                applyWallpaperOptions(theme);
+                applyOverlays(theme, callback);
+            }
+
+            @Override
+            public void onError(@Nullable Throwable throwable) {
+                callback.onError(throwable);
+            }
+        };
+    }
+
+    protected void applyWallpaperOptions(ThemeBundle theme) {
+        //Do nothing.
+    }
+
+    private void applyWallpaper(ThemeBundle theme, int destination,
+            SetWallpaperCallback callback) {
+        Point defaultCropSurfaceSize = WallpaperCropUtils.getDefaultCropSurfaceSize(
+                mActivity.getResources(),
+                mActivity.getWindowManager().getDefaultDisplay());
+        Asset wallpaperAsset = theme.getWallpaperInfo().getAsset(mActivity);
+        if (wallpaperAsset != null) {
+            wallpaperAsset.decodeRawDimensions(mActivity,
+                    dimensions -> {
+                        float scale = 1f;
+                        // Calculate scale to fit the screen height
+                        if (dimensions != null && dimensions.y > 0) {
+                            scale = (float) defaultCropSurfaceSize.y / dimensions.y;
+                        }
+                        mWallpaperSetter.setCurrentWallpaper(mActivity,
+                                theme.getWallpaperInfo(),
+                                wallpaperAsset,
+                                destination,
+                                scale, null, callback);
+                    });
+        } else {
+            mWallpaperSetter.setCurrentWallpaper(mActivity,
+                    theme.getWallpaperInfo(),
+                    null,
+                    destination,
+                    1f, null, callback);
+        }
     }
 
     private void applyOverlays(ThemeBundle theme, Callback callback) {
diff --git a/src/com/android/customization/model/theme/custom/CustomTheme.java b/src/com/android/customization/model/theme/custom/CustomTheme.java
index 9c14f01..8a7ca05 100644
--- a/src/com/android/customization/model/theme/custom/CustomTheme.java
+++ b/src/com/android/customization/model/theme/custom/CustomTheme.java
@@ -41,7 +41,7 @@
 
     public CustomTheme(@NonNull String id, String title, Map<String, String> overlayPackages,
             @Nullable PreviewInfo previewInfo) {
-        super(title, overlayPackages, false, previewInfo);
+        super(title, overlayPackages, false, null, null, previewInfo);
         mId = id;
     }
 
@@ -76,6 +76,11 @@
     }
 
     @Override
+    public boolean shouldUseThemeWallpaper() {
+        return false;
+    }
+
+    @Override
     public boolean isActive(CustomizationManager<ThemeBundle> manager) {
         return isDefined() && super.isActive(manager);
     }
diff --git a/src/com/android/customization/module/CustomizationInjector.java b/src/com/android/customization/module/CustomizationInjector.java
index 85853de..d25f925 100644
--- a/src/com/android/customization/module/CustomizationInjector.java
+++ b/src/com/android/customization/module/CustomizationInjector.java
@@ -23,11 +23,13 @@
 import com.android.customization.model.theme.ThemeBundleProvider;
 import com.android.customization.model.theme.ThemeManager;
 import com.android.wallpaper.module.Injector;
+import com.android.wallpaper.module.WallpaperSetter;
 
 public interface CustomizationInjector extends Injector {
 
     CustomizationPreferences getCustomizationPreferences(Context context);
 
     ThemeManager getThemeManager(ThemeBundleProvider provider, FragmentActivity activity,
-            OverlayManagerCompat overlayManagerCompat, ThemesUserEventLogger logger);
+            WallpaperSetter wallpaperSetter, OverlayManagerCompat overlayManagerCompat,
+            ThemesUserEventLogger logger);
 }
diff --git a/src/com/android/customization/module/DefaultCustomizationInjector.java b/src/com/android/customization/module/DefaultCustomizationInjector.java
index 7358d8f..a0d435b 100644
--- a/src/com/android/customization/module/DefaultCustomizationInjector.java
+++ b/src/com/android/customization/module/DefaultCustomizationInjector.java
@@ -30,6 +30,7 @@
 import com.android.wallpaper.module.LoggingOptInStatusProvider;
 import com.android.wallpaper.module.WallpaperPreferences;
 import com.android.wallpaper.module.WallpaperRotationRefresher;
+import com.android.wallpaper.module.WallpaperSetter;
 import com.android.wallpaper.monitor.PerformanceMonitor;
 import com.android.wallpaper.picker.PreviewFragment;
 
@@ -49,6 +50,7 @@
         return mPrefs;
     }
 
+
     @Override
     public CustomizationPreferences getCustomizationPreferences(Context context) {
         return (CustomizationPreferences) getPreferences(context);
@@ -113,8 +115,9 @@
 
     @Override
     public ThemeManager getThemeManager(ThemeBundleProvider provider, FragmentActivity activity,
-            OverlayManagerCompat overlayManagerCompat, ThemesUserEventLogger logger) {
-        return new ThemeManager(provider, activity, overlayManagerCompat, logger);
+            WallpaperSetter wallpaperSetter, OverlayManagerCompat overlayManagerCompat,
+            ThemesUserEventLogger logger) {
+        return new ThemeManager(provider, activity, wallpaperSetter, overlayManagerCompat, logger);
     }
 
 }
diff --git a/src/com/android/customization/picker/CustomizationPickerActivity.java b/src/com/android/customization/picker/CustomizationPickerActivity.java
index f1e9a47..86a0bf6 100644
--- a/src/com/android/customization/picker/CustomizationPickerActivity.java
+++ b/src/com/android/customization/picker/CustomizationPickerActivity.java
@@ -25,7 +25,6 @@
 import android.view.Gravity;
 import android.view.Menu;
 import android.view.MenuItem;
-import android.view.View;
 
 import androidx.annotation.IdRes;
 import androidx.annotation.NonNull;
@@ -66,7 +65,7 @@
 import com.android.wallpaper.module.Injector;
 import com.android.wallpaper.module.InjectorProvider;
 import com.android.wallpaper.module.UserEventLogger;
-import com.android.wallpaper.picker.BottomActionBarFragment;
+import com.android.wallpaper.module.WallpaperSetter;
 import com.android.wallpaper.picker.CategoryFragment;
 import com.android.wallpaper.picker.CategoryFragment.CategoryFragmentHost;
 import com.android.wallpaper.picker.MyPhotosStarter;
@@ -74,8 +73,6 @@
 import com.android.wallpaper.picker.TopLevelPickerActivity;
 import com.android.wallpaper.picker.WallpaperPickerDelegate;
 import com.android.wallpaper.picker.WallpapersUiContainer;
-import com.android.wallpaper.widget.BottomActionBar;
-import com.android.wallpaper.widget.BottomActionBar.BottomActionBarHost;
 
 import com.google.android.material.bottomnavigation.BottomNavigationView;
 
@@ -87,8 +84,7 @@
  *  Fragments providing customization options.
  */
 public class CustomizationPickerActivity extends FragmentActivity implements WallpapersUiContainer,
-        CategoryFragmentHost, ThemeFragmentHost, GridFragmentHost, ClockFragmentHost,
-        BottomActionBarHost {
+        CategoryFragmentHost, ThemeFragmentHost, GridFragmentHost, ClockFragmentHost {
 
     private static final String TAG = "CustomizationPickerActivity";
     @VisibleForTesting static final String WALLPAPER_FLAVOR_EXTRA =
@@ -102,7 +98,7 @@
 
     private static final Map<Integer, CustomizationSection> mSections = new HashMap<>();
     private CategoryFragment mWallpaperCategoryFragment;
-    private BottomActionBar mBottomActionBar;
+    private WallpaperSetter mWallpaperSetter;
 
     private boolean mWallpaperCategoryInitialized;
 
@@ -139,15 +135,6 @@
                                 ? R.id.nav_wallpaper : R.id.nav_theme);
             }
         }
-
-        mBottomActionBar = findViewById(R.id.bottom_actionbar);
-        mBottomActionBar.getViewTreeObserver().addOnGlobalLayoutListener(() -> {
-            // Only update the visibility of mBottomNav when mBottomActionBar visibility changes.
-            // Since the listener will be triggered by mBottomActionBar and its child views.
-            if (mBottomActionBar.getVisibility() == mBottomNav.getVisibility()) {
-                mBottomNav.setVisibility(mBottomActionBar.isVisible() ? View.GONE : View.VISIBLE);
-            }
-        });
     }
 
     @Override
@@ -207,11 +194,14 @@
         }
         //Theme
         CustomizationInjector injector = (CustomizationInjector) InjectorProvider.getInjector();
+        mWallpaperSetter = new WallpaperSetter(injector.getWallpaperPersister(this),
+                injector.getPreferences(this), mUserEventLogger, false);
         ThemesUserEventLogger eventLogger = (ThemesUserEventLogger) injector.getUserEventLogger(
                 this);
         ThemeManager themeManager = injector.getThemeManager(
                 new DefaultThemeProvider(this, injector.getCustomizationPreferences(this)),
-                this, new OverlayManagerCompat(this), eventLogger);
+                this,
+                mWallpaperSetter, new OverlayManagerCompat(this), eventLogger);
         if (themeManager.isAvailable()) {
             mSections.put(R.id.nav_theme, new ThemeSection(R.id.nav_theme, themeManager));
         } else {
@@ -302,17 +292,9 @@
 
     @Override
     public void onBackPressed() {
-        Fragment fragment = getSupportFragmentManager().findFragmentById(R.id.fragment_container);
-        if (fragment instanceof BottomActionBarFragment
-                && ((BottomActionBarFragment) fragment).onBackPressed()) {
-            return;
-        }
-
-        // For wallpaper tab, since it had child fragment.
         if (mWallpaperCategoryFragment != null && mWallpaperCategoryFragment.popChildFragment()) {
             return;
         }
-
         if (getSupportFragmentManager().popBackStackImmediate()) {
             return;
         }
@@ -409,6 +391,14 @@
     }
 
     @Override
+    protected void onDestroy() {
+        super.onDestroy();
+        if (mWallpaperSetter != null) {
+            mWallpaperSetter.cleanUp();
+        }
+    }
+
+    @Override
     protected void onActivityResult(int requestCode, int resultCode, @Nullable Intent data) {
         super.onActivityResult(requestCode, resultCode, data);
         if (mDelegate.handleActivityResult(requestCode, resultCode, data)) {
@@ -422,11 +412,6 @@
         finish();
     }
 
-    @Override
-    public BottomActionBar getBottomActionBar() {
-        return mBottomActionBar;
-    }
-
     /**
      * Represents a section of the Picker (eg "ThemeBundle", "Clock", etc).
      * There should be a concrete subclass per available section, providing the corresponding
diff --git a/src/com/android/customization/picker/clock/ClockFragment.java b/src/com/android/customization/picker/clock/ClockFragment.java
index 47fad5f..48b4c0d 100644
--- a/src/com/android/customization/picker/clock/ClockFragment.java
+++ b/src/com/android/customization/picker/clock/ClockFragment.java
@@ -43,7 +43,7 @@
 import com.android.wallpaper.R;
 import com.android.wallpaper.asset.Asset;
 import com.android.wallpaper.module.InjectorProvider;
-import com.android.wallpaper.picker.AppbarFragment;
+import com.android.wallpaper.picker.ToolbarFragment;
 import com.android.wallpaper.widget.PreviewPager;
 
 import java.util.List;
@@ -51,7 +51,7 @@
 /**
  * Fragment that contains the main UI for selecting and applying a Clockface.
  */
-public class ClockFragment extends AppbarFragment {
+public class ClockFragment extends ToolbarFragment {
 
     private static final String TAG = "ClockFragment";
 
@@ -64,7 +64,7 @@
 
     public static ClockFragment newInstance(CharSequence title) {
         ClockFragment fragment = new ClockFragment();
-        fragment.setArguments(AppbarFragment.createArguments(title));
+        fragment.setArguments(ToolbarFragment.createArguments(title));
         return fragment;
     }
 
diff --git a/src/com/android/customization/picker/grid/GridFragment.java b/src/com/android/customization/picker/grid/GridFragment.java
index 43eb82f..7a01f26 100644
--- a/src/com/android/customization/picker/grid/GridFragment.java
+++ b/src/com/android/customization/picker/grid/GridFragment.java
@@ -17,16 +17,12 @@
 
 import static android.content.res.Configuration.ORIENTATION_LANDSCAPE;
 
-import static com.android.wallpaper.widget.BottomActionBar.BottomAction.APPLY;
-
 import android.app.Activity;
 import android.content.Context;
 import android.content.res.Resources;
 import android.graphics.drawable.BitmapDrawable;
 import android.net.Uri;
 import android.os.Bundle;
-import android.os.Message;
-import android.os.RemoteException;
 import android.util.DisplayMetrics;
 import android.util.Log;
 import android.view.LayoutInflater;
@@ -36,7 +32,6 @@
 import android.view.View.OnLayoutChangeListener;
 import android.view.ViewGroup;
 import android.widget.ImageView;
-import android.widget.Toast;
 
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
@@ -52,15 +47,14 @@
 import com.android.customization.picker.BasePreviewAdapter;
 import com.android.customization.picker.BasePreviewAdapter.PreviewPage;
 import com.android.customization.widget.OptionSelectorController;
+import com.android.systemui.shared.system.SurfaceViewRequestUtils;
 import com.android.wallpaper.R;
 import com.android.wallpaper.asset.Asset;
 import com.android.wallpaper.asset.ContentUriAsset;
 import com.android.wallpaper.model.WallpaperInfo;
 import com.android.wallpaper.module.CurrentWallpaperInfoFactory;
 import com.android.wallpaper.module.InjectorProvider;
-import com.android.wallpaper.picker.AppbarFragment;
-import com.android.wallpaper.util.SurfaceViewUtils;
-import com.android.wallpaper.widget.BottomActionBar;
+import com.android.wallpaper.picker.ToolbarFragment;
 import com.android.wallpaper.widget.PreviewPager;
 
 import com.bumptech.glide.Glide;
@@ -71,7 +65,7 @@
 /**
  * Fragment that contains the UI for selecting and applying a GridOption.
  */
-public class GridFragment extends AppbarFragment {
+public class GridFragment extends ToolbarFragment {
 
     private static final int PREVIEW_FADE_DURATION_MS = 100;
 
@@ -86,7 +80,7 @@
 
     public static GridFragment newInstance(CharSequence title) {
         GridFragment fragment = new GridFragment();
-        fragment.setArguments(AppbarFragment.createArguments(title));
+        fragment.setArguments(ToolbarFragment.createArguments(title));
         return fragment;
     }
 
@@ -104,43 +98,7 @@
     private ContentLoadingProgressBar mLoading;
     private View mContent;
     private View mError;
-    private BottomActionBar mBottomActionBar;
     private ThemesUserEventLogger mEventLogger;
-    private boolean mReloadOptionsAfterApplying;
-
-    private final Callback mApplyGridCallback = new Callback() {
-        @Override
-        public void onSuccess() {
-            mGridManager.fetchOptions(new OptionsFetchedListener<GridOption>() {
-                @Override
-                public void onOptionsLoaded(List<GridOption> options) {
-                    mOptionsController.resetOptions(options);
-                    mSelectedOption = getSelectedOption(options);
-                    mReloadOptionsAfterApplying = true;
-                    // It will trigger OptionSelectedListener#onOptionSelected.
-                    mOptionsController.setSelectedOption(mSelectedOption);
-                    Toast.makeText(getContext(), R.string.applied_grid_msg, Toast.LENGTH_SHORT)
-                            .show();
-                    // Since we disabled it when clicked apply button.
-                    mBottomActionBar.enableActions();
-                    mBottomActionBar.hide();
-                }
-
-                @Override
-                public void onError(@Nullable Throwable throwable) {
-                    if (throwable != null) {
-                        Log.e(TAG, "Error loading grid options", throwable);
-                    }
-                    showError();
-                }
-            }, true);
-        }
-
-        @Override
-        public void onError(@Nullable Throwable throwable) {
-            //TODO(chihhangchuang): handle
-        }
-    };
 
     @Override
     public void onAttach(Context context) {
@@ -181,7 +139,20 @@
         // Clear memory cache whenever grid fragment view is being loaded.
         Glide.get(getContext()).clearMemory();
         setUpOptions();
+        view.findViewById(R.id.apply_button).setOnClickListener(v -> {
+            mGridManager.apply(mSelectedOption,  new Callback() {
+                @Override
+                public void onSuccess() {
+                    getActivity().finish();
+                }
 
+                @Override
+                public void onError(@Nullable Throwable throwable) {
+                    //TODO(santie): handle
+                }
+            });
+
+        });
         CurrentWallpaperInfoFactory factory = InjectorProvider.getInjector()
                 .getCurrentWallpaperFactory(getContext().getApplicationContext());
 
@@ -204,16 +175,6 @@
         return view;
     }
 
-    @Override
-    protected void onBottomActionBarReady(BottomActionBar bottomActionBar) {
-        mBottomActionBar = bottomActionBar;
-        mBottomActionBar.showActionsOnly(APPLY);
-        mBottomActionBar.setActionClickListener(APPLY, unused -> {
-            mBottomActionBar.disableActions();
-            mGridManager.apply(mSelectedOption, mApplyGridCallback);
-        });
-    }
-
     private void loadWallpaperBackground() {
         if (mHomeWallpaper != null && mCardHeight > 0 && mCardWidth > 0) {
             mHomeWallpaper.getThumbAsset(getContext()).decodeBitmap(mCardWidth,
@@ -244,16 +205,19 @@
 
                 mOptionsController.addListener(selected -> {
                     mSelectedOption = (GridOption) selected;
-                    if (mReloadOptionsAfterApplying) {
-                        mReloadOptionsAfterApplying = false;
-                        return;
-                    }
-                    mBottomActionBar.show();
                     mEventLogger.logGridSelected(mSelectedOption);
                     createAdapter();
                 });
                 mOptionsController.initOptions(mGridManager);
-                mSelectedOption = getSelectedOption(options);
+                for (GridOption option : options) {
+                    if (option.isActive(mGridManager)) {
+                        mSelectedOption = option;
+                    }
+                }
+                // For development only, as there should always be a grid set.
+                if (mSelectedOption == null) {
+                    mSelectedOption = options.get(0);
+                }
                 createAdapter();
             }
 
@@ -267,14 +231,6 @@
         }, false);
     }
 
-    private GridOption getSelectedOption(List<GridOption> options) {
-        return options.stream()
-                .filter(option -> option.isActive(mGridManager))
-                .findAny()
-                // For development only, as there should always be a grid set.
-                .orElse(options.get(0));
-    }
-
     private void hideError() {
         mContent.setVisibility(View.VISIBLE);
         mError.setVisibility(View.GONE);
@@ -326,18 +282,11 @@
             if (usesSurfaceViewForPreview) {
                 mPreviewSurface.setZOrderOnTop(true);
                 mPreviewSurface.getHolder().addCallback(new SurfaceHolder.Callback() {
-
-                    private Message mCallback;
-
                     @Override
                     public void surfaceCreated(SurfaceHolder holder) {
-                        Bundle result = mGridManager.renderPreview(
-                                SurfaceViewUtils.createSurfaceViewRequest(mPreviewSurface), mName);
-                        if (result != null) {
-                            mPreviewSurface.setChildSurfacePackage(
-                                    SurfaceViewUtils.getSurfacePackage(result));
-                            mCallback = SurfaceViewUtils.getCallback(result);
-                        }
+                        Bundle bundle = SurfaceViewRequestUtils.createSurfaceBundle(
+                                mPreviewSurface);
+                        mGridManager.renderPreview(bundle, mName);
                     }
 
                     @Override
@@ -345,17 +294,7 @@
                             int height) {}
 
                     @Override
-                    public void surfaceDestroyed(SurfaceHolder holder) {
-                        if (mCallback != null) {
-                            try {
-                                mCallback.replyTo.send(mCallback);
-                            } catch (RemoteException e) {
-                                e.printStackTrace();
-                            } finally {
-                                mCallback = null;
-                            }
-                        }
-                    }
+                    public void surfaceDestroyed(SurfaceHolder holder) {}
                 });
             } else {
                 mPreviewAsset.loadDrawableWithTransition(mActivity,
diff --git a/src/com/android/customization/picker/theme/CustomThemeActivity.java b/src/com/android/customization/picker/theme/CustomThemeActivity.java
index 6173b3d..fe537ba 100644
--- a/src/com/android/customization/picker/theme/CustomThemeActivity.java
+++ b/src/com/android/customization/picker/theme/CustomThemeActivity.java
@@ -53,6 +53,7 @@
 import com.android.customization.picker.theme.CustomThemeStepFragment.CustomThemeComponentStepHost;
 import com.android.wallpaper.R;
 import com.android.wallpaper.module.InjectorProvider;
+import com.android.wallpaper.module.WallpaperSetter;
 
 import org.json.JSONException;
 
@@ -105,6 +106,8 @@
         mThemeManager = injector.getThemeManager(
                 new DefaultThemeProvider(this, injector.getCustomizationPreferences(this)),
                 this,
+                new WallpaperSetter(injector.getWallpaperPersister(this),
+                        injector.getPreferences(this), mUserEventLogger, false),
                 new OverlayManagerCompat(this),
                 mUserEventLogger);
         mThemeManager.fetchOptions(null, false);
diff --git a/src/com/android/customization/picker/theme/CustomThemeComponentFragment.java b/src/com/android/customization/picker/theme/CustomThemeComponentFragment.java
index b2217aa..f6471f0 100644
--- a/src/com/android/customization/picker/theme/CustomThemeComponentFragment.java
+++ b/src/com/android/customization/picker/theme/CustomThemeComponentFragment.java
@@ -15,20 +15,27 @@
  */
 package com.android.customization.picker.theme;
 
+import android.app.AlertDialog;
+import android.content.Context;
+import android.graphics.drawable.Drawable;
 import android.os.Bundle;
 import android.view.LayoutInflater;
+import android.view.MenuItem;
 import android.view.View;
 import android.view.ViewGroup;
+import android.widget.TextView;
 
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
+import androidx.annotation.StringRes;
 import androidx.recyclerview.widget.RecyclerView;
 
+import com.android.customization.model.theme.custom.CustomThemeManager;
 import com.android.customization.model.theme.custom.ThemeComponentOption;
 import com.android.customization.model.theme.custom.ThemeComponentOptionProvider;
 import com.android.customization.widget.OptionSelectorController;
 import com.android.wallpaper.R;
-import com.android.wallpaper.picker.AppbarFragment;
+import com.android.wallpaper.picker.ToolbarFragment;
 
 public class CustomThemeComponentFragment extends CustomThemeStepFragment {
     private static final String ARG_USE_GRID_LAYOUT = "CustomThemeComponentFragment.use_grid";;
@@ -41,7 +48,7 @@
     public static CustomThemeComponentFragment newInstance(CharSequence toolbarTitle, int position,
             int titleResId, boolean allowGridLayout) {
         CustomThemeComponentFragment fragment = new CustomThemeComponentFragment();
-        Bundle arguments = AppbarFragment.createArguments(toolbarTitle);
+        Bundle arguments = ToolbarFragment.createArguments(toolbarTitle);
         arguments.putInt(ARG_KEY_POSITION, position);
         arguments.putInt(ARG_KEY_TITLE_RES_ID, titleResId);
         arguments.putBoolean(ARG_USE_GRID_LAYOUT, allowGridLayout);
diff --git a/src/com/android/customization/picker/theme/CustomThemeNameFragment.java b/src/com/android/customization/picker/theme/CustomThemeNameFragment.java
index ed59eb1..e21fc80 100644
--- a/src/com/android/customization/picker/theme/CustomThemeNameFragment.java
+++ b/src/com/android/customization/picker/theme/CustomThemeNameFragment.java
@@ -42,14 +42,14 @@
 import com.android.wallpaper.asset.BitmapCachingAsset;
 import com.android.wallpaper.module.CurrentWallpaperInfoFactory;
 import com.android.wallpaper.module.InjectorProvider;
-import com.android.wallpaper.picker.AppbarFragment;
+import com.android.wallpaper.picker.ToolbarFragment;
 
 public class CustomThemeNameFragment extends CustomThemeStepFragment {
 
     public static CustomThemeNameFragment newInstance(CharSequence toolbarTitle, int position,
             int titleResId) {
         CustomThemeNameFragment fragment = new CustomThemeNameFragment();
-        Bundle arguments = AppbarFragment.createArguments(toolbarTitle);
+        Bundle arguments = ToolbarFragment.createArguments(toolbarTitle);
         arguments.putInt(ARG_KEY_POSITION, position);
         arguments.putInt(ARG_KEY_TITLE_RES_ID, titleResId);
         fragment.setArguments(arguments);
@@ -104,6 +104,7 @@
         mNameEditor = view.findViewById(R.id.custom_theme_name);
         mNameEditor.setText(mCustomThemeManager.getOriginalTheme().getTitle());
         bindCover(view.findViewById(R.id.component_preview_content));
+
         return view;
     }
 
diff --git a/src/com/android/customization/picker/theme/CustomThemeStepFragment.java b/src/com/android/customization/picker/theme/CustomThemeStepFragment.java
index b05ebc4..ebf55e3 100644
--- a/src/com/android/customization/picker/theme/CustomThemeStepFragment.java
+++ b/src/com/android/customization/picker/theme/CustomThemeStepFragment.java
@@ -18,9 +18,9 @@
 import com.android.customization.model.theme.custom.ThemeComponentOption;
 import com.android.customization.model.theme.custom.ThemeComponentOptionProvider;
 import com.android.wallpaper.R;
-import com.android.wallpaper.picker.AppbarFragment;
+import com.android.wallpaper.picker.ToolbarFragment;
 
-abstract class CustomThemeStepFragment extends AppbarFragment {
+abstract class CustomThemeStepFragment extends ToolbarFragment {
     protected static final String ARG_KEY_POSITION = "CustomThemeStepFragment.position";
     protected static final String ARG_KEY_TITLE_RES_ID = "CustomThemeStepFragment.title_res";
     protected CustomThemeComponentStepHost mHost;
diff --git a/src/com/android/customization/picker/theme/ThemeFragment.java b/src/com/android/customization/picker/theme/ThemeFragment.java
index b56c402..b5c68e5 100644
--- a/src/com/android/customization/picker/theme/ThemeFragment.java
+++ b/src/com/android/customization/picker/theme/ThemeFragment.java
@@ -35,6 +35,7 @@
 import android.view.View.OnClickListener;
 import android.view.View.OnLayoutChangeListener;
 import android.view.ViewGroup;
+import android.widget.CheckBox;
 import android.widget.CompoundButton;
 import android.widget.ImageView;
 import android.widget.SeekBar;
@@ -62,11 +63,10 @@
 import com.android.wallpaper.R;
 import com.android.wallpaper.asset.Asset;
 import com.android.wallpaper.asset.Asset.CenterCropBitmapTask;
-import com.android.wallpaper.asset.BitmapCachingAsset;
 import com.android.wallpaper.model.WallpaperInfo;
 import com.android.wallpaper.module.CurrentWallpaperInfoFactory;
 import com.android.wallpaper.module.InjectorProvider;
-import com.android.wallpaper.picker.AppbarFragment;
+import com.android.wallpaper.picker.ToolbarFragment;
 import com.android.wallpaper.widget.PreviewPager;
 
 import java.util.List;
@@ -74,7 +74,7 @@
 /**
  * Fragment that contains the main UI for selecting and applying a ThemeBundle.
  */
-public class ThemeFragment extends AppbarFragment {
+public class ThemeFragment extends ToolbarFragment {
 
     private static final String TAG = "ThemeFragment";
     private static final String KEY_SELECTED_THEME = "ThemeFragment.SelectedThemeBundle";
@@ -87,11 +87,12 @@
     }
     public static ThemeFragment newInstance(CharSequence title) {
         ThemeFragment fragment = new ThemeFragment();
-        fragment.setArguments(AppbarFragment.createArguments(title));
+        fragment.setArguments(ToolbarFragment.createArguments(title));
         return fragment;
     }
 
     private RecyclerView mOptionsContainer;
+    private CheckBox mUseMyWallpaperButton;
     private OptionSelectorController<ThemeBundle> mOptionsController;
     private ThemeManager mThemeManager;
     private ThemesUserEventLogger mEventLogger;
@@ -101,9 +102,8 @@
     private ContentLoadingProgressBar mLoading;
     private View mContent;
     private View mError;
-    private boolean mUseMyWallpaper = true;
+    private boolean mUseMyWallpaper;
     private WallpaperInfo mCurrentHomeWallpaper;
-    private Asset mCurrentWallpaperThumbAsset;
     private CurrentWallpaperInfoFactory mCurrentWallpaperFactory;
     private TimeTicker mTicker;
 
@@ -145,6 +145,8 @@
         view.findViewById(R.id.apply_button).setOnClickListener(v -> {
             applyTheme();
         });
+        mUseMyWallpaperButton = view.findViewById(R.id.use_my_wallpaper);
+        mUseMyWallpaperButton.setOnCheckedChangeListener(this::onUseMyWallpaperCheckChanged);
         setUpOptions(savedInstanceState);
 
         return view;
@@ -218,23 +220,32 @@
         super.onActivityResult(requestCode, resultCode, data);
     }
 
+    private void onUseMyWallpaperCheckChanged(CompoundButton checkbox, boolean checked) {
+        mUseMyWallpaper = checked;
+        reloadWallpaper();
+    }
+
     private void reloadWallpaper() {
         mCurrentWallpaperFactory.createCurrentWallpaperInfos(
                 (homeWallpaper, lockWallpaper, presentationMode) -> {
                     mCurrentHomeWallpaper = homeWallpaper;
-                    mCurrentWallpaperThumbAsset = new BitmapCachingAsset(getContext(),
-                            mCurrentHomeWallpaper.getThumbAsset(getContext()));
-                    if (mSelectedTheme != null && mAdapter != null) {
-                        mAdapter.setWallpaperAsset(mCurrentWallpaperThumbAsset);
-                        mAdapter.rebindWallpaperIfAvailable();
+                    if (mSelectedTheme != null) {
+                        if (mUseMyWallpaper || (mSelectedTheme instanceof CustomTheme)) {
+                            mSelectedTheme.setOverrideThemeWallpaper(homeWallpaper);
+                        } else {
+                            mSelectedTheme.setOverrideThemeWallpaper(null);
+                        }
+                        if (mAdapter != null) {
+                            mAdapter.rebindWallpaperIfAvailable();
+                        }
                     }
         }, false);
     }
 
     private void createAdapter(List<ThemeBundle> options) {
         mAdapter = new ThemePreviewAdapter(getActivity(), mSelectedTheme,
-                mCurrentWallpaperThumbAsset,
-                mSelectedTheme instanceof CustomTheme ? this::onEditClicked : null);
+                mSelectedTheme instanceof CustomTheme ? this::onEditClicked : null,
+                new PreloadWallpapersLayoutListener(options));
         mPreviewPager.setAdapter(mAdapter);
     }
 
@@ -244,6 +255,11 @@
         }
     }
 
+    private void updateButtonsVisibility() {
+        mUseMyWallpaperButton.setVisibility(mSelectedTheme instanceof CustomTheme
+                ? View.INVISIBLE : View.VISIBLE);
+    }
+
     private void hideError() {
         mContent.setVisibility(View.VISIBLE);
         mError.setVisibility(View.GONE);
@@ -276,6 +292,7 @@
                         mEventLogger.logThemeSelected(mSelectedTheme,
                                 selected instanceof CustomTheme);
                         createAdapter(options);
+                        updateButtonsVisibility();
                     }
                 });
                 mOptionsController.initOptions(mThemeManager);
@@ -367,12 +384,11 @@
                 R.id.shape_preview_icon_0, R.id.shape_preview_icon_1, R.id.shape_preview_icon_2,
                 R.id.shape_preview_icon_3, R.id.shape_preview_icon_4, R.id.shape_preview_icon_5
         };
-        private Asset mWallpaperAsset;
 
-        ThemePreviewAdapter(Activity activity, ThemeBundle theme, @Nullable Asset wallpaperAsset,
-                @Nullable OnClickListener editClickListener) {
+        ThemePreviewAdapter(Activity activity, ThemeBundle theme,
+                @Nullable OnClickListener editClickListener,
+                @Nullable OnLayoutChangeListener coverCardLayoutListener) {
             super(activity, R.layout.theme_preview_card);
-            mWallpaperAsset = wallpaperAsset;
             final Resources res = activity.getResources();
             final PreviewInfo previewInfo = theme.getPreviewInfo();
 
@@ -381,14 +397,14 @@
                     : null;
 
             WallpaperPreviewLayoutListener wallpaperListener = new WallpaperPreviewLayoutListener(
-                    () -> mWallpaperAsset, previewInfo, coverScrim, true);
+                    theme, previewInfo, coverScrim, true);
 
             addPage(new ThemeCoverPage(activity, theme.getTitle(),
                     previewInfo.resolveAccentColor(res), previewInfo.icons,
                     previewInfo.headlineFontFamily, previewInfo.bottomSheeetCornerRadius,
                     previewInfo.shapeDrawable, previewInfo.shapeAppIcons, editClickListener,
                     mColorButtonIds, mColorTileIds, mColorTileIconIds, mShapeIconIds,
-                    wallpaperListener));
+                    wallpaperListener, coverCardLayoutListener));
             addPage(new ThemePreviewPage(activity, R.string.preview_name_font, R.drawable.ic_font,
                     R.layout.preview_card_font_content,
                     previewInfo.resolveAccentColor(res)) {
@@ -485,6 +501,31 @@
                     }
                 });
             }
+            if (previewInfo.wallpaperAsset != null) {
+                addPage(new ThemePreviewPage(activity, R.string.preview_name_wallpaper,
+                        R.drawable.ic_nav_wallpaper, R.layout.preview_card_wallpaper_content,
+                        previewInfo.resolveAccentColor(res)) {
+
+                    private final WallpaperPreviewLayoutListener mListener =
+                            new WallpaperPreviewLayoutListener(theme, previewInfo, null, false);
+
+                    @Override
+                    protected boolean containsWallpaper() {
+                        return true;
+                    }
+
+                    @Override
+                    protected void bindBody(boolean forceRebind) {
+                        if (card == null) {
+                            return;
+                        }
+                        card.addOnLayoutChangeListener(mListener);
+                        if (forceRebind) {
+                            card.requestLayout();
+                        }
+                    }
+                });
+            }
         }
 
         public void rebindWallpaperIfAvailable() {
@@ -503,23 +544,15 @@
             }
         }
 
-        public void setWallpaperAsset(Asset wallpaperAsset) {
-            mWallpaperAsset = wallpaperAsset;
-        }
-
         private static class WallpaperPreviewLayoutListener implements OnLayoutChangeListener {
-            interface WallpaperPreviewAssetProvider {
-                Asset getAsset();
-            }
-            private final WallpaperPreviewAssetProvider mWallpaperPreviewAssetProvider;
+            private final ThemeBundle mTheme;
             private final PreviewInfo mPreviewInfo;
             private final Drawable mScrim;
             private final boolean mIsTranslucent;
 
-            WallpaperPreviewLayoutListener(
-                    WallpaperPreviewAssetProvider wallpaperPreviewAssetProvider,
-                    PreviewInfo previewInfo, Drawable scrim, boolean translucent) {
-                mWallpaperPreviewAssetProvider = wallpaperPreviewAssetProvider;
+            public WallpaperPreviewLayoutListener(ThemeBundle theme, PreviewInfo previewInfo,
+                    Drawable scrim, boolean translucent) {
+                mTheme = theme;
                 mPreviewInfo = previewInfo;
                 mScrim = scrim;
                 mIsTranslucent = translucent;
@@ -531,7 +564,8 @@
                 int targetWidth = right - left;
                 int targetHeight = bottom - top;
                 if (targetWidth > 0 && targetHeight > 0) {
-                    Asset wallpaperPreviewAsset = mWallpaperPreviewAssetProvider.getAsset();
+                    Asset wallpaperPreviewAsset = mTheme.getWallpaperPreviewAsset(
+                            view.getContext());
                     if (wallpaperPreviewAsset != null) {
                         wallpaperPreviewAsset.decodeBitmap(
                                 targetWidth, targetHeight,
@@ -577,4 +611,41 @@
             }
         }
     }
+
+    /**
+     * Runs only once after the card size is known, and requests decoding wallpaper bitmaps
+     * for all the options, to warm-up the bitmap cache.
+     */
+    private static class PreloadWallpapersLayoutListener implements OnLayoutChangeListener {
+        private static boolean alreadyRunOnce;
+        private final List<ThemeBundle> mOptions;
+
+        public PreloadWallpapersLayoutListener(List<ThemeBundle> options) {
+            mOptions = options;
+        }
+
+        @Override
+        public void onLayoutChange(View view, int left, int top, int right,
+                int bottom, int oldLeft, int oldTop, int oldRight, int oldBottom) {
+            if (alreadyRunOnce) {
+                view.removeOnLayoutChangeListener(this);
+                return;
+            }
+            int targetWidth = right - left;
+            int targetHeight = bottom - top;
+            if (targetWidth > 0 && targetHeight > 0) {
+                for (ThemeBundle theme : mOptions) {
+                    if (theme instanceof CustomTheme && !((CustomTheme) theme).isDefined()) {
+                        continue;
+                    }
+                    Asset wallpaperAsset = theme.getWallpaperPreviewAsset(view.getContext());
+                    if (wallpaperAsset != null) {
+                        wallpaperAsset.decodeBitmap(targetWidth, targetHeight, bitmap -> {});
+                    }
+                }
+                view.removeOnLayoutChangeListener(this);
+                alreadyRunOnce = true;
+            }
+        }
+    }
 }
diff --git a/src/com/android/customization/picker/theme/ThemePreviewPage.java b/src/com/android/customization/picker/theme/ThemePreviewPage.java
index cbcbc78..c3af91c 100644
--- a/src/com/android/customization/picker/theme/ThemePreviewPage.java
+++ b/src/com/android/customization/picker/theme/ThemePreviewPage.java
@@ -179,6 +179,7 @@
             }
 
             // Shape preview icons:
+
             for (int i = 0; i < 3 && i < mShapeAppIcons.size(); i++) {
                 ImageView iconView = card.findViewById(mShapeIconIds[i]);
                 iconView.setBackground(mShapeAppIcons.get(i));
diff --git a/tests/src/com/android/customization/testing/TestCustomizationInjector.java b/tests/src/com/android/customization/testing/TestCustomizationInjector.java
index dbbdb74..10b5934 100644
--- a/tests/src/com/android/customization/testing/TestCustomizationInjector.java
+++ b/tests/src/com/android/customization/testing/TestCustomizationInjector.java
@@ -13,6 +13,7 @@
 import com.android.wallpaper.module.DrawableLayerResolver;
 import com.android.wallpaper.module.PackageStatusNotifier;
 import com.android.wallpaper.module.UserEventLogger;
+import com.android.wallpaper.module.WallpaperSetter;
 import com.android.wallpaper.testing.TestInjector;
 
 /**
@@ -37,10 +38,12 @@
     public ThemeManager getThemeManager(
             ThemeBundleProvider provider,
             FragmentActivity activity,
+            WallpaperSetter wallpaperSetter,
             OverlayManagerCompat overlayManagerCompat,
             ThemesUserEventLogger logger) {
         if (mThemeManager == null) {
-            mThemeManager = new TestThemeManager(provider, activity, overlayManagerCompat, logger);
+            mThemeManager = new TestThemeManager(provider, activity, wallpaperSetter,
+                    overlayManagerCompat, logger);
         }
         return mThemeManager;
     }
diff --git a/tests/src/com/android/customization/testing/TestThemeManager.java b/tests/src/com/android/customization/testing/TestThemeManager.java
index c4d25fb..cc8567f 100644
--- a/tests/src/com/android/customization/testing/TestThemeManager.java
+++ b/tests/src/com/android/customization/testing/TestThemeManager.java
@@ -6,6 +6,7 @@
 import com.android.customization.model.theme.ThemeBundleProvider;
 import com.android.customization.model.theme.ThemeManager;
 import com.android.customization.module.ThemesUserEventLogger;
+import com.android.wallpaper.module.WallpaperSetter;
 
 /**
  * Test implementation of {@link ThemeManager}.
@@ -17,9 +18,10 @@
     public TestThemeManager(
             ThemeBundleProvider provider,
             FragmentActivity activity,
+            WallpaperSetter wallpaperSetter,
             OverlayManagerCompat overlayManagerCompat,
             ThemesUserEventLogger logger) {
-        super(provider, activity, overlayManagerCompat, logger);
+        super(provider, activity, wallpaperSetter, overlayManagerCompat, logger);
     }
 
     @Override