diff --git a/Android.bp b/Android.bp
index 66a1e14..53439ae 100644
--- a/Android.bp
+++ b/Android.bp
@@ -23,6 +23,14 @@
 }
 
 filegroup {
+    name: "ThemePickerLeaf_src_overrides",
+    srcs: [
+        "src_override/**/*.java",
+        "src_override/**/*.kt",
+    ],
+}
+
+filegroup {
     name: "ThemePickerLeaf_Manifest",
     srcs: [
         "AndroidManifest.xml",
@@ -53,28 +61,37 @@
         + "$(location soong_zip) -o $(out) -C $$RES_DIR -D $$RES_DIR"
 }
 
-android_library {
-    name: "ThemePickerLeaf-core",
-    defaults: ["ThemePicker_defaults"],
-}
+java_defaults {
+    name: "ThemePickerLeaf_defaults",
 
-android_app {
-    name: "ThemePickerLeaf",
+    static_libs: [
+        "guava",
+        "monet",
+        "renderscript_toolkit",
+        "wallpaper-common-deps",
+        "SettingsLibSettingsTheme",
+        "SystemUI-statsd",
+        "styleprotoslite",
+        "androidx.lifecycle_lifecycle-livedata-ktx",
+        "androidx.lifecycle_lifecycle-runtime-ktx",
+        "androidx.lifecycle_lifecycle-viewmodel-ktx",
+        "androidx.recyclerview_recyclerview",
+        "SystemUICustomizationLib",
+        "hilt_android",
+    ],
 
-    manifest: ":ThemePickerLeaf_Manifest",
-    additional_manifests: [
-        ":ThemePicker_Manifest",
-        ":WallpaperPicker2_Manifest",
+    jni_libs: [
+        "librenderscript-toolkit",
     ],
 
     srcs: [
         ":WallpaperPicker2_srcs",
         ":ThemePicker_srcs",
         ":ThemePickerLeaf_srcs",
+        ":ThemePickerLeaf_src_overrides",
     ],
-    static_libs: [
-        "ThemePickerLeaf-core",
-    ],
+
+    use_embedded_native_libs: true,
 
     resource_zips: [
         ":WallpaperPicker2_res",
@@ -88,12 +105,22 @@
         enabled: false,
     },
     kotlincflags: ["-Xjvm-default=enable"],
-    use_embedded_native_libs: true,
-
     certificate: "",
+
     privileged: true,
     system_ext_specific: true,
+}
+
+android_app {
+    name: "ThemePickerLeaf",
+    defaults: ["ThemePickerLeaf_defaults"],
+
     platform_apis: true,
+    manifest: ":ThemePickerLeaf_Manifest",
+    additional_manifests: [
+        ":ThemePicker_Manifest",
+        ":WallpaperPicker2_Manifest",
+    ],
 
     required: [
         "privapp_whitelist_com.android.wallpaper.xml",
diff --git a/AndroidManifest.xml b/AndroidManifest.xml
index 116eac6..bbcd1af 100644
--- a/AndroidManifest.xml
+++ b/AndroidManifest.xml
@@ -29,9 +29,4 @@
 
     <uses-permission android:name="com.android.launcher3.permission.READ_SETTINGS" />
     <uses-permission android:name="com.android.launcher3.permission.WRITE_SETTINGS" />
-
-    <application
-        tools:replace="android:name"
-        android:name="org.leafos.customization.picker.LeafCustomizationPickerApplication">
-    </application>
 </manifest>
diff --git a/src/org/leafos/customization/module/LeafCustomizationSections.java b/src/org/leafos/customization/module/LeafCustomizationSections.java
index eaeb955..2a4ed10 100644
--- a/src/org/leafos/customization/module/LeafCustomizationSections.java
+++ b/src/org/leafos/customization/module/LeafCustomizationSections.java
@@ -31,9 +31,6 @@
 import org.leafos.customization.model.iconshape.IconShapeSectionController;
 
 import com.android.customization.model.theme.OverlayManagerCompat;
-import com.android.customization.picker.quickaffordance.domain.interactor.KeyguardQuickAffordancePickerInteractor;
-import com.android.customization.picker.quickaffordance.ui.section.KeyguardQuickAffordanceSectionController;
-import com.android.customization.picker.quickaffordance.ui.viewmodel.KeyguardQuickAffordancePickerViewModel;
 import com.android.wallpaper.model.CustomizationSectionController;
 import com.android.wallpaper.model.CustomizationSectionController.CustomizationSectionNavigationController;
 import com.android.wallpaper.model.PermissionRequester;
@@ -49,22 +46,13 @@
 
 public class LeafCustomizationSections implements CustomizationSections {
     private CustomizationSections mDefaultCustomizationSections;
-    private final KeyguardQuickAffordancePickerInteractor mKeyguardQuickAffordancePickerInteractor;
-    private final KeyguardQuickAffordancePickerViewModel.Factory
-            mKeyguardQuickAffordancePickerViewModelFactory;
 
-    public LeafCustomizationSections(CustomizationSections defaultCustomizationSections,
-            KeyguardQuickAffordancePickerInteractor keyguardQuickAffordancePickerInteractor,
-            KeyguardQuickAffordancePickerViewModel.Factory
-                    keyguardQuickAffordancePickerViewModelFactory) {
+    public LeafCustomizationSections(CustomizationSections defaultCustomizationSections) {
         mDefaultCustomizationSections = defaultCustomizationSections;
-        mKeyguardQuickAffordancePickerInteractor = keyguardQuickAffordancePickerInteractor;
-        mKeyguardQuickAffordancePickerViewModelFactory =
-                keyguardQuickAffordancePickerViewModelFactory;
     }
 
     @Override
-    public List<CustomizationSectionController<?>> getRevampedUISectionControllersForScreen(
+    public List<CustomizationSectionController<?>> getSectionControllersForScreen(
             Screen screen,
             FragmentActivity activity,
             LifecycleOwner lifecycleOwner,
@@ -80,7 +68,7 @@
             WallpaperManager wallpaperManager,
             boolean isTwoPaneAndSmallWidth) {
         List<CustomizationSectionController<?>> sections = mDefaultCustomizationSections
-                .getRevampedUISectionControllersForScreen(screen, activity, lifecycleOwner,
+                .getSectionControllersForScreen(screen, activity, lifecycleOwner,
                         wallpaperColorsViewModel, permissionRequester, wallpaperPreviewNavigator,
                         sectionNavigationController, savedInstanceState, wallpaperInfoFactory,
                         displayUtils, customizationPickerViewModel,
@@ -101,45 +89,4 @@
         return sections;
     }
 
-    @Override
-    public List<CustomizationSectionController<?>> getAllSectionControllers(
-            FragmentActivity activity,
-            LifecycleOwner lifecycleOwner,
-            WallpaperColorsViewModel wallpaperColorsViewModel,
-            PermissionRequester permissionRequester,
-            WallpaperPreviewNavigator wallpaperPreviewNavigator,
-            CustomizationSectionNavigationController sectionNavigationController,
-            @Nullable Bundle savedInstanceState,
-            DisplayUtils displayUtils) {
-        List<CustomizationSectionController<?>> sections = mDefaultCustomizationSections
-                .getAllSectionControllers(
-                        activity, lifecycleOwner, wallpaperColorsViewModel, permissionRequester,
-                        wallpaperPreviewNavigator, sectionNavigationController, savedInstanceState,
-                        displayUtils);
-
-        // Icon pack selection section.
-        sections.add(new IconPackSectionController(
-                IconPackManager.getInstance(activity, new OverlayManagerCompat(activity)), sectionNavigationController));
-
-        // Font selection section.
-        sections.add(new FontSectionController(
-                FontManager.getInstance(activity, new OverlayManagerCompat(activity)), sectionNavigationController));
-
-        // Icon shape selection section.
-        sections.add(new IconShapeSectionController(
-                IconShapeManager.getInstance(activity, new OverlayManagerCompat(activity)), sectionNavigationController));
-
-        // Lock screen quick affordances section.
-        sections.add(
-                new KeyguardQuickAffordanceSectionController(
-                        sectionNavigationController,
-                        mKeyguardQuickAffordancePickerInteractor,
-                        new ViewModelProvider(
-                                activity,
-                                mKeyguardQuickAffordancePickerViewModelFactory)
-                                .get(KeyguardQuickAffordancePickerViewModel.class),
-                        lifecycleOwner));
-
-        return sections;
-    }
 }
diff --git a/src/org/leafos/customization/module/LeafThemePickerInjector.java b/src/org/leafos/customization/module/LeafThemePickerInjector.java
index d526b9e..f8c509c 100644
--- a/src/org/leafos/customization/module/LeafThemePickerInjector.java
+++ b/src/org/leafos/customization/module/LeafThemePickerInjector.java
@@ -18,18 +18,31 @@
 import androidx.activity.ComponentActivity;
 
 import com.android.customization.module.ThemePickerInjector;
+import com.android.wallpaper.dispatchers.BackgroundDispatcher;
+import com.android.wallpaper.dispatchers.MainDispatcher;
 import com.android.wallpaper.module.CustomizationSections;
 
+import javax.inject.Inject;
+import javax.inject.Singleton;
+import kotlinx.coroutines.CoroutineDispatcher;
+import kotlinx.coroutines.CoroutineScope;
+
+@Singleton
 public class LeafThemePickerInjector extends ThemePickerInjector {
     private CustomizationSections mCustomizationSections;
 
+    @Inject
+    public LeafThemePickerInjector(@MainDispatcher CoroutineScope mainScope,
+            @MainDispatcher CoroutineDispatcher mainDispatcher,
+            @BackgroundDispatcher CoroutineDispatcher bgDispatcher) {
+        super(mainScope, mainDispatcher, bgDispatcher);
+    }
+
     @Override
     public CustomizationSections getCustomizationSections(ComponentActivity activity) {
         if (mCustomizationSections == null) {
             mCustomizationSections = new LeafCustomizationSections(
-                    super.getCustomizationSections(activity),
-                    getKeyguardQuickAffordancePickerInteractor(activity),
-                    getKeyguardQuickAffordancePickerViewModelFactory(activity));
+                    super.getCustomizationSections(activity));
         }
         return mCustomizationSections;
     }
diff --git a/src/org/leafos/customization/picker/LeafCustomizationPickerApplication.java b/src_override/com/android/customization/picker/CustomizationPickerApplication.java
similarity index 63%
rename from src/org/leafos/customization/picker/LeafCustomizationPickerApplication.java
rename to src_override/com/android/customization/picker/CustomizationPickerApplication.java
index 8970ddf..461605d 100644
--- a/src/org/leafos/customization/picker/LeafCustomizationPickerApplication.java
+++ b/src_override/com/android/customization/picker/CustomizationPickerApplication.java
@@ -13,18 +13,26 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.leafos.customization.picker;
+package com.android.customization.picker;
 
-import com.android.customization.picker.CustomizationPickerApplication;
+import android.app.Application;
+
+import com.android.customization.module.CustomizationInjector;
 import com.android.wallpaper.module.InjectorProvider;
 
-import org.leafos.customization.module.LeafThemePickerInjector;
+import dagger.hilt.android.HiltAndroidApp;
 
-public class LeafCustomizationPickerApplication extends CustomizationPickerApplication {
+import javax.inject.Inject;
+
+@HiltAndroidApp(Application.class)
+public class CustomizationPickerApplication extends Hilt_CustomizationPickerApplication {
+
+    @Inject CustomizationInjector mInjector;
+
     @Override
     public void onCreate() {
         super.onCreate();
 
-        InjectorProvider.setInjector(new LeafThemePickerInjector());
+        InjectorProvider.setInjector(mInjector);
     }
 }
diff --git a/src_override/com/android/wallpaper/module/AppModule.java b/src_override/com/android/wallpaper/module/AppModule.java
new file mode 100644
index 0000000..bfb0993
--- /dev/null
+++ b/src_override/com/android/wallpaper/module/AppModule.java
@@ -0,0 +1,49 @@
+/*
+ * Copyright (C) 2023 The LeafOS Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.wallpaper.module;
+
+import android.content.Context;
+
+import com.android.customization.module.CustomizationInjector;
+import com.android.customization.module.DefaultCustomizationPreferences;
+
+import org.leafos.customization.module.LeafThemePickerInjector;
+
+import dagger.Binds;
+import dagger.Module;
+import dagger.Provides;
+import dagger.hilt.InstallIn;
+import dagger.hilt.android.qualifiers.ApplicationContext;
+import dagger.hilt.components.SingletonComponent;
+import javax.inject.Singleton;
+
+@Module
+@InstallIn(SingletonComponent.class)
+public abstract class AppModule {
+    @Binds @Singleton public abstract CustomizationInjector bindInjector(LeafThemePickerInjector impl);
+
+    @Module
+    @InstallIn(SingletonComponent.class)
+    public static class Companion {
+        @Provides
+        @Singleton
+        public static WallpaperPreferences provideWallpaperPreferences(
+            @ApplicationContext Context context
+        ) {
+            return new DefaultCustomizationPreferences(context);
+        }
+    }
+}