diff --git a/packages/SystemUI/customization/src/com/android/systemui/shared/customization/data/content/CustomizationProviderContract.kt b/packages/SystemUI/customization/src/com/android/systemui/shared/customization/data/content/CustomizationProviderContract.kt
index 0d880759bd09a..f9e8aafcfa479 100644
--- a/packages/SystemUI/customization/src/com/android/systemui/shared/customization/data/content/CustomizationProviderContract.kt
+++ b/packages/SystemUI/customization/src/com/android/systemui/shared/customization/data/content/CustomizationProviderContract.kt
@@ -184,6 +184,9 @@ object CustomizationProviderContract {
         /** Flag denoting whether the Monochromatic Theme is enabled. */
         const val FLAG_NAME_MONOCHROMATIC_THEME = "is_monochromatic_theme_enabled"
 
+        /** Flag denoting AI Wallpapers are enabled in wallpaper picker. */
+        const val FLAG_NAME_WALLPAPER_PICKER_UI_FOR_AIWP = "wallpaper_picker_ui_for_aiwp"
+
         object Columns {
             /** String. Unique ID for the flag. */
             const val NAME = "name"
diff --git a/packages/SystemUI/src/com/android/systemui/flags/Flags.kt b/packages/SystemUI/src/com/android/systemui/flags/Flags.kt
index 20e079e28eb57..f5e8f45362bd5 100644
--- a/packages/SystemUI/src/com/android/systemui/flags/Flags.kt
+++ b/packages/SystemUI/src/com/android/systemui/flags/Flags.kt
@@ -221,6 +221,15 @@ object Flags {
             teamfood = true,
         )
 
+    /** Enables UI updates for AI wallpapers in the wallpaper picker. */
+    // TODO(b/267722622): Tracking Bug
+    @JvmField
+    val WALLPAPER_PICKER_UI_FOR_AIWP =
+            releasedFlag(
+                    229,
+                    "wallpaper_picker_ui_for_aiwp"
+            )
+
     /** Whether to inflate the bouncer view on a background thread. */
     // TODO(b/272091103): Tracking Bug
     @JvmField
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardQuickAffordanceInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardQuickAffordanceInteractor.kt
index 66f87bade308d..e31771a8cecd3 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardQuickAffordanceInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardQuickAffordanceInteractor.kt
@@ -407,6 +407,10 @@ constructor(
             KeyguardPickerFlag(
                 name = Contract.FlagsTable.FLAG_NAME_MONOCHROMATIC_THEME,
                 value = featureFlags.isEnabled(Flags.MONOCHROMATIC_THEME)
+            ),
+            KeyguardPickerFlag(
+                name = Contract.FlagsTable.FLAG_NAME_WALLPAPER_PICKER_UI_FOR_AIWP,
+                value = featureFlags.isEnabled(Flags.WALLPAPER_PICKER_UI_FOR_AIWP)
             )
         )
     }