New Version: {new_version_name}'
+ part = part + '
'
+
+ # description
+ if description:
+ part = part + ''+ description + '
' + + # preview image + if img_url: + part = part + f"
" + + + output = output + part + + util.printD(f"Done. Find {count} models have new version. Check UI for detail.") + + return output + + +# 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 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 + + 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] + + # get model type + if "name" not in model_info.keys(): + util.printD("model name is not in model_info") + return + + model_name = model_info["name"] + if not model_name: + util.printD("model name is Empty") + model_name = "" + + # 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 + + version_strs = [] + for version in modelVersions: + # 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_strs.append(version_str) + + # 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}") + + 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 + + modelVersions = model_info["modelVersions"] + if not modelVersions: + util.printD("modelVersions is Empty") + return + + # find version by version_str + version = None + for ver in modelVersions: + # 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"]) + if ver_str == version_str: + # find version + version = ver + + 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 + + modelVersions = model_info["modelVersions"] + if not modelVersions: + util.printD("modelVersions is Empty") + return + + # find version by version_str + version = None + for ver in modelVersions: + # 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"]) + if ver_str == version_str: + # find version + version = ver + + 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 + + # 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) + + +# 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: + + output = "" + + if not model_info: + output = "model_info is None" + util.printD(output) + return output + + if not model_type: + output = "model_type is None" + util.printD(output) + return output + + if not subfolder_str: + output = "subfolder string is None" + util.printD(output) + return output + + if not version_str: + output = "version_str is None" + util.printD(output) + return output + + # get model root folder + if model_type not in model.folders.keys(): + output = "Unknow model type: "+model_type + util.printD(output) + return output + + model_root_folder = model.folders[model_type] + + + # get subfolder + subfolder = "" + if subfolder_str == "/" or subfolder_str == "\\": + subfolder = "" + elif subfolder_str[:1] == "/" or subfolder_str[:1] == "\\": + 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 + util.printD(output) + return output + + # 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" + util.printD(output) + return output + + version_id = ver_info["id"] + + + 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"]) + + if not len(download_urls): + if "downloadUrl" in ver_info.keys(): + download_urls.append(ver_info["downloadUrl"]) + + + # 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 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 + util.printD(output) + return output diff --git a/Stable-Diffusion-Webui-Civitai-Helper/scripts/ch_lib/msg_handler.py b/Stable-Diffusion-Webui-Civitai-Helper/scripts/ch_lib/msg_handler.py new file mode 100644 index 0000000000000000000000000000000000000000..1132d775f0f03547f0291138dbe53d970f2fe869 --- /dev/null +++ b/Stable-Diffusion-Webui-Civitai-Helper/scripts/ch_lib/msg_handler.py @@ -0,0 +1,64 @@ +# -*- 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") +py_actions = ("open_url") + + +# handle request from javascript +# parameter: msg - msg from js as string in a hidden textbox +# return: dict for result +def parse_js_msg(msg): + util.printD("Start parse js msg") + msg_dict = json.loads(msg) + + # in case client side run JSON.stringify twice + if (type(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"] + if not action: + util.printD("Action from js request is None") + return + + if action not in js_actions: + util.printD("Unknow action: " + action) + return + + 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): + 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 + + msg = { + "action" : action, + "content": content + } + + + util.printD("End build_msg") + return json.dumps(msg) \ No newline at end of file diff --git a/Stable-Diffusion-Webui-Civitai-Helper/scripts/ch_lib/setting.py b/Stable-Diffusion-Webui-Civitai-Helper/scripts/ch_lib/setting.py new file mode 100644 index 0000000000000000000000000000000000000000..aabff852d5fb989840f2aefc05867bce8d673994 --- /dev/null +++ b/Stable-Diffusion-Webui-Civitai-Helper/scripts/ch_lib/setting.py @@ -0,0 +1,113 @@ +# -*- 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, + "always_display": False, + "show_btn_on_thumb": 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 "always_display" not in data["general"].keys(): + data["general"]["always_display"] = False + + if "show_btn_on_thumb" not in data["general"].keys(): + data["general"]["show_btn_on_thumb"] = True + + 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, always_display, show_btn_on_thumb, 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, + "always_display": always_display, + "show_btn_on_thumb": show_btn_on_thumb, + "proxy": proxy, + }, + "tool":{ + } + } + + output = save() + + if not output: + output = "" + + return output + diff --git a/Stable-Diffusion-Webui-Civitai-Helper/scripts/ch_lib/util.py b/Stable-Diffusion-Webui-Civitai-Helper/scripts/ch_lib/util.py new file mode 100644 index 0000000000000000000000000000000000000000..3e43c98d01c4c62a64fbdcbeaafdce051ddd4d36 --- /dev/null +++ b/Stable-Diffusion-Webui-Civitai-Helper/scripts/ch_lib/util.py @@ -0,0 +1,105 @@ +# -*- coding: UTF-8 -*- +import os +import io +import hashlib +import requests +import shutil + + +version = "1.6.4" + +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'} + + +proxies = None + + +# print for debugging +def printD(msg): + 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.""" + 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) + if not folder: + printD("folder can not be None") + return + + if not os.path.isdir(folder): + printD("path is not a folder") + return + + prefix_len = len(folder) + 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) + + return subfolders + + +# 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: + return "" + if not item_path.startswith(parent_path): + return item_path + + relative = item_path[len(parent_path):] + if relative[:1] == "/" or relative[:1] == "\\": + relative = relative[1:] + + # printD("relative:"+relative) + return relative \ No newline at end of file diff --git a/Stable-Diffusion-Webui-Civitai-Helper/scripts/civitai_helper.py b/Stable-Diffusion-Webui-Civitai-Helper/scripts/civitai_helper.py new file mode 100644 index 0000000000000000000000000000000000000000..ad9e70a1cff66b38085691765753c5dd5dc0677d --- /dev/null +++ b/Stable-Diffusion-Webui-Civitai-Helper/scripts/civitai_helper.py @@ -0,0 +1,218 @@ +# -*- 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.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 setting +from scripts.ch_lib import civitai +from scripts.ch_lib import util + + +# init + +# root path +root_path = os.getcwd() + +# extension path +extension_path = scripts.basedir() + +model.get_custom_model_folder() +setting.load() + +# set proxy +if setting.data["general"]["proxy"]: + util.printD("Set Proxy: "+setting.data["general"]["proxy"]) + util.proxies = { + "http": setting.data["general"]["proxy"], + "https": setting.data["general"]["proxy"], + } + + + + +def on_ui_tabs(): + # init + # init_py_msg = { + # # relative extension path + # "extension_path": util.get_relative_path(extension_path, root_path), + # } + # init_py_msg_str = json.dumps(init_py_msg) + + + # get prompt textarea + # 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) + + + 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)] + + # ====UI==== + with gr.Blocks(analytics_enabled=False) as civitai_helper: + # with gr.Blocks(css=".block.padded {padding: 10px !important}") as civitai_helper: + + # init + 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"] + always_display = setting.data["general"]["always_display"] + show_btn_on_thumb = setting.data["general"]["show_btn_on_thumb"] + proxy = setting.data["general"]["proxy"] + + 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.Box(elem_classes="ch_box"): + with gr.Column(): + gr.Markdown("### Scan Models for Civitai") + with gr.Row(): + 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") + scan_model_types_ckbg = gr.CheckboxGroup(choices=model_types, label="Model Types", value=model_types) + + # 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("") + + 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") + + 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") + + with gr.Box(elem_classes="ch_box"): + with gr.Column(): + gr.Markdown("### Other Setting") + 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") + always_display_ckb = gr.Checkbox(label="Always Display Buttons", value=always_display, elem_id="ch_always_display_ckb") + show_btn_on_thumb_ckb = gr.Checkbox(label="Show Button On Thumb Mode", value=show_btn_on_thumb, elem_id="ch_show_btn_on_thumb_ckb") + + proxy_txtbox = gr.Textbox(label="Proxy", interactive=True, lines=1, value=proxy, info="format: http://127.0.0.1:port") + + save_setting_btn = gr.Button(value="Save Setting") + general_log_md = gr.Markdown(value="") + + + # ====Footer==== + gr.Markdown(f"