diff --git a/extensions/gallery/script.js b/extensions/gallery/script.js
index b80785b6..4ff23afc 100644
--- a/extensions/gallery/script.js
+++ b/extensions/gallery/script.js
@@ -1,29 +1,33 @@
 let gallery_element = document.getElementById('gallery-extension');
 let chat_mode_element = document.getElementById('chat-mode');
 
-let extensions_block = gallery_element.parentElement;
+let extensions_block = document.getElementById('extensions');
 let extensions_block_size = extensions_block.childNodes.length;
 let gallery_only = (extensions_block_size == 5);
 
-main_parent.addEventListener('click', function(e) {
-    let chat_visible = (chat_tab.offsetHeight > 0 && chat_tab.offsetWidth > 0);
-    let chat_mode_visible = (chat_mode_element.offsetHeight > 0 && chat_mode_element.offsetWidth > 0);
-    let notebook_visible = (notebook_tab.offsetHeight > 0 && notebook_tab.offsetWidth > 0);
-    let default_visible = (default_tab.offsetHeight > 0 && default_tab.offsetWidth > 0);
+document.querySelector('.header_bar').addEventListener('click', function(event) {
+    if (event.target.tagName === 'BUTTON') {
+        const buttonText = event.target.textContent.trim();
 
-    // Only show this extension in the Chat tab
-    if (chat_visible) {
-        if (chat_mode_visible) {
-            gallery_element.style.display = 'block';
-            extensions_block.style.display = '';
+        let chat_visible = (buttonText == 'Chat');
+        let default_visible = (buttonText == 'Default');
+        let notebook_visible = (buttonText == 'Notebook');
+        let chat_mode_visible = (chat_mode_element.offsetHeight > 0 && chat_mode_element.offsetWidth > 0);
+
+        // Only show this extension in the Chat tab
+        if (chat_visible) {
+            if (chat_mode_visible) {
+                gallery_element.style.display = 'block';
+                extensions_block.style.display = '';
+            } else {
+                gallery_element.style.display = 'none';
+                extensions_block.style.display = 'none';
+            }
         } else {
             gallery_element.style.display = 'none';
-            extensions_block.style.display = 'none';
-        }
-    } else {
-        gallery_element.style.display = 'none';
-        if (gallery_only) {
-            extensions_block.style.display = 'none';
+            if (gallery_only) {
+                extensions_block.style.display = 'none';
+            }
         }
     }
 });
diff --git a/js/main.js b/js/main.js
index 3428c62c..6a27c3b4 100644
--- a/js/main.js
+++ b/js/main.js
@@ -1,32 +1,31 @@
-let chat_tab = document.getElementById('chat-tab');
-let notebook_tab = document.getElementById('notebook-tab');
-let default_tab = document.getElementById('default-tab');
-
-let main_parent = chat_tab.parentNode;
+let main_parent = document.getElementById('chat-tab').parentNode;
 let extensions = document.getElementById('extensions');
 
 main_parent.childNodes[0].classList.add("header_bar");
 main_parent.style = "padding: 0; margin: 0";
 main_parent.parentNode.parentNode.style = "padding: 0";
 
-// Add an event listener to the generation tabs
-main_parent.addEventListener('click', function(e) {
-    let chat_visible = (chat_tab.offsetHeight > 0 && chat_tab.offsetWidth > 0);
-    let notebook_visible = (notebook_tab.offsetHeight > 0 && notebook_tab.offsetWidth > 0);
-    let default_visible = (default_tab.offsetHeight > 0 && default_tab.offsetWidth > 0);
+document.querySelector('.header_bar').addEventListener('click', function(event) {
+    if (event.target.tagName === 'BUTTON') {
+        const buttonText = event.target.textContent.trim();
 
-    // Check if one of the generation tabs is visible
-    if (chat_visible || notebook_visible || default_visible) {
-        extensions.style.display = 'flex';
-        if (chat_visible) {
-            extensions.style.maxWidth = "800px";
-            extensions.style.padding = "0px";
+        let chat_visible = (buttonText == 'Chat');
+        let default_visible = (buttonText == 'Default');
+        let notebook_visible = (buttonText == 'Notebook');
+
+        // Check if one of the generation tabs is visible
+        if (chat_visible || notebook_visible || default_visible) {
+            extensions.style.display = 'flex';
+            if (chat_visible) {
+                extensions.style.maxWidth = "800px";
+                extensions.style.padding = "0px";
+            } else {
+                extensions.style.maxWidth = "none";
+                extensions.style.padding = "15px";
+            }
         } else {
-            extensions.style.maxWidth = "none";
-            extensions.style.padding = "15px";
+            extensions.style.display = 'none';
         }
-    } else {
-        extensions.style.display = 'none';
     }
 });