move token counter to the location of the prompt, add token counting for the negative prompt

4 files changed, 87 insertions, 59 deletions

diff --git a/extensions-builtin/prompt-bracket-checker/javascript/prompt-bracket-checker.js b/extensions-builtin/prompt-bracket-checker/javascript/prompt-bracket-checker.js
index 251a1f57..4a85c8eb 100644
--- a/extensions-builtin/prompt-bracket-checker/javascript/prompt-bracket-checker.js
+++ b/extensions-builtin/prompt-bracket-checker/javascript/prompt-bracket-checker.js
@@ -4,16 +4,10 @@
 // Counts open and closed brackets (round, square, curly) in the prompt and negative prompt text boxes in the txt2img and img2img tabs.
 // If there's a mismatch, the keyword counter turns red and if you hover on it, a tooltip tells you what's wrong.
 
-function checkBrackets(evt) {
-  textArea = evt.target;
-  tabName = evt.target.parentElement.parentElement.id.split("_")[0];
-  counterElt = document.querySelector('gradio-app').shadowRoot.querySelector('#' + tabName + '_token_counter');
-
-  promptName = evt.target.parentElement.parentElement.id.includes('neg') ? ' negative' : '';
-
-  errorStringParen  = '(' + tabName + promptName + ' prompt) - Different number of opening and closing parentheses detected.\n';
-  errorStringSquare = '[' + tabName + promptName + ' prompt] - Different number of opening and closing square brackets detected.\n';
-  errorStringCurly  = '{' + tabName + promptName + ' prompt} - Different number of opening and closing curly brackets detected.\n';
+function checkBrackets(evt, textArea, counterElt) {
+  errorStringParen  = '(...) - Different number of opening and closing parentheses detected.\n';
+  errorStringSquare = '[...] - Different number of opening and closing square brackets detected.\n';
+  errorStringCurly  = '{...} - Different number of opening and closing curly brackets detected.\n';
 
   openBracketRegExp = /\(/g;
   closeBracketRegExp = /\)/g;
@@ -86,24 +80,31 @@ function checkBrackets(evt) {
   }
 
   if(counterElt.title != '') {
-    counterElt.style = 'color: #FF5555;';
+    counterElt.classList.add('error');
   } else {
-    counterElt.style = '';
+    counterElt.classList.remove('error');
   }
 }
 
-var shadowRootLoaded = setInterval(function() {
-  var sahdowRoot = document.querySelector('gradio-app').shadowRoot;
-  if(! sahdowRoot)  return false;
+function setupBracketChecking(id_prompt, id_counter){
+    var textarea = gradioApp().querySelector("#" + id_prompt + " > label > textarea");
+    var counter = gradioApp().getElementById(id_counter)
+    textarea.addEventListener("input", function(evt){
+        checkBrackets(evt, textarea, counter)
+    });
+}
 
-  var shadowTextArea = sahdowRoot.querySelectorAll('#txt2img_prompt > label > textarea');
-  if(shadowTextArea.length < 1)  return false;
+var shadowRootLoaded = setInterval(function() {
+    var shadowRoot = document.querySelector('gradio-app').shadowRoot;
+    if(! shadowRoot)  return false;
 
+    var shadowTextArea = shadowRoot.querySelectorAll('#txt2img_prompt > label > textarea');
+    if(shadowTextArea.length < 1)  return false;
 
-  clearInterval(shadowRootLoaded);
+    clearInterval(shadowRootLoaded);
 
-  document.querySelector('gradio-app').shadowRoot.querySelector('#txt2img_prompt').onkeyup = checkBrackets;
-  document.querySelector('gradio-app').shadowRoot.querySelector('#txt2img_neg_prompt').onkeyup = checkBrackets;
-  document.querySelector('gradio-app').shadowRoot.querySelector('#img2img_prompt').onkeyup = checkBrackets;
-  document.querySelector('gradio-app').shadowRoot.querySelector('#img2img_neg_prompt').onkeyup = checkBrackets;
+    setupBracketChecking('txt2img_prompt', 'txt2img_token_counter')
+    setupBracketChecking('txt2img_neg_prompt', 'txt2img_negative_token_counter')
+    setupBracketChecking('img2img_prompt', 'imgimg_token_counter')
+    setupBracketChecking('img2img_neg_prompt', 'img2img_negative_token_counter')
 }, 1000);
diff --git a/javascript/ui.js b/javascript/ui.js
index 37788a3e..3ba90ca8 100644
--- a/javascript/ui.js
+++ b/javascript/ui.js
@@ -230,14 +230,26 @@ onUiUpdate(function(){
 
     json_elem.parentElement.style.display="none"
 
-	if (!txt2img_textarea) {
-		txt2img_textarea = gradioApp().querySelector("#txt2img_prompt > label > textarea");
-		txt2img_textarea?.addEventListener("input", () => update_token_counter("txt2img_token_button"));
-	}
-	if (!img2img_textarea) {
-		img2img_textarea = gradioApp().querySelector("#img2img_prompt > label > textarea");
-		img2img_textarea?.addEventListener("input", () => update_token_counter("img2img_token_button"));
-	}
+    function registerTextarea(id, id_counter, id_button){
+        var prompt = gradioApp().getElementById(id)
+        var counter = gradioApp().getElementById(id_counter)
+        var textarea = gradioApp().querySelector("#" + id + " > label > textarea");
+
+        if(counter.parentElement == prompt.parentElement){
+            return
+        }
+
+        prompt.parentElement.insertBefore(counter, prompt)
+        counter.classList.add("token-counter")
+        prompt.parentElement.style.position = "relative"
+
+		textarea.addEventListener("input", () => update_token_counter(id_button));
+    }
+
+    registerTextarea('txt2img_prompt', 'txt2img_token_counter', 'txt2img_token_button')
+    registerTextarea('txt2img_neg_prompt', 'txt2img_negative_token_counter', 'txt2img_negative_token_button')
+    registerTextarea('img2img_prompt', 'img2img_token_counter', 'img2img_token_button')
+    registerTextarea('img2img_neg_prompt', 'img2img_negative_token_counter', 'img2img_negative_token_button')
 
     show_all_pages = gradioApp().getElementById('settings_show_all_pages')
     settings_tabs = gradioApp().querySelector('#settings div')
@@ -249,6 +261,7 @@ onUiUpdate(function(){
             })
         }
     }
+
 })
 
 
diff --git a/modules/ui.py b/modules/ui.py
index eb45a128..06c11848 100644
--- a/modules/ui.py
+++ b/modules/ui.py
@@ -335,28 +335,23 @@ def update_token_counter(text, steps):
     flat_prompts = reduce(lambda list1, list2: list1+list2, prompt_schedules)

     prompts = [prompt_text for step, prompt_text in flat_prompts]

     token_count, max_length = max([model_hijack.get_prompt_lengths(prompt) for prompt in prompts], key=lambda args: args[0])

-    style_class = ' class="red"' if (token_count > max_length) else ""

-    return f"<span {style_class}>{token_count}/{max_length}</span>"

+    return f"<span class='gr-box gr-text-input'>{token_count}/{max_length}</span>"

 

 

 def create_toprow(is_img2img):

     id_part = "img2img" if is_img2img else "txt2img"

 

-    with gr.Row(elem_id="toprow"):

-        with gr.Column(scale=6):

+    with gr.Row(elem_id=f"{id_part}_toprow", variant="compact"):

+        with gr.Column(elem_id=f"{id_part}_prompt_container", scale=6):

             with gr.Row():

                 with gr.Column(scale=80):

                     with gr.Row():

-                        prompt = gr.Textbox(label="Prompt", elem_id=f"{id_part}_prompt", show_label=False, lines=2,

-                            placeholder="Prompt (press Ctrl+Enter or Alt+Enter to generate)"

-                        )

+                        prompt = gr.Textbox(label="Prompt", elem_id=f"{id_part}_prompt", show_label=False, lines=2, placeholder="Prompt (press Ctrl+Enter or Alt+Enter to generate)")

 

             with gr.Row():

                 with gr.Column(scale=80):

                     with gr.Row():

-                        negative_prompt = gr.Textbox(label="Negative prompt", elem_id=f"{id_part}_neg_prompt", show_label=False, lines=2,

-                            placeholder="Negative prompt (press Ctrl+Enter or Alt+Enter to generate)"

-                        )

+                        negative_prompt = gr.Textbox(label="Negative prompt", elem_id=f"{id_part}_neg_prompt", show_label=False, lines=2, placeholder="Negative prompt (press Ctrl+Enter or Alt+Enter to generate)")

 

         with gr.Column(scale=1, elem_id="roll_col"):

             paste = gr.Button(value=paste_symbol, elem_id="paste")

@@ -365,6 +360,8 @@ def create_toprow(is_img2img):
             clear_prompt_button = gr.Button(value=clear_prompt_symbol, elem_id=f"{id_part}_clear_prompt")

             token_counter = gr.HTML(value="<span></span>", elem_id=f"{id_part}_token_counter")

             token_button = gr.Button(visible=False, elem_id=f"{id_part}_token_button")

+            negative_token_counter = gr.HTML(value="<span></span>", elem_id=f"{id_part}_negative_token_counter")

+            negative_token_button = gr.Button(visible=False, elem_id=f"{id_part}_negative_token_button")

 

             clear_prompt_button.click(

                 fn=lambda *x: x,

@@ -402,7 +399,7 @@ def create_toprow(is_img2img):
                 prompt_styles = gr.Dropdown(label="Styles", elem_id=f"{id_part}_styles", choices=[k for k, v in shared.prompt_styles.styles.items()], value=[], multiselect=True)

                 create_refresh_button(prompt_styles, shared.prompt_styles.reload, lambda: {"choices": [k for k, v in shared.prompt_styles.styles.items()]}, f"refresh_{id_part}_styles")

 

-    return prompt, prompt_styles, negative_prompt, submit, button_interrogate, button_deepbooru, prompt_style_apply, save_style, paste, token_counter, token_button

+    return prompt, prompt_styles, negative_prompt, submit, button_interrogate, button_deepbooru, prompt_style_apply, save_style, paste, token_counter, token_button, negative_token_counter, negative_token_button

 

 

 def setup_progressbar(*args, **kwargs):

@@ -619,7 +616,7 @@ def create_ui():
     modules.scripts.scripts_txt2img.initialize_scripts(is_img2img=False)

 

     with gr.Blocks(analytics_enabled=False) as txt2img_interface:

-        txt2img_prompt, txt2img_prompt_styles, txt2img_negative_prompt, submit, _, _,txt2img_prompt_style_apply, txt2img_save_style, txt2img_paste, token_counter, token_button = create_toprow(is_img2img=False)

+        txt2img_prompt, txt2img_prompt_styles, txt2img_negative_prompt, submit, _, _, txt2img_prompt_style_apply, txt2img_save_style, txt2img_paste, token_counter, token_button, negative_token_counter, negative_token_button = create_toprow(is_img2img=False)

 

         dummy_component = gr.Label(visible=False)

         txt_prompt_img = gr.File(label="", elem_id="txt2img_prompt_image", file_count="single", type="binary", visible=False)

@@ -795,12 +792,13 @@ def create_ui():
             ]

 

             token_button.click(fn=wrap_queued_call(update_token_counter), inputs=[txt2img_prompt, steps], outputs=[token_counter])

+            negative_token_button.click(fn=wrap_queued_call(update_token_counter), inputs=[txt2img_negative_prompt, steps], outputs=[negative_token_counter])

 

     modules.scripts.scripts_current = modules.scripts.scripts_img2img

     modules.scripts.scripts_img2img.initialize_scripts(is_img2img=True)

 

     with gr.Blocks(analytics_enabled=False) as img2img_interface:

-        img2img_prompt, img2img_prompt_styles, img2img_negative_prompt, submit, img2img_interrogate, img2img_deepbooru, img2img_prompt_style_apply, img2img_save_style, img2img_paste,token_counter, token_button = create_toprow(is_img2img=True)

+        img2img_prompt, img2img_prompt_styles, img2img_negative_prompt, submit, img2img_interrogate, img2img_deepbooru, img2img_prompt_style_apply, img2img_save_style, img2img_paste, token_counter, token_button, negative_token_counter, negative_token_button = create_toprow(is_img2img=True)

 

         img2img_prompt_img = gr.File(label="", elem_id="img2img_prompt_image", file_count="single", type="binary", visible=False)

 

@@ -1064,6 +1062,7 @@ def create_ui():
                 )

 

             token_button.click(fn=update_token_counter, inputs=[img2img_prompt, steps], outputs=[token_counter])

