diff --git a/frontend/playwright/ui/visual-specs/workspace.spec.js b/frontend/playwright/ui/visual-specs/workspace.spec.js
index 6eba9b5bd5..a9be81bf91 100644
--- a/frontend/playwright/ui/visual-specs/workspace.spec.js
+++ b/frontend/playwright/ui/visual-specs/workspace.spec.js
@@ -110,6 +110,53 @@ test.describe("Design tab", () => {
     await expect(workspace.page).toHaveScreenshot();
   });
 
+  test("Shows the wrap and padding modes of a flex layout", async ({
+    page,
+  }) => {
+    const workspace = new WorkspacePage(page);
+    await workspace.setupEmptyFile();
+    await workspace.mockRPC(/get\-file\?/, "workspace/get-file-7760.json");
+    await workspace.mockRPC(
+      "get-file-fragment?file-id=*&fragment-id=*",
+      "workspace/get-file-fragment-7760.json",
+    );
+    await workspace.mockRPC(
+      "update-file?id=*",
+      "workspace/update-file-create-rect.json",
+    );
+    await workspace.goToWorkspace({
+      fileId: "cd90e028-326a-80b4-8004-7cdec16ffad5",
+      pageId: "cd90e028-326a-80b4-8004-7cdec16ffad6",
+    });
+
+    await workspace.clickLeafLayer("Flex Board");
+
+    const section = workspace.rightSidebar.getByRole("region", {
+      name: "Layout container section",
+    });
+    const wrap = section.getByRole("button", { name: /^(No w|W)rap$/ });
+    const padding = section.getByRole("button", {
+      name: "Show 4 sided padding options",
+    });
+    await expect(wrap).toBeVisible();
+    // The buttons are a small part of the section: compare every pixel.
+    await expect(section).toHaveScreenshot("flex-layout-rest.png", {
+      maxDiffPixelRatio: 0,
+      threshold: 0,
+    });
+
+    await wrap.click();
+    await padding.click();
+    await padding.blur();
+    await page.mouse.move(0, 0);
+    await expect(section.getByRole("textbox", { name: /top padding/i }))
+      .toBeVisible();
+    await expect(section).toHaveScreenshot("flex-layout-pressed.png", {
+      maxDiffPixelRatio: 0,
+      threshold: 0,
+    });
+  });
+
   test("Shows the align options when selecting two shapes", async ({
     page,
   }) => {
diff --git a/frontend/src/app/main/ui/workspace/sidebar/options/menus/layout_container.cljs b/frontend/src/app/main/ui/workspace/sidebar/options/menus/layout_container.cljs
index 4fb0ebde64..5e043118d2 100644
--- a/frontend/src/app/main/ui/workspace/sidebar/options/menus/layout_container.cljs
+++ b/frontend/src/app/main/ui/workspace/sidebar/options/menus/layout_container.cljs
@@ -192,13 +192,12 @@
 (mf/defc wrap-row
   {::mf/props :obj}
   [{:keys [wrap-type on-click]}]
-  [:button {:class (stl/css-case :wrap-button true
-                                 :selected (= wrap-type :wrap))
-            :title (if (= :wrap wrap-type)
-                     "No wrap"
-                     "Wrap")
-            :on-click on-click}
-   deprecated-icon/wrap])
+  [:> icon-button* {:variant "ghost"
+                    :tooltip-placement "top-left"
+                    :aria-pressed (= wrap-type :wrap)
+                    :aria-label "Wrap"
+                    :icon i/wrap
+                    :on-click on-click}])
 
 (mf/defc align-row
   {::mf/props :obj}
@@ -662,14 +661,13 @@
         (= type :multiple)
         [:> multiple-padding-selection* props])]
 
-     [:button {:class (stl/css-case
-                       :padding-toggle true
-                       :selected (= type :multiple))
-               :title (tr "workspace.layout-grid.editor.padding.expand")
-               :aria-label (tr "workspace.layout-grid.editor.padding.expand")
-               :data-type (d/name type)
-               :on-click on-type-change'}
-      deprecated-icon/padding-extended]]))
+     [:> icon-button* {:variant "ghost"
+                       :tooltip-placement "top-left"
+                       :aria-pressed (= type :multiple)
+                       :aria-label (tr "workspace.layout-grid.editor.padding.expand")
+                       :data-type (d/name type)
+                       :icon i/padding-extended
+                       :on-click on-type-change'}]]))
 
 ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
 ;; GAP
diff --git a/frontend/src/app/main/ui/workspace/sidebar/options/menus/layout_container.scss b/frontend/src/app/main/ui/workspace/sidebar/options/menus/layout_container.scss
index 7182dc5729..123efabf29 100644
--- a/frontend/src/app/main/ui/workspace/sidebar/options/menus/layout_container.scss
+++ b/frontend/src/app/main/ui/workspace/sidebar/options/menus/layout_container.scss
@@ -61,25 +61,6 @@
   grid-column: span 4;
 }
 
-// TODO: Replace this buttons with DS buttons
-.wrap-button {
-  @extend %button-tertiary;
-
-  border-radius: $br-8;
-  block-size: $sz-32;
-  inline-size: $sz-32;
-
-  svg {
-    @extend %button-icon;
-
-    stroke: var(--color-foreground-secondary);
-  }
-
-  &.selected {
-    @extend %button-icon-selected;
-  }
-}
-
 .second-row,
 .third-row {
   grid-column: 1 / -1;
@@ -155,25 +136,6 @@
   gap: var(--sp-xs);
 }
 
-// TODO: Replace this buttons with DS buttons
-.padding-toggle {
-  @extend %button-tertiary;
-
-  block-size: $sz-32;
-  inline-size: $sz-32;
-  border-radius: $br-8;
-
-  svg {
-    @extend %button-icon;
-
-    stroke: var(--color-foreground-secondary);
-  }
-
-  &.selected {
-    @extend %button-icon-selected;
-  }
-}
-
 .grid-layout-menu {
   @include sidebar.option-grid-structure;
 }