From 0643aff59d40c8f671797b3df4d0787ca182df48 Mon Sep 17 00:00:00 2001
From: sn0w12 <lucasskold49@gmail.com>
Date: Wed, 31 Jan 2024 15:05:43 +0100
Subject: [PATCH] Get font size now takes an image input, added get_all_styles

---
 __init__.py           |  3 +++
 src/get_all_styles.py | 41 +++++++++++++++++++++++++++++++++++++++++
 src/get_font_size.py  | 13 +++++++++++--
 src/lora_tester.py    |  4 ++--
 4 files changed, 57 insertions(+), 4 deletions(-)
 create mode 100644 src/get_all_styles.py

diff --git a/__init__.py b/__init__.py
index 22f0393..725ff3c 100644
--- a/__init__.py
+++ b/__init__.py
@@ -11,6 +11,7 @@ from .src.load_lora_xl import LoraLoraXLNode
 from .src.load_lora_15 import LoraLora15Node
 from .src.load_lora_character import LoadLoraCharacterNode
 from .src.get_font_size import GetFontSizeNode
+from .src.get_all_styles import GetStylesNode
 
 NODE_CLASS_MAPPINGS = {
     "Find SDXL Resolution": FindResolutionNode,
@@ -26,6 +27,7 @@ NODE_CLASS_MAPPINGS = {
     "Load Lora 1.5": LoraLora15Node,
     "Load Lora Character": LoadLoraCharacterNode,
     "Get Font Size": GetFontSizeNode,
+    "Get All Styles": GetStylesNode,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
@@ -42,4 +44,5 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "Load Lora 1.5": "Load Lora 1.5",
     "Load Lora Character": "Load Lora Character",
     "Get Font Size": "Get Font Size",
+    "Get All Styles": "Get All Styles",
 }
diff --git a/src/get_all_styles.py b/src/get_all_styles.py
new file mode 100644
index 0000000..d1a506b
--- /dev/null
+++ b/src/get_all_styles.py
@@ -0,0 +1,41 @@
+import folder_paths
+
+class GetStylesNode:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "xl": ("BOOLEAN", {"default": False},),
+                "lora_strength": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}),
+                "add_default_generation": ("BOOLEAN", {"default": False},),
+            },
+        }
+
+    RETURN_TYPES = ("STRING", "INT",)
+    RETURN_NAMES = ("LORA_INFO", "TOTAL_LORAS",)
+    FUNCTION = "get_styles"
+    CATEGORY = "sn0w"
+        
+    def get_styles(self, xl, lora_strength, add_default_generation):
+        # Select the appropriate lora list based on 'xl'
+        if xl:
+            lora_paths = folder_paths.get_filename_list("loras_xl")
+        else:
+            lora_paths = folder_paths.get_filename_list("loras_15")
+
+        loras = []
+        for lora in lora_paths:
+            if "style" in lora.lower():
+                # Extract the lora string from the file path
+                lora_string = lora.split('\\')[-1].replace('.safetensors', '')
+                loras.append(f"<lora:{lora_string}:{lora_strength:.1f}>")
+
+        total_loras = len(loras)
+        if add_default_generation:
+            total_loras += 1
+            loras.append("Nothing")
+
+        # Join the lora strings into a single string separated by semicolons
+        lora_output = ';'.join(loras)
+
+        return (lora_output, total_loras, )
diff --git a/src/get_font_size.py b/src/get_font_size.py
index 6f092c0..40e5272 100644
--- a/src/get_font_size.py
+++ b/src/get_font_size.py
@@ -1,9 +1,11 @@
+import torch
+
 class GetFontSizeNode:
     @classmethod
     def INPUT_TYPES(cls):
         return {
             "required": {
-                "width": ("INT", {"default": 0, "min": 0}),
+                "image": ("IMAGE", ),
                 "lora_info": ("STRING", {"default": ""}),
             },
         }
@@ -13,7 +15,7 @@ class GetFontSizeNode:
     FUNCTION = "estimate_font_size"
     CATEGORY = "sn0w"
         
-    def estimate_font_size(self, width, lora_info):
+    def estimate_font_size(self, image, lora_info):
         loras = lora_info.split(";")
         longest_string = max(loras, key=len)
         text_length = len(longest_string)
@@ -21,6 +23,13 @@ class GetFontSizeNode:
         tolerance = 50
         estimated_char_width = lambda font_size: font_size * 0.5
 
+        # Attempt to determine if the image is a batch or a single image
+        try:
+            # Assume image is a batch and try to get the width of the first image
+            width = image[0].shape[1]  # Image batch format: [N, C, H, W]
+        except TypeError:
+            # If it's not a batch, get the width as a single image
+            width = image.shape[1]  # Single image format: [C, H, W]
         # Start with a guess that is proportionate to the width
         font_size = width // (text_length * 0.5)
         step_size = max(1, font_size // 2)  # Start with a larger step size
diff --git a/src/lora_tester.py b/src/lora_tester.py
index fe6977b..c3fee5e 100644
--- a/src/lora_tester.py
+++ b/src/lora_tester.py
@@ -64,9 +64,9 @@ class LoraTestNode:
                 upscaled_image = upscaler.upscale(upscale_model, image)[0]
                 upscaled_latent = vae_encode.encode(vae, upscaled_image)[0]
                 upscaled_samples = k_sampleradvanced_node.sample(modified_model, "enable", seed, 30, cfg, sampler_name, scheduler, positive_prompt, negative_prompt, upscaled_latent, 17, 1000, "disable")[0]
-                image = vae_decode.decode(vae, upscaled_samples)
+                image = vae_decode.decode(vae, upscaled_samples)[0]
             else:
-                image = vae_decode.decode(vae, samples)
+                image = vae_decode.decode(vae, samples)[0]
 
             images.append(image)