diff --git a/src/com/android/customization/module/ThemePickerInjector.kt b/src/com/android/customization/module/ThemePickerInjector.kt
index da25950..a9aa182 100644
--- a/src/com/android/customization/module/ThemePickerInjector.kt
+++ b/src/com/android/customization/module/ThemePickerInjector.kt
@@ -57,14 +57,11 @@
 import com.android.customization.picker.grid.ui.viewmodel.GridScreenViewModel
 import com.android.customization.picker.notifications.domain.interactor.NotificationsSnapshotRestorer
 import com.android.customization.picker.notifications.ui.viewmodel.NotificationSectionViewModel
-import com.android.customization.picker.quickaffordance.data.repository.KeyguardQuickAffordancePickerRepository
 import com.android.customization.picker.quickaffordance.domain.interactor.KeyguardQuickAffordancePickerInteractor
 import com.android.customization.picker.quickaffordance.domain.interactor.KeyguardQuickAffordanceSnapshotRestorer
 import com.android.customization.picker.quickaffordance.ui.viewmodel.KeyguardQuickAffordancePickerViewModel
 import com.android.customization.picker.settings.ui.viewmodel.ColorContrastSectionViewModel
 import com.android.systemui.shared.clocks.ClockRegistry
-import com.android.systemui.shared.customization.data.content.CustomizationProviderClient
-import com.android.systemui.shared.customization.data.content.CustomizationProviderClientImpl
 import com.android.systemui.shared.notifications.data.repository.NotificationSettingsRepository
 import com.android.systemui.shared.notifications.domain.interactor.NotificationSettingsInteractor
 import com.android.wallpaper.config.BaseFlags
@@ -97,15 +94,10 @@
 ) : WallpaperPicker2Injector(mainScope, bgDispatcher), CustomizationInjector {
     private var customizationSections: CustomizationSections? = null
     private var wallpaperInteractor: WallpaperInteractor? = null
-    private var keyguardQuickAffordancePickerInteractor: KeyguardQuickAffordancePickerInteractor? =
-        null
     private var keyguardQuickAffordancePickerViewModelFactory:
         KeyguardQuickAffordancePickerViewModel.Factory? =
         null
-    private var customizationProviderClient: CustomizationProviderClient? = null
     private var fragmentFactory: FragmentFactory? = null
-    private var keyguardQuickAffordanceSnapshotRestorer: KeyguardQuickAffordanceSnapshotRestorer? =
-        null
     private var notificationsSnapshotRestorer: NotificationsSnapshotRestorer? = null
     private var clockPickerInteractor: ClockPickerInteractor? = null
     private var clockCarouselViewModelFactory: ClockCarouselViewModel.Factory? = null
@@ -129,6 +121,12 @@
     // Injected objects, sorted by type
     @Inject
     lateinit var colorContrastSectionViewModelFactory: Lazy<ColorContrastSectionViewModel.Factory>
+    @Inject
+    lateinit var keyguardQuickAffordancePickerInteractor:
+        Lazy<KeyguardQuickAffordancePickerInteractor>
+    @Inject
+    lateinit var keyguardQuickAffordanceSnapshotRestorer:
+        Lazy<KeyguardQuickAffordanceSnapshotRestorer>
     @Inject lateinit var themesUserEventLogger: Lazy<ThemesUserEventLogger>
 
     override fun getCustomizationSections(activity: ComponentActivity): CustomizationSections {
@@ -185,7 +183,7 @@
     ): Map<Int, SnapshotRestorer> {
         return super<WallpaperPicker2Injector>.getSnapshotRestorers(context).toMutableMap().apply {
             this[KEY_QUICK_AFFORDANCE_SNAPSHOT_RESTORER] =
-                getKeyguardQuickAffordanceSnapshotRestorer(context)
+                keyguardQuickAffordanceSnapshotRestorer.get()
             // TODO(b/285047815): Enable after adding wallpaper id for default static wallpaper
             if (getFlags().isWallpaperRestorerEnabled()) {
                 this[KEY_WALLPAPER_SNAPSHOT_RESTORER] = getWallpaperSnapshotRestorer(context)
@@ -237,10 +235,7 @@
     override fun getKeyguardQuickAffordancePickerInteractor(
         context: Context
     ): KeyguardQuickAffordancePickerInteractor {
-        return keyguardQuickAffordancePickerInteractor
-            ?: getKeyguardQuickAffordancePickerInteractorImpl(context).also {
-                keyguardQuickAffordancePickerInteractor = it
-            }
+        return keyguardQuickAffordancePickerInteractor.get()
     }
 
     fun getKeyguardQuickAffordancePickerViewModelFactory(
@@ -257,39 +252,6 @@
                 .also { keyguardQuickAffordancePickerViewModelFactory = it }
     }
 
-    private fun getKeyguardQuickAffordancePickerInteractorImpl(
-        context: Context
-    ): KeyguardQuickAffordancePickerInteractor {
-        val client = getKeyguardQuickAffordancePickerProviderClient(context)
-        val appContext = context.applicationContext
-        return KeyguardQuickAffordancePickerInteractor(
-            KeyguardQuickAffordancePickerRepository(client, getApplicationCoroutineScope()),
-            client
-        ) {
-            getKeyguardQuickAffordanceSnapshotRestorer(appContext)
-        }
-    }
-
-    private fun getKeyguardQuickAffordancePickerProviderClient(
-        context: Context
-    ): CustomizationProviderClient {
-        return customizationProviderClient
-            ?: CustomizationProviderClientImpl(context.applicationContext, bgDispatcher).also {
-                customizationProviderClient = it
-            }
-    }
-
-    private fun getKeyguardQuickAffordanceSnapshotRestorer(
-        context: Context
-    ): KeyguardQuickAffordanceSnapshotRestorer {
-        return keyguardQuickAffordanceSnapshotRestorer
-            ?: KeyguardQuickAffordanceSnapshotRestorer(
-                    getKeyguardQuickAffordancePickerInteractor(context),
-                    getKeyguardQuickAffordancePickerProviderClient(context)
-                )
-                .also { keyguardQuickAffordanceSnapshotRestorer = it }
-    }
-
     fun getNotificationSectionViewModelFactory(
         context: Context,
     ): NotificationSectionViewModel.Factory {
diff --git a/src/com/android/customization/picker/quickaffordance/data/repository/KeyguardQuickAffordancePickerRepository.kt b/src/com/android/customization/picker/quickaffordance/data/repository/KeyguardQuickAffordancePickerRepository.kt
index 6bfe348..ff5f828 100644
--- a/src/com/android/customization/picker/quickaffordance/data/repository/KeyguardQuickAffordancePickerRepository.kt
+++ b/src/com/android/customization/picker/quickaffordance/data/repository/KeyguardQuickAffordancePickerRepository.kt
@@ -20,7 +20,10 @@
 import com.android.customization.picker.quickaffordance.shared.model.KeyguardQuickAffordancePickerAffordanceModel as AffordanceModel
 import com.android.customization.picker.quickaffordance.shared.model.KeyguardQuickAffordancePickerSelectionModel as SelectionModel
 import com.android.customization.picker.quickaffordance.shared.model.KeyguardQuickAffordancePickerSlotModel as SlotModel
-import com.android.systemui.shared.customization.data.content.CustomizationProviderClient as Client
+import com.android.systemui.shared.customization.data.content.CustomizationProviderClient
+import com.android.wallpaper.picker.di.modules.MainDispatcher
+import javax.inject.Inject
+import javax.inject.Singleton
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.SharingStarted
@@ -31,10 +34,10 @@
  * Abstracts access to application state related to functionality for selecting, picking, or setting
  * lock screen quick affordances.
  */
-class KeyguardQuickAffordancePickerRepository(
-    private val client: Client,
-    private val scope: CoroutineScope
-) {
+@Singleton
+class KeyguardQuickAffordancePickerRepository
+@Inject
+constructor(client: CustomizationProviderClient, @MainDispatcher mainScope: CoroutineScope) {
     /** List of slots available on the device. */
     val slots: Flow<List<SlotModel>> =
         client.observeSlots().map { slots -> slots.map { slot -> slot.toModel() } }
@@ -44,23 +47,23 @@
         client
             .observeAffordances()
             .map { affordances -> affordances.map { affordance -> affordance.toModel() } }
-            .shareIn(scope, replay = 1, started = SharingStarted.Lazily)
+            .shareIn(mainScope, replay = 1, started = SharingStarted.Lazily)
 
     /** List of slot-affordance pairs, modeling what the user has currently chosen for each slot. */
     val selections: Flow<List<SelectionModel>> =
         client
             .observeSelections()
             .map { selections -> selections.map { selection -> selection.toModel() } }
-            .shareIn(scope, replay = 1, started = SharingStarted.Lazily)
+            .shareIn(mainScope, replay = 1, started = SharingStarted.Lazily)
 
-    private fun Client.Slot.toModel(): SlotModel {
+    private fun CustomizationProviderClient.Slot.toModel(): SlotModel {
         return SlotModel(
             id = id,
             maxSelectedQuickAffordances = capacity,
         )
     }
 
-    private fun Client.Affordance.toModel(): AffordanceModel {
+    private fun CustomizationProviderClient.Affordance.toModel(): AffordanceModel {
         return AffordanceModel(
             id = id,
             name = name,
@@ -73,7 +76,7 @@
         )
     }
 
-    private fun Client.Selection.toModel(): SelectionModel {
+    private fun CustomizationProviderClient.Selection.toModel(): SelectionModel {
         return SelectionModel(
             slotId = slotId,
             affordanceId = affordanceId,
diff --git a/src/com/android/customization/picker/quickaffordance/domain/interactor/KeyguardQuickAffordancePickerInteractor.kt b/src/com/android/customization/picker/quickaffordance/domain/interactor/KeyguardQuickAffordancePickerInteractor.kt
index 3eca624..b17b939 100644
--- a/src/com/android/customization/picker/quickaffordance/domain/interactor/KeyguardQuickAffordancePickerInteractor.kt
+++ b/src/com/android/customization/picker/quickaffordance/domain/interactor/KeyguardQuickAffordancePickerInteractor.kt
@@ -23,18 +23,22 @@
 import com.android.customization.picker.quickaffordance.shared.model.KeyguardQuickAffordancePickerAffordanceModel as AffordanceModel
 import com.android.customization.picker.quickaffordance.shared.model.KeyguardQuickAffordancePickerSelectionModel as SelectionModel
 import com.android.customization.picker.quickaffordance.shared.model.KeyguardQuickAffordancePickerSlotModel as SlotModel
-import com.android.systemui.shared.customization.data.content.CustomizationProviderClient as Client
-import javax.inject.Provider
+import com.android.systemui.shared.customization.data.content.CustomizationProviderClient
+import javax.inject.Inject
+import javax.inject.Singleton
 import kotlinx.coroutines.flow.Flow
 
 /**
  * Single entry-point for all application state and business logic related to quick affordances on
  * the lock screen.
  */
-class KeyguardQuickAffordancePickerInteractor(
-    private val repository: KeyguardQuickAffordancePickerRepository,
-    private val client: Client,
-    private val snapshotRestorer: Provider<KeyguardQuickAffordanceSnapshotRestorer>,
+@Singleton
+class KeyguardQuickAffordancePickerInteractor
+@Inject
+constructor(
+    repository: KeyguardQuickAffordancePickerRepository,
+    private val client: CustomizationProviderClient,
+    private val snapshotRestorer: KeyguardQuickAffordanceSnapshotRestorer,
 ) {
     /** List of slots available on the device. */
     val slots: Flow<List<SlotModel>> = repository.slots
@@ -60,7 +64,7 @@
             affordanceId = affordanceId,
         )
 
-        snapshotRestorer.get().storeSnapshot()
+        snapshotRestorer.storeSnapshot()
     }
 
     /** Unselects all affordances from the slot with the given ID. */
@@ -69,7 +73,7 @@
             slotId = slotId,
         )
 
-        snapshotRestorer.get().storeSnapshot()
+        snapshotRestorer.storeSnapshot()
     }
 
     /** Unselects all affordances from all slots. */
diff --git a/src/com/android/customization/picker/quickaffordance/domain/interactor/KeyguardQuickAffordanceSnapshotRestorer.kt b/src/com/android/customization/picker/quickaffordance/domain/interactor/KeyguardQuickAffordanceSnapshotRestorer.kt
index fee0cb5..f467989 100644
--- a/src/com/android/customization/picker/quickaffordance/domain/interactor/KeyguardQuickAffordanceSnapshotRestorer.kt
+++ b/src/com/android/customization/picker/quickaffordance/domain/interactor/KeyguardQuickAffordanceSnapshotRestorer.kt
@@ -21,10 +21,14 @@
 import com.android.wallpaper.picker.undo.domain.interactor.SnapshotRestorer
 import com.android.wallpaper.picker.undo.domain.interactor.SnapshotStore
 import com.android.wallpaper.picker.undo.shared.model.RestorableSnapshot
+import javax.inject.Inject
+import javax.inject.Singleton
 
 /** Handles state restoration for the quick affordances system. */
-class KeyguardQuickAffordanceSnapshotRestorer(
-    private val interactor: KeyguardQuickAffordancePickerInteractor,
+@Singleton
+class KeyguardQuickAffordanceSnapshotRestorer
+@Inject
+constructor(
     private val client: CustomizationProviderClient,
 ) : SnapshotRestorer {
 
@@ -43,7 +47,7 @@
 
     override suspend fun restoreToSnapshot(snapshot: RestorableSnapshot) {
         // reset all current selections
-        interactor.unselectAll()
+        client.querySlots().forEach { client.deleteAllSelections(it.id) }
 
         val allSelections = checkNotNull(snapshot.args[KEY_SELECTIONS])
         if (allSelections.isEmpty()) return
@@ -55,9 +59,9 @@
             }
 
         selections.forEach { (slotId, affordanceId) ->
-            interactor.select(
-                slotId,
-                affordanceId,
+            client.insertSelection(
+                slotId = slotId,
+                affordanceId = affordanceId,
             )
         }
     }
diff --git a/src_override/com/android/wallpaper/modules/ThemePickerAppModule.kt b/src_override/com/android/wallpaper/modules/ThemePickerAppModule.kt
index ab1541c..8f144d8 100644
--- a/src_override/com/android/wallpaper/modules/ThemePickerAppModule.kt
+++ b/src_override/com/android/wallpaper/modules/ThemePickerAppModule.kt
@@ -23,12 +23,15 @@
 import com.android.customization.module.ThemePickerInjector
 import com.android.customization.module.logging.ThemesUserEventLogger
 import com.android.customization.module.logging.ThemesUserEventLoggerImpl
+import com.android.systemui.shared.customization.data.content.CustomizationProviderClient
+import com.android.systemui.shared.customization.data.content.CustomizationProviderClientImpl
 import com.android.wallpaper.customization.ui.binder.ThemePickerCustomizationOptionsBinder
 import com.android.wallpaper.module.DefaultPartnerProvider
 import com.android.wallpaper.module.PartnerProvider
 import com.android.wallpaper.module.WallpaperPreferences
 import com.android.wallpaper.module.logging.UserEventLogger
 import com.android.wallpaper.picker.customization.ui.binder.CustomizationOptionsBinder
+import com.android.wallpaper.picker.di.modules.BackgroundDispatcher
 import com.android.wallpaper.picker.preview.data.util.DefaultLiveWallpaperDownloader
 import com.android.wallpaper.picker.preview.data.util.LiveWallpaperDownloader
 import com.android.wallpaper.picker.preview.ui.util.DefaultImageEffectDialogUtil
@@ -42,6 +45,7 @@
 import dagger.hilt.android.qualifiers.ApplicationContext
 import dagger.hilt.components.SingletonComponent
 import javax.inject.Singleton
+import kotlinx.coroutines.CoroutineDispatcher
 
 @Module
 @InstallIn(SingletonComponent::class)
@@ -100,5 +104,14 @@
         ): ColorCustomizationManager {
             return ColorCustomizationManager.getInstance(context, OverlayManagerCompat(context))
         }
+
+        @Provides
+        @Singleton
+        fun provideCustomizationProviderClient(
+            @ApplicationContext context: Context,
+            @BackgroundDispatcher bgDispatcher: CoroutineDispatcher,
+        ): CustomizationProviderClient {
+            return CustomizationProviderClientImpl(context, bgDispatcher)
+        }
     }
 }
diff --git a/tests/module/src/com/android/wallpaper/ThemePickerTestModule.kt b/tests/module/src/com/android/wallpaper/ThemePickerTestModule.kt
index a12190f..ca89c01 100644
--- a/tests/module/src/com/android/wallpaper/ThemePickerTestModule.kt
+++ b/tests/module/src/com/android/wallpaper/ThemePickerTestModule.kt
@@ -15,6 +15,7 @@
  */
 package com.android.wallpaper
 
+import android.content.Context
 import androidx.test.core.app.ApplicationProvider
 import com.android.customization.model.color.ColorCustomizationManager
 import com.android.customization.model.theme.OverlayManagerCompat
@@ -24,6 +25,8 @@
 import com.android.customization.module.logging.ThemesUserEventLogger
 import com.android.customization.testing.TestCustomizationInjector
 import com.android.customization.testing.TestDefaultCustomizationPreferences
+import com.android.systemui.shared.customization.data.content.CustomizationProviderClient
+import com.android.systemui.shared.customization.data.content.CustomizationProviderClientImpl
 import com.android.wallpaper.effects.EffectsController
 import com.android.wallpaper.effects.FakeEffectsController
 import com.android.wallpaper.module.Injector
@@ -35,6 +38,7 @@
 import com.android.wallpaper.network.Requester
 import com.android.wallpaper.picker.customization.ui.binder.CustomizationOptionsBinder
 import com.android.wallpaper.picker.customization.ui.binder.DefaultCustomizationOptionsBinder
+import com.android.wallpaper.picker.di.modules.BackgroundDispatcher
 import com.android.wallpaper.picker.di.modules.EffectsModule
 import com.android.wallpaper.picker.preview.data.util.LiveWallpaperDownloader
 import com.android.wallpaper.picker.preview.ui.util.DefaultImageEffectDialogUtil
@@ -47,9 +51,11 @@
 import dagger.Binds
 import dagger.Module
 import dagger.Provides
+import dagger.hilt.android.qualifiers.ApplicationContext
 import dagger.hilt.components.SingletonComponent
 import dagger.hilt.testing.TestInstallIn
 import javax.inject.Singleton
+import kotlinx.coroutines.CoroutineDispatcher
 
 @Module
 @TestInstallIn(
@@ -126,5 +132,14 @@
                 OverlayManagerCompat(ApplicationProvider.getApplicationContext())
             )
         }
+
+        @Provides
+        @Singleton
+        fun provideCustomizationProviderClient(
+            @ApplicationContext context: Context,
+            @BackgroundDispatcher bgDispatcher: CoroutineDispatcher,
+        ): CustomizationProviderClient {
+            return CustomizationProviderClientImpl(context, bgDispatcher)
+        }
     }
 }
diff --git a/tests/robotests/src/com/android/customization/model/picker/quickaffordance/data/repository/KeyguardQuickAffordancePickerRepositoryTest.kt b/tests/robotests/src/com/android/customization/model/picker/quickaffordance/data/repository/KeyguardQuickAffordancePickerRepositoryTest.kt
index 8687b30..55fb2cb 100644
--- a/tests/robotests/src/com/android/customization/model/picker/quickaffordance/data/repository/KeyguardQuickAffordancePickerRepositoryTest.kt
+++ b/tests/robotests/src/com/android/customization/model/picker/quickaffordance/data/repository/KeyguardQuickAffordancePickerRepositoryTest.kt
@@ -53,7 +53,7 @@
         underTest =
             KeyguardQuickAffordancePickerRepository(
                 client = client,
-                scope = testScope.backgroundScope,
+                mainScope = testScope.backgroundScope,
             )
     }
 
diff --git a/tests/robotests/src/com/android/customization/model/picker/quickaffordance/domain/interactor/KeyguardQuickAffordancePickerInteractorTest.kt b/tests/robotests/src/com/android/customization/model/picker/quickaffordance/domain/interactor/KeyguardQuickAffordancePickerInteractorTest.kt
index 4b4790a..2b84ee4 100644
--- a/tests/robotests/src/com/android/customization/model/picker/quickaffordance/domain/interactor/KeyguardQuickAffordancePickerInteractorTest.kt
+++ b/tests/robotests/src/com/android/customization/model/picker/quickaffordance/domain/interactor/KeyguardQuickAffordancePickerInteractorTest.kt
@@ -24,12 +24,10 @@
 import com.android.customization.picker.quickaffordance.shared.model.KeyguardQuickAffordancePickerSelectionModel
 import com.android.systemui.shared.customization.data.content.FakeCustomizationProviderClient
 import com.android.systemui.shared.keyguard.shared.model.KeyguardQuickAffordanceSlots
-import com.android.wallpaper.testing.FakeSnapshotStore
 import com.android.wallpaper.testing.collectLastValue
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.ExperimentalCoroutinesApi
-import kotlinx.coroutines.runBlocking
 import kotlinx.coroutines.test.StandardTestDispatcher
 import kotlinx.coroutines.test.TestScope
 import kotlinx.coroutines.test.resetMain
@@ -62,16 +60,10 @@
                 repository =
                     KeyguardQuickAffordancePickerRepository(
                         client = client,
-                        scope = testScope.backgroundScope,
+                        mainScope = testScope.backgroundScope,
                     ),
                 client = client,
-                snapshotRestorer = {
-                    KeyguardQuickAffordanceSnapshotRestorer(
-                            interactor = underTest,
-                            client = client,
-                        )
-                        .apply { runBlocking { setUpSnapshotRestorer(FakeSnapshotStore()) } }
-                },
+                snapshotRestorer = KeyguardQuickAffordanceSnapshotRestorer(client),
             )
     }
 
diff --git a/tests/robotests/src/com/android/customization/model/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModelTest.kt b/tests/robotests/src/com/android/customization/model/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModelTest.kt
index 53ade86..870d9f5 100644
--- a/tests/robotests/src/com/android/customization/model/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModelTest.kt
+++ b/tests/robotests/src/com/android/customization/model/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModelTest.kt
@@ -38,7 +38,6 @@
 import com.android.wallpaper.picker.customization.data.repository.WallpaperRepository
 import com.android.wallpaper.picker.customization.domain.interactor.WallpaperInteractor
 import com.android.wallpaper.picker.option.ui.viewmodel.OptionItemViewModel
-import com.android.wallpaper.testing.FakeSnapshotStore
 import com.android.wallpaper.testing.FakeWallpaperClient
 import com.android.wallpaper.testing.TestCurrentWallpaperInfoFactory
 import com.android.wallpaper.testing.TestInjector
@@ -48,7 +47,6 @@
 import com.google.common.truth.Truth.assertWithMessage
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.ExperimentalCoroutinesApi
-import kotlinx.coroutines.runBlocking
 import kotlinx.coroutines.test.StandardTestDispatcher
 import kotlinx.coroutines.test.TestScope
 import kotlinx.coroutines.test.resetMain
@@ -89,16 +87,10 @@
                 repository =
                     KeyguardQuickAffordancePickerRepository(
                         client = client,
-                        scope = testScope.backgroundScope,
+                        mainScope = testScope.backgroundScope,
                     ),
                 client = client,
-                snapshotRestorer = {
-                    KeyguardQuickAffordanceSnapshotRestorer(
-                            interactor = quickAffordanceInteractor,
-                            client = client,
-                        )
-                        .apply { runBlocking { setUpSnapshotRestorer(FakeSnapshotStore()) } }
-                },
+                snapshotRestorer = KeyguardQuickAffordanceSnapshotRestorer(client),
             )
         wallpaperInteractor =
             WallpaperInteractor(