From 6e3b261ce94325968a08f3dc78e7fdf175b47731 Mon Sep 17 00:00:00 2001 From: filtered <176114999+webfiltered@users.noreply.github.com> Date: Sun, 26 Jan 2025 02:55:13 +1100 Subject: [PATCH] Allow complete reset of default directories under a new base_path. --- folder_paths.py | 206 ++++++++++++++++++++++++++++-------------------- 1 file changed, 121 insertions(+), 85 deletions(-) diff --git a/folder_paths.py b/folder_paths.py index b8c1926db..370354a4d 100644 --- a/folder_paths.py +++ b/folder_paths.py @@ -7,88 +7,11 @@ import logging from typing import Literal, TypedDict from collections.abc import Collection +# Path initialisation occurs on import at the end of this file. + supported_pt_extensions: set[str] = {'.ckpt', '.pt', '.bin', '.pth', '.safetensors', '.pkl', '.sft'} folder_names_and_paths: dict[str, tuple[list[str], set[str] | list[str]]] = {} - -base_path = os.path.dirname(os.path.realpath(__file__)) -models_dir = os.path.join(base_path, "models") - - -class DefaultModelSubdirOptions(TypedDict, total=False): - extensions: set[str] | list[str] - alternate_names: list[str] - - -# Default models subdirs -default_paths: dict[str, DefaultModelSubdirOptions] = { - "checkpoints": {}, - "configs": { - "extensions": [".yaml"], - }, - "loras": {}, - "vae": {}, - "text_encoders": { - "alternate_names": ["clip"], - }, - "clip_vision": {}, - "style_models": {}, - "embeddings": {}, - "diffusers": { - "extensions": ["folder"], - }, - "vae_approx": {}, - "controlnet": { - "alternate_names": ["t2i_adapter"], - }, - "gligen": {}, - "upscale_models": {}, - "hypernetworks": {}, - "photomaker": {}, - "classifiers": { - "extensions": {""}, - }, -} - -def add_default_model_paths(models_path: str, is_default: bool = False) -> None: - """ - Adds all built-in model paths under a specified directory to the known folder paths. - - Args: - models_path: The path to the models directory that contains default folders, e.g. ``checkpoints``, ``loras``. - is_default: If the created paths should be the new default paths. - """ - - for [name, options] in default_paths.items(): - valid_extensions = options.get("extensions", supported_pt_extensions) - subdirs = [name, *options.get("alternate_names", [])] - - for subdir in subdirs: - add_model_folder_path( - name, - os.path.join(models_path, subdir), - is_default=is_default, - is_legacy=True, - valid_extensions=valid_extensions, - ) - - # diffusion_models config shipped with the reverse order of other paths. - # Maintain original order, as this may impact ordering in extensions / UI. - add_model_folder_path("diffusion_models", os.path.join(models_path, "unet"), is_default=is_default, is_legacy=True, valid_extensions=supported_pt_extensions) - add_model_folder_path("diffusion_models", os.path.join(models_path, "diffusion_models"), is_default=is_default, is_legacy=True, valid_extensions=supported_pt_extensions) - - -# Default configuration: add model paths relative to this file. -add_default_model_paths(base_path) - -# TODO: custom_nodes is now added out of original order; confirm no impact. -folder_names_and_paths["custom_nodes"] = ([os.path.join(base_path, "custom_nodes")], set()) - -output_directory = os.path.join(base_path, "output") -temp_directory = os.path.join(base_path, "temp") -input_directory = os.path.join(base_path, "input") -user_directory = os.path.join(base_path, "user") - filename_list_cache: dict[str, tuple[list[str], dict[str, float], float]] = {} class CacheHelper: @@ -130,12 +53,6 @@ def map_legacy(folder_name: str) -> str: "clip": "text_encoders"} return legacy.get(folder_name, folder_name) -if not os.path.exists(input_directory): - try: - os.makedirs(input_directory) - except: - logging.error("Failed to create input directory") - def set_output_directory(output_dir: str) -> None: global output_directory output_directory = output_dir @@ -436,3 +353,122 @@ def get_save_image_path(filename_prefix: str, output_dir: str, image_width=0, im os.makedirs(full_output_folder, exist_ok=True) counter = 1 return full_output_folder, filename, counter, subfolder, filename_prefix + + +class DefaultModelSubdirOptions(TypedDict, total=False): + extensions: set[str] | list[str] + alternate_names: list[str] + + +# Default models subdirs +default_paths: dict[str, DefaultModelSubdirOptions] = { + "checkpoints": {}, + "configs": { + "extensions": [".yaml"], + }, + "loras": {}, + "vae": {}, + "text_encoders": { + "alternate_names": ["clip"], + }, + "clip_vision": {}, + "style_models": {}, + "embeddings": {}, + "diffusers": { + "extensions": ["folder"], + }, + "vae_approx": {}, + "controlnet": { + "alternate_names": ["t2i_adapter"], + }, + "gligen": {}, + "upscale_models": {}, + "hypernetworks": {}, + "photomaker": {}, + "classifiers": { + "extensions": {""}, + }, +} + + +def reset_all_paths(new_base_path: str) -> None: + """ + Internal use only. Designed for use immediately after startup. + + Removes all existing known paths, clears the filename cache, and creates the defaults under a new base path. + + Paths: + - All default model paths + - input + - output + - temp + - user + - custom_nodes + + Also creates the input directory if missing. + + Args: + new_base_path: The base path to prepend to all default relative paths. + """ + + global base_path + global models_dir + global output_directory + global temp_directory + global input_directory + global user_directory + + filename_list_cache.clear() + folder_names_and_paths.clear() + + base_path = new_base_path + models_dir = os.path.join(base_path, "models") + + add_default_model_paths(base_path) + + output_directory = os.path.join(base_path, "output") + temp_directory = os.path.join(base_path, "temp") + input_directory = os.path.join(base_path, "input") + user_directory = os.path.join(base_path, "user") + + # custom_nodes uses a different path format. + add_model_folder_path("custom_nodes", os.path.join(base_path, "custom_nodes")) + + # Create input dir if it does not already exist. + if not os.path.exists(input_directory): + try: + os.makedirs(input_directory) + except: + logging.error("Failed to create input directory") + + +def add_default_model_paths(models_path: str, is_default: bool = False) -> None: + """ + Adds all built-in model paths under a specified directory to the known folder paths. + + Args: + models_path: The path to the models directory that contains default folders, e.g. ``checkpoints``, ``loras``. + is_default: If the created paths should be the new default paths. + """ + + for [name, options] in default_paths.items(): + valid_extensions = options.get("extensions", supported_pt_extensions) + subdirs = [name, *options.get("alternate_names", [])] + + for subdir in subdirs: + add_model_folder_path( + name, + os.path.join(models_path, subdir), + is_default=is_default, + is_legacy=True, + valid_extensions=valid_extensions, + ) + + # diffusion_models config shipped with the reverse order of other paths. + # Maintain original order, as this may impact ordering in extensions / UI. + add_model_folder_path("diffusion_models", os.path.join(models_path, "unet"), is_default=is_default, is_legacy=True, valid_extensions=supported_pt_extensions) + add_model_folder_path("diffusion_models", os.path.join(models_path, "diffusion_models"), is_default=is_default, is_legacy=True, valid_extensions=supported_pt_extensions) + + +# Default configuration: add all paths relative to this file. +reset_all_paths(os.path.dirname(os.path.realpath(__file__)))