brunvelop/ComfyUI

1import os2import time3 4supported_pt_extensions = set(['.ckpt', '.pt', '.bin', '.pth', '.safetensors'])5 6folder_names_and_paths = {}7 8base_path = os.path.dirname(os.path.realpath(__file__))9models_dir = os.path.join(base_path, "models")10folder_names_and_paths["checkpoints"] = ([os.path.join(models_dir, "checkpoints")], supported_pt_extensions)11folder_names_and_paths["configs"] = ([os.path.join(models_dir, "configs")], [".yaml"])12 13folder_names_and_paths["loras"] = ([os.path.join(models_dir, "loras")], supported_pt_extensions)14folder_names_and_paths["vae"] = ([os.path.join(models_dir, "vae")], supported_pt_extensions)15folder_names_and_paths["clip"] = ([os.path.join(models_dir, "clip")], supported_pt_extensions)16folder_names_and_paths["unet"] = ([os.path.join(models_dir, "unet")], supported_pt_extensions)17folder_names_and_paths["clip_vision"] = ([os.path.join(models_dir, "clip_vision")], supported_pt_extensions)18folder_names_and_paths["style_models"] = ([os.path.join(models_dir, "style_models")], supported_pt_extensions)19folder_names_and_paths["embeddings"] = ([os.path.join(models_dir, "embeddings")], supported_pt_extensions)20folder_names_and_paths["diffusers"] = ([os.path.join(models_dir, "diffusers")], ["folder"])21folder_names_and_paths["vae_approx"] = ([os.path.join(models_dir, "vae_approx")], supported_pt_extensions)22 23folder_names_and_paths["controlnet"] = ([os.path.join(models_dir, "controlnet"), os.path.join(models_dir, "t2i_adapter")], supported_pt_extensions)24folder_names_and_paths["gligen"] = ([os.path.join(models_dir, "gligen")], supported_pt_extensions)25 26folder_names_and_paths["upscale_models"] = ([os.path.join(models_dir, "upscale_models")], supported_pt_extensions)27 28folder_names_and_paths["custom_nodes"] = ([os.path.join(base_path, "custom_nodes")], [])29 30folder_names_and_paths["hypernetworks"] = ([os.path.join(models_dir, "hypernetworks")], supported_pt_extensions)31 32folder_names_and_paths["classifiers"] = ([os.path.join(models_dir, "classifiers")], {""})33 34output_directory = os.path.join(os.path.dirname(os.path.realpath(__file__)), "output")35temp_directory = os.path.join(os.path.dirname(os.path.realpath(__file__)), "temp")36input_directory = os.path.join(os.path.dirname(os.path.realpath(__file__)), "input")37 38filename_list_cache = {}39 40if not os.path.exists(input_directory):41    os.makedirs(input_directory)42 43def set_output_directory(output_dir):44    global output_directory45    output_directory = output_dir46 47def set_temp_directory(temp_dir):48    global temp_directory49    temp_directory = temp_dir50 51def set_input_directory(input_dir):52    global input_directory53    input_directory = input_dir54 55def get_output_directory():56    global output_directory57    return output_directory58 59def get_temp_directory():60    global temp_directory61    return temp_directory62 63def get_input_directory():64    global input_directory65    return input_directory66 67 68#NOTE: used in http server so don't put folders that should not be accessed remotely69def get_directory_by_type(type_name):70    if type_name == "output":71        return get_output_directory()72    if type_name == "temp":73        return get_temp_directory()74    if type_name == "input":75        return get_input_directory()76    return None77 78 79# determine base_dir rely on annotation if name is 'filename.ext [annotation]' format80# otherwise use default_path as base_dir81def annotated_filepath(name):82    if name.endswith("[output]"):83        base_dir = get_output_directory()84        name = name[:-9]85    elif name.endswith("[input]"):86        base_dir = get_input_directory()87        name = name[:-8]88    elif name.endswith("[temp]"):89        base_dir = get_temp_directory()90        name = name[:-7]91    else:92        return name, None93 94    return name, base_dir95 96 97def get_annotated_filepath(name, default_dir=None):98    name, base_dir = annotated_filepath(name)99 100    if base_dir is None:101        if default_dir is not None:102            base_dir = default_dir103        else:104            base_dir = get_input_directory()  # fallback path105 106    return os.path.join(base_dir, name)107 108 109def exists_annotated_filepath(name):110    name, base_dir = annotated_filepath(name)111 112    if base_dir is None:113        base_dir = get_input_directory()  # fallback path114 115    filepath = os.path.join(base_dir, name)116    return os.path.exists(filepath)117 118 119def add_model_folder_path(folder_name, full_folder_path):120    global folder_names_and_paths121    if folder_name in folder_names_and_paths:122        folder_names_and_paths[folder_name][0].append(full_folder_path)123    else:124        folder_names_and_paths[folder_name] = ([full_folder_path], set())125 126def get_folder_paths(folder_name):127    return folder_names_and_paths[folder_name][0][:]128 129def recursive_search(directory, excluded_dir_names=None):130    if not os.path.isdir(directory):131        return [], {}132 133    if excluded_dir_names is None:134        excluded_dir_names = []135 136    result = []137    dirs = {directory: os.path.getmtime(directory)}138    for dirpath, subdirs, filenames in os.walk(directory, followlinks=True, topdown=True):139        subdirs[:] = [d for d in subdirs if d not in excluded_dir_names]140        for file_name in filenames:141            relative_path = os.path.relpath(os.path.join(dirpath, file_name), directory)142            result.append(relative_path)143        for d in subdirs:144            path = os.path.join(dirpath, d)145            dirs[path] = os.path.getmtime(path)146    return result, dirs147 148def filter_files_extensions(files, extensions):149    return sorted(list(filter(lambda a: os.path.splitext(a)[-1].lower() in extensions or len(extensions) == 0, files)))150 151 152 153def get_full_path(folder_name, filename):154    global folder_names_and_paths155    if folder_name not in folder_names_and_paths:156        return None157    folders = folder_names_and_paths[folder_name]158    filename = os.path.relpath(os.path.join("/", filename), "/")159    for x in folders[0]:160        full_path = os.path.join(x, filename)161        if os.path.isfile(full_path):162            return full_path163 164    return None165 166def get_filename_list_(folder_name):167    global folder_names_and_paths168    output_list = set()169    folders = folder_names_and_paths[folder_name]170    output_folders = {}171    for x in folders[0]:172        files, folders_all = recursive_search(x, excluded_dir_names=[".git"])173        output_list.update(filter_files_extensions(files, folders[1]))174        output_folders = {**output_folders, **folders_all}175 176    return (sorted(list(output_list)), output_folders, time.perf_counter())177 178def cached_filename_list_(folder_name):179    global filename_list_cache180    global folder_names_and_paths181    if folder_name not in filename_list_cache:182        return None183    out = filename_list_cache[folder_name]184    if time.perf_counter() < (out[2] + 0.5):185        return out186    for x in out[1]:187        time_modified = out[1][x]188        folder = x189        if os.path.getmtime(folder) != time_modified:190            return None191 192    folders = folder_names_and_paths[folder_name]193    for x in folders[0]:194        if os.path.isdir(x):195            if x not in out[1]:196                return None197 198    return out199 200def get_filename_list(folder_name):201    out = cached_filename_list_(folder_name)202    if out is None:203        out = get_filename_list_(folder_name)204        global filename_list_cache205        filename_list_cache[folder_name] = out206    return list(out[0])207 208def get_save_image_path(filename_prefix, output_dir, image_width=0, image_height=0):209    def map_filename(filename):210        prefix_len = len(os.path.basename(filename_prefix))211        prefix = filename[:prefix_len + 1]212        try:213            digits = int(filename[prefix_len + 1:].split('_')[0])214        except:215            digits = 0216        return (digits, prefix)217 218    def compute_vars(input, image_width, image_height):219        input = input.replace("%width%", str(image_width))220        input = input.replace("%height%", str(image_height))221        return input222 223    filename_prefix = compute_vars(filename_prefix, image_width, image_height)224 225    subfolder = os.path.dirname(os.path.normpath(filename_prefix))226    filename = os.path.basename(os.path.normpath(filename_prefix))227 228    full_output_folder = os.path.join(output_dir, subfolder)229 230    if os.path.commonpath((output_dir, os.path.abspath(full_output_folder))) != output_dir:231        print("Saving image outside the output folder is not allowed.")232        return {}233 234    try:235        counter = max(filter(lambda a: a[1][:-1] == filename and a[1][-1] == "_", map(map_filename, os.listdir(full_output_folder))))[0] + 1236    except ValueError:237        counter = 1238    except FileNotFoundError:239        os.makedirs(full_output_folder, exist_ok=True)240        counter = 1241    return full_output_folder, filename, counter, subfolder, filename_prefix242