diff --git a/README.jp.md b/README.jp.md index ab82cac..15203e8 100644 --- a/README.jp.md +++ b/README.jp.md @@ -28,4 +28,9 @@ git clone https://github.com/nkchocoai/ComfyUI-SizeFromPresets.git ### Empty Latent Image From Presets (SDXL) - SDXL用の画像サイズのプリセットを選択できます。 - 選択したサイズの空のLatent Imageと幅と高さを出力します。 -- プリセットは [presets/sdxl.csv](presets/sdxl.csv) から設定できます。 \ No newline at end of file +- プリセットは [presets/sdxl.csv](presets/sdxl.csv) から設定できます。 + +### Random ... From Presets (SD...) +- CSVファイルに記載したプリセットからランダムに選びます。 +- Sizeの場合は、選ばれたサイズの幅と高さを出力します。 +- Latent Imageの場合は、選ばれたサイズの空のLatent Imageと幅と高さを出力します。 \ No newline at end of file diff --git a/README.md b/README.md index 849cbc8..d8cc6c5 100644 --- a/README.md +++ b/README.md @@ -30,4 +30,9 @@ git clone https://github.com/nkchocoai/ComfyUI-SizeFromPresets.git ### Empty Latent Image From Presets (SDXL) - Select image size presets for SDXL. - Output empty latent image, width and height of the selected preset. -- Presets can be set from [presets/sdxl.csv](presets/sdxl.csv). \ No newline at end of file +- Presets can be set from [presets/sdxl.csv](presets/sdxl.csv). + +### Random ... From Presets (SD...) +- Select randomly from the presets listed in the CSV file. +- In the case of Size, output width and height of the selected preset. +- In the case of Empty Latent Image, output empty latent image, width and height of the selected preset. \ No newline at end of file diff --git a/__init__.py b/__init__.py index e9dd978..50b5792 100644 --- a/__init__.py +++ b/__init__.py @@ -1,10 +1,15 @@ from .nodes.node import * +from .nodes.random import * NODE_CLASS_MAPPINGS = { "SizeFromPresetsSD15": SizeFromPresetsSD15, "SizeFromPresetsSDXL": SizeFromPresetsSDXL, "EmptyLatentImageFromPresetsSD15": EmptyLatentImageFromPresetsSD15, "EmptyLatentImageFromPresetsSDXL": EmptyLatentImageFromPresetsSDXL, + "RandomSizeFromPresetsSD15": RandomSizeFromPresetsSD15, + "RandomSizeFromPresetsSDXL": RandomSizeFromPresetsSDXL, + "RandomEmptyLatentImageFromPresetsSD15": RandomEmptyLatentImageFromPresetsSD15, + "RandomEmptyLatentImageFromPresetsSDXL": RandomEmptyLatentImageFromPresetsSDXL, } NODE_DISPLAY_NAME_MAPPINGS = { @@ -12,4 +17,8 @@ NODE_DISPLAY_NAME_MAPPINGS = { "SizeFromPresetsSDXL": "Size From Presets (SDXL)", "EmptyLatentImageFromPresetsSD15": "Empty Latent Image From Presets (SD1.5)", "EmptyLatentImageFromPresetsSDXL": "Empty Latent Image From Presets (SDXL)", + "RandomSizeFromPresetsSD15": "Random Size From Presets (SD1.5)", + "RandomSizeFromPresetsSDXL": "Random Size From Presets (SDXL)", + "RandomEmptyLatentImageFromPresetsSD15": "Random Empty Latent Image From Presets (SD1.5)", + "RandomEmptyLatentImageFromPresetsSDXL": "Random Empty Latent Image From Presets (SDXL)", } \ No newline at end of file diff --git a/nodes/random.py b/nodes/random.py new file mode 100644 index 0000000..349e744 --- /dev/null +++ b/nodes/random.py @@ -0,0 +1,82 @@ +import csv +import os + +import numpy as np +import torch + +import folder_paths + +import comfy.model_management + +custom_nodes_dir = folder_paths.get_folder_paths("custom_nodes")[0] +presets_dir = os.path.join(custom_nodes_dir, "ComfyUI-SizeFromPresets", "presets") + +def load_size_presets_input(file_name,digit): + with open(os.path.join(presets_dir, file_name),'r') as f: + reader = csv.reader(f) + data = [row for row in reader] + size_presets = [[int(v.strip()) for v in row] for row in data[1:]] + + return [f'{w: >{digit}} x {h: >{digit}}' for w,h in size_presets] + +class RandomSizeFromPresetsBase: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), + } + } + + RETURN_TYPES = ("INT","INT") + RETURN_NAMES = ("w","h") + + FUNCTION = "get_size" + + CATEGORY = "SizeFromPresets" + + def get_size(self, seed): + random_gen = np.random.default_rng(seed) + preset = random_gen.choice(self.SIZE_PRESETS_INPUT) + w, h = [int(v.strip()) for v in preset.split('x')] + return (w, h) + +class RandomSizeFromPresetsSD15(RandomSizeFromPresetsBase): + SIZE_PRESETS_INPUT = load_size_presets_input('sd15.csv',3) + +class RandomSizeFromPresetsSDXL(RandomSizeFromPresetsBase): + SIZE_PRESETS_INPUT = load_size_presets_input('sdxl.csv',4) + +class RandomEmptyLatentImageFromPresetsBase: + def __init__(self): + self.device = comfy.model_management.intermediate_device() + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "batch_size": ("INT", {"default": 1, "min": 1, "max": 4096}), + "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), + } + } + RETURN_TYPES = ("LATENT","INT","INT") + RETURN_NAMES = ("latent","w","h") + FUNCTION = "generate" + + CATEGORY = "SizeFromPresets" + + def generate(self, seed, batch_size=1): + random_gen = np.random.default_rng(seed) + preset = random_gen.choice(self.SIZE_PRESETS_INPUT) + w,h = [int(v.strip()) for v in preset.split('x')] + latent = torch.zeros([batch_size, 4, h // 8, w // 8], device=self.device) + return ({"samples":latent}, w, h) + +class RandomEmptyLatentImageFromPresetsSD15(RandomEmptyLatentImageFromPresetsBase): + SIZE_PRESETS_INPUT = load_size_presets_input('sd15.csv',3) + +class RandomEmptyLatentImageFromPresetsSDXL(RandomEmptyLatentImageFromPresetsBase): + SIZE_PRESETS_INPUT = load_size_presets_input('sdxl.csv',4) \ No newline at end of file