From 27663f3366f24966d3da2d508a468f19323d9e84 Mon Sep 17 00:00:00 2001
From: zahand <48160881+muerrilla@users.noreply.github.com>
Date: Wed, 5 Apr 2023 05:09:54 +0330
Subject: [PATCH] Put an end to all the tab shenanigans

---
 javascript/setup.js | 31 +++++++++++++------------------
 scripts/npw.py      | 16 +++++++++-------
 2 files changed, 22 insertions(+), 25 deletions(-)

diff --git a/javascript/setup.js b/javascript/setup.js
index 306040b..89d43c4 100644
--- a/javascript/setup.js
+++ b/javascript/setup.js
@@ -1,31 +1,26 @@
 function setupNPW() {
-  fixAccordion();  
-  fixInputs();
+  fixAccordion('tab_txt2img');  
+  fixAccordion('tab_img2img');  
+  fixInputs('tab_txt2img');
+  fixInputs('tab_img2img');
 }
 
-function fixInputs() {
-  const npwSlider_t2i = document.querySelector("#tab_txt2img #npw-slider");
-  const npwSlider_i2i = document.querySelector("#tab_img2img #npw-slider");
+function fixInputs(tab) {
+  const npwSlider = document.querySelector(`#${tab} #npw-slider`);
 
-  npwSlider_t2i.querySelector('.head').remove(); 
-  npwSlider_i2i.querySelector('.head').remove(); 
+  npwSlider.querySelector('.head').remove(); 
 
   const newSpan = document.createElement("span");
   newSpan.innerHTML = "Negative Prompt Weight";
-  const ancestor_t2i = npwSlider_t2i.parentNode.parentNode.parentNode;
-  ancestor_t2i.insertBefore(newSpan, ancestor_t2i.firstChild);
-  const ancestor_i2i = npwSlider_i2i.parentNode.parentNode.parentNode;
-  ancestor_i2i.insertBefore(newSpan.cloneNode(true), ancestor_i2i.firstChild);
+  const ancestor = npwSlider.parentNode.parentNode.parentNode;
+  ancestor.insertBefore(newSpan, ancestor.firstChild);
 
-  document.querySelector('#tab_txt2img #npw-number input[type="number"]').setAttribute("step", "0.01");
-  document.querySelector('#tab_img2img #npw-number input[type="number"]').setAttribute("step", "0.01");
+  document.querySelector(`#${tab} #npw-number input[type="number"]`).setAttribute("step", "0.01");
 }
 
-function fixAccordion() {
-  document.querySelector('#tab_txt2img #npw .icon').remove();
-  document.querySelector('#tab_img2img #npw .icon').remove();
-  document.querySelector('#tab_txt2img #npw .open').remove();  
-  document.querySelector('#tab_img2img #npw .open').remove();  
+function fixAccordion(tab) {
+  document.querySelector(`#${tab} #npw .icon`).remove();
+  document.querySelector(`#${tab} #npw .open`).remove();  
 }
 
 onUiLoaded(setupNPW);
diff --git a/scripts/npw.py b/scripts/npw.py
index 470e2f6..3ce3500 100644
--- a/scripts/npw.py
+++ b/scripts/npw.py
@@ -26,13 +26,15 @@ class Script(scripts.Script):
                         reset_but = gr.Button(value='✕', elem_id='npw-x').style(full_width=False, size='sm')               
 
             js = """(v) => {
-                let t=document.querySelector('#txt2img_negative_token_counter'),
-                    i=document.querySelector('#img2img_negative_token_counter');
-                t.style.cssText+=`outline:4px solid rgba(255,0,128,${Math.sqrt(Math.abs(v-1))}); border-radius: 0.4em !important;`
-                i.style.cssText+=`outline:4px solid rgba(255,0,128,${Math.sqrt(Math.abs(v-1))}); border-radius: 0.4em !important;`
-                return v;
-                }"""                
-
+              ['#txt2img_negative_token_counter', '#img2img_negative_token_counter'].forEach((selector, index) => {
+                const element = document.querySelector(selector);
+                if (document.querySelector(`#tab_${index ? 'img2img' : 'txt2img'}`).style.display === 'block') {
+                  element.style.cssText += `outline:4px solid rgba(255,0,128,${Math.sqrt(Math.abs(v-1))}); border-radius: 0.4em !important;`;
+                }
+              });
+              return v;
+            }"""
+               
             weight_input.change(None, [weight_input], weight_input_slider, _js=js)
             weight_input_slider.release(None, weight_input_slider, weight_input, _js="(x) => x")
             reset_but.click(None, [], [weight_input,weight_input_slider], _js="(x) => [1,1]")