diff --git a/packages/SystemUI/aconfig/systemui.aconfig b/packages/SystemUI/aconfig/systemui.aconfig
index 16bc7f2..968879a 100644
--- a/packages/SystemUI/aconfig/systemui.aconfig
+++ b/packages/SystemUI/aconfig/systemui.aconfig
@@ -525,6 +525,13 @@
 }
 
 flag {
+   name: "smartspace_lockscreen_viewmodel"
+   namespace: "systemui"
+   description: "Indicate Smartspace lockscreen viewmodel"
+   bug: "331451011"
+}
+
+flag {
    name: "pin_input_field_styled_focus_state"
    namespace: "systemui"
    description: "Enables styled focus states on pin input field if keyboard is connected"
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/lockscreen/LockscreenSmartspaceController.kt b/packages/SystemUI/src/com/android/systemui/statusbar/lockscreen/LockscreenSmartspaceController.kt
index c29a64e..5c88da2 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/lockscreen/LockscreenSmartspaceController.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/lockscreen/LockscreenSmartspaceController.kt
@@ -40,6 +40,7 @@
 import com.android.keyguard.KeyguardUpdateMonitor
 import com.android.settingslib.Utils
 import com.android.systemui.Dumpable
+import com.android.systemui.Flags.smartspaceLockscreenViewmodel
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Background
 import com.android.systemui.dagger.qualifiers.Main
@@ -466,7 +467,9 @@
         configurationController.addCallback(configChangeListener)
         statusBarStateController.addCallback(statusBarStateListener)
         bypassController.registerOnBypassStateChangedListener(bypassStateChangedListener)
-        wakefulnessLifecycle.addObserver(wakefulnessLifecycleObserver)
+        if (!smartspaceLockscreenViewmodel()) {
+            wakefulnessLifecycle.addObserver(wakefulnessLifecycleObserver)
+        }
 
         datePlugin?.registerSmartspaceEventNotifier { e -> session?.notifySmartspaceEvent(e) }
         weatherPlugin?.registerSmartspaceEventNotifier { e -> session?.notifySmartspaceEvent(e) }
@@ -509,7 +512,9 @@
         configurationController.removeCallback(configChangeListener)
         statusBarStateController.removeCallback(statusBarStateListener)
         bypassController.unregisterOnBypassStateChangedListener(bypassStateChangedListener)
-        wakefulnessLifecycle.removeObserver(wakefulnessLifecycleObserver)
+        if (!smartspaceLockscreenViewmodel()) {
+            wakefulnessLifecycle.removeObserver(wakefulnessLifecycleObserver)
+        }
         session = null
 
         datePlugin?.registerSmartspaceEventNotifier(null)
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/lockscreen/LockscreenSmartspaceControllerTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/lockscreen/LockscreenSmartspaceControllerTest.kt
index 5abad61..3757fa3 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/lockscreen/LockscreenSmartspaceControllerTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/lockscreen/LockscreenSmartspaceControllerTest.kt
@@ -30,6 +30,7 @@
 import android.os.Bundle
 import android.os.Handler
 import android.os.UserHandle
+import android.platform.test.annotations.DisableFlags
 import android.provider.Settings
 import android.view.View
 import android.widget.FrameLayout
@@ -785,6 +786,7 @@
     }
 
     @Test
+    @DisableFlags(com.android.systemui.Flags.FLAG_SMARTSPACE_LOCKSCREEN_VIEWMODEL)
     fun testWakefulnessLifecycleDispatch_wake_setsSmartspaceScreenOnTrue() {
         // Connect session
         connectSession()
@@ -801,6 +803,7 @@
     }
 
     @Test
+    @DisableFlags(com.android.systemui.Flags.FLAG_SMARTSPACE_LOCKSCREEN_VIEWMODEL)
     fun testWakefulnessLifecycleDispatch_sleep_setsSmartspaceScreenOnFalse() {
         // Connect session
         connectSession()