diff --git a/lb_ux.module b/lb_ux.module
index 8e1dc19..e56cee1 100644
--- a/lb_ux.module
+++ b/lb_ux.module
@@ -12,10 +12,29 @@ use Drupal\lb_ux\Form\FormAlter;
 use Drupal\lb_ux\Plugin\Block\InlineBlockUX;
 
 /**
- * Implements hook_element_plugin_alter().
+ * Implements hook_module_implements_alter().
  */
-function lb_ux_element_plugin_alter(array &$definitions) {
-  $definitions['layout_builder']['class'] = LayoutBuilderUX::class;
+function lb_ux_module_implements_alter(&$implementations, $hook) {
+  // Move our hook_element_info_alter() implementation to the end of the list.
+  // it must run after layout_builder_modal_element_info_alter().
+  if ($hook == 'element_info_alter') {
+    $group = $implementations['lb_ux'];
+    unset($implementations['lb_ux']);
+    $implementations['lb_ux'] = $group;
+  }
+}
+
+
+/**
+ * Implements hook_element_info_alter().
+ */
+function lb_ux_element_info_alter(array &$info) {
+  if (isset($info['layout_builder'])) {
+    $info['layout_builder']['#pre_render'][] = [
+      LayoutBuilderUX::class,
+      'preRender',
+    ];
+  }
 }
 
 /**
diff --git a/src/Element/LayoutBuilderUX.php b/src/Element/LayoutBuilderUX.php
index ab88cf5..19ad7dd 100644
--- a/src/Element/LayoutBuilderUX.php
+++ b/src/Element/LayoutBuilderUX.php
@@ -3,143 +3,152 @@
 namespace Drupal\lb_ux\Element;
 
 use Drupal\Core\Render\Element;
+use Drupal\Core\Security\TrustedCallbackInterface;
 use Drupal\Core\Url;
-use Drupal\layout_builder\Element\LayoutBuilder;
-use Drupal\layout_builder\SectionStorageInterface;
 
 /**
  * Alters the Layout Builder UI element.
  */
-class LayoutBuilderUX extends LayoutBuilder {
+class LayoutBuilderUX implements TrustedCallbackInterface {
 
   /**
    * {@inheritdoc}
    */
-  protected function buildAdministrativeSection(SectionStorageInterface $section_storage, $delta) {
-    $build = parent::buildAdministrativeSection($section_storage, $delta);
+  public static function preRender(array $element) {
+    if (!array_key_exists('layout_builder', $element)) {
+      return;
+    }
 
-    $section_label = $build['#attributes']['aria-label'];
+    foreach (Element::children($element['layout_builder']) as $child) {
+      if (!array_key_exists('layout-builder__section', $element['layout_builder'][$child])) {
+        continue;
+      }
+      $build = &$element['layout_builder'][$child];
+      $section_label = $build['#attributes']['aria-label'];
 
-    $build['configure']['#title'] = $this->t('<span class="visually-hidden">Configure @section</span>', ['@section' => $section_label]);
-    $build['configure']['#url'] = Url::fromRoute('layout_builder.configure_section_form', $build['configure']['#url']->getRouteParameters());
+      $build['configure']['#title'] = t('<span class="visually-hidden">Configure @section</span>', ['@section' => $section_label]);
+      $build['configure']['#url'] = Url::fromRoute('layout_builder.configure_section_form', $build['configure']['#url']->getRouteParameters());
 
-    $build['remove']['#title'] = $this->t('<span class="visually-hidden">Remove @section</span>', ['@section' => $section_label]);
+      $build['remove']['#title'] = t('<span class="visually-hidden">Remove @section</span>', ['@section' => $section_label]);
 
-    $build['actions'] = [
-      '#type' => 'container',
-      '#weight' => -100,
-      '#attributes' => [
-        'class' => [
-          'layout-builder__actions',
-          'layout-builder__actions__section',
-        ],
-      ],
-      'label' => [
-        '#type' => 'html_tag',
-        '#tag' => 'span',
+      $build['actions'] = [
+        '#type' => 'container',
+        '#weight' => -100,
         '#attributes' => [
-          'class' => ['layout-builder__section-label'],
+          'class' => [
+            'layout-builder__actions',
+            'layout-builder__actions__section',
+          ],
         ],
-        'content' => ['#markup' => $section_label],
-      ],
-      'configure' => $build['configure'],
-      'remove' => $build['remove'],
-    ];
-    unset($build['configure'], $build['remove'], $build['section_label']);
+        'label' => [
+          '#type' => 'html_tag',
+          '#tag' => 'span',
+          '#attributes' => [
+            'class' => ['layout-builder__section-label'],
+          ],
+          'content' => ['#markup' => $section_label],
+        ],
+        'configure' => $build['configure'],
+        'remove' => $build['remove'],
+      ];
+      unset($build['configure'], $build['remove'], $build['section_label']);
 
-    foreach (Element::children($build['layout-builder__section']) as $region) {
-      foreach (Element::children($build['layout-builder__section'][$region]) as $uuid) {
-        if (in_array($uuid, ['layout_builder_add_block', 'region_label'])) {
-          continue 1;
-        }
+      foreach (Element::children($build['layout-builder__section']) as $region) {
+        foreach (Element::children($build['layout-builder__section'][$region]) as $uuid) {
+          if (in_array($uuid, ['layout_builder_add_block', 'region_label'])) {
+            continue 1;
+          }
 
-        $preview_fallback_string = $build['layout-builder__section'][$region][$uuid]['#attributes']['data-layout-content-preview-placeholder-label'];
-        $route_parameters = $build['layout-builder__section'][$region][$uuid]['#contextual_links']['layout_builder_block']['route_parameters'];
+          $preview_fallback_string = $build['layout-builder__section'][$region][$uuid]['#attributes']['data-layout-content-preview-placeholder-label'];
+          $route_parameters = $build['layout-builder__section'][$region][$uuid]['#contextual_links']['layout_builder_block']['route_parameters'];
 
-        // Remove default contextual links.
-        unset($build['layout-builder__section'][$region][$uuid]['#contextual_links']['layout_builder_block']);
+          // Remove default contextual links.
+          unset($build['layout-builder__section'][$region][$uuid]['#contextual_links']['layout_builder_block']);
 
-        // Ensure the 'content' key is present, as set by
-        // \Drupal\layout_builder\EventSubscriber\BlockComponentRenderArray.
-        assert(isset($build['layout-builder__section'][$region][$uuid]['content']));
+          // Ensure the 'content' key is present, as set by
+          // \Drupal\layout_builder\EventSubscriber\BlockComponentRenderArray.
+          assert(isset($build['layout-builder__section'][$region][$uuid]['content']));
 
-        // Replace content with actions and previous content.
-        $build['layout-builder__section'][$region][$uuid]['content'] = [
-          'actions' => [
-            '#type' => 'container',
-            '#attributes' => [
-              'class' => [
-                'layout-builder__actions',
-                'layout-builder__actions__block',
-              ],
-              'tabindex' => 0,
-            ],
-            'label' => [
-              '#type' => 'html_tag',
-              '#tag' => 'span',
-              '#attributes' => [
-                'class' => ['layout-builder__block-label'],
-              ],
-              'content' => ['#markup' => $preview_fallback_string],
-            ],
-            'move' => [
-              '#type' => 'link',
-              '#title' => $this->t('<span class="visually-hidden">Move @block</span>', ['@block' => $preview_fallback_string]),
-              '#url' => Url::fromRoute('layout_builder.move_block_form', $route_parameters),
+          // Replace content with actions and previous content.
+          $build['layout-builder__section'][$region][$uuid]['content'] = [
+            'actions' => [
+              '#type' => 'container',
               '#attributes' => [
                 'class' => [
-                  'use-ajax',
-                  'layout-builder__link',
-                  'layout-builder__link--move',
+                  'layout-builder__actions',
+                  'layout-builder__actions__block',
                 ],
-                'data-dialog-type' => 'dialog',
-                'data-dialog-renderer' => 'off_canvas',
+                'tabindex' => 0,
               ],
-            ],
-            'configure' => [
-              '#type' => 'link',
-              '#title' => $this->t('<span class="visually-hidden">Configure @block</span>', ['@block' => $preview_fallback_string]),
-              '#url' => Url::fromRoute('layout_builder.update_block', $route_parameters),
-              '#attributes' => [
-                'class' => [
-                  'use-ajax',
-                  'layout-builder__link',
-                  'layout-builder__link--configure',
+              'label' => [
+                '#type' => 'html_tag',
+                '#tag' => 'span',
+                '#attributes' => [
+                  'class' => ['layout-builder__block-label'],
                 ],
-                'data-dialog-type' => 'dialog',
-                'data-dialog-renderer' => 'off_canvas',
+                'content' => ['#markup' => $preview_fallback_string],
               ],
-            ],
-            'remove' => [
-              '#type' => 'link',
-              '#title' => $this->t('<span class="visually-hidden">Remove @block</span>', ['@block' => $preview_fallback_string]),
-              '#url' => Url::fromRoute('layout_builder.remove_block', $route_parameters),
-              '#attributes' => [
-                'class' => [
-                  'use-ajax',
-                  'layout-builder__link',
-                  'layout-builder__link--remove',
+              'move' => [
+                '#type' => 'link',
+                '#title' => t('<span class="visually-hidden">Move @block</span>', ['@block' => $preview_fallback_string]),
+                '#url' => Url::fromRoute('layout_builder.move_block_form', $route_parameters),
+                '#attributes' => [
+                  'class' => [
+                    'use-ajax',
+                    'layout-builder__link',
+                    'layout-builder__link--move',
+                  ],
+                  'data-dialog-type' => 'dialog',
+                  'data-dialog-renderer' => 'off_canvas',
+                ],
+              ],
+              'configure' => [
+                '#type' => 'link',
+                '#title' => t('<span class="visually-hidden">Configure @block</span>', ['@block' => $preview_fallback_string]),
+                '#url' => Url::fromRoute('layout_builder.update_block', $route_parameters),
+                '#attributes' => [
+                  'class' => [
+                    'use-ajax',
+                    'layout-builder__link',
+                    'layout-builder__link--configure',
+                  ],
+                  'data-dialog-type' => 'dialog',
+                  'data-dialog-renderer' => 'off_canvas',
+                ],
+              ],
+              'remove' => [
+                '#type' => 'link',
+                '#title' => t('<span class="visually-hidden">Remove @block</span>', ['@block' => $preview_fallback_string]),
+                '#url' => Url::fromRoute('layout_builder.remove_block', $route_parameters),
+                '#attributes' => [
+                  'class' => [
+                    'use-ajax',
+                    'layout-builder__link',
+                    'layout-builder__link--remove',
+                  ],
+                  'data-dialog-type' => 'dialog',
+                  'data-dialog-renderer' => 'off_canvas',
                 ],
-                'data-dialog-type' => 'dialog',
-                'data-dialog-renderer' => 'off_canvas',
               ],
             ],
-          ],
-          'content' => $build['layout-builder__section'][$region][$uuid]['content'],
-        ];
+            'content' => $build['layout-builder__section'][$region][$uuid]['content'],
+          ];
+        }
       }
     }
-    return $build;
+
+    $element['layout_builder']['#attached']['library'][] = 'lb_ux/drupal.lb_ux';
+    $element['layout_builder']['#attached']['library'][] = 'lb_ux/drupal.lb_ux_message';
+    return $element;
   }
 
   /**
    * {@inheritdoc}
    */
-  protected function layout(SectionStorageInterface $section_storage) {
-    $build = parent::layout($section_storage);
-    $build['#attached']['library'][] = 'lb_ux/drupal.lb_ux';
-    $build['#attached']['library'][] = 'lb_ux/drupal.lb_ux_message';
-    return $build;
+  public static function trustedCallbacks() {
+    return [
+      'preRender',
+    ];
   }
 
 }