initial commit

scripts/reactor_helpers.py

@@ -0,0 +1,235 @@
+import os, glob, random
+from collections import Counter
+from PIL import Image
+from math import isqrt, ceil
+from typing import List
+import logging
+import hashlib
+import torch
+from safetensors.torch import save_file, safe_open
+from insightface.app.common import Face
+
+from modules.images import FilenameGenerator, get_next_sequence_number
+from modules import shared, script_callbacks
+from scripts.reactor_globals import DEVICE, BASE_PATH, FACE_MODELS_PATH, IS_SDNEXT
+
+try:
+    from modules.paths_internal import models_path
+except:
+    try:
+        from modules.paths import models_path
+    except:
+        model_path = os.path.abspath("models")
+
+MODELS_PATH = None
+
+def set_Device(value):
+    global DEVICE
+    DEVICE = value
+    with open(os.path.join(BASE_PATH, "last_device.txt"), "w") as txt:
+        txt.write(DEVICE)
+
+def get_Device():
+    global DEVICE
+    return DEVICE
+
+def set_SDNEXT():
+    global IS_SDNEXT
+    IS_SDNEXT = True
+
+def get_SDNEXT():
+    global IS_SDNEXT
+    return IS_SDNEXT
+
+def make_grid(image_list: List):
+    
+    # Count the occurrences of each image size in the image_list
+    size_counter = Counter(image.size for image in image_list)
+    
+    # Get the most common image size (size with the highest count)
+    common_size = size_counter.most_common(1)[0][0]
+    
+    # Filter the image_list to include only images with the common size
+    image_list = [image for image in image_list if image.size == common_size]
+    
+    # Get the dimensions (width and height) of the common size
+    size = common_size
+    
+    # If there are more than one image in the image_list
+    if len(image_list) > 1:
+        num_images = len(image_list)
+        
+        # Calculate the number of rows and columns for the grid
+        rows = isqrt(num_images)
+        cols = ceil(num_images / rows)
+
+        # Calculate the size of the square image
+        square_size = (cols * size[0], rows * size[1])
+
+        # Create a new RGB image with the square size
+        square_image = Image.new("RGB", square_size)
+
+        # Paste each image onto the square image at the appropriate position
+        for i, image in enumerate(image_list):
+            row = i // cols
+            col = i % cols
+
+            square_image.paste(image, (col * size[0], row * size[1]))
+
+        # Return the resulting square image
+        return square_image
+    
+    # Return None if there are no images or only one image in the image_list
+    return None
+
+def get_image_path(image, path, basename, seed=None, prompt=None, extension='png', p=None, suffix=""):
+    
+    namegen = FilenameGenerator(p, seed, prompt, image)
+
+    save_to_dirs = shared.opts.save_to_dirs
+
+    if save_to_dirs:
+        dirname = namegen.apply(shared.opts.directories_filename_pattern or "[prompt_words]").lstrip(' ').rstrip('\\ /')
+        path = os.path.join(path, dirname)
+
+    os.makedirs(path, exist_ok=True)
+
+    if seed is None:
+        file_decoration = ""
+    elif shared.opts.save_to_dirs:
+        file_decoration = shared.opts.samples_filename_pattern or "[seed]"
+    else:
+        file_decoration = shared.opts.samples_filename_pattern or "[seed]-[prompt_spaces]"
+
+    file_decoration = namegen.apply(file_decoration) + suffix
+
+    add_number = shared.opts.save_images_add_number or file_decoration == ''
+
+    if file_decoration != "" and add_number:
+        file_decoration = f"-{file_decoration}"
+
+    if add_number:
+        basecount = get_next_sequence_number(path, basename)
+        fullfn = None
+        for i in range(500):
+            fn = f"{basecount + i:05}" if basename == '' else f"{basename}-{basecount + i:04}"
+            fullfn = os.path.join(path, f"{fn}{file_decoration}.{extension}")
+            if not os.path.exists(fullfn):
+                break
+    else:
+        fullfn = os.path.join(path, f"{file_decoration}.{extension}")
+
+    pnginfo = {}
+
+    params = script_callbacks.ImageSaveParams(image, p, fullfn, pnginfo)
+    # script_callbacks.before_image_saved_callback(params)
+
+    fullfn = params.filename
+
+    fullfn_without_extension, extension = os.path.splitext(params.filename)
+    if hasattr(os, 'statvfs'):
+        max_name_len = os.statvfs(path).f_namemax
+        fullfn_without_extension = fullfn_without_extension[:max_name_len - max(4, len(extension))]
+        params.filename = fullfn_without_extension + extension
+        fullfn = params.filename
+
+    return fullfn
+
+def addLoggingLevel(levelName, levelNum, methodName=None):
+    if not methodName:
+        methodName = levelName.lower()
+
+    def logForLevel(self, message, *args, **kwargs):
+        if self.isEnabledFor(levelNum):
+            self._log(levelNum, message, args, **kwargs)
+
+    def logToRoot(message, *args, **kwargs):
+        logging.log(levelNum, message, *args, **kwargs)
+
+    logging.addLevelName(levelNum, levelName)
+    setattr(logging, levelName, levelNum)
+    setattr(logging.getLoggerClass(), methodName, logForLevel)
+    setattr(logging, methodName, logToRoot)
+
+def get_image_md5hash(image: Image.Image):
+    md5hash = hashlib.md5(image.tobytes())
+    return md5hash.hexdigest()
+
+def save_face_model(face: Face, filename: str) -> None:
+    try:
+        tensors = {
+            "bbox": torch.tensor(face["bbox"]),
+            "kps": torch.tensor(face["kps"]),
+            "det_score": torch.tensor(face["det_score"]),
+            "landmark_3d_68": torch.tensor(face["landmark_3d_68"]),
+            "pose": torch.tensor(face["pose"]),
+            "landmark_2d_106": torch.tensor(face["landmark_2d_106"]),
+            "embedding": torch.tensor(face["embedding"]),
+            "gender": torch.tensor(face["gender"]),
+            "age": torch.tensor(face["age"]),
+        }
+        save_file(tensors, filename)
+        # print(f"Face model has been saved to '{filename}'")
+    except Exception as e:
+        print(f"Error: {e}")
+
+def get_models():
+    global MODELS_PATH
+    models_path_init = os.path.join(models_path, "insightface/*")
+    models = glob.glob(models_path_init)
+    models = [x for x in models if x.endswith(".onnx") or x.endswith(".pth")]
+    models_names = []
+    for model in models:
+        model_path = os.path.split(model)
+        if MODELS_PATH is None:
+            MODELS_PATH = model_path[0]
+        model_name = model_path[1]
+        models_names.append(model_name)
+    return models_names
+
+def load_face_model(filename: str):
+    face = {}
+    model_path = os.path.join(FACE_MODELS_PATH, filename)
+    with safe_open(model_path, framework="pt") as f:
+        for k in f.keys():
+            face[k] = f.get_tensor(k).numpy()
+    return Face(face)
+
+def get_facemodels():
+    models_path = os.path.join(FACE_MODELS_PATH, "*")
+    models = glob.glob(models_path)
+    models = [x for x in models if x.endswith(".safetensors")]
+    return models
+
+def get_model_names(get_models):
+    models = get_models()
+    names = ["None"]
+    for x in models:
+        names.append(os.path.basename(x))
+    return names
+
+def get_images_from_folder(path: str):
+    files_path = os.path.join(path, "*")
+    files = glob.glob(files_path)
+    images = []
+    images_names = []
+    for x in files:
+        if x.endswith(('jpg', 'png', 'jpeg', 'webp', 'bmp')):
+            images.append(Image.open(x))
+            images_names.append(os.path.basename(x))
+    return images,images_names
+    # return [Image.open(x) for x in images if x.endswith(('jpg', 'png', 'jpeg', 'webp', 'bmp'))],[os.path.basename(x) for x in images if x.endswith(('jpg', 'png', 'jpeg', 'webp', 'bmp'))]
+
+def get_random_image_from_folder(path: str):
+    images,names = get_images_from_folder(path)
+    random_image_index = random.randint(0, len(images) - 1)
+    return [images[random_image_index]],[names[random_image_index]]
+
+def get_images_from_list(imgs: List):
+    images = []
+    images_names = []
+    for x in imgs:
+        images.append(Image.open(os.path.abspath(x.name)))
+        images_names.append(os.path.basename(x.name))
+    return images,images_names
+    # return [Image.open(os.path.abspath(x.name)) for x in imgs],[os.path.basename(x.name) for x in imgs]