surena26 commited on
Commit
614e660
·
verified ·
1 Parent(s): badcd5b

Upload ComfyUI/folder_paths.py with huggingface_hub

Browse files
Files changed (1) hide show
  1. ComfyUI/folder_paths.py +267 -0
ComfyUI/folder_paths.py ADDED
@@ -0,0 +1,267 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import time
3
+ import logging
4
+
5
+ supported_pt_extensions = set(['.ckpt', '.pt', '.bin', '.pth', '.safetensors', '.pkl'])
6
+
7
+ folder_names_and_paths = {}
8
+
9
+ base_path = os.path.dirname(os.path.realpath(__file__))
10
+ models_dir = os.path.join(base_path, "models")
11
+ folder_names_and_paths["checkpoints"] = ([os.path.join(models_dir, "checkpoints")], supported_pt_extensions)
12
+ folder_names_and_paths["configs"] = ([os.path.join(models_dir, "configs")], [".yaml"])
13
+
14
+ folder_names_and_paths["loras"] = ([os.path.join(models_dir, "loras")], supported_pt_extensions)
15
+ folder_names_and_paths["vae"] = ([os.path.join(models_dir, "vae")], supported_pt_extensions)
16
+ folder_names_and_paths["clip"] = ([os.path.join(models_dir, "clip")], supported_pt_extensions)
17
+ folder_names_and_paths["unet"] = ([os.path.join(models_dir, "unet")], supported_pt_extensions)
18
+ folder_names_and_paths["clip_vision"] = ([os.path.join(models_dir, "clip_vision")], supported_pt_extensions)
19
+ folder_names_and_paths["style_models"] = ([os.path.join(models_dir, "style_models")], supported_pt_extensions)
20
+ folder_names_and_paths["embeddings"] = ([os.path.join(models_dir, "embeddings")], supported_pt_extensions)
21
+ folder_names_and_paths["diffusers"] = ([os.path.join(models_dir, "diffusers")], ["folder"])
22
+ folder_names_and_paths["vae_approx"] = ([os.path.join(models_dir, "vae_approx")], supported_pt_extensions)
23
+
24
+ folder_names_and_paths["controlnet"] = ([os.path.join(models_dir, "controlnet"), os.path.join(models_dir, "t2i_adapter")], supported_pt_extensions)
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
+
31
+ folder_names_and_paths["hypernetworks"] = ([os.path.join(models_dir, "hypernetworks")], supported_pt_extensions)
32
+
33
+ folder_names_and_paths["photomaker"] = ([os.path.join(models_dir, "photomaker")], supported_pt_extensions)
34
+
35
+ folder_names_and_paths["classifiers"] = ([os.path.join(models_dir, "classifiers")], {""})
36
+
37
+ output_directory = os.path.join(os.path.dirname(os.path.realpath(__file__)), "output")
38
+ temp_directory = os.path.join(os.path.dirname(os.path.realpath(__file__)), "temp")
39
+ input_directory = os.path.join(os.path.dirname(os.path.realpath(__file__)), "input")
40
+ user_directory = os.path.join(os.path.dirname(os.path.realpath(__file__)), "user")
41
+
42
+ filename_list_cache = {}
43
+
44
+ if not os.path.exists(input_directory):
45
+ try:
46
+ os.makedirs(input_directory)
47
+ except:
48
+ logging.error("Failed to create input directory")
49
+
50
+ def set_output_directory(output_dir):
51
+ global output_directory
52
+ output_directory = output_dir
53
+
54
+ def set_temp_directory(temp_dir):
55
+ global temp_directory
56
+ temp_directory = temp_dir
57
+
58
+ def set_input_directory(input_dir):
59
+ global input_directory
60
+ input_directory = input_dir
61
+
62
+ def get_output_directory():
63
+ global output_directory
64
+ return output_directory
65
+
66
+ def get_temp_directory():
67
+ global temp_directory
68
+ return temp_directory
69
+
70
+ def get_input_directory():
71
+ global input_directory
72
+ return input_directory
73
+
74
+
75
+ #NOTE: used in http server so don't put folders that should not be accessed remotely
76
+ def get_directory_by_type(type_name):
77
+ if type_name == "output":
78
+ return get_output_directory()
79
+ if type_name == "temp":
80
+ return get_temp_directory()
81
+ if type_name == "input":
82
+ return get_input_directory()
83
+ return None
84
+
85
+
86
+ # determine base_dir rely on annotation if name is 'filename.ext [annotation]' format
87
+ # otherwise use default_path as base_dir
88
+ def annotated_filepath(name):
89
+ if name.endswith("[output]"):
90
+ base_dir = get_output_directory()
91
+ name = name[:-9]
92
+ elif name.endswith("[input]"):
93
+ base_dir = get_input_directory()
94
+ name = name[:-8]
95
+ elif name.endswith("[temp]"):
96
+ base_dir = get_temp_directory()
97
+ name = name[:-7]
98
+ else:
99
+ return name, None
100
+
101
+ return name, base_dir
102
+
103
+
104
+ def get_annotated_filepath(name, default_dir=None):
105
+ name, base_dir = annotated_filepath(name)
106
+
107
+ if base_dir is None:
108
+ if default_dir is not None:
109
+ base_dir = default_dir
110
+ else:
111
+ base_dir = get_input_directory() # fallback path
112
+
113
+ return os.path.join(base_dir, name)
114
+
115
+
116
+ def exists_annotated_filepath(name):
117
+ name, base_dir = annotated_filepath(name)
118
+
119
+ if base_dir is None:
120
+ base_dir = get_input_directory() # fallback path
121
+
122
+ filepath = os.path.join(base_dir, name)
123
+ return os.path.exists(filepath)
124
+
125
+
126
+ def add_model_folder_path(folder_name, full_folder_path):
127
+ global folder_names_and_paths
128
+ if folder_name in folder_names_and_paths:
129
+ folder_names_and_paths[folder_name][0].append(full_folder_path)
130
+ else:
131
+ folder_names_and_paths[folder_name] = ([full_folder_path], set())
132
+
133
+ def get_folder_paths(folder_name):
134
+ return folder_names_and_paths[folder_name][0][:]
135
+
136
+ def recursive_search(directory, excluded_dir_names=None):
137
+ if not os.path.isdir(directory):
138
+ return [], {}
139
+
140
+ if excluded_dir_names is None:
141
+ excluded_dir_names = []
142
+
143
+ result = []
144
+ dirs = {}
145
+
146
+ # Attempt to add the initial directory to dirs with error handling
147
+ try:
148
+ dirs[directory] = os.path.getmtime(directory)
149
+ except FileNotFoundError:
150
+ logging.warning(f"Warning: Unable to access {directory}. Skipping this path.")
151
+
152
+ logging.debug("recursive file list on directory {}".format(directory))
153
+ for dirpath, subdirs, filenames in os.walk(directory, followlinks=True, topdown=True):
154
+ subdirs[:] = [d for d in subdirs if d not in excluded_dir_names]
155
+ for file_name in filenames:
156
+ relative_path = os.path.relpath(os.path.join(dirpath, file_name), directory)
157
+ result.append(relative_path)
158
+
159
+ for d in subdirs:
160
+ path = os.path.join(dirpath, d)
161
+ try:
162
+ dirs[path] = os.path.getmtime(path)
163
+ except FileNotFoundError:
164
+ logging.warning(f"Warning: Unable to access {path}. Skipping this path.")
165
+ continue
166
+ logging.debug("found {} files".format(len(result)))
167
+ return result, dirs
168
+
169
+ def filter_files_extensions(files, extensions):
170
+ return sorted(list(filter(lambda a: os.path.splitext(a)[-1].lower() in extensions or len(extensions) == 0, files)))
171
+
172
+
173
+
174
+ def get_full_path(folder_name, filename):
175
+ global folder_names_and_paths
176
+ if folder_name not in folder_names_and_paths:
177
+ return None
178
+ folders = folder_names_and_paths[folder_name]
179
+ filename = os.path.relpath(os.path.join("/", filename), "/")
180
+ for x in folders[0]:
181
+ full_path = os.path.join(x, filename)
182
+ if os.path.isfile(full_path):
183
+ return full_path
184
+ elif os.path.islink(full_path):
185
+ logging.warning("WARNING path {} exists but doesn't link anywhere, skipping.".format(full_path))
186
+
187
+ return None
188
+
189
+ def get_filename_list_(folder_name):
190
+ global folder_names_and_paths
191
+ output_list = set()
192
+ folders = folder_names_and_paths[folder_name]
193
+ output_folders = {}
194
+ for x in folders[0]:
195
+ files, folders_all = recursive_search(x, excluded_dir_names=[".git"])
196
+ output_list.update(filter_files_extensions(files, folders[1]))
197
+ output_folders = {**output_folders, **folders_all}
198
+
199
+ return (sorted(list(output_list)), output_folders, time.perf_counter())
200
+
201
+ def cached_filename_list_(folder_name):
202
+ global filename_list_cache
203
+ global folder_names_and_paths
204
+ if folder_name not in filename_list_cache:
205
+ return None
206
+ out = filename_list_cache[folder_name]
207
+
208
+ for x in out[1]:
209
+ time_modified = out[1][x]
210
+ folder = x
211
+ if os.path.getmtime(folder) != time_modified:
212
+ return None
213
+
214
+ folders = folder_names_and_paths[folder_name]
215
+ for x in folders[0]:
216
+ if os.path.isdir(x):
217
+ if x not in out[1]:
218
+ return None
219
+
220
+ return out
221
+
222
+ def get_filename_list(folder_name):
223
+ out = cached_filename_list_(folder_name)
224
+ if out is None:
225
+ out = get_filename_list_(folder_name)
226
+ global filename_list_cache
227
+ filename_list_cache[folder_name] = out
228
+ return list(out[0])
229
+
230
+ def get_save_image_path(filename_prefix, output_dir, image_width=0, image_height=0):
231
+ def map_filename(filename):
232
+ prefix_len = len(os.path.basename(filename_prefix))
233
+ prefix = filename[:prefix_len + 1]
234
+ try:
235
+ digits = int(filename[prefix_len + 1:].split('_')[0])
236
+ except:
237
+ digits = 0
238
+ return (digits, prefix)
239
+
240
+ def compute_vars(input, image_width, image_height):
241
+ input = input.replace("%width%", str(image_width))
242
+ input = input.replace("%height%", str(image_height))
243
+ return input
244
+
245
+ filename_prefix = compute_vars(filename_prefix, image_width, image_height)
246
+
247
+ subfolder = os.path.dirname(os.path.normpath(filename_prefix))
248
+ filename = os.path.basename(os.path.normpath(filename_prefix))
249
+
250
+ full_output_folder = os.path.join(output_dir, subfolder)
251
+
252
+ if os.path.commonpath((output_dir, os.path.abspath(full_output_folder))) != output_dir:
253
+ err = "**** ERROR: Saving image outside the output folder is not allowed." + \
254
+ "\n full_output_folder: " + os.path.abspath(full_output_folder) + \
255
+ "\n output_dir: " + output_dir + \
256
+ "\n commonpath: " + os.path.commonpath((output_dir, os.path.abspath(full_output_folder)))
257
+ logging.error(err)
258
+ raise Exception(err)
259
+
260
+ try:
261
+ counter = max(filter(lambda a: os.path.normcase(a[1][:-1]) == os.path.normcase(filename) and a[1][-1] == "_", map(map_filename, os.listdir(full_output_folder))))[0] + 1
262
+ except ValueError:
263
+ counter = 1
264
+ except FileNotFoundError:
265
+ os.makedirs(full_output_folder, exist_ok=True)
266
+ counter = 1
267
+ return full_output_folder, filename, counter, subfolder, filename_prefix