diff --git a/ui/easydiffusion/app.py b/ui/easydiffusion/app.py index 106db364..83bb08c1 100644 --- a/ui/easydiffusion/app.py +++ b/ui/easydiffusion/app.py @@ -5,6 +5,7 @@ import json import traceback import logging import shlex +import urllib from rich.logging import RichHandler from sdkit.utils import log as sdkit_log # hack, so we can overwrite the log config @@ -54,6 +55,10 @@ APP_CONFIG_DEFAULTS = { }, } +IMAGE_EXTENSIONS = [".png", ".apng", ".jpg", ".jpeg", ".jfif", ".pjpeg", ".pjp", ".jxl", ".gif", ".webp", ".avif", ".svg"] +CUSTOM_MODIFIERS_DIR = os.path.abspath(os.path.join(SD_DIR, "..", "modifiers")) +CUSTOM_MODIFIERS_PORTRAIT_EXTENSIONS=[".portrait", "_portrait", " portrait", "-portrait"] +CUSTOM_MODIFIERS_LANDSCAPE_EXTENSIONS=[".landscape", "_landscape", " landscape", "-landscape"] def init(): os.makedirs(USER_UI_PLUGINS_DIR, exist_ok=True) @@ -234,3 +239,90 @@ def open_browser(): import webbrowser webbrowser.open(f"http://localhost:{port}") + +def get_image_modifiers(): + modifiers_json_path = os.path.join(SD_UI_DIR, "modifiers.json") + + modifier_categories = {} + original_category_order=[] + with open(modifiers_json_path, "r", encoding="utf-8") as f: + modifiers_file = json.load(f) + + # The trailing slash is needed to support symlinks + if not os.path.isdir(f"{CUSTOM_MODIFIERS_DIR}/"): + return modifiers_file + + # convert modifiers from a list of objects to a dict of dicts + for category_item in modifiers_file: + category_name = category_item['category'] + original_category_order.append(category_name) + category = {} + for modifier_item in category_item['modifiers']: + modifier = {} + for preview_item in modifier_item['previews']: + modifier[preview_item['name']] = preview_item['path'] + category[modifier_item['modifier']] = modifier + modifier_categories[category_name] = category + + def scan_directory(directory_path: str, category_name="Modifiers"): + for entry in os.scandir(directory_path): + if entry.is_file(): + file_extension = list(filter(lambda e: entry.name.endswith(e), IMAGE_EXTENSIONS)) + if len(file_extension) == 0: + continue + + modifier_name = entry.name[: -len(file_extension[0])] + modifier_path = f"custom/{entry.path[len(CUSTOM_MODIFIERS_DIR) + 1:]}" + # URL encode path segments + modifier_path = "/".join(map(lambda segment: urllib.parse.quote(segment), modifier_path.split("/"))) + is_portrait = True + is_landscape = True + + portrait_extension = list(filter(lambda e: modifier_name.lower().endswith(e), CUSTOM_MODIFIERS_PORTRAIT_EXTENSIONS)) + landscape_extension = list(filter(lambda e: modifier_name.lower().endswith(e), CUSTOM_MODIFIERS_LANDSCAPE_EXTENSIONS)) + + if len(portrait_extension) > 0: + is_landscape = False + modifier_name = modifier_name[: -len(portrait_extension[0])] + elif len(landscape_extension) > 0: + is_portrait = False + modifier_name = modifier_name[: -len(landscape_extension[0])] + + if (category_name not in modifier_categories): + modifier_categories[category_name] = {} + + category = modifier_categories[category_name] + + if (modifier_name not in category): + category[modifier_name] = {} + + if (is_portrait or "portrait" not in category[modifier_name]): + category[modifier_name]["portrait"] = modifier_path + + if (is_landscape or "landscape" not in category[modifier_name]): + category[modifier_name]["landscape"] = modifier_path + elif entry.is_dir(): + scan_directory( + entry.path, + entry.name if directory_path==CUSTOM_MODIFIERS_DIR else f"{category_name}/{entry.name}", + ) + + scan_directory(CUSTOM_MODIFIERS_DIR) + + custom_categories = sorted( + [cn for cn in modifier_categories.keys() if cn not in original_category_order], + key=str.casefold, + ) + + # convert the modifiers back into a list of objects + modifier_categories_list = [] + for category_name in [*original_category_order, *custom_categories]: + category = { 'category': category_name, 'modifiers': [] } + for modifier_name in sorted(modifier_categories[category_name].keys(), key=str.casefold): + modifier = { 'modifier': modifier_name, 'previews': [] } + for preview_name, preview_path in modifier_categories[category_name][modifier_name].items(): + modifier['previews'].append({ 'name': preview_name, 'path': preview_path }) + category['modifiers'].append(modifier) + modifier_categories_list.append(category) + + return modifier_categories_list diff --git a/ui/easydiffusion/server.py b/ui/easydiffusion/server.py index c04189c1..77abd41f 100644 --- a/ui/easydiffusion/server.py +++ b/ui/easydiffusion/server.py @@ -25,6 +25,13 @@ NOCACHE_HEADERS = {"Cache-Control": "no-cache, no-store, must-revalidate", "Prag class NoCacheStaticFiles(StaticFiles): + def __init__(self, directory: str): + # follow_symlink is only available on fastapi >= 0.92.0 + if (os.path.islink(directory)): + super().__init__(directory = os.path.realpath(directory)) + else: + super().__init__(directory = directory) + def is_not_modified(self, response_headers, request_headers) -> bool: if "content-type" in response_headers and ( "javascript" in response_headers["content-type"] or "css" in response_headers["content-type"] @@ -45,6 +52,13 @@ class SetAppConfigRequest(BaseModel): def init(): + if os.path.isdir(app.CUSTOM_MODIFIERS_DIR): + server_api.mount( + "/media/modifier-thumbnails/custom", + NoCacheStaticFiles(directory=app.CUSTOM_MODIFIERS_DIR), + name="custom-thumbnails", + ) + server_api.mount("/media", NoCacheStaticFiles(directory=os.path.join(app.SD_UI_DIR, "media")), name="media") for plugins_dir, dir_prefix in app.UI_PLUGINS_SOURCES: @@ -156,7 +170,7 @@ def read_web_data_internal(key: str = None): elif key == "models": return JSONResponse(model_manager.getModels(), headers=NOCACHE_HEADERS) elif key == "modifiers": - return FileResponse(os.path.join(app.SD_UI_DIR, "modifiers.json"), headers=NOCACHE_HEADERS) + return JSONResponse(app.get_image_modifiers(), headers=NOCACHE_HEADERS) elif key == "ui_plugins": return JSONResponse(app.getUIPlugins(), headers=NOCACHE_HEADERS) else: