Hack around Lora and Lycoris being treated as the same by webui

pull/232/head
zixaphir 2023-07-28 15:28:02 -07:00
parent 06c7dbef02
commit ccb16925b3
2 changed files with 82 additions and 51 deletions

View File

@ -207,11 +207,22 @@ def load_model_info_by_search_term(model_type, search_term):
if base[:1] == "/": if base[:1] == "/":
model_info_base = base[1:] model_info_base = base[1:]
model_folder = model.folders[model_type] if model_type == "lora" and model.folders['lycoris']:
model_info_filename = model_info_base + suffix + model.info_ext model_folders = [model.folders[model_type], model.folders['lycoris']]
model_info_filepath = os.path.join(model_folder, model_info_filename) else:
model_folders = [model.folders[model_type]]
if not os.path.isfile(model_info_filepath): #model_folder = model.folders[model_type]
for model_folder in model_folders:
model_info_filename = model_info_base + suffix + model.info_ext
model_info_filepath = os.path.join(model_folder, model_info_filename)
found = os.path.isfile(model_info_filepath)
if found:
break;
if not found:
util.printD("Can not find model info file: " + model_info_filepath) util.printD("Can not find model info file: " + model_info_filepath)
return return
@ -227,7 +238,10 @@ def load_model_info_by_search_term(model_type, search_term):
# return: model name list # return: model name list
def get_model_names_by_type_and_filter(model_type:str, filter:dict) -> list: def get_model_names_by_type_and_filter(model_type:str, filter:dict) -> list:
model_folder = model.folders[model_type] if model_type == "lora" and model.folders['lycoris']:
model_folders = [model.folders[model_type], model.folders['lycoris']]
else:
model_folders = [model.folders[model_type]]
# set filter # set filter
# only get models don't have a civitai info file # only get models don't have a civitai info file
@ -245,35 +259,35 @@ def get_model_names_by_type_and_filter(model_type:str, filter:dict) -> list:
# get information from filter # get information from filter
# only get those model names don't have a civitai model info file # only get those model names don't have a civitai model info file
model_names = [] model_names = []
for root, dirs, files in os.walk(model_folder, followlinks=True): for model_folder in model_folders:
for filename in files: for root, dirs, files in os.walk(model_folder, followlinks=True):
item = os.path.join(root, filename) for filename in files:
# check extension item = os.path.join(root, filename)
base, ext = os.path.splitext(item) # check extension
if ext in model.exts: base, ext = os.path.splitext(item)
# find a model if ext in model.exts:
# find a model
# check filter # check filter
if no_info_only: if no_info_only:
# check model info file # check model info file
info_file = base + suffix + model.info_ext info_file = base + suffix + model.info_ext
if os.path.isfile(info_file): if os.path.isfile(info_file):
continue continue
if empty_info_only: if empty_info_only:
# check model info file # check model info file
info_file = base + suffix + model.info_ext info_file = base + suffix + model.info_ext
if os.path.isfile(info_file): if os.path.isfile(info_file):
# load model info # load model info
model_info = model.load_model_info(info_file) model_info = model.load_model_info(info_file)
# check content # check content
if model_info: if model_info:
if "id" in model_info.keys(): if "id" in model_info.keys():
# find a non-empty model info file # find a non-empty model info file
continue continue
model_names.append(filename)
model_names.append(filename)
return model_names return model_names

View File

@ -44,10 +44,18 @@ def get_custom_model_folder():
folders["lora"] = shared.cmd_opts.lora_dir folders["lora"] = shared.cmd_opts.lora_dir
try: try:
# pre-1.5.0
if shared.cmd_opts.lyco_dir and os.path.isdir(shared.cmd_opts.lyco_dir): if shared.cmd_opts.lyco_dir and os.path.isdir(shared.cmd_opts.lyco_dir):
folders["lycoris"] = shared.cmd_opts.lyco_dir folders["lycoris"] = shared.cmd_opts.lyco_dir
except: except:
pass # XXX sd-webui v1.5.0 handles the lyco directory differently. try:
# sd-webui v1.5.1 added a backcompat option for lyco.
if shared.cmd_opts.lyco_dir_backcompat and os.path.isdir(shared.cmd_opts.lyco_dir_backcompat):
folders["lycoris"] = shared.cmd_opts.lyco_dir_backcompat
except:
# XXX v1.5.0 has no options for the Lyco dir: it is hardcoded as 'os.path.join(paths.models_path, "LyCORIS")'
pass
# write model info to file # write model info to file
@ -75,20 +83,24 @@ def load_model_info(path):
# parameter: model_type - string # parameter: model_type - string
# return: model name list # return: model name list
def get_model_names_by_type(model_type:str) -> list: def get_model_names_by_type(model_type:str) -> list:
model_folder = folders[model_type] if model_type == "lora" and folders['lycoris']:
model_folders = [folders[model_type], folders['lycoris']]
else:
model_folders = [folders[model_type]]
# get information from filter # get information from filter
# only get those model names don't have a civitai model info file # only get those model names don't have a civitai model info file
model_names = [] model_names = []
for root, dirs, files in os.walk(model_folder, followlinks=True): for model_folder in model_folders:
for filename in files: for root, dirs, files in os.walk(model_folder, followlinks=True):
item = os.path.join(root, filename) for filename in files:
# check extension item = os.path.join(root, filename)
base, ext = os.path.splitext(item) # check extension
if ext in exts: base, ext = os.path.splitext(item)
# find a model if ext in exts:
model_names.append(filename) # find a model
model_names.append(filename)
return model_names return model_names
@ -104,19 +116,24 @@ def get_model_path_by_type_and_name(model_type:str, model_name:str) -> str:
if not model_name: if not model_name:
util.printD("model name can not be empty") util.printD("model name can not be empty")
return return
folder = folders[model_type] if model_type == "lora" and folders['lycoris']:
model_folders = [folders[model_type], folders['lycoris']]
else:
model_folders = [folders[model_type]]
# model could be in subfolder, need to walk. # model could be in subfolder, need to walk.
model_root = "" model_root = ""
model_path = "" model_path = ""
for root, dirs, files in os.walk(folder, followlinks=True): for folder in model_folders:
for filename in files: for root, dirs, files in os.walk(folder, followlinks=True):
if filename == model_name: for filename in files:
# find model if filename == model_name:
model_root = root # find model
model_path = os.path.join(root, filename) model_root = root
return (model_root, model_path) model_path = os.path.join(root, filename)
return (model_root, model_path)
return return