From ccb16925b320e695387666ffbbe058fcb15fbcf0 Mon Sep 17 00:00:00 2001 From: zixaphir Date: Fri, 28 Jul 2023 15:28:02 -0700 Subject: [PATCH] Hack around Lora and Lycoris being treated as the same by webui --- scripts/ch_lib/civitai.py | 76 +++++++++++++++++++++++---------------- scripts/ch_lib/model.py | 57 ++++++++++++++++++----------- 2 files changed, 82 insertions(+), 51 deletions(-) diff --git a/scripts/ch_lib/civitai.py b/scripts/ch_lib/civitai.py index 9002dcd..0ca3150 100644 --- a/scripts/ch_lib/civitai.py +++ b/scripts/ch_lib/civitai.py @@ -207,11 +207,22 @@ def load_model_info_by_search_term(model_type, search_term): if base[:1] == "/": model_info_base = base[1:] - model_folder = model.folders[model_type] - model_info_filename = model_info_base + suffix + model.info_ext - model_info_filepath = os.path.join(model_folder, model_info_filename) + if model_type == "lora" and model.folders['lycoris']: + model_folders = [model.folders[model_type], model.folders['lycoris']] + 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) return @@ -227,7 +238,10 @@ def load_model_info_by_search_term(model_type, search_term): # return: model name 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 # 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 # only get those model names don't have a civitai model info file model_names = [] - for root, dirs, files in os.walk(model_folder, followlinks=True): - 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 + for model_folder in model_folders: + for root, dirs, files in os.walk(model_folder, followlinks=True): + 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 + # 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) + 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 diff --git a/scripts/ch_lib/model.py b/scripts/ch_lib/model.py index 6945978..22bd0ac 100644 --- a/scripts/ch_lib/model.py +++ b/scripts/ch_lib/model.py @@ -44,10 +44,18 @@ def get_custom_model_folder(): folders["lora"] = shared.cmd_opts.lora_dir try: + # pre-1.5.0 if shared.cmd_opts.lyco_dir and os.path.isdir(shared.cmd_opts.lyco_dir): folders["lycoris"] = shared.cmd_opts.lyco_dir + 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 @@ -75,20 +83,24 @@ def load_model_info(path): # parameter: model_type - string # return: model name 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 # only get those model names don't have a civitai model info file model_names = [] - for root, dirs, files in os.walk(model_folder, followlinks=True): - 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) + for model_folder in model_folders: + for root, dirs, files in os.walk(model_folder, followlinks=True): + 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 @@ -104,19 +116,24 @@ def get_model_path_by_type_and_name(model_type:str, model_name:str) -> str: if not model_name: util.printD("model name can not be empty") 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_root = "" model_path = "" - for root, dirs, files in os.walk(folder, followlinks=True): - 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) + for folder in model_folders: + for root, dirs, files in os.walk(folder, followlinks=True): + 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