folder_paths.py 8.63 KB
Newer Older
1
import os
2
import time
3

4
supported_pt_extensions = set(['.ckpt', '.pt', '.bin', '.pth', '.safetensors'])
5
6
7

folder_names_and_paths = {}

8
9
base_path = os.path.dirname(os.path.realpath(__file__))
models_dir = os.path.join(base_path, "models")
10
folder_names_and_paths["checkpoints"] = ([os.path.join(models_dir, "checkpoints")], supported_pt_extensions)
11
12
13
14
15
folder_names_and_paths["configs"] = ([os.path.join(models_dir, "configs")], [".yaml"])

folder_names_and_paths["loras"] = ([os.path.join(models_dir, "loras")], supported_pt_extensions)
folder_names_and_paths["vae"] = ([os.path.join(models_dir, "vae")], supported_pt_extensions)
folder_names_and_paths["clip"] = ([os.path.join(models_dir, "clip")], supported_pt_extensions)
16
folder_names_and_paths["unet"] = ([os.path.join(models_dir, "unet")], supported_pt_extensions)
17
18
folder_names_and_paths["clip_vision"] = ([os.path.join(models_dir, "clip_vision")], supported_pt_extensions)
folder_names_and_paths["style_models"] = ([os.path.join(models_dir, "style_models")], supported_pt_extensions)
19
folder_names_and_paths["embeddings"] = ([os.path.join(models_dir, "embeddings")], supported_pt_extensions)
20
folder_names_and_paths["diffusers"] = ([os.path.join(models_dir, "diffusers")], ["folder"])
21
folder_names_and_paths["vae_approx"] = ([os.path.join(models_dir, "vae_approx")], supported_pt_extensions)
22
23

folder_names_and_paths["controlnet"] = ([os.path.join(models_dir, "controlnet"), os.path.join(models_dir, "t2i_adapter")], supported_pt_extensions)
24
25
folder_names_and_paths["gligen"] = ([os.path.join(models_dir, "gligen")], supported_pt_extensions)

26
27
folder_names_and_paths["upscale_models"] = ([os.path.join(models_dir, "upscale_models")], supported_pt_extensions)

28
29
folder_names_and_paths["custom_nodes"] = ([os.path.join(base_path, "custom_nodes")], [])

30
folder_names_and_paths["hypernetworks"] = ([os.path.join(models_dir, "hypernetworks")], supported_pt_extensions)
31

32
33
folder_names_and_paths["classifiers"] = ([os.path.join(models_dir, "classifiers")], {""})

34
35
36
37
output_directory = os.path.join(os.path.dirname(os.path.realpath(__file__)), "output")
temp_directory = os.path.join(os.path.dirname(os.path.realpath(__file__)), "temp")
input_directory = os.path.join(os.path.dirname(os.path.realpath(__file__)), "input")

38
39
filename_list_cache = {}

40
if not os.path.exists(input_directory):
Enrico Fasoli's avatar
Enrico Fasoli committed
41
42
43
44
    try:
        os.makedirs(input_directory)
    except:
        print("Failed to create input directory")
45
46
47
48
49

def set_output_directory(output_dir):
    global output_directory
    output_directory = output_dir

50
51
52
53
def set_temp_directory(temp_dir):
    global temp_directory
    temp_directory = temp_dir

Jairo Correa's avatar
Jairo Correa committed
54
55
56
57
def set_input_directory(input_dir):
    global input_directory
    input_directory = input_dir

58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
def get_output_directory():
    global output_directory
    return output_directory

def get_temp_directory():
    global temp_directory
    return temp_directory

def get_input_directory():
    global input_directory
    return input_directory


#NOTE: used in http server so don't put folders that should not be accessed remotely
def get_directory_by_type(type_name):
    if type_name == "output":
        return get_output_directory()
    if type_name == "temp":
        return get_temp_directory()
    if type_name == "input":
        return get_input_directory()
    return None

81

ltdrdata's avatar
ltdrdata committed
82
83
# determine base_dir rely on annotation if name is 'filename.ext [annotation]' format
# otherwise use default_path as base_dir
84
def annotated_filepath(name):
ltdrdata's avatar
ltdrdata committed
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
    if name.endswith("[output]"):
        base_dir = get_output_directory()
        name = name[:-9]
    elif name.endswith("[input]"):
        base_dir = get_input_directory()
        name = name[:-8]
    elif name.endswith("[temp]"):
        base_dir = get_temp_directory()
        name = name[:-7]
    else:
        return name, None

    return name, base_dir


def get_annotated_filepath(name, default_dir=None):
101
    name, base_dir = annotated_filepath(name)
ltdrdata's avatar
ltdrdata committed
102
103
104
105
106
107
108
109
110
111
112

    if base_dir is None:
        if default_dir is not None:
            base_dir = default_dir
        else:
            base_dir = get_input_directory()  # fallback path

    return os.path.join(base_dir, name)


def exists_annotated_filepath(name):
113
    name, base_dir = annotated_filepath(name)
