auto

commit.bat

@@ -0,0 +1,30 @@
+@echo off
+
+git config user.name "anonymous"
+git config user.email my@gmail.com
+
+set /P comment=Comment?
+
+if not "%comment%" == "" goto :next
+set comment=auto
+:next
+
+git reset
+
+::git add *.* --force
+::git add *
+git add --all
+
+echo ----- git add done ------
+
+git status
+
+echo ready. press any key to commit START.
+::timeout /t 1
+pause
+
+git commit -m "%comment%"
+
+git push -u origin main
+
+pause

extensions/Stable-Diffusion-Webui-Civitai-Helper/about_version2.md

@@ -1,30 +0,0 @@
-## About Civitai Helper2: Model Info Helper
-The first version of "Civitai Helper2" is already finished and I have used it for a while on my local machine. It works exactly as in the demo video I have released before(with some more functions).   
-
-[https://www.youtube.com/watch?v=mPcKwQDDH8s](https://www.youtube.com/watch?v=mPcKwQDDH8s)  
-
-But the developing process is extremely painful. Its maintenance will also be difficult. If I release it, with all those requests and issues from users, my life gonna be in hell.  
-
-The reason for this is, gradio is not suitable for a professional level's management system. But under gradio's framework and SD webui's design, a lot of javascript things won't work well.  
-
-So, it is too painful that I don't want to do any further maintenance and development on it. In the meantime, a lot of alternatives become popular. For example, ComfyUI and InvokeAI 3.x.
-
-Especially InvokeAI 3.x, with its Unified Canvas and build-in node system, now we have a photoshop level's SD tool for free.  
-
-If you are not familiar with InvokeAI 3.x, you can check this video:  
-[https://www.youtube.com/watch?v=1Iz4F7o6hgQ](https://www.youtube.com/watch?v=1Iz4F7o6hgQ)  
-And its unified canvas system:  
-[https://www.youtube.com/watch?v=aU0jGZpDIVc](https://www.youtube.com/watch?v=aU0jGZpDIVc)  
-
-So, for Model Info Helper, I'm re-designing it into a stand alone model and workflow management system, which should work with most popular SD UI, not just for SD webui.  
-
-And I'll take down all workflow management functions so this stand alone version can be released sooner.   
-
-Since it won't be released as estimated, I have updated the old civitai helper 1.x to work with latest SD webui.  
-
-
-
-
-
-
-

extensions/Stable-Diffusion-Webui-Civitai-Helper/javascript/civitai_helper.js

@@ -1,38 +1,61 @@
 "use strict";
 
 
-function ch_convert_file_path_to_url(path){
+function ch_convert_file_path_to_url(path) {
     let prefix = "file=";
     let path_to_url = path.replaceAll('\\', '/');
     return prefix+path_to_url;
 }
 
-function ch_img_node_str(path){
+function ch_img_node_str(path) {
     return `<img src='${ch_convert_file_path_to_url(path)}' style="width:24px"/>`;
 }
 
-
-function ch_gradio_version(){
+function ch_gradio_version() {
     let foot = gradioApp().getElementById("footer");
-    if (!foot){return null;}
+    if (!foot) { return null; }
 
     let versions = foot.querySelector(".versions");
-    if (!versions){return null;}
+    if (!versions) { return null; }
 
-    if (versions.innerHTML.indexOf("gradio: 3.16.2")>0) {
+    if (versions.textContent.indexOf("gradio: 3.16.2")>0) {
         return "3.16.2";
     } else {
         return "3.23.0";
     }
-    
+}
+
+/*
+ * Functions for scan for duplicates elements.
+ */
+
+let ch_path_el = null;
+function display_ch_path(event, path) {
+    if (!ch_path_el) {
+        ch_path_el = document.createElement("div");
+        ch_path_el.id = "ch_path_el";
+        document.body.appendChild(ch_path_el);
+    }
+
+    ch_path_el.textContent = path;
+    ch_path_el.style.display = "block";
+}
+
+function move_ch_path(event) {
+    ch_path_el.style.top = `calc(${event.clientY}px - 2em)`;
+    ch_path_el.style.left = `calc(${event.clientX}px + 2em)`;
+}
+
+function hide_ch_path(event) {
+    ch_path_el.style.display = "none";
 }
 
 
 // send msg to python side by filling a hidden text box
 // then will click a button to trigger an action
 // msg is an object, not a string, will be stringify in this function
-function send_ch_py_msg(msg){
-    console.log("run send_ch_py_msg")
+function send_ch_py_msg(msg) {
+    console.log("run send_ch_py_msg");
     let js_msg_txtbox = gradioApp().querySelector("#ch_js_msg_txtbox textarea");
     if (js_msg_txtbox && msg) {
         // fill to msg box
@@ -44,16 +67,16 @@ function send_ch_py_msg(msg){
 
 // get msg from python side from a hidden textbox
 // normally this is an old msg, need to wait for a new msg
-function get_ch_py_msg(){
-    console.log("run get_ch_py_msg")
+function get_ch_py_msg() {
+    console.log("run get_ch_py_msg");
     const py_msg_txtbox = gradioApp().querySelector("#ch_py_msg_txtbox textarea");
     if (py_msg_txtbox && py_msg_txtbox.value) {
         console.log("find py_msg_txtbox");
         console.log("py_msg_txtbox value: ");
-        console.log(py_msg_txtbox.value)
-        return py_msg_txtbox.value
+        console.log(py_msg_txtbox.value);
+        return py_msg_txtbox.value;
     } else {
-        return ""
+        return "";
     }
 }
 
@@ -61,7 +84,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=5) => new Promise((resolve, reject) => {
-    console.log("run get_new_ch_py_msg")
+    console.log("run get_new_ch_py_msg");
 
     let count = 0;
     let new_msg = "";
@@ -73,11 +96,11 @@ const get_new_ch_py_msg = (max_count=5) => new Promise((resolve, reject) => {
         if (py_msg_txtbox && py_msg_txtbox.value) {
             console.log("find py_msg_txtbox");
             console.log("py_msg_txtbox value: ");
-            console.log(py_msg_txtbox.value)
+            console.log(py_msg_txtbox.value);
 
-            new_msg = py_msg_txtbox.value
+            new_msg = py_msg_txtbox.value;
             if (new_msg != "") {
-                find_msg=true
+                find_msg = true;
             }
         }
 
@@ -98,7 +121,7 @@ const get_new_ch_py_msg = (max_count=5) => new Promise((resolve, reject) => {
         }
 
     }, 1000);
-})
+});
 
 
 function getActiveTabType() {
@@ -112,7 +135,9 @@ function getActiveTabType() {
     return null;
 }
 
-
+function getExtraTabs(prefix) {
+    return gradioApp().getElementById(prefix + "_extra_tabs");
+}
 
 function getActivePrompt() {
     const currentTab = get_uiCurrentTabContent();
@@ -138,16 +163,20 @@ function getActiveNegativePrompt() {
 
 
 //button's click function
-async function open_model_url(event, model_type, search_term){
+async function open_model_url(event, model_type, search_term) {
     console.log("start open_model_url");
 
-    //get hidden components of extension 
+    // stop parent event
+    event.stopPropagation();
+    event.preventDefault();
+
+    //get hidden components of extension
     let js_open_url_btn = gradioApp().getElementById("ch_js_open_url_btn");
     if (!js_open_url_btn) {
-        return
+        console.log("Failed to find js_open_url_btn");
+        return;
     }
 
-
     //msg to python side
     let msg = {
         "action": "",
@@ -155,8 +184,7 @@ async function open_model_url(event, model_type, search_term){
         "search_term": "",
         "prompt": "",
         "neg_prompt": "",
-    }
-
+    };
 
     msg["action"] = "open_url";
     msg["model_type"] = model_type;
@@ -165,24 +193,19 @@ async function open_model_url(event, model_type, search_term){
     msg["neg_prompt"] = "";
 
     // fill to msg box
-    send_ch_py_msg(msg)
+    send_ch_py_msg(msg);
 
     //click hidden button
     js_open_url_btn.click();
 
-    // stop parent event
-    event.stopPropagation()
-    event.preventDefault()
-
     //check response msg from python
-    let new_py_msg = "";
+    let new_py_msg = null;
     try {
         new_py_msg = await get_new_ch_py_msg();
-        
     } catch (error) {
-        console.log(error);        
+        console.log(error);
     }
-    
+
     console.log("new_py_msg:");
     console.log(new_py_msg);
 
@@ -194,28 +217,25 @@ async function open_model_url(event, model_type, search_term){
             if (py_msg_json.content.url) {
                 window.open(py_msg_json.content.url, "_blank");
             }
-
         }
-
-
     }
 
-    
     console.log("end open_model_url");
 
-
 }
 
-function add_trigger_words(event, model_type, search_term){
+function add_trigger_words(event, model_type, search_term) {
     console.log("start add_trigger_words");
 
-    //get hidden components of extension 
+    event.stopPropagation();
+    event.preventDefault();
+
+    //get hidden components of extension
     let js_add_trigger_words_btn = gradioApp().getElementById("ch_js_add_trigger_words_btn");
     if (!js_add_trigger_words_btn) {
-        return
+        return;
     }
 
-
     //msg to python side
     let msg = {
         "action": "",
@@ -223,7 +243,7 @@ function add_trigger_words(event, model_type, search_term){
         "search_term": "",
         "prompt": "",
         "neg_prompt": "",
-    }
+    };
 
     msg["action"] = "add_trigger_words";
     msg["model_type"] = model_type;
@@ -235,26 +255,25 @@ function add_trigger_words(event, model_type, search_term){
     msg["prompt"] = act_prompt.value;
 
     // fill to msg box
-    send_ch_py_msg(msg)
+    send_ch_py_msg(msg);
 
     //click hidden button
     js_add_trigger_words_btn.click();
 
     console.log("end add_trigger_words");
 
-    event.stopPropagation()
-    event.preventDefault()
-
-    
 }
 
-function use_preview_prompt(event, model_type, search_term){
+function use_preview_prompt(event, model_type, search_term) {
     console.log("start use_preview_prompt");
 
-    //get hidden components of extension 
+    event.stopPropagation();
+    event.preventDefault();
+
+    //get hidden components of extension
     let js_use_preview_prompt_btn = gradioApp().getElementById("ch_js_use_preview_prompt_btn");
     if (!js_use_preview_prompt_btn) {
-        return
+        return;
     }
 
     //msg to python side
@@ -264,7 +283,7 @@ function use_preview_prompt(event, model_type, search_term){
         "search_term": "",
         "prompt": "",
         "neg_prompt": "",
-    }
+    };
 
     msg["action"] = "use_preview_prompt";
     msg["model_type"] = model_type;
@@ -279,35 +298,64 @@ function use_preview_prompt(event, model_type, search_term){
     msg["neg_prompt"] = neg_prompt.value;
 
     // fill to msg box
-    send_ch_py_msg(msg)
+    send_ch_py_msg(msg);
 
     //click hidden button
     js_use_preview_prompt_btn.click();
 
     console.log("end use_preview_prompt");
 
-    event.stopPropagation()
-    event.preventDefault()
+}
+
+
+async function remove_dup_card(event, model_type, search_term) {
+    event.stopPropagation();
+    event.preventDefault();
+
+    let el = event.currentTarget;
+
+    let success = await remove_card(event, model_type, search_term);
+
+    if (success === true) {
+        let parent = el.parentElement;
 
+        let sha256 = search_term.split(" ").pop().toUpperCase();
+        let row_id = `ch_${sha256}`;
+        let cards_id = `${row_id}_cards`;
+
+        let cards = document.getElementById(cards_id);
+
+        cards.removeChild(parent);
+
+        if (cards.children.length < 2) {
+            let row = document.getElementById(row_id);
+            row.parentElement.removeChild(row);
+        }
+    }
 }
 
 
-async function remove_card(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 status = false
+
+    // stop parent event
+    event.stopPropagation()
+    event.preventDefault()
+
+    //get hidden components of extension
     let js_remove_card_btn = gradioApp().getElementById("ch_js_remove_card_btn");
     if (!js_remove_card_btn) {
-        return
+        return status;
     }
 
     // must confirm before removing
-    let rm_confirm = "\nConfirm to remove this model.\n\nCheck console log for detail.";
+    let rm_confirm = "\nConfirm to remove this model and all related files. This process is irreversible.";
     if (!confirm(rm_confirm)) {
-        return
+        return status;
     }
 
-
     //msg to python side
     let msg = {
         "action": "",
@@ -315,12 +363,9 @@ async function remove_card(event, model_type, search_term){
         "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)
@@ -328,10 +373,76 @@ async function remove_card(event, model_type, search_term){
     //click hidden button
     js_remove_card_btn.click();
 
+    //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;
+    }
+
+    if (result == "Done") {
+        status = true;
+        refresh_cards_list(model_type);
+    }
+
+    console.log("end remove_card");
+
+    return status;
+
+}
+
+
+async function rename_card(event, model_type, search_term, model_name) {
+    console.log("start rename_card");
+
     // stop parent event
     event.stopPropagation()
     event.preventDefault()
 
+    //get hidden components of extension
+    let js_rename_card_btn = gradioApp().getElementById("ch_js_rename_card_btn");
+    if (!js_rename_card_btn) {
+        return;
+    }
+
+    // must confirm before removing
+    let rename_prompt = "\nRename this model to:";
+    let new_name = prompt(rename_prompt, model_name);
+    if (!new_name) {
+        return;
+    }
+
+    //msg to python side
+    let msg = {
+        "action": "",
+        "model_type": "",
+        "search_term": "",
+        "new_name": "",
+    }
+
+    msg["action"] = "rename_card";
+    msg["model_type"] = model_type;
+    msg["search_term"] = search_term;
+    msg["new_name"] = new_name;
+
+    // fill to msg box
+    send_ch_py_msg(msg)
+
+    //click hidden button
+    js_rename_card_btn.click();
+
     //check response msg from python
     let new_py_msg = "";
     try {
@@ -340,7 +451,7 @@ async function remove_card(event, model_type, search_term){
         console.log(error);
         new_py_msg = error;
     }
-    
+
     console.log("new_py_msg:");
     console.log(new_py_msg);
 
@@ -351,44 +462,49 @@ async function remove_card(event, model_type, search_term){
         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();
-            }
-        }
+    if (result == "Done") {
+        refresh_cards_list(model_type);
     }
-    
-    console.log("end remove_card");
+
+    console.log("end rename_card");
+
+}
 
 
+function replace_preview(e, page, type, name) {
+    // we have to create a whole hidden editor window to access preview replace functionality
+    extraNetworksEditUserMetadata(e, page, type, name);
+
+    // the editor window takes quite some time to populate
+    waitForEditor(page, type, name).then(editor => {
+        // Gather the buttons we need to both replace the preview and close the editor
+        let cancel_button = editor.querySelector('.edit-user-metadata-buttons button:first-of-type');
+        let replace_preview_button = editor.querySelector('.edit-user-metadata-buttons button:nth-of-type(2)');
+
+        replace_preview_button.click();
+        cancel_button.click();
+    });
 }
 
 
 // download model's new version into SD at python side
-function ch_dl_model_new_version(event, model_path, version_id, download_url){
+function ch_dl_model_new_version(event, model_path, version_id, download_url, model_type) {
     console.log("start ch_dl_model_new_version");
 
+    event.stopPropagation();
+    event.preventDefault();
+
     // must confirm before downloading
     let dl_confirm = "\nConfirm to download.\n\nCheck Download Model Section's log and console log for detail.";
+
     if (!confirm(dl_confirm)) {
-        return
+        return;
     }
 
-    //get hidden components of extension 
+    //get hidden components of extension
     let js_dl_model_new_version_btn = gradioApp().getElementById("ch_js_dl_model_new_version_btn");
     if (!js_dl_model_new_version_btn) {
-        return
+        return;
     }
 
     //msg to python side
@@ -397,58 +513,463 @@ function ch_dl_model_new_version(event, model_path, version_id, download_url){
         "model_path": "",
         "version_id": "",
         "download_url": "",
-    }
+    };
 
     msg["action"] = "dl_model_new_version";
     msg["model_path"] = model_path;
     msg["version_id"] = version_id;
     msg["download_url"] = download_url;
+    msg["model_type"] = model_type;
 
     // fill to msg box
-    send_ch_py_msg(msg)
+    send_ch_py_msg(msg);
 
     //click hidden button
     js_dl_model_new_version_btn.click();
 
     console.log("end dl_model_new_version");
 
-    event.stopPropagation()
-    event.preventDefault()
 
+}
+
+
+function refresh_cards_list(model_type) {
+    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) {
+            // webui 1.8
+            refresh_btn_id = `${active_tab}_${model_type}_extra_refresh`;
+            refresh_btn = gradioApp().getElementById(refresh_btn_id);
+        }
+        if (refresh_btn) {
+            console.log(`click button: ${refresh_btn_id}`);
+            refresh_btn.click();
+        }
+    }
+}
+
+function processCards(tab, extra_tab_els) {
+    if (!(opts && "ch_always_display" in opts)) {
+        // Lobe theme can cause a race condition.
+        console.log("Waiting for webui settings to become available");
+        console.log(opts);
+        const try_again = function () {
+            processCards(tab, extra_tab_els);
+        }
+        return setTimeout(try_again, 500);
+    }
+
+    const prefix_length = tab.length + 1;
+    for (const el of extra_tab_els) {
+        const model_type = el.id.slice(prefix_length, -6);
+        const cards = el.querySelectorAll('.card');
+        for (const card of cards) {
+            processSingleCard(tab, getShortModelTypeFromFull(model_type), card);
+        }
+    }
+}
+
+
+function getModelCardsEl(prefix, model_type) {
+    const id = prefix + "_" + model_type + "_cards";
+    return gradioApp().getElementById(id);
+}
+
+
+function waitForExtraTabs(tab, extra_tabs) {
+    function findTabs() {
+        const tab_elements = [];
+        for (const extra_tab of extra_tabs) {
+            const extra_tab_el = getModelCardsEl(tab, extra_tab);
+
+            if (extra_tab_el == null) {
+
+                // XXX lycoris models do not have their own tab in sdwebui 1.5
+                // most of the time. In the case that there is a LyCoris tab,
+                // it would have been added at the same time as the others,
+                // making it almost impossible to be null by the time we're at
+                // this point in the code if the other tabs are loaded.
+                if (extra_tab == 'lycoris') { continue; }
+
+                return null;
+            }
+
+            tab_elements.push(extra_tab_el);
+        }
+        return tab_elements;
+    }
+
+    const tab_elements = findTabs(tab, extra_tabs);
+    if (tab_elements) {
+        processCards(tab, tab_elements);
+    }
+
+    const observer = new MutationObserver(records => {
+        let tab_elements;
+        for (const record of records) {
+            if (record.type != "childList") {
+                continue;
+            }
+
+            tab_elements = findTabs(tab, extra_tabs);
+            if (!tab_elements) {
+                return;
+            }
+
+            processCards(tab, tab_elements);
+            return;
+        }
+    });
+
+    const extra_networks = getExtraTabs(tab);
+
+    const options = {
+        subtree: true,
+        childList: true,
+    };
+
+    observer.observe(extra_networks, options);
+
+}
+
+
+function waitForEditor(page, type, name) {
+    const id = page + '_' + type + '_edit_user_metadata';
+
+    return new Promise(resolve => {
+        let name_field;
+        const gradio = gradioApp();
 
+        const editor = gradio.getElementById(id);
+        const popup = gradio.querySelector(".global-popup");
+
+        if (popup != null) {
+            // hide the editor window so it doesn't get in the user's
+            // way while we wait for the replace preview functionality
+            // to become available.
+            popup.style.display = "none";
+        }
+
+        // not only do we need to wait for the editor,
+        // but also for it to populate with the model metadata.
+        if (editor != null) {
+            name_field = editor.querySelector('.extra-network-name');
+            if (name_field.textContent.trim() == name) {
+                return resolve(editor);
+            }
+        }
+
+        const observer = new MutationObserver(() => {
+            const editor = gradioApp().getElementById(id);
+            let name_field;
+            if (editor != null) {
+                name_field = editor.querySelector('.extra-network-name');
+                if (name_field.textContent.trim() == name) {
+                    resolve(editor);
+                    observer.disconnect();
+                }
+            }
+        });
+
+        observer.observe(document.body, {
+            subtree: true,
+            childList: true,
+        });
+    });
 }
 
 
+function getShortModelTypeFromFull(model_type_full) {
+    switch (model_type_full) {
+        case "textual_inversion":
+            return "ti";
+        case "hypernetworks":
+            return "hyper";
+        case "checkpoints":
+            return "ckp";
+        case "lora":
+        case "lycoris":
+            return model_type_full;
+    }
+}
+
+
+function getLongModelTypeFromShort(model_type_short) {
+    switch (model_type_short) {
+        case "ti":
+            return "textual_inversion";
+        case "hyper":
+            return "hypernetworks";
+        case "ckp":
+            return "checkpoints";
+        case "lora":
+        case "lycoris":
+            return model_type_short;
+    }
+}
+
+
+function isThumbMode(extra_network_node) {
+    if (extra_network_node?.className == "extra-network-thumbs") {
+        return true;
+    }
+    return false;
+}
+
+
+function processSingleCard(active_tab_type, active_extra_tab_type, card) {
+    let metadata_button = null;
+    let additional_node = null;
+    let replace_preview_btn = null;
+    let ul_node = null;
+    let model_name = "";
+    let search_term = "";
+    let model_type = active_extra_tab_type;
+    let js_model_type = getLongModelTypeFromShort(model_type);
+    let addedNodes = [];
+
+    let is_thumb_mode = isThumbMode(getModelCardsEl(active_tab_type, js_model_type));
+
+    //metadata_buttoncard
+    metadata_button = card.querySelector(".metadata-button");
+    //additional node
+    additional_node = card.querySelector(".actions .additional");
+    //get ul node, which is the parent of all buttons
+    ul_node = card.querySelector(".actions .additional ul");
+
+    // check thumb mode
+    if (is_thumb_mode) {
+        additional_node.style.display = null;
+
+        if (!ul_node) {
+            // nothing to do.
+            return;
+        }
+
+        if (opts["ch_show_btn_on_thumb"]) {
+            ul_node.style.background = btn_thumb_background;
+        } else {
+            let ch_btn_txts = ["💡", "🌐", "🏷️", "✏️", "❌"];
+
+            // remove existed buttons
+            //reset
+            ul_node.style.background = null;
+            // find all .a child nodes
+            let atags = ul_node.querySelectorAll("a");
+
+            for (let atag of atags) {
+                //reset display
+                atag.style.display = null;
+                //remove extension's button
+                if (ch_btn_txts.indexOf(atag.textContent)>=0) {
+                    //need to remove
+                    ul_node.removeChild(atag);
+                } else {
+                    //do not remove, just reset
+                    atag.textContent = replace_preview_text;
+                    atag.style.display = null;
+                    atag.style.fontSize = null;
+                    atag.style.position = null;
+                    atag.style.backgroundImage = null;
+                }
+            }
+
+            //also remove br tag in ul
+            let brtag = ul_node.querySelector("br");
+            if (brtag) {
+                ul_node.removeChild(brtag);
+            }
+
+            //just reset and remove nodes, do nothing else
+            return;
+
+        }
+
+    } else {
+        // full preview mode
+
+        if (opts["ch_always_display"]) {
+            additional_node.style.display = "block";
+        } else {
+            additional_node.style.display = null;
+        }
+
+        if (!ul_node) {
+            ul_node = document.createElement("ul");
+        } else {
+            // remove br tag
+            let brtag = ul_node.querySelector("br");
+            if (brtag) {
+                ul_node.removeChild(brtag);
+            }
+        }
+
+    }
+
+    if (ul_node.dataset.ch_helper) {
+        return;
+    }
+
+    ul_node.dataset.ch_helper = true;
+
+    model_name = card.dataset.name;
+
+    // replace preview text button
+    replace_preview_btn = card.querySelector(".actions .additional a");
+
+    if ((replace_preview_btn == null) && !("replace_preview_button" in opts["ch_hide_buttons"])) {
+        /*
+        * in sdwebui 1.5, the replace preview button has been
+        * moved to a hard to reach location, so we have to do
+        * quite a lot to get to its functionality.
+        */
+
+        // waste memory by keeping all of this in scope, per card.
+        let page = active_tab_type;
+        let type = js_model_type;
+        let name = card.dataset.name.replace("'", "\\'");
+
+        // create the replace_preview_btn, as it no longer exists
+        replace_preview_btn = document.createElement("a");
+
+        // create an event handler to redirect a click to the real replace_preview_button
+        replace_preview_btn.setAttribute("onclick", `replace_preview(event, '${page}', '${type}', '${model_name}')`);
+    }
+
+    // change replace preview text button into icon
+    if (!opts["ch_hide_buttons"].includes("replace_preview_button")) {
+        if (replace_preview_btn.textContent !== "🖼️") {
+            replace_preview_btn.textContent = "🖼️";
+            addedNodes.push(replace_preview_btn);
+        }
+
+        replace_preview_btn.classList.add("card-button", "removecard");
+
+    } else if (replace_preview_btn.parentElement) {
+        replace_preview_btn.parentElement.removeChild(replace_preview_btn);
+    }
+
+    // search_term node
+    // search_term: /[subfolder path]/[model name].[ext] [hash]
+    // get search_term
+    let search_term_nodes = card.querySelectorAll(".actions .additional .search_term, .actions .additional .search_terms");
+    if (!search_term_nodes) {
+        console.log("can not find search_term node for cards in " + active_tab_type + "_" + active_extra_tab_type + "_cards");
+        return;
+    }
+
+    if (search_term_nodes.length > 1) {
+        let search_terms = [];
+        for (let search_term_node of search_term_nodes) {
+            search_terms.push(search_term_node.textContent);
+        }
+
+        let model_path = search_terms.join(" ");
+        let separator = model_path.match(/[\/\\]/)[0];
+        model_path = model_path.split(separator).slice(1).join(separator);
+
+        search_term = model_path;
+    } else {
+        let search_term_node = search_term_nodes[0];
+        search_term = search_term_node.textContent;
+
+        // for whatever reason, sometimes search_terms will not include hashes.
+        if (search_term_node.classList.contains("search_terms")) {
+            let separator = search_term.match(/[\/\\]/)[0];
+            search_term = search_term.split(separator).slice(1).join(
+                separator === "\\" ? "\\\\" : "/"
+            );
+        }
+    }
+
+    search_term = search_term.replaceAll("\\", "\\\\").replace("'", "\\'");
+
+    if (!search_term) {
+        console.log("search_term is empty for cards in " + active_tab_type + "_" + active_extra_tab_type + "_cards");
+        return;
+    }
+
+    // then we need to add buttons to each ul node:
+    if (!opts["ch_hide_buttons"].includes("open_url_button")) {
+        let open_url_node = document.createElement("a");
+        open_url_node.href = "#";
+        open_url_node.textContent = "🌐";
+        open_url_node.classList.add("card-button", "openurl");
+        open_url_node.title = "Open this model's civitai url";
+        open_url_node.setAttribute("onclick", `open_model_url(event, '${model_type}', '${search_term}')`);
+        addedNodes.push(open_url_node);
+    }
+
+    // add br if metadata_button exists
+    if (is_thumb_mode && metadata_button) {
+        addedNodes.push(document.createElement("br"));
+    }
+
+    if (!opts["ch_hide_buttons"].includes("add_trigger_words_button")) {
+        let add_trigger_words_node = document.createElement("a");
+        add_trigger_words_node.href = "#";
+        add_trigger_words_node.textContent = "💡";
+        add_trigger_words_node.classList.add("card-button", "addtriggerwords");
+        add_trigger_words_node.title = "Add trigger words to prompt";
+        add_trigger_words_node.setAttribute("onclick", `add_trigger_words(event, '${model_type}', '${search_term}')`);
+        addedNodes.push(add_trigger_words_node);
+    }
+
+    if (!opts["ch_hide_buttons"].includes("add_preview_prompt_button")) {
+        let use_preview_prompt_node = document.createElement("a");
+        use_preview_prompt_node.href = "#";
+        use_preview_prompt_node.textContent = "🏷️";
+        use_preview_prompt_node.classList.add("card-button", "usepreviewprompt");
+        use_preview_prompt_node.title = "Use prompt from preview image";
+        use_preview_prompt_node.setAttribute("onclick", `use_preview_prompt(event, '${model_type}', '${search_term}')`);
+        addedNodes.push(use_preview_prompt_node);
+    }
+
+    if (!opts["ch_hide_buttons"].includes("rename_model_button")) {
+        let rename_card_node = document.createElement("a");
+        rename_card_node.href = "#";
+        rename_card_node.innerHTML = "✏️";
+        rename_card_node.classList.add("card-button", "renamecard");
+        rename_card_node.title = "Rename this model";
+        rename_card_node.setAttribute("onclick", `rename_card(event, '${model_type}', '${search_term}', '${model_name}')`);
+        addedNodes.push(rename_card_node);
+    }
+
+    if (!opts["ch_hide_buttons"].includes("remove_model_button")) {
+        let remove_card_node = document.createElement("a");
+        remove_card_node.href = "#";
+        remove_card_node.innerHTML = "❌";
+        remove_card_node.classList.add("card-button", "removecard");
+        remove_card_node.title = "Remove this model";
+        remove_card_node.setAttribute("onclick", `remove_card(event, '${model_type}', '${search_term}')`);
+        addedNodes.push(remove_card_node);
+    }
+
+    // add to buttons row
+    for (const node of addedNodes) {
+        ul_node.appendChild(node);
+    }
+
+    // add buttons to card
+    if (!ul_node.parentElement && ul_node.children) {
+        additional_node.appendChild(ul_node);
+    }
+}
+
 onUiLoaded(() => {
 
     //get gradio version
-    let gradio_ver = ch_gradio_version();
-    console.log("gradio_ver:" + gradio_ver);
+    const gradio_ver = ch_gradio_version();
+    console.log("Running Stable-Diffusion-Webui-Civitai-Helper on Gradio Version: " + gradio_ver);
 
     // get all extra network tabs
-    let tab_prefix_list = ["txt2img", "img2img"];
-    let model_type_list = ["textual_inversion", "hypernetworks", "checkpoints", "lora"];
-    let cardid_suffix = "cards";
-
-    //get init py msg
-    // let init_py_msg_str = get_ch_py_msg();
-    // let extension_path = "";
-    // if (!init_py_msg_str) {
-    //     console.log("Can not get init_py_msg");
-    // } else {
-    //     init_py_msg = JSON.parse(init_py_msg_str);
-    //     if (init_py_msg) {
-    //         extension_path = init_py_msg.extension_path;
-    //         console.log("get extension path: " + extension_path);
-    //     }
-    // }
-
-    // //icon image node as string
-    // function icon(icon_name){
-    //     let icon_path = extension_path+"/icon/"+icon_name;
-    //     return ch_img_node_str(icon_path);
-    // }
-
+    const tab_prefix_list = ["txt2img", "img2img"];
+    const model_type_list = ["textual_inversion", "hypernetworks", "checkpoints", "lora", "lycoris"];
 
     // update extra network tab pages' cards
     // * replace "replace preview" text button into an icon
@@ -459,87 +980,39 @@ onUiLoaded(() => {
     // notice: javascript can not get response from python side
     // so, these buttons just sent request to python
     // then, python side gonna open url and update prompt text box, without telling js side.
-    function update_card_for_civitai(){
-
-        //css
-        let btn_margin = "0px 5px";
-        let btn_fontSize = "200%";
-        let btn_thumb_fontSize = "100%";
-        let btn_thumb_display = "inline";
-        let btn_thumb_pos = "static";
-        let btn_thumb_backgroundImage = "none";
-        let btn_thumb_background = "rgba(0, 0, 0, 0.8)";
-
-        let ch_btn_txts = ['🌐', '💡', '🏷️'];
+    function update_card_for_civitai() {
+        if (!(opts && "ch_always_display" in opts)) {
+            // Lobe theme can cause a race condition.
+            return setTimeout(update_card_for_civitai, 500);
+        }
+
         let replace_preview_text = getTranslation("replace preview");
         if (!replace_preview_text) {
             replace_preview_text = "replace preview";
         }
 
-
-        //change all "replace preview" into an icon
-        let extra_network_id = "";
         let extra_network_node = null;
-        let button_row = null;
-        let search_term_node = null;
-        let search_term = "";
         let model_type = "";
         let cards = null;
-        let need_to_add_buttons = false;
 
         //get current tab
         let active_tab_type = getActiveTabType();
-        if (!active_tab_type){active_tab_type = "txt2img";}
+        if (!active_tab_type) {active_tab_type = "txt2img";}
 
         for (const tab_prefix of tab_prefix_list) {
             if (tab_prefix != active_tab_type) {continue;}
 
-
-            //find out current selected model type tab
-            let active_extra_tab_type = "";
-            let extra_tabs = gradioApp().getElementById(tab_prefix+"_extra_tabs");
-            if (!extra_tabs) {console.log("can not find extra_tabs: " + tab_prefix+"_extra_tabs");}
-
             //get active extratab
+            const re = new RegExp(tab_prefix + "_(.+)_cards");
             const active_extra_tab = Array.from(get_uiCurrentTabContent().querySelectorAll('.extra-network-cards,.extra-network-thumbs'))
                 .find(el => el.closest('.tabitem').style.display === 'block')
-                ?.id.match(/^(txt2img|img2img)_(.+)_cards$/)[2]
-
-                
-            console.log("found active tab: " + active_extra_tab);
-
-            switch (active_extra_tab) {
-                case "textual_inversion":
-                    active_extra_tab_type = "ti";
-                    break;
-                case "hypernetworks":
-                    active_extra_tab_type = "hyper";
-                    break;
-                case "checkpoints":
-                    active_extra_tab_type = "ckp";
-                    break;
-                case "lora":
-                    active_extra_tab_type = "lora";
-                    break;
-            }
+                ?.id.match(re)[1];
 
+            const active_extra_tab_type = getShortModelTypeFromFull(active_extra_tab);
 
             for (const js_model_type of model_type_list) {
                 //get model_type for python side
-                switch (js_model_type) {
-                    case "textual_inversion":
-                        model_type = "ti";
-                        break;
-                    case "hypernetworks":
-                        model_type = "hyper";
-                        break;
-                    case "checkpoints":
-                        model_type = "ckp";
-                        break;
-                    case "lora":
-                        model_type = "lora";
-                        break;
-                }
+                model_type = getShortModelTypeFromFull(js_model_type);
 
                 if (!model_type) {
                     console.log("can not get model_type from: " + js_model_type);
@@ -552,142 +1025,75 @@ onUiLoaded(() => {
                     continue;
                 }
 
-                console.log("handle active extra tab");
-
+                //console.log("handle active extra tab");
+                extra_network_id = tab_prefix + "_" + js_model_type + "_cards";
 
-                extra_network_id = tab_prefix+"_"+js_model_type+"_"+cardid_suffix;
                 // console.log("searching extra_network_node: " + extra_network_id);
-                extra_network_node = gradioApp().getElementById(extra_network_id);
-
-                // console.log("find extra_network_node: " + extra_network_id);
+                extra_network_node = getModelCardsEl(tab_prefix, js_model_type);
 
                 // get all card nodes
                 cards = extra_network_node.querySelectorAll(".card");
-                for (let card of cards) {
-                    //get button row
-                    button_row = card.querySelector(".button-row");
-                    
-                    if (!button_row){
-                        console.log("can not find button_row");
-                        continue;
-                    }
-
-                    let atags = button_row.querySelectorAll("a");
-                    if (atags && atags.length) {
-                        console.log("find atags: " + atags.length);
-                    } else {
-                        console.log("no atags");
-                        need_to_add_buttons = true;
-                    }
-
-                    if (!need_to_add_buttons) {
-                        console.log("do not need to add buttons");
-                        continue;
-                    }
-
-
-                    // search_term node
-                    // search_term = subfolder path + model name + ext
-                    search_term_node = card.querySelector(".actions .additional .search_term");
-                    if (!search_term_node){
-                        console.log("can not find search_term node for cards in " + extra_network_id);
-                        continue;
-                    }
-
-                    // get search_term
-                    search_term = search_term_node.innerHTML.trim();
-                    if (!search_term) {
-                        console.log("search_term is empty for cards in " + extra_network_id);
-                        continue;
+                for (const card of cards) {
+                    // don't let an issue with a single card kill functionality for following cards
+                    try {
+                        processSingleCard(active_tab_type, active_extra_tab_type, card);
+                    } catch(err) {
+                        console.log(err);
                     }
-
-                    console.log("adding buttons");
-                    // then we need to add 3 buttons to each ul node:
-                    let open_url_node = document.createElement("a");
-                    open_url_node.href = "#";
-                    open_url_node.innerHTML = "🌐";
-                    open_url_node.className = "card-button";
-
-                    open_url_node.title = "Open this model's civitai url";
-                    open_url_node.setAttribute("onclick","open_model_url(event, '"+model_type+"', '"+search_term+"')");
-
-                    let add_trigger_words_node = document.createElement("a");
-                    add_trigger_words_node.href = "#";
-                    add_trigger_words_node.innerHTML = "💡";
-                    add_trigger_words_node.className = "card-button";
-
-                    add_trigger_words_node.title = "Add trigger words to prompt";
-                    add_trigger_words_node.setAttribute("onclick","add_trigger_words(event, '"+model_type+"', '"+search_term+"')");
-
-                    let use_preview_prompt_node = document.createElement("a");
-                    use_preview_prompt_node.href = "#";
-                    use_preview_prompt_node.innerHTML = "🏷️";
-                    use_preview_prompt_node.className = "card-button";
-
-                    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);
-
-
-
-
                 }
 
-                
             }
-        }
 
+        }
 
     }
 
-
-    let tab_id = ""
-    let extra_tab = null;
-    let extra_toolbar = null;
+    /*
     let extra_network_refresh_btn = null;
+    */
+    let extra_networks_btn = null;
+
     //add refresh button to extra network's toolbar
-    for (let prefix of tab_prefix_list) {
-        tab_id = prefix + "_extra_tabs";
-        extra_tab = gradioApp().getElementById(tab_id);
+    for (const prefix of tab_prefix_list) {
+        // load extra networks button
+        extra_networks_btn = gradioApp().getElementById(prefix + "_extra_networks");
 
-        //get toolbar
-        //get Refresh button
-        extra_network_refresh_btn = gradioApp().getElementById(prefix+"_extra_refresh");
 
+        // pre-1.6
+        if (extra_networks_btn) {
+            function extraNetworksClick() {
+                waitForExtraTabs(prefix, model_type_list);
+                extra_networks_btn.removeEventListener("click", extraNetworksClick);
+            }
 
-        if (!extra_network_refresh_btn){
-            console.log("can not get extra network refresh button for " + tab_id);
+            // add listener to extra_networks_btn
+            extra_networks_btn.addEventListener("click", extraNetworksClick);
             continue;
+
         }
 
-        // add refresh button to toolbar
-        let ch_refresh = document.createElement("button");
-        ch_refresh.innerHTML = "🔁";
-        ch_refresh.title = "Refresh Civitai Helper's additional buttons";
-        ch_refresh.className = "lg secondary gradio-button";
-        ch_refresh.style.fontSize = "200%";
-        ch_refresh.onclick = update_card_for_civitai;
+        // 1.6 and higher
+        const extra_tab = getExtraTabs(prefix);
+        const headers = extra_tab.firstChild.children;
 
-        extra_network_refresh_btn.parentNode.appendChild(ch_refresh);
+        for (const header of headers) {
+            const model_type = header.textContent.trim().replace(" ", "_").toLowerCase();
 
-    }
+            function extraNetworksClick() {
+                waitForExtraTabs(prefix, [model_type]);
+                header.removeEventListener("click", extraNetworksClick);
+            }
 
+            header.addEventListener("click", extraNetworksClick);
+        }
+
+        //get toolbar
+        extra_networks_btn = gradioApp().getElementById(prefix + "_extra_networks");
+
+    }
 
     //run it once
-    update_card_for_civitai();
+    // update_card_for_civitai();
 
 
 });

extensions/Stable-Diffusion-Webui-Civitai-Helper/scripts/ch_lib/civitai.py

@@ -1,617 +1,726 @@
-# -*- coding: UTF-8 -*-
-# handle msg between js and python side
+""" -*- coding: UTF-8 -*-
+handle msg between js and python side
+"""
+
 import os
-import time
-import json
 import re
-import requests
 from . import util
 from . import model
-from . import setting
+from . import downloader
 
-suffix = ".civitai"
+SUFFIX = ".civitai"
 
-url_dict = {
+URLS = {
     "modelPage":"https://civitai.com/models/",
     "modelId": "https://civitai.com/api/v1/models/",
     "modelVersionId": "https://civitai.com/api/v1/model-versions/",
     "hash": "https://civitai.com/api/v1/model-versions/by-hash/"
 }
 
-model_type_dict = {
+MODEL_TYPES = {
     "Checkpoint": "ckp",
     "TextualInversion": "ti",
     "Hypernetwork": "hyper",
     "LORA": "lora",
-    "LoCon": "lora",
+    "LoCon": "lycoris",
 }
 
+FILE_TYPES = [
+    "Model", "Training Data", "Config", "VAE"
+]
+
+NSFW_LEVELS = ["None", "Soft", "Mature", "X", "Allow All"]
+
+def civitai_get(civitai_url:str):
+    """
+    Gets JSON from Civitai.
+    return: dict:json or None
+    """
+
+    success, response = downloader.request_get(
+        civitai_url
+    )
+
+    if not success:
+        return None
+
+    # try to get content
+    content = None
+    try:
+        content = response.json()
+    except ValueError as e:
+        util.printD(util.indented_msg(
+            f"""
+            Parse response json failed
+            Error: {str(e)}
+            Response: {response.text}
+            """
+        ))
+        return None
+
+    return content
+
+
+def append_parent_model_metadata(content):
+    """
+    Some model metadata is stored in a "parent" context.
+    When we're fething a model by its hash, we're getting
+    the metadata for that model *file*, not the model entry
+    on Civitai, which may contain multiple versions.
 
+    This method gets the parent metadata and appends it to
+    our model file metadata.
 
-# get image with full size
-# width is in number, not string
-# return: url str
-def get_full_size_image_url(image_url, width):
-    return re.sub('/width=\d+/', '/width=' + str(width) + '/', image_url)
+    return: model metadata with parent description, creator,
+    and permissions appended.
+    """
+    util.printD("Fetching Parent Model Information")
+    parent_model = get_model_info_by_id(content["modelId"])
+
+    metadatas = [
+        "description", "tags", "allowNoCredit",
+        "allowCommercialUse", "allowDerivatives",
+        "allowDifferentLicense"
+    ]
+
+    content["creator"] = parent_model.get("creator", "{}")
+
+    model_metadata = content["model"]
+    for metadata in metadatas:
+        model_metadata[metadata] = parent_model.get(metadata, "")
+
+    return content
 
 
 # use this sha256 to get model info from civitai
 # return: model info dict
-def get_model_info_by_hash(hash:str):
+def get_model_info_by_hash(model_hash:str):
+    """
+    use this sha256 to get model info from civitai's api
+
+    return:
+        model info dict if a model is found
+        {} if civitai does not have the model
+        None if an error occurs.
+    """
     util.printD("Request model info from civitai")
 
-    if not hash:
+    if not model_hash:
         util.printD("hash is empty")
-        return
+        return None
 
-    r = requests.get(url_dict["hash"]+hash, headers=util.def_headers, proxies=util.proxies)
-    if not r.ok:
-        if r.status_code == 404:
-            # this is not a civitai model
-            util.printD("Civitai does not have this model")
-            return {}
-        else:
-            util.printD("Get error code: " + str(r.status_code))
-            util.printD(r.text)
-            return
+    content = civitai_get(f'{URLS["hash"]}{model_hash}')
 
-    # try to get content
-    content = None
-    try:
-        content = r.json()
-    except Exception as e:
-        util.printD("Parse response json failed")
-        util.printD(str(e))
-        util.printD("response:")
-        util.printD(r.text)
-        return
-    
     if not content:
-        util.printD("error, content from civitai is None")
-        return
-    
+        return None
+
+    content = append_parent_model_metadata(content)
+
     return content
 
 
+def get_model_info_by_id(model_id:str) -> dict:
+    """
+    Fetches model info by its model id.
+    returns: dict:model_info
+    """
 
-def get_model_info_by_id(id:str) -> dict:
-    util.printD("Request model info from civitai")
+    util.printD(f"Request model info from civitai: {model_id}")
 
-    if not id:
-        util.printD("id is empty")
-        return
+    if not model_id:
+        util.printD("model_id is empty")
+        return False
 
-    r = requests.get(url_dict["modelId"]+str(id), headers=util.def_headers, proxies=util.proxies)
-    if not r.ok:
-        if r.status_code == 404:
-            # this is not a civitai model
-            util.printD("Civitai does not have this model")
-            return {}
-        else:
-            util.printD("Get error code: " + str(r.status_code))
-            util.printD(r.text)
-            return
+    content = civitai_get(f'{URLS["modelId"]}{model_id}')
 
-    # try to get content
-    content = None
-    try:
-        content = r.json()
-    except Exception as e:
-        util.printD("Parse response json failed")
-        util.printD(str(e))
-        util.printD("response:")
-        util.printD(r.text)
-        return
-    
-    if not content:
-        util.printD("error, content from civitai is None")
-        return
-    
     return content
 
 
-def get_version_info_by_version_id(id:str) -> dict:
+def get_version_info_by_version_id(version_id:str) -> dict:
+    """
+    Gets model version info from Civitai by version id
+    return: dict:model_info
+    """
     util.printD("Request version info from civitai")
 
-    if not id:
-        util.printD("id is empty")
-        return
+    if not version_id:
+        util.printD("version_id is empty")
+        return None
 
-    r = requests.get(url_dict["modelVersionId"]+str(id), headers=util.def_headers, proxies=util.proxies)
-    if not r.ok:
-        if r.status_code == 404:
-            # this is not a civitai model
-            util.printD("Civitai does not have this model version")
-            return {}
-        else:
-            util.printD("Get error code: " + str(r.status_code))
-            util.printD(r.text)
-            return
+    content = civitai_get(f'{URLS["modelVersionId"]}{version_id}')
+
+    if content:
+        content = append_parent_model_metadata(content)
 
-    # try to get content
-    content = None
-    try:
-        content = r.json()
-    except Exception as e:
-        util.printD("Parse response json failed")
-        util.printD(str(e))
-        util.printD("response:")
-        util.printD(r.text)
-        return
-    
-    if not content:
-        util.printD("error, content from civitai is None")
-        return
-    
     return content
 
 
-def get_version_info_by_model_id(id:str) -> dict:
+def get_version_info_by_model_id(model_id:str) -> dict:
+    """
+    Fetches version info by model id.
+    returns: dict:version_info
+    """
 
-    model_info = get_model_info_by_id(id)
+    model_info = get_model_info_by_id(model_id)
     if not model_info:
-        util.printD(f"Failed to get model info by id: {id}")
-        return
-    
+        util.printD(f"Failed to get model info by id: {model_id}")
+        return None
+
     # check content to get version id
-    if "modelVersions" not in model_info.keys():
-        util.printD("There is no modelVersions in this model_info")
-        return
-    
-    if not model_info["modelVersions"]:
-        util.printD("modelVersions is None")
-        return
-    
-    if len(model_info["modelVersions"])==0:
-        util.printD("modelVersions is Empty")
-        return
-    
-    def_version = model_info["modelVersions"][0]
+    versions = model_info.get("modelVersions", [])
+    if len(versions) == 0:
+        util.printD("Found no model versions")
+        return None
+
+    def_version = versions[0]
     if not def_version:
         util.printD("default version is None")
-        return
-    
-    if "id" not in def_version.keys():
-        util.printD("default version has no id")
-        return
-    
-    version_id = def_version["id"]
-    
+        return None
+
+    version_id = def_version.get("id", "")
+
     if not version_id:
-        util.printD("default version's id is None")
-        return
+        util.printD("Could not get valid version id")
+        return None
 
     # get version info
-    version_info = get_version_info_by_version_id(str(version_id))
+    version_info = get_version_info_by_version_id(f"{version_id}")
     if not version_info:
         util.printD(f"Failed to get version info by version_id: {version_id}")
-        return
+        return None
 
     return version_info
 
 
-
-
-# get model info file's content by model type and search_term
-# parameter: model_type, search_term
-# return: model_info
 def load_model_info_by_search_term(model_type, search_term):
+    """
+    get model info file's content by model type and search_term
+    parameter: model_type, search_term
+    return: model_info
+    """
     util.printD(f"Load model info of {search_term} in {model_type}")
-    if model_type not in model.folders.keys():
-        util.printD("unknow model type: " + model_type)
-        return
-    
-    # search_term = subfolderpath + model name + ext. And it always start with a / even there is no sub folder
-    base, ext = os.path.splitext(search_term)
+    if model.folders.get(model_type, None) is None:
+        util.printD(f"unknown model type: {model_type}")
+        return None
+
+    # search_term = f"{subfolderpath}{model name}{ext}"
+    # And it always start with a / even when there is no sub folder
+    base, _ = os.path.splitext(search_term)
     model_info_base = base
     if base[:1] == "/":
         model_info_base = base[1:]
 
-    model_folder = model.folders[model_type]
-    model_info_filename = model_info_base + suffix + model.info_ext
-    model_info_filepath = os.path.join(model_folder, model_info_filename)
+    if model_type == "lora" and model.folders['lycoris']:
+        model_folders = [model.folders[model_type], model.folders['lycoris']]
+    else:
+        model_folders = [model.folders[model_type]]
 
-    if not os.path.isfile(model_info_filepath):
-        util.printD("Can not find model info file: " + model_info_filepath)
-        return
-    
-    return model.load_model_info(model_info_filepath)
+    for model_folder in model_folders:
+        model_info_filename = f"{model_info_base}{SUFFIX}{model.CIVITAI_EXT}"
+        model_info_filepath = os.path.join(model_folder, model_info_filename)
+
+        found = os.path.isfile(model_info_filepath)
+
+        if found:
+            break
 
+    if not found:
+        util.printD(f"Can not find model info file: {model_info_filepath}")
+        return None
 
+    return model.load_model_info(model_info_filepath)
 
 
+def get_model_names_by_type_and_filter(model_type:str, metadata_filter:dict) -> list:
+    """
+    get model file names by model type
+    parameter: model_type - string
+    parameter: filter - dict, which kind of model you need
+    return: model name list
+    """
 
-# get model file names by model type
-# parameter: model_type - string
-# parameter: filter - dict, which kind of model you need
-# return: model name list
-def get_model_names_by_type_and_filter(model_type:str, filter:dict) -> list:
-    
-    model_folder = model.folders[model_type]
+    if model_type == "lora" and model.folders['lycoris']:
+        model_folders = [model.folders[model_type], model.folders['lycoris']]
+    else:
+        model_folders = [model.folders[model_type]]
 
-    # set filter
+    # set metadata_filter
     # only get models don't have a civitai info file
     no_info_only = False
     empty_info_only = False
 
-    if filter:
-        if "no_info_only" in filter.keys():
-            no_info_only = filter["no_info_only"]
-        if "empty_info_only" in filter.keys():
-            empty_info_only = filter["empty_info_only"]
-
-
+    if metadata_filter:
+        no_info_only = metadata_filter.get("no_info_only", False)
+        empty_info_only = metadata_filter.get("empty_info_only", False)
 
     # get information from filter
     # only get those model names don't have a civitai model info file
     model_names = []
-    for root, dirs, files in os.walk(model_folder, followlinks=True):
-        for filename in files:
-            item = os.path.join(root, filename)
-            # check extension
-            base, ext = os.path.splitext(item)
-            if ext in model.exts:
-                # find a model
-
-                # check filter
-                if no_info_only:
-                    # check model info file
-                    info_file = base + suffix + model.info_ext
-                    if os.path.isfile(info_file):
-                        continue
-
-                if empty_info_only:
-                    # check model info file
-                    info_file = base + suffix + model.info_ext
-                    if os.path.isfile(info_file):
-                        # load model info
-                        model_info = model.load_model_info(info_file)
-                        # check content
-                        if model_info:
-                            if "id" in model_info.keys():
-                                # find a non-empty model info file
-                                continue
-
-                model_names.append(filename)
-
+    for model_folder in model_folders:
+        for root, _, files in os.walk(model_folder, followlinks=True):
+            for filename in files:
+                if is_valid_file(root, filename, no_info_only, empty_info_only):
+                    model_names.append(filename)
 
     return model_names
 
+
+def is_valid_file(root, filename, no_info_only, empty_info_only):
+    """
+    Filters through model files to determine if they are
+    valid targets for downloading new metadata.
+
+    return: bool
+    """
+    item = os.path.join(root, filename)
+    # check extension
+    base, ext = os.path.splitext(item)
+    if ext not in model.EXTS:
+        return False
+
+    # find a model
+    info_file = f"{base}{SUFFIX}{model.CIVITAI_EXT}"
+
+    # check filter
+    if os.path.isfile(info_file):
+        if no_info_only:
+            return False
+
+        if empty_info_only:
+            # load model info
+            model_info = model.load_model_info(info_file)
+            # check content
+            if model_info and not model_info.get("id", "") == "":
+                # find a non-empty model info file
+                return False
+
+    return True
+
+
 def get_model_names_by_input(model_type, empty_info_only):
+    """ return: list of model filenames with empty civitai info files """
     return get_model_names_by_type_and_filter(model_type, {"empty_info_only":empty_info_only})
-    
+
 
 # get id from url
 def get_model_id_from_url(url:str) -> str:
+    """ return: model_id from civitai url """
     util.printD("Run get_model_id_from_url")
-    id = ""
+    model_id = ""
 
     if not url:
         util.printD("url or model id can not be empty")
         return ""
 
     if url.isnumeric():
-        # is already an id
-        id = str(url)
-        return id
-    
-    s = re.sub("\\?.+$", "", url).split("/")
-    if len(s) < 2:
+        # is already an model_id
+        model_id = f"{url}"
+        return model_id
+
+    split_url = re.sub("\\?.+$", "", url).split("/")
+    if len(split_url) < 2:
         util.printD("url is not valid")
         return ""
-    
-    if s[-2].isnumeric():
-        id  = s[-2]
-    elif s[-1].isnumeric():
-        id  = s[-1]
+
+    if split_url[-2].isnumeric():
+        model_id  = split_url[-2]
+    elif split_url[-1].isnumeric():
+        model_id  = split_url[-1]
     else:
         util.printD("There is no model id in this url")
         return ""
-    
-    return id
+
+    return model_id
+
+
+def preview_exists(model_path):
+    """ Search for existing preview image. return True if it exists, else false """
+
+    previews = model.get_potential_model_preview_files(model_path)
+
+    for prev in previews:
+        if os.path.isfile(prev):
+            return True
+
+    return False
+
+
+def should_skip(user_rating, image_rating):
+    """ return: True if preview_nsfw level higher than user threshold """
+    order = NSFW_LEVELS
+    if user_rating == "Skip":
+        # Old config
+        return False
+    if isinstance(image_rating, bool):
+        # Image using old NSFW system?
+        if image_rating:
+            image_rating = order[-1]
+        else:
+            image_rating = order[0]
+    return order.index(image_rating) >= order.index(user_rating)
+
+
+def show_only_nsfw(user_rating, image_rating):
+    """ return: True if image is NSFW """
+    order = NSFW_LEVELS
+    if user_rating == "Skip":
+        # Old config
+        return True
+    return order.index(image_rating) <= order.index(user_rating)
+
+
+def get_image_url(img_dict, max_size_preview):
+    """
+    Create the image download URL
+
+    return: url:str
+    """
+
+    url = img_dict["url"]
+    if max_size_preview:
+        # use max width
+        width = img_dict.get("width", False)
+        if width:
+            url = re.sub(r'/width=\d+/', '/width=' + str(width) + '/', url)
+
+    return url
+
+
+def verify_preview(path, img_dict, max_size_preview, nsfw_preview_threshold):
+    """
+    Downloads a preview image if it meets the user's requirements.
+    """
+
+    img_url = img_dict.get("url", None)
+    if img_url is None:
+        yield (False, None)
+
+    image_rating = img_dict.get("nsfw", "None")
+    if image_rating != "None":
+        util.printD(f"This image is NSFW: {image_rating}")
+        if should_skip(nsfw_preview_threshold, image_rating):
+            util.printD("Skip NSFW image")
+            yield (False, None)
+
+    preview_type = img_dict.get("type")
+    if preview_type != "image":
+        util.printD(f"Preview is not an image. Found {preview_type} instead. Skipping.")
+        yield (False, None)
+
+    img_url = get_image_url(img_dict, max_size_preview)
+
+    success = False
+    preview_path = ""
+    for result in downloader.dl_file(img_url, file_path=path):
+        if not isinstance(result, str):
+            success, preview_path = result
+            break
+
+        yield result
+
+
+    if not success:
+        yield (False, None)
+
+    # we only need 1 preview image
+    yield (True, preview_path)
 
 
 # get preview image by model path
 # image will be saved to file, so no return
-def get_preview_image_by_model_path(model_path:str, max_size_preview, skip_nsfw_preview):
+def get_preview_image_by_model_path(model_path:str, max_size_preview, nsfw_preview_threshold, preferred_preview=None):
+    """
+    Downloads a preview image for a model if one doesn't already exist.
+    Skips images that are more NSFW than the user's NSFW threshold
+    """
+    util.printD("Downloading model image.")
+
     if not model_path:
         util.printD("model_path is empty")
         return
 
     if not os.path.isfile(model_path):
-        util.printD("model_path is not a file: "+model_path)
+        util.printD(f"model_path is not a file: {model_path}")
         return
 
-    base, ext = os.path.splitext(model_path)
-    first_preview = base+".png"
-    sec_preview = base+".preview.png"
-    info_file = base + suffix + model.info_ext
-
-    # check preview image
-    if not os.path.isfile(sec_preview):
-        # need to download preview image
-        util.printD("Checking preview image for model: " + model_path)
-        # load model_info file
-        if os.path.isfile(info_file):
-            model_info = model.load_model_info(info_file)
-            if not model_info:
-                util.printD("Model Info is empty")
+    base, _ = os.path.splitext(model_path)
+    preview_path =  f"{base}.preview.png" # TODO png not strictly required
+    info_file = f"{base}{SUFFIX}{model.CIVITAI_EXT}"
+
+    # need to download preview image
+    util.printD(f"Checking preview image for model: {model_path}")
+
+    if preview_exists(model_path):
+        output = "Existing model image found. Skipping."
+        util.printD(output)
+        yield output
+        return
+
+    # load model_info file
+    if not os.path.isfile(info_file):
+        return
+
+    try:
+        images = model.load_model_info(info_file)["images"]
+
+    except (KeyError, TypeError):
+        return
+
+    if preferred_preview:
+        img_url = preferred_preview
+        for img_dict in images:
+            if img_dict["url"] == preferred_preview:
+                img_url = get_image_url(img_dict, max_size_preview)
+                break
+
+        for result in downloader.dl_file(img_url, file_path=preview_path):
+            if isinstance(result, str):
+                yield result
+                continue
+
+            success, msg = result
+
+            if success:
                 return
 
-            if "images" in model_info.keys():
-                if model_info["images"]:
-                    for img_dict in model_info["images"]:
-                        if "nsfw" in img_dict.keys():
-                            if img_dict["nsfw"] and img_dict["nsfw"] != "None":
-                                util.printD("This image is NSFW")
-                                if skip_nsfw_preview:
-                                    util.printD("Skip NSFW image")
-                                    continue
-
-                        preview_type = img_dict.get("type")
-                        if preview_type != "image":
-                            util.printD(f"Unsupported preview type: {preview_type}, ignore.")
-                            continue
-                        
-                        if "url" in img_dict.keys():
-                            img_url = img_dict["url"]
-                            if max_size_preview:
-                                # use max width
-                                if "width" in img_dict.keys():
-                                    if img_dict["width"]:
-                                        img_url = get_full_size_image_url(img_url, img_dict["width"])
-
-                            util.download_file(img_url, sec_preview)
-                            # we only need 1 preview image
-                            break
+            util.printD(msg)
+            util.printD("Failed to download preferred preview. Trying to find another")
 
+            break
+
+    for img_dict in images:
+        for result in verify_preview(
+            preview_path, img_dict, max_size_preview, nsfw_preview_threshold
+        ):
+            if not isinstance(result, str):
+                success, _ = result
+                # Only download one image
+                if success:
+                    return
+
+                break
+
+            yield result
+
+    util.printD(f"Could not find any valid preview images for model: {model_path}")
+    yield
 
 
 # search local model by version id in 1 folder, no subfolder
 # return - model_info
 def search_local_model_info_by_version_id(folder:str, version_id:int) -> dict:
+    """ Searches a folder for model_info files,
+        returns the model_info from a file if its id matches the model id.
+    """
     util.printD("Searching local model by version id")
-    util.printD("folder: " + folder)
-    util.printD("version_id: " + str(version_id))
+    util.printD(f"folder: {folder}")
+    util.printD(f"version_id: {version_id}")
 
     if not folder:
         util.printD("folder is none")
-        return
+        return None
 
     if not os.path.isdir(folder):
         util.printD("folder is not a dir")
-        return
-    
+        return None
+
     if not version_id:
         util.printD("version_id is none")
-        return
-    
+        return None
+
     # search civitai model info file
     for filename in os.listdir(folder):
         # check ext
         base, ext = os.path.splitext(filename)
-        if ext == model.info_ext:
+        if ext == model.CIVITAI_EXT:
             # find info file
-            if len(base) < 9:
+            if not (len(base) > 8 and base[-8:] == SUFFIX):
                 # not a civitai info file
                 continue
 
-            if base[-8:] == suffix:
-                # find a civitai info file
-                path = os.path.join(folder, filename)
-                model_info = model.load_model_info(path)
-                if not model_info:
-                    continue
+            # find a civitai info file
+            path = os.path.join(folder, filename)
+            model_info = model.load_model_info(path)
+            if not model_info:
+                continue
 
-                if "id" not in model_info.keys():
-                    continue
+            model_id = model_info.get("id", None)
+            if not model_id:
+                continue
 
-                id = model_info["id"]
-                if not id:
-                    continue
+            # util.printD(f"Compare version id, src: {model_id}, target:{version_id}")
+            if f"{model_id}" == f"{version_id}":
+                # find the one
+                return model_info
 
-                # util.printD(f"Compare version id, src: {id}, target:{version_id}")
-                if str(id) == str(version_id):
-                    # find the one
-                    return model_info
-                    
+    return None
 
-    return
 
+def get_model_id_from_model_path(model_path:str):
+    """ return model_id using model_path """
+    # get model info file name
+    base, _ = os.path.splitext(model_path)
+    info_file = f"{base}{SUFFIX}{model.CIVITAI_EXT}"
 
+    if not os.path.isfile(info_file):
+        return None
 
+    # get model info
+    model_info_file = model.load_model_info(info_file)
+    local_version_id = model_info_file.get("id", None)
+    model_id = model_info_file.get("modelId", None)
 
+    if None in [model_id, local_version_id]:
+        return None
 
-# check new version for a model by model path
-# return (model_path, model_id, model_name, new_verion_id, new_version_name, description, download_url, img_url)
-def check_model_new_version_by_path(model_path:str, delay:float=1) -> tuple:
-    if not model_path:
-        util.printD("model_path is empty")
-        return
+    return (model_id, local_version_id)
 
-    if not os.path.isfile(model_path):
-        util.printD("model_path is not a file: "+model_path)
-        return
-    
-    # get model info file name
-    base, ext = os.path.splitext(model_path)
-    info_file = base + suffix + model.info_ext
-    
-    if not os.path.isfile(info_file):
-        return
-    
-    # get model info
-    model_info_file = model.load_model_info(info_file)
-    if not model_info_file:
-        return
 
-    if "id" not in model_info_file.keys():
-        return
-    
-    local_version_id = model_info_file["id"]
-    if not local_version_id:
-        return
+def check_model_new_version_by_path(model_path:str, delay:float=0.2) -> tuple:
+    """
+    check new version for a model by model path
+    return (
+        model_path, model_id, model_name, new_verion_id,
+        new_version_name, description, download_url, img_url
+    )
+    """
+
+    if not (model_path and os.path.isfile(model_path)):
+        util.printD(f"model_path is not a file: {model_path}")
+        return None
+
+    result = get_model_id_from_model_path(model_path)
+    if not result:
+        return None
+
+    model_id, local_version_id = result
 
-    if "modelId" not in model_info_file.keys():
-        return
-    
-    model_id = model_info_file["modelId"]
-    if not model_id:
-        return
-    
     # get model info by id from civitai
     model_info = get_model_info_by_id(model_id)
-    # delay before next request, to prevent to be treat as DDoS 
-    util.printD(f"delay:{delay} second")
-    time.sleep(delay)
+
+    util.delay(delay)
 
     if not model_info:
-        return
-    
-    if "modelVersions" not in model_info.keys():
-        return
-    
-    modelVersions = model_info["modelVersions"]
-    if not modelVersions:
-        return
-    
-    if not len(modelVersions):
-        return
-    
-    current_version = modelVersions[0]
+        return None
+
+    model_versions = model_info.get("modelVersions", [])
+
+    if len(model_versions) == 0:
+        return None
+
+    current_version = model_versions[0]
     if not current_version:
-        return
-    
-    if "id" not in current_version.keys():
-        return
-    
-    current_version_id = current_version["id"]
-    if not current_version_id:
-        return
+        return None
 
-    util.printD(f"Compare version id, local: {local_version_id}, remote: {current_version_id} ")
-    if current_version_id == local_version_id:
-        return
+    current_version_id = current_version.get("id", False)
+
+    util.printD(f"Compare version id, local: {local_version_id}, remote: {current_version_id}")
+
+    if not (current_version_id and current_version_id != local_version_id):
+        return None
 
-    model_name = ""
-    if "name" in model_info.keys():
-        model_name = model_info["name"]
-    
-    if not model_name:
-        model_name = ""
-
-
-    new_version_name = ""
-    if "name" in current_version.keys():
-        new_version_name = current_version["name"]
-    
-    if not new_version_name:
-        new_version_name = ""
-
-    description = ""
-    if "description" in current_version.keys():
-        description = current_version["description"]
-    
-    if not description:
-        description = ""
-
-    downloadUrl = ""
-    if "downloadUrl" in current_version.keys():
-        downloadUrl = current_version["downloadUrl"]
-    
-    if not downloadUrl:
-        downloadUrl = ""
+    model_name = model_info.get("name", "")
+    new_version_name = current_version.get("name", "")
+    description = current_version.get("description", "")
+    download_url = current_version.get("downloadUrl", "")
 
     # get 1 preview image
-    img_url = ""
-    if "images" in current_version.keys():
-        if current_version["images"]:
-            if current_version["images"][0]:
-                if "url" in current_version["images"][0].keys():
-                    img_url = current_version["images"][0]["url"]
-                    if not img_url:
-                        img_url = ""
+    try:
+        img_url = current_version["images"][0]["url"]
+    except (IndexError, KeyError):
+        img_url = ""
+
+    return (
+        model_path, model_id, model_name, current_version_id,
+        new_version_name, description, download_url, img_url
+    )
+
 
+def check_single_model_new_version(root, filename, model_type, delay):
+    """
+    return: True if a valid model has a new version.
+    """
+    # check ext
+    item = os.path.join(root, filename)
+    _, ext = os.path.splitext(item)
 
-    
-    return (model_path, model_id, model_name, current_version_id, new_version_name, description, downloadUrl, img_url)
+    if ext not in model.EXTS:
+        return False
 
+    # find a model
+    request = check_model_new_version_by_path(item, delay)
 
+    if not request:
+        return False
 
+    request = request + (model_type,)
 
-# check model's new version
-# parameter: delay - float, how many seconds to delay between each request to civitai
-# return: new_versions - a list for all new versions, each one is (model_path, model_id, model_name, new_verion_id, new_version_name, description, download_url, img_url)
-def check_models_new_version_by_model_types(model_types:list, delay:float=1) -> list:
+    # model_path, model_id, model_name, version_id, new_version_name, description, downloadUrl, img_url = request
+    version_id = request[3]
+
+    # check exist
+    if not version_id:
+        return False
+
+    # search this new version id to check if this model is already downloaded
+    target_model_info = search_local_model_info_by_version_id(root, version_id)
+    if target_model_info:
+        util.printD("New version is already exists")
+        return False
+
+    return request
+
+
+def check_models_new_version_by_model_types(model_types:list, delay:float=0.2) -> list:
+    """
+    check all models of model_types for new version
+    parameter: delay - float, how many seconds to delay between each request to civitai
+    return: new_versions
+        a list for all new versions, each one is
+        (model_path, model_id, model_name, new_verion_id,
+        new_version_name, description, download_url, img_url)
+    """
     util.printD("Checking models' new version")
 
     if not model_types:
         return []
 
-    # check model types, which cloud be a string as 1 type
+    # check model types, which could be a string as 1 type
     mts = []
-    if type(model_types) == str:
+    if isinstance(model_types, str):
         mts.append(model_types)
-    elif type(model_types) == list:
+    elif isinstance(model_types, list):
         mts = model_types
     else:
-        util.printD("Unknow model types:")
+        util.printD("Unknown model types:")
         util.printD(model_types)
         return []
 
-    # output is a markdown document string to show a list of new versions on UI
-    output = ""
     # new version list
     new_versions = []
+    new_version_ids = []
 
     # walk all models
     for model_type, model_folder in model.folders.items():
         if model_type not in mts:
             continue
 
-        util.printD("Scanning path: " + model_folder)
-        for root, dirs, files in os.walk(model_folder, followlinks=True):
+        util.printD(f"Scanning path: {model_folder}")
+        for root, _, files in os.walk(model_folder, followlinks=True):
             for filename in files:
-                # check ext
-                item = os.path.join(root, filename)
-                base, ext = os.path.splitext(item)
-                if ext in model.exts:
-                    # find a model
-                    r = check_model_new_version_by_path(item, delay)
-
-                    if not r:
-                        continue
-
-                    model_path, model_id, model_name, current_version_id, new_version_name, description, downloadUrl, img_url = r
-                    # check exist
-                    if not current_version_id:
-                        continue
-
-                    # check this version id in list
-                    is_already_in_list = False
-                    for new_version in new_versions:
-                        if current_version_id == new_version[3]:
-                            # already in list
-                            is_already_in_list = True
-                            break
-
-                    if is_already_in_list:
-                        util.printD("New version is already in list")
-                        continue
+                version = check_single_model_new_version(root, filename, model_type, delay)
 
-                    # search this new version id to check if this model is already downloaded
-                    target_model_info = search_local_model_info_by_version_id(root, current_version_id)
-                    if target_model_info:
-                        util.printD("New version is already existed")
-                        continue
-
-                    # add to list
-                    new_versions.append(r)
+                if not version:
+                    continue
 
+                # model_path, model_id, model_name, version_id, new_version_name, description, downloadUrl, img_url = version
+                version_id = version[3]
 
+                if version_id in new_version_ids:
+                    continue
 
+                # add to list
+                new_versions.append(version)
+                new_version_ids.append(version_id)
 
     return new_versions
-
-
-

extensions/Stable-Diffusion-Webui-Civitai-Helper/scripts/ch_lib/downloader.py

@@ -1,125 +1,368 @@
-# -*- coding: UTF-8 -*-
-import sys
-import requests
+""" -*- coding: UTF-8 -*-
+Used for downloading files
+"""
+from __future__ import annotations
+from collections.abc import Generator
 import os
+import platform
+import time
+from typing import cast, Literal
+from tqdm import tqdm
+import requests
+import urllib3
 from . import util
 
 
-dl_ext = ".downloading"
+DL_EXT = ".downloading"
+MAX_RETRIES = 3
 
 # disable ssl warning info
-requests.packages.urllib3.disable_warnings()
+urllib3.disable_warnings()
 
-# output is downloaded file path
-def dl(url, folder, filename, filepath):
-    util.printD("Start downloading from: " + url)
-    # get file_path
-    file_path = ""
-    if filepath:
-        file_path = filepath
-    else:
-        # if file_path is not in parameter, then folder must be in parameter
-        if not folder:
-            util.printD("folder is none")
-            return
 
-        if not os.path.isdir(folder):
-            util.printD("folder does not exist: "+folder)
-            return
+def request_get(
+    url:str,
+    headers:dict | None=None,
+    retries=0
+) -> tuple[Literal[True], requests.Response] | tuple[Literal[False], str]:
+    """
+    Performs a GET request
 
-        if filename:
-            file_path = os.path.join(folder, filename)
+    returns: tuple(success:bool, response:Response or failure message:str)
+    """
 
-    # first request for header
-    rh = requests.get(url, stream=True, verify=False, headers=util.def_headers, proxies=util.proxies)
-    # get file size
-    total_size = int(rh.headers.get('Content-Length', -1))
-    if (total_size < 0):
-        util.printD('This model requires an API key to download. More info: https://github.com/butaixianran/Stable-Diffusion-Webui-Civitai-Helper#civitai-api-key')
-        return
-    util.printD(f"File size: {total_size}")
-
-    # if file_path is empty, need to get file name from download url's header
-    if not file_path:
-        filename = ""
-        if "Content-Disposition" in rh.headers.keys():
-            # headers default is decoded with latin1, so need to re-decode it with utf-8
-            cd = rh.headers["Content-Disposition"].encode('latin1').decode('utf-8', errors='ignore')
-            # Extract the filename from the header
-            # content of a CD: "attachment;filename=FileName.txt"
-            # in case "" is in CD filename's start and end, need to strip them out
-            filename = cd.split("=")[1].strip('"')
-            if not filename:
-                util.printD("Fail to get file name from Content-Disposition: " + cd)
-                return
+    headers = util.append_default_headers(headers or {})
 
-        if not filename:
-            util.printD("Can not get file name from download url's header")
-            return
+    try:
+        response = requests.get(
+            url,
+            stream=True,
+            verify=False,
+            headers=headers,
+            proxies=util.PROXIES,
+            timeout=util.REQUEST_TIMEOUT
+        )
+
+    except TimeoutError:
+        output = f"GET Request timed out for {url}"
+        print(output)
+        return (False, output)
+
+    if not response.ok:
+        status_code = response.status_code
+        reason = response.reason
+        util.printD(util.indented_msg(
+            f"""
+            GET Request failed with error code:
+            {status_code}: {reason}
+            """
+        ))
+
+        if status_code == 401:
+            return (
+                False,
+                "This download requires Authentication. Please add an API Key to Civitai Helper's settings to continue this download. See [Wiki](https://github.com/zixaphir/Stable-Diffusion-Webui-Civitai-Helper/wiki/Civitai-API-Key) for details on how to create an API Key."
+            )
+
+        if status_code == 416:
+            response.raise_for_status()
+
+        if status_code != 404 and retries < MAX_RETRIES:
+            util.printD("Retrying")
+            return request_get(url, headers, retries + 1)
+
+        return (False, reason)
+
+    return (True, response)
+
+
+def visualize_progress(percent:int, downloaded:int, total:int, speed:int | float, show_bar=True) -> str:
+    """ Used to display progress in webui """
+
+    s_total = f"{total}"
+    s_downloaded = f"{downloaded:>{len(s_total)}}"
+    s_percent = f"{percent:>3}"
+    s_speed = f'{human_readable_filesize(speed)}Bps'
 
-        # with folder and filename, now we have the full file path
-        file_path = os.path.join(folder, filename)
+    snippet = f"`{s_percent}%: {s_downloaded} / {s_total} @ {s_speed}`"
 
+    if not show_bar:
+        # Unfortunately showing a progress bar in webui
+        # is very weird on mobile with limited horizontal
+        # space
+        return snippet.replace(" ", "\u00a0")
 
-    util.printD("Target file path: " + file_path)
-    base, ext = os.path.splitext(file_path)
+    progress = "\u2588" * percent
 
-    # check if file is already exist
-    count = 2
-    new_base = base
-    while os.path.isfile(file_path):
-        util.printD("Target file already exist.")
-        # re-name
-        new_base = base + "_" + str(count)
-        file_path = new_base + ext
-        count += 1
+    return f"`[{progress:<100}] {snippet}`".replace(" ", "\u00a0")
 
+
+def download_progress(
+    url:str,
+    file_path:str,
+    total_size:int,
+    headers:dict | None=None,
+    response_without_range:requests.Response | None=None
+) -> Generator[tuple[bool, str] | str, None, None]:
+    """
+    Performs a file download.
+
+    yields: tuple(success:bool, filepath or failure message:str) or progress:str
+    """
     # use a temp file for downloading
-    dl_file_path = new_base+dl_ext
 
+    if not headers:
+        headers = {}
+
+    dl_path = f"{file_path}{DL_EXT}"
 
-    util.printD(f"Downloading to temp file: {dl_file_path}")
+    util.printD(f"Downloading to temp file: {dl_path}")
 
-    # check if downloading file is exsited
+    # check if downloading file exists
     downloaded_size = 0
-    if os.path.exists(dl_file_path):
-        downloaded_size = os.path.getsize(dl_file_path)
+    if os.path.exists(dl_path):
+        downloaded_size = os.path.getsize(dl_path)
+        util.printD(f"Resuming partially downloaded file from progress: {downloaded_size}")
+
+    # use response without range or create request with range
+    if response_without_range and downloaded_size == 0:
+        response = response_without_range
+    else:
+        if response_without_range:
+            response_without_range.close()
 
-    util.printD(f"Downloaded size: {downloaded_size}")
+        # create header range
+        headers_with_range = util.append_default_headers({
+            **headers,
+            "Range": f"bytes={downloaded_size:d}-",
+        })
 
-    # create header range
-    headers = {'Range': 'bytes=%d-' % downloaded_size}
-    headers['User-Agent'] = util.def_headers['User-Agent']
-    if util.civitai_api_key:
-        headers["Authorization"] = f"Bearer {util.civitai_api_key}"
+        # download with header
+        try:
+            success, response_or_error = request_get(
+                url,
+                headers=headers_with_range,
+            )
 
-    # download with header
-    r = requests.get(url, stream=True, verify=False, headers=headers, proxies=util.proxies)
+        except requests.HTTPError as dl_error:
+            # 416 - Range Not Satisfiable
+            response = dl_error.response
+            if not response or response.status_code != 416:
+                raise
+
+            util.printD("Could not resume download from existing temporary file. Restarting download")
+
+            os.remove(dl_path)
+
+            yield from download_progress(url, file_path, total_size, headers)
+            return
+
+        if not success:
+            yield (False, cast(str, response_or_error))
+            return
+
+        response = cast(requests.Response, response_or_error)
+
+    last_tick = 0
+    start = time.time()
+
+    downloaded_this_session = 0
 
     # write to file
-    with open(dl_file_path, "ab") as f:
-        for chunk in r.iter_content(chunk_size=1024):
+    with open(dl_path, 'ab') as target, tqdm(
+        initial=target.tell(),
+        total=total_size,
+        unit='iB',
+        unit_scale=True,
+        unit_divisor=1024
+    ) as progress_bar:
+        for chunk in response.iter_content(chunk_size=256*1024):
             if chunk:
+                downloaded_this_session += len(chunk)
                 downloaded_size += len(chunk)
-                f.write(chunk)
-                # force to write to disk
-                f.flush()
+                written = target.write(chunk)
+
+                # write to disk
+                target.flush()
+
+                progress_bar.update(written)
+
+                percent = int(100 * (downloaded_size / total_size))
+                timer = time.time()
+
+                # Gradio output is a *slooowwwwwwww* asynchronous FIFO queue
+                if timer - last_tick > 0.2 or percent == 100:
+
+                    last_tick = timer
+                    elapsed = timer - start
+                    speed = downloaded_this_session // elapsed if elapsed >= 1 \
+                        else downloaded_this_session
 
-                # progress
-                progress = int(50 * downloaded_size / total_size)
-                sys.stdout.reconfigure(encoding='utf-8')
-                sys.stdout.write("\r[%s%s] %d%%" % ('-' * progress, ' ' * (50 - progress), 100 * downloaded_size / total_size))
-                sys.stdout.flush()
+                    text_progress = visualize_progress(
+                        percent,
+                        downloaded_size,
+                        total_size,
+                        speed,
+                        False
+                    )
 
-    print()
+                    yield text_progress
 
     # check file size
-    downloaded_size = os.path.getsize(dl_file_path)
-    if downloaded_size < total_size:
-        util.printD(f"Download failed due to insufficient file size. Try again later or download it manually: {url}")
-        return
+    downloaded_size = os.path.getsize(dl_path)
+    if downloaded_size != total_size:
+        warning = util.indented_msg(
+            f"""
+            File is not the correct size: {file_path}.
+            Expected {total_size:d}, got {downloaded_size:d}.
+            The file may be corrupt. If you encounter issues,
+            you can try again later or download the file manually: {url}
+            """
+        )
+        util.warning(warning)
+        util.printD(warning)
 
     # rename file
-    os.rename(dl_file_path, file_path)
-    util.printD(f"File Downloaded to: {file_path}")
-    return file_path
+    os.rename(dl_path, file_path)
+    output = f"File Downloaded to: {file_path}"
+    util.printD(output)
+
+    yield (True, file_path)
+
+
+def get_file_path_from_service_headers(response:requests.Response, folder:str) -> str | None:
+    """
+    Parses a response header to get a filename
+    then builds a file_path.
+
+    return: file_path:str
+    """
+
+    content_disposition = response.headers.get("Content-Disposition", None)
+
+    if content_disposition is None:
+        util.printD("Can not get file name from download url's header")
+        return None
+
+    # Extract the filename from the header
+    # content of a CD: "attachment;filename=FileName.txt"
+    # in case "" is in CD filename's start and end, need to strip them out
+    filename = content_disposition.split("=")[1].strip('"')
+    filename = filename.encode('iso8859-1').decode('utf-8')
+    if not filename:
+        util.printD(f"Fail to get file name from Content-Disposition: {content_disposition}")
+        return None
+
+    # with folder and filename, now we have the full file path
+    return os.path.join(folder, filename)
+
+
+# output is downloaded file path
+def dl_file(
+    url:str,
+    folder:str | None=None,
+    filename:str | None=None,
+    file_path:str | None=None,
+    headers:dict | None=None,
+    duplicate:str | None=None
+) -> Generator[tuple[bool, str] | str, None, None]:
+    """
+    Perform a download.
+
+    yields: tuple(success:bool, filepath or failure message:str) or progress:str
+    """
+
+    if not headers:
+        headers = {}
+
+    success, response_or_error = request_get(url, headers=headers)
+
+    if not success:
+        yield (False, cast(str, response_or_error))
+        return
+
+    response = cast(requests.Response, response_or_error)
+
+    util.printD(f"Start downloading from: {url}")
+
+    # close the response when the function ends
+    with response:
+
+        # get file_path
+        if not file_path:
+            if not (folder and os.path.isdir(folder)):
+                yield (
+                    False,
+                    "No directory to save model to."
+                )
+                return
+
+            if filename:
+                file_path = os.path.join(folder, filename)
+            else:
+                file_path = get_file_path_from_service_headers(response, folder)
+
+            if not file_path:
+                yield (
+                    False,
+                    "Could not get a file_path to place saved file."
+                )
+                return
+
+        util.printD(f"Target file path: {file_path}")
+        base, ext = os.path.splitext(file_path)
+
+        # duplicate handling
+        if os.path.isfile(file_path):
+            if duplicate == "Rename New":
+                # check if file is already exist
+                count = 2
+                new_base = base
+                while os.path.isfile(file_path):
+                    util.printD("Target file already exist.")
+                    # rename duplicate
+                    new_base = f"{base}_{count}"
+                    file_path = f"{new_base}{ext}"
+                    count += 1
+
+            elif duplicate != "Overwrite":
+                yield (
+                    False,
+                    f"File {file_path} already exists! Download will not proceed."
+                )
+                return
+
+        # get file size
+        total_size = int(response.headers['Content-Length'])
+        util.printD(f"File size: {total_size} ({human_readable_filesize(total_size)})")
+
+        yield from download_progress(url, file_path, total_size, headers, response)
+
+
+def human_readable_filesize(size:int | float) -> str:
+    """ Convert file size to human readable text """
+    prefixes = ["", "K", "M", "G"]
+
+    # Mac reports filesizes in multiples of 1000
+    # We should respect platform differences
+    unit = 1000 if platform.system() == "Darwin" else 1024
+
+    i = 0
+    while size > unit and i < len(prefixes) - 1:
+        i += 1
+        size = size / unit
+
+    return f"{round(size, 2)}{prefixes[i]}"
+
+
+def error(download_url:str, msg:str) -> str:
+    """ Display a download error """
+    output = util.indented_msg(
+        f"""
+        Download failed.
+        {msg}
+        Download url: {download_url}
+        """
+    )
+    util.printD(output)
+    return output

extensions/Stable-Diffusion-Webui-Civitai-Helper/scripts/ch_lib/duplicate_check.py

@@ -0,0 +1,392 @@
+""" -*- coding: UTF-8 -*-
+Check for duplicate models.
+"""
+
+import os
+import html
+import json
+import traceback
+import gradio as gr
+from . import util
+from . import model
+from . import civitai
+from . import templates
+
+
+def scan_for_dups(scan_model_types, cached_hash, progress=gr.Progress()):
+    """ Scans model metadata to detect duplicates
+        by using the model hash.
+    """
+
+    util.printD("Start scan_for_dups")
+    output = ""
+
+    # check model types
+    if not scan_model_types:
+        output = "Model Types is None. Will not scan."
+        util.printD(output)
+        return output
+
+    model_types = []
+
+    # check type if it is a string
+    if isinstance(scan_model_types, str):
+        model_types.append(scan_model_types)
+    else:
+        model_types = scan_model_types
+
+    result = None
+    for result in gather_model_data(model_types, cached_hash):
+        if isinstance(result, tuple):
+            percent, status = result
+            progress(percent, desc=status)
+    models = result
+
+    dups = check_for_dups(models)
+
+    output = create_dups_html(dups)
+
+    return f"<section class=extra-network-cards>{output}</section>"
+
+
+def gather_model_data(model_types, cached_hash):
+    """ Collects model metadata files and parses them for metadata """
+
+    models = {}
+
+    for model_type, model_folder in model.folders.items():
+        if model_type not in model_types:
+            continue
+
+        for result in scan_dir(model_folder, model_type, cached_hash):
+            yield result
+
+        models[model_type] = result
+
+    yield models
+
+def scan_dir(model_folder, model_type, cached_hash):
+    """
+        Scans dir for models and their metadata
+    """
+
+    suffix = f"{civitai.SUFFIX}{model.CIVITAI_EXT}"
+    suffix_len = -len(suffix)
+
+    metadata = []
+    util.printD(f"Scanning path: {model_folder}")
+    for root, _, files in os.walk(model_folder, followlinks=True):
+        for filename in files:
+            try:
+                if filename[suffix_len:] == suffix:
+                    for result in parse_metadata(model_folder, root, filename, suffix, model_type, cached_hash):
+                        yield result
+                    data = result
+                    if data:
+                        metadata.append(data)
+
+            except (IndexError, KeyError, ValueError):
+                util.printD(f"Error occurred on file `{root}/{filename}`")
+                traceback.print_exc()
+                util.printD("You can probably ignore this")
+                continue
+
+    yield metadata
+
+
+def parse_metadata(model_folder, root, filename, suffix, model_type, cached_hash):
+    """ Parses model JSON file for hash / other metadata """
+
+    metadata = {}
+
+    filepath = f"{root}/{filename}"
+    model_name = filename[:-len(suffix)]
+
+    util.printD(f"Processing {model_name}")
+
+    with open(filepath) as file:
+        try:
+            model_info = json.load(file)
+        except json.JSONDecodeError:
+            yield None
+            return
+
+    model_file = model_info["files"][0]
+    model_ext = model_file["name"].split(".").pop()
+
+    description = None
+
+    # Backwards compatibility with older model info files
+    # from pre civitai helper 1.7
+    try:
+        description = model_info["model"]["description"]
+    except (ValueError, KeyError):
+        description = model_info.get("description", None)
+
+    if not description:
+        description = ""
+
+    model_path = f"{root}/{model_name}.{model_ext}"
+
+    if not os.path.isfile(model_path):
+        model_path = locate_model_from_partial(root, model_name)
+
+        if not (model_path and os.path.isfile(model_path)):
+            util.printD(f"No model path found for {filepath}")
+            yield None
+            return
+
+    for result in get_hash(model_path, model_file, model_type, cached_hash):
+        yield result
+    sha256 = result
+
+    metadata = {
+        "model_name": model_name,
+        "civitai_name": model_info["model"]["name"],
+        "description": description,
+        "model_path": model_path,
+        "subpath": model_path[len(model_folder):],
+        "model_type": model_type,
+        "hash": sha256,
+        "search_term": make_search_term(model_type, model_path, sha256)
+    }
+
+    yield metadata
+
+
+def get_hash(model_path, model_file, model_type, cached_hash):
+    """
+        Get or calculate hash of `model_path/model_file`
+    """
+
+    sha256 = None
+    if cached_hash:
+        try:
+            sha256 = model_file["hashes"]["SHA256"].upper()
+
+        except (KeyError, ValueError):
+            pass
+
+        if sha256:
+            yield sha256
+            return
+
+        util.printD(f"No sha256 hash in metadata for {model_file}. \
+                \n\tGenerating one. This will be slower")
+
+    model_hash_type = {
+        "ckp": "checkpoint",
+        "ti": "textual_inversion",
+        "hyper": "hypernet",
+        "lora": "lora",
+        "lycoris": "lycoris"
+    }[model_type]
+
+    result = None
+    for result in util.gen_file_sha256(
+        model_path,
+        model_type=model_hash_type,
+        use_addnet_hash=False
+    ):
+        yield result
+
+    sha256 = result.upper()
+
+    yield sha256
+
+
+def make_search_term(model_type, model_path, sha256):
+    """ format search term for the correct model type """
+
+    folder = model.folders[model_type]
+    subpath = model_path[len(folder):]
+    sha256 = sha256.lower()
+
+    if not subpath.startswith("/"):
+        subpath = f"/{subpath}"
+
+    if model_type == "hyper":
+        snippet = subpath.split(".")
+        snippet.pop()
+        subpath = ".".join(snippet)
+        return f"{subpath}"
+
+    # All other supported model types seem to use this format
+    return f"{subpath} {sha256}"
+
+
+def locate_model_from_partial(root, model_name):
+    """
+        Tries to locate a model if the extension
+        doesn't match the metadata
+    """
+
+    for ext in model.EXTS:
+        filename = f"{root}/{model_name}{ext}"
+        if os.path.isfile(filename):
+            return filename
+
+    return None
+
+
+def check_for_dups(models):
+    """
+        returns all models that share the a stored
+        sha256 hash with another model
+    """
+
+    scanned = {}
+    dups = {}
+
+    for model_type, models_of_type in models.items():
+        scanned[model_type] = {}
+        scanned_type = scanned[model_type]
+        for model_data in models_of_type:
+            sha256 = model_data["hash"]
+
+            if model_type == "lycoris":
+                if is_lycoris_lora(model_data, scanned):
+                    continue
+
+            if not scanned_type.get(sha256, None):
+                scanned_type[sha256] = [model_data]
+                continue
+
+            scanned_type[sha256].append(model_data)
+
+    for model_type, models_of_type in scanned.items():
+        dups_of_type = dups[model_type] = {}
+        for key, model_data in models_of_type.items():
+            if len(model_data) > 1:
+                dups_of_type[key] = model_data
+
+    return dups
+
+
+def get_preview(model_path):
+    """ Finds the appropriate preview image for a model """
+
+    prevs = model.get_potential_model_preview_files(model_path, True)
+
+    bg_image = None
+    for prev in prevs:
+        if os.path.isfile(prev):
+            bg_image = prev
+            break
+
+    if not bg_image:
+        return ""
+
+    return templates.duplicate_preview.substitute(
+        bg_image=bg_image
+    )
+
+
+def make_model_card(model_data):
+    """ Creates the HTML for a single model card """
+
+    card_t = templates.duplicate_card
+
+    bg_image = get_preview(model_data["model_path"])
+    style = "font-size:100%"
+    model_name = model_data["model_name"]
+    subpath = model_data["subpath"].replace("'", "\\'")
+    description = html.escape(model_data["description"])
+    search_term = model_data["search_term"].replace("'", "\\'")
+    model_type = model_data["model_type"]
+
+    util.printD(subpath)
+
+    card = card_t.substitute(
+        style=style,
+        name=model_name,
+        path=subpath,
+        background_image=bg_image,
+        description=description,
+        search_term=search_term,
+        model_type=model_type
+    )
+
+    return card
+
+
+def create_dups_html(dups):
+    """ creates an HTML snippet containing duplicate models """
+
+    article_t = templates.duplicate_article
+    row_t = templates.duplicate_row
+    column_t = templates.duplicate_column
+
+    articles = []
+
+    for model_type, models_of_type in dups.items():
+        rows = []
+        for dup_data in models_of_type.values():
+            civitai_name = ""
+            sha256 = ""
+
+            columns = []
+            for count, model_data in enumerate(dup_data):
+                card = make_model_card(model_data)
+
+                column = column_t.substitute(
+                    count=count,
+                    card=card
+                )
+
+                columns.append(column)
+
+                if model_data["civitai_name"] and not civitai_name:
+                    civitai_name = model_data["civitai_name"]
+                    sha256 = model_data["hash"]
+
+            rows.append(
+                row_t.substitute(
+                    civitai_name=civitai_name,
+                    hash=sha256,
+                    columns="".join(columns)
+                )
+            )
+
+        content = ""
+        if len(rows) > 0:
+            content = "".join(rows)
+        else:
+            content = f"No duplicate {model_type}s found!"
+
+        articles.append(
+            article_t.substitute(
+                section_name=model_type,
+                contents=content
+            )
+        )
+
+    if len(articles) < 1:
+        return "Found no duplicate models!"
+
+    return "".join(articles)
+
+
+def is_lycoris_lora(lyco, models):
+    """
+        Compares a lycoris model to scanned loras to ensure that they're not the same file.
+    """
+    loras = None
+    try:
+        loras = models["lora"][lyco["hash"]]
+
+    except (KeyError, ValueError):
+        return False
+
+    try:
+        lyco_path = os.path.realpath(lyco["model_path"], strict=True)
+
+        for lora in loras:
+            lora_path = os.path.realpath(lora["model_path"], strict=True)
+            if lyco_path == lora_path:
+                return True
+
+    except OSError:
+        return False
+
+    return False

extensions/Stable-Diffusion-Webui-Civitai-Helper/scripts/ch_lib/js_action_civitai.py

@@ -1,8 +1,8 @@
-# -*- coding: UTF-8 -*-
-# handle msg between js and python side
+""" -*- coding: UTF-8 -*-
+handle msg between js and python side
+"""
 import os
-import json
-import requests
+from pathlib import Path
 import webbrowser
 from . import util
 from . import model
@@ -11,19 +11,21 @@ from . import msg_handler
 from . import downloader
 
 
-
-# get civitai's model url and open it in browser
-# parameter: model_type, search_term
-# output: python msg - will be sent to hidden textbox then picked by js side
-def open_model_url(msg, open_url_with_js):
+def open_model_url(msg):
+    """
+    get civitai's model url and open it in browser
+    parameter: model_type, search_term
+    output: python msg
+        - will be sent to hidden textbox then picked by js side
+    """
     util.printD("Start open_model_url")
 
     output = ""
     result = msg_handler.parse_js_msg(msg)
     if not result:
         util.printD("Parsing js ms failed")
-        return
-    
+        return None
+
     model_type = result["model_type"]
     search_term = result["search_term"]
 
@@ -41,16 +43,15 @@ def open_model_url(msg, open_url_with_js):
         util.printD(f"model id from info file of {model_type} {search_term} is None")
         return ""
 
-    url = civitai.url_dict["modelPage"]+str(model_id)
-
+    url = f'{civitai.URLS["modelPage"]}{model_id}'
 
     # msg content for js
     content = {
-        "url":""
+        "url": ""
     }
 
-    if not open_url_with_js:
-        util.printD("Open Url: " + url)
+    if not util.get_opts("ch_open_url_with_js"):
+        util.printD(f"Open Url: {url}")
         # open url
         webbrowser.open_new_tab(url)
     else:
@@ -62,50 +63,50 @@ def open_model_url(msg, open_url_with_js):
     return output
 
 
-
-# add trigger words to prompt
-# parameter: model_type, search_term, prompt
-# return: [new_prompt, new_prompt] - new prompt with trigger words, return twice for txt2img and img2img
 def add_trigger_words(msg):
+    """
+    add trigger words to prompt
+    parameter: model_type, search_term, prompt
+    return: [new_prompt, new_prompt]
+        - new prompt with trigger words, return twice for txt2img and img2img
+    """
     util.printD("Start add_trigger_words")
 
     result = msg_handler.parse_js_msg(msg)
     if not result:
         util.printD("Parsing js ms failed")
-        return
-    
+        return None
+
     model_type = result["model_type"]
     search_term = result["search_term"]
     prompt = result["prompt"]
 
-
     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():
         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:
+
+    trained_words = model_info.get("trainedWords", [])
+    if len(trained_words) == 0:
         util.printD(f"trainedWords from info file for {model_type} {search_term} is empty")
         return [prompt, prompt]
-    
-    # get ful trigger words
-    trigger_words = ""
-    for word in trainedWords:
-        trigger_words = trigger_words + word + ", "
 
-    new_prompt = prompt + " " + trigger_words
-    util.printD("trigger_words: " + trigger_words)
-    util.printD("prompt: " + prompt)
-    util.printD("new_prompt: " + new_prompt)
+    # guess if trained words are a list of words or list of prompts
+    prompt_list = ',' in trained_words[0]
+
+    # if a list of prompts, join with a newline, else a comma and a space
+    separator = "\n" if prompt_list else ", "
+    trigger_words = separator.join(trained_words)
+
+    new_prompt = f"{prompt} {trigger_words}"
+
+    util.printD(f"trigger_words: {trigger_words}")
+    util.printD(f"prompt: {prompt}")
+    util.printD(f"new_prompt: {new_prompt}")
 
     util.printD("End add_trigger_words")
 
@@ -113,18 +114,20 @@ def add_trigger_words(msg):
     return [new_prompt, new_prompt]
 
 
-
-# use preview image's prompt as prompt
-# parameter: model_type, model_name, prompt, neg_prompt
-# return: [new_prompt, new_neg_prompt, new_prompt, new_neg_prompt,] - return twice for txt2img and img2img
 def use_preview_image_prompt(msg):
+    """
+    use preview image's prompt as prompt
+    parameter: model_type, model_name, prompt, neg_prompt
+    return: [new_prompt, new_neg_prompt, new_prompt, new_neg_prompt,]
+        - return twice for txt2img and img2img
+    """
     util.printD("Start use_preview_image_prompt")
 
     result = msg_handler.parse_js_msg(msg)
     if not result:
         util.printD("Parsing js ms failed")
-        return
-    
+        return None
+
     model_type = result["model_type"]
     search_term = result["search_term"]
     prompt = result["prompt"]
@@ -135,188 +138,244 @@ 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}")
-        return [prompt, neg_prompt, prompt, neg_prompt]
-    
+
+    images = model_info.get("images", [])
     if len(images) == 0:
-        util.printD(f"images from info file for {model_type} {search_term} is empty")
+        util.printD(f"No images from info file for {model_type} {search_term}")
         return [prompt, neg_prompt, prompt, neg_prompt]
-    
+
     # get prompt from preview images' meta data
     preview_prompt = ""
     preview_neg_prompt = ""
     for img in images:
-        if "meta" in img.keys():
-            if img["meta"]:
-                if "prompt" in img["meta"].keys():
-                    if img["meta"]["prompt"]:
-                        preview_prompt = img["meta"]["prompt"]
-                
-                if "negativePrompt" in img["meta"].keys():
-                    if img["meta"]["negativePrompt"]:
-                        preview_neg_prompt = img["meta"]["negativePrompt"]
-
-                # we only need 1 prompt
-                if preview_prompt:
-                    break
-            
+        meta = img.get("meta", {})
+        preview_prompt = meta.get("prompt", "")
+        preview_neg_prompt = meta.get("negativePrompt", "")
+
+        # we only need 1 prompt
+        if preview_prompt:
+            break
+
     if not preview_prompt:
         util.printD(f"There is no prompt of {model_type} {search_term} in its preview image")
         return [prompt, neg_prompt, prompt, neg_prompt]
-    
+
     util.printD("End use_preview_image_prompt")
-    
+
     return [preview_prompt, preview_neg_prompt, preview_prompt, preview_neg_prompt]
 
 
-# download model's new verson by model path, version id and download url
-# output is a md log
-def dl_model_new_version(msg, max_size_preview, skip_nsfw_preview):
+def dl_model_new_version(msg, nsfw_preview_threshold):
+    """
+    download model's new verson by model path, version id and download url
+    output is a md log
+
+    This method is triggered by a click event on the client/js
+    side that sends a signal to download a single new model
+    version. The actual check for new models is in
+    `model_action_civitai.check_models_new_version_to_md`.
+    return: output:str
+    """
     util.printD("Start dl_model_new_version")
 
     output = ""
 
+    max_size_preview = util.get_opts("ch_max_size_preview")
+
     result = msg_handler.parse_js_msg(msg)
     if not result:
-        output = "Parsing js ms failed"
+        output = "Parsing js msg failed"
         util.printD(output)
-        return output
-    
+        yield output
+        return
+
     model_path = result["model_path"]
     version_id = result["version_id"]
     download_url = result["download_url"]
-
-    util.printD("model_path: " + model_path)
-    util.printD("version_id: " + str(version_id))
-    util.printD("download_url: " + download_url)
+    model_type = result["model_type"]
 
     # check data
-    if not model_path:
-        output = "model_path is empty"
+    if not (model_path and version_id and download_url):
+        output = util.indented_msg(f"""
+            Missing parameter:
+            {model_path=}
+            {version_id=}
+            {download_url=}
+        """)
         util.printD(output)
-        return output
+        yield output
+        return
 
-    if not version_id:
-        output = "version_id is empty"
-        util.printD(output)
-        return output
-    
-    if not download_url:
-        output = "download_url is empty"
-        util.printD(output)
-        return output
+    util.printD(f"model_path: {model_path}")
+    util.printD(f"version_id: {version_id}")
+    util.printD(f"download_url: {download_url}")
 
     if not os.path.isfile(model_path):
-        output = "model_path is not a file: "+ model_path
+        output = f"model_path is not a file: {model_path}"
         util.printD(output)
-        return output
+        yield output
+        return
 
     # get model folder from model path
     model_folder = os.path.dirname(model_path)
 
-    # no need to check when downloading new version, since checking new version is already checked
-    # check if this model is already existed
-    # r = civitai.search_local_model_info_by_version_id(model_folder, version_id)
-    # if r:
-    #     output = "This model version is already existed"
-    #     util.printD(output)
-    #     return output
-
-    # download file
-    new_model_path = downloader.dl(download_url, model_folder, None, None)
-    if not new_model_path:
-        output = "Download failed, check console log for detail. Download url: " + download_url
+    success = False
+    # download file + webui visible progress bar
+    for result in downloader.dl_file(download_url, folder=model_folder):
+        if not isinstance(result, str):
+            success, output = result
+            break
+
+        yield result
+
+    if not success:
         util.printD(output)
-        return output
+        yield "Model download failed. See console for more details."
+        return
 
     # get version info
     version_info = civitai.get_version_info_by_version_id(version_id)
-    if not version_info:
-        output = "Model downloaded, but failed to get version info, check console log for detail. Model saved to: " + new_model_path
-        util.printD(output)
-        return output
 
-    # now write version info to file
-    base, ext = os.path.splitext(new_model_path)
-    info_file = base + civitai.suffix + model.info_ext
-    model.write_model_info(info_file, version_info)
+    # now write version info to files
+    model.process_model_info(output, version_info, model_type)
 
     # then, get preview image
-    civitai.get_preview_image_by_model_path(new_model_path, max_size_preview, skip_nsfw_preview)
-    
-    output = "Done. Model downloaded to: " + new_model_path
+    for result in civitai.get_preview_image_by_model_path(
+        output,
+        max_size_preview,
+        nsfw_preview_threshold
+    ):
+        yield result
+
+    output = f"Done. Model downloaded to: {output}"
     util.printD(output)
-    return output
+    yield output
 
 
+def get_model_path_from_js_msg(result):
+    """
+    Gets a model path based on the webui js_msg.
 
-# 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)
+    return: model_path
+    """
     if not result:
         output = "Parsing js ms failed"
+        util.error(output)
         util.printD(output)
-        return output
-    
+        return None
+
     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.error(output)
         util.printD(output)
-        return output
-
+        return None
 
     if not os.path.isfile(model_path):
         output = f"Model {model_type} {search_term} does not exist, no need to remove"
+        util.error(output)
         util.printD(output)
-        return output
-    
-    # all files need to be removed
-    related_paths = []
-    related_paths.append(model_path)
+        return None
+
+    return model_path
+
+
+def make_new_filename(candidate_file, model_name, new_name):
+    """
+    Substitutes and old model name for a new model name.
+    return: new_path:str or None
+    """
+    path, filename = os.path.split(candidate_file)
+
+    if filename.index(model_name) != 0:
+        output = util.indented_msg(f"""
+            Could not find model_name in candidate file
+            {model_name=}
+            {candidate_file=}
+            {new_name=}
+        """)
+        util.error(output)
+        util.printD(output)
+        return None
 
+    # handles [model_name].civitai.info and [model_name].preview.[ext]
+    new_filename = filename.replace(model_name, new_name, 1)
 
-    # 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
+    new_path = os.path.join(path, new_filename)
 
-    if os.path.isfile(civitai_info_path):
-        related_paths.append(civitai_info_path)
+    return new_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)
+def rename_model_by_path(msg):
+    """
+    Rename a model file and all related ch_helper/
+    preview image files.
+    """
+    util.printD("Start rename_model_by_path")
 
-    if os.path.isfile(info_path):
-        related_paths.append(info_path)
+    output = ""
+    result = msg_handler.parse_js_msg(msg)
 
-    # remove files
-    for rp in related_paths:
-        if os.path.isfile(rp):
-            util.printD(f"Removing file {rp}")
-            os.remove(rp)
+    model_path = get_model_path_from_js_msg(result)
 
-    util.printD(f"{len(related_paths)} file removed")
+    if model_path is None:
+        output = "Could not rename model."
+        return output
 
-    util.printD("End remove_model_by_path")
+    # all files need to be renamed
+    model_files = model.get_model_files_from_model_path(model_path)
+    model_name = Path(model_path).stem
+    new_name = util.bash_filename(result["new_name"])
+
+    renamed = []
+    for candidate_file in model_files:
+        new_path = make_new_filename(candidate_file, model_name, new_name)
+        if new_path is None:
+            continue
+
+        renamed.append(f"* {candidate_file} to {new_path}")
+        util.printD(f"Renaming file {candidate_file} to {new_path}")
+        os.rename(candidate_file, new_path)
+
+    renamed = "\n".join(renamed)
+    status = f"The following files were renamed: \n{renamed}"
+    util.info(status)
+
+    util.printD("End rename_model_by_path")
     return output
 
 
+def remove_model_by_path(msg):
+    """
+    Remove a model file and all related ch_helper/
+    preview image files.
+    """
+    output = ""
+    util.printD("Start remove_model_by_path")
+
+    result = msg_handler.parse_js_msg(msg)
+
+    model_path = get_model_path_from_js_msg(result)
+
+    if model_path is None:
+        output = "Could not remove model."
+        return output
+
+    # all files need to be renamed
+    model_files = model.get_model_files_from_model_path(model_path)
+
+    removed = []
+    for candidate_file in model_files:
+        util.printD(f"* Removing file {candidate_file}")
+        removed.append(candidate_file)
+        os.remove(candidate_file)
+
+    removed = "\n".join(removed)
+    status = f"The following files were removed: \n{removed}"
+    util.info(status)
 
+    util.printD("End remove_model_by_path")
+    return output

extensions/Stable-Diffusion-Webui-Civitai-Helper/scripts/ch_lib/model.py

@@ -1,35 +1,73 @@
-# -*- coding: UTF-8 -*-
-# handle msg between js and python side
+""" -*- coding: UTF-8 -*-
+Handle model operations
+"""
 import os
 import json
-from . import util
+import re
+from PIL import Image
+import piexif
+import piexif.helper
 from modules import shared
+from modules import paths_internal
+from . import civitai
+from . import downloader
+from . import util
 
 
 # this is the default root path
-root_path = os.getcwd()
-
-# if command line arguement is used to change model folder, 
-# then model folder is in absolute path, not based on this root path anymore.
-# so to make extension work with those absolute model folder paths, model folder also need to be in absolute path
+ROOT_PATH = paths_internal.data_path
+
+EXTS = (".bin", ".pt", ".safetensors", ".ckpt")
+CIVITAI_EXT = ".info"
+SDWEBUI_EXT = ".json"
+
+"""
+If command line arguement is used to change model folder,
+then model folder is in absolute path, not based on this root path anymore.
+so to make extension work with those absolute model folder paths, model
+folder also need to be in absolute path
+"""
 folders = {
-    "ti": os.path.join(root_path, "embeddings"),
-    "hyper": os.path.join(root_path, "models", "hypernetworks"),
-    "ckp": os.path.join(root_path, "models", "Stable-diffusion"),
-    "lora": os.path.join(root_path, "models", "Lora"),
+    "ti": os.path.join(ROOT_PATH, "embeddings"),
+    "hyper": os.path.join(ROOT_PATH, "models", "hypernetworks"),
+    "ckp": os.path.join(ROOT_PATH, "models", "Stable-diffusion"),
+    "lora": os.path.join(ROOT_PATH, "models", "Lora"),
+    "lycoris": os.path.join(ROOT_PATH, "models", "LyCORIS"),
 }
 
-exts = (".bin", ".pt", ".safetensors", ".ckpt")
-info_ext = ".info"
-vae_suffix = ".vae"
+# Separate because the above is used for detecting supported models
+# in other features
+vae_folder = os.path.join(ROOT_PATH, "models", "VAE")
+
+
+class VersionMismatchException(Exception):
+    """ Used for version comarison failures """
+
+    def __init__(self, value):
+        self.value = value
+
+    def __str__(self):
+        return repr(self.value)
+
 
+def get_model_info_paths(model_path):
+    """
+    Retrieve model info paths
+    return: (info_file:str, sd15_file:str)
+    """
+    base, _ = os.path.splitext(model_path)
+    info_file = f"{base}{civitai.SUFFIX}{CIVITAI_EXT}"
+    sd15_file = f"{base}{SDWEBUI_EXT}"
+    return (info_file, sd15_file)
 
-# get cusomter model path
+
+# get custom model path
 def get_custom_model_folder():
+    """
+    Update extra network directories with user-specified values.
+    """
     util.printD("Get Custom Model Folder")
 
-    global folders
-
     if shared.cmd_opts.embeddings_dir and os.path.isdir(shared.cmd_opts.embeddings_dir):
         folders["ti"] = shared.cmd_opts.embeddings_dir
 
@@ -42,131 +80,602 @@ def get_custom_model_folder():
     if shared.cmd_opts.lora_dir and os.path.isdir(shared.cmd_opts.lora_dir):
         folders["lora"] = shared.cmd_opts.lora_dir
 
+    if shared.cmd_opts.vae_dir and os.path.isdir(shared.cmd_opts.vae_dir):
+        vae_folder = shared.cmd_opts.vae_dir
+
+    if util.get_opts("ch_dl_lyco_to_lora"):
+        folders["lycoris"] = folders["lora"]
+
+    try:
+        # pre-1.5.0
+        if os.path.isdir(shared.cmd_opts.lyco_dir):
+            folders["lycoris"] = shared.cmd_opts.lyco_dir
+
+    except AttributeError:
+        try:
+            # sd-webui v1.5.1 added a backcompat option for lyco.
+            if os.path.isdir(shared.cmd_opts.lyco_dir_backcompat):
+                folders["lycoris"] = shared.cmd_opts.lyco_dir_backcompat
+
+        except AttributeError:
+            # v1.5.0 has no options for the Lyco dir:
+            # it is hardcoded as 'os.path.join(paths.models_path, "LyCORIS")'
+            return
+
+
+def metadata_needed(info_file, sd15_file, refetch_old):
+    """ return True if metadata is needed
+    """
+
+    need_civitai = metadata_needed_for_type(info_file, "civitai", refetch_old)
+    need_sdwebui = metadata_needed_for_type(sd15_file, "sdwebui", refetch_old)
+
+    return need_civitai or need_sdwebui
+
+
+def metadata_needed_for_type(path, meta_type, refetch_old):
+    """ return True if metadata is needed for path
+    """
+
+    if meta_type == "sdwebui" and not util.get_opts("ch_dl_webui_metadata"):
+        return False
+
+    if not os.path.isfile(path):
+        return True
+
+    if refetch_old:
+        metadata = None
+        with open(path) as file:
+            metadata = json.load(file)
+
+        metadata_version = util.metadata_version(metadata)
+
+        if not metadata_version:
+            return True
+
+        if meta_type == "civitai":
+            compat_version = util.COMPAT_VERSION_CIVITAI
+        else:
+            compat_version = util.COMPAT_VERSION_SDWEBUI
+
+        return util.newer_version(compat_version, metadata_version)
+
+    return False
+
 
+def verify_overwrite_eligibility(path, new_data):
+    """
+    Verifies a file is valid to be overwritten
+    Throws an error if the model ID does not match the new version's model ID
+    return: True if valid, False if not.
+    """
+    if not os.path.isfile(path):
+        return True
 
+    with open(path, "r") as file:
+        old_data = json.load(file)
 
+    if "civitai" in path:
+        new_id = new_data.get("id", "")
+        old_id = old_data.get("id", "")
+        if new_id != old_id:
+            if old_id != "":
+                raise VersionMismatchException(
+                    f"New metadata id ({new_id}) does not match old metadata id ({old_id})"
+                )
 
-# write model info to file
-def write_model_info(path, model_info):
-    util.printD("Write model info to file: " + path)
-    with open(os.path.realpath(path), 'w') as f:
-        f.write(json.dumps(model_info, indent=4))
+    new_description = new_data.get("description", "")
+    old_description = old_data.get("description", "")
+    if new_description == "" and old_description != "":
+        util.printD(
+            f"New description is blank while old description contains data. Skipping {path}"
+        )
+        return False
+
+    return True
+
+
+def write_info(data, path, info_type):
+    """ Writes model info to a file """
+    util.printD(f"Write model {info_type} info to file: {path}")
+    with open(os.path.realpath(path), 'w') as info_file:
+        info_file.write(json.dumps(data, indent=4))
+
+
+def process_model_info(model_path, model_info, model_type="ckp", refetch_old=False):
+    """
+    Write model info to file
+
+    SD1.5 Webui added saving model information to JSON files.
+    Much of this extension's metadata management is replicated
+    by this new functionality, including automatically adding
+    activator keywords to the prompt. It also provides a much
+    cleaner UI than civitai (not a high bar to clear) to
+    simply read a model's description.
+
+    So why not populate it with useful information?
+
+    Returns True if successful, otherwise an error message.
+    """
+
+    if model_info is None:
+        util.printD("Failed to get model info.")
+        return
+
+    info_file, sd15_file = get_model_info_paths(model_path)
+
+    parent = model_info["model"]
+
+    description = parent.get("description", "")
+    if description:
+        description = util.trim_html(description)
+    parent["description"] = description
+
+    version_description = model_info.get("description", "")
+    if version_description:
+        version_description = util.trim_html(version_description)
+    model_info["description"] = version_description
+
+    # Create extension versioning information so that users
+    # can replace stale info files without newer entries.
+    model_info["extensions"] = util.create_extension_block(
+        model_info.get("extensions", None),
+        model_info.get("skeleton_file", False)
+    )
+
+    # civitai model info file
+    if metadata_needed_for_type(info_file, "civitai", refetch_old):
+        if refetch_old:
+            try:
+                if verify_overwrite_eligibility(info_file, model_info):
+                    write_info(model_info, info_file, "civitai")
+
+            except VersionMismatchException as e:
+                util.printD(f"{e}, aborting")
+                return
+
+        else:
+            write_info(model_info, info_file, "civitai")
+
+    if not util.get_opts("ch_dl_webui_metadata"):
+        return
+
+    # Do not overwrite user-created files!
+    # TODO: maybe populate empty fields in existing files?
+    if not metadata_needed_for_type(sd15_file, "sdwebui", refetch_old):
+        util.printD(f"Metadata not needed for: {sd15_file}.")
+        return
+
+    process_sd15_info(sd15_file, model_info, parent, model_type, refetch_old)
+
+
+def process_sd15_info(sd15_file, model_info, parent, model_type, refetch_old):
+    """ Creates/Processes [model_name].json """
+
+    # sd v1.5 model info file
+    sd_data = {}
+
+    sd_data["description"] = parent.get("description", "")
+
+    # I suppose notes are more for user notes, but populating it
+    # with potentially useful information about this particular
+    # version of the model is fine too, right? The user can
+    # always replace these if they're unneeded or add to them
+    version_info = model_info.get("description", None)
+    if version_info is not None:
+        sd_data["notes"] = version_info
+
+    # AFAIK civitai model versions are currently:
+    # SD 1.4, SD 1.5, SD 2.0, SD 2.0 786, SD 2.1, SD 2.1 786
+    # SD 2.1 Unclip, SDXL 0.9, SDXL 1.0, and Other.
+    # Conveniently, the 4th character is all we need for webui.
+    #
+    # INFO: On Civitai, all models list base model/"sd version".
+    # The SD WebUI interface only displays them for Lora/Lycoris.
+    #
+    # I'm populating the field anyways in hopes it eventually gets
+    # added.
+    base_model = model_info.get("baseModel", None)
+    sd_version = 'Unknown'
+    if base_model:
+        version = base_model[3]
+
+        sd_version = {
+            "1": 'SD1',
+            "2": 'SD2',
+            "L": 'SDXL',
+        }.get(version, 'Unknown')
+
+    sd_data["sd version"] = sd_version
+
+    for filedata in model_info["files"]:
+        if filedata["type"] == "VAE":
+            sd_data["vae"] = filedata["name"]
+
+    # INFO: On Civitai, all non-checkpoint models can have trained words.
+    # The SD WebUI interface only displays them for Lora/Lycoris.
+    # I'm populating the field anyways in hopes it eventually gets
+    # added.
+    #
+    # "trained words" usage is inconsistent among model authors.
+    # Some use each entry as an individual activator, while others
+    # use them as entire prompts
+    activator = model_info.get("trainedWords", [])
+    if (activator and activator[0]):
+        if "," in activator[0]:
+            # assume trainedWords is a prompt list
+
+            # webui does not support newlines in activator text
+            # so this is the best hinting I can give the user at the
+            # moment that these are mutually-exclusive prompts.
+            sd_data["activation text"] = " || ".join(activator)
+        else:
+            # assume trainedWords are single keywords
+            sd_data["activation text"] = ", ".join(activator)
+
+    # Sadly, Civitai does not provide default weight information,
+    # So 0 disables this functionality on webui's end and uses
+    # the user's global setting
+    if model_type in ["lora", "lycoris"]:
+        sd_data["preferred weight"] = 0
+
+    sd_data["extensions"] = util.create_extension_block(
+        model_info.get("extensions", None),
+        model_info.get("skeleton_file", False)
+    )
+
+    if refetch_old:
+        if verify_overwrite_eligibility(sd15_file, sd_data):
+            write_info(sd_data, sd15_file, "webui")
+    else:
+        write_info(sd_data, sd15_file, "webui")
 
 
 def load_model_info(path):
-    # util.printD("Load model info from file: " + path)
+    """ Opens a JSON file and loads its JSON """
     model_info = None
-    with open(os.path.realpath(path), 'r') as f:
+    with open(os.path.realpath(path), 'r') as json_file:
         try:
-            model_info = json.load(f)
-        except Exception as e:
-            util.printD("Selected file is not json: " + path)
-            util.printD(e)
-            return
-        
+            model_info = json.load(json_file)
+        except ValueError:
+            util.printD(f"Selected file is not json: {path}")
+            return None
+
     return model_info
 
 
-# get model file names by model type
-# parameter: model_type - string
-# return: model name list
+def get_potential_model_preview_files(model_path, all_prevs=False):
+    """
+    Find existing preview images, if any.
+
+    Extensions from `find_preview` method in webui `modules/ui_extra_networks.py`
+    gif added in webui commit c602471b85d270e8c36707817d9bad92b0ff991e
+
+    return: preview_files
+    """
+    preview_exts = ["png", "jpg", "jpeg", "webp", "gif"]
+    preview_files = []
+
+    base, _ = os.path.splitext(model_path)
+
+    for ext in preview_exts:
+        if all_prevs:
+            preview_files.append(f"{base}.{ext}")
+        preview_files.append(f"{base}.preview.{ext}")
+
+    return preview_files
+
+
+def get_model_files_from_model_path(model_path):
+    """ return: list of paths """
+
+    base, _ = os.path.splitext(model_path)
+
+    info_file, sd15_file = get_model_info_paths(model_path)
+    user_preview_path = f"{base}.png"
+
+    paths = [model_path, info_file, sd15_file, user_preview_path]
+    preview_paths = get_potential_model_preview_files(model_path)
+
+    paths = paths + preview_paths
+
+    return [path for path in paths if os.path.isfile(path)]
+
+
 def get_model_names_by_type(model_type:str) -> list:
-    
-    model_folder = folders[model_type]
+    """
+    get model file names by model type
+    parameter: model_type - string
+    return: model name list
+    """
+
+    if model_type == "lora" and folders['lycoris']:
+        model_folders = [folders[model_type], folders['lycoris']]
+    else:
+        model_folders = [folders[model_type]]
 
     # get information from filter
     # only get those model names don't have a civitai model info file
     model_names = []
-    for root, dirs, files in os.walk(model_folder, followlinks=True):
-        for filename in files:
-            item = os.path.join(root, filename)
-            # check extension
-            base, ext = os.path.splitext(item)
-            if ext in exts:
-                # find a model
-                model_names.append(filename)
-
+    for model_folder in model_folders:
+        for root, _, files in os.walk(model_folder, followlinks=True):
+            for filename in files:
+                item = os.path.join(root, filename)
+                # check extension
+                _, ext = os.path.splitext(item)
+                if ext in EXTS:
+                    # find a model
+                    model_names.append(filename)
 
     return model_names
 
 
 # return 2 values: (model_root, model_path)
 def get_model_path_by_type_and_name(model_type:str, model_name:str) -> str:
+    """ return: model_path:str matching model_name and model_type """
     util.printD("Run get_model_path_by_type_and_name")
-    if model_type not in folders.keys():
-        util.printD("unknown model_type: " + model_type)
-        return
-    
     if not model_name:
         util.printD("model name can not be empty")
-        return
-    
-    folder = folders[model_type]
+        return None
 
-    # model could be in subfolder, need to walk.
-    model_root = ""
-    model_path = ""
-    for root, dirs, files in os.walk(folder, followlinks=True):
-        for filename in files:
-            if filename == model_name:
-                # find model
-                model_root = root
-                model_path = os.path.join(root, filename)
-                return (model_root, model_path)
+    model_folders = [folders.get(model_type, None)]
+
+    if model_folders[0] is None:
+        util.printD(f"unknown model_type: {model_type}")
+        return None
 
-    return
+    if model_type == "lora" and folders['lycoris']:
+        model_folders.append(folders['lycoris'])
 
+    # model could be in subfolder, need to walk.
+    model_path = util.find_file_in_folders(model_folders, model_name)
+
+    msg = util.indented_msg(f"""
+        Got following info:
+        {model_path=}
+    """)
+    util.printD(msg)
 
+    # May return `None`
+    return model_path
 
 
 # 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:str, search_term:str):
+def get_model_path_by_search_term(model_type, search_term):
+    """
+    Gets a model path based on the webui search term.
+
+    return: model_path
+    """
     util.printD(f"Search model of {search_term} in {model_type}")
-    if model_type not in folders.keys():
+    if folders.get(model_type, None) is None:
         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
+        return None
+
+    # 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
-    has_hash = True
-    if model_type == "hyper":
-        has_hash = False
-    elif search_term.endswith(".pt") or search_term.endswith(".bin") or search_term.endswith(".safetensors") or search_term.endswith(".ckpt"):
-        has_hash = False
-
-    # remove hash
-    # model name may have multiple spaces
-    splited_path = search_term.split()
-    model_sub_path = splited_path[0]
-    if has_hash and len(splited_path) > 1:
-        model_sub_path = ""
-        for i in range(0, len(splited_path)-1):
-            model_sub_path += splited_path[i] + " "
-        
-        model_sub_path = model_sub_path.strip()
 
+    # this used to be
+    # `model_sub_path = search_term.split()[0]`
+    # but it was failing on models containing spaces.
+    model_hash = search_term.split()[-1]
+    model_sub_path = search_term.replace(f" {model_hash}", "")
+
+    if model_type == "hyper":
+        model_sub_path = f"{search_term}.pt"
 
     if model_sub_path[:1] == "/":
         model_sub_path = model_sub_path[1:]
 
-    if model_type == "hyper":
-        model_sub_path = model_sub_path+".pt"
+    if model_type == "lora" and folders['lycoris']:
+        model_folders = [folders[model_type], folders['lycoris']]
+    else:
+        model_folders = [folders[model_type]]
 
-    model_folder = folders[model_type]
+    for folder in model_folders:
+        model_folder = folder
+        model_path = os.path.join(model_folder, model_sub_path)
 
-    model_path = os.path.join(model_folder, model_sub_path)
+        if os.path.isfile(model_path):
+            break
 
-    print("model_folder: " + model_folder)
-    print("model_sub_path: " + model_sub_path)
-    print("model_path: " + model_path)
+    msg = util.indented_msg(f"""
+        Got following info:
+        {model_folder=}
+        {model_sub_path=}
+        {model_path=}
+    """)
+    util.printD(msg)
 
     if not os.path.isfile(model_path):
-        util.printD("Can not find model file: " + model_path)
-        return
-    
+        util.printD(f"Can not find model file: {model_path}")
+        return None
+
     return model_path
 
+
+pattern = re.compile(r"\s*([^:,]+):\s*([^,]+)")
+
+def sd_format(data):
+    """
+    Parse image exif data for image creation parameters.
+
+    return parameters:dict or None
+    """
+
+    if not data:
+        return None
+
+    prompt = ""
+    negative = ""
+    setting = ""
+
+    steps_index = data.find("\nSteps:")
+
+    if steps_index != -1:
+        prompt = data[:steps_index].strip()
+        setting = data[steps_index:].strip()
+
+    if "Negative prompt:" in data:
+        prompt_index = data.find("\nNegative prompt:")
+
+        if steps_index != -1:
+            negative = data[
+                prompt_index + len("Negative prompt:") + 1 : steps_index
+            ].strip()
+
+        else:
+            negative = data[
+                prompt_index + len("Negative prompt:") + 1 :
+            ].strip()
+
+        prompt = data[:prompt_index].strip()
+
+    elif steps_index == -1:
+        prompt = data
+
+    setting_dict = dict(re.findall(pattern, setting))
+
+    data = {
+         "prompt": prompt,
+         "negative": negative,
+         "Steps": setting_dict.get("Steps", ""),
+         "Sampler": setting_dict.get("Sampler", ""),
+         "CFG_scale": setting_dict.get("CFG scale", ""),
+         "Seed": setting_dict.get("Seed", ""),
+         "Size": setting_dict.get("Size", ""),
+    }
+
+    return data
+
+
+def parse_image(image_file):
+    """
+    Read image exif for userComment entry.
+    return: userComment:str
+    """
+    data = None
+    with Image.open(image_file) as image:
+        if image.format == "PNG":
+            # However, unlike other image formats, EXIF data is not
+            # guaranteed to be present in info until load() has been called.
+            # https://pillow.readthedocs.io/en/stable/handbook/image-file-formats.html#png
+            image.load()
+            data = image.info.get("parameters")
+
+        elif image.format in ["JPEG", "WEBP"]:
+            try:
+                usercomment = piexif.ExifIFD.UserComment
+                exif = image.info.get("exif")
+                if not exif:
+                    return None
+                jpegexif = piexif.load(exif) or {}
+                data = piexif.helper.UserComment.load(
+                    jpegexif.get("Exif", {}).get(usercomment, None)
+                )
+
+            except (ValueError, TypeError):
+                util.printD("Failed to parse image exif.")
+                return None
+
+    return data
+
+
+def get_remote_image_info(img_src):
+    """
+    Download a remote image and parse out its creation parameters
+
+    return parameters:dict or None
+    """
+    # anti-DDOS protection
+    util.delay(0.2)
+
+    success, response = downloader.request_get(img_src)
+
+    if not success:
+        return None
+
+    image_file = response.raw
+    try:
+        data = parse_image(image_file)
+
+    except OSError: #, UnidentifiedImageError
+        util.printD("Failed to open image.")
+        return None
+
+    if not data:
+        return None
+
+    sd_data = sd_format(data)
+    return sd_data
+
+
+def update_civitai_info_image_meta(filename):
+    """
+    Read model metadata and update missing image creation parameters,
+    if available.
+    """
+    need_update = False
+    data = {}
+
+    if not os.path.isfile(filename):
+        return
+
+    with open(filename, 'r') as model_json:
+        data = json.load(model_json)
+
+    for image in data.get('images', []):
+        metadata = image.get('meta', None)
+        if not metadata and metadata != {}:
+            url = image.get("url", "")
+            if not url:
+                continue
+
+            util.printD(f"{filename} missing generation info for {url}. Processing {url}.")
+
+            image_data = get_remote_image_info(url)
+            if not image_data:
+                util.printD(f"Failed to find generation info on remote image at {url}.")
+
+                # "mark" image so additional runs will skip it.
+                image["meta"] = {}
+                need_update = True
+                continue
+
+            util.printD(
+                "The following information will be added to "
+                f"{filename} for {url}:\n{image_data}"
+            )
+            metadata = image_data
+            image["meta"] = metadata
+
+            need_update = True
+
+    if need_update:
+        with open(filename, 'w') as info_file:
+            json.dump(data, info_file, indent=4)
+
+
+def scan_civitai_info_image_meta():
+    """ Search for *.civitai.info files """
+    util.printD("Start Scan_civitai_info_image_meta")
+    output = ""
+    count = 0
+
+    directories = [y for x, y in folders.items() if os.path.isdir(y)]
+    util.printD(f"{directories=}")
+    for directory in directories:
+        for root, _, files in os.walk(directory):
+            for filename in files:
+                if filename.endswith('.civitai.info'):
+                    update_civitai_info_image_meta(os.path.join(root, filename))
+                    count = count + 1
+
+    output = f"Done. Scanned {count} files."
+    util.printD(output)
+    return output

extensions/Stable-Diffusion-Webui-Civitai-Helper/scripts/ch_lib/model_action_civitai.py

@@ -1,511 +1,828 @@
-# -*- coding: UTF-8 -*-
-# handle msg between js and python side
+""" -*- coding: UTF-8 -*-
+handle msg between js and python side
+"""
 import os
 import time
+import re
+import gradio as gr
+from modules import sd_models
 from . import util
 from . import model
 from . import civitai
 from . import downloader
+from . import templates
+
+
+def get_metadata_skeleton():
+    """
+    Used to generate at least something when model is not on civitai.
+    """
+    metadata = {
+        "id": "",
+        "modelId": "",
+        "name": "",
+        "trainedWords": [],
+        "baseModel": "Unknown",
+        "description": "",
+        "model": {
+            "name": "",
+            "type": "",
+            "nsfw": "",
+            "poi": ""
+        },
+        "files": [
+            {
+                "name": "",
+                "sizeKB": 0,
+                "type": "Model",
+                "hashes": {
+                    "AutoV2": "",
+                    "SHA256": ""
+                }
+            }
+        ],
+        "tags": [],
+        "downloadUrl": "",
+        "skeleton_file": True
+    }
+
+    return metadata
+
+
+def scan_single_model(filepath, model_type, refetch_old, delay):
+    """
+    Gets model info for a model by feeding its sha256 hash into civitai's api
+
+    return: success:bool
+    """
+
+    filename = os.path.basename(filepath)
+
+    # find a model, get info file
+    info_file, sd15_file = model.get_model_info_paths(filepath)
 
+    output = ""
+
+    # check info file
+    if model.metadata_needed(info_file, sd15_file, refetch_old):
+        output = f"Creating model info for: {filename}"
+        util.printD(output)
+        yield output
+
+        # get model's sha256
+        result = None
+        for result in util.gen_file_sha256(filepath):
+            if isinstance(result, tuple):
+                yield result
+        sha256_hash = result
+
+        util.printD(f"model action sha256: {sha256_hash}")
+
+        if not sha256_hash:
+            output = f"failed generating SHA256 for model: {filename}"
+            util.printD(output)
+            yield output
+            time.sleep(delay)
+            yield False
+
+        yield "Requesting model information from Civitai"
+        # use this sha256 to get model info from civitai
+        model_info = civitai.get_model_info_by_hash(sha256_hash)
+
+        if not model_info:
+            model_info = dummy_model_info(filepath, sha256_hash, model_type)
+
+        model.process_model_info(filepath, model_info, model_type, refetch_old=refetch_old)
+
+        # delay before next request, to prevent being treated as a DDoS attack
+        time.sleep(delay)
+
+    else:
+        util.printD(f"Model metadata not needed for {filename}")
+
+    yield True
+
+
+def scan_model(scan_model_types, nsfw_preview_threshold, refetch_old, progress=gr.Progress()):
+    """ Scan model to generate SHA256, then use this SHA256 to get model info from civitai
+        return output msg
+    """
+
+    delay = 0.2
 
-# scan model to generate SHA256, then use this SHA256 to get model info from civitai
-# return output msg
-def scan_model(scan_model_types, max_size_preview, skip_nsfw_preview):
     util.printD("Start scan_model")
     output = ""
 
+    max_size_preview = util.get_opts("ch_max_size_preview")
+
     # check model types
     if not scan_model_types:
         output = "Model Types is None, can not scan."
         util.printD(output)
-        return output
-    
-    model_types = []
-    # check type if it is a string
-    if type(scan_model_types) == str:
-        model_types.append(scan_model_types)
-    else:
-        model_types = scan_model_types
-    
-    model_count = 0
-    image_count = 0
-    # scan_log = ""
+        yield output
+        return
+
+    model_types = scan_model_types
+    if isinstance(scan_model_types, str):
+        # check if type is a string
+        model_types = [scan_model_types]
+
+    models = []
     for model_type, model_folder in model.folders.items():
         if model_type not in model_types:
             continue
 
-        util.printD("Scanning path: " + model_folder)
-        for root, dirs, files in os.walk(model_folder, followlinks=True):
+        util.printD(f"Scanning path: {model_folder}")
+        for root, _, files in os.walk(model_folder, followlinks=True):
             for filename in files:
+
                 # check ext
-                item = os.path.join(root, filename)
-                base, ext = os.path.splitext(item)
-                if ext in model.exts:
-                    # ignore vae file
-                    if len(base) > 4:
-                        if base[-4:] == model.vae_suffix:
-                            # find .vae
-                            util.printD("This is a vae file: " + filename)
-                            continue
-
-                    # find a model
-                    # get info file
-                    info_file = base + civitai.suffix + model.info_ext
-                    # check info file
-                    if not os.path.isfile(info_file):
-                        util.printD("Creating model info for: " + filename)
-                        # get model's sha256
-                        hash = util.gen_file_sha256(item)
-
-                        if not hash:
-                            output = "failed generating SHA256 for model:" + filename
-                            util.printD(output)
-                            return output
-                        
-                        # use this sha256 to get model info from civitai
-                        model_info = civitai.get_model_info_by_hash(hash)
-                        # delay 1 second for ti
-                        if model_type == "ti":
-                            util.printD("Delay 1 second for TI")
-                            time.sleep(1)
-
-                        if model_info is None:
-                            output = "Connect to Civitai API service failed. Wait a while and try again"
-                            util.printD(output)
-                            return output+", check console log for detail"
-
-                        # write model info to file
-                        model.write_model_info(info_file, model_info)
-
-                    # set model_count
-                    model_count = model_count+1
-
-                    # check preview image
-                    civitai.get_preview_image_by_model_path(item, max_size_preview, skip_nsfw_preview)
-                    image_count = image_count+1
-
-
-    # scan_log = "Done"
-
-    output = f"Done. Scanned {model_count} models, checked {image_count} images"
+                filepath = os.path.join(root, filename)
+                _, ext = os.path.splitext(filepath)
+                if ext not in model.EXTS:
+                    continue
 
-    util.printD(output)
+                models.append((filepath, model_type))
 
-    return output
+    count = [0, 0]
+    total = len(models)
+    for filepath, model_type in models:
+        success = None
 
+        tracker = (count[0], total)
 
+        progress(
+            tracker,
+            desc="Scanning...",
+            unit="models"
+        )
+
+        count[0] = count[0] + 1
+
+        for result in scan_single_model(filepath, model_type, refetch_old, delay):
+            if isinstance(result, str):
+                progress(tracker, desc=result, unit="models")
+                continue
+            if isinstance(result, tuple):
+                percent, status = result
+                progress(percent, desc=status)
+                continue
+
+            success = result
+            break
+
+        if not success:
+            continue
 
-# Get model info by model type, name and url
-# output is log info to display on markdown component
-def get_model_info_by_input(model_type, model_name, model_url_or_id, max_size_preview, skip_nsfw_preview):
+        # set model_count
+        count[1] = count[1] + 1
+
+        # check preview image
+        for _ in civitai.get_preview_image_by_model_path(
+            filepath,
+            max_size_preview,
+            nsfw_preview_threshold
+        ):
+            pass
+
+    # this previously had an image count, but it always matched the model count.
+    output = f"Done. Successfully scanned {count[1]} of {len(models)} models."
+
+    util.printD(output)
+
+    yield output
+
+
+def dummy_model_info(path, sha256_hash, model_type):
+    """
+    Fills model metadata with information we can get locally.
+    """
+    if not sha256_hash:
+        return {}
+
+    model_info = get_metadata_skeleton()
+
+    autov2 = sha256_hash[:10]
+    filename = os.path.basename(path)
+    filesize = os.path.getsize(path) // 1024
+
+    model_metadata = model_info["model"]
+    file_metadata = model_info["files"][0]
+
+    model_metadata["name"] = filename
+    model_metadata["type"] = model_type
+
+    file_metadata["name"] = filename
+    file_metadata["sizeKB"] = filesize
+    file_metadata["hashes"]["SHA256"] = sha256_hash
+    file_metadata["hashes"]["AutoV2"] = autov2
+
+    # We can't get data on the model from civitai, but some models
+    # do store their training data.
+    trained_words = model_info["trainedWords"]
+    tags = model_info["tags"]
+
+    try:
+        file_metadata = sd_models.read_metadata_from_safetensors(path)
+    except AssertionError:
+        # model is not a safetensors file. This is fine,
+        # it just doesn't have metadata we can read
+        pass
+
+    tag_frequency = file_metadata.get("ss_tag_frequency", {})
+
+    for trained_word in tag_frequency.keys():
+        # kohya training scripts use
+        # `{iterations}_{trained_word}`
+        # for training finetune concepts.
+        word = re.sub(r"^\d+_", "", trained_word)
+        trained_words.append(word)
+
+        # "tags" in this case are just words used in image captions
+        # when training the finetune model.
+        # They may or may not be useful for prompting
+        for tag in tag_frequency[trained_word].keys():
+            tag = tag.replace(",", "").strip()
+            if tag == "":
+                continue
+            tags.append(tag)
+
+    return model_info
+
+
+def get_model_info_by_input(
+    model_type, model_name, model_url_or_id, nsfw_preview_threshold
+):
+    """
+    Get model info by model type, name and url
+    output is log info to display on markdown component
+    """
     output = ""
+
+    max_size_preview = util.get_opts("ch_max_size_preview")
+
     # parse model id
     model_id = civitai.get_model_id_from_url(model_url_or_id)
     if not model_id:
-        output = "failed to parse model id from url: " + model_url_or_id
+        output = f"failed to parse model id from url: {model_url_or_id}"
         util.printD(output)
-        return output
+        yield output
+        return
 
     # get model file path
     # model could be in subfolder
-    result = model.get_model_path_by_type_and_name(model_type, model_name)
-    if not result:
-        output = "failed to get model file path"
-        util.printD(output)
-        return output
-    
-    model_root, model_path = result
-    if not model_path:
-        output = "model path is empty"
+    model_path = model.get_model_path_by_type_and_name(model_type, model_name)
+
+    if model_path is None:
+        output = "Could not get Model Path"
         util.printD(output)
-        return output
-    
-    # get info file path
-    base, ext = os.path.splitext(model_path)
-    info_file = base + civitai.suffix + model.info_ext
+        yield output
+        return
 
-    # get model info    
+    # get model info
     #we call it model_info, but in civitai, it is actually version info
     model_info = civitai.get_version_info_by_model_id(model_id)
 
-    if not model_info:
-        output = "failed to get model info from url: " + model_url_or_id
-        util.printD(output)
-        return output
-    
-    # write model info to file
-    model.write_model_info(info_file, model_info)
+    model.process_model_info(model_path, model_info, model_type)
 
-    util.printD("Saved model info to: "+ info_file)
+    # check preview image + webui-visible progress bar
+    yield from civitai.get_preview_image_by_model_path(
+        model_path,
+        max_size_preview,
+        nsfw_preview_threshold
+    )
 
-    # check preview image
-    civitai.get_preview_image_by_model_path(model_path, max_size_preview, skip_nsfw_preview)
+    yield output
 
-    output = "Model Info saved to: " + info_file
-    return output
 
+def build_article_from_version(version):
+    """
+    Builds the HTML for displaying new model versions to the user.
+
+    return: html:str
+    """
+    (
+        model_path, model_id, model_name, new_version_id,
+        new_version_name, description, download_url,
+        img_url, model_type
+    ) = version
+
+    thumbnail = ""
+    if img_url:
+        thumbnail = templates.thumbnail.substitute(
+            img_url=img_url,
+        )
+
+    if download_url:
+        # replace "\" to "/" in model_path for windows
+        download_model_path = model_path.replace('\\', '\\\\')
+
+        download_section = templates.download.substitute(
+            new_version_id=new_version_id,
+            new_version_name=new_version_name,
+            model_path=download_model_path,
+            model_type=model_type,
+            download_url=download_url
+        )
+
+    else:
+        download_section = templates.no_download.substitute(
+            new_version_name=new_version_name,
+        )
+
+    description_section = ""
+    if description:
+        description_section = templates.description.substitute(
+            description=util.safe_html(download_section),
+        )
+
+    article = templates.article.substitute(
+        url=f'{civitai.URLS["modelPage"]}{model_id}',
+        thumbnail=thumbnail,
+        download=download_section,
+        description=description_section,
+        model_name=model_name,
+        model_path=model_path
+    )
+
+    return article
 
 
-# check models' new version and output to UI as markdown doc
 def check_models_new_version_to_md(model_types:list) -> str:
-    new_versions = civitai.check_models_new_version_by_model_types(model_types, 1)
+    """
+    check models' new version and output to UI as html doc
+    return: html:str
+    """
+    new_versions = civitai.check_models_new_version_by_model_types(model_types, 0.2)
 
-    count = 0
-    output = ""
     if not new_versions:
-        output = "No model has new version"
+        util.printD("Done: no new versions found.")
+        return "No models have new versions"
+
+    articles = []
+    count = 0
+    for count, new_version in enumerate(new_versions):
+        article = build_article_from_version(new_version)
+        articles.append(article)
+
+    output = f"Found new versions for following models: <section>{''.join(articles)}</section>"
+
+    if count != 1:
+        util.printD(f"Done. Found {count} models that have new versions. Check UI for detail")
     else:
-        output = "Found new version for following models:  <br>"
-        for new_version in new_versions:
-            count = count+1
-            model_path, model_id, model_name, new_verion_id, new_version_name, description, download_url, img_url = new_version
-            # in md, each part is something like this:
-            # [model_name](model_url)
-            # [version_name](download_url)
-            # version description
-            url = civitai.url_dict["modelPage"]+str(model_id)
-
-            part = f'<div style="font-size:20px;margin:6px 0px;"><b>Model: <a href="{url}" target="_blank"><u>{model_name}</u></a></b></div>'
-            part = part + f'<div style="font-size:16px">File: {model_path}</div>'
-            if download_url:
-                # replace "\" to "/" in model_path for windows
-                model_path = model_path.replace('\\', '\\\\')
-                part = part + f'<div style="font-size:16px;margin:6px 0px;">New Version: <u><a href="{download_url}" target="_blank" style="margin:0px 10px;">{new_version_name}</a></u>'
-                # add js function to download new version into SD webui by python
-                part = part + "    "
-                # in embed HTML, onclick= will also follow a ", never a ', so have to write it as following
-                part = part + f"<u><a href='#' style='margin:0px 10px;' onclick=\"ch_dl_model_new_version(event, '{model_path}', '{new_verion_id}', '{download_url}')\">[Download into SD]</a></u>"
-                
-            else:
-                part = part + f'<div style="font-size:16px;margin:6px 0px;">New Version: {new_version_name}'
-            part = part + '</div>'
+        util.printD(f"Done. Found {count} model that has a new version. Check UI for detail.")
 
-            # description
-            if description:
-                part = part + '<blockquote style="font-size:16px;margin:6px 0px;">'+ description + '</blockquote><br>'
+    return output
 
-            # preview image            
-            if img_url:
-                part = part + f"<img src='{img_url}'><br>"
-                
 
-            output = output + part
+def get_model_info_by_url(model_url_or_id:str) -> dict:
+    """
+    Retrieves model information necessary to populate HTML
+    with Model Name, Model Type, valid saving directories,
+    and available model versions.
 
-    util.printD(f"Done. Find {count} models have new version. Check UI for detail.")
+    return: tuple or None
+    """
+    util.printD(f"Getting model info by: {model_url_or_id}")
 
-    return output
+    try:
+        # parse model id
+        model_id = civitai.get_model_id_from_url(model_url_or_id)
 
+        # download model info
+        model_info = civitai.get_model_info_by_id(model_id)
 
-# get model info by url
-def get_model_info_by_url(model_url_or_id:str) -> tuple:
-    util.printD("Getting model info by: " + model_url_or_id)
+        # parse model type, model name, subfolder, version from this model info
+        # get model type
+        civitai_model_type = model_info["type"]
 
-    # parse model id
-    model_id = civitai.get_model_id_from_url(model_url_or_id)
-    if not model_id:
-        util.printD("failed to parse model id from url or id")
-        return    
+        if civitai_model_type not in civitai.MODEL_TYPES:
+            util.printD(f"This model type is not supported: {civitai_model_type}")
+            return None
 
-    model_info = civitai.get_model_info_by_id(model_id)
-    if model_info is None:
-        util.printD("Connect to Civitai API service failed. Wait a while and try again")
-        return
-    
-    if not model_info:
-        util.printD("failed to get model info from url or id")
-        return
-    
-    # parse model type, model name, subfolder, version from this model info
-    # get model type
-    if "type" not in model_info.keys():
-        util.printD("model type is not in model_info")
-        return
-    
-    civitai_model_type = model_info["type"]
-    if civitai_model_type not in civitai.model_type_dict.keys():
-        util.printD("This model type is not supported:"+civitai_model_type)
-        return
-    
-    model_type = civitai.model_type_dict[civitai_model_type]
+        model_type = civitai.MODEL_TYPES[civitai_model_type]
 
-    # get model type
-    if "name" not in model_info.keys():
-        util.printD("model name is not in model_info")
-        return
+        # get model type
+        model_name = model_info["name"]
 
-    model_name = model_info["name"]
-    if not model_name:
-        util.printD("model name is Empty")
-        model_name = ""
+        # get version lists
+        model_versions = model_info["modelVersions"]
 
-    # get version list
-    if "modelVersions" not in model_info.keys():
-        util.printD("modelVersions is not in model_info")
-        return
-    
-    modelVersions = model_info["modelVersions"]
-    if not modelVersions:
-        util.printD("modelVersions is Empty")
-        return
-    
+    except (KeyError, ValueError, TypeError) as e:
+        util.printD(f"An error occurred while attempting to process model info: \n\t{e}")
+        return None
+
+    filenames = []
+    files = []
     version_strs = []
-    for version in modelVersions:
+    previews = {}
+    for version in model_versions:
         # version name can not be used as id
         # version id is not readable
         # so , we use name_id as version string
-        version_str = version["name"]+"_"+str(version["id"])
+        version_str = f'{version["name"]}_{version["id"]}'
+
+        filename = ""
+        try:
+            for filedata in version["files"]:
+                if filedata["type"] == "Model":
+                    filename = filedata["name"]
+
+        except (ValueError, KeyError):
+            pass
+
+        files.append(version["files"])
+        filenames.append(filename)
         version_strs.append(version_str)
+        previews[version_str] = version["images"]
 
     # get folder by model type
     folder = model.folders[model_type]
-    # get subfolders
-    subfolders = util.get_subfolders(folder)
-    if not subfolders:
-        subfolders = []
 
-    # add default root folder
-    subfolders.append("/")
-
-    util.printD("Get following info for downloading:")
-    util.printD(f"model_name:{model_name}")
-    util.printD(f"model_type:{model_type}")
-    util.printD(f"subfolders:{subfolders}")
-    util.printD(f"version_strs:{version_strs}")
+    # get subfolders
+    subfolders = ["/"] + util.get_subfolders(folder)
+
+    msg = util.indented_msg(f"""
+        Got following info for downloading:
+        {model_name=}
+        {model_type=}
+        {version_strs=}
+        {subfolders=}
+        {previews=}
+    """)
+    util.printD(msg)
+
+    return {
+        "model_info": model_info,
+        "model_name": model_name,
+        "model_type": model_type,
+        "files": files,
+        "filenames": filenames,
+        "subfolders": subfolders,
+        "version_strs": version_strs,
+        "previews": previews
+    }
 
-    return (model_info, model_name, model_type, subfolders, version_strs)
 
-# get version info by version string
 def get_ver_info_by_ver_str(version_str:str, model_info:dict) -> dict:
-    if not version_str:
-        util.printD("version_str is empty")
-        return
-    
-    if not model_info:
-        util.printD("model_info is None")
-        return
-    
-    # get version list
-    if "modelVersions" not in model_info.keys():
-        util.printD("modelVersions is not in model_info")
-        return
+    """
+    get version info by version string
+
+    return: version_info:dict
+    """
+
+    if not (version_str and model_info):
+        output = util.indented_msg(
+            f"""
+            Missing Parameter:
+            {model_info=}
+             {version_str=}
+            """
+            )
+        util.printD(output)
+        return None
 
-    modelVersions = model_info["modelVersions"]
-    if not modelVersions:
+    # get version list
+    model_versions = model_info.get("modelVersions", None)
+    if model_versions is None:
         util.printD("modelVersions is Empty")
-        return
-    
+        return None
+
     # find version by version_str
     version = None
-    for ver in modelVersions:
+    for ver in model_versions:
         # version name can not be used as id
         # version id is not readable
         # so , we use name_id as version string
-        ver_str = ver["name"]+"_"+str(ver["id"])
+        ver_str = f'{ver["name"]}_{ver["id"]}'
         if ver_str == version_str:
             # find version
             version = ver
+            break
+
+    if not (version and ("id" in version)):
+        util.printD(f"can not find version or id by version string: {version_str}")
+        return None
 
-    if not version:
-        util.printD("can not find version by version string: " + version_str)
-        return
-    
-    # get version id
-    if "id" not in version.keys():
-        util.printD("this version has no id")
-        return
-    
     return version
 
 
-# get download url from model info by version string
-# return - (version_id, download_url)
 def get_id_and_dl_url_by_version_str(version_str:str, model_info:dict) -> tuple:
-    if not version_str:
-        util.printD("version_str is empty")
-        return
-    
-    if not model_info:
-        util.printD("model_info is None")
-        return
-    
-    # get version list
-    if "modelVersions" not in model_info.keys():
-        util.printD("modelVersions is not in model_info")
-        return
+    """
+    get download url from model info by version string
+    return - (version_id, download_url)
+    """
+    if not (version_str and model_info):
+        output = util.indented_msg(f"""
+            Missing Parameter:
+            {model_info=}
+            {version_str=}
+        """)
+        util.printD(output)
+        return (False, output)
 
-    modelVersions = model_info["modelVersions"]
-    if not modelVersions:
+    # get version list
+    model_versions = model_info.get("modelVersions", None)
+    if model_versions is None:
         util.printD("modelVersions is Empty")
-        return
-    
+        return (False, output)
+
     # find version by version_str
     version = None
-    for ver in modelVersions:
+    for ver in model_versions:
         # version name can not be used as id
         # version id is not readable
         # so , we use name_id as version string
-        ver_str = ver["name"]+"_"+str(ver["id"])
+        ver_str = f'{ver["name"]}_{ver["id"]}'
         if ver_str == version_str:
             # find version
             version = ver
+            break
+
+    version_id = None
+    download_url = None
+    if version:
+        download_url = version.get("downloadUrl", None)
+        version_id = version.get("id", None)
+
+    if None in [version, version_id, download_url]:
+        output = util.indented_msg(f"""
+            Invalid Version Information:
+            {version=}
+            {version_id=}
+            {download_url=}
+        """)
+        util.printD(output)
+        return (False, output)
 
-    if not version:
-        util.printD("can not find version by version string: " + version_str)
-        return
-    
-    # get version id
-    if "id" not in version.keys():
-        util.printD("this version has no id")
-        return
-    
-    version_id = version["id"]
-    if not version_id:
-        util.printD("version id is Empty")
-        return
+    util.printD(f"Get Download Url: {download_url}")
 
-    # get download url
-    if "downloadUrl" not in version.keys():
-        util.printD("downloadUrl is not in this version")
-        return
-    
-    downloadUrl = version["downloadUrl"]
-    if not downloadUrl:
-        util.printD("downloadUrl is Empty")
-        return
-    
-    util.printD("Get Download Url: " + downloadUrl)
-    
-    return (version_id, downloadUrl)
-    
+    return (version_id, download_url)
 
-# download model from civitai by input
-# output to markdown log
-def dl_model_by_input(model_info:dict, model_type:str, subfolder_str:str, version_str:str, dl_all_bool:bool, max_size_preview:bool, skip_nsfw_preview:bool) -> str:
+def parse_file_info(file_info, basename):
+    """
+        returns data required to download a file from civitai
+    """
 
-    output = ""
+    download_url = file_info.get("downloadUrl", None)
+    if download_url is None:
+        return None
 
-    if not model_info:
-        output = "model_info is None"
-        util.printD(output)
-        return output
-    
-    if not model_type:
-        output = "model_type is None"
+    filetype = file_info["type"]
+    filename = file_info["name"]
+    if basename and not filetype == "VAE":
+        filename = f"{basename}.{filename.split('.')[-1]}"
+
+    return {
+        "url": download_url,
+        "filename": filename,
+        "type": filetype
+    }
+
+
+def download_files(filename, model_folder, ver_info, headers, filetypes, dl_all, duplicate):
+    """
+    get download urls from files info
+    some model versions have multiple files
+    """
+
+    version_id = ver_info["id"]
+
+    downloads = []
+
+    # check if this model already exists
+    result = civitai.search_local_model_info_by_version_id(model_folder, version_id)
+    if result:
+        output = "This model version already exists"
         util.printD(output)
-        return output
-    
-    if not subfolder_str:
-        output = "subfolder string is None"
+        yield (False, output)
+
+    for file_info in ver_info.get("files", {}):
+        if not dl_all:
+            if not file_info["type"] in filetypes:
+                continue
+
+        dl_info = parse_file_info(file_info, filename)
+
+        if dl_info:
+            downloads.append(dl_info)
+
+    if len(downloads) == 0:
+        dl_info = parse_file_info(ver_info, filename)
+        if dl_info:
+            downloads.append(dl_info)
+
+    # download
+    success = False
+    output = ""
+    filepath = None
+    total = len(downloads)
+    errors = []
+    errors_count = 0
+    snippet = None
+
+    for count, dl_info in enumerate(downloads):
+
+        url = dl_info["url"]
+        if errors_count > 0:
+            snippet = f"{errors_count}/{total} files failed"
+
+        dl_folder = model_folder
+        if dl_info["type"] == "VAE":
+            dl_folder = model.vae_folder
+
+        # webui visible progress bar
+        for result in downloader.dl_file(
+            url, filename=dl_info["filename"], folder=dl_folder, duplicate=duplicate,
+            headers=headers
+        ):
+            if not isinstance(result, str):
+                success, output = result
+                break
+
+            output = f"{result} | {count}/{total} files"
+            if snippet:
+                " | ".join([output, snippet])
+
+            yield output
+
+        if not success:
+            errors.append(downloader.error(url, output))
+            errors_count += 1
+            continue
+
+        if dl_info["type"] == "Model":
+            filepath = output
+
+    additional = None
+    if errors_count > 0:
+        additional = "\n\n".join(errors)
+
+        if errors_count == total:
+            yield (False, additional)
+            return
+
+    yield (True, filepath, additional)
+
+
+def download_one(filename, model_folder, ver_info, headers, duplicate):
+    """
+    only download one file
+    get download url
+    """
+
+    download_url = ver_info["downloadUrl"]
+
+    output = ""
+    if not download_url:
+        output = "Failed to find a download url"
         util.printD(output)
-        return output
-    
-    if not version_str:
-        output = "version_str is None"
+        yield (False, output)
+
+    # download
+    success = False
+    for result in downloader.dl_file(
+        download_url, filename=filename, folder=model_folder,
+        duplicate=duplicate, headers=headers
+    ):
+        if not isinstance(result, str):
+            success, output = result
+            break
+
+        yield result
+
+    if not success:
+        downloader.error(download_url, output)
+        yield (False, output)
+
+    yield (True, output)
+
+
+def dl_model_by_input(
+    ch_state:dict,
+    model_type:str,
+    subfolder_str:str,
+    version_str:str,
+    filename:str,
+    file_ext:str,
+    dl_all:bool,
+    nsfw_preview_threshold:bool,
+    duplicate:str,
+    preview:str,
+    *args
+) -> str:
+    """ download model from civitai by input
+        output to markdown log
+    """
+
+    model_info = ch_state["model_info"]
+    max_size_preview = util.get_opts("ch_max_size_preview")
+
+    if not (model_info and model_type and subfolder_str and version_str):
+        output = util.indented_msg(f"""
+            Missing Required Parameter in dl_model_by_input. Parameters given:
+            {model_type=}*
+            {subfolder_str=}*
+            {version_str=}*
+            {filename=}
+            {file_ext=}
+            {max_size_preview=}
+            {nsfw_preview_threshold=}
+            {duplicate=}
+        """)
+
+        # Keep model info away from util.indented_msg
+        # which can screw with complex strings
+        output = f"{output}\n    {model_info=}*\n    * Required"
         util.printD(output)
-        return output
-    
+        yield output
+        return
+
     # get model root folder
-    if model_type not in model.folders.keys():
-        output = "Unknow model type: "+model_type
+    if model_type not in model.folders:
+        output = f"Unsupported model type: {model_type}"
         util.printD(output)
-        return output
-    
+        yield output
+        return
+
+    folder = ""
+    subfolder = ""
+    output = ""
+    version_info = None
+
+    filetypes = []
+    for filetype, will_dl in zip(civitai.FILE_TYPES, args):
+        if will_dl:
+            filetypes.append(filetype)
+
     model_root_folder = model.folders[model_type]
 
+    if not os.path.exists(model_root_folder):
+        # Model directories may not exist by default
+        os.mkdir(model_root_folder)
 
     # get subfolder
-    subfolder = ""
-    if subfolder_str == "/" or subfolder_str == "\\":
+    if subfolder_str in ["/", "\\"]:
         subfolder = ""
-    elif subfolder_str[:1] == "/" or subfolder_str[:1] == "\\":
+    elif subfolder_str[:1] in ["/", "\\"]:
         subfolder = subfolder_str[1:]
     else:
         subfolder = subfolder_str
 
     # get model folder for downloading
-    model_folder = os.path.join(model_root_folder, subfolder)
-    if not os.path.isdir(model_folder):
-        output = "Model folder is not a dir: "+ model_folder
+    folder = os.path.join(model_root_folder, subfolder)
+    if not os.path.isdir(folder):
+        output = f"Model folder is not a dir: {folder}"
         util.printD(output)
-        return output
-    
+        yield output
+        return
+
     # get version info
     ver_info = get_ver_info_by_ver_str(version_str, model_info)
     if not ver_info:
-        output = "Fail to get version info, check console log for detail"
+        output = "Failed to get version info, check console log for detail"
         util.printD(output)
-        return output
-    
-    version_id = ver_info["id"]
-
+        yield output
+        return
 
-    if dl_all_bool:
-        # get all download url from files info
-        # some model versions have multiple files
-        download_urls = []
-        if "files" in ver_info.keys():
-            for file_info in ver_info["files"]:
-                if "downloadUrl" in file_info.keys():
-                    download_urls.append(file_info["downloadUrl"])
+    headers = {
+        "content-type": "application/json"
+    }
+    api_key = util.get_opts("ch_civiai_api_key")
+    if api_key:
+        headers["Authorization"] = f"Bearer {api_key}"
+
+    additional = None
+    for result in download_files(filename, folder, ver_info, headers, filetypes, dl_all, duplicate):
+        if not isinstance(result, str):
+            if len(result) > 2:
+                success, output, additional = result
+            else:
+                success, output = result
 
-        if not len(download_urls):
-            if "downloadUrl" in ver_info.keys():
-                download_urls.append(ver_info["downloadUrl"])
+            break
 
+        yield result
 
-        # check if this model is already existed
-        r = civitai.search_local_model_info_by_version_id(model_folder, version_id)
-        if r:
-            output = "This model version is already existed"
-            util.printD(output)
-            return output
-        
-        # download
-        filepath = ""
-        for url in download_urls:
-            model_filepath = downloader.dl(url, model_folder, None, None)
-            if not model_filepath:
-                output = "Downloading failed, check console log for detail"
-                util.printD(output)
-                return output
-            
-            if url == ver_info["downloadUrl"]:
-                filepath = model_filepath
-    else:
-        # only download one file
-        # get download url
-        url = ver_info["downloadUrl"]
-        if not url:
-            output = "Fail to get download url, check console log for detail"
-            util.printD(output)
-            return output
-        
-        # download
-        filepath = downloader.dl(url, model_folder, None, None)
-        if not filepath:
-            output = "Downloading failed, check console log for detail"
-            util.printD(output)
-            return output
-
+    if not success:
+        yield output
+        return
 
-    if not filepath:
-        filepath = model_filepath
-    
     # get version info
-    version_info = civitai.get_version_info_by_version_id(version_id)
-    if not version_info:
-        output = "Model downloaded, but failed to get version info, check console log for detail. Model saved to: " + filepath
-        util.printD(output)
-        return output
-
-    # write version info to file
-    base, ext = os.path.splitext(filepath)
-    info_file = base + civitai.suffix + model.info_ext
-    model.write_model_info(info_file, version_info)
-
-    # then, get preview image
-    civitai.get_preview_image_by_model_path(filepath, max_size_preview, skip_nsfw_preview)
-    
-    output = "Done. Model downloaded to: " + filepath
+    version_info = civitai.get_version_info_by_version_id(ver_info["id"])
+    model.process_model_info(output, version_info, model_type)
+
+    # then, get preview image + webui-visible progress
+    for result in civitai.get_preview_image_by_model_path(
+        output,
+        max_size_preview,
+        nsfw_preview_threshold,
+        preferred_preview=preview
+    ):
+        yield f"Downloading model preview:\n{result}"
+
+    output = f"Done. Downloaded to: {output}"
+    if additional:
+        output = f"{output}. Additionally, the following failures occurred: \n{additional}"
     util.printD(output)
-    return output
+    yield output

extensions/Stable-Diffusion-Webui-Civitai-Helper/scripts/ch_lib/msg_handler.py

@@ -1,64 +1,74 @@
-# -*- coding: UTF-8 -*-
-# handle msg between js and python side
+""" -*- coding: UTF-8 -*-
+handle msg between js and python side
+"""
 import json
 from . import util
 
 # action list
-js_actions = ("open_url", "add_trigger_words", "use_preview_prompt", "dl_model_new_version", "remove_card")
-py_actions = ("open_url", "remove_card")
+JS_ACTIONS = (
+    "open_url",
+    "add_trigger_words",
+    "use_preview_prompt",
+    "dl_model_new_version",
+    "rename_card",
+    "remove_card"
+)
+
+PY_ACTIONS = (
+    "open_url",
+    "rename_card",
+    "remove_card"
+)
 
 
-# handle request from javascript
-# parameter: msg - msg from js as string in a hidden textbox
-# return: dict for result
 def parse_js_msg(msg):
+    """
+    handle request from javascript
+    parameter: msg - msg from js as string in a hidden textbox
+    return: dict for result
+    """
     util.printD("Start parse js msg")
     msg_dict = json.loads(msg)
 
     # in case client side run JSON.stringify twice
-    if (type(msg_dict) == str):
+    if isinstance(msg_dict, str):
         msg_dict = json.loads(msg_dict)
 
-    if "action" not in msg_dict.keys():
-        util.printD("Can not find action from js request")
-        return
-
-    action = msg_dict["action"]
+    action = msg_dict.get("action", "")
     if not action:
-        util.printD("Action from js request is None")
-        return
+        util.printD("No action from js request")
+        return None
 
-    if action not in js_actions:
-        util.printD("Unknow action: " + action)
-        return
+    if action not in JS_ACTIONS:
+        util.printD(f"Unknown action: {action}")
+        return None
 
     util.printD("End parse js msg")
 
     return msg_dict
 
 
-# build python side msg for sending to js
-# parameter: content dict
-# return: msg as string, to fill into a hidden textbox
 def build_py_msg(action:str, content:dict):
+    """
+    build python side msg for sending to js
+    parameter: content dict
+    return: msg as string, to fill into a hidden textbox
+    """
     util.printD("Start build_msg")
-    if not content:
-        util.printD("Content is None")
-        return
-    
-    if not action:
-        util.printD("Action is None")
-        return
-
-    if action not in py_actions:
-        util.printD("Unknow action: " + action)
-        return
+    if not (content and action and action in PY_ACTIONS):
+        util.indented_msg(
+            f"""
+            Could not run action on content:
+            {action=}
+            {content=}
+            """
+        )
+        return None
 
     msg = {
         "action" : action,
         "content": content
     }
 
-
     util.printD("End build_msg")
-    return json.dumps(msg)
+    return json.dumps(msg)

extensions/Stable-Diffusion-Webui-Civitai-Helper/scripts/ch_lib/setting.py

@@ -1,103 +0,0 @@
-# -*- coding: UTF-8 -*-
-# collecting settings to here
-import json
-import os
-import modules.scripts as scripts
-from . import util
-
-
-name = "setting.json"
-path = os.path.join(scripts.basedir(), name)
-
-data = {
-    "model":{
-        "max_size_preview": True,
-        "skip_nsfw_preview": False
-    },
-    "general":{
-        "open_url_with_js": True,
-        "proxy": "",
-    },
-    "tool":{
-    }
-}
-
-
-
-# save setting
-# return output msg for log
-def save():
-    print("Saving setting to: " + path)
-
-    json_data = json.dumps(data, indent=4)
-
-    output = ""
-
-    #write to file
-    try:
-        with open(path, 'w') as f:
-            f.write(json_data)
-    except Exception as e:
-        util.printD("Error when writing file:"+path)
-        output = str(e)
-        util.printD(str(e))
-        return output
-
-    output = "Setting saved to: " + path
-    util.printD(output)
-
-    return output
-
-
-# load setting to global data
-def load():
-    # load data into globel data
-    global data
-
-    util.printD("Load setting from: " + path)
-
-    if not os.path.isfile(path):
-        util.printD("No setting file, use default")
-        return
-
-    json_data = None
-    with open(path, 'r') as f:
-        json_data = json.load(f)
-
-    # check error
-    if not json_data:
-        util.printD("load setting file failed")
-        return
-
-    data = json_data
-
-    # check for new key
-    if "proxy" not in data["general"].keys():
-        data["general"]["proxy"] = ""
-
-
-    return
-
-# save setting from parameter
-def save_from_input(max_size_preview, skip_nsfw_preview, open_url_with_js, proxy):
-    global data
-    data = {
-        "model":{
-            "max_size_preview": max_size_preview,
-            "skip_nsfw_preview": skip_nsfw_preview
-        },
-        "general":{
-            "open_url_with_js": open_url_with_js,
-            "proxy": proxy,
-        },
-        "tool":{
-        }
-    }
-
-    output = save()
-
-    if not output:
-        output = ""
-
-    return output
-

extensions/Stable-Diffusion-Webui-Civitai-Helper/scripts/ch_lib/templates.py

@@ -0,0 +1,94 @@
+"""
+HTML templates
+"""
+
+from string import Template
+from . import util
+
+article = Template(util.dedent("""
+    <article style="margin: 5px; clear: both;">
+        $thumbnail
+        <div style="font-size:16px">
+            File: $model_path
+        </div>
+        $download
+        <div style="font-size:20px;margin:6px 0px;">
+            <b>Model: <a href="$url" target="_blank"><u>$model_name</u></a></b>
+        </div>
+    </article>
+""").strip())
+
+thumbnail = Template(util.dedent("""
+    <img src='$img_url' style='float: left; margin: 5px;'>
+""").strip())
+
+description = Template(util.dedent("""
+    <blockquote style="font-size:16px;margin:6px 0px;">
+        $description
+    </blockquote>
+    <br>
+""").strip())
+
+# add js function to download new version into SD webui by python
+# in embed HTML, onclick= will also follow a ", never a ',
+# so have to write it as following
+download = Template(util.dedent("""
+    <div style="font-size:16px;margin:6px 0px;">
+        New Version: <u><a href="$download_url" target="_blank" style="margin:0px 10px;">$new_version_name</a></u>
+        <u><a href='#' style='margin:0px 10px;' onclick="ch_dl_model_new_version(event, '$model_path', '$new_version_id', '$download_url', '$model_type')">[Download into SD]</a></u>
+    </div>
+""").strip())
+
+no_download = Template(util.dedent("""
+    <div style="font-size:16px;margin:6px 0px;">
+        New Version: $new_version_name
+    </div>
+""").strip())
+
+duplicate_card = Template(util.dedent("""
+    <a href='#' onclick="remove_dup_card(event, '$model_type', '$search_term')" onmouseover="display_ch_path(event, '$path')" onmouseout="hide_ch_path(event)" onmousemove="move_ch_path(event)">
+        <div class='card' style=$style data-name="$name">
+            $background_image
+            <div class='actions'>
+                <div class='additional'>
+                    <span style="display:none">$search_term</span>
+                </div>
+                <span class='name'>$name</span>
+                <span class='description'>$description</span>
+            </div>
+        </div>
+    </a>
+""").strip())
+
+duplicate_preview = Template(util.dedent("""
+    <img src="./sd_extra_networks/thumb?filename=$bg_image" class="preview" loading="lazy">
+""").strip())
+
+duplicate_article = Template(util.dedent("""
+    <article>
+        <h1>
+            $section_name
+        </h1>
+        $contents
+    </article>
+""").strip())
+
+duplicate_row = Template(util.dedent("""
+    <div id="ch_$hash">
+        <div class=civitai_name>
+            <h2>
+                $civitai_name
+            </h2>
+        </div>
+        <div class=duplicate_model id="ch_${hash}_cards">
+            $columns
+        </div>
+    </div>
+""").strip())
+
+
+duplicate_column = Template(util.dedent("""
+        <div class=dup_$count>
+            $card
+        </div>
+""").strip())

extensions/Stable-Diffusion-Webui-Civitai-Helper/scripts/ch_lib/util.py

@@ -1,108 +1,493 @@
-# -*- coding: UTF-8 -*-
+""" -*- coding: UTF-8 -*-
+Utility functions for Stable Diffusion Civitai Helper
+"""
+from __future__ import annotations
 import os
 import io
+import re
 import hashlib
-import requests
-import shutil
+import textwrap
+import time
+import subprocess
+import gradio as gr
+from modules import shared
+from modules.shared import opts
+from modules import hashes
+import launch
+from packaging.version import parse as parse_version
 
+try:
+    # Automatic1111 SD WebUI
+    import modules.cache as sha256_cache
+except ModuleNotFoundError:
+    # Vladmandic "SD.Next"
+    import modules.hashes as sha256_cache
 
-version = "1.9.1"
+# used to append extension information to JSON/INFO files
+SHORT_NAME = "sd_civitai_helper"
 
-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',
-               "Authorization": ""}
+# current version of the exension
+VERSION = "1.8.1"
 
+# Civitai INFO files below this version will regenerated
+COMPAT_VERSION_CIVITAI = "1.7.2"
 
-proxies = None
+# SD webui model info JSON below this version will be regenerated
+COMPAT_VERSION_SDWEBUI = "1.8.0"
 
-civitai_api_key = ""
+DEFAULT_HEADERS = {
+    "User-Agent": (
+        "Mozilla/5.0 (iPad; CPU OS 12_2 like Mac OS X) "
+        "AppleWebKit/605.1.15 (KHTML, like Gecko) Mobile/15E148"
+    )
+}
+
+PROXIES = {
+    "http": None,
+    "https": None,
+}
+
+REQUEST_TIMEOUT = 300  # 5 minutes
+REQUEST_RETRIES = 5
+
+_MINUTE = 60
+_HOUR = _MINUTE * 60
+_DAY = _HOUR * 24
 
 
 # print for debugging
-def printD(msg):
+def printD(msg:any) -> str:
+    """ Print a message to stderr """
     print(f"Civitai Helper: {msg}")
 
 
-def read_chunks(file, size=io.DEFAULT_BUFFER_SIZE):
-    """Yield pieces of data from a file-like object until EOF."""
+def append_default_headers(headers:dict) -> dict:
+    """ Append extension default values to customized headers where missing """
+
+    for key, val in DEFAULT_HEADERS.items():
+        if key not in headers:
+            headers[key] = val
+    return headers
+
+
+def indented_msg(msg:str) -> str:
+    """
+    Clean up an indented message in the format of
+    [header]
+    var1=var1
+    var2=var2
+    var3=var3
+
+    and print the results in the format of:
+
+    Civitai Helper: [header]
+        var1: var1
+        var2: var2
+        var3: var3
+
+    return: msg:str
+    """
+
+    msg_parts = textwrap.dedent(
+        msg
+    ).strip().split('\n')
+    msg = [msg_parts.pop(0)]
+    for part in msg_parts:
+        part = ": ".join(part.split("="))
+        msg.append(f"   {part}")
+    msg = "\n".join(msg)
+
+    return msg
+
+
+def delay(seconds:float) -> None:
+    """ delay before next request, mostly to prevent being treated as a DDoS """
+    time.sleep(seconds)
+
+
+def is_stale(timestamp:float) -> bool:
+    """ Returns if a timestamp was more than a day ago. """
+    cur_time = ch_time()
+    elapsed = cur_time - timestamp
+
+    if elapsed > _DAY:
+        return True
+
+    return False
+
+
+def info(msg:str) -> None:
+    """ Display an info smessage on the client DOM """
+    gr.Info(msg)
+
+
+def warning(msg:str) -> None:
+    """ Display a warning message on the client DOM """
+    gr.Warning(msg)
+
+
+def error(msg:str) -> None:
+    """ Display an error message on the client DOM """
+    gr.Error(msg)
+
+
+def ch_time() -> int:
+    """ Unix timestamp """
+    return int(time.time())
+
+
+def dedent(text:str) -> str:
+    """ alias for textwrap.dedent """
+    return textwrap.dedent(text)
+
+
+def get_name(model_path:str, model_type:str) -> str:
+    """ return: {model_type}/{model_name} """
+
+    model_name = os.path.splitext(os.path.basename(model_path))[0]
+    return f"{model_type}/{model_name}"
+
+
+def get_opts(key):
+    """ return: option value """
+    return opts.data.get(key, None)
+
+
+def gen_file_sha256(filename:str, model_type="lora", use_addnet_hash=False) -> str:
+    """ return a sha256 hash for a file """
+
+    cache = sha256_cache.cache
+    dump_cache = sha256_cache.dump_cache
+    model_name = get_name(filename, model_type)
+
+    sha256_hashes = None
+    if use_addnet_hash:
+        sha256_hashes = cache("hashes-addnet")
+    else:
+        sha256_hashes = cache("hashes")
+
+    sha256_value = hashes.sha256_from_cache(filename, model_name, use_addnet_hash)
+    if sha256_value is not None:
+        yield sha256_value
+        return
+
+    if shared.cmd_opts.no_hashing:
+        printD("SD WebUI Civitai Helper requires hashing functions for this feature. \
+            Please remove the commandline argument `--no-hashing` for this functionality.")
+        yield None
+        return
+
+    with open(filename, "rb") as model_file:
+        result = None
+        for result in calculate_sha256(model_file, use_addnet_hash):
+            yield result
+        sha256_value = result
+
+    printD(f"sha256: {sha256_value}")
+
+    sha256_hashes[model_name] = {
+        "mtime": os.path.getmtime(filename),
+        "sha256": sha256_value,
+    }
+
+    dump_cache()
+
+    yield sha256_value
+
+def calculate_sha256(model_file, use_addnet_hash=False):
+    """ calculate the sha256 hash for a model file """
+
+    blocksize= 1 << 20
+    sha256_hash = hashlib.sha256()
+
+    size = os.fstat(model_file.fileno()).st_size
+
+    offset = 0
+    if use_addnet_hash:
+        model_file.seek(0)
+        header= model_file.read(8)
+        offset = int.from_bytes(header, "little") + 8
+        model_file.seek(offset)
+
+    pos = 0
+    for block in read_chunks(model_file, size=blocksize):
+        pos += len(block)
+
+        percent = (pos - offset) / (size - offset)
+
+        yield (percent, f"hashing model {model_file.name}")
+
+        sha256_hash.update(block)
+
+    hash_value =  sha256_hash.hexdigest()
+    yield hash_value
+
+
+def read_chunks(file, size=io.DEFAULT_BUFFER_SIZE) -> bytes:
+    """ Yield pieces of data from a file-like object until EOF. """
     while True:
         chunk = file.read(size)
         if not chunk:
             break
         yield chunk
 
-# Now, hashing use the same way as pip's source code.
-def gen_file_sha256(filname):
-    printD("Use Memory Optimized SHA256")
-    blocksize=1 << 20
-    h = hashlib.sha256()
-    length = 0
-    with open(os.path.realpath(filname), 'rb') as f:
-        for block in read_chunks(f, size=blocksize):
-            length += len(block)
-            h.update(block)
-
-    hash_value =  h.hexdigest()
-    printD("sha256: " + hash_value)
-    printD("length: " + str(length))
-    return hash_value
-
-
-
-# get preview image
-def download_file(url, path):
-    printD("Downloading file from: " + url)
-    # get file
-    r = requests.get(url, stream=True, headers=def_headers, proxies=proxies)
-    if not r.ok:
-        printD("Get error code: " + str(r.status_code))
-        printD(r.text)
-        return
-    
-    # write to file
-    with open(os.path.realpath(path), 'wb') as f:
-        r.raw.decode_content = True
-        shutil.copyfileobj(r.raw, f)
-
-    printD("File downloaded to: " + path)
 
-# get subfolder list
-def get_subfolders(folder:str) -> list:
-    printD("Get subfolder for: " + folder)
+def get_subfolders(folder:str) -> list[str]:
+    """ return: list of subfolders """
+    printD(f"Get subfolder for: {folder}")
     if not folder:
         printD("folder can not be None")
-        return
-    
+        return []
+
     if not os.path.isdir(folder):
         printD("path is not a folder")
-        return
-    
+        return []
+
     prefix_len = len(folder)
+    full_dirs_searched = []
     subfolders = []
-    for root, dirs, files in os.walk(folder, followlinks=True):
-        for dir in dirs:
-            full_dir_path = os.path.join(root, dir)
-            # get subfolder path from it
-            subfolder = full_dir_path[prefix_len:]
-            subfolders.append(subfolder)
+    for root, dirs, _ in os.walk(folder, followlinks=True):
+        if root == folder:
+            continue
+
+        # Prevent following recursive symlinks
+        follow = []
+        for directory in dirs:
+            full_dir_path = os.path.join(root, directory)
+            try:
+                canonical_dir = os.path.realpath(full_dir_path, strict=True)
+                if canonical_dir not in full_dirs_searched:
+                    full_dirs_searched.append(canonical_dir)
+                    follow.append(directory)
+
+            except OSError:
+                printD(f"Symlink loop: {directory}")
+                continue
+
+        # Get subfolder path
+        subfolder = root[prefix_len:]
+        subfolders.append(subfolder)
+
+        # Update dirs parameter to prevent following recursive symlinks
+        dirs[:] = follow
 
     return subfolders
 
 
+def find_file_in_folders(folders:list, filename:str) -> str:
+    """
+    Searches a directory for a filename,
+
+    return: filename:str or None
+    """
+    for folder in folders:
+        for root, _, files in os.walk(folder, followlinks=True):
+            if filename in files:
+                # found file
+                return os.path.join(root, filename)
+
+    return None
+
+
 # get relative path
 def get_relative_path(item_path:str, parent_path:str) -> str:
-    # printD("item_path:"+item_path)
-    # printD("parent_path:"+parent_path)
-    # item path must start with parent_path
-    if not item_path:
-        return ""
-    if not parent_path:
+    """
+    Gets a relative path from an absolute path and its parent_path
+    item path must start with parent_path
+    return: relative_path:str
+    """
+
+    if not (item_path and parent_path):
         return ""
+
     if not item_path.startswith(parent_path):
+        # return absolute path
         return item_path
 
     relative = item_path[len(parent_path):]
     if relative[:1] == "/" or relative[:1] == "\\":
         relative = relative[1:]
 
-    # printD("relative:"+relative)
-    return relative
+    # printD(f"relative: {relative}")
+    return relative
+
+
+# Allowed HTML tags
+whitelist = re.compile(r"</?(a|img|br|p|b|strong|i|h[0-9]|code)[^>]*>")
+
+# Allowed HTML attributes
+attrs = re.compile(r"""(?:href|src|target)=['"]?[^\s'"]*['"]?""")
+
+def safe_html_replace(match) -> str:
+    """ Given a block of text, returns that block with most HTML removed
+        and unneeded attributes pruned.
+    """
+    tag = None
+    attr = None
+    close = False
+
+    match = whitelist.match(match.group(0))
+    if match is not None:
+        html_elem = match.group(0)
+        tag = match.group(1)
+        close = html_elem[1] == "/"
+        if (tag in ["a", "img"]) and not close:
+            sub_match = attrs.findall(html_elem)
+            if sub_match is not None:
+                attr = " ".join(sub_match)
+
+        if close:
+            return f"</{tag}>"
+
+        return f"<{tag} {attr}>" if attr else f"<{tag}>"
+
+    return ""
+
+def safe_html(html:str) -> str:
+    """ whitelist only HTML I"m comfortable displaying in webui """
+
+    return re.sub("<[^<]+?>", safe_html_replace, html)
+
+
+def trim_html(html:str) -> str:
+    """ Remove any HTML for a given string and, if needed, replace it with
+        a comparable plain-text alternative.
+    """
+
+    def sub_tag(match):
+        tag = match.group(1)
+        if tag == "/p":
+            return "\n\n"
+        if tag == "br":
+            return "\n"
+        if tag == "li":
+            return "* "
+        if tag in ["code", "/code"]:
+            return "`"
+        return ''
+
+    def sub_escaped(match):
+        escaped = match.group(1)
+        unescaped = {
+            "gt": ">",
+            "lt": "<",
+            "quot": '"',
+            "amp": "&"
+        }
+        return unescaped.get(escaped, "")
+
+    # non-breaking space. Useless unstyled content
+    html = html.replace("\u00a0", "")
+
+    # remove non-whitelisted HTML tags,
+    # replace whitelisted tags with text-equivalents
+    html = re.sub(r"<(/?[a-zA-Z]+)(?:[^>]+)?>", sub_tag, html)
+
+    # Replace HTML-escaped characters with displayables.
+    html = re.sub(r"\&(gt|lt|quot|amp)\;", sub_escaped, html)
+
+    # Because we encapsulate the description in HTML comment,
+    # We have to prevent those comments from being cancelled.
+    html.replace("-->", "→")
+
+    # https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/13241
+    return f"<!--\n{html.strip()}\n-->"
+
+
+def newer_version(ver1:str, ver2:str, allow_equal=False) -> bool:
+    """ Returns ver1 > ver2
+        if allow_equal, returns ver1 >= ver2
+    """
+    if allow_equal:
+        return parse_version(ver1) >= parse_version(ver2)
+
+    return parse_version(ver1) > parse_version(ver2)
+
+
+def metadata_version(metadata:dict) -> str | bool:
+    """ Attempts retrieve the extension version used to create
+        to create the object block
+    """
+    try:
+        return metadata["extensions"][SHORT_NAME]["version"]
+    except KeyError:
+        return False
+
+
+def create_extension_block(data=None, skeleton=False) -> dict:
+    """ Creates or edits an extensions block for usage in JSON files
+        created or edited by this extension.
+
+        Adds the current version of this extension to the extensions block
+    """
+
+    cur_time = ch_time()
+
+    block = {
+        SHORT_NAME: {
+            "version": VERSION,
+            "last_update": cur_time,
+            "skeleton_file": skeleton
+        }
+    }
+
+    if not data:
+        return block
+
+    data[SHORT_NAME] = block[SHORT_NAME]
+
+    return data
+
+
+def webui_version() -> str:
+    ''' Gets the current webui version using webui's launch tools
+
+        The version is expected to be in the format `v1.6.0-128-g792589fd`,
+        tho all that is explicitly required is `vX`.
+
+        returns the version in the form 'X.Y.Z'
+    '''
+    version = None
+    try:
+        tag = launch.git_tag()
+        match = re.match(r"v([\d.]+)", tag)
+        if match:
+            version = match.group(1)
+        else:
+            # XXX assume a modern SD Webui version if one cannot be found.
+            version = "1.6.0"
+
+    except AttributeError:
+        try:
+            return subprocess.check_output(
+                ["git", "describe", "--tags"],
+                shell=False,
+                encoding='utf8'
+            ).strip()
+
+        except subprocess.SubprocessError:
+            try:
+                changelog_md = os.path.join(
+                    os.path.dirname(os.path.dirname(__file__)),
+                    "CHANGELOG.md"
+                )
+                with open(changelog_md, "r", encoding="utf-8") as file:
+                    line = next((line.strip() for line in file if line.strip()), "<none>")
+                    line = line.replace("## ", "")
+                    version = line
+
+            except OSError:
+                version = "1.6.0"
+
+    return version
+
+
+filename_re = re.compile(r"[^A-Za-z\d\s\^\-\+_.\(\)\[\]]")
+def bash_filename(filename:str) -> str:
+    """
+    Bashes a filename with a large fish until I'm comfortable using it.
+
+    Keeps a limited set of valid characters, but does not account for
+    reserved names like COM.
+    """
+    return re.sub(filename_re, "", filename)

extensions/Stable-Diffusion-Webui-Civitai-Helper/scripts/civitai_helper.py

@@ -1,59 +1,68 @@
-# -*- coding: UTF-8 -*-
-# This extension can help you manage your models from civitai. It can download preview, add trigger words, open model page and use the prompt from preview image
-# repo: https://github.com/butaixianran/
+""" -*- coding: UTF-8 -*-
+This extension can help you manage your models from civitai.
+ It can download preview, add trigger words, open model page and use the prompt from preview image
+repo: https://github.com/butaixianran/
+"""
 
-
-
-import modules.scripts as scripts
-import gradio as gr
 import os
-import webbrowser
-import requests
-import random
-import hashlib
-import json
-import shutil
-import re
+import gradio as gr
 import modules
-from modules import script_callbacks
+from modules import scripts
 from modules import shared
+from modules import script_callbacks
 from scripts.ch_lib import model
 from scripts.ch_lib import js_action_civitai
-from scripts.ch_lib import model_action_civitai
 from scripts.ch_lib import civitai
 from scripts.ch_lib import util
-
+from scripts import sections
 
 # init
-
 # root path
-root_path = os.getcwd()
+ROOT_PATH = os.getcwd()
 
 # extension path
-extension_path = scripts.basedir()
+EXTENSION_PATH = scripts.basedir()
+
+# default hidden values for civitai helper buttons
+BUTTONS = {
+    "replace_preview_button": False,
+    "open_url_button": False,
+    "add_trigger_words_button": False,
+    "add_preview_prompt_button": False,
+    "rename_model_button": False,
+    "remove_model_button": False,
+}
 
 model.get_custom_model_folder()
 
+def update_proxy():
+    """ Set proxy, allow for changes at runtime """
+    proxy = util.get_opts("ch_proxy")
+
+    util.printD(f"Set Proxy: {proxy}")
+    if proxy:
+        util.PROXIES["http"] = proxy
+        util.PROXIES["https"] = proxy
+        return
+
+    util.PROXIES["http"] = None
+    util.PROXIES["https"] = None
 
-# Setting now can not be saved from extension tab
-# All settings now must be saved from setting page.
-def on_ui_settings():
-    ch_section = ("civitai_helper", "Civitai Helper")
-    # settings
-    shared.opts.add_option("ch_max_size_preview", shared.OptionInfo(True, "Download Max Size Preview", gr.Checkbox, {"interactive": True}, section=ch_section))
-    shared.opts.add_option("ch_skip_nsfw_preview", shared.OptionInfo(False, "Skip NSFW Preview Images", gr.Checkbox, {"interactive": True}, section=ch_section))
-    shared.opts.add_option("ch_open_url_with_js", shared.OptionInfo(True, "Open Url At Client Side", gr.Checkbox, {"interactive": True}, section=ch_section))
-    shared.opts.add_option("ch_proxy", shared.OptionInfo("", "Civitai Helper Proxy", gr.Textbox, {"interactive": True, "lines":1, "info":"format: socks5h://127.0.0.1:port"}, section=ch_section))
-    shared.opts.add_option("ch_civiai_api_key", shared.OptionInfo("", "Civitai API Key", gr.Textbox, {"interactive": True, "lines":1, "info":"check doc:https://github.com/zixaphir/Stable-Diffusion-Webui-Civitai-Helper/tree/master#api-key"}, section=ch_section))
 
 def on_ui_tabs():
     # init
     # init_py_msg = {
     #     # relative extension path
-    #     "extension_path": util.get_relative_path(extension_path, root_path),
+    #     "EXTENSION_PATH": util.get_relative_path(EXTENSION_PATH, ROOT_PATH),
     # }
     # init_py_msg_str = json.dumps(init_py_msg)
 
+    try:
+        BUTTONS["add_trigger_words_button"] = util.newer_version(
+            util.webui_version(), '1.5.0', allow_equal=True
+        )
+    except ValueError: # packaging.version.InvalidVersion
+        BUTTONS["add_trigger_words_button"] = False
 
     # get prompt textarea
     # check modules/ui.py, search for txt2img_paste_fields
@@ -63,190 +72,245 @@ def on_ui_tabs():
     img2img_prompt = modules.ui.img2img_paste_fields[0][0]
     img2img_neg_prompt = modules.ui.img2img_paste_fields[1][0]
 
-
-    # get settings
-    max_size_preview = shared.opts.data.get("ch_max_size_preview", True)
-    skip_nsfw_preview = shared.opts.data.get("ch_skip_nsfw_preview", False)
-    open_url_with_js = shared.opts.data.get("ch_open_url_with_js", True)
-    proxy = shared.opts.data.get("ch_proxy", "")
-    civitai_api_key = shared.opts.data.get("ch_civiai_api_key", "")
-
-    util.printD("Settings:")
-    util.printD("max_size_preview: " + str(max_size_preview))
-    util.printD("skip_nsfw_preview: " + str(skip_nsfw_preview))
-    util.printD("open_url_with_js: " + str(open_url_with_js))
-    util.printD("proxy: " + str(proxy))
-
-    # set civitai_api_key
-    has_api_key = False
-    if civitai_api_key:
-        has_api_key = True
-        util.civitai_api_key = civitai_api_key
-        util.def_headers["Authorization"] = f"Bearer {civitai_api_key}"
-
-    util.printD(f"use civitai api key: {has_api_key}")
-
-    # set proxy
-    if proxy:
-        util.proxies = {
-            "http": proxy,
-            "https": proxy,
-        }
-
-
-    # ====Event's function====
-    def scan_model(scan_model_types):
-        return model_action_civitai.scan_model(scan_model_types, max_size_preview, skip_nsfw_preview)
-    
-    def get_model_info_by_input(model_type_drop, model_name_drop, model_url_or_id_txtbox):
-        return model_action_civitai.get_model_info_by_input(model_type_drop, model_name_drop, model_url_or_id_txtbox, max_size_preview, skip_nsfw_preview)
-
-    def dl_model_by_input(dl_model_info, dl_model_type_txtbox, dl_subfolder_drop, dl_version_drop, dl_all_ckb):
-        return model_action_civitai.dl_model_by_input(dl_model_info, dl_model_type_txtbox, dl_subfolder_drop, dl_version_drop, dl_all_ckb, max_size_preview, skip_nsfw_preview)
-
-    def check_models_new_version_to_md(dl_model_info, dl_model_type_txtbox, dl_subfolder_drop, dl_version_drop, dl_all_ckb):
-        return model_action_civitai.check_models_new_version_to_md(dl_model_info, dl_model_type_txtbox, dl_subfolder_drop, dl_version_drop, dl_all_ckb, max_size_preview, skip_nsfw_preview)
-
-    def open_model_url(js_msg_txtbox):
-        return js_action_civitai.open_model_url(js_msg_txtbox, open_url_with_js)
-
-    def dl_model_new_version(js_msg_txtbox, max_size_preview):
-        return js_action_civitai.dl_model_new_version(js_msg_txtbox, max_size_preview, skip_nsfw_preview)
-
-
-    def get_model_names_by_input(model_type, empty_info_only):
-        names = civitai.get_model_names_by_input(model_type, empty_info_only)
-        return model_name_drop.update(choices=names)
-
-    def get_model_info_by_url(url):
-        r = model_action_civitai.get_model_info_by_url(url)
-
-        model_info = {}
-        model_name = ""
-        model_type = ""
-        subfolders = []
-        version_strs = []
-        if r:
-            model_info, model_name, model_type, subfolders, version_strs = r
-
-        return [model_info, model_name, model_type, dl_subfolder_drop.update(choices=subfolders), dl_version_drop.update(choices=version_strs)]
+    # Used by some elements to pass messages to python
+    js_msg_txtbox = gr.Textbox(
+        label="Request Msg From Js",
+        visible=False,
+        lines=1,
+        value="",
+        elem_id="ch_js_msg_txtbox"
+    )
 
     # ====UI====
-    with gr.Blocks(analytics_enabled=False) as civitai_helper:
-
-        model_types = list(model.folders.keys())
-        no_info_model_names = civitai.get_model_names_by_input("ckp", False)
-
-        # session data
-        dl_model_info = gr.State({})
-
-
+    with gr.Blocks(
+        analytics_enabled=False
+    ) as civitai_helper:
+    # with gr.Blocks(css=".block.padded {padding: 10px !important}") as civitai_helper:
 
+        # init
         with gr.Box(elem_classes="ch_box"):
-            with gr.Column():
-                gr.Markdown("### Scan Models for Civitai")
-                with gr.Row():
-                    scan_model_types_ckbg = gr.CheckboxGroup(choices=model_types, label="Model Types", value=model_types)
+            sections.scan_models_section()
 
-                # with gr.Row():
-                scan_model_civitai_btn = gr.Button(value="Scan", variant="primary", elem_id="ch_scan_model_civitai_btn")
-                # with gr.Row():
-                scan_model_log_md = gr.Markdown(value="Scanning takes time, just wait. Check console log for detail", elem_id="ch_scan_model_log_md")
-
-        
         with gr.Box(elem_classes="ch_box"):
-            with gr.Column():
-                gr.Markdown("### Get Model Info from Civitai by URL")
-                gr.Markdown("Use this when scanning can not find a local model on civitai")
-                with gr.Row():
-                    model_type_drop = gr.Dropdown(choices=model_types, label="Model Type", value="ckp", multiselect=False)
-                    empty_info_only_ckb = gr.Checkbox(label="Only Show Models have no Info", value=False, elem_id="ch_empty_info_only_ckb", elem_classes="ch_vpadding")
-                    model_name_drop = gr.Dropdown(choices=no_info_model_names, label="Model", value="ckp", multiselect=False)
-
-                model_url_or_id_txtbox = gr.Textbox(label="Civitai URL", lines=1, value="")
-                get_civitai_model_info_by_id_btn = gr.Button(value="Get Model Info from Civitai", variant="primary")
-                get_model_by_id_log_md = gr.Markdown("")
+            sections.get_model_info_by_url_section()
 
         with gr.Box(elem_classes="ch_box"):
-            with gr.Column():
-                gr.Markdown("### Download Model")
-                with gr.Row():
-                    dl_model_url_or_id_txtbox = gr.Textbox(label="Civitai URL", lines=1, value="")
-                    dl_model_info_btn = gr.Button(value="1. Get Model Info by Civitai Url", variant="primary")
-
-                gr.Markdown(value="2. Pick Subfolder and Model Version")
-                with gr.Row():
-                    dl_model_name_txtbox = gr.Textbox(label="Model Name", interactive=False, lines=1, value="")
-                    dl_model_type_txtbox = gr.Textbox(label="Model Type", interactive=False, lines=1, value="")
-                    dl_subfolder_drop = gr.Dropdown(choices=[], label="Sub-folder", value="", interactive=True, multiselect=False)
-                    dl_version_drop = gr.Dropdown(choices=[], label="Model Version", value="", interactive=True, multiselect=False)
-                    dl_all_ckb = gr.Checkbox(label="Download All files", value=False, elem_id="ch_dl_all_ckb", elem_classes="ch_vpadding")
-                
-                dl_civitai_model_by_id_btn = gr.Button(value="3. Download Model", variant="primary")
-                dl_log_md = gr.Markdown(value="Check Console log for Downloading Status")
+            sections.download_section()
 
         with gr.Box(elem_classes="ch_box"):
-            with gr.Column():
-                gr.Markdown("### Check models' new version")
-                with gr.Row():
-                    model_types_ckbg = gr.CheckboxGroup(choices=model_types, label="Model Types", value=["lora"])
-                    check_models_new_version_btn = gr.Button(value="Check New Version from Civitai", variant="primary")
-
-                check_models_new_version_log_md = gr.HTML("It takes time, just wait. Check console log for detail")
+            sections.scan_for_duplicates_section()
 
         with gr.Box(elem_classes="ch_box"):
-            with gr.Column():
-                gr.Markdown("### Other")
-                # save_setting_btn = gr.Button(value="Save Setting")
-                gr.Markdown(value="Settings are moved into Settings Tab->Civitai Helper section")
-
+            sections.check_new_versions_section(js_msg_txtbox)
 
         # ====Footer====
-        gr.Markdown(f"<center>version:{util.version}</center>")
+        gr.HTML(f"<center>{util.SHORT_NAME} version: {util.VERSION}</center>")
 
         # ====hidden component for js, not in any tab====
-        js_msg_txtbox = gr.Textbox(label="Request Msg From Js", visible=False, lines=1, value="", elem_id="ch_js_msg_txtbox")
-        py_msg_txtbox = gr.Textbox(label="Response Msg From Python", visible=False, lines=1, value="", elem_id="ch_py_msg_txtbox")
-
-        js_open_url_btn = gr.Button(value="Open Model Url", visible=False, elem_id="ch_js_open_url_btn")
-        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")
+        js_msg_txtbox.render()
+        py_msg_txtbox = gr.Textbox(
+            label="Response Msg From Python",
+            visible=False,
+            lines=1,
+            value="",
+            elem_id="ch_py_msg_txtbox"
+        )
+
+        js_open_url_btn = gr.Button(
+            value="Open Model Url",
+            visible=False,
+            elem_id="ch_js_open_url_btn"
+        )
+        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_rename_card_btn = gr.Button(
+            value="Rename Card",
+            visible=False,
+            elem_id="ch_js_rename_card_btn"
+        )
+        js_remove_card_btn = gr.Button(
+            value="Remove Card",
+            visible=False,
+            elem_id="ch_js_remove_card_btn"
+        )
 
         # ====events====
-        # Scan Models for Civitai
-        scan_model_civitai_btn.click(scan_model, inputs=[scan_model_types_ckbg], outputs=scan_model_log_md)
-
-        # Get Civitai Model Info by Model Page URL
-        model_type_drop.change(get_model_names_by_input, inputs=[model_type_drop, empty_info_only_ckb], outputs=model_name_drop)
-        empty_info_only_ckb.change(get_model_names_by_input, inputs=[model_type_drop, empty_info_only_ckb], outputs=model_name_drop)
-
-        get_civitai_model_info_by_id_btn.click(get_model_info_by_input, inputs=[model_type_drop, model_name_drop, model_url_or_id_txtbox], outputs=get_model_by_id_log_md)
-
-        # Download Model
-        dl_model_info_btn.click(get_model_info_by_url, inputs=dl_model_url_or_id_txtbox, outputs=[dl_model_info, dl_model_name_txtbox, dl_model_type_txtbox, dl_subfolder_drop, dl_version_drop])
-        dl_civitai_model_by_id_btn.click(dl_model_by_input, inputs=[dl_model_info, dl_model_type_txtbox, dl_subfolder_drop, dl_version_drop, dl_all_ckb], outputs=dl_log_md)
-
-        # Check models' new version
-        check_models_new_version_btn.click(model_action_civitai.check_models_new_version_to_md, inputs=model_types_ckbg, outputs=check_models_new_version_log_md)
-
         # js action
-        js_open_url_btn.click(open_model_url, inputs=[js_msg_txtbox], 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_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)
+        js_open_url_btn.click(
+            js_action_civitai.open_model_url,
+            inputs=[js_msg_txtbox],
+            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_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_rename_card_btn.click(
+            js_action_civitai.rename_model_by_path,
+            inputs=[js_msg_txtbox],
+            outputs=py_msg_txtbox
+        )
+        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"),
-
+    return ((civitai_helper, "Civitai Helper", "civitai_helper"),)
 
 
+def on_ui_settings():
+    section = ('civitai_helper', "Civitai Helper")
+    shared.opts.add_option(
+        "ch_civiai_api_key",
+        shared.OptionInfo(
+            "",
+            (
+                "API key for authenticating with Civitai. "
+                "This is required to download some models. "
+                "See Wiki for more details."
+            ),
+            gr.Textbox,
+            {"interactive": True, "max_lines": 1},
+            section=section
+        ).link(
+            "Wiki",
+            "https://github.com/zixaphir/Stable-Diffusion-Webui-Civitai-Helper/wiki/Civitai-API-Key"
+        )
+    )
+    shared.opts.add_option(
+        "ch_dl_lyco_to_lora",
+        shared.OptionInfo(
+            False,
+            (
+                "Save LyCORIS models to Lora directory. Do not use this if you are on "
+                "older versions of webui or you use an extension that handles LyCORIS "
+                "models."
+            ),
+            gr.Checkbox,
+            {"interactive": True},
+            section=section)
+    )
+    shared.opts.add_option(
+        "ch_open_url_with_js",
+        shared.OptionInfo(
+            True,
+            (
+                "Open model Url on the user's client side, rather than server side. "
+                "If you are running WebUI locally, disabling this may open URLs in your "
+                "default internet browser if it is different than the one you are running "
+                "WebUI in"
+            ),
+            gr.Checkbox,
+            {"interactive": True},
+            section=section
+        )
+    )
+    shared.opts.add_option(
+        "ch_hide_buttons",
+        shared.OptionInfo(
+           [x for x, y in BUTTONS.items() if y],
+           "Hide checked Civitai Helper buttons on model cards",
+           gr.CheckboxGroup,
+           {"choices": list(BUTTONS)},
+           section=section
+        )
+   )
+    shared.opts.add_option(
+        "ch_always_display",
+        shared.OptionInfo(
+            False,
+            "Always Display Buttons on model cards",
+            gr.Checkbox,
+            {"interactive": True},
+            section=section
+        )
+    )
+    shared.opts.add_option(
+        "ch_show_btn_on_thumb",
+        shared.OptionInfo(
+            True,
+            "Show Button On Thumb Mode in SD webui versions before 1.5.0",
+            gr.Checkbox,
+            {"interactive": True},
+            section=section
+        )
+    )
+    shared.opts.add_option(
+        "ch_max_size_preview",
+        shared.OptionInfo(
+            True,
+            "Download Max Size Preview",
+            gr.Checkbox,
+            {"interactive": True},
+            section=section
+        )
+    )
+    shared.opts.add_option(
+        "ch_nsfw_preview_threshold",
+        shared.OptionInfo(
+            civitai.NSFW_LEVELS[-1], # Allow all
+            util.dedent(
+                """
+                Block NSFW images of a certain threshold and higher.
+                Civitai marks all images for NSFW models as also being NSFW.
+                These ratings do not seem to be explicitly defined on Civitai's
+                end, but "Soft" seems to be suggestive, with NSFW elements but
+                not explicit nudity, "Mature" seems to include nudity but not
+                always, and "X" seems to be explicitly adult content.
+                """
+            ).strip().replace("\n", " "),
+            gr.Dropdown,
+            {
+                "choices": civitai.NSFW_LEVELS[1:],
+                "interactive": True
+            },
+            section=section
+        )
+    )
+    shared.opts.add_option(
+        "ch_dl_webui_metadata",
+        shared.OptionInfo(
+            True,
+            "Also add data for WebUI metadata editor",
+            gr.Checkbox,
+            {"interactive": True},
+            section=section)
+    )
+    shared.opts.add_option(
+        "ch_proxy",
+        shared.OptionInfo(
+            "",
+            "Proxy to use for fetching models and model data. Format:  http://127.0.0.1:port",
+            gr.Textbox,
+            {"interactive": True, "max_lines": 1},
+            section=section)
+    )
+    shared.opts.onchange(
+        "ch_proxy",
+        update_proxy
+    )
 
 
 script_callbacks.on_ui_settings(on_ui_settings)
 script_callbacks.on_ui_tabs(on_ui_tabs)
-
-
-

extensions/Stable-Diffusion-Webui-Civitai-Helper/scripts/sections.py

@@ -0,0 +1,693 @@
+""" -*- coding: UTF-8 -*-
+Sections for civitai_helper tab.
+"""
+
+import gradio as gr
+from scripts.ch_lib import model
+from scripts.ch_lib import js_action_civitai
+from scripts.ch_lib import model_action_civitai
+from scripts.ch_lib import civitai
+from scripts.ch_lib import duplicate_check
+from scripts.ch_lib import util
+
+model_types = list(model.folders.keys())
+
+def scan_models_section():
+    """ Scan Models Section """
+    with gr.Row():
+        gr.Markdown("### Scan Models for Civitai")
+    with gr.Row():
+        with gr.Column():
+            scan_model_types_drop = gr.CheckboxGroup(
+                choices=model_types,
+                label="Model Types",
+                value=model_types
+            )
+            nsfw_preview_scan_drop = gr.Dropdown(
+                label="Block NSFW Level Above",
+                choices=civitai.NSFW_LEVELS[1:],
+                value=util.get_opts("ch_nsfw_preview_threshold"),
+                elem_id="ch_nsfw_preview_scan_drop"
+            )
+    with gr.Row():
+        with gr.Column():
+            refetch_old_ckb = gr.Checkbox(
+                label="Replace Old Metadata Formats*",
+                value=False,
+                elem_id="ch_refetch_old_ckb"
+            )
+            gr.HTML("""
+                * [<a href=https://github.com/zixaphir/Stable-Diffusion-Webui-Civitai-Helper/wiki/Metadata-Format-Changes>wiki</a>] Do not use this option if you have made changes with the metadata editor without backing up your data!!<br><br>
+                """)
+
+        with gr.Column():
+            scan_model_civitai_btn = gr.Button(
+                value="Scan",
+                variant="primary",
+                elem_id="ch_scan_model_civitai_btn"
+            )
+
+            scan_civitai_info_image_meta_btn = gr.Button(
+                value="Update image generation information (Experimental)",
+                variant="primary",
+                elem_id="ch_Scan_civitai_info_image_meta_btn"
+            )
+
+    with gr.Row():
+        scan_model_log_md = gr.Markdown(
+            value="Scanning takes time, just wait. Check console log for details",
+            elem_id="ch_scan_model_log_md"
+        )
+
+    # ====events====
+    scan_model_civitai_btn.click(
+        model_action_civitai.scan_model,
+        inputs=[
+            scan_model_types_drop,
+            nsfw_preview_scan_drop, refetch_old_ckb
+        ],
+        outputs=scan_model_log_md
+    )
+
+    scan_civitai_info_image_meta_btn.click(
+        model.scan_civitai_info_image_meta,
+        outputs=scan_model_log_md
+    )
+
+def get_model_info_by_url_section():
+    """ Get Civitai Model Info by Model Page URL Section """
+
+    def get_model_names_by_input(model_type, empty_info_only):
+        names = civitai.get_model_names_by_input(model_type, empty_info_only)
+        return model_name_drop.update(choices=names)
+
+    no_info_model_names = civitai.get_model_names_by_input("ckp", False)
+
+    with gr.Column():
+        gr.Markdown("### Get Model Info from Civitai by URL")
+        gr.Markdown("Use this when scanning can not find a local model on civitai")
+        with gr.Row():
+            with gr.Column(scale=2):
+                model_type_drop = gr.Dropdown(
+                    choices=model_types,
+                    label="Model Type",
+                    value="ckp",
+                    multiselect=False,
+                    elem_classes="ch_vpadding"
+                )
+            with gr.Column(scale=1):
+                empty_info_only_ckb = gr.Checkbox(
+                    label="Only Show Models have no Info",
+                    value=False,
+                    elem_id="ch_empty_info_only_ckb",
+                    elem_classes="ch_vpadding"
+                )
+        with gr.Row():
+            with gr.Column(scale=2):
+                model_name_drop = gr.Dropdown(
+                    choices=no_info_model_names,
+                    label="Model",
+                    value="",
+                    multiselect=False
+                )
+            with gr.Column(scale=1):
+                nsfw_preview_url_drop = gr.Dropdown(
+                    label="Block NSFW Level Above",
+                    choices=civitai.NSFW_LEVELS[1:],
+                    value=util.get_opts("ch_nsfw_preview_threshold"),
+                    elem_id="ch_nsfw_preview_url_drop"
+                )
+        with gr.Row():
+            with gr.Column(scale=2, elem_classes="justify-bottom"):
+                model_url_or_id_txtbox = gr.Textbox(
+                    label="Civitai URL",
+                    lines=1,
+                    value=""
+                )
+            with gr.Column(scale=1, elem_classes="justify-bottom"):
+                get_civitai_model_info_by_id_btn = gr.Button(
+                    value="Get Model Info from Civitai",
+                    variant="primary"
+                )
+
+    get_model_by_id_log_md = gr.Markdown("")
+
+    # ====events====
+    model_type_drop.change(
+        get_model_names_by_input,
+        inputs=[
+            model_type_drop, empty_info_only_ckb
+        ],
+        outputs=model_name_drop
+    )
+    empty_info_only_ckb.change(
+        get_model_names_by_input,
+        inputs=[
+            model_type_drop, empty_info_only_ckb
+        ],
+        outputs=model_name_drop
+    )
+
+    get_civitai_model_info_by_id_btn.click(
+        model_action_civitai.get_model_info_by_input,
+        inputs=[
+            model_type_drop, model_name_drop,
+            model_url_or_id_txtbox, nsfw_preview_url_drop
+        ],
+        outputs=get_model_by_id_log_md
+    )
+
+def download_section():
+    """ Download Models Section """
+
+    model_filetypes = civitai.FILE_TYPES
+    file_elems = {}
+
+    dl_state = gr.State({
+        "model_info": {},
+        "filenames": {
+            # dl_version_str: filename,
+        },
+        "previews": {
+            # dl_version_str: [{url: url, nsfw: nsfw}],
+        },
+        "files": {
+            # dl_version_str: {
+            #   Model: bool,
+            #   Config: bool,
+            #   ...
+            # }
+        },
+        "files_count": {
+            # dl_version_str: int
+        },
+        "filtered_previews": []
+    })
+
+    def get_model_info_by_url(url, subfolder, state):
+        data = model_action_civitai.get_model_info_by_url(url)
+
+        if not data:
+            return None
+
+        state = {
+            "model_info": {},
+            "filenames": {},
+            "previews": {},
+            "files": {},
+            "files_count": {},
+            "filtered_previews": []
+        }
+
+        subfolders = sorted(data["subfolders"])
+        version_strs = data["version_strs"]
+        filenames = data["filenames"]
+
+        if subfolder == "" or subfolder not in subfolders:
+            subfolder = "/"
+
+        state["model_info"] = data["model_info"]
+        state["previews"] = data["previews"]
+
+        for filename, version in zip(filenames, version_strs):
+            state["filenames"][version] = filename
+
+        for version_files, version in zip(data["files"], version_strs):
+            filetypes = state["files"][version] = {}
+            files_count = 0
+            unhandled_files = []
+            for filedata in version_files:
+                files_count += 1
+                ch_filedata = {
+                    "id": filedata["id"],
+                    "name": filedata["name"],
+                }
+
+                if filedata["type"] in model_filetypes:
+                    filetypes[filedata["type"]] = (True, ch_filedata)
+                    continue
+
+                unhandled_files.append(f"{ch_filedata['id']}: {ch_filedata['name']}")
+
+            if len(unhandled_files) > 0:
+                filetypes["unhandled_files"] = "\n".join(unhandled_files)
+            else:
+                filetypes["unhandled_files"] = None
+
+            state["files_count"][version] = files_count
+
+        return [
+            state, data["model_name"], data["model_type"],
+            dl_subfolder_drop.update(
+                choices=subfolders,
+                value=subfolder
+            ),
+            dl_version_drop.update(
+                choices=version_strs,
+                value=version_strs[0]
+            ),
+            files_row.update(
+                visible=True
+            )
+        ]
+
+    def filter_previews(previews, nsfw_preview_url_drop):
+        images = []
+        for preview in previews:
+            if civitai.should_skip(nsfw_preview_url_drop, preview["nsfw"]):
+                continue
+            if preview["type"] == "image":
+                # Civitai added videos as previews, and webui does not like it
+                images.append(preview["url"])
+
+        return images
+
+    def update_dl_inputs(state, dl_version, nsfw_threshold, dl_preview_index):
+        filename = state["filenames"][dl_version]
+
+        if not filename:
+            filename = dl_filename_txtbox.value
+
+        file_parts = filename.split(".")
+        ext = file_parts.pop()
+        base = ".".join(file_parts)
+
+        previews = filter_previews(state["previews"][dl_version], nsfw_threshold)
+        state["filtered_previews"] = previews
+
+        preview = None
+        if len(previews) > dl_preview_index:
+            preview = previews[dl_preview_index]
+        elif len(previews) > 0:
+            preview = previews[0]
+
+        output_add = []
+
+        for key, elems in file_elems.items():
+            filedata = state["files"][dl_version].get(key, False)
+
+            visible = False
+            filename = ""
+            if filedata:
+                visible = True
+                if isinstance(filedata, str):
+                    filename = ", ".join([
+                        filedata_str.split(":")[1].strip()
+                        for filedata_str in filedata.split("\n")
+                    ])
+                elif isinstance(filedata, tuple):
+                    _, data = filedata
+                    filename = data["name"]
+                else:
+                    raise ValueError(f"Invalid filedata: {filedata}")
+
+            output_add.append(elems["txtbx"].update(value=filename))
+            output_add.append(elems["row"].update(visible=visible))
+
+        return [
+            state,
+            dl_filename_txtbox.update(
+                value=base
+            ),
+            dl_extension_txtbox.update(
+                value=ext
+            ),
+            dl_preview_img.update(
+                value=previews
+            ),
+            dl_preview_url.update(
+                value=preview
+            ),
+            download_all_row.update(
+                visible=(state["files_count"][dl_version] > 1)
+            )
+        ] + output_add
+
+    def update_dl_files_visibility(dl_all):
+        files_chkboxes = []
+        for chkbox in ch_dl_model_types_visibility:
+            files_chkboxes.append(
+                chkbox.update(
+                    visible=not dl_all
+                )
+            )
+
+        return files_chkboxes
+
+    def update_dl_preview_url(state, dl_preview_index):
+        preview_url = state["filtered_previews"][dl_preview_index]
+
+        return dl_preview_url.update(
+            value=preview_url
+        )
+
+    def update_dl_preview_index(evt: gr.SelectData):
+        # For some reason, you can't pass gr.SelectData and
+        # inputs at the same time. :/
+
+        return dl_preview_index.update(
+            value=evt.index
+        )
+
+    with gr.Row():
+        gr.Markdown("### Download Model")
+
+    with gr.Row():
+        with gr.Column(scale=2, elem_id="ch_dl_model_inputs"):
+
+            gr.Markdown(value="**1. Add URL and retrieve Model Info**")
+
+            with gr.Row():
+                with gr.Column(scale=2, elem_classes="justify-bottom"):
+                    dl_model_url_or_id_txtbox = gr.Textbox(
+                        label="Civitai URL",
+                        lines=1,
+                        max_lines=1,
+                        value="",
+                        placeholder="Model URL or Model ID"
+                    )
+                with gr.Column(elem_classes="justify-bottom"):
+                    dl_model_info_btn = gr.Button(
+                        value="Get Model Info by Civitai Url",
+                        variant="primary"
+                    )
+
+            gr.Markdown(value="**2. Pick Subfolder and Model Version**")
+
+            with gr.Row(elem_classes="ch_grid"):
+                dl_model_name_txtbox = gr.Textbox(
+                    label="Model Name",
+                    interactive=False,
+                    lines=1,
+                    max_lines=1,
+                    min_width=320,
+                    value=""
+                )
+                dl_model_type_txtbox = gr.Textbox(
+                    label="Model Type",
+                    interactive=False,
+                    lines=1,
+                    max_lines=1,
+                    min_width=320,
+                    value=""
+                )
+                dl_version_drop = gr.Dropdown(
+                    choices=[],
+                    label="Model Version",
+                    value="",
+                    min_width=320,
+                    multiselect=False
+                )
+                dl_subfolder_drop = gr.Dropdown(
+                    choices=[],
+                    label="Sub-folder",
+                    value="",
+                    min_width=320,
+                    multiselect=False
+                )
+                dl_duplicate_drop = gr.Dropdown(
+                    choices=["Skip", "Overwrite", "Rename New"],
+                    label="Duplicate File Behavior",
+                    value="Skip",
+                    min_width=320,
+                    multiselect=False
+                )
+                nsfw_preview_dl_drop = gr.Dropdown(
+                    label="Block NSFW Level Above",
+                    choices=civitai.NSFW_LEVELS[1:],
+                    value=util.get_opts("ch_nsfw_preview_threshold"),
+                    min_width=320,
+                    elem_id="ch_nsfw_preview_dl_drop"
+                )
+
+            with gr.Column(
+                visible=False,
+                variant="panel"
+            ) as files_row:
+
+                with gr.Row(variant="compact"):
+                    gr.Markdown("**Files**")
+
+                ch_output_add = []
+                ch_dl_model_types = []
+                ch_dl_model_types_visibility = []
+
+                for filetype in model_filetypes:
+                    with gr.Row(
+                        visible=False,
+                        equal_height=True,
+                    ) as row:
+                        file_elems[filetype] = elems = {}
+                        elems["row"] = row
+
+                        with gr.Column(scale=0, min_width=24, elem_classes="flex-center") as ckb_column:
+                            elems["ckb"] = filetype_ckb = gr.Checkbox(
+                                label="",
+                                value=True,
+                                min_width=0,
+                                interactive=(not filetype == "Model")
+                            )
+                        with gr.Column(scale=1, min_width=0):
+                            elems["txtbx"] = gr.Textbox(
+                                value="",
+                                interactive=False,
+                                label=filetype,
+                                max_lines=1,
+                                min_width=0
+                            )
+
+                        ch_dl_model_types_visibility.append(ckb_column)
+                        ch_dl_model_types.append(filetype_ckb)
+
+                        ch_output_add.append(elems["txtbx"])
+                        ch_output_add.append(row)
+
+                with gr.Row(visible=False) as unhandled_files_row:
+                    file_elems["unhandled_files"] = elems = {}
+                    elems["row"] = row
+
+                    elems["txtbx"] = gr.Textbox(
+                        value="",
+                        interactive=False,
+                        label="Unhandled Files (Files that we don't know what to do with but will still downloaded with \"Download All Files\")",
+                    )
+
+                    ch_output_add.append(elems["txtbx"])
+                    ch_output_add.append(unhandled_files_row)
+
+                with gr.Row(visible=False) as download_all_row:
+                    dl_all_ckb = gr.Checkbox(
+                        label="Download All Files",
+                        value=False,
+                        elem_id="ch_dl_all_ckb",
+                        elem_classes="ch_vpadding"
+                    )
+
+        with gr.Column(scale=1, elem_id="ch_preview_col", min_width=512):
+            with gr.Row(elem_classes="flex-center"):
+                dl_preview_img = gr.Gallery(
+                    show_label=True,
+                    label="Preview Image",
+                    value=None,
+                    elem_id="ch_dl_preview_img",
+                    allow_preview=True,
+                    preview=False,
+                    object_fit="scale-down"
+                )
+                dl_preview_url = gr.Textbox(
+                    value="",
+                    visible=False,
+                    elem_id="ch_dl_preview_url"
+                )
+                dl_preview_index = gr.Number(
+                    value=0,
+                    visible=False,
+                    elem_id="ch_dl_preview_index",
+                    precision=0
+                )
+
+    with gr.Row():
+        with gr.Column(scale=2, elem_classes="justify-bottom"):
+            dl_filename_txtbox = gr.Textbox(
+                label="Rename Model",
+                value="",
+                lines=1,
+                max_lines=1,
+                elem_id="ch_dl_filename_txtbox",
+                elem_classes="ch_vpadding"
+            )
+            dl_extension_txtbox = gr.Textbox(
+                label="Model extension",
+                value="",
+                elem_id="ch_dl_extension_txtbox",
+                visible=False
+            )
+
+        with gr.Column(elem_classes="justify-bottom"):
+            dl_civitai_model_by_id_btn = gr.Button(
+                value="3. Download Model",
+                elem_classes="ch_vmargin",
+                variant="primary"
+            )
+
+    with gr.Row():
+        dl_log_md = gr.Markdown(
+            value="Check Console log for Downloading Status"
+        )
+
+    # ====events====
+    dl_model_info_btn.click(
+        get_model_info_by_url,
+        inputs=[
+            dl_model_url_or_id_txtbox, dl_subfolder_drop, dl_state
+        ],
+        outputs=[
+            dl_state, dl_model_name_txtbox,
+            dl_model_type_txtbox, dl_subfolder_drop,
+            dl_version_drop, files_row
+        ]
+    )
+
+    dl_inputs = [
+            dl_state, dl_model_type_txtbox,
+            dl_subfolder_drop, dl_version_drop,
+            dl_filename_txtbox, dl_extension_txtbox,
+            dl_all_ckb, nsfw_preview_dl_drop,
+            dl_duplicate_drop, dl_preview_url
+        ] + ch_dl_model_types
+
+    dl_civitai_model_by_id_btn.click(
+        model_action_civitai.dl_model_by_input,
+        inputs=dl_inputs,
+        outputs=dl_log_md
+    )
+
+    ver_outputs = [
+        dl_state, dl_filename_txtbox, dl_extension_txtbox,
+        dl_preview_img, dl_preview_url, download_all_row,
+    ] + ch_output_add
+
+    dl_version_drop.change(
+        update_dl_inputs,
+        inputs=[dl_state, dl_version_drop, nsfw_preview_dl_drop, dl_preview_index],
+        outputs=ver_outputs
+    )
+    dl_all_ckb.change(
+        update_dl_files_visibility,
+        inputs=dl_all_ckb,
+        outputs=ch_dl_model_types_visibility
+    )
+    nsfw_preview_dl_drop.change(
+        update_dl_inputs,
+        inputs=[dl_version_drop, dl_state, nsfw_preview_dl_drop],
+        outputs=ver_outputs
+    )
+    # Gradio has so many issues with Gradio.Gallery...
+    dl_preview_img.select(
+        update_dl_preview_index,
+        None,
+        dl_preview_index
+    )
+    dl_preview_index.change(
+        update_dl_preview_url,
+        [dl_state, dl_preview_index],
+        dl_preview_url
+    )
+
+
+def scan_for_duplicates_section():
+    """ Scan Duplicate Models Section """
+    with gr.Column():
+        gr.Markdown("### Scan for duplicate models")
+        with gr.Row():
+            with gr.Column():
+                scan_dup_model_types_drop = gr.CheckboxGroup(
+                    choices=model_types,
+                    label="Model Types",
+                    value=model_types
+                )
+        with gr.Row():
+            with gr.Column(scale=2):
+                cached_hash_ckb = gr.Checkbox(
+                    label="Use Hash from Metadata (May have false-positives but can be useful if you've pruned models)",
+                    value=False,
+                    elem_id="ch_cached_hash_ckb"
+                )
+            with gr.Column():
+                scan_dup_model_btn = gr.Button(
+                    value="Scan",
+                    variant="primary",
+                    elem_id="ch_scan_dup_model_civitai_btn"
+                )
+
+        # with gr.Row():
+        scan_dup_model_log_md = gr.HTML(
+            value="Scanning takes time, just wait. Check console log for details",
+            elem_id="ch_scan_dup_model_log_md"
+        )
+
+    # ====events====
+    scan_dup_model_btn.click(
+        duplicate_check.scan_for_dups,
+        inputs=[
+            scan_dup_model_types_drop,
+            cached_hash_ckb
+        ],
+        outputs=scan_dup_model_log_md
+    )
+
+def check_new_versions_section(js_msg_txtbox):
+    """ Check models' new version section """
+
+    with gr.Column():
+        gr.Markdown("### Check models' new version")
+        with gr.Row():
+            with gr.Column(scale=2):
+                model_types_ckbg = gr.CheckboxGroup(
+                    choices=model_types,
+                    label="Model Types",
+                    value=[
+                        "ti", "hyper", "ckp", "lora", "lycoris"
+                    ]
+                )
+                nsfw_preview_update_drop = gr.Dropdown(
+                    label="Block NSFW Level Above",
+                    choices=civitai.NSFW_LEVELS[1:],
+                    value=util.get_opts("ch_nsfw_preview_threshold"),
+                    elem_id="ch_nsfw_preview_dl_drop"
+                )
+        with gr.Row():
+            with gr.Column(scale=2):
+                check_models_new_version_btn = gr.Button(
+                    value="Check New Version from Civitai",
+                    variant="primary"
+                )
+
+        with gr.Row():
+            with gr.Column():
+                dl_new_version_log_md = gr.Markdown()
+                check_models_new_version_log_md = gr.HTML(
+                    "It takes time, just wait. Check console log for details"
+                )
+
+    # ====events====
+    check_models_new_version_btn.click(
+        model_action_civitai.check_models_new_version_to_md,
+        inputs=model_types_ckbg,
+        outputs=check_models_new_version_log_md
+    )
+
+    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_dl_model_new_version_btn.click(
+        js_action_civitai.dl_model_new_version,
+        inputs=[
+            js_msg_txtbox,
+            nsfw_preview_update_drop
+        ],
+        outputs=dl_new_version_log_md
+    )

extensions/Stable-Diffusion-Webui-Civitai-Helper/style.css

@@ -5,7 +5,7 @@ blockquote ul {
 
 blockquote ol {
     list-style:decimal;
-    margin:4px 40px;
+    margin: 4px 40px;
 }
 
 .block.padded.ch_box {
@@ -16,3 +16,110 @@ blockquote ol {
     padding: 10px 0 !important;
 }
 
+#tab_civitai_helper .ch_vmargin {
+    margin-top: 10px;
+    margin-bottom: 10px;
+}
+
+#tab_civitai_helper .extra-network-cards {
+    min-height: 0;
+    max-height: 50vh;
+}
+
+.card-button {
+    margin: 0px 2px;
+    padding: 2px;
+    font-size: 150%;
+}
+
+#tab_civitai_helper .justify-bottom {
+    justify-content: flex-end;
+}
+
+#tab_civitai_helper .ch_grid div {
+    flex-basis: 30%;
+}
+
+#tab_civitai_helper .flex-center {
+    justify-content: center;
+    align-content: center;
+}
+@media only screen and (min-width: 600px) {
+    #ch_dl_preview_img {
+        max-width: 512px;
+        min-width: 512px;
+        max-height: 768px;
+    }
+}
+
+/* SD WebUI has not line height limit for model names.
+ * This can result in the model name filling the entire card,
+ * hiding our buttons. This is honestly a WebUI issue more
+ * than an issue on this end, but we can fix it anyways, so
+ * why not?
+ */
+.extra-network-cards .card .name:not(:hover) {
+    display: block;
+    overflow: hidden;
+    max-height: 3em;
+}
+
+.extra-network-thumbs .card-button {
+    font-size: 100%;
+    display: inline;
+    position: static;
+    background-image: none;
+    background: rgba(0, 0, 0, 0.8);
+}
+
+.duplicate_model {
+    display: flex;
+    flex-flow: row wrap;
+    justify-content: flex-start;
+    align-items: flex-start;
+    gap: 10px;
+}
+
+#ch_path_el {
+    position: fixed;
+    border: var(--input-border-width) solid var(--input-border-color);
+    border-radius: var(--input-radius);
+    background: var(--input-background-fill);
+    padding: var(--input-padding);
+    color: var(--body-text-color);
+    font-weight: var(--input-text-weight);
+    font-size: var(--input-text-size);
+    line-height: var(--line-sm);
+    box-shadow: var(--input-shadow);
+    font-family:var(--font);
+    pointer-events: none;
+}
+
+/* Lobe theme */
+.actions .additional > a {
+    display: none;
+}
+
+/*
+ * Gross, ugly, and I hate it, but the only way to override an
+ * !important is with a more granular !important.
+ */
+.acss-culdhq .extra-networks .actions .additional > ul a.card-button {
+    margin: 0px 2px !important;
+    padding: 2px !important;
+    font-size: 150% !important;
+}
+
+/* Work around a display bug in Lobe that causes lowers opacity to 0.2 */
+#txt2img_extra_tabs [style*="_cards_html"] .pending,
+#img2img_extra_tabs [style*="_cards_html"] .pending {
+    opacity: 1;
+}
+
+/* Work around a display bug in Lobe that causes an infinite counter */
+#txt2img_extra_tabs [id*="_cards_html"] .progress-text ,
+#img2img_extra_tabs [id*="_cards_html"] .progress-text  {
+    display: none;
+}
+
+/* end of Lobe theme */

extensions/sd-webui-infinite-image-browsing/.env

@@ -0,0 +1,3 @@
+IIB_SECRET_KEY=
+IIB_ACCESS_CONTROL=disable
+IIB_ACCESS_CONTROL_PERMISSION=read-write

extensions/sd-webui-infinite-image-browsing/.gitignore

@@ -1,17 +0,0 @@
-
-*.log
-__pycache__
-iib.db
-tags-translate.csv
-launch.sh
-conf.json
-iib.db-journal
-.env
-standalone.cmd
-.vscode
-build/**/*
-dist/**/*
-*.spec
-out/**/*
-venv/**/*
-zip_temp/*.zip

forcePull.bat

@@ -0,0 +1,11 @@
+@echo off
+
+git fetch --all
+git reset --hard origin/main
+git pull
+git repack -a -d --depth=250 --window=250
+cd .git
+rmdir /S /Q lfs
+cd ..
+
+pause