diff --git a/docs/_includes/head-custom.html b/docs/_includes/head-custom.html
index b6fdac7..4b65776 100644
--- a/docs/_includes/head-custom.html
+++ b/docs/_includes/head-custom.html
@@ -4,10 +4,14 @@
   document.addEventListener("DOMContentLoaded", async function() {
     const elements = document.querySelectorAll('.language-mermaid');
     elements.forEach(el => {
-      const codeBlock = el.querySelector('code');
-      if (codeBlock) {
-        el.textContent = codeBlock.textContent;
-        el.classList.add('mermaid');
+      if (el.tagName === 'CODE') {
+        const pre = el.parentElement;
+        if (pre && pre.tagName === 'PRE') {
+          const mermaidDiv = document.createElement('div');
+          mermaidDiv.className = 'mermaid';
+          mermaidDiv.textContent = el.textContent;
+          pre.replaceWith(mermaidDiv);
+        }
       }
     });
     await mermaid.run({