ltdrdata's avatar
ltdrdata committed
114
115
116
117
118
119
120
121

    if base_dir is None:
        base_dir = get_input_directory()  # fallback path

    filepath = os.path.join(base_dir, name)
    return os.path.exists(filepath)


122
def add_model_folder_path(folder_name, full_folder_path):
123
    global folder_names_and_paths
124
125
    if folder_name in folder_names_and_paths:
        folder_names_and_paths[folder_name][0].append(full_folder_path)
126
127
    else:
        folder_names_and_paths[folder_name] = ([full_folder_path], set())
128

129
130
def get_folder_paths(folder_name):
    return folder_names_and_paths[folder_name][0][:]
131

132
def recursive_search(directory, excluded_dir_names=None):
133
134
    if not os.path.isdir(directory):
        return [], {}
135
136
137
138

    if excluded_dir_names is None:
        excluded_dir_names = []

139
    result = []
140
    dirs = {directory: os.path.getmtime(directory)}
141
142
143
144
145
146
147
    for dirpath, subdirs, filenames in os.walk(directory, followlinks=True, topdown=True):
        subdirs[:] = [d for d in subdirs if d not in excluded_dir_names]
        for file_name in filenames:
            relative_path = os.path.relpath(os.path.join(dirpath, file_name), directory)
            result.append(relative_path)
        for d in subdirs:
            path = os.path.join(dirpath, d)
148
149
            dirs[path] = os.path.getmtime(path)
    return result, dirs
150
151

def filter_files_extensions(files, extensions):
152
    return sorted(list(filter(lambda a: os.path.splitext(a)[-1].lower() in extensions or len(extensions) == 0, files)))
153
154
155
156
157



def get_full_path(folder_name, filename):
    global folder_names_and_paths
158
159
    if folder_name not in folder_names_and_paths:
        return None
160
    folders = folder_names_and_paths[folder_name]
161
    filename = os.path.relpath(os.path.join("/", filename), "/")
162
163
164
165
166
    for x in folders[0]:
        full_path = os.path.join(x, filename)
        if os.path.isfile(full_path):
            return full_path

167
    return None
168

169
def get_filename_list_(folder_name):
170
    global folder_names_and_paths
171
    output_list = set()
172
173
174
    folders = folder_names_and_paths[folder_name]
    output_folders = {}
    for x in folders[0]:
175
        files, folders_all = recursive_search(x, excluded_dir_names=[".git"])
176
177
178
        output_list.update(filter_files_extensions(files, folders[1]))
        output_folders = {**output_folders, **folders_all}

179
    return (sorted(list(output_list)), output_folders, time.perf_counter())
180
181
182
183
184
185
186

def cached_filename_list_(folder_name):
    global filename_list_cache
    global folder_names_and_paths
    if folder_name not in filename_list_cache:
        return None
    out = filename_list_cache[folder_name]
187
188
    if time.perf_counter() < (out[2] + 0.5):
        return out
189
190
191
192
193
194
    for x in out[1]:
        time_modified = out[1][x]
        folder = x
        if os.path.getmtime(folder) != time_modified:
            return None

195
196
    folders = folder_names_and_paths[folder_name]
    for x in folders[0]:
197
198
199
        if os.path.isdir(x):
            if x not in out[1]:
                return None
200
201
202
203
204
205
206
207
208

    return out

def get_filename_list(folder_name):
    out = cached_filename_list_(folder_name)
    if out is None:
        out = get_filename_list_(folder_name)
        global filename_list_cache
        filename_list_cache[folder_name] = out
209
    return list(out[0])
210

211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
def get_save_image_path(filename_prefix, output_dir, image_width=0, image_height=0):
    def map_filename(filename):
        prefix_len = len(os.path.basename(filename_prefix))
        prefix = filename[:prefix_len + 1]
        try:
            digits = int(filename[prefix_len + 1:].split('_')[0])
        except:
            digits = 0
        return (digits, prefix)

    def compute_vars(input, image_width, image_height):
        input = input.replace("%width%", str(image_width))
        input = input.replace("%height%", str(image_height))
        return input

    filename_prefix = compute_vars(filename_prefix, image_width, image_height)

    subfolder = os.path.dirname(os.path.normpath(filename_prefix))
    filename = os.path.basename(os.path.normpath(filename_prefix))

    full_output_folder = os.path.join(output_dir, subfolder)

    if os.path.commonpath((output_dir, os.path.abspath(full_output_folder))) != output_dir:
        print("Saving image outside the output folder is not allowed.")
        return {}

    try:
        counter = max(filter(lambda a: a[1][:-1] == filename and a[1][-1] == "_", map(map_filename, os.listdir(full_output_folder))))[0] + 1
    except ValueError:
        counter = 1
    except FileNotFoundError:
        os.makedirs(full_output_folder, exist_ok=True)
        counter = 1
    return full_output_folder, filename, counter, subfolder, filename_prefix