From b5ac40fbdcd15d12580581457c0d2600a25b186b Mon Sep 17 00:00:00 2001 From: Sourcery AI <> Date: Mon, 22 May 2023 22:53:33 +0000 Subject: [PATCH] 'Refactored by Sourcery' --- scripts/ch_lib/civitai.py | 137 ++++++++++----------- scripts/ch_lib/downloader.py | 20 ++-- scripts/ch_lib/js_action_civitai.py | 42 +++---- scripts/ch_lib/model.py | 12 +- scripts/ch_lib/model_action_civitai.py | 160 ++++++++++++------------- scripts/ch_lib/msg_handler.py | 6 +- scripts/ch_lib/setting.py | 8 +- scripts/ch_lib/util.py | 26 ++-- 8 files changed, 201 insertions(+), 210 deletions(-) diff --git a/scripts/ch_lib/civitai.py b/scripts/ch_lib/civitai.py index 2a15ae7..ab4af45 100644 --- a/scripts/ch_lib/civitai.py +++ b/scripts/ch_lib/civitai.py @@ -32,7 +32,7 @@ model_type_dict = { # width is in number, not string # return: url str def get_full_size_image_url(image_url, width): - return re.sub('/width=\d+/', '/width=' + str(width) + '/', image_url) + return re.sub('/width=\d+/', f'/width={str(width)}/', image_url) # use this sha256 to get model info from civitai @@ -51,7 +51,7 @@ def get_model_info_by_hash(hash:str): util.printD("Civitai does not have this model") return {} else: - util.printD("Get error code: " + str(r.status_code)) + util.printD(f"Get error code: {r.status_code}") util.printD(r.text) return @@ -65,11 +65,11 @@ def get_model_info_by_hash(hash:str): util.printD("response:") util.printD(r.text) return - + if not content: util.printD("error, content from civitai is None") return - + return content @@ -81,14 +81,18 @@ def get_model_info_by_id(id:str) -> dict: util.printD("id is empty") return - r = requests.get(url_dict["modelId"]+str(id), headers=util.def_headers, proxies=util.proxies) + r = requests.get( + url_dict["modelId"] + id, + headers=util.def_headers, + proxies=util.proxies, + ) if not r.ok: if r.status_code == 404: # this is not a civitai model util.printD("Civitai does not have this model") return {} else: - util.printD("Get error code: " + str(r.status_code)) + util.printD(f"Get error code: {r.status_code}") util.printD(r.text) return @@ -102,11 +106,11 @@ def get_model_info_by_id(id:str) -> dict: util.printD("response:") util.printD(r.text) return - + if not content: util.printD("error, content from civitai is None") return - + return content @@ -117,14 +121,18 @@ def get_version_info_by_version_id(id:str) -> dict: util.printD("id is empty") return - r = requests.get(url_dict["modelVersionId"]+str(id), headers=util.def_headers, proxies=util.proxies) + r = requests.get( + url_dict["modelVersionId"] + id, + headers=util.def_headers, + proxies=util.proxies, + ) if not r.ok: if r.status_code == 404: # this is not a civitai model util.printD("Civitai does not have this model version") return {} else: - util.printD("Get error code: " + str(r.status_code)) + util.printD(f"Get error code: {r.status_code}") util.printD(r.text) return @@ -138,11 +146,11 @@ def get_version_info_by_version_id(id:str) -> dict: util.printD("response:") util.printD(r.text) return - + if not content: util.printD("error, content from civitai is None") return - + return content @@ -198,9 +206,9 @@ def get_version_info_by_model_id(id:str) -> dict: def load_model_info_by_search_term(model_type, search_term): util.printD(f"Load model info of {search_term} in {model_type}") if model_type not in model.folders.keys(): - util.printD("unknow model type: " + model_type) + util.printD(f"unknow model type: {model_type}") return - + # search_term = subfolderpath + model name + ext. And it always start with a / even there is no sub folder base, ext = os.path.splitext(search_term) model_info_base = base @@ -212,9 +220,9 @@ def load_model_info_by_search_term(model_type, search_term): model_info_filepath = os.path.join(model_folder, model_info_filename) if not os.path.isfile(model_info_filepath): - util.printD("Can not find model info file: " + model_info_filepath) + util.printD(f"Can not find model info file: {model_info_filepath}") return - + return model.load_model_info(model_info_filepath) @@ -235,9 +243,9 @@ def get_model_names_by_type_and_filter(model_type:str, filter:dict) -> list: empty_info_only = False if filter: - if "no_info_only" in filter.keys(): + if "no_info_only" in filter: no_info_only = filter["no_info_only"] - if "empty_info_only" in filter.keys(): + if "empty_info_only" in filter: empty_info_only = filter["empty_info_only"] @@ -264,10 +272,7 @@ def get_model_names_by_type_and_filter(model_type:str, filter:dict) -> list: # 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 model_info := model.load_model_info(info_file): if "id" in model_info.keys(): # find a non-empty model info file continue @@ -291,15 +296,12 @@ def get_model_id_from_url(url:str) -> str: return "" if url.isnumeric(): - # is already an id - id = str(url) - return id - + return url s = re.sub("\\?.+$", "", url).split("/") if len(s) < 2: util.printD("url is not valid") return "" - + if s[-2].isnumeric(): id = s[-2] elif s[-1].isnumeric(): @@ -307,7 +309,7 @@ def get_model_id_from_url(url:str) -> str: else: util.printD("There is no model id in this url") return "" - + return id @@ -319,18 +321,18 @@ def get_preview_image_by_model_path(model_path:str, max_size_preview, skip_nsfw_ return if not os.path.isfile(model_path): - util.printD("model_path is not a file: "+model_path) + util.printD(f"model_path is not a file: {model_path}") return base, ext = os.path.splitext(model_path) - first_preview = base+".png" - sec_preview = base+".preview.png" + first_preview = f"{base}.png" + sec_preview = f"{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) + util.printD(f"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) @@ -347,7 +349,7 @@ def get_preview_image_by_model_path(model_path:str, max_size_preview, skip_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: @@ -366,8 +368,8 @@ def get_preview_image_by_model_path(model_path:str, max_size_preview, skip_nsfw_ # return - model_info def search_local_model_info_by_version_id(folder:str, version_id:int) -> dict: util.printD("Searching local model by version id") - util.printD("folder: " + folder) - util.printD("version_id: " + str(version_id)) + util.printD(f"folder: {folder}") + util.printD(f"version_id: {version_id}") if not folder: util.printD("folder is none") @@ -376,11 +378,11 @@ def search_local_model_info_by_version_id(folder:str, version_id:int) -> dict: if not os.path.isdir(folder): util.printD("folder is not a dir") return - + if not version_id: util.printD("version_id is none") return - + # search civitai model info file for filename in os.listdir(folder): # check ext @@ -409,7 +411,7 @@ def search_local_model_info_by_version_id(folder:str, version_id:int) -> dict: if str(id) == str(version_id): # find the one return model_info - + return @@ -425,16 +427,16 @@ def check_model_new_version_by_path(model_path:str, delay:float=1) -> tuple: return if not os.path.isfile(model_path): - util.printD("model_path is not a file: "+model_path) + util.printD(f"model_path is not a file: {model_path}") return - + # get model info file name base, ext = os.path.splitext(model_path) info_file = base + suffix + model.info_ext - + if not os.path.isfile(info_file): return - + # get model info model_info_file = model.load_model_info(info_file) if not model_info_file: @@ -442,18 +444,18 @@ def check_model_new_version_by_path(model_path:str, delay:float=1) -> tuple: if "id" not in model_info_file.keys(): return - + local_version_id = model_info_file["id"] if not local_version_id: return if "modelId" not in model_info_file.keys(): return - + model_id = model_info_file["modelId"] if not model_id: return - + # get model info by id from civitai model_info = get_model_info_by_id(model_id) # delay before next request, to prevent to be treat as DDoS @@ -462,24 +464,24 @@ def check_model_new_version_by_path(model_path:str, delay:float=1) -> tuple: if not model_info: return - + if "modelVersions" not in model_info.keys(): return - + modelVersions = model_info["modelVersions"] if not modelVersions: return - + if not len(modelVersions): return - + current_version = modelVersions[0] if not current_version: return - + if "id" not in current_version.keys(): return - + current_version_id = current_version["id"] if not current_version_id: return @@ -488,10 +490,7 @@ def check_model_new_version_by_path(model_path:str, delay:float=1) -> tuple: if current_version_id == local_version_id: return - model_name = "" - if "name" in model_info.keys(): - model_name = model_info["name"] - + model_name = model_info["name"] if "name" in model_info.keys() else "" if not model_name: model_name = "" @@ -499,21 +498,21 @@ def check_model_new_version_by_path(model_path:str, delay:float=1) -> tuple: new_version_name = "" if "name" in current_version.keys(): new_version_name = current_version["name"] - + if not new_version_name: new_version_name = "" description = "" if "description" in current_version.keys(): description = current_version["description"] - + if not description: description = "" downloadUrl = "" if "downloadUrl" in current_version.keys(): downloadUrl = current_version["downloadUrl"] - + if not downloadUrl: downloadUrl = "" @@ -528,7 +527,7 @@ def check_model_new_version_by_path(model_path:str, delay:float=1) -> tuple: img_url = "" - + return (model_path, model_id, model_name, current_version_id, new_version_name, description, downloadUrl, img_url) @@ -564,7 +563,7 @@ def check_models_new_version_by_model_types(model_types:list, delay:float=1) -> if model_type not in mts: continue - util.printD("Scanning path: " + model_folder) + util.printD(f"Scanning path: {model_folder}") for root, dirs, files in os.walk(model_folder, followlinks=True): for filename in files: # check ext @@ -582,21 +581,17 @@ def check_models_new_version_by_model_types(model_types:list, delay:float=1) -> if not current_version_id: continue - # check this version id in list - is_already_in_list = False - for new_version in new_versions: - if current_version_id == new_version[3]: - # already in list - is_already_in_list = True - break - + is_already_in_list = any( + current_version_id == new_version[3] + for new_version in new_versions + ) if is_already_in_list: util.printD("New version is already in list") continue - # search this new version id to check if this model is already downloaded - target_model_info = search_local_model_info_by_version_id(root, current_version_id) - if target_model_info: + if target_model_info := search_local_model_info_by_version_id( + root, current_version_id + ): util.printD("New version is already existed") continue diff --git a/scripts/ch_lib/downloader.py b/scripts/ch_lib/downloader.py index 2b14ebc..28e5caa 100644 --- a/scripts/ch_lib/downloader.py +++ b/scripts/ch_lib/downloader.py @@ -12,7 +12,7 @@ requests.packages.urllib3.disable_warnings() # output is downloaded file path def dl(url, folder, filename, filepath): - util.printD("Start downloading from: " + url) + util.printD(f"Start downloading from: {url}") # get file_path file_path = "" if filepath: @@ -22,11 +22,11 @@ def dl(url, folder, filename, filepath): if not folder: util.printD("folder is none") return - + if not os.path.isdir(folder): - util.printD("folder does not exist: "+folder) + util.printD(f"folder does not exist: {folder}") return - + if filename: file_path = os.path.join(folder, filename) @@ -40,25 +40,25 @@ def dl(url, folder, filename, filepath): # if file_path is empty, need to get file name from download url's header if not file_path: filename = "" - if "Content-Disposition" in rh.headers.keys(): + if "Content-Disposition" in rh.headers: cd = rh.headers["Content-Disposition"] # Extract the filename from the header # content of a CD: "attachment;filename=FileName.txt" # in case "" is in CD filename's start and end, need to strip them out filename = cd.split("=")[1].strip('"') if not filename: - util.printD("Fail to get file name from Content-Disposition: " + cd) + util.printD(f"Fail to get file name from Content-Disposition: {cd}") return - + if not filename: util.printD("Can not get file name from download url's header") return - + # with folder and filename, now we have the full file path file_path = os.path.join(folder, filename) - util.printD("Target file path: " + file_path) + util.printD(f"Target file path: {file_path}") base, ext = os.path.splitext(file_path) # check if file is already exist @@ -67,7 +67,7 @@ def dl(url, folder, filename, filepath): while os.path.isfile(file_path): util.printD("Target file already exist.") # re-name - new_base = base + "_" + str(count) + new_base = f"{base}_{str(count)}" file_path = new_base + ext count += 1 diff --git a/scripts/ch_lib/js_action_civitai.py b/scripts/ch_lib/js_action_civitai.py index 0c81853..a9afd42 100644 --- a/scripts/ch_lib/js_action_civitai.py +++ b/scripts/ch_lib/js_action_civitai.py @@ -23,7 +23,7 @@ def open_model_url(msg, open_url_with_js): if not result: util.printD("Parsing js ms failed") return - + model_type = result["model_type"] search_term = result["search_term"] @@ -50,7 +50,7 @@ def open_model_url(msg, open_url_with_js): } if not open_url_with_js: - util.printD("Open Url: " + url) + util.printD(f"Open Url: {url}") # open url webbrowser.open_new_tab(url) else: @@ -73,7 +73,7 @@ def add_trigger_words(msg): if not result: util.printD("Parsing js ms failed") return - + model_type = result["model_type"] search_term = result["search_term"] prompt = result["prompt"] @@ -83,29 +83,29 @@ def add_trigger_words(msg): if not model_info: util.printD(f"Failed to get model info for {model_type} {search_term}") return [prompt, prompt] - + if "trainedWords" not in model_info.keys(): util.printD(f"Failed to get trainedWords from info file for {model_type} {search_term}") return [prompt, prompt] - + trainedWords = model_info["trainedWords"] if not trainedWords: util.printD(f"No trainedWords from info file for {model_type} {search_term}") return [prompt, prompt] - + if len(trainedWords) == 0: util.printD(f"trainedWords from info file for {model_type} {search_term} is empty") return [prompt, prompt] - + # get ful trigger words trigger_words = "" for word in trainedWords: trigger_words = trigger_words + word + ", " - new_prompt = prompt + " " + trigger_words - util.printD("trigger_words: " + trigger_words) - util.printD("prompt: " + prompt) - util.printD("new_prompt: " + new_prompt) + new_prompt = f"{prompt} {trigger_words}" + util.printD(f"trigger_words: {trigger_words}") + util.printD(f"prompt: {prompt}") + util.printD(f"new_prompt: {new_prompt}") util.printD("End add_trigger_words") @@ -188,14 +188,14 @@ def dl_model_new_version(msg, max_size_preview, skip_nsfw_preview): output = "Parsing js ms failed" util.printD(output) return output - + model_path = result["model_path"] version_id = result["version_id"] download_url = result["download_url"] - util.printD("model_path: " + model_path) - util.printD("version_id: " + str(version_id)) - util.printD("download_url: " + download_url) + util.printD(f"model_path: {model_path}") + util.printD(f"version_id: {str(version_id)}") + util.printD(f"download_url: {download_url}") # check data if not model_path: @@ -207,14 +207,14 @@ def dl_model_new_version(msg, max_size_preview, skip_nsfw_preview): output = "version_id is empty" util.printD(output) return output - + if not download_url: output = "download_url is empty" util.printD(output) return output if not os.path.isfile(model_path): - output = "model_path is not a file: "+ model_path + output = f"model_path is not a file: {model_path}" util.printD(output) return output @@ -232,14 +232,14 @@ def dl_model_new_version(msg, max_size_preview, skip_nsfw_preview): # download file new_model_path = downloader.dl(download_url, model_folder, None, None) if not new_model_path: - output = "Download failed, check console log for detail. Download url: " + download_url + output = f"Download failed, check console log for detail. Download url: {download_url}" util.printD(output) return output # get version info version_info = civitai.get_version_info_by_version_id(version_id) if not version_info: - output = "Model downloaded, but failed to get version info, check console log for detail. Model saved to: " + new_model_path + output = f"Model downloaded, but failed to get version info, check console log for detail. Model saved to: {new_model_path}" util.printD(output) return output @@ -250,7 +250,7 @@ def dl_model_new_version(msg, max_size_preview, skip_nsfw_preview): # then, get preview image civitai.get_preview_image_by_model_path(new_model_path, max_size_preview, skip_nsfw_preview) - - output = "Done. Model downloaded to: " + new_model_path + + output = f"Done. Model downloaded to: {new_model_path}" util.printD(output) return output diff --git a/scripts/ch_lib/model.py b/scripts/ch_lib/model.py index 729012d..1cdc012 100644 --- a/scripts/ch_lib/model.py +++ b/scripts/ch_lib/model.py @@ -48,7 +48,7 @@ def get_custom_model_folder(): # write model info to file def write_model_info(path, model_info): - util.printD("Write model info to file: " + path) + util.printD(f"Write model info to file: {path}") with open(os.path.realpath(path), 'w') as f: f.write(json.dumps(model_info, indent=4)) @@ -60,10 +60,10 @@ def load_model_info(path): try: model_info = json.load(f) except Exception as e: - util.printD("Selected file is not json: " + path) + util.printD(f"Selected file is not json: {path}") util.printD(e) return - + return model_info @@ -94,13 +94,13 @@ def get_model_names_by_type(model_type:str) -> list: 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) + util.printD(f"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. diff --git a/scripts/ch_lib/model_action_civitai.py b/scripts/ch_lib/model_action_civitai.py index f09f075..640ca35 100644 --- a/scripts/ch_lib/model_action_civitai.py +++ b/scripts/ch_lib/model_action_civitai.py @@ -19,14 +19,14 @@ def scan_model(scan_model_types, max_size_preview, skip_nsfw_preview): output = "Model Types is None, can not scan." util.printD(output) return output - + model_types = [] # check type if it is a string if type(scan_model_types) == str: model_types.append(scan_model_types) else: model_types = scan_model_types - + model_count = 0 image_count = 0 # scan_log = "" @@ -34,7 +34,7 @@ def scan_model(scan_model_types, max_size_preview, skip_nsfw_preview): if model_type not in model_types: continue - util.printD("Scanning path: " + model_folder) + util.printD(f"Scanning path: {model_folder}") for root, dirs, files in os.walk(model_folder, followlinks=True): for filename in files: # check ext @@ -45,7 +45,7 @@ def scan_model(scan_model_types, max_size_preview, skip_nsfw_preview): if len(base) > 4: if base[-4:] == model.vae_suffix: # find .vae - util.printD("This is a vae file: " + filename) + util.printD(f"This is a vae file: {filename}") continue # find a model @@ -53,15 +53,15 @@ def scan_model(scan_model_types, max_size_preview, skip_nsfw_preview): info_file = base + civitai.suffix + model.info_ext # check info file if not os.path.isfile(info_file): - util.printD("Creating model info for: " + filename) + util.printD(f"Creating model info for: {filename}") # get model's sha256 hash = util.gen_file_sha256(item) if not hash: - output = "failed generating SHA256 for model:" + filename + output = f"failed generating SHA256 for model:{filename}" util.printD(output) return output - + # use this sha256 to get model info from civitai model_info = civitai.get_model_info_by_hash(hash) # delay 1 second for ti @@ -72,7 +72,7 @@ def scan_model(scan_model_types, max_size_preview, skip_nsfw_preview): if model_info is None: output = "Connect to Civitai API service failed. Wait a while and try again" util.printD(output) - return output+", check console log for detail" + return f"{output}, check console log for detail" # write model info to file model.write_model_info(info_file, model_info) @@ -102,7 +102,7 @@ def get_model_info_by_input(model_type, model_name, model_url_or_id, max_size_pr # parse model id model_id = civitai.get_model_id_from_url(model_url_or_id) if not model_id: - output = "failed to parse model id from url: " + model_url_or_id + output = f"failed to parse model id from url: {model_url_or_id}" util.printD(output) return output @@ -113,13 +113,13 @@ def get_model_info_by_input(model_type, model_name, model_url_or_id, max_size_pr 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 @@ -129,19 +129,19 @@ def get_model_info_by_input(model_type, model_name, model_url_or_id, max_size_pr 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 + output = f"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) + util.printD(f"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 + output = f"Model Info saved to: {info_file}" return output @@ -166,28 +166,28 @@ def check_models_new_version_to_md(model_types:list) -> str: url = civitai.url_dict["modelPage"]+str(model_id) part = f'
Model: {model_name}
' - part = part + f'
File: {model_path}
' + part = f'{part}
File: {model_path}
' if download_url: # replace "\" to "/" in model_path for windows model_path = model_path.replace('\\', '\\\\') - part = part + f'
New Version: {new_version_name}' + part = f'{part}
New Version: {new_version_name}' # add js function to download new version into SD webui by python - part = part + " " + part = f"{part} " # in embed HTML, onclick= will also follow a ", never a ', so have to write it as following - part = part + f"[Download into SD]" - + part = f"""{part}[Download into SD]""" + else: - part = part + f'
New Version: {new_version_name}' - part = part + '
' + part = f'{part}
New Version: {new_version_name}' + part = f'{part}
' # description if description: - part = part + '
'+ description + '

' + part = f'{part}
{description}

' - # preview image + # preview image if img_url: - part = part + f"
" - + part = f"{part}
" + output = output + part @@ -198,7 +198,7 @@ def check_models_new_version_to_md(model_types:list) -> str: # 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) + util.printD(f"Getting model info by: {model_url_or_id}") # parse model id model_id = civitai.get_model_id_from_url(model_url_or_id) @@ -210,22 +210,22 @@ def get_model_info_by_url(model_url_or_id:str) -> tuple: 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) + util.printD(f"This model type is not supported:{civitai_model_type}") return - + model_type = civitai.model_type_dict[civitai_model_type] # get model type @@ -242,20 +242,15 @@ def get_model_info_by_url(model_url_or_id:str) -> tuple: 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) + version_strs = [ + version["name"] + "_" + str(version["id"]) for version in modelVersions + ] # get folder by model type folder = model.folders[model_type] # get subfolders @@ -279,11 +274,11 @@ 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") @@ -293,7 +288,7 @@ def get_ver_info_by_ver_str(version_str:str, model_info:dict) -> dict: if not modelVersions: util.printD("modelVersions is Empty") return - + # find version by version_str version = None for ver in modelVersions: @@ -306,14 +301,14 @@ def get_ver_info_by_ver_str(version_str:str, model_info:dict) -> dict: version = ver if not version: - util.printD("can not find version by version string: " + version_str) + util.printD(f"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 @@ -323,11 +318,11 @@ 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") @@ -337,7 +332,7 @@ def get_id_and_dl_url_by_version_str(version_str:str, model_info:dict) -> tuple: if not modelVersions: util.printD("modelVersions is Empty") return - + # find version by version_str version = None for ver in modelVersions: @@ -350,14 +345,14 @@ def get_id_and_dl_url_by_version_str(version_str:str, model_info:dict) -> tuple: version = ver if not version: - util.printD("can not find version by version string: " + version_str) + util.printD(f"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") @@ -367,14 +362,14 @@ def get_id_and_dl_url_by_version_str(version_str:str, model_info:dict) -> tuple: 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) - + + util.printD(f"Get Download Url: {downloadUrl}") + return (version_id, downloadUrl) @@ -388,36 +383,36 @@ def dl_model_by_input(model_info:dict, model_type:str, subfolder_str:str, versio 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 + output = f"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 == "\\": + if subfolder_str in {"/", "\\"}: subfolder = "" - elif subfolder_str[:1] == "/" or subfolder_str[:1] == "\\": + elif subfolder_str[:1] in ["/", "\\"]: subfolder = subfolder_str[1:] else: subfolder = subfolder_str @@ -425,17 +420,17 @@ def dl_model_by_input(model_info:dict, model_type:str, subfolder_str:str, versio # 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 + output = f"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"] @@ -444,22 +439,23 @@ def dl_model_by_input(model_info:dict, model_type:str, subfolder_str:str, versio # 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"]) - + download_urls.extend( + file_info["downloadUrl"] + for file_info in ver_info["files"] + if "downloadUrl" in file_info.keys() + ) 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: + if r := civitai.search_local_model_info_by_version_id( + model_folder, version_id + ): output = "This model version is already existed" util.printD(output) return output - + # download filepath = "" for url in download_urls: @@ -468,7 +464,7 @@ def dl_model_by_input(model_info:dict, model_type:str, subfolder_str:str, versio output = "Downloading failed, check console log for detail" util.printD(output) return output - + if url == ver_info["downloadUrl"]: filepath = model_filepath else: @@ -479,7 +475,7 @@ def dl_model_by_input(model_info:dict, model_type:str, subfolder_str:str, versio 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: @@ -490,11 +486,11 @@ def dl_model_by_input(model_info:dict, model_type:str, subfolder_str:str, versio 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 + output = f"Model downloaded, but failed to get version info, check console log for detail. Model saved to: {filepath}" util.printD(output) return output @@ -505,7 +501,7 @@ def dl_model_by_input(model_info:dict, model_type:str, subfolder_str:str, versio # then, get preview image civitai.get_preview_image_by_model_path(filepath, max_size_preview, skip_nsfw_preview) - - output = "Done. Model downloaded to: " + filepath + + output = f"Done. Model downloaded to: {filepath}" util.printD(output) return output diff --git a/scripts/ch_lib/msg_handler.py b/scripts/ch_lib/msg_handler.py index 1132d77..fcb6bf3 100644 --- a/scripts/ch_lib/msg_handler.py +++ b/scripts/ch_lib/msg_handler.py @@ -29,7 +29,7 @@ def parse_js_msg(msg): return if action not in js_actions: - util.printD("Unknow action: " + action) + util.printD(f"Unknow action: {action}") return util.printD("End parse js msg") @@ -45,13 +45,13 @@ def build_py_msg(action:str, content:dict): 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) + util.printD(f"Unknow action: {action}") return msg = { diff --git a/scripts/ch_lib/setting.py b/scripts/ch_lib/setting.py index aabff85..ec82a03 100644 --- a/scripts/ch_lib/setting.py +++ b/scripts/ch_lib/setting.py @@ -29,7 +29,7 @@ data = { # save setting # return output msg for log def save(): - print("Saving setting to: " + path) + print(f"Saving setting to: {path}") json_data = json.dumps(data, indent=4) @@ -40,12 +40,12 @@ def save(): with open(path, 'w') as f: f.write(json_data) except Exception as e: - util.printD("Error when writing file:"+path) + util.printD(f"Error when writing file:{path}") output = str(e) util.printD(str(e)) return output - output = "Setting saved to: " + path + output = f"Setting saved to: {path}" util.printD(output) return output @@ -56,7 +56,7 @@ def load(): # load data into globel data global data - util.printD("Load setting from: " + path) + util.printD(f"Load setting from: {path}") if not os.path.isfile(path): util.printD("No setting file, use default") diff --git a/scripts/ch_lib/util.py b/scripts/ch_lib/util.py index 3e43c98..acc0ec4 100644 --- a/scripts/ch_lib/util.py +++ b/scripts/ch_lib/util.py @@ -22,10 +22,10 @@ def printD(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: + if chunk := file.read(size): + yield chunk + else: break - yield chunk # Now, hashing use the same way as pip's source code. def gen_file_sha256(filname): @@ -39,40 +39,40 @@ def gen_file_sha256(filname): h.update(block) hash_value = h.hexdigest() - printD("sha256: " + hash_value) - printD("length: " + str(length)) + printD(f"sha256: {hash_value}") + printD(f"length: {str(length)}") return hash_value # get preview image def download_file(url, path): - printD("Downloading file from: " + url) + printD(f"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(f"Get error code: {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) + printD(f"File downloaded to: {path}") # get subfolder list def get_subfolders(folder:str) -> list: - printD("Get subfolder for: " + folder) + printD(f"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): @@ -98,7 +98,7 @@ def get_relative_path(item_path:str, parent_path:str) -> str: return item_path relative = item_path[len(parent_path):] - if relative[:1] == "/" or relative[:1] == "\\": + if relative[:1] in ["/", "\\"]: relative = relative[1:] # printD("relative:"+relative)