diff --git a/__init__.py b/__init__.py index ccc653f..fc42c54 100644 --- a/__init__.py +++ b/__init__.py @@ -1 +1,12 @@ -# comfyui-aspect-ratio-presets custom node package \ No newline at end of file +# comfyui-aspect-ratio-presets/__init__.py + +from .nodes import empty_latent_preset, empty_latent_by_axis + +NODE_CLASS_MAPPINGS = { + # Display name : Python callable + "Empty Latent (Aspect Ratio Preset)": empty_latent_preset, + "Empty Latent (Aspect Ratio by Axis)": empty_latent_by_axis, +} + +# This is the category under which they appear in the ComfyUI UI +NODE_DISPLAY_NAME = "ComfyUI Aspect Ratio Preset" \ No newline at end of file diff --git a/aspect_ratio_nodes.py b/aspect_ratio_nodes.py deleted file mode 100644 index 409c1d0..0000000 --- a/aspect_ratio_nodes.py +++ /dev/null @@ -1,104 +0,0 @@ -# aspect_ratio_nodes.py - -import numpy as np -import modules.scripts as scripts -from modules import shared -from modules.ui import create_node, create_node_callbacks -from presets import PRESETS - -def _validate_dim(v): - if v <= 0: - raise ValueError("Dimension must be > 0") - if v % 8 != 0: - raise ValueError("Dimension must be divisible by 8") - -class EmptyLatentAspectPreset(scripts.Script): - title = "Empty Latent (Aspect Ratio Preset)" - ui_category = "ComfyUI Aspect Ratio Preset" - - def ui(self): - self.resolution = self.ui.dropdown( - label="Resolution Preset", - choices=[f"{m} {w}×{h} ({lbl})" for m, lbl, w, h in PRESETS], - default=0, - tooltip="Choose a predefined resolution + aspect ratio" - ) - self.batch_size = self.ui.int_field( - label="Batch Size", - default=1, - minimum=1, - tooltip="How many blank latents to produce" - ) - - def run(self, p): - model, lbl, W, H = PRESETS[self.resolution] - _validate_dim(W); _validate_dim(H) - B = self.batch_size - C = shared.sd_model.cond_stage_model.latent_channels # usually 4 - shape = (B, C, H // 8, W // 8) - return np.zeros(shape, dtype=np.float32) - - -class EmptyLatentAspectByAxis(scripts.Script): - title = "Empty Latent (Aspect Ratio by Axis)" - ui_category = "ComfyUI Aspect Ratio Preset" - - 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)), - ] - - def ui(self): - self.primary = self.ui.int_field( - label="Primary Dimension", - default=512, - tooltip="Your value for the chosen Reference Axis" - ) - self.axis = self.ui.dropdown( - label="Reference Axis", - choices=["Width","Height"], - default="Width", - tooltip="Which side does the Primary Dimension refer to?" - ) - self.aspect = self.ui.dropdown( - label="Aspect Ratio", - choices=[lbl for lbl,_ in self.ASPECT_CHOICES], - default=0, - tooltip="Select X:Y ratio and orientation" - ) - self.batch_size = self.ui.int_field( - label="Batch Size", - default=1, - minimum=1, - tooltip="How many blank latents to produce" - ) - - def run(self, p): - val = self.primary - _validate_dim(val) - w_ratio, h_ratio = self.ASPECT_CHOICES[self.aspect][1] - - if self.axis == "Width": - W = val - H = round(val * (h_ratio / w_ratio)) - else: - H = val - W = round(val * (w_ratio / h_ratio)) - - _validate_dim(W); _validate_dim(H) - - B = self.batch_size - C = shared.sd_model.cond_stage_model.latent_channels - shape = (B, C, H // 8, W // 8) - return np.zeros(shape, dtype=np.float32) \ No newline at end of file diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..e43afe8 --- /dev/null +++ b/nodes.py @@ -0,0 +1,73 @@ +# nodes.py + +import numpy as np +from .presets import PRESETS + +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 ----------------------------------------------------------- + +def empty_latent_preset(resolution_index:int, batch_size:int) -> np.ndarray: + model, label, W, H = PRESETS[resolution_index] + _validate_dim(W); _validate_dim(H) + return np.zeros((batch_size, 4, H//8, W//8), dtype=np.float32) + +def _preset_input_types(): + return { + "required": { + "resolution_index": ( + "INT", { + "default": 0, + "choices": [ + f"{m} {w}×{h} ({lbl})" + for m,lbl,w,h in PRESETS + ] + } + ), + "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" + +# --- Node 2 ----------------------------------------------------------- + +ASPECT_CHOICES = [ + ("1:1 Square",(1,1)), ("3:2 Landscape",(3,2)), ("2:3 Portrait",(2,3)), + # ... etc ... +] + +def empty_latent_by_axis(primary_dim:int, reference:str, aspect_index:int, batch_size:int) -> np.ndarray: + w_ratio,h_ratio = ASPECT_CHOICES[aspect_index][1] + _validate_dim(primary_dim) + if reference == "Width": + W,H = primary_dim, round(primary_dim * h_ratio / w_ratio) + else: + H,W = primary_dim, round(primary_dim * w_ratio / h_ratio) + _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}), + "reference": ("STRING", {"default":"Width", "choices":["Width","Height"]}), + "aspect_index": ("INT", { + "default":0, + "choices":[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 ────────────────────────────────────────────────