fix random seed to max length of the sizes, and add custom sizes folder to the sizes searching function

README.md

@@ -1,6 +1,6 @@
 # ComfyUI-RandomSize
 
-A ComfyUI custom node that randomly selects a height and width pair from a list in a config file. 
+A ComfyUI custom node that randomly selects a height and width pair from a list in a config file.
 The core use is to provide variety in generations of large numbers of images at once.
 
 Presets included:
@@ -20,7 +20,7 @@ git clone https://github.com/JerryOrbachJr/ComfyUI-RandomSize.git
 ```
 
 ## Usage
-* Select a seed. The same seed should produce the same size every time
+* Select a seed. The seed is the index of a size in the yaml file, startting with 0. The same seed will always give the same output. Set control after generate to fix to fix the output.
 * Control After(/Before) Generate: set to increment, decrement, or random to get a random size from the preset each generation.
 * Select a preset. The default preset is SD 1.5 if none is chosen.
 * Run your queue.

js/show-sizes.js

@@ -30,6 +30,63 @@ const ext = {
     async beforeRegisterNodeDef(nodeType, nodeData, app) {
         if (nodeType.comfyClass=="JOJR_RandomSize" ) {
             nodeType.prototype.displayMessage = displayMessage;
+            const originalNodeCreated = nodeType.prototype.onNodeCreated;
+            nodeType.prototype.onNodeCreated = async function () {
+              if (originalNodeCreated) {
+                originalNodeCreated.apply(this, arguments);
+              }
+
+              const presetWidget = this.widgets.find((w) => w.name === "preset");
+              const sizeIndexWidget = this.widgets.find((w) => w.name === "seed");
+
+              const fetchSizes = async (preset) => {
+                try {
+                  const response = await fetch("/JOJR_RandomSize/get_sizes", {
+                    method: "POST",
+                    headers: {
+                      "Content-Type": "application/json",
+                    },
+                    body: JSON.stringify({
+                      preset: preset,
+                    }),
+                  });
+
+                  if (response.ok) {
+                    const sizes = await response.json();
+                    console.log("Fetched sizes:", sizes);
+                    return sizes;
+                  } else {
+                    console.error(`Failed to fetch sizes: ${response.status}`);
+                    return [];
+                  }
+                } catch (error) {
+                  console.error(`Error fetching sizes for preset ${preset}:`, error);
+                  return [];
+                }
+              };
+
+              const updateSizeIndex = async () => {
+                const preset = presetWidget.value;
+
+                console.log(`Selected preset: ${preset}`);
+
+                const sizes = await fetchSizes(preset);
+
+                if (!sizes || !sizes.length) return;
+
+                sizeIndexWidget.options.max = sizes.length;
+                if(sizeIndexWidget.options.value > sizes.length) {
+                    sizeIndexWidget.options.value = sizes.length;
+                }
+
+                this.triggerSlot(0);
+              };
+
+              presetWidget.callback = updateSizeIndex;
+
+              // Initial update
+              await updateSizeIndex();
+            };
         }
     }
 }

random_size.py

@@ -1,5 +1,6 @@
 import yaml, os, random
 from server import PromptServer
+from aiohttp import web
 sizes_dir = os.path.join(os.path.dirname(__file__), "sizes")
 sizes_dir = os.path.abspath(sizes_dir)
 custom_sizes_dir = os.path.join(os.path.dirname(__file__), "sizes","custom")
@@ -7,18 +8,36 @@ custom_sizes_dir = os.path.abspath(custom_sizes_dir)
 
 
 def get_sizes_from_preset_file(preset_file):
-    if preset_file == "Preset":
-        preset_file_path = os.path.join(sizes_dir,"SD1.5.yaml")
+    if not preset_file or preset_file == "Preset":
+        preset_file_path = os.path.join(sizes_dir, "SD1.5.yaml")
     else:
-        preset_file_path = os.path.join(sizes_dir,preset_file)
-    with open( preset_file_path, 'r') as file:
+        preset_file_path = os.path.join(custom_sizes_dir, preset_file)
+
+        if not os.path.exists(preset_file_path):
+            preset_file_path = os.path.join(sizes_dir, preset_file)
+
+            if not os.path.exists(preset_file_path):
+                raise FileNotFoundError(f"Preset file '{preset_file}' not found in either {sizes_dir} or {custom_sizes_dir}.")
+
+    with open(preset_file_path, 'r') as file:
         return yaml.safe_load(file)['sizes']
 
+
 def load_presets():
     default_list = [dfile for dfile in os.listdir(sizes_dir) if dfile.endswith(".yaml")]
     user_list = [ufile for ufile in os.listdir(custom_sizes_dir) if ufile.endswith(".yaml")]
     return default_list + user_list
 
+
+@PromptServer.instance.routes.post("/JOJR_RandomSize/get_sizes")
+async def get_sizes_endpoint(request):
+    data = await request.json()
+    preset = data.get("preset")
+
+    sizes = get_sizes_from_preset_file(preset)
+    return web.json_response(sizes)
+
+
 class RandomSize:
     CATEGORY = "utils"
     @classmethod
@@ -26,7 +45,7 @@ class RandomSize:
         preset = ["Preset"]
         preset += load_presets()
         return {
-            "required": { 
+            "required": {
                 "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
                 "preset": (preset,)
             },
@@ -37,13 +56,13 @@ class RandomSize:
     FUNCTION = "func"
     def func(self:str,seed, preset,id):
         sizes = get_sizes_from_preset_file(preset)
-        
+
         rand_obj = random.Random(seed)
         size = rand_obj.choice(sizes)
         for i in range(len(sizes)):
             if sizes[i] == size:
                 sizes[i] = f"*{size}*"
-    
+
         PromptServer.instance.send_sync("jojr.random-sizes.sendmessage", {"id": id, "message":sizes})
         w, h = [int(i) for i in size.split('x')]
         return (w,h)