diff --git a/fasthtml/js.py b/fasthtml/js.py
index d379e435..4b5d54e0 100644
--- a/fasthtml/js.py
+++ b/fasthtml/js.py
@@ -1,3 +1,4 @@
+import re
 from fastcore.utils import *
 from fasthtml.xtend import Script,jsd,Style,Link
 
@@ -12,17 +13,42 @@ def MarkdownJS(sel='.marked'):
     src = "proc_htmx('%s', e => e.innerHTML = marked.parse(e.textContent));" % sel
     return Script(marked_imp+src, type='module')
 
-def KatexMarkdownJS(sel='.marked', katex_tags='$'):
-    right_tags = '\\$' if katex_tags=='$' else '\\]'
-    src = """
-    import katex from "https://cdn.jsdelivr.net/npm/katex/dist/katex.mjs";
-    const renderMath = tex => katex.renderToString(tex, {throwOnError: false, displayMode: false});
+def KatexMarkdownJS(sel='.marked', inline_delim='$', display_delim='$$', math_envs=None):
+    math_envs = math_envs or ['equation', 'align', 'gather', 'multline']
+    env_list = ','.join(f"'{env}'" for env in math_envs)
 
-    proc_htmx('%s', e => {
-    e.innerHTML = marked.parse(e.textContent).replace(/%s{1,2}\\n*(.+?)\\n*%s{1,2}/g, (_, tex) => renderMath(tex));
+    src = r"""
+    import katex from "https://cdn.jsdelivr.net/npm/katex/dist/katex.mjs";
+    const renderMath = (tex, displayMode) => { return katex.renderToString(tex, {
+            throwOnError: false,
+            displayMode: displayMode,
+            output: 'html',
+            trust: true
+        });
+    };
+    const processLatexEnvironments = (content) => {
+        return content.replace(/\\begin{(\w+)}([\s\S]*?)\\end{\1}/g, (match, env, innerContent) => {
+            if ([%(env_list)s].includes(env)) { return `%(display_delim)s${match}%(display_delim)s`; }
+            return match;
+        });
+    };
+    proc_htmx('%(sel)s', e => {
+        let content = processLatexEnvironments(e.textContent);
+        // Handle display math (including environments)
+        content = content.replace(/%(display_delim)s([\s\S]+?)%(display_delim)s/gm, (_, tex) => renderMath(tex.trim(), true));
+        // Handle inline math
+        content = content.replace(/(?<!\w)%(inline_delim)s([^%(inline_delim)s\s](?:[^%(inline_delim)s]*[^%(inline_delim)s\s])?)%(inline_delim)s(?!\w)/g, (_, tex) => renderMath(tex.trim(), false));
+        e.innerHTML = marked.parse(content);
     });
-    """ % (sel, "\\"+katex_tags, right_tags)
-    return (Script(marked_imp+src, type='module'),
+    """
+    format_dict = {
+        'env_list': env_list,
+        'sel': sel,
+        'display_delim': re.escape(display_delim),
+        'inline_delim': re.escape(inline_delim)
+    }
+    formatted_src = src % format_dict
+    return (Script(marked_imp + formatted_src, type='module'),
             Link(rel="stylesheet", href="https://cdn.jsdelivr.net/npm/katex@0.16.11/dist/katex.min.css"))