From 3c6a049fc3c6b54ada3736710a7e86663ea7f3d9 Mon Sep 17 00:00:00 2001
From: Liam <liamthekerr@gmail.com>
Date: Fri, 30 Sep 2022 12:12:44 -0400
Subject: [PATCH] consolidated token counter functions

---
 javascript/ui.js | 21 +++++++++------------
 modules/ui.py    |  6 +++---
 2 files changed, 12 insertions(+), 15 deletions(-)

diff --git a/javascript/ui.js b/javascript/ui.js
index 88fd45ae..f94ed081 100644
--- a/javascript/ui.js
+++ b/javascript/ui.js
@@ -199,21 +199,18 @@ let txt2img_textarea, img2img_textarea = undefined;
 let wait_time = 800
 let token_timeout;
 
-function roll_artist_txt2img(prompt_text) {
+function update_txt2img_tokens(...args) {
 	update_token_counter("txt2img_token_button")
-	return prompt_text;
+	if (args.length == 2)
+		return args[0]
+	return args;
 }
-function roll_artist_img2img(prompt_text) {
+
+function update_img2img_tokens(...args) {
 	update_token_counter("img2img_token_button")
-	return prompt_text;
-}
-function update_style_txt2img(prompt_text, negative_prompt, style1, style2) {
-	update_token_counter("txt2img_token_button")
-	return [prompt_text, negative_prompt, style1, style2]
-}
-function update_style_img2img(prompt_text, negative_prompt, style1, style2) {
-	update_token_counter("img2img_token_button")
-	return [prompt_text, negative_prompt, style1, style2]
+	if (args.length == 2)
+		return args[0]
+	return args;
 }
 
 function update_token_counter(button_id) {
diff --git a/modules/ui.py b/modules/ui.py
index 6bf28562..40c08984 100644
--- a/modules/ui.py
+++ b/modules/ui.py
@@ -543,7 +543,7 @@ def create_ui(txt2img, img2img, run_extras, run_pnginfo, run_modelmerger):
 
             roll.click(
                 fn=roll_artist,
-                _js="roll_artist_txt2img",
+                _js="update_txt2img_tokens",
                 inputs=[
                     txt2img_prompt,
                 ],
@@ -749,7 +749,7 @@ def create_ui(txt2img, img2img, run_extras, run_pnginfo, run_modelmerger):
 
             roll.click(
                 fn=roll_artist,
-                _js="roll_artist_img2img",
+                _js="update_img2img_tokens",
                 inputs=[
                     img2img_prompt,
                 ],
@@ -760,7 +760,7 @@ def create_ui(txt2img, img2img, run_extras, run_pnginfo, run_modelmerger):
 
             prompts = [(txt2img_prompt, txt2img_negative_prompt), (img2img_prompt, img2img_negative_prompt)]
             style_dropdowns = [(txt2img_prompt_style, txt2img_prompt_style2), (img2img_prompt_style, img2img_prompt_style2)]
-            style_js_funcs = ["update_style_txt2img", "update_style_img2img"]
+            style_js_funcs = ["update_txt2img_tokens", "update_img2img_tokens"]
 
             for button, (prompt, negative_prompt) in zip([txt2img_save_style, img2img_save_style], prompts):
                 button.click(