+            negative_token_button.click(fn=wrap_queued_call(update_token_counter), inputs=[txt2img_negative_prompt, steps], outputs=[negative_token_counter])

 

             img2img_paste_fields = [

                 (img2img_prompt, "Prompt"),

diff --git a/style.css b/style.css
index c10e32a1..994932de 100644
--- a/style.css
+++ b/style.css
@@ -2,12 +2,26 @@
     max-width: 100%;

 }

 

-#txt2img_token_counter {

-    height: 0px;

+.token-counter{

+    position: absolute;

+    display: inline-block;

+    right: 2em;

+    min-width: 0 !important;

+    width: auto;

+    z-index: 100;

+}

+

+.token-counter.error span{

+    box-shadow: 0 0 0.0 0.3em rgba(255,0,0,0.15), inset 0 0 0.6em rgba(255,0,0,0.075);

+    border: 2px solid rgba(255,0,0,0.4) !important;

 }

 

-#img2img_token_counter {

-    height: 0px;

+.token-counter div{

+    display: inline;

+}

+

+.token-counter span{

+    padding: 0.1em 0.75em;

 }

 

 #sh{

@@ -113,7 +127,7 @@
 #roll_col{

     min-width: unset !important;

     flex-grow: 0 !important;

-    padding: 0.4em 0;

+    padding: 0 1em 0 0;

     gap: 0;

 }

 

@@ -160,16 +174,6 @@
     margin-bottom: 0;

 }

 

-#toprow div.gr-box, #toprow div.gr-form{

-    border: none;

-    gap: 0;

-    background: transparent;

-    box-shadow: none;

-}

-#toprow div{

-    gap: 0;

-}

-

 #resize_mode{

     flex: 1.5;

 }

@@ -706,6 +710,14 @@ footer {
     opacity: 0.5;

 }

 

+[id*='_prompt_container']{

+    gap: 0;

+}

+

+[id*='_prompt_container'] > div{

+    margin: -0.4em 0 0 0;

+}

+

 .gr-compact {

     border: none;

 }

@@ -715,8 +727,11 @@ footer {
     margin-left: 0.8em;

 }

 

+.gr-compact{

+    overflow: visible;

+}

+

 .gr-compact > *{

-    margin-top: 0.5em !important;

 }

 

 .gr-compact .gr-block, .gr-compact .gr-form{