1580158
/
Stable-Diffusion-Webui-Civitai-Helper
mirror of https://github.com/butaixianran/Stable-Diffusion-Webui-Civitai-Helper.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

'Refactored by Sourcery'

pull/232/head^2^2
Sourcery AI 2023-08-05 23:18:45 +00:00
parent 2d95b6fcf9
commit b43ab281f0
2 changed files with 71 additions and 76 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

77
scripts/ch_lib/civitai.py
View File

@ -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
@ -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
@ -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
@ -198,7 +206,7 @@ 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
@ -223,7 +231,7 @@ def load_model_info_by_search_term(model_type, search_term):
break;
if not found:
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)
@ -249,9 +257,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"]
@ -279,10 +287,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
@ -305,10 +310,7 @@ 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")
@ -333,18 +335,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)
@ -380,8 +382,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")
@ -439,7 +441,7 @@ 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
@ -502,10 +504,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 = ""
@ -578,7 +577,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
@ -596,21 +595,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

6
scripts/ch_lib/model.py
View File

@ -60,7 +60,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))
@ -72,7 +72,7 @@ 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
@ -110,7 +110,7 @@ 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: