diff --git a/res/layout-land/fragment_grid_picker.xml b/res/layout-land/fragment_grid_picker.xml
index dab9308..bebcfef 100644
--- a/res/layout-land/fragment_grid_picker.xml
+++ b/res/layout-land/fragment_grid_picker.xml
@@ -42,18 +42,34 @@
             <include layout="@layout/grid_preview_card"/>
         </FrameLayout>
 
-        <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/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 280c187..f813fa9 100644
--- a/res/values/dimens.xml
+++ b/res/values/dimens.xml
@@ -21,13 +21,13 @@
     <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_icon_size">16dp</dimen>
diff --git a/res/values/strings.xml b/res/values/strings.xml
index 55de455..4add810 100755
--- a/res/values/strings.xml
+++ b/res/values/strings.xml
@@ -81,7 +81,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] -->
@@ -106,20 +106,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 cf55147..0f322c1 100644
--- a/src/com/android/customization/model/grid/GridOptionsManager.java
+++ b/src/com/android/customization/model/grid/GridOptionsManager.java
@@ -57,24 +57,22 @@
 
     @Override
     public void fetchOptions(OptionsFetchedListener<GridOption> callback, boolean reload) {
-        new FetchTask(mProvider, callback, reload).execute();
+        new FetchTask(mProvider, callback).execute();
     }
 
     /** 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, List<GridOption>> {
         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
diff --git a/src/com/android/customization/model/grid/LauncherGridOptionsProvider.java b/src/com/android/customization/model/grid/LauncherGridOptionsProvider.java
index cdab1a6..a189245 100644
--- a/src/com/android/customization/model/grid/LauncherGridOptionsProvider.java
+++ b/src/com/android/customization/model/grid/LauncherGridOptionsProvider.java
@@ -27,6 +27,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;
 
@@ -50,9 +51,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;
@@ -93,9 +91,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));
             }
@@ -110,11 +106,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 ddadc7d..5c670db 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,10 +29,15 @@
 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.ApplicationInfo;
 import android.content.pm.PackageManager.NameNotFoundException;
+import android.content.pm.ResolveInfo;
 import android.content.res.Resources.NotFoundException;
+import android.service.wallpaper.WallpaperService;
 import android.graphics.drawable.Drawable;
 import android.text.TextUtils;
 import android.util.Log;
@@ -45,6 +52,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;
@@ -52,7 +60,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;
@@ -76,6 +86,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";
@@ -162,6 +178,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),
@@ -172,6 +190,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
@@ -188,18 +270,16 @@
             builder.setTitle(mContext.getString(R.string.default_theme_title));
         }
 
-        String colorOverlayPackage = getOverlayPackage(COLOR_PREFIX, DEFAULT_THEME_NAME);
-
         try {
+            String colorOverlayPackage = getOverlayPackage(COLOR_PREFIX, DEFAULT_THEME_NAME);
             mOverlayProvider.addColorOverlay(builder, colorOverlayPackage);
         } catch (NameNotFoundException | NotFoundException e) {
             Log.d(TAG, "Didn't find color overlay for default theme, will use system default");
             mOverlayProvider.addSystemDefaultColor(builder);
         }
 
-        String fontOverlayPackage = getOverlayPackage(FONT_PREFIX, DEFAULT_THEME_NAME);
-
         try {
+            String fontOverlayPackage = getOverlayPackage(FONT_PREFIX, DEFAULT_THEME_NAME);
             mOverlayProvider.addFontOverlay(builder, fontOverlayPackage);
         } catch (NameNotFoundException | NotFoundException e) {
             Log.d(TAG, "Didn't find font overlay for default theme, will use system default");
@@ -253,6 +333,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 4ede59e..5208eae 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;
@@ -49,6 +52,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;
@@ -56,10 +60,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;
 import java.util.stream.Collectors;
@@ -79,15 +85,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(removeNullValues(overlayPackages));
     }
 
@@ -160,6 +171,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,
@@ -168,6 +189,15 @@
         return mOverrideWallpaperAsset;
     }
 
+    public WallpaperInfo getWallpaperInfo() {
+        return mWallpaperInfo;
+    }
+
+    @Nullable
+    public String getWallpaperOptions() {
+        return mWallpaperOptions;
+    }
+
     boolean isDefault() {
         return mIsDefault;
     }
@@ -259,6 +289,7 @@
         @ColorInt public final int colorAccentDark;
         public final List<Drawable> icons;
         public final Drawable shapeDrawable;
+        @Nullable public final Asset wallpaperAsset;
         public final List<ShapeAppIcon> shapeAppIcons;
         @Dimension public final int bottomSheeetCornerRadius;
 
@@ -294,6 +325,8 @@
             this.icons = icons;
             this.shapeDrawable = shapeDrawable;
             this.bottomSheeetCornerRadius = cornerRadius;
+            this.wallpaperAsset = wallpaperAsset == null
+                    ? null : new BitmapCachingAsset(context, wallpaperAsset);
             this.shapeAppIcons = shapeAppIcons;
         }
 
@@ -320,11 +353,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<ShapeAppIcon> 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) {
@@ -355,7 +392,8 @@
                 }
             }
             return new PreviewInfo(context, mBodyFontFamily, mHeadlineFontFamily, mColorAccentLight,
-                    mColorAccentDark, mIcons, shapeDrawable, mCornerRadius, shapeIcons);
+                    mColorAccentDark, mIcons, shapeDrawable, mCornerRadius,
+                    mWallpaperAsset, shapeIcons);
         }
 
         public Map<String, String> getPackages() {
@@ -411,6 +449,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 a1ee106..22bcc81 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 @@
 
     private 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 e18cb2d..24c4813 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);
@@ -114,8 +116,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 240f879..87dc2a1 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,6 +65,7 @@
 import com.android.wallpaper.module.Injector;
 import com.android.wallpaper.module.InjectorProvider;
 import com.android.wallpaper.module.UserEventLogger;
+import com.android.wallpaper.module.WallpaperSetter;
 import com.android.wallpaper.module.WallpaperPreferences;
 import com.android.wallpaper.picker.BottomActionBarFragment;
 import com.android.wallpaper.picker.CategoryFragment;
@@ -75,8 +75,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;
 
@@ -88,8 +86,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 =
@@ -103,7 +100,7 @@
 
     private static final Map<Integer, CustomizationSection> mSections = new HashMap<>();
     private CategoryFragment mWallpaperCategoryFragment;
-    private BottomActionBar mBottomActionBar;
+    private WallpaperSetter mWallpaperSetter;
 
     private boolean mWallpaperCategoryInitialized;
 
@@ -211,11 +208,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 {
@@ -306,17 +306,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;
         }
@@ -423,6 +415,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)) {
@@ -436,11 +436,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 bc02ae3..c2ab626 100644
--- a/src/com/android/customization/picker/clock/ClockFragment.java
+++ b/src/com/android/customization/picker/clock/ClockFragment.java
@@ -42,7 +42,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;
@@ -50,7 +50,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";
 
@@ -63,7 +63,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 ad37fbf..e5f5df7 100644
--- a/src/com/android/customization/picker/grid/GridFragment.java
+++ b/src/com/android/customization/picker/grid/GridFragment.java
@@ -50,7 +50,7 @@
 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.BottomActionBar;
 
 import com.bumptech.glide.Glide;
@@ -60,7 +60,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 FULL_PREVIEW_REQUEST_CODE = 1000;
     private static final String KEY_STATE_SELECTED_OPTION = "GridFragment.selectedOption";
@@ -78,7 +78,7 @@
 
     public static GridFragment newInstance(CharSequence title) {
         GridFragment fragment = new GridFragment();
-        fragment.setArguments(AppbarFragment.createArguments(title));
+        fragment.setArguments(ToolbarFragment.createArguments(title));
         return fragment;
     }
 
diff --git a/src/com/android/customization/picker/theme/CustomThemeActivity.java b/src/com/android/customization/picker/theme/CustomThemeActivity.java
index a91faeb..9e184ba 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 a6fdb1c..07bee02 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 f36c80b..afc2713 100644
--- a/src/com/android/customization/picker/theme/CustomThemeNameFragment.java
+++ b/src/com/android/customization/picker/theme/CustomThemeNameFragment.java
@@ -43,7 +43,7 @@
 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 org.json.JSONArray;
 import org.json.JSONException;
@@ -56,7 +56,7 @@
     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);
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 de8aa69..b6fdc12 100644
--- a/src/com/android/customization/picker/theme/ThemeFragment.java
+++ b/src/com/android/customization/picker/theme/ThemeFragment.java
@@ -65,7 +65,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";
@@ -81,7 +81,7 @@
     }
     public static ThemeFragment newInstance(CharSequence title) {
         ThemeFragment fragment = new ThemeFragment();
-        fragment.setArguments(AppbarFragment.createArguments(title));
+        fragment.setArguments(ToolbarFragment.createArguments(title));
         return fragment;
     }
 
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