diff --git a/src/com/android/customization/module/DefaultCustomizationSections.java b/src/com/android/customization/module/DefaultCustomizationSections.java
index 49062d1..46f0334 100644
--- a/src/com/android/customization/module/DefaultCustomizationSections.java
+++ b/src/com/android/customization/module/DefaultCustomizationSections.java
@@ -4,7 +4,6 @@
 
 import androidx.annotation.Nullable;
 import androidx.fragment.app.FragmentActivity;
-import androidx.lifecycle.AbstractSavedStateViewModelFactory;
 import androidx.lifecycle.LifecycleOwner;
 import androidx.lifecycle.ViewModelProvider;
 
@@ -43,13 +42,13 @@
     private final KeyguardQuickAffordancePickerInteractor mKeyguardQuickAffordancePickerInteractor;
     private final KeyguardQuickAffordancePickerViewModel.Factory
             mKeyguardQuickAffordancePickerViewModelFactory;
-    private final AbstractSavedStateViewModelFactory mNotificationSectionViewModelFactory;
+    private final NotificationSectionViewModel.Factory mNotificationSectionViewModelFactory;
 
     public DefaultCustomizationSections(
             KeyguardQuickAffordancePickerInteractor keyguardQuickAffordancePickerInteractor,
             KeyguardQuickAffordancePickerViewModel.Factory
                     keyguardQuickAffordancePickerViewModelFactory,
-            AbstractSavedStateViewModelFactory notificationSectionViewModelFactory) {
+            NotificationSectionViewModel.Factory notificationSectionViewModelFactory) {
         mKeyguardQuickAffordancePickerInteractor = keyguardQuickAffordancePickerInteractor;
         mKeyguardQuickAffordancePickerViewModelFactory =
                 keyguardQuickAffordancePickerViewModelFactory;
diff --git a/src/com/android/customization/module/ThemePickerInjector.kt b/src/com/android/customization/module/ThemePickerInjector.kt
index 3de1b43..b448283 100644
--- a/src/com/android/customization/module/ThemePickerInjector.kt
+++ b/src/com/android/customization/module/ThemePickerInjector.kt
@@ -81,15 +81,14 @@
     private var clockRegistry: ClockRegistry? = null
     private var pluginManager: PluginManager? = null
     private var notificationsInteractor: NotificationsInteractor? = null
+    private var notificationSectionViewModelFactory: NotificationSectionViewModel.Factory? = null
 
     override fun getCustomizationSections(activity: ComponentActivity): CustomizationSections {
         return customizationSections
             ?: DefaultCustomizationSections(
                     getKeyguardQuickAffordancePickerInteractor(activity),
                     getKeyguardQuickAffordancePickerViewModelFactory(activity),
-                    NotificationSectionViewModel.newFactory(
-                        owner = activity,
-                        defaultArgs = null,
+                    NotificationSectionViewModel.Factory(
                         interactor = getNotificationsInteractor(activity),
                     ),
                 )
@@ -191,6 +190,16 @@
                 .also { keyguardQuickAffordancePickerViewModelFactory = it }
     }
 
+    fun getNotificationSectionViewModelFactory(
+        context: Context,
+    ): NotificationSectionViewModel.Factory {
+        return notificationSectionViewModelFactory
+            ?: NotificationSectionViewModel.Factory(
+                    interactor = getNotificationsInteractor(context),
+                )
+                .also { notificationSectionViewModelFactory = it }
+    }
+
     private fun getKeyguardQuickAffordancePickerInteractorImpl(
         context: Context
     ): KeyguardQuickAffordancePickerInteractor {
diff --git a/src/com/android/customization/picker/notifications/ui/viewmodel/NotificationSectionViewModel.kt b/src/com/android/customization/picker/notifications/ui/viewmodel/NotificationSectionViewModel.kt
index f33b10b..97b0448 100644
--- a/src/com/android/customization/picker/notifications/ui/viewmodel/NotificationSectionViewModel.kt
+++ b/src/com/android/customization/picker/notifications/ui/viewmodel/NotificationSectionViewModel.kt
@@ -17,14 +17,11 @@
 
 package com.android.customization.picker.notifications.ui.viewmodel
 
-import android.os.Bundle
 import androidx.annotation.StringRes
 import androidx.annotation.VisibleForTesting
-import androidx.lifecycle.AbstractSavedStateViewModelFactory
-import androidx.lifecycle.SavedStateHandle
 import androidx.lifecycle.ViewModel
+import androidx.lifecycle.ViewModelProvider
 import androidx.lifecycle.viewModelScope
-import androidx.savedstate.SavedStateRegistryOwner
 import com.android.customization.picker.notifications.domain.interactor.NotificationsInteractor
 import com.android.wallpaper.R
 import kotlinx.coroutines.flow.Flow
@@ -57,25 +54,15 @@
         viewModelScope.launch { interactor.toggleShowNotificationsOnLockScreenEnabled() }
     }
 
-    companion object {
-        @JvmStatic
-        fun newFactory(
-            owner: SavedStateRegistryOwner,
-            defaultArgs: Bundle? = null,
-            interactor: NotificationsInteractor,
-        ): AbstractSavedStateViewModelFactory =
-            object : AbstractSavedStateViewModelFactory(owner, defaultArgs) {
-                @Suppress("UNCHECKED_CAST")
-                override fun <T : ViewModel> create(
-                    key: String,
-                    modelClass: Class<T>,
-                    handle: SavedStateHandle,
-                ): T {
-                    return NotificationSectionViewModel(
-                        interactor = interactor,
-                    )
-                        as T
-                }
-            }
+    class Factory(
+        private val interactor: NotificationsInteractor,
+    ) : ViewModelProvider.Factory {
+        @Suppress("UNCHECKED_CAST")
+        override fun <T : ViewModel> create(modelClass: Class<T>): T {
+            return NotificationSectionViewModel(
+                interactor = interactor,
+            )
+                as T
+        }
     }
 }