Agents-MCP-Hackathon/AI-QR-code-generator
0
1from __future__ import annotations2 3import os4import time5import mimetypes6import logging7from typing import Literal, List8from collections.abc import Collection9 10from comfy.cli_args import args11 12supported_pt_extensions: set[str] = {'.ckpt', '.pt', '.pt2', '.bin', '.pth', '.safetensors', '.pkl', '.sft'}13 14folder_names_and_paths: dict[str, tuple[list[str], set[str]]] = {}15 16# --base-directory - Resets all default paths configured in folder_paths with a new base path17if args.base_directory:18 base_path = os.path.abspath(args.base_directory)19else:20 base_path = os.path.dirname(os.path.realpath(__file__))21 22models_dir = os.path.join(base_path, "models")23folder_names_and_paths["checkpoints"] = ([os.path.join(models_dir, "checkpoints")], supported_pt_extensions)24folder_names_and_paths["configs"] = ([os.path.join(models_dir, "configs")], [".yaml"])25 26folder_names_and_paths["loras"] = ([os.path.join(models_dir, "loras")], supported_pt_extensions)27folder_names_and_paths["vae"] = ([os.path.join(models_dir, "vae")], supported_pt_extensions)28folder_names_and_paths["text_encoders"] = ([os.path.join(models_dir, "text_encoders"), os.path.join(models_dir, "clip")], supported_pt_extensions)29folder_names_and_paths["diffusion_models"] = ([os.path.join(models_dir, "unet"), os.path.join(models_dir, "diffusion_models")], supported_pt_extensions)30folder_names_and_paths["clip_vision"] = ([os.path.join(models_dir, "clip_vision")], supported_pt_extensions)31folder_names_and_paths["style_models"] = ([os.path.join(models_dir, "style_models")], supported_pt_extensions)32folder_names_and_paths["embeddings"] = ([os.path.join(models_dir, "embeddings")], supported_pt_extensions)33folder_names_and_paths["diffusers"] = ([os.path.join(models_dir, "diffusers")], ["folder"])34folder_names_and_paths["vae_approx"] = ([os.path.join(models_dir, "vae_approx")], supported_pt_extensions)35 36folder_names_and_paths["controlnet"] = ([os.path.join(models_dir, "controlnet"), os.path.join(models_dir, "t2i_adapter")], supported_pt_extensions)37folder_names_and_paths["gligen"] = ([os.path.join(models_dir, "gligen")], supported_pt_extensions)38 39folder_names_and_paths["upscale_models"] = ([os.path.join(models_dir, "upscale_models")], supported_pt_extensions)40 41folder_names_and_paths["custom_nodes"] = ([os.path.join(base_path, "custom_nodes")], set())42 43folder_names_and_paths["hypernetworks"] = ([os.path.join(models_dir, "hypernetworks")], supported_pt_extensions)44 45folder_names_and_paths["photomaker"] = ([os.path.join(models_dir, "photomaker")], supported_pt_extensions)46 47folder_names_and_paths["classifiers"] = ([os.path.join(models_dir, "classifiers")], {""})48 49output_directory = os.path.join(base_path, "output")50temp_directory = os.path.join(base_path, "temp")51input_directory = os.path.join(base_path, "input")52user_directory = os.path.join(base_path, "user")53 54filename_list_cache: dict[str, tuple[list[str], dict[str, float], float]] = {}55 56class CacheHelper:57 """58 Helper class for managing file list cache data.59 """60 def __init__(self):61 self.cache: dict[str, tuple[list[str], dict[str, float], float]] = {}62 self.active = False63 64 def get(self, key: str, default=None) -> tuple[list[str], dict[str, float], float]:65 if not self.active:66 return default67 return self.cache.get(key, default)68 69 def set(self, key: str, value: tuple[list[str], dict[str, float], float]) -> None:70 if self.active:71 self.cache[key] = value72 73 def clear(self):74 self.cache.clear()75 76 def __enter__(self):77 self.active = True78 return self79 80 def __exit__(self, exc_type, exc_value, traceback):81 self.active = False82 self.clear()83 84cache_helper = CacheHelper()85 86extension_mimetypes_cache = {87 "webp" : "image",88 "fbx" : "model",89}90 91def map_legacy(folder_name: str) -> str:92 legacy = {"unet": "diffusion_models",93 "clip": "text_encoders"}94 return legacy.get(folder_name, folder_name)95 96if not os.path.exists(input_directory):97 try:98 os.makedirs(input_directory)99 except:100 logging.error("Failed to create input directory")101 102def set_output_directory(output_dir: str) -> None:103 global output_directory104 output_directory = output_dir105 106def set_temp_directory(temp_dir: str) -> None:107 global temp_directory108 temp_directory = temp_dir109 110def set_input_directory(input_dir: str) -> None:111 global input_directory112 input_directory = input_dir113 114def get_output_directory() -> str:115 global output_directory116 return output_directory117 118def get_temp_directory() -> str:119 global temp_directory120 return temp_directory121 122def get_input_directory() -> str:123 global input_directory124 return input_directory125 126def get_user_directory() -> str:127 return user_directory128 129def set_user_directory(user_dir: str) -> None:130 global user_directory131 user_directory = user_dir132 133 134#NOTE: used in http server so don't put folders that should not be accessed remotely135def get_directory_by_type(type_name: str) -> str | None:136 if type_name == "output":137 return get_output_directory()138 if type_name == "temp":139 return get_temp_directory()140 if type_name == "input":141 return get_input_directory()142 return None143 144def filter_files_content_types(files: list[str], content_types: List[Literal["image", "video", "audio", "model"]]) -> list[str]:145 """146 Example:147 files = os.listdir(folder_paths.get_input_directory())148 videos = filter_files_content_types(files, ["video"])149 150 Note:151 - 'model' in MIME context refers to 3D models, not files containing trained weights and parameters152 """153 global extension_mimetypes_cache154 result = []155 for file in files:156 extension = file.split('.')[-1]157 if extension not in extension_mimetypes_cache:158 mime_type, _ = mimetypes.guess_type(file, strict=False)159 if not mime_type:160 continue161 content_type = mime_type.split('/')[0]162 extension_mimetypes_cache[extension] = content_type163 else:164 content_type = extension_mimetypes_cache[extension]165 166 if content_type in content_types:167 result.append(file)168 return result169 170# determine base_dir rely on annotation if name is 'filename.ext [annotation]' format171# otherwise use default_path as base_dir172def annotated_filepath(name: str) -> tuple[str, str | None]:173 if name.endswith("[output]"):174 base_dir = get_output_directory()175 name = name[:-9]176 elif name.endswith("[input]"):177 base_dir = get_input_directory()178 name = name[:-8]179 elif name.endswith("[temp]"):180 base_dir = get_temp_directory()181 name = name[:-7]182 else:183 return name, None184 185 return name, base_dir186 187 188def get_annotated_filepath(name: str, default_dir: str | None=None) -> str:189 name, base_dir = annotated_filepath(name)190 191 if base_dir is None:192 if default_dir is not None:193 base_dir = default_dir194 else:195 base_dir = get_input_directory() # fallback path196 197 return os.path.join(base_dir, name)198 199 200def exists_annotated_filepath(name) -> bool:201 name, base_dir = annotated_filepath(name)202 203 if base_dir is None:204 base_dir = get_input_directory() # fallback path205 206 filepath = os.path.join(base_dir, name)207 return os.path.exists(filepath)208 209 210def add_model_folder_path(folder_name: str, full_folder_path: str, is_default: bool = False) -> None:211 global folder_names_and_paths212 folder_name = map_legacy(folder_name)213 if folder_name in folder_names_and_paths:214 paths, _exts = folder_names_and_paths[folder_name]215 if full_folder_path in paths:216 if is_default and paths[0] != full_folder_path:217 # If the path to the folder is not the first in the list, move it to the beginning.218 paths.remove(full_folder_path)219 paths.insert(0, full_folder_path)220 else:221 if is_default:222 paths.insert(0, full_folder_path)223 else:224 paths.append(full_folder_path)225 else:226 folder_names_and_paths[folder_name] = ([full_folder_path], set())227 228def get_folder_paths(folder_name: str) -> list[str]:229 folder_name = map_legacy(folder_name)230 return folder_names_and_paths[folder_name][0][:]231 232def recursive_search(directory: str, excluded_dir_names: list[str] | None=None) -> tuple[list[str], dict[str, float]]:233 if not os.path.isdir(directory):234 return [], {}235 236 if excluded_dir_names is None:237 excluded_dir_names = []238 239 result = []240 dirs = {}241 242 # Attempt to add the initial directory to dirs with error handling243 try:244 dirs[directory] = os.path.getmtime(directory)245 except FileNotFoundError:246 logging.warning(f"Warning: Unable to access {directory}. Skipping this path.")247 248 logging.debug("recursive file list on directory {}".format(directory))249 dirpath: str250 subdirs: list[str]251 filenames: list[str]252 253 for dirpath, subdirs, filenames in os.walk(directory, followlinks=True, topdown=True):254 subdirs[:] = [d for d in subdirs if d not in excluded_dir_names]255 for file_name in filenames:256 try:257 relative_path = os.path.relpath(os.path.join(dirpath, file_name), directory)258 result.append(relative_path)259 except:260 logging.warning(f"Warning: Unable to access {file_name}. Skipping this file.")261 continue262 263 for d in subdirs:264 path: str = os.path.join(dirpath, d)265 try:266 dirs[path] = os.path.getmtime(path)267 except FileNotFoundError:268 logging.warning(f"Warning: Unable to access {path}. Skipping this path.")269 continue270 logging.debug("found {} files".format(len(result)))271 return result, dirs272 273def filter_files_extensions(files: Collection[str], extensions: Collection[str]) -> list[str]:274 return sorted(list(filter(lambda a: os.path.splitext(a)[-1].lower() in extensions or len(extensions) == 0, files)))275 276 277 278def get_full_path(folder_name: str, filename: str) -> str | None:279 global folder_names_and_paths280 folder_name = map_legacy(folder_name)281 if folder_name not in folder_names_and_paths:282 return None283 folders = folder_names_and_paths[folder_name]284 filename = os.path.relpath(os.path.join("/", filename), "/")285 for x in folders[0]:286 full_path = os.path.join(x, filename)287 if os.path.isfile(full_path):288 return full_path289 elif os.path.islink(full_path):290 logging.warning("WARNING path {} exists but doesn't link anywhere, skipping.".format(full_path))291 292 return None293 294 295def get_full_path_or_raise(folder_name: str, filename: str) -> str:296 full_path = get_full_path(folder_name, filename)297 if full_path is None:298 raise FileNotFoundError(f"Model in folder '{folder_name}' with filename '{filename}' not found.")299 return full_path300 301 302def get_filename_list_(folder_name: str) -> tuple[list[str], dict[str, float], float]:303 folder_name = map_legacy(folder_name)304 global folder_names_and_paths305 output_list = set()306 folders = folder_names_and_paths[folder_name]307 output_folders = {}308 for x in folders[0]:309 files, folders_all = recursive_search(x, excluded_dir_names=[".git"])310 output_list.update(filter_files_extensions(files, folders[1]))311 output_folders = {**output_folders, **folders_all}312 313 return sorted(list(output_list)), output_folders, time.perf_counter()314 315def cached_filename_list_(folder_name: str) -> tuple[list[str], dict[str, float], float] | None:316 strong_cache = cache_helper.get(folder_name)317 if strong_cache is not None:318 return strong_cache319 320 global filename_list_cache321 global folder_names_and_paths322 folder_name = map_legacy(folder_name)323 if folder_name not in filename_list_cache:324 return None325 out = filename_list_cache[folder_name]326 327 for x in out[1]:328 time_modified = out[1][x]329 folder = x330 if os.path.getmtime(folder) != time_modified:331 return None332 333 folders = folder_names_and_paths[folder_name]334 for x in folders[0]:335 if os.path.isdir(x):336 if x not in out[1]:337 return None338 339 return out340 341def get_filename_list(folder_name: str) -> list[str]:342 folder_name = map_legacy(folder_name)343 out = cached_filename_list_(folder_name)344 if out is None:345 out = get_filename_list_(folder_name)346 global filename_list_cache347 filename_list_cache[folder_name] = out348 cache_helper.set(folder_name, out)349 return list(out[0])350 351def get_save_image_path(filename_prefix: str, output_dir: str, image_width=0, image_height=0) -> tuple[str, str, int, str, str]:352 def map_filename(filename: str) -> tuple[int, str]:353 prefix_len = len(os.path.basename(filename_prefix))354 prefix = filename[:prefix_len + 1]355 try:356 digits = int(filename[prefix_len + 1:].split('_')[0])357 except:358 digits = 0359 return digits, prefix360 361 def compute_vars(input: str, image_width: int, image_height: int) -> str:362 input = input.replace("%width%", str(image_width))363 input = input.replace("%height%", str(image_height))364 now = time.localtime()365 input = input.replace("%year%", str(now.tm_year))366 input = input.replace("%month%", str(now.tm_mon).zfill(2))367 input = input.replace("%day%", str(now.tm_mday).zfill(2))368 input = input.replace("%hour%", str(now.tm_hour).zfill(2))369 input = input.replace("%minute%", str(now.tm_min).zfill(2))370 input = input.replace("%second%", str(now.tm_sec).zfill(2))371 return input372 373 if "%" in filename_prefix:374 filename_prefix = compute_vars(filename_prefix, image_width, image_height)375 376 subfolder = os.path.dirname(os.path.normpath(filename_prefix))377 filename = os.path.basename(os.path.normpath(filename_prefix))378 379 full_output_folder = os.path.join(output_dir, subfolder)380 381 if os.path.commonpath((output_dir, os.path.abspath(full_output_folder))) != output_dir:382 err = "**** ERROR: Saving image outside the output folder is not allowed." + \383 "\n full_output_folder: " + os.path.abspath(full_output_folder) + \384 "\n output_dir: " + output_dir + \385 "\n commonpath: " + os.path.commonpath((output_dir, os.path.abspath(full_output_folder)))386 logging.error(err)387 raise Exception(err)388 389 try:390 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] + 1391 except ValueError:392 counter = 1393 except FileNotFoundError:394 os.makedirs(full_output_folder, exist_ok=True)395 counter = 1396 return full_output_folder, filename, counter, subfolder, filename_prefix397