add remove button

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

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);
 
 
 

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
+
+
+

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
+

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

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'}
 

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"),