diff --git a/.gitignore b/.gitignore index b1c59cb..0467a1d 100644 --- a/.gitignore +++ b/.gitignore @@ -1,2 +1,3 @@ scripts/__pycache__/ scripts/lib/__pycache__/ +setting.json diff --git a/README.cn.md b/README.cn.md index 46de42b..4a2d6b5 100644 --- a/README.cn.md +++ b/README.cn.md @@ -1,6 +1,9 @@ # Stable-Diffusion-Webui-Civitai-Helper Stable Diffusion Webui 扩展Civitai助手，用于更轻松的管理和使用Civitai模型。 +## 注意 +最新的v1.3版本，强大，但是是实验性质的。如果碰到问题，可以去Civitai下载1.2.1老版本: [Civitai Url](https://civitai.com/models/16768/civitai-helper-sd-webui-civitai-extension) + # 功能 * 扫描所有模型，从Civitai下载模型信息和预览图 * 修改了内置的"Extra Network"模型卡片，每个卡片增加了如下功能按钮: diff --git a/README.md b/README.md index 93ac0b5..ec63df7 100644 --- a/README.md +++ b/README.md @@ -4,6 +4,9 @@ # Stable-Diffusion-Webui-Civitai-Helper Stable Diffusion Webui Extension for Civitai, to handle your models much more easily. +## Notice +Latest v1.3 is powerful but experimental. If you have issue, you can go back to 1.2.1 by download old version from Civitai: [Civitai Url](https://civitai.com/models/16768/civitai-helper-sd-webui-civitai-extension) + # Features * Scans all models to download model information and preview images from Civitai. * Modified Built-in "Extra Network" cards, to add the following buttons on each card: @@ -78,6 +81,13 @@ Enjoy! * It cannot force a model link to Civitai by model ID for now. This will be added later. # Change Log +## v1.3 +* Open url at client side +* Link selected model to civitai by url or model id +* Save and load extension setting to file +* Show button action's output to UI +* Code refactoring + ## v1.2.1 * Add more error checking to work with different versions of SD webui. @@ -91,4 +101,4 @@ Enjoy! * Support subfolders * Check if refresh is needed when clicking "Refresh Civitai Helper" * Add space when adding trigger words -* Add memory optimised sha256 as an option +* Add memory Optimized sha256 as an option diff --git a/javascript/civitai_helper.js b/javascript/civitai_helper.js index 3ceee20..7ef21d2 100644 --- a/javascript/civitai_helper.js +++ b/javascript/civitai_helper.js @@ -1,5 +1,55 @@ "use strict"; +// 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") + 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 + } else { + return "" + } +} + + +// get msg from python side from a hidden textbox +// if textbox's value is different from old value then it will consider it is a new msg +// it will try once in every sencond, until it reach the max try times +const get_new_ch_py_msg = (old_value, max_count=3) => new Promise((resolve, reject) => { + console.log("run get_new_ch_py_msg") + + let count = 0; + let new_msg = ""; + let find_msg = false; + const interval = setInterval(() => { + const py_msg_txtbox = gradioApp().querySelector("#ch_py_msg_txtbox textarea"); + count++; + + 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) + + new_msg = py_msg_txtbox.value + if (new_msg != old_value) { + find_msg=true + } + } + + if (find_msg) { + resolve(new_msg); + clearInterval(interval); + } else if (count > max_count) { + reject(''); + clearInterval(interval); + } + + }, 1000); +}) function getActivePrompt() { const currentTab = get_uiCurrentTabContent(); @@ -25,7 +75,7 @@ function getActiveNegativePrompt() { //button's click function -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 @@ -53,15 +103,39 @@ function open_model_url(event, model_type, search_term){ js_msg_txtbox.value = JSON.stringify(msg); js_msg_txtbox.dispatchEvent(new Event("input")); + //get old py msg + let py_msg = get_ch_py_msg(); + //click hidden button js_open_url_btn.click(); - console.log("end open_model_url"); - - + // stop parent event event.stopPropagation() event.preventDefault() + //check response msg from python + let new_py_msg = await get_new_ch_py_msg(""); + console.log("new_py_msg:"); + console.log(new_py_msg); + + //check msg + if (new_py_msg) { + let py_msg_json = JSON.parse(new_py_msg); + //check for url + if (py_msg_json && py_msg_json.content) { + 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){ diff --git a/scripts/civitai_helper.py b/scripts/civitai_helper.py deleted file mode 100644 index 23247d9..0000000 --- a/scripts/civitai_helper.py +++ /dev/null @@ -1,526 +0,0 @@ -# -*- 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 modules -from modules import script_callbacks -from modules import shared - -# from modules import images -# from modules.processing import process_images, Processed -# from modules.processing import Processed -# from modules.shared import opts, cmd_opts, state - -# print for debugging -def printD(msg): - print(f"Civitai Helper: {msg}") - - -# 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 -model_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"), -} - - -# get cusomter model path -# will be modified when refactoring -if shared.cmd_opts.embeddings_dir and os.path.isdir(shared.cmd_opts.embeddings_dir): - model_folders["ti"] = shared.cmd_opts.embeddings_dir - -if shared.cmd_opts.hypernetwork_dir and os.path.isdir(shared.cmd_opts.hypernetwork_dir): - model_folders["hyper"] = shared.cmd_opts.hypernetwork_dir - -if shared.cmd_opts.ckpt_dir and os.path.isdir(shared.cmd_opts.ckpt_dir): - model_folders["ckp"] = shared.cmd_opts.ckpt_dir - -if shared.cmd_opts.lora_dir and os.path.isdir(shared.cmd_opts.lora_dir): - model_folders["lora"] = shared.cmd_opts.lora_dir - - -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") - - - -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 - -def gen_file_sha256_low_memory(filname): - printD("Using Memory Optimised SHA256") - hash_sha256 = hashlib.sha256() - with open(filname, "rb") as f: - for chunk in iter(lambda: f.read(4096), b""): - hash_sha256.update(chunk) - - hash_value = hash_sha256.hexdigest() - printD("sha256: " + hash_value) - return hash_value - - -# get image with full size -# width is in number, not string -def get_full_size_image_url(image_url, width): - return re.sub('/width=\d+/', '/width=' + str(width) + '/', image_url) - - -# scan model to generate SHA256, then use this SHA256 to get model info from civitai -def scan_model(low_memory_sha, max_size_preview, readable_model_info, skip_nsfw_preview): - printD("Start scan_model") - - - count = 0 - scan_log = "" - for model_type, model_folder in model_folders.items(): - printD("Scanning path: " + model_folder) - for root, dirs, files in os.walk(model_folder): - 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 - # set a Progress log - scan_log = "Scanned: " + str(count) + ", Scanning: "+ filename - # try to update to UI Here - # Is still trying to find a way - - # 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 = "" - if low_memory_sha: - hash = gen_file_sha256_low_memory(item) - else: - 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 = None - if readable_model_info: - data = json.dumps(content, indent=4) - else: - 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(): - 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"]) - - printD("Sending request for image: " + img_url) - # get image - img_r = requests.get(img_url, stream=True) - if not img_r.ok: - printD("Get error code: " + 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 - - # set counter - count = count+1 - - # for testing, we only check 1 model for each type - # break - - scan_log = "Done" - - printD("End scan_model") - - - -# handle request from javascript -# parameter: msg - msg from js -# return: (action, model_type, search_term, 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 "search_term" not in msg_dict.keys(): - printD("Can not find search_term 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"] - search_term = msg_dict["search_term"] - 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 search_term: - printD("search_term 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, search_term, prompt, neg_prompt) - - -# get model info file's content by model type and model name -# parameter: model_type, search_term -# return: model_info_dict -def get_model_info(model_type, search_term): - if model_type not in model_folders.keys(): - printD("unknow model type: " + model_type) - return None - - # 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) - model_info_base = base - if base[:1] == "/": - model_info_base = base[1:] - - - model_folder = model_folders[model_type] - model_info_filename = model_info_base + civitai_info_suffix + model_info_exts - model_info_filepath = os.path.join(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, search_term -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, search_term, prompt, neg_prompt = result - - model_info = get_model_info(model_type, search_term) - if not model_info: - printD(f"Failed to get model info for {model_type} {search_term}") - return - - if "modelId" not in model_info.keys(): - printD(f"Failed to get model id from info file for {model_type} {search_term}") - return - - model_id = model_info["modelId"] - if not model_id: - printD(f"model id from info file of {model_type} {search_term} 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, search_term, 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, search_term, prompt, neg_prompt = result - - - model_info = get_model_info(model_type, search_term) - if not model_info: - printD(f"Failed to get model info for {model_type} {search_term}") - return [prompt, prompt] - - if "trainedWords" not in model_info.keys(): - printD(f"Failed to get trainedWords from info file for {model_type} {search_term}") - return [prompt, prompt] - - trainedWords = model_info["trainedWords"] - if not trainedWords: - printD(f"No trainedWords from info file for {model_type} {search_term}") - return [prompt, prompt] - - if len(trainedWords) == 0: - 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 - 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, search_term, prompt, neg_prompt = result - - - model_info = get_model_info(model_type, search_term) - if not model_info: - 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(): - 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: - printD(f"No images from info file for {model_type} {search_term}") - return [prompt, neg_prompt, prompt, neg_prompt] - - if len(images) == 0: - printD(f"images from info file for {model_type} {search_term} 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 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 - - if not preview_prompt: - printD(f"There is no prompt of {model_type} {search_term} 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: - with gr.Blocks() as civitai_helper: - # UI will have 3 tabs: - # Model Info: Scan model or force a model link to civitai model info by model id or url - # Settging: Setting for general use, also can save setting for all tabs - # Tool: handy functions, like making all model info readable. - with gr.Tab("Model"): - with gr.Row(): - low_memory_sha_ckb = gr.Checkbox(label="Memory Optimised SHA256", value=True, elem_id="ch_low_memory_sha_ckb") - max_size_preview_ckb = gr.Checkbox(label="Download Max Size Preview", value=True, elem_id="ch_max_size_preview_ckb") - readable_model_info_ckb = gr.Checkbox(label="Readable Model Info file", value=True, elem_id="ch_readable_model_info_ckb") - 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") - - gr.Markdown("Check console log window for detail, after clicking Scan button") - - - # with gr.Tab("Settging"): - - # with gr.Tab("Tool"): - - # 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") - 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=[low_memory_sha_ckb, max_size_preview_ckb, readable_model_info_ckb, 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) - - diff --git a/scripts/lib/civitai.py b/scripts/lib/civitai.py index 4d19ddf..c12dc26 100644 --- a/scripts/lib/civitai.py +++ b/scripts/lib/civitai.py @@ -6,7 +6,7 @@ import re import requests from . import util from . import model - +from . import setting suffix = ".civitai" @@ -63,6 +63,124 @@ def get_model_info_by_hash(hash:str): +def get_model_info_by_id(id:str) -> dict: + util.printD("Request model info from civitai") + + if not id: + util.printD("id is empty") + return + + r = requests.get(url_dict["modelId"]+str(id)) + 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 + + # 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: + util.printD("Request version info from civitai") + + if not id: + util.printD("id is empty") + return + + r = requests.get(url_dict["modelVersionId"]+id) + 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 + + # 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: + + model_info = get_model_info_by_id(id) + if not model_info: + util.printD(f"Failed to get model info by id: {id}") + return + + # 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] + 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"] + + if not version_id: + util.printD("default version's id is None") + return + + # get version info + version_info = get_version_info_by_version_id(str(version_id)) + if not version_info: + util.printD(f"Failed to get version info by version_id: {version_id}") + return + + return version_info + + + + # get model info file's content by model type and search_term # parameter: model_type, search_term # return: model_info @@ -86,4 +204,148 @@ def load_model_info_by_search_term(model_type, search_term): util.printD("Can not find model info file: " + model_info_filepath) return - return model.load_model_info(model_info_filepath) \ No newline at end of file + return model.load_model_info(model_info_filepath) + + + + + +# 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] + + # set 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"] + + + + # 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): + 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) + + + return model_names + +def get_model_names_by_input(model_type, empty_info_only): + 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: + util.printD("Run get_model_id_from_url") + id = "" + + if not url: + util.printD("url or model id can not be empty") + return "" + + if url.isnumeric(): + # is already an id + id = url + return "" + + s = url.split("/") + if len(s) < 2: + util.printD("url is not valid") + return "" + + if s[-2].isnumeric(): + id = s[-2] + elif s[-1].isnumeric(): + id = s[-1] + else: + util.printD("There is no model id in this url") + return "" + + return id + + +# 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) -> str: + 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") + 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") + 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"]: + util.printD("This image is NSFW") + if skip_nsfw_preview: + util.printD("Skip NSFW image") + 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 + diff --git a/scripts/lib/general_action.py b/scripts/lib/general_action.py new file mode 100644 index 0000000..72a5047 --- /dev/null +++ b/scripts/lib/general_action.py @@ -0,0 +1,11 @@ +# -*- coding: UTF-8 -*- +# handle msg between js and python side +import os +import json +import requests +import shutil +import webbrowser +from . import util +from . import model +from . import civitai +from . import msg_handler diff --git a/scripts/lib/operator.py b/scripts/lib/js_action_civitai.py similarity index 51% rename from scripts/lib/operator.py rename to scripts/lib/js_action_civitai.py index 3c7245a..7311cc1 100644 --- a/scripts/lib/operator.py +++ b/scripts/lib/js_action_civitai.py @@ -3,110 +3,23 @@ import os import json import requests -import shutil import webbrowser from . import util from . import model from . import civitai -from . import msg +from . import msg_handler -# scan model to generate SHA256, then use this SHA256 to get model info from civitai -def scan_model(low_memory_sha, max_size_preview, skip_nsfw_preview): - util.printD("Start scan_model") - - model_count = 0 - image_count = 0 - scan_log = "" - for model_type, model_folder in model.folders.items(): - util.printD("Scanning path: " + model_folder) - for root, dirs, files in os.walk(model_folder): - 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 - # set a Progress log - scan_log = "Scanned: " + str(model_count) + ", Scanning: "+ filename - # try to update to UI Here - # Is still trying to find a way - - # get preview image - first_preview = base+".png" - sec_preview = base+".preview.png" - # get info file - info_file = base + civitai.suffix + model.info_ext - # check info file - if not os.path.isfile(info_file): - # get model's sha256 - util.printD("Generate SHA256 for model: " + filename) - hash = util.gen_file_sha256(item, low_memory_sha) - - if not hash: - util.printD("failed generate SHA256 for this file.") - return - - # use this sha256 to get model info from civitai - model_info = civitai.get_model_info_by_hash(hash) - if model_info is None: - util.printD("Fail to get model_info") - return - - # write model info to file - model.write_model_info(info_file, model_info) - - # set model_count - model_count = model_count+1 - - # check preview image - if not os.path.isfile(sec_preview): - # need to download preview image - util.printD("Need preview image for this model") - # 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") - continue - - 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"]: - util.printD("This image is NSFW") - if skip_nsfw_preview: - util.printD("Skip NSFW image") - 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 = civitai.get_full_size_image_url(img_url, img_dict["width"]) - - util.download_file(img_url, sec_preview) - image_count = image_count + 1 - # we only need 1 preview image - break - - - scan_log = "Done" - - util.printD("End scan_model") - - # get civitai's model url and open it in browser # parameter: model_type, search_term -def open_model_url(msg): +# output: python msg - will be sent to hidden textbox then picked by js side +def open_model_url(msg, open_url_with_js): util.printD("Start open_model_url") - result = msg.parse_js_msg(msg) + output = "" + result = msg_handler.parse_js_msg(msg) if not result: util.printD("Parsing js ms failed") return @@ -116,24 +29,36 @@ def open_model_url(msg): 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 + return "" if "modelId" not in model_info.keys(): util.printD(f"Failed to get model id from info file for {model_type} {search_term}") - return + return "" model_id = model_info["modelId"] if not model_id: util.printD(f"model id from info file of {model_type} {search_term} is None") - return + return "" url = civitai.url_dict["modelPage"]+str(model_id) - util.printD("Open Url: " + url) - # open url - webbrowser.open_new_tab(url) + + # msg content for js + content = { + "url":"" + } + + if not open_url_with_js: + util.printD("Open Url: " + url) + # open url + webbrowser.open_new_tab(url) + else: + util.printD("Send Url to js") + content["url"] = url + output = msg_handler.build_py_msg("open_url", content) util.printD("End open_model_url") + return output @@ -143,7 +68,7 @@ def open_model_url(msg): def add_trigger_words(msg): util.printD("Start add_trigger_words") - result = msg.parse_js_msg(msg) + result = msg_handler.parse_js_msg(msg) if not result: util.printD("Parsing js ms failed") return @@ -192,7 +117,7 @@ def add_trigger_words(msg): def use_preview_image_prompt(msg): util.printD("Start use_preview_image_prompt") - result = msg.parse_js_msg(msg) + result = msg_handler.parse_js_msg(msg) if not result: util.printD("Parsing js ms failed") return diff --git a/scripts/lib/model.py b/scripts/lib/model.py index 9594de4..72926c4 100644 --- a/scripts/lib/model.py +++ b/scripts/lib/model.py @@ -26,6 +26,8 @@ info_ext = ".info" # get cusomter model path def get_custom_model_folder(): + util.printD("Get Custom Model Folder") + global folders if shared.cmd_opts.embeddings_dir and os.path.isdir(shared.cmd_opts.embeddings_dir): @@ -40,7 +42,7 @@ 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 -get_custom_model_folder() + @@ -64,3 +66,55 @@ def load_model_info(path): return model_info + +# get model file names by model type +# parameter: model_type - string +# return: model name list +def get_model_names_by_type(model_type:str) -> list: + + model_folder = 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): + 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) + + + 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: + 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] + + # model could be in subfolder, need to walk. + model_root = "" + model_path = "" + for root, dirs, files in os.walk(folder): + 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) + + return + + + diff --git a/scripts/lib/model_action_civitai.py b/scripts/lib/model_action_civitai.py new file mode 100644 index 0000000..4b7e33b --- /dev/null +++ b/scripts/lib/model_action_civitai.py @@ -0,0 +1,115 @@ +# -*- coding: UTF-8 -*- +# handle msg between js and python side +import os +from . import util +from . import model +from . import civitai + + + +# scan model to generate SHA256, then use this SHA256 to get model info from civitai +# return output msg +def scan_model(low_memory_sha, max_size_preview, skip_nsfw_preview): + util.printD("Start scan_model") + + output = "" + model_count = 0 + image_count = 0 + # scan_log = "" + for model_type, model_folder in model.folders.items(): + util.printD("Scanning path: " + model_folder) + for root, dirs, files in os.walk(model_folder): + 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 + # get info file + info_file = base + civitai.suffix + model.info_ext + # check info file + if not os.path.isfile(info_file): + # get model's sha256 + hash = util.gen_file_sha256(item, low_memory_sha) + + 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) + if model_info is None: + output = "Failed to get model_info" + 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" + + util.printD(output) + + return output + + + +# Get model info by model type, name and url +# output is log info to display on markdown component +def get_model_info_by_id(model_type, model_name, model_url_or_id, max_size_preview, skip_nsfw_preview): + output = "" + # 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 + util.printD(output) + return output + + # 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" + util.printD(output) + return output + + # get info file path + base, ext = os.path.splitext(model_path) + info_file = base + civitai.suffix + model.info_ext + + # 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) + + util.printD("Saved model info to: "+ info_file) + + # check preview image + civitai.get_preview_image_by_model_path(model_path, max_size_preview, skip_nsfw_preview) + + output = "Model Info saved to: " + info_file + return output \ No newline at end of file diff --git a/scripts/lib/msg.py b/scripts/lib/msg_handler.py similarity index 100% rename from scripts/lib/msg.py rename to scripts/lib/msg_handler.py diff --git a/scripts/lib/setting.py b/scripts/lib/setting.py index 8ec239a..b42c860 100644 --- a/scripts/lib/setting.py +++ b/scripts/lib/setting.py @@ -16,7 +16,7 @@ data = { "skip_nsfw_preview": False }, "general":{ - "open_url_with_js": False, + "open_url_with_js": True, "check_model_version_at_startup": False, }, "tool":{ @@ -26,15 +26,13 @@ data = { # save setting +# return output msg for log def save(): - print("Saving tranlation service setting...") - # write data into globel trans_setting - global trans_setting + print("Saving setting to: " + path) - + json_data = json.dumps(data, indent=4) - # to json - json_data = json.dumps(data) + output = "" #write to file try: @@ -42,10 +40,14 @@ def save(): f.write(json_data) except Exception as e: util.printD("Error when writing file:"+path) + output = str(e) util.printD(str(e)) - return + return output - util.printD("Setting saved to: " + path) + output = "Setting saved to: " + path + util.printD(output) + + return output # load setting to global data @@ -89,7 +91,12 @@ def save_from_input(low_memory_sha, max_size_preview, skip_nsfw_preview, open_ur } } - save() + output = save() + + if not output: + output = "" + + return output # load to output def load_to_output(): diff --git a/scripts/lib/util.py b/scripts/lib/util.py index 3e31e59..3b4bc01 100644 --- a/scripts/lib/util.py +++ b/scripts/lib/util.py @@ -13,7 +13,7 @@ def gen_file_sha256(filname, is_low_memory=True): hash_sha256 = hashlib.sha256() with open(filname, "rb") as f: if is_low_memory: - printD("Using Memory Optimised SHA256") + printD("Using Memory Optimized SHA256") for chunk in iter(lambda: f.read(4096), b""): hash_sha256.update(chunk) else: @@ -27,7 +27,7 @@ def gen_file_sha256(filname, is_low_memory=True): # get preview image def download_file(url, path): - printD("Download file from: " + url) + printD("Downloading file from: " + url) # get file r = requests.get(url, stream=True) if not r.ok: diff --git a/scripts/ui.py b/scripts/ui.py new file mode 100644 index 0000000..5fb39e8 --- /dev/null +++ b/scripts/ui.py @@ -0,0 +1,138 @@ +# -*- 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 modules +from modules import script_callbacks +from modules import shared +from scripts.lib import model +from scripts.lib import js_action_civitai +from scripts.lib import model_action_civitai +from scripts.lib import setting +from scripts.lib import civitai + +# init +model.get_custom_model_folder() +setting.load() + + + +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] + + # ====Event's function==== + 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) + + + + # ====UI==== + # with gr.Blocks(analytics_enabled=False) as civitai_helper: + with gr.Blocks(css="button {background-color: #228be6}") as civitai_helper: + + # init + low_memory_sha = setting.data["model"]["low_memory_sha"] + max_size_preview = setting.data["model"]["max_size_preview"] + skip_nsfw_preview = setting.data["model"]["skip_nsfw_preview"] + open_url_with_js = setting.data["general"]["open_url_with_js"] + check_model_version_at_startup = setting.data["general"]["check_model_version_at_startup"] + + model_types = list(model.folders.keys()) + no_info_model_names = civitai.get_model_names_by_input("ckp", False) + + + # UI will have 3 tabs: + # Model Info: Scan model or force a model link to civitai model info by model id or url + # General: Setting for general use, also can save setting for all tabs + # Tool: handy functions, like making all model info readable. + with gr.Tab("Model"): + with gr.Box(): + with gr.Column(): + gr.Markdown("### Scan Models for Civitai") + with gr.Row(): + low_memory_sha_ckb = gr.Checkbox(label="Memory Optimized SHA256", value=low_memory_sha, elem_id="ch_low_memory_sha_ckb") + max_size_preview_ckb = gr.Checkbox(label="Download Max Size Preview", value=max_size_preview, elem_id="ch_max_size_preview_ckb") + skip_nsfw_preview_ckb = gr.Checkbox(label="SKip NSFW Preview images", value=skip_nsfw_preview, elem_id="ch_skip_nsfw_preview_ckb") + + # 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(): + with gr.Column(): + gr.Markdown("### Get Civitai Model Info by Model ID") + 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 file", value=False, elem_id="cn_empty_info_only_ckb") + model_name_drop = gr.Dropdown(choices=no_info_model_names, label="Model", value="ckp", multiselect=False) + + model_url_or_id = gr.Textbox(label="Civitai URL or Model ID", lines=1, value="") + get_civitai_model_info_by_id_btn = gr.Button(value="Get 1 Model Info from Civitai") + get_model_by_id_log_md = gr.Markdown("") + + + with gr.Tab("General"): + with gr.Row(): + open_url_with_js_ckb = gr.Checkbox(label="Open Url At Client Side", value=open_url_with_js, elem_id="ch_open_url_with_js_ckb") + check_model_version_at_startup_ckb = gr.Checkbox(label="Check Model Version At Startup", value=open_url_with_js, visible=False, elem_id="ch_check_model_version_at_startup_ckb") + + save_setting_btn = gr.Button(value="Save Setting", variant="primary", elem_id="ch_save_setting_btn") + general_log_md = gr.Markdown(value="", elem_id="ch_general_log_md") + + # with gr.Tab("Tool"): + + # 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") + + # ====events==== + # Model + scan_model_civitai_btn.click(model_action_civitai.scan_model, inputs=[low_memory_sha_ckb, max_size_preview_ckb, skip_nsfw_preview_ckb], outputs=scan_model_log_md) + + 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_id, inputs=[model_type_drop, model_name_drop, model_url_or_id, max_size_preview_ckb, skip_nsfw_preview_ckb], outputs=get_model_by_id_log_md) + + + # General + save_setting_btn.click(setting.save_from_input, inputs=[low_memory_sha_ckb, max_size_preview_ckb, skip_nsfw_preview_ckb, open_url_with_js_ckb, check_model_version_at_startup_ckb], outputs=general_log_md) + + # js action + js_open_url_btn.click(js_action_civitai.open_model_url, inputs=[js_msg_txtbox, open_url_with_js_ckb], outputs=py_msg_txtbox) + js_add_trigger_words_btn.click(js_action_civitai.add_trigger_words, inputs=[js_msg_txtbox], outputs=[txt2img_prompt, img2img_prompt]) + js_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]) + + # 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) + +