From 9ba3896f61ebe5bed012f1538f58facbc0fd6875 Mon Sep 17 00:00:00 2001 From: Zhang Yan <18758266469@163.com> Date: Sat, 8 Apr 2023 05:37:08 +0800 Subject: [PATCH] negWords support; The info file needs to manually add the negWords field. --- javascript/civitai_helper.js | 5 ++- scripts/ch_lib/js_action_civitai.py | 65 +++++++++++++++++++---------- scripts/civitai_helper.py | 2 +- 3 files changed, 49 insertions(+), 23 deletions(-) diff --git a/javascript/civitai_helper.js b/javascript/civitai_helper.js index 4454d0d..991a67a 100644 --- a/javascript/civitai_helper.js +++ b/javascript/civitai_helper.js @@ -221,12 +221,15 @@ function add_trigger_words(event, model_type, search_term){ msg["action"] = "add_trigger_words"; msg["model_type"] = model_type; msg["search_term"] = search_term; - msg["neg_prompt"] = ""; // get active prompt let act_prompt = getActivePrompt(); msg["prompt"] = act_prompt.value; + // get active neg prompt + let neg_prompt = getActiveNegativePrompt(); + msg["neg_prompt"] = neg_prompt.value; + // fill to msg box send_ch_py_msg(msg) diff --git a/scripts/ch_lib/js_action_civitai.py b/scripts/ch_lib/js_action_civitai.py index 0c81853..bfe6044 100644 --- a/scripts/ch_lib/js_action_civitai.py +++ b/scripts/ch_lib/js_action_civitai.py @@ -77,40 +77,63 @@ def add_trigger_words(msg): model_type = result["model_type"] search_term = result["search_term"] prompt = result["prompt"] + neg_prompt = result["neg_prompt"] + trainedWords = [] + negWords = [] model_info = civitai.load_model_info_by_search_term(model_type, search_term) if not model_info: util.printD(f"Failed to get model info for {model_type} {search_term}") - return [prompt, prompt] - - if "trainedWords" not in model_info.keys(): + return [prompt, prompt, neg_prompt, neg_prompt] + + if "trainedWords" in model_info.keys(): + trainedWords = model_info["trainedWords"] + if not trainedWords: + util.printD(f"No trainedWords from info file for {model_type} {search_term}") + trainedWords = [] + + if len(trainedWords) == 0: + util.printD(f"trainedWords from info file for {model_type} {search_term} is empty") + trainedWords = [] + else: util.printD(f"Failed to get trainedWords from info file for {model_type} {search_term}") - return [prompt, prompt] - - trainedWords = model_info["trainedWords"] - if not trainedWords: - util.printD(f"No trainedWords from info file for {model_type} {search_term}") - return [prompt, prompt] - - if len(trainedWords) == 0: - util.printD(f"trainedWords from info file for {model_type} {search_term} is empty") - return [prompt, prompt] - + + if "negWords" in model_info.keys(): + negWords = model_info["negWords"] + if not negWords: + util.printD(f"No negWords from info file for {model_type} {search_term}") + negWords = [] + + if len(negWords) == 0: + util.printD(f"negWords from info file for {model_type} {search_term} is empty") + negWords = [] + else: + util.printD(f"Failed to get negWords from info file for {model_type} {search_term}") + # get ful trigger words trigger_words = "" - for word in trainedWords: - trigger_words = trigger_words + word + ", " + if trainedWords: + for word in trainedWords: + trigger_words = trigger_words + word + ", " - new_prompt = prompt + " " + trigger_words + neg_words = "" + if negWords: + for word in negWords: + neg_words = neg_words + word + ", " + + new_prompt = prompt + " " + trigger_words if trigger_words != "" else prompt + new_neg_prompt = neg_prompt + " " + neg_words if neg_words != "" else neg_prompt util.printD("trigger_words: " + trigger_words) util.printD("prompt: " + prompt) util.printD("new_prompt: " + new_prompt) + util.printD("new_prompt: " + new_prompt) + util.printD("new_neg_prompt: " + new_neg_prompt) util.printD("End add_trigger_words") # add to prompt - return [new_prompt, new_prompt] + return [new_prompt, new_neg_prompt, new_prompt, new_neg_prompt] @@ -124,7 +147,7 @@ def use_preview_image_prompt(msg): if not result: util.printD("Parsing js ms failed") return - + model_type = result["model_type"] search_term = result["search_term"] prompt = result["prompt"] @@ -135,11 +158,11 @@ def use_preview_image_prompt(msg): if not model_info: util.printD(f"Failed to get model info for {model_type} {search_term}") return [prompt, neg_prompt, prompt, neg_prompt] - + if "images" not in model_info.keys(): util.printD(f"Failed to get images from info file for {model_type} {search_term}") return [prompt, neg_prompt, prompt, neg_prompt] - + images = model_info["images"] if not images: util.printD(f"No images from info file for {model_type} {search_term}") diff --git a/scripts/civitai_helper.py b/scripts/civitai_helper.py index c5d84a5..7304863 100644 --- a/scripts/civitai_helper.py +++ b/scripts/civitai_helper.py @@ -204,7 +204,7 @@ def on_ui_tabs(): # js action js_open_url_btn.click(js_action_civitai.open_model_url, inputs=[js_msg_txtbox, open_url_with_js_ckb], outputs=py_msg_txtbox) - js_add_trigger_words_btn.click(js_action_civitai.add_trigger_words, inputs=[js_msg_txtbox], outputs=[txt2img_prompt, img2img_prompt]) + js_add_trigger_words_btn.click(js_action_civitai.add_trigger_words, inputs=[js_msg_txtbox], outputs=[txt2img_prompt, txt2img_neg_prompt, img2img_prompt, img2img_neg_prompt]) js_use_preview_prompt_btn.click(js_action_civitai.use_preview_image_prompt, inputs=[js_msg_txtbox], outputs=[txt2img_prompt, txt2img_neg_prompt, img2img_prompt, img2img_neg_prompt]) js_dl_model_new_version_btn.click(js_action_civitai.dl_model_new_version, inputs=[js_msg_txtbox, max_size_preview_ckb, skip_nsfw_preview_ckb], outputs=dl_log_md)