diff --git a/javascript/00_modified_gallery.js b/javascript/00_modified_gallery.js
index a8359b6..0711ef4 100644
--- a/javascript/00_modified_gallery.js
+++ b/javascript/00_modified_gallery.js
@@ -8,8 +8,10 @@ class DTEModifiedGallery{
 
     setElement(elem){
         this.#elem = elem;
-        this.#items_grid = this.#elem.querySelectorAll('div.grid button.gallery-item')
-        this.#items_selector = this.#elem.querySelectorAll('div.absolute.overflow-x-scroll button.gallery-item')
+        this.#items_grid = this.#elem.querySelectorAll('div.grid-wrap > div.grid-container > button.thumbnail-item')
+        this.#items_selector = this.#elem.querySelectorAll('div.preview > div.thumbnails > button.thumbnail-item')
+        console.log(this.#items_grid)
+        console.log(this.#items_selector)
     }
 
     updateFilter(){
@@ -51,7 +53,7 @@ class DTEModifiedGallery{
     getVisibleSelectedIndex(){
         if (!this.#elem || !this.#items_selector) return -1;
 
-        let button = this.#elem.querySelector('.gallery-item.\\!ring-2')
+        let button = this.#elem.querySelector('.gradio-gallery .thumbnail-item.selected')
         
         for (let i = 0; i < this.#items_selector.length; ++i){
             if (this.#items_selector[i] == button){
@@ -108,7 +110,7 @@ class DTEModifiedGallery{
                 }
             case 'Escape':
                 {
-                    let imgPreview_close = this.#elem.querySelector('div.modify-upload button')
+                    let imgPreview_close = this.#elem.querySelector('div.preview > div > button[class^="svelte"]')
                     if (imgPreview_close != null) {
                         imgPreview_close.click()
                     }
@@ -184,10 +186,10 @@ class DTEModifiedGallery{
     addClickNextHandler(callback_clicked){
         if (!this.#elem) return;
         
-        let fullImg_preview = this.#elem.querySelectorAll('img.w-full.object-contain')
+        let fullImg_preview = this.#elem.querySelectorAll('div.preview > img')
         if (fullImg_preview != null) {
             fullImg_preview.forEach(function (e) {
-                if (e && e.parentElement.tagName == 'DIV') {
+                if (e) {
                     e.addEventListener('click', callback_clicked, false);
                 }
             });
@@ -198,10 +200,10 @@ class DTEModifiedGallery{
     addClickCloseHandler(callback_clicked){
         if (!this.#elem) return;
         
-        let imgPreview_close = this.#elem.querySelectorAll('div.modify-upload button')
+        let imgPreview_close = this.#elem.querySelectorAll('div.preview > div > button[class^="svelte"]')
         if (imgPreview_close != null) {
             imgPreview_close.forEach(function (e) {
-                if (e && e.parentElement.tagName == 'DIV') {
+                if (e) {
                     e.addEventListener('click', callback_clicked, false);
                 }
             });
@@ -212,10 +214,10 @@ class DTEModifiedGallery{
     clickClose(){
         if (!this.#elem) return;
         
-        let imgPreview_close = this.#elem.querySelectorAll('div.modify-upload button')
+        let imgPreview_close = this.#elem.querySelectorAll('div.preview > div > button[class^="svelte"]')
         if (imgPreview_close != null) {
             imgPreview_close.forEach(function (e) {
-                if (e && e.parentElement.tagName == 'DIV') {
+                if (e) {
                     e.click()
                 }
             });
diff --git a/javascript/99_main.js b/javascript/99_main.js
index f8fb415..92ad487 100644
--- a/javascript/99_main.js
+++ b/javascript/99_main.js
@@ -148,7 +148,6 @@ document.addEventListener("DOMContentLoaded", function () {
             }
     
             prompt.parentElement.insertBefore(counter, prompt)
-            counter.classList.add("token-counter-dte")
             prompt.parentElement.style.position = "relative"
             counter.style.width = "auto"
         }
diff --git a/scripts/ui/block_tag_filter.py b/scripts/ui/block_tag_filter.py
index 99e7fc7..1c696d4 100644
--- a/scripts/ui/block_tag_filter.py
+++ b/scripts/ui/block_tag_filter.py
@@ -108,7 +108,9 @@ class TagFilterUI():
         return self.cbg_tags_update()
 
 
-    def cbg_tags_changed(self, cbg_tags: List[str]):
+    def cbg_tags_changed(self,
+                         cbg_tags#: List[str]
+                         ):
         self.selected_tags = dte_instance.cleanup_tagset(set(dte_instance.read_tags(cbg_tags)))
         return self.cbg_tags_update()
 
diff --git a/scripts/ui/block_tag_select.py b/scripts/ui/block_tag_select.py
index d95e1c9..d372aa7 100644
--- a/scripts/ui/block_tag_select.py
+++ b/scripts/ui/block_tag_select.py
@@ -85,7 +85,9 @@ class TagSelectUI():
         return self.cbg_tags_update()
 
 
-    def cbg_tags_changed(self, cbg_tags: List[str]):
+    def cbg_tags_changed(self,
+                         cbg_tags#: List[str]
+                         ):
         self.selected_tags = set(dte_instance.read_tags(cbg_tags))
         return self.cbg_tags_update()
 
diff --git a/scripts/ui/tab_edit_caption_of_selected_image.py b/scripts/ui/tab_edit_caption_of_selected_image.py
index 188a25a..47d493d 100644
--- a/scripts/ui/tab_edit_caption_of_selected_image.py
+++ b/scripts/ui/tab_edit_caption_of_selected_image.py
@@ -27,7 +27,7 @@ class EditCaptionOfSelectedImageUI(UIBase):
             self.btn_hidden_save_caption = gr.Button(elem_id="dataset_tag_editor_btn_hidden_save_caption")
         with gr.Tab(label='Read Caption from Selected Image'):
             self.tb_caption = gr.Textbox(label='Caption of Selected Image', interactive=False, lines=6, elem_id='dte_caption')
-            self.token_counter_caption = gr.HTML(value='<span></span>', elem_id='dte_caption_counter')
+            self.token_counter_caption = gr.HTML(value='<span>0/75</span>', elem_id='dte_caption_counter', elem_classes=["token-counter-dte"])
             with gr.Row():
                 self.btn_copy_caption = gr.Button(value='Copy and Overwrite')
                 self.btn_prepend_caption = gr.Button(value='Prepend')
@@ -53,7 +53,7 @@ class EditCaptionOfSelectedImageUI(UIBase):
             self.cb_ask_save_when_caption_changed = gr.Checkbox(value=cfg_edit_selected.warn_change_not_saved, label='Warn if changes in caption is not saved')
         with gr.Column():
             self.tb_edit_caption = gr.Textbox(label='Edit Caption', interactive=True, lines=6, elem_id= 'dte_edit_caption')
-            self.token_counter_edit_caption = gr.HTML(value='<span></span>', elem_id='dte_edit_caption_counter')
+            self.token_counter_edit_caption = gr.HTML(value='<span>0/75</span>', elem_id='dte_edit_caption_counter', elem_classes=["token-counter-dte"])
         self.btn_apply_changes_selected_image = gr.Button(value='Apply changes to selected image', variant='primary')
         self.btn_apply_changes_all_images = gr.Button(value='Apply changes to ALL displayed images', variant='primary')
 
diff --git a/style.css b/style.css
index 290108c..7a1ed3e 100644
--- a/style.css
+++ b/style.css
@@ -13,4 +13,11 @@
 
 .token-counter-dte span{
     padding: 0.1em 0.75em;
+}
+
+.token-counter-dte span{
+    background: var(--input-background-fill) !important;
+    box-shadow: 0 0 0.0 0.3em rgba(192,192,192,0.15), inset 0 0 0.6em rgba(192,192,192,0.075);
+    border: 2px solid rgba(192,192,192,0.4) !important;
+    border-radius: 0.4em;
 }
\ No newline at end of file