diff --git a/apps/builder/app/builder/features/style-panel/sections/layout/layout.tsx b/apps/builder/app/builder/features/style-panel/sections/layout/layout.tsx
index def38ae42f..9556d477f2 100644
--- a/apps/builder/app/builder/features/style-panel/sections/layout/layout.tsx
+++ b/apps/builder/app/builder/features/style-panel/sections/layout/layout.tsx
@@ -1,5 +1,7 @@
 import { useState, type ReactNode } from "react";
+import { isFeatureEnabled } from "@webstudio-is/feature-flags";
 import {
+  theme,
   Box,
   EnhancedTooltip,
   Flex,
@@ -38,21 +40,18 @@ import {
   ACSpaceBetweenIcon,
   ACStretchIcon,
 } from "@webstudio-is/icons";
-import type { SectionProps } from "../shared/section";
 import { FlexGrid } from "./shared/flex-grid";
 import { MenuControl, SelectControl } from "../../controls";
 import { styleConfigByName } from "../../shared/configs";
-import { StyleSection } from "../../shared/style-section";
 import { createBatchUpdate, deleteProperty } from "../../shared/use-style-data";
+import { StyleSection } from "../../shared/style-section";
 import {
   type IntermediateStyleValue,
   CssValueInput,
 } from "../../shared/css-value-input";
-import { theme } from "@webstudio-is/design-system";
-import { isFeatureEnabled } from "@webstudio-is/feature-flags";
 import { ToggleControl } from "../../controls/toggle/toggle-control";
 import { PropertyInfo, PropertyLabel } from "../../property-label";
-import { useComputedStyles } from "../../shared/model";
+import { useComputedStyles, useComputedStyleDecl } from "../../shared/model";
 import type { ComputedStyleDecl } from "~/shared/style-object-model";
 
 const GapLinked = ({
@@ -357,27 +356,14 @@ const FlexGap = () => {
   );
 };
 
-const LayoutSectionFlex = ({
-  currentStyle,
-  createBatchUpdate,
-}: {
-  currentStyle: SectionProps["currentStyle"];
-  createBatchUpdate: SectionProps["createBatchUpdate"];
-}) => {
-  const batchUpdate = createBatchUpdate();
-
-  const flexWrapValue = currentStyle.flexWrap?.value;
-
-  // From design: Notice that the align-content icon button is not visible by default.
-  // This property only applies when flex-wrap is set to "wrap".
-  const showAlignContent =
-    flexWrapValue?.type === "keyword" &&
-    (flexWrapValue.value === "wrap" || flexWrapValue.value === "wrap-reverse");
+const LayoutSectionFlex = () => {
+  const flexWrap = useComputedStyleDecl("flexWrap");
+  const flexWrapValue = toValue(flexWrap.cascadedValue);
 
   return (
     <Flex css={{ flexDirection: "column", gap: theme.spacing[5] }}>
       <Flex css={{ gap: theme.spacing[7] }} align="stretch">
-        <FlexGrid currentStyle={currentStyle} batchUpdate={batchUpdate} />
+        <FlexGrid />
         <Flex direction="column" justify="between">
           <Flex css={{ gap: theme.spacing[7] }}>
             <MenuControl
@@ -434,7 +420,7 @@ const LayoutSectionFlex = ({
                 { name: "end", label: "End", icon: JCEndIcon },
               ]}
             />
-            {showAlignContent && (
+            {(flexWrapValue === "wrap" || flexWrapValue === "wrap-reverse") && (
               <MenuControl
                 property="alignContent"
                 items={[
@@ -488,9 +474,9 @@ export const properties = [
   "columnGap",
 ] satisfies Array<StyleProperty>;
 
-export const Section = ({ currentStyle, createBatchUpdate }: SectionProps) => {
-  const value = toValue(currentStyle.display?.value);
-
+export const Section = () => {
+  const display = useComputedStyleDecl("display");
+  const displayValue = toValue(display.cascadedValue);
   return (
     <StyleSection label="Layout" properties={properties}>
       <Flex direction="column" gap="2">
@@ -511,12 +497,8 @@ export const Section = ({ currentStyle, createBatchUpdate }: SectionProps) => {
             items={orderedDisplayValues.map((name) => ({ name, label: name }))}
           />
         </Grid>
-
-        {(value === "flex" || value === "inline-flex") && (
-          <LayoutSectionFlex
-            currentStyle={currentStyle}
-            createBatchUpdate={createBatchUpdate}
-          />
+        {(displayValue === "flex" || displayValue === "inline-flex") && (
+          <LayoutSectionFlex />
         )}
       </Flex>
     </StyleSection>
diff --git a/apps/builder/app/builder/features/style-panel/sections/layout/shared/flex-grid.stories.tsx b/apps/builder/app/builder/features/style-panel/sections/layout/shared/flex-grid.stories.tsx
index 2c9b9d90cd..2bb26bf3c0 100644
--- a/apps/builder/app/builder/features/style-panel/sections/layout/shared/flex-grid.stories.tsx
+++ b/apps/builder/app/builder/features/style-panel/sections/layout/shared/flex-grid.stories.tsx
@@ -1,14 +1,8 @@
 import { Box, Flex } from "@webstudio-is/design-system";
-import { FlexGrid } from "./flex-grid";
+import { FlexGrid as FlexGridComponent } from "./flex-grid";
 
 export default {
-  component: FlexGrid,
-};
-
-const batchUpdate = {
-  setProperty: () => () => {},
-  deleteProperty: () => {},
-  publish: () => {},
+  component: FlexGridComponent,
 };
 
 const alignItems = ["normal", "stretch", "baseline", "center", "start", "end"];
@@ -21,7 +15,7 @@ const justifyContent = [
   "end",
 ];
 
-const Base = ({ flexDirection }: { flexDirection: string }) => {
+export const FlexGrid = () => {
   return (
     <Flex>
       <Box
@@ -62,20 +56,7 @@ const Base = ({ flexDirection }: { flexDirection: string }) => {
             </Box>
             {justifyContent.map((justifyContent) => (
               <Box key={justifyContent} css={{ position: "relative" }}>
-                <FlexGrid
-                  currentStyle={{
-                    flexDirection: {
-                      value: { type: "keyword", value: flexDirection },
-                    },
-                    justifyContent: {
-                      value: { type: "keyword", value: justifyContent },
-                    },
-                    alignItems: {
-                      value: { type: "keyword", value: alignItems },
-                    },
-                  }}
-                  batchUpdate={batchUpdate}
-                />
+                <FlexGridComponent />
               </Box>
             ))}
           </Flex>
@@ -84,19 +65,3 @@ const Base = ({ flexDirection }: { flexDirection: string }) => {
     </Flex>
   );
 };
-
-export const Row = () => {
-  return <Base flexDirection="row" />;
-};
-
-export const RowReverse = () => {
-  return <Base flexDirection="row-reverse" />;
-};
-
-export const Column = () => {
-  return <Base flexDirection="column" />;
-};
-
-export const ColumnReverse = () => {
-  return <Base flexDirection="column-reverse" />;
-};
diff --git a/apps/builder/app/builder/features/style-panel/sections/layout/shared/flex-grid.tsx b/apps/builder/app/builder/features/style-panel/sections/layout/shared/flex-grid.tsx
index d212345f5e..395f95537f 100644
--- a/apps/builder/app/builder/features/style-panel/sections/layout/shared/flex-grid.tsx
+++ b/apps/builder/app/builder/features/style-panel/sections/layout/shared/flex-grid.tsx
@@ -2,9 +2,10 @@ import type { CSSProperties } from "react";
 import { Box, Flex, Grid, IconButton } from "@webstudio-is/design-system";
 import { toValue } from "@webstudio-is/css-engine";
 import { DotIcon } from "@webstudio-is/icons";
-import type { CreateBatchUpdate } from "../../../shared/use-style-data";
-import { getStyleSource, type StyleInfo } from "../../../shared/style-info";
 import { theme } from "@webstudio-is/design-system";
+import { useComputedStyles } from "../../../shared/model";
+import { getPriorityStyleValueSource } from "../../../property-label";
+import { createBatchUpdate } from "../../../shared/use-style-data";
 
 // Sometimes we need to hide a dot that ends up at an end
 // of a line and visually extends it
@@ -31,36 +32,30 @@ const shouldHideDot = ({
   return false;
 };
 
-export const FlexGrid = ({
-  currentStyle,
-  batchUpdate,
-}: {
-  currentStyle: StyleInfo;
-  batchUpdate: ReturnType<CreateBatchUpdate>;
-}) => {
-  const styleSource = getStyleSource(
-    currentStyle.flexDirection,
-    currentStyle.justifyContent,
-    currentStyle.alignItems
-  );
-  const flexDirection = toValue(currentStyle.flexDirection?.value);
-  const justifyContent = toValue(currentStyle.justifyContent?.value);
-  const alignItems = toValue(currentStyle.alignItems?.value);
-  const setAlignItems = batchUpdate.setProperty("alignItems");
-  const setJustifyContent = batchUpdate.setProperty("justifyContent");
+export const FlexGrid = () => {
+  const styles = useComputedStyles([
+    "flexDirection",
+    "justifyContent",
+    "alignItems",
+  ]);
+  const styleValueSourceColor = getPriorityStyleValueSource(styles);
+  const [flexDirection, justifyContent, alignItems] = styles;
+  const flexDirectionValue = toValue(flexDirection.cascadedValue);
+  const justifyContentValue = toValue(justifyContent.cascadedValue);
+  const alignItemsValue = toValue(alignItems.cascadedValue);
   const alignment = ["start", "center", "end"];
   const gridSize = alignment.length;
   const isFlexDirectionColumn =
-    flexDirection === "column" || flexDirection === "column-reverse";
+    flexDirectionValue === "column" || flexDirectionValue === "column-reverse";
 
   let color = theme.colors.foregroundFlexUiMain;
-  if (styleSource === "local") {
+  if (styleValueSourceColor === "local") {
     color = theme.colors.borderLocalFlexUi;
   }
-  if (styleSource === "overwritten") {
+  if (styleValueSourceColor === "overwritten") {
     color = theme.colors.borderOverwrittenFlexUi;
   }
-  if (styleSource === "remote") {
+  if (styleValueSourceColor === "remote") {
     color = theme.colors.borderRemoteFlexUi;
   }
 
@@ -88,8 +83,8 @@ export const FlexGrid = ({
         gridTemplateRows: "repeat(3, 1fr)",
         color,
         outlineStyle: "solid",
-        outlineWidth: styleSource === "default" ? 1 : 2,
-        outlineOffset: styleSource === "default" ? -1 : -2,
+        outlineWidth: styleValueSourceColor === "default" ? 1 : 2,
+        outlineOffset: styleValueSourceColor === "default" ? -1 : -2,
         outlineColor: color,
         "&:focus-within": {
           outlineWidth: 2,
@@ -138,13 +133,24 @@ export const FlexGrid = ({
               onClick={() => {
                 const justifyContent = alignment[x];
                 const alignItems = alignment[y];
-                setAlignItems({ type: "keyword", value: alignItems });
-                setJustifyContent({ type: "keyword", value: justifyContent });
-                batchUpdate.publish();
+                const batch = createBatchUpdate();
+                batch.setProperty("alignItems")({
+                  type: "keyword",
+                  value: alignItems,
+                });
+                batch.setProperty("justifyContent")({
+                  type: "keyword",
+                  value: justifyContent,
+                });
+                batch.publish();
               }}
             >
-              {shouldHideDot({ x, y, justifyContent, alignItems }) ===
-                false && <DotIcon />}
+              {shouldHideDot({
+                x,
+                y,
+                justifyContent: justifyContentValue,
+                alignItems: alignItemsValue,
+              }) === false && <DotIcon />}
             </IconButton>
           </Flex>
         );
@@ -160,13 +166,13 @@ export const FlexGrid = ({
           pointerEvents: "none",
         }}
         style={{
-          flexDirection: flexDirection as CSSProperties["flexDirection"],
-          justifyContent,
-          alignItems,
+          flexDirection: flexDirectionValue as CSSProperties["flexDirection"],
+          justifyContent: justifyContentValue,
+          alignItems: alignItemsValue,
           ...addjustLinesPadding(
-            justifyContent === "space-between"
+            justifyContentValue === "space-between"
               ? 8
-              : justifyContent === "space-around"
+              : justifyContentValue === "space-around"
                 ? 14.5
                 : undefined
           ),