diff --git a/services/core/java/com/android/server/wm/Transition.java b/services/core/java/com/android/server/wm/Transition.java
index 84df11a..9ff9e34 100644
--- a/services/core/java/com/android/server/wm/Transition.java
+++ b/services/core/java/com/android/server/wm/Transition.java
@@ -1337,7 +1337,16 @@
         // Commit wallpaper visibility after activity, because usually the wallpaper target token is
         // an activity, and wallpaper's visibility depends on activity's visibility.
         for (int i = mParticipants.size() - 1; i >= 0; --i) {
-            final WallpaperWindowToken wt = mParticipants.valueAt(i).asWallpaperToken();
+            final WindowContainer<?> wc = mParticipants.valueAt(i);
+            WallpaperWindowToken wt = wc.asWallpaperToken();
+            if (!Flags.ensureWallpaperInTransitions()) {
+                if (wt == null) {
+                    final WindowState windowState = wc.asWindowState();
+                    if (windowState != null) {
+                        wt = windowState.mToken.asWallpaperToken();
+                    }
+                }
+            }
             if (wt == null) continue;
             final WindowState target = wt.mDisplayContent.mWallpaperController.getWallpaperTarget();
             final boolean isTargetInvisible = target == null || !target.mToken.isVisible();