diff --git a/packages/SystemUI/compose/core/src/com/android/compose/animation/Expandable.kt b/packages/SystemUI/compose/core/src/com/android/compose/animation/Expandable.kt
index ac1ef15094151..d4a81f9c765d2 100644
--- a/packages/SystemUI/compose/core/src/com/android/compose/animation/Expandable.kt
+++ b/packages/SystemUI/compose/core/src/com/android/compose/animation/Expandable.kt
@@ -70,10 +70,8 @@ import androidx.compose.ui.platform.ComposeView
 import androidx.compose.ui.platform.LocalContext
 import androidx.compose.ui.unit.Density
 import androidx.compose.ui.unit.dp
-import androidx.lifecycle.findViewTreeLifecycleOwner
-import androidx.lifecycle.findViewTreeViewModelStoreOwner
-import androidx.lifecycle.setViewTreeLifecycleOwner
-import androidx.lifecycle.setViewTreeViewModelStoreOwner
+import androidx.lifecycle.ViewTreeLifecycleOwner
+import androidx.lifecycle.ViewTreeViewModelStoreOwner
 import com.android.systemui.animation.Expandable
 import com.android.systemui.animation.LaunchAnimator
 import kotlin.math.max
@@ -370,10 +368,13 @@ private fun AnimatedContentInOverlay(
                     context,
                     overlay,
                 )
-
-            overlayViewGroup.setViewTreeLifecycleOwner(composeViewRoot.findViewTreeLifecycleOwner())
-            overlayViewGroup.setViewTreeViewModelStoreOwner(
-                composeViewRoot.findViewTreeViewModelStoreOwner()
+            ViewTreeLifecycleOwner.set(
+                overlayViewGroup,
+                ViewTreeLifecycleOwner.get(composeViewRoot),
+            )
+            ViewTreeViewModelStoreOwner.set(
+                overlayViewGroup,
+                ViewTreeViewModelStoreOwner.get(composeViewRoot),
             )
             ViewTreeSavedStateRegistryOwner.set(
                 overlayViewGroup,