From 8103bf98ad53ec410d810a65f38528dd44cbbbaa Mon Sep 17 00:00:00 2001 From: Budi Hartono Date: Sat, 7 Jun 2025 22:19:10 +0700 Subject: [PATCH] Works for the first time --- __init__.py | 8 +- nodes.py | 206 ++++++++++++++++++++++------------------------------ 2 files changed, 91 insertions(+), 123 deletions(-) diff --git a/__init__.py b/__init__.py index 09fe93f..9b235d4 100644 --- a/__init__.py +++ b/__init__.py @@ -1,10 +1,8 @@ -# __init__.py - -from .nodes import empty_latent_preset, empty_latent_by_axis +from .nodes import EmptyLatentAspectPreset, EmptyLatentAspectByAxis NODE_CLASS_MAPPINGS = { - "Empty Latent (Aspect Ratio Preset)": empty_latent_preset, - "Empty Latent (Aspect Ratio by Axis)": empty_latent_by_axis, + "Empty Latent (Aspect Ratio Preset)": EmptyLatentAspectPreset, + "Empty Latent (Aspect Ratio by Axis)": EmptyLatentAspectByAxis, } NODE_DISPLAY_NAME = "ComfyUI Aspect Ratio Preset" \ No newline at end of file diff --git a/nodes.py b/nodes.py index 8ad28f7..4421779 100644 --- a/nodes.py +++ b/nodes.py @@ -1,130 +1,100 @@ -# nodes.py - -import numpy as np +import torch from .presets import PRESETS -# ─── Helpers ────────────────────────────────────────────────────────── - def _validate_dim(v: int): - if v <= 0: - raise ValueError("Dimension must be > 0") - if v % 8 != 0: - raise ValueError("Dimension must be divisible by 8") - -# ─── Node 1: Preset ─────────────────────────────────────────────────── - -def empty_latent_preset( - preset: str, - batch_size: int -) -> np.ndarray: - """ - Empty Latent (Aspect Ratio Preset) - - preset: e.g. "512x512 - 1:1 SD 1.5" - batch_size: how many latents to produce - """ - # find matching entry - for model, lbl, W, H in PRESETS: - key = f"{W}x{H} - {lbl.split()[0]} {model}" - if key == preset: - break - else: - raise ValueError(f"Unknown preset: {preset}") - - _validate_dim(W); _validate_dim(H) - return np.zeros((batch_size, 4, H // 8, W // 8), dtype=np.float32) + if v <= 0 or v % 8 != 0: + raise ValueError("Dimension must be >0 and divisible by 8") -def _preset_input_types(): - # Build the list of dropdown labels: - choices = [ - f"{W}x{H} - {lbl.split()[0]} {model}" - for model, lbl, W, H in PRESETS +class EmptyLatentAspectPreset: + """Creates a blank latent using one of the predefined presets.""" + def __init__(self): + # Build a mapping from the dropdown key to (W, H) + self._map = {} + for model, lbl, w, h in PRESETS: + # "512x512 - 1:1 SD 1.5" + key = f"{w}x{h} - {lbl.split()[0]} {model}" + self._map[key] = (w, h) + + @classmethod + def INPUT_TYPES(cls): + choices = [f"{w}x{h} - {lbl.split()[0]} {model}" for model, lbl, w, h in PRESETS] + return { + "required": { + "preset": (choices,), + "batch_size": ("INT", {"default": 1, "min": 1}) + } + } + + RETURN_TYPES = ("LATENT",) + RETURN_NAMES = ("empty_latent",) + FUNCTION = "generate" + CATEGORY = "ComfyUI Aspect Ratio Preset" + + def generate(self, preset: str, batch_size: int): + if preset not in self._map: + raise ValueError(f"Unknown preset: {preset}") + w, h = self._map[preset] + + _validate_dim(w) + _validate_dim(h) + + latent = torch.zeros([batch_size, 4, h // 8, w // 8], dtype=torch.float32) + # return a single‐element tuple of a dict + return ({"samples": latent},) + + +class EmptyLatentAspectByAxis: + """Creates a blank latent by fixing one axis and computing the other from an aspect ratio.""" + # Precompute aspect choices + ASPECT_CHOICES = [ + ("1:1 Square", (1, 1)), + ("3:2 Landscape", (3, 2)), + ("2:3 Portrait", (2, 3)), + ("4:3 Landscape", (4, 3)), + ("3:4 Portrait", (3, 4)), + ("16:9 Landscape", (16, 9)), + ("9:16 Portrait", (9, 16)), + ("5:4 Landscape", (5, 4)), + ("4:5 Portrait", (4, 5)), + ("21:9 Widescreen",(21, 9)), + ("9:21 Portrait", (9, 21)), + ("7:5 Landscape", (7, 5)), + ("5:7 Portrait", (5, 7)), ] - return { - "required": { - # A true dropdown of strings: - "preset": (choices,), - "batch_size": ("INT", {"default": 1, "min": 1}) + REFERENCE_CHOICES = ["Width", "Height"] + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "primary_dim": ("INT", {"default": 512, "min": 8}), + "reference": (cls.REFERENCE_CHOICES,), + "aspect_ratio": ([lbl for lbl, _ in cls.ASPECT_CHOICES],), + "batch_size": ("INT", {"default": 1, "min": 1}) + } } - } -empty_latent_preset.INPUT_TYPES = _preset_input_types -empty_latent_preset.RETURN_TYPES = ("LATENT",) -empty_latent_preset.CATEGORY = "ComfyUI Aspect Ratio Preset" + RETURN_TYPES = ("LATENT",) + RETURN_NAMES = ("empty_latent",) + FUNCTION = "generate" + CATEGORY = "ComfyUI Aspect Ratio Preset" + def generate(self, primary_dim: int, reference: str, aspect_ratio: str, batch_size: int): + # Lookup numeric aspect tuple + ratio_map = dict(self.ASPECT_CHOICES) + if aspect_ratio not in ratio_map: + raise ValueError(f"Unknown aspect ratio: {aspect_ratio}") + wr, hr = ratio_map[aspect_ratio] -# ─── Node 2: By Axis ────────────────────────────────────────────────── + _validate_dim(primary_dim) + if reference == "Width": + w, h = primary_dim, round(primary_dim * hr / wr) + else: + h, w = primary_dim, round(primary_dim * wr / hr) -ASPECT_CHOICES = [ - ("1:1 Square", (1, 1)), - ("3:2 Landscape", (3, 2)), - ("2:3 Portrait", (2, 3)), - ("4:3 Landscape", (4, 3)), - ("3:4 Portrait", (3, 4)), - ("16:9 Landscape", (16, 9)), - ("9:16 Portrait", (9, 16)), - ("5:4 Landscape", (5, 4)), - ("4:5 Portrait", (4, 5)), - ("21:9 Widescreen",(21, 9)), - ("9:21 Portrait", (9, 21)), - ("7:5 Landscape", (7, 5)), - ("5:7 Portrait", (5, 7)), -] + _validate_dim(w) + _validate_dim(h) -def empty_latent_by_axis( - primary_dim: int, - reference: str, - aspect_ratio: str, - batch_size: int -) -> np.ndarray: - """ - Empty Latent (Aspect Ratio by Axis) - - primary_dim: numeric value for Width or Height - reference: "Width" or "Height" - aspect_ratio: e.g. "3:2 Landscape" - batch_size: how many latents to produce - """ - # Look up numeric ratio - for lbl, (wr, hr) in ASPECT_CHOICES: - if lbl == aspect_ratio: - break - else: - raise ValueError(f"Unknown aspect ratio: {aspect_ratio}") - - _validate_dim(primary_dim) - if reference == "Width": - W = primary_dim - H = round(primary_dim * hr / wr) - else: - H = primary_dim - W = round(primary_dim * wr / hr) - - _validate_dim(W); _validate_dim(H) - return np.zeros((batch_size, 4, H // 8, W // 8), dtype=np.float32) - - -def _byaxis_input_types(): - return { - "required": { - "primary_dim": ("INT", {"default": 512, "min": 8}), - # A true dropdown of strings: - "reference": (["Width", "Height"],), - "aspect_ratio": ([lbl for lbl, _ in ASPECT_CHOICES],), - "batch_size": ("INT", {"default": 1, "min": 1}) - } - } - -empty_latent_by_axis.INPUT_TYPES = _byaxis_input_types -empty_latent_by_axis.RETURN_TYPES = ("LATENT",) -empty_latent_by_axis.CATEGORY = "ComfyUI Aspect Ratio Preset" - - -# ─── Expose for ComfyUI ─────────────────────────────────────────────── - -NODE_CLASS_MAPPINGS = { - "Empty Latent (Aspect Ratio Preset)": empty_latent_preset, - "Empty Latent (Aspect Ratio by Axis)": empty_latent_by_axis, -} -NODE_DISPLAY_NAME = "ComfyUI Aspect Ratio Preset" \ No newline at end of file + latent = torch.zeros([batch_size, 4, h // 8, w // 8], dtype=torch.float32) + return ({"samples": latent},) \ No newline at end of file