From a2c73de56d6c9e1c545fa6a466cd4e8f2468d28e Mon Sep 17 00:00:00 2001 From: geetesha Date: Fri, 9 Aug 2024 23:40:16 -0700 Subject: [PATCH] fix random seed to max length of the sizes, and add custom sizes folder to the sizes searching function --- README.md | 4 ++-- js/show-sizes.js | 57 ++++++++++++++++++++++++++++++++++++++++++++++++ random_size.py | 33 ++++++++++++++++++++++------ 3 files changed, 85 insertions(+), 9 deletions(-) diff --git a/README.md b/README.md index 31a992a..4fb96bd 100644 --- a/README.md +++ b/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. diff --git a/js/show-sizes.js b/js/show-sizes.js index e0fed8d..5b1dae1 100644 --- a/js/show-sizes.js +++ b/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(); + }; } } } diff --git a/random_size.py b/random_size.py index e537c40..2af1ab9 100644 --- a/random_size.py +++ b/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) \ No newline at end of file