diff --git a/packages/SystemUI/aconfig/systemui.aconfig b/packages/SystemUI/aconfig/systemui.aconfig
index a745ab5..a9dc145 100644
--- a/packages/SystemUI/aconfig/systemui.aconfig
+++ b/packages/SystemUI/aconfig/systemui.aconfig
@@ -143,8 +143,17 @@
 }
 
 flag {
+   name: "theme_overlay_controller_wakefulness_deprecation"
+   namespace: "systemui"
+   description: "Replacing WakefulnessLifecycle by KeyguardTransitionInteractor in "
+        "ThemOverlayController to mitigate flickering when locking the device"
+   bug: "308676488"
+}
+
+flag {
    name: "media_in_scene_container"
    namespace: "systemui"
    description: "Enable media in the scene container framework"
    bug: "296122467"
 }
+
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionInteractor.kt
index 4da48f6..706aba3c 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionInteractor.kt
@@ -302,4 +302,11 @@
     fun isFinishedInState(state: KeyguardState): Flow<Boolean> {
         return finishedKeyguardState.map { it == state }.distinctUntilChanged()
     }
+
+    /**
+     * Whether we've FINISHED a transition to a state that matches the given predicate. Consider
+     * using [isFinishedInStateWhere] whenever possible instead
+     */
+    fun isFinishedInStateWhereValue(stateMatcher: (KeyguardState) -> Boolean) =
+        stateMatcher(finishedKeyguardState.replayCache.last())
 }
diff --git a/packages/SystemUI/src/com/android/systemui/theme/ThemeOverlayController.java b/packages/SystemUI/src/com/android/systemui/theme/ThemeOverlayController.java
index 5a9f5d5..886fa70 100644
--- a/packages/SystemUI/src/com/android/systemui/theme/ThemeOverlayController.java
+++ b/packages/SystemUI/src/com/android/systemui/theme/ThemeOverlayController.java
@@ -19,6 +19,7 @@
 import static android.util.TypedValue.TYPE_INT_COLOR_ARGB8;
 
 import static com.android.systemui.keyguard.WakefulnessLifecycle.WAKEFULNESS_ASLEEP;
+import static com.android.systemui.Flags.themeOverlayControllerWakefulnessDeprecation;
 import static com.android.systemui.theme.ThemeOverlayApplier.COLOR_SOURCE_HOME;
 import static com.android.systemui.theme.ThemeOverlayApplier.COLOR_SOURCE_LOCK;
 import static com.android.systemui.theme.ThemeOverlayApplier.COLOR_SOURCE_PRESET;
@@ -71,12 +72,15 @@
 import com.android.systemui.flags.FeatureFlags;
 import com.android.systemui.flags.Flags;
 import com.android.systemui.keyguard.WakefulnessLifecycle;
+import com.android.systemui.keyguard.domain.interactor.KeyguardTransitionInteractor;
+import com.android.systemui.keyguard.shared.model.KeyguardState;
 import com.android.systemui.monet.ColorScheme;
 import com.android.systemui.monet.Style;
 import com.android.systemui.monet.TonalPalette;
 import com.android.systemui.settings.UserTracker;
 import com.android.systemui.statusbar.policy.DeviceProvisionedController;
 import com.android.systemui.statusbar.policy.DeviceProvisionedController.DeviceProvisionedListener;
+import com.android.systemui.util.kotlin.JavaAdapter;
 import com.android.systemui.util.settings.SecureSettings;
 
 import com.google.ux.material.libmonet.dynamiccolor.MaterialDynamicColors;
@@ -127,7 +131,6 @@
     private final SecureSettings mSecureSettings;
     private final Executor mMainExecutor;
     private final Handler mBgHandler;
-    private final boolean mIsMonochromaticEnabled;
     private final Context mContext;
     private final boolean mIsMonetEnabled;
     private final boolean mIsFidelityEnabled;
@@ -161,6 +164,8 @@
     private final SparseArray<WallpaperColors> mDeferredWallpaperColors = new SparseArray<>();
     private final SparseIntArray mDeferredWallpaperColorsFlags = new SparseIntArray();
     private final WakefulnessLifecycle mWakefulnessLifecycle;
+    private final JavaAdapter mJavaAdapter;
+    private final KeyguardTransitionInteractor mKeyguardTransitionInteractor;
     private final UiModeManager mUiModeManager;
     private DynamicScheme mDynamicSchemeDark;
     private DynamicScheme mDynamicSchemeLight;
@@ -200,8 +205,12 @@
                 return;
             }
             boolean currentUser = userId == mUserTracker.getUserId();
-            if (currentUser && !mAcceptColorEvents
-                    && mWakefulnessLifecycle.getWakefulness() != WAKEFULNESS_ASLEEP) {
+            boolean isAsleep = themeOverlayControllerWakefulnessDeprecation()
+                    ? mKeyguardTransitionInteractor.isFinishedInStateWhereValue(
+                        state -> KeyguardState.Companion.deviceIsAsleepInState(state))
+                    : mWakefulnessLifecycle.getWakefulness() != WAKEFULNESS_ASLEEP;
+
+            if (currentUser && !mAcceptColorEvents && isAsleep) {
                 mDeferredWallpaperColors.put(userId, wallpaperColors);
                 mDeferredWallpaperColorsFlags.put(userId, which);
                 Log.i(TAG, "colors received; processing deferred until screen off: "
@@ -395,9 +404,10 @@
             FeatureFlags featureFlags,
             @Main Resources resources,
             WakefulnessLifecycle wakefulnessLifecycle,
+            JavaAdapter javaAdapter,
+            KeyguardTransitionInteractor keyguardTransitionInteractor,
             UiModeManager uiModeManager) {
         mContext = context;
-        mIsMonochromaticEnabled = featureFlags.isEnabled(Flags.MONOCHROMATIC_THEME);
         mIsMonetEnabled = featureFlags.isEnabled(Flags.MONET);
         mIsFidelityEnabled = featureFlags.isEnabled(Flags.COLOR_FIDELITY);
         mDeviceProvisionedController = deviceProvisionedController;
@@ -412,6 +422,8 @@
         mUserTracker = userTracker;
         mResources = resources;
         mWakefulnessLifecycle = wakefulnessLifecycle;
+        mJavaAdapter = javaAdapter;
+        mKeyguardTransitionInteractor = keyguardTransitionInteractor;
         mUiModeManager = uiModeManager;
         dumpManager.registerDumpable(TAG, this);
     }
@@ -494,21 +506,34 @@
         }
         mWallpaperManager.addOnColorsChangedListener(mOnColorsChangedListener, null,
                 UserHandle.USER_ALL);
-        mWakefulnessLifecycle.addObserver(new WakefulnessLifecycle.Observer() {
-            @Override
-            public void onFinishedGoingToSleep() {
-                final int userId = mUserTracker.getUserId();
-                final WallpaperColors colors = mDeferredWallpaperColors.get(userId);
-                if (colors != null) {
-                    int flags = mDeferredWallpaperColorsFlags.get(userId);
 
-                    mDeferredWallpaperColors.put(userId, null);
-                    mDeferredWallpaperColorsFlags.put(userId, 0);
+        Runnable whenAsleepHandler = () -> {
+            final int userId = mUserTracker.getUserId();
+            final WallpaperColors colors = mDeferredWallpaperColors.get(userId);
+            if (colors != null) {
+                int flags = mDeferredWallpaperColorsFlags.get(userId);
 
-                    handleWallpaperColors(colors, flags, userId);
-                }
+                mDeferredWallpaperColors.put(userId, null);
+                mDeferredWallpaperColorsFlags.put(userId, 0);
+
+                handleWallpaperColors(colors, flags, userId);
             }
-        });
+        };
+
+        if (themeOverlayControllerWakefulnessDeprecation()) {
+            mJavaAdapter.alwaysCollectFlow(
+                    mKeyguardTransitionInteractor.isFinishedInState(KeyguardState.DOZING),
+                    isFinishedInDozing -> {
+                        if (isFinishedInDozing) whenAsleepHandler.run();
+                    });
+        } else {
+            mWakefulnessLifecycle.addObserver(new WakefulnessLifecycle.Observer() {
+                @Override
+                public void onFinishedGoingToSleep() {
+                    whenAsleepHandler.run();
+                }
+            });
+        }
     }
 
     private void reevaluateSystemTheme(boolean forceReload) {
diff --git a/packages/SystemUI/tests/src/com/android/systemui/theme/ThemeOverlayControllerTest.java b/packages/SystemUI/tests/src/com/android/systemui/theme/ThemeOverlayControllerTest.java
index c454b45..1123688 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/theme/ThemeOverlayControllerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/theme/ThemeOverlayControllerTest.java
@@ -61,10 +61,12 @@
 import com.android.systemui.flags.FeatureFlags;
 import com.android.systemui.flags.Flags;
 import com.android.systemui.keyguard.WakefulnessLifecycle;
+import com.android.systemui.keyguard.domain.interactor.KeyguardTransitionInteractor;
 import com.android.systemui.monet.Style;
 import com.android.systemui.settings.UserTracker;
 import com.android.systemui.statusbar.policy.DeviceProvisionedController;
 import com.android.systemui.statusbar.policy.DeviceProvisionedController.DeviceProvisionedListener;
+import com.android.systemui.util.kotlin.JavaAdapter;
 import com.android.systemui.util.settings.SecureSettings;
 
 import com.google.common.util.concurrent.MoreExecutors;
@@ -88,7 +90,10 @@
 
     private static final int USER_SYSTEM = UserHandle.USER_SYSTEM;
     private static final int USER_SECONDARY = 10;
-
+    @Mock
+    private JavaAdapter mJavaAdapter;
+    @Mock
+    private KeyguardTransitionInteractor mKeyguardTransitionInteractor;
     private ThemeOverlayController mThemeOverlayController;
     @Mock
     private Executor mBgExecutor;
@@ -150,11 +155,12 @@
                 .thenReturn(Color.YELLOW);
         when(mResources.getColor(eq(android.R.color.system_neutral2_500), any()))
                 .thenReturn(Color.BLACK);
+
         mThemeOverlayController = new ThemeOverlayController(mContext,
                 mBroadcastDispatcher, mBgHandler, mMainExecutor, mBgExecutor, mThemeOverlayApplier,
                 mSecureSettings, mWallpaperManager, mUserManager, mDeviceProvisionedController,
                 mUserTracker, mDumpManager, mFeatureFlags, mResources, mWakefulnessLifecycle,
-                mUiModeManager) {
+                mJavaAdapter, mKeyguardTransitionInteractor, mUiModeManager) {
             @VisibleForTesting
             protected boolean isNightMode() {
                 return false;
@@ -736,7 +742,7 @@
                 mBroadcastDispatcher, mBgHandler, executor, executor, mThemeOverlayApplier,
                 mSecureSettings, mWallpaperManager, mUserManager, mDeviceProvisionedController,
                 mUserTracker, mDumpManager, mFeatureFlags, mResources, mWakefulnessLifecycle,
-                mUiModeManager) {
+                mJavaAdapter, mKeyguardTransitionInteractor, mUiModeManager) {
             @VisibleForTesting
             protected boolean isNightMode() {
                 return false;
@@ -776,7 +782,7 @@
                 mBroadcastDispatcher, mBgHandler, executor, executor, mThemeOverlayApplier,
                 mSecureSettings, mWallpaperManager, mUserManager, mDeviceProvisionedController,
                 mUserTracker, mDumpManager, mFeatureFlags, mResources, mWakefulnessLifecycle,
-                mUiModeManager) {
+                mJavaAdapter, mKeyguardTransitionInteractor, mUiModeManager) {
             @VisibleForTesting
             protected boolean isNightMode() {
                 return false;