From b5a661eec8c8d581562e7c27f72ed1f62e8d4b73 Mon Sep 17 00:00:00 2001 From: Olivia Godone-Maresca Date: Sun, 9 Jul 2023 22:13:20 -0400 Subject: [PATCH 1/2] Add embeddings to metadata --- ui/easydiffusion/utils/save_utils.py | 28 +++++++++++++++++++++++++--- 1 file changed, 25 insertions(+), 3 deletions(-) diff --git a/ui/easydiffusion/utils/save_utils.py b/ui/easydiffusion/utils/save_utils.py index ff2906a6..78857e92 100644 --- a/ui/easydiffusion/utils/save_utils.py +++ b/ui/easydiffusion/utils/save_utils.py @@ -1,6 +1,7 @@ import os import re import time +import regex from datetime import datetime from functools import reduce @@ -30,11 +31,12 @@ TASK_TEXT_MAPPING = { "lora_alpha": "LoRA Strength", "use_hypernetwork_model": "Hypernetwork model", "hypernetwork_strength": "Hypernetwork Strength", + "use_embedding_models": "Embedding models", "tiling": "Seamless Tiling", "use_face_correction": "Use Face Correction", "use_upscale": "Use Upscaling", "upscale_amount": "Upscale By", - "latent_upscaler_steps": "Latent Upscaler Steps" + "latent_upscaler_steps": "Latent Upscaler Steps", } time_placeholders = { @@ -202,6 +204,9 @@ def get_printable_request(req: GenerateImageRequest, task_data: TaskData): req_metadata = req.dict() task_data_metadata = task_data.dict() + app_config = app.getConfig() + using_diffusers = app_config.get("test_diffusers", False) + # Save the metadata in the order defined in TASK_TEXT_MAPPING metadata = {} for key in TASK_TEXT_MAPPING.keys(): @@ -209,6 +214,24 @@ def get_printable_request(req: GenerateImageRequest, task_data: TaskData): metadata[key] = req_metadata[key] elif key in task_data_metadata: metadata[key] = task_data_metadata[key] + elif key is "use_embedding_models" and using_diffusers: + embeddings_extensions = {".pt", ".bin", ".safetensors"} + def scan_directory(directory_path: str): + used_embeddings = [] + for entry in os.scandir(directory_path): + if entry.is_file(): + entry_extension = os.path.splitext(entry.name)[1] + if entry_extension not in embeddings_extensions: + continue + + embedding_name_regex = regex.compile(r"(^|[\s,])" + regex.escape(os.path.splitext(entry.name)[0]) + r"([+-]*$|[\s,]|[+-]+[\s,])") + if embedding_name_regex.search(req.prompt) or embedding_name_regex.search(req.negative_prompt): + used_embeddings.append(entry.path) + elif entry.is_dir(): + used_embeddings.extend(scan_directory(entry.path)) + return used_embeddings + used_embeddings = scan_directory(os.path.join(app.MODELS_DIR, "embeddings")) + metadata["use_embedding_models"] = ", ".join(used_embeddings) if len(used_embeddings) > 0 else None # Clean up the metadata if req.init_image is None and "prompt_strength" in metadata: @@ -222,8 +245,7 @@ def get_printable_request(req: GenerateImageRequest, task_data: TaskData): if task_data.use_upscale != "latent_upscaler" and "latent_upscaler_steps" in metadata: del metadata["latent_upscaler_steps"] - app_config = app.getConfig() - if not app_config.get("test_diffusers", False): + if not using_diffusers: for key in (x for x in ["use_lora_model", "lora_alpha", "clip_skip", "tiling", "latent_upscaler_steps"] if x in metadata): del metadata[key] From ab75527df248f57d0580ffcb0a13b97af214b90d Mon Sep 17 00:00:00 2001 From: Olivia Godone-Maresca Date: Sun, 9 Jul 2023 22:14:07 -0400 Subject: [PATCH 2/2] Formatting --- ui/easydiffusion/utils/save_utils.py | 1 + 1 file changed, 1 insertion(+) diff --git a/ui/easydiffusion/utils/save_utils.py b/ui/easydiffusion/utils/save_utils.py index 78857e92..75d35dc8 100644 --- a/ui/easydiffusion/utils/save_utils.py +++ b/ui/easydiffusion/utils/save_utils.py @@ -2,6 +2,7 @@ import os import re import time import regex + from datetime import datetime from functools import reduce