From fe030dcb1b1f61ee413208933cb5c8f2835ca8c8 Mon Sep 17 00:00:00 2001 From: butaixianran Date: Wed, 27 Sep 2023 08:43:32 +0800 Subject: [PATCH] add remove button --- README.md | 5 +- javascript/civitai_helper.js | 102 +++++++++++++++++++++++++++- scripts/ch_lib/js_action_civitai.py | 66 ++++++++++++++++++ scripts/ch_lib/model.py | 38 +++++++++++ scripts/ch_lib/msg_handler.py | 4 +- scripts/ch_lib/util.py | 2 +- scripts/civitai_helper.py | 2 + 7 files changed, 213 insertions(+), 6 deletions(-) diff --git a/README.md b/README.md index 91fa22f..754cc56 100644 --- a/README.md +++ b/README.md @@ -206,8 +206,11 @@ Since v1.5.5, we've already optimized the SHA256 function to the top. So the onl # Change Log +## v1.8.0 +* Add remove model button + ## v1.7.0 -* Fix all issues for latest sd webui +* Catch up with latest sd webui * Move all settings to Setting tab ## v1.6.4 diff --git a/javascript/civitai_helper.js b/javascript/civitai_helper.js index f9280be..52d7883 100644 --- a/javascript/civitai_helper.js +++ b/javascript/civitai_helper.js @@ -60,7 +60,7 @@ function get_ch_py_msg(){ // get msg from python side from a hidden textbox // it will try once in every sencond, until it reach the max try times -const get_new_ch_py_msg = (max_count=3) => new Promise((resolve, reject) => { +const get_new_ch_py_msg = (max_count=5) => new Promise((resolve, reject) => { console.log("run get_new_ch_py_msg") let count = 0; @@ -175,7 +175,14 @@ async function open_model_url(event, model_type, search_term){ event.preventDefault() //check response msg from python - let new_py_msg = await get_new_ch_py_msg(); + let new_py_msg = ""; + try { + new_py_msg = await get_new_ch_py_msg(); + + } catch (error) { + console.log(error); + } + console.log("new_py_msg:"); console.log(new_py_msg); @@ -285,6 +292,88 @@ function use_preview_prompt(event, model_type, search_term){ } +async function remove_card(event, model_type, search_term){ + console.log("start remove_card"); + + //get hidden components of extension + let js_remove_card_btn = gradioApp().getElementById("ch_js_remove_card_btn"); + if (!js_remove_card_btn) { + return + } + + // must confirm before removing + let rm_confirm = "\nConfirm to remove this model.\n\nCheck console log for detail."; + if (!confirm(rm_confirm)) { + return + } + + + //msg to python side + let msg = { + "action": "", + "model_type": "", + "search_term": "", + } + + + msg["action"] = "remove_card"; + msg["model_type"] = model_type; + msg["search_term"] = search_term; + msg["prompt"] = ""; + msg["neg_prompt"] = ""; + + // fill to msg box + send_ch_py_msg(msg) + + //click hidden button + js_remove_card_btn.click(); + + // stop parent event + event.stopPropagation() + event.preventDefault() + + //check response msg from python + let new_py_msg = ""; + try { + new_py_msg = await get_new_ch_py_msg(); + } catch (error) { + console.log(error); + new_py_msg = error; + } + + console.log("new_py_msg:"); + console.log(new_py_msg); + + //check msg + let result = "Done"; + //check msg + if (new_py_msg) { + result = new_py_msg; + } + + // alert result + alert(result); + + if (result=="Done"){ + console.log("refresh card list"); + //refresh card list + let active_tab = getActiveTabType(); + console.log("get active tab id: " + active_tab); + if (active_tab){ + let refresh_btn_id = active_tab + "_extra_refresh"; + let refresh_btn = gradioApp().getElementById(refresh_btn_id); + if (refresh_btn){ + console.log("click button: "+refresh_btn_id); + refresh_btn.click(); + } + } + } + + console.log("end remove_card"); + + +} + // download model's new version into SD at python side function ch_dl_model_new_version(event, model_path, version_id, download_url){ @@ -538,10 +627,19 @@ onUiLoaded(() => { use_preview_prompt_node.title = "Use prompt from preview image"; use_preview_prompt_node.setAttribute("onclick","use_preview_prompt(event, '"+model_type+"', '"+search_term+"')"); + let remove_card_node = document.createElement("a"); + remove_card_node.href = "#"; + remove_card_node.innerHTML = "❌"; + remove_card_node.className = "card-button"; + + remove_card_node.title = "Remove this model"; + remove_card_node.setAttribute("onclick","remove_card(event, '"+model_type+"', '"+search_term+"')"); + //add to card button_row.appendChild(open_url_node); button_row.appendChild(add_trigger_words_node); button_row.appendChild(use_preview_prompt_node); + button_row.appendChild(remove_card_node); diff --git a/scripts/ch_lib/js_action_civitai.py b/scripts/ch_lib/js_action_civitai.py index 0c81853..3caef33 100644 --- a/scripts/ch_lib/js_action_civitai.py +++ b/scripts/ch_lib/js_action_civitai.py @@ -254,3 +254,69 @@ def dl_model_new_version(msg, max_size_preview, skip_nsfw_preview): output = "Done. Model downloaded to: " + new_model_path util.printD(output) return output + + + +# remove a model and all related files +def remove_model_by_path(msg): + util.printD("Start remove_model_by_path") + + output = "" + result = msg_handler.parse_js_msg(msg) + if not result: + output = "Parsing js ms failed" + util.printD(output) + return output + + model_type = result["model_type"] + search_term = result["search_term"] + + model_path = model.get_model_path_by_search_term(model_type, search_term) + if not model_path: + output = f"Fail to get model for {model_type} {search_term}" + util.printD(output) + return output + + + if not os.path.isfile(model_path): + output = f"Model {model_type} {search_term} does not exist, no need to remove" + util.printD(output) + return output + + # all files need to be removed + related_paths = [] + related_paths.append(model_path) + + + # get info file + base, ext = os.path.splitext(model_path) + info_path = base + model.info_ext + first_preview_path = base+".png" + sec_preview_path = base+".preview.png" + civitai_info_path = base + civitai.suffix + model.info_ext + + if os.path.isfile(civitai_info_path): + related_paths.append(civitai_info_path) + + if os.path.isfile(first_preview_path): + related_paths.append(first_preview_path) + + if os.path.isfile(sec_preview_path): + related_paths.append(sec_preview_path) + + if os.path.isfile(info_path): + related_paths.append(info_path) + + # remove files + for rp in related_paths: + if os.path.isfile(rp): + util.printD(f"Removing file {rp}") + os.remove(rp) + + util.printD(f"{len(related_paths)} file removed") + + util.printD("End remove_model_by_path") + return output + + + diff --git a/scripts/ch_lib/model.py b/scripts/ch_lib/model.py index 729012d..85c6698 100644 --- a/scripts/ch_lib/model.py +++ b/scripts/ch_lib/model.py @@ -118,3 +118,41 @@ def get_model_path_by_type_and_name(model_type:str, model_name:str) -> str: + +# get model path by model type and search_term +# parameter: model_type, search_term +# return: model_path +def get_model_path_by_search_term(model_type, search_term): + util.printD(f"Search model of {search_term} in {model_type}") + if model_type not in folders.keys(): + util.printD("unknow model type: " + model_type) + return + + # for lora: search_term = subfolderpath + model name + ext + " " + hash. And it always start with a / even there is no sub folder + # for ckp: search_term = subfolderpath + model name + ext + " " + hash + # for ti: search_term = subfolderpath + model name + ext + " " + hash + # for hyper: search_term = subfolderpath + model name + + + + model_sub_path = search_term.split()[0] + if model_sub_path[:1] == "/": + model_sub_path = model_sub_path[1:] + + if model_type == "hyper": + model_sub_path = model_sub_path+".pt" + + model_folder = folders[model_type] + + model_path = os.path.join(model_folder, model_sub_path) + + print("model_folder: " + model_folder) + print("model_sub_path: " + model_sub_path) + print("model_path: " + model_path) + + if not os.path.isfile(model_path): + util.printD("Can not find model file: " + model_path) + return + + return model_path + diff --git a/scripts/ch_lib/msg_handler.py b/scripts/ch_lib/msg_handler.py index 1132d77..80d3406 100644 --- a/scripts/ch_lib/msg_handler.py +++ b/scripts/ch_lib/msg_handler.py @@ -4,8 +4,8 @@ import json from . import util # action list -js_actions = ("open_url", "add_trigger_words", "use_preview_prompt", "dl_model_new_version") -py_actions = ("open_url") +js_actions = ("open_url", "add_trigger_words", "use_preview_prompt", "dl_model_new_version", "remove_card") +py_actions = ("open_url", "remove_card") # handle request from javascript diff --git a/scripts/ch_lib/util.py b/scripts/ch_lib/util.py index 60fe71a..fb171bf 100644 --- a/scripts/ch_lib/util.py +++ b/scripts/ch_lib/util.py @@ -6,7 +6,7 @@ import requests import shutil -version = "1.7.0" +version = "1.8.0" def_headers = {'User-Agent': 'Mozilla/5.0 (iPad; CPU OS 12_2 like Mac OS X) AppleWebKit/605.1.15 (KHTML, like Gecko) Mobile/15E148'} diff --git a/scripts/civitai_helper.py b/scripts/civitai_helper.py index b3dde37..2b0b6cb 100644 --- a/scripts/civitai_helper.py +++ b/scripts/civitai_helper.py @@ -200,6 +200,7 @@ def on_ui_tabs(): js_add_trigger_words_btn = gr.Button(value="Add Trigger Words", visible=False, elem_id="ch_js_add_trigger_words_btn") js_use_preview_prompt_btn = gr.Button(value="Use Prompt from Preview Image", visible=False, elem_id="ch_js_use_preview_prompt_btn") js_dl_model_new_version_btn = gr.Button(value="Download Model's new version", visible=False, elem_id="ch_js_dl_model_new_version_btn") + js_remove_card_btn = gr.Button(value="Remove Card", visible=False, elem_id="ch_js_remove_card_btn") # ====events==== # Scan Models for Civitai @@ -223,6 +224,7 @@ def on_ui_tabs(): js_add_trigger_words_btn.click(js_action_civitai.add_trigger_words, inputs=[js_msg_txtbox], outputs=[txt2img_prompt, img2img_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(dl_model_new_version, inputs=[js_msg_txtbox], outputs=dl_log_md) + js_remove_card_btn.click(js_action_civitai.remove_model_by_path, inputs=[js_msg_txtbox], outputs=py_msg_txtbox) # the third parameter is the element id on html, with a "tab_" as prefix return (civitai_helper , "Civitai Helper", "civitai_helper"),