diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..e30f246 --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +scripts/__pycache__/ diff --git a/README.md b/README.md index 61af757..fd374fa 100644 --- a/README.md +++ b/README.md @@ -1,2 +1,64 @@ # Stable-Diffusion-Webui-Civitai-Helper -Stable Diffusion Webui Extension for Civitai, to manage your model much more easily. +Stable Diffusion Webui Extension for Civitai, to handle your models much more easily. + +# Feature +* Scan all models to download model information and preview images from Civitai. +* Modified Build-in "Extra Network" cards, to add following buttons on each card: + - 🖼: Modified "replace preview" text into this icon + - 🌐: Open this model's Civitai url in a new tab + - 💡: Add this model's trigger words to prompt + - 🏷: Use this model's preview image's prompt + +# Install +Go to SD webui's extension tab, go to `Install from url` sub-tab. +Copy this project's url into it, click install. + +Or, just download this project as zip file, unzip it to `Your SD webui/extensions`. + +Then reload UI with "Reload UI" Button in Setting tab. + +Done. + +# How ot use +## Scan model +Go to extension tab "Civitai Helper". There is a button called "Scan Model". + +![](img/extension_tab.jpg) + +Click it, extension will scan all your models to generate SHA256 hash, and use this hash, to get model information and preview images from civitai. + +For each model, it will create a json file to save all model info from civitai. This model info file will be "Your_model_name.civitai.info" in your model folder. + +![](img/model_info_file.jpg) + +If a model info file is already exists, it will skip this model. If a model can not be find in civitai, it will create an empty model info file, so it won't scan this model twice. + +### Add new models +When you have some new models, just click this button again, to get new model's information and preview images. + +## Model Card +Open SD webui's build-in "Extra Network" tab, to show model cards. + +![](img/extra_network.jpg) + + +Move your mouse on to the bottom of a model card. It will show 4 icon buttons: + - 🖼: Modified "replace preview" text into this icon + - 🌐: Open this model's Civitai url in a new tab + - 💡: Add this model's trigger words to prompt + - 🏷: Use this model's preview image's prompt + +![](img/model_card.jpg) + +If you click Refresh Button of extra network, those additional buttons will be removed. You can click `Refresh Civitai Helper` button to bring them back. + +![](img/refresh_ch.jpg) + + + +Enjoy! + + + + + diff --git a/img/extension_tab.jpg b/img/extension_tab.jpg new file mode 100644 index 0000000..d49a3a1 Binary files /dev/null and b/img/extension_tab.jpg differ diff --git a/img/extra_network.jpg b/img/extra_network.jpg new file mode 100644 index 0000000..25b2acb Binary files /dev/null and b/img/extra_network.jpg differ diff --git a/img/model_card.jpg b/img/model_card.jpg new file mode 100644 index 0000000..315e359 Binary files /dev/null and b/img/model_card.jpg differ diff --git a/img/model_info_file.jpg b/img/model_info_file.jpg new file mode 100644 index 0000000..d9e36c8 Binary files /dev/null and b/img/model_info_file.jpg differ diff --git a/img/refresh_ch.jpg b/img/refresh_ch.jpg new file mode 100644 index 0000000..33f616c Binary files /dev/null and b/img/refresh_ch.jpg differ diff --git a/javascript/civitai_helper.js b/javascript/civitai_helper.js new file mode 100644 index 0000000..555ee52 --- /dev/null +++ b/javascript/civitai_helper.js @@ -0,0 +1,312 @@ +"use strict"; + + +function getActivePrompt() { + const currentTab = get_uiCurrentTabContent(); + switch (currentTab.id) { + case "tab_txt2img": + return currentTab.querySelector("#txt2img_prompt textarea"); + case "tab_img2img": + return currentTab.querySelector("#img2img_prompt textarea"); + } + return null; +} + +function getActiveNegativePrompt() { + const currentTab = get_uiCurrentTabContent(); + switch (currentTab.id) { + case "tab_txt2img": + return currentTab.querySelector("#txt2img_neg_prompt textarea"); + case "tab_img2img": + return currentTab.querySelector("#img2img_neg_prompt textarea"); + } + return null; +} + + +//button's click function +function open_model_url(model_type, model_name){ + console.log("start open_model_url"); + + //get hidden components of extension + let js_msg_txtbox = gradioApp().querySelector("#ch_js_msg_txtbox textarea"); + let js_open_url_btn = gradioApp().getElementById("ch_js_open_url_btn"); + + + //msg to python side + let msg = { + "action": "", + "model_type": "", + "model_name": "", + "prompt": "", + "neg_prompt": "", + } + + + msg["action"] = "open_url"; + msg["model_type"] = model_type; + msg["model_name"] = model_name; + msg["prompt"] = ""; + msg["neg_prompt"] = ""; + + // fill to msg box + js_msg_txtbox.value = JSON.stringify(msg); + js_msg_txtbox.dispatchEvent(new Event("input")); + + //click hidden button + js_open_url_btn.click(); + + console.log("end open_model_url"); + + +} + +function add_trigger_words(model_type, model_name){ + console.log("start add_trigger_words"); + + //get hidden components of extension + let js_msg_txtbox = gradioApp().querySelector("#ch_js_msg_txtbox textarea"); + let js_add_trigger_words_btn = gradioApp().getElementById("ch_js_add_trigger_words_btn"); + + + + //msg to python side + let msg = { + "action": "", + "model_type": "", + "model_name": "", + "prompt": "", + "neg_prompt": "", + } + + msg["action"] = "add_trigger_words"; + msg["model_type"] = model_type; + msg["model_name"] = model_name; + msg["neg_prompt"] = ""; + + // get active prompt + let prompt = getActivePrompt(); + msg["prompt"] = prompt.value; + + // fill to msg box + js_msg_txtbox.value = JSON.stringify(msg); + js_msg_txtbox.dispatchEvent(new Event("input")); + + //click hidden button + js_add_trigger_words_btn.click(); + + console.log("end add_trigger_words"); + +} + +function use_preview_prompt(model_type, model_name){ + console.log("start use_preview_prompt"); + + //get hidden components of extension + let js_msg_txtbox = gradioApp().querySelector("#ch_js_msg_txtbox textarea"); + let js_use_preview_prompt_btn = gradioApp().getElementById("ch_js_use_preview_prompt_btn"); + + + + //msg to python side + let msg = { + "action": "", + "model_type": "", + "model_name": "", + "prompt": "", + "neg_prompt": "", + } + + msg["action"] = "use_preview_prompt"; + msg["model_type"] = model_type; + msg["model_name"] = model_name; + + // get active prompt + prompt = getActivePrompt(); + msg["prompt"] = prompt.value; + + // get active neg prompt + let neg_prompt = getActiveNegativePrompt(); + msg["neg_prompt"] = neg_prompt.value; + + // fill to msg box + js_msg_txtbox.value = JSON.stringify(msg); + js_msg_txtbox.dispatchEvent(new Event("input")); + + //click hidden button + js_use_preview_prompt_btn.click(); + + console.log("end use_preview_prompt"); + +} + + +onUiLoaded(() => { + + + + // get all extra network tabs + let tab_prefix_list = ["txt2img", "img2img"]; + let model_type_list = ["textual_inversion", "hypernetworks", "checkpoints", "lora"]; + let cardid_suffix = "cards"; + + // update extra network tab pages' cards + // * replace "replace preview" text button into an icon + // * add 3 button to each card: + // - open model url 🌐 + // - add trigger words 💡 + // - use preview image's prompt 🏷 + // 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(){ + + + //change all "replace preview" into an icon + let extra_network_id = ""; + let extra_network_node = null; + let addtional_nodes = null; + let replace_preview_btn = null; + let ul_node = null; + let model_name_node = null; + let model_name = ""; + let model_type = ""; + let cards = null; + for (const tab_prefix of tab_prefix_list) { + 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; + } + + if (!model_type) { + console.log("can not get model_type from: " + js_model_type); + continue; + } + + 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); + if (!extra_network_node) { + console.log("can not find extra_network_node: " + extra_network_id); + continue; + } + // console.log("find extra_network_node: " + extra_network_id); + + // get all card nodes + cards = extra_network_node.querySelectorAll(".card"); + for (let card of cards) { + // replace preview text button into icon + replace_preview_btn = card.querySelector(".actions .additional a"); + replace_preview_btn.style.margin = "0px 10px"; + if (replace_preview_btn) { + if (replace_preview_btn.innerHTML == "replace preview") { + replace_preview_btn.innerHTML = "🖼"; + } + } + + //get model name node + model_name_node = card.querySelector(".actions .name"); + if (!model_name_node){ + console.log("can not find model name node for cards in " + extra_network_id); + continue; + } + + // get model name + model_name = model_name_node.innerHTML; + if (!model_name) { + console.log("model name is empty for cards in " + extra_network_id); + continue; + } + + + //get ul node, which is the parent of all buttons + ul_node = card.querySelector(".actions .additional ul"); + + // then we need to add 3 buttons to each ul node: + let open_url_node = document.createElement("button"); + open_url_node.innerHTML = "🌐"; + open_url_node.title = "Open this model's civitai url"; + open_url_node.style.margin = "0px 10px"; + open_url_node.setAttribute("onclick","open_model_url('"+model_type+"', '"+model_name+"')"); + + let add_trigger_words_node = document.createElement("button"); + add_trigger_words_node.innerHTML = "💡"; + add_trigger_words_node.title = "Add trigger words to prompt"; + add_trigger_words_node.style.margin = "0px 10px"; + add_trigger_words_node.setAttribute("onclick","add_trigger_words('"+model_type+"', '"+model_name+"')"); + + let use_preview_prompt_node = document.createElement("button"); + use_preview_prompt_node.innerHTML = "🏷"; + use_preview_prompt_node.title = "Use promt from preview image"; + use_preview_prompt_node.style.margin = "0px 10px"; + use_preview_prompt_node.setAttribute("onclick","use_preview_prompt('"+model_type+"', '"+model_name+"')"); + + //add to card + ul_node.appendChild(open_url_node); + ul_node.appendChild(add_trigger_words_node); + ul_node.appendChild(use_preview_prompt_node); + + + + + + + + } + + + } + } + + + } + + //run it once + update_card_for_civitai(); + + let tab_id = "" + let extra_tab = null; + let extra_toolbar = 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); + + //get toolbar + extra_toolbar = extra_tab.querySelector("div.flex.border-b-2.flex-wrap"); + + if (!extra_toolbar){ + console.log("can not get extra network toolbar for " + tab_id); + continue; + } + + // add refresh button to toolbar + let ch_refresh = document.createElement("button"); + ch_refresh.innerHTML = "Refresh Civitai Helper"; + ch_refresh.title = "Refresh Civitai Helper's model card buttons"; + ch_refresh.className = "gr-button gr-button-lg gr-button-secondary"; + ch_refresh.onclick = update_card_for_civitai; + + extra_toolbar.appendChild(ch_refresh); + + } + + + + +}); + + + diff --git a/scripts/civitai_helper.py b/scripts/civitai_helper.py new file mode 100644 index 0000000..cb501ba --- /dev/null +++ b/scripts/civitai_helper.py @@ -0,0 +1,434 @@ +# -*- 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 modules +from modules import script_callbacks + +# from modules import images +# from modules.processing import process_images, Processed +# from modules.processing import Processed +# from modules.shared import opts, cmd_opts, state + + +# init +model_folders = { + "ti": "embeddings", + "hyper": os.path.join("models", "hypernetworks"), + "ckp": os.path.join("models", "Stable-diffusion"), + "lora": os.path.join("models", "Lora"), +} + +model_exts = (".bin", ".pt", ".safetensors", ".ckpt") +model_info_exts = ".info" +civitai_info_suffix = ".civitai" +civitai_hash_api_url = "https://civitai.com/api/v1/model-versions/by-hash/" + +# js action list +js_actions = ("open_url", "add_trigger_words", "use_preview_prompt") + +root_path = os.getcwd() + +# print for debugging +def printD(msg): + print(f"Civitai Helper: {msg}") + + +def gen_file_sha256(filname): + ''' calculate file sha256 ''' + hash_value = "" + with open(filname, "rb") as f: + sha256obj = hashlib.sha256() + sha256obj.update(f.read()) + hash_value = sha256obj.hexdigest() + + printD("sha256: " + hash_value) + return hash_value + +# scan model to generate SHA256, then use this SHA256 to get model info from civitai +def scan_model(skip_nsfw_preview): + printD("Start scan_model") + + for model_type, model_folder in model_folders.items(): + folder_path = os.path.join(root_path, model_folder) + printD("Scanning path: " + folder_path) + for filename in os.listdir(folder_path): + # check ext + item = os.path.join(folder_path, filename) + base, ext = os.path.splitext(item) + if ext in model_exts: + # find a model + # get preview image + first_preview = base+".png" + sec_preview = base+".preview.png" + # get info file + info_file = base + civitai_info_suffix + model_info_exts + # check info file + if not os.path.isfile(info_file): + # get model's sha256 + printD("Generate SHA256 for model: " + filename) + hash = gen_file_sha256(item) + + if not hash: + printD("failed generate SHA256 for this file.") + return + + # use this sha256 to get model info from civitai + printD("Request model info from civitai") + r = requests.get(civitai_hash_api_url+hash) + if not r.ok: + if r.status_code == 404: + # this is not a civitai model + printD("Civitai does not have this model") + printD("Write empty model info file") + empty_info = {} + with open(info_file, 'w') as f: + data = json.dumps(empty_info) + f.write(data) + # go to next file + continue + else: + printD("Get errorcode: " + str(r.status_code)) + printD(r.text) + return + + # try to get content + content = None + try: + content = r.json() + except Exception as e: + printD("Parse response json failed") + printD(str(e)) + printD("response:") + printD(r.text) + return + + if not content: + printD("error, content from civitai is None") + return + + # write model info to file + printD("Write model info to file: " + info_file) + with open(info_file, 'w') as f: + data = json.dumps(content) + f.write(data) + + # check preview image + if not os.path.isfile(sec_preview): + # need to download preview image + printD("Need preview image for this model") + if content["images"]: + for img_dict in content["images"]: + if "nsfw" in img_dict.keys(): + if img_dict["nsfw"]: + printD("This image is NSFW") + if skip_nsfw_preview: + printD("Skip NSFW image") + continue + + if "url" in img_dict.keys(): + printD("Sending request for image: " + img_dict["url"]) + # get image + img_r = requests.get(img_dict["url"], stream=True) + if not img_r.ok: + printD("Get errorcode: " + str(r.status_code)) + printD(r.text) + return + + # write to file + with open(sec_preview, 'wb') as f: + img_r.raw.decode_content = True + shutil.copyfileobj(img_r.raw, f) + + printD("Created Preview image: " + sec_preview) + + # we only need 1 preview image + break + + # for testing, we only check 1 model for each type + # break + + printD("End scan_model") + + + +# handle request from javascript +# parameter: msg - msg from js +# return: (action, model_type, model_name, prompt, neg_prompt) +def parse_js_msg(msg): + printD("Start parse js msg") + msg_dict = json.loads(msg) + + if "action" not in msg_dict.keys(): + printD("Can not find action from js request") + return + + if "model_type" not in msg_dict.keys(): + printD("Can not find model type from js request") + return + + if "model_name" not in msg_dict.keys(): + printD("Can not find model name from js request") + return + + if "prompt" not in msg_dict.keys(): + printD("Can not find prompt from js request") + return + + if "neg_prompt" not in msg_dict.keys(): + printD("Can not find neg_prompt from js request") + return + + action = msg_dict["action"] + model_type = msg_dict["model_type"] + model_name = msg_dict["model_name"] + prompt = msg_dict["prompt"] + neg_prompt = msg_dict["neg_prompt"] + + if not action: + printD("Action from js request is None") + return + + if not model_type: + printD("model_type from js request is None") + return + + if not model_name: + printD("model_name from js request is None") + return + + + if action not in js_actions: + printD("Unknow action: " + action) + return + + if model_type not in model_folders.keys(): + printD("Unknow model_type: " + model_type) + return + + printD("End parse js msg") + + return (action, model_type, model_name, prompt, neg_prompt) + + + + +# get model info file's content by model type and model name +# parameter: model_type, model_name +# return: model_info_dict +def get_model_info(model_type, model_name): + if model_type not in model_folders.keys(): + printD("unknow model type: " + model_type) + return None + + model_folder = model_folders[model_type] + model_info_filename = model_name + civitai_info_suffix + model_info_exts + model_info_filepath = os.path.join(root_path, model_folder, model_info_filename) + + if not os.path.isfile(model_info_filepath): + printD("Can not find model info file: " + model_info_filepath) + return None + + model_info = None + with open(model_info_filepath, 'r') as f: + try: + model_info = json.load(f) + except Exception as e: + printD("Selected file is not json: " + model_info_filepath) + printD(e) + return None + + return model_info + + +# get civitai's model url and open it in browser +# parameter: model_type, model_name +def open_model_url(msg): + printD("Start open_model_url") + + result = parse_js_msg(msg) + if not result: + printD("Parsing js ms failed") + return + + action, model_type, model_name, prompt, neg_prompt = result + + model_info = get_model_info(model_type, model_name) + if not model_info: + printD(f"Failed to get model info for {model_type} {model_name}") + return + + if "modelId" not in model_info.keys(): + printD(f"Failed to get model id from info file for {model_type} {model_name}") + return + + model_id = model_info["modelId"] + if not model_id: + printD(f"model id from info file of {model_type} {model_name} is None") + return + + url = "https://civitai.com/models/"+str(model_id) + + printD("Open Url: " + url) + # open url + webbrowser.open_new_tab(url) + + printD("End open_model_url") + + +# add trigger words to prompt +# parameter: model_type, model_name, prompt +# return: [new_prompt, new_prompt] - new prompt with trigger words, return twice for txt2img and img2img +def add_trigger_words(msg): + printD("Start add_trigger_words") + + result = parse_js_msg(msg) + if not result: + printD("Parsing js ms failed") + return + + action, model_type, model_name, prompt, neg_prompt = result + + + model_info = get_model_info(model_type, model_name) + if not model_info: + printD(f"Failed to get model info for {model_type} {model_name}") + return [prompt, prompt] + + if "trainedWords" not in model_info.keys(): + printD(f"Failed to get trainedWords from info file for {model_type} {model_name}") + return [prompt, prompt] + + trainedWords = model_info["trainedWords"] + if not trainedWords: + printD(f"No trainedWords from info file for {model_type} {model_name}") + return [prompt, prompt] + + if len(trainedWords) == 0: + printD(f"trainedWords from info file for {model_type} {model_name} 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 + printD("trigger_words: " + trigger_words) + printD("prompt: " + prompt) + printD("new_prompt: " + new_prompt) + + printD("End add_trigger_words") + + # add to prompt + 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): + printD("Start use_preview_image_prompt") + + result = parse_js_msg(msg) + if not result: + printD("Parsing js ms failed") + return + + action, model_type, model_name, prompt, neg_prompt = result + + + model_info = get_model_info(model_type, model_name) + if not model_info: + printD(f"Failed to get model info for {model_type} {model_name}") + return [prompt, neg_prompt, prompt, neg_prompt] + + if "images" not in model_info.keys(): + printD(f"Failed to get images from info file for {model_type} {model_name}") + return [prompt, neg_prompt, prompt, neg_prompt] + + images = model_info["images"] + if not images: + printD(f"No images from info file for {model_type} {model_name}") + return [prompt, neg_prompt, prompt, neg_prompt] + + if len(images) == 0: + printD(f"images from info file for {model_type} {model_name} is empty") + 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 "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 + + if not preview_prompt: + printD(f"There is no prompt of {model_type} {model_name} in its preview image") + return [prompt, neg_prompt, prompt, neg_prompt] + + printD("End use_preview_image_prompt") + + return [preview_prompt, preview_neg_prompt, preview_prompt, preview_neg_prompt] + + + + +def on_ui_tabs(): + # init + + # get prompt textarea + # UI structure + # check modules/ui.py, search for txt2img_paste_fields + # Negative prompt is the second element + txt2img_prompt = modules.ui.txt2img_paste_fields[0][0] + txt2img_neg_prompt = modules.ui.txt2img_paste_fields[1][0] + img2img_prompt = modules.ui.img2img_paste_fields[0][0] + img2img_neg_prompt = modules.ui.img2img_paste_fields[1][0] + + # ====UI==== + with gr.Blocks(analytics_enabled=False) as civitai_helper: + # info + gr.Markdown("Civitai Helper's extension tab") + + skip_nsfw_preview_ckb = gr.Checkbox(label="SKip NSFW Preview images", value=False, elem_id="ch_skip_nsfw_preview_ckb") + scan_model_btn = gr.Button(value="Scan model", elem_id="ch_scan_model_btn") + + # hidden component for js + js_msg_txtbox = gr.Textbox(label="Request Msg From Js", visible=False, lines=1, value="", elem_id="ch_js_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") + + # ====events==== + scan_model_btn.click(scan_model, inputs=[skip_nsfw_preview_ckb]) + js_open_url_btn.click(open_model_url, inputs=[js_msg_txtbox]) + js_add_trigger_words_btn.click(add_trigger_words, inputs=[js_msg_txtbox], outputs=[txt2img_prompt, img2img_prompt]) + js_use_preview_prompt_btn.click(use_preview_image_prompt, inputs=[js_msg_txtbox], outputs=[txt2img_prompt, txt2img_neg_prompt, img2img_prompt, img2img_neg_prompt]) + + # the third parameter is the element id on html, with a "tab_" as prefix + return (civitai_helper , "Civitai Helper", "civitai_helper"), + +script_callbacks.on_ui_tabs(on_ui_tabs)