diff --git a/js/global_scope_js.js b/js/global_scope_js.js
index 983d60f1..f4d9c673 100644
--- a/js/global_scope_js.js
+++ b/js/global_scope_js.js
@@ -21,3 +21,27 @@ function copyToClipboard(element) {
 function regenerateClick() {
   document.getElementById("Regenerate").click();
 }
+
+function handleMorphdomUpdate(text) {
+  console.log("Morphing!");
+  morphdom(
+    document.getElementById("chat").parentNode,
+    "<div class=\"prose svelte-1ybaih5\">" + text + "</div>",
+    {
+      onBeforeElUpdated: function(fromEl, toEl) {
+        if (fromEl.tagName === "PRE" && fromEl.querySelector("code[data-highlighted]")) {
+          const fromCode = fromEl.querySelector("code");
+          const toCode = toEl.querySelector("code");
+
+          if (fromCode && toCode && fromCode.textContent === toCode.textContent) {
+            // If the <code> content is the same, preserve the entire <pre> element
+            toEl.className = fromEl.className;
+            toEl.innerHTML = fromEl.innerHTML;
+            return false; // Skip updating the <pre> element
+          }
+        }
+        return !fromEl.isEqualNode(toEl); // Update only if nodes differ
+      }
+    }
+  );
+}
diff --git a/modules/ui_chat.py b/modules/ui_chat.py
index 8497f7df..e80fa33b 100644
--- a/modules/ui_chat.py
+++ b/modules/ui_chat.py
@@ -181,32 +181,7 @@ def create_event_handlers():
     shared.reload_inputs = gradio(reload_arr)
 
     # Morph HTML updates instead of updating everything
-    shared.gradio['display'].change(None, gradio('display'), None,
-      js="""
-      (text) => {
-        morphdom(
-          document.getElementById('chat').parentNode,
-          '<div class="prose svelte-1ybaih5">' + text + '</div>',
-          {
-            onBeforeElUpdated: function(fromEl, toEl) {
-              if (fromEl.tagName === 'PRE' && fromEl.querySelector('code[data-highlighted]')) {
-                const fromCode = fromEl.querySelector('code');
-                const toCode = toEl.querySelector('code');
-
-                if (fromCode && toCode && fromCode.textContent === toCode.textContent) {
-                  // If the <code> content is the same, preserve the entire <pre> element
-                  toEl.className = fromEl.className;
-                  toEl.innerHTML = fromEl.innerHTML;
-                  return false; // Skip updating the <pre> element
-                }
-              }
-              return !fromEl.isEqualNode(toEl); // Update only if nodes differ
-            }
-          }
-        );
-      }
-      """
-    );
+    shared.gradio['display'].change(None, gradio('display'), None, js="(text) => handleMorphdomUpdate(text)")
 
     shared.gradio['Generate'].click(
         ui.gather_interface_values, gradio(shared.input_elements), gradio('interface_state')).then(