From f3bd016ba4c8f2a71859eba2d6b045c7f5fcf001 Mon Sep 17 00:00:00 2001 From: malkuthro <166633396+malkuthro@users.noreply.github.com> Date: Sun, 28 Sep 2025 22:36:49 +0200 Subject: [PATCH] system-A --- __init__.py | 145 ++++++++++++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 140 insertions(+), 5 deletions(-) diff --git a/__init__.py b/__init__.py index 07967b5..72c965e 100644 --- a/__init__.py +++ b/__init__.py @@ -1,5 +1,8 @@ import json import os +import torch +from comfy.utils import common_upscale +from nodes import MAX_RESOLUTION class Wan22EasyPrompt: """ @@ -23,7 +26,7 @@ class Wan22EasyPrompt: except (FileNotFoundError, json.JSONDecodeError) as e: print(f"Error loading config.json: {e}. Using dummy config.") config = { - "category": "Wan/Parameters", + "category": "Koolook/Wan_Video", "node_name": "Wan_2.2", "fields": [], "output_type": "STRING" @@ -65,7 +68,7 @@ class Wan22EasyPrompt: RETURN_TYPES = ("STRING", "STRING", "CONDITIONING",) RETURN_NAMES = ("combined_prompt", "fields_only", "conditioning",) FUNCTION = "execute" - CATEGORY = "Wan/Parameters" + CATEGORY = "Koolook/Wan_Video" OUTPUT_NODE = False def check_lazy_status(self, **kwargs): @@ -107,11 +110,143 @@ class Wan22EasyPrompt: except FileNotFoundError: return "dummy" -# Export the node +class ImageAspect: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "image": ("IMAGE",), + "base_on": (["Width", "Height"], {"default": "Width"}), + "base_size": ("INT", {"default": 512, "min": 1, "max": MAX_RESOLUTION, "step": 1}), + "aspect_ratio": ("STRING", {"default": "16:9"}), + "upscale_method": (["nearest-exact", "bilinear", "area", "bicubic", "lanczos"], {"default": "nearest-exact"}), + "keep_proportion": (["stretch", "letterbox", "pillarbox"], {"default": "stretch"}), + "pad_red": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}), + "pad_green": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}), + "pad_blue": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}), + "crop_position": (["top", "bottom", "left", "right", "center"], {"default": "center"}), + "divisible_by": ("INT", {"default": 32, "min": 1, "max": 128, "step": 1}), + "device": (["cpu", "cuda"], {"default": "cpu"}), + }, + "optional": { + "mask": ("MASK",), + } + } + + RETURN_TYPES = ("IMAGE", "MASK", "INT", "INT") + RETURN_NAMES = ("IMAGE", "MASK", "width", "height") + FUNCTION = "adjust_to_aspect" + CATEGORY = "Koolook/Image" + DESCRIPTION = """ +Adjusts the input image to a target resolution based on a base dimension (width or height), +an aspect ratio (e.g., '16:9' for landscape, '9:16' for portrait), ensuring both dimensions +are divisible by the specified value. Supports stretch or pad/crop modes for proportion control. +Outputs the resized image, optional mask, and final width/height for VFX pipeline integration. +""" + + def adjust_to_aspect(self, image, base_on, base_size, aspect_ratio, upscale_method, keep_proportion, pad_red, pad_green, pad_blue, crop_position, divisible_by, device, mask=None): + def round_to_nearest_multiple(number, multiple): + return multiple * round(number / multiple) + + # Parse aspect ratio (order matters: '16:9' implies width > height if base allows) + try: + ar_parts = [int(part.strip()) for part in aspect_ratio.split(':')] + if len(ar_parts) != 2 or ar_parts[0] <= 0 or ar_parts[1] <= 0: + raise ValueError("Aspect ratio must be 'w:h' with positive integers, e.g., '16:9'.") + ar_width, ar_height = ar_parts + except ValueError as e: + raise ValueError(str(e)) + + ratio = ar_width / ar_height # w/h ratio + + # Make base_size divisible by 'divisible_by' + base_rounded = max(divisible_by, round_to_nearest_multiple(base_size, divisible_by)) + + # Calculate target dimensions + if base_on == "Width": + target_width = base_rounded + computed_height = target_width / ratio + target_height = max(divisible_by, round_to_nearest_multiple(computed_height, divisible_by)) + else: # Height + target_height = base_rounded + computed_width = target_height * ratio + target_width = max(divisible_by, round_to_nearest_multiple(computed_width, divisible_by)) + + # Move to specified device + image = image.to(device) + if mask is not None: + mask = mask.to(device) + + # Image processing (inspired by ResizeImage logic in image_nodes.py) + B, H, W, C = image.shape + image = image.movedim(-1, 1) # Channels-first + + if mask is not None: + mask = mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])).movedim(1, -1).movedim(-1, 1) # Adjust for interpolation + + if keep_proportion == "stretch": + # Simple resize + out_image = common_upscale(image, target_width, target_height, upscale_method, "disabled") + if mask is not None: + out_mask = torch.nn.functional.interpolate(mask, size=(target_height, target_width), mode="nearest") + else: + out_mask = None + else: + # Preserve aspect with pad/crop (letterbox or pillarbox) + scale = min(target_width / W, target_height / H) + new_w = round_to_nearest_multiple(W * scale, divisible_by) + new_h = round_to_nearest_multiple(H * scale, divisible_by) + + resized_image = common_upscale(image, new_w, new_h, upscale_method, "disabled") + if mask is not None: + resized_mask = torch.nn.functional.interpolate(mask, size=(new_h, new_w), mode="nearest") + + # Determine padding based on mode and position + pad_l = pad_r = pad_t = pad_b = 0 + if keep_proportion == "letterbox": # Add bars top/bottom (for wider target) + total_pad_h = target_height - new_h + if crop_position == "top": + pad_b = total_pad_h + elif crop_position == "bottom": + pad_t = total_pad_h + else: # center, left, right (left/right not affecting vertical) + pad_t = total_pad_h // 2 + pad_b = total_pad_h - pad_t + elif keep_proportion == "pillarbox": # Add bars left/right (for taller target) + total_pad_w = target_width - new_w + if crop_position == "left": + pad_r = total_pad_w + elif crop_position == "right": + pad_l = total_pad_w + else: # center, top, bottom + pad_l = total_pad_w // 2 + pad_r = total_pad_w - pad_l + + # Create background with pad color (assume C=3 for RGB) + pad_color = [pad_red, pad_green, pad_blue] + pad_color_tensor = torch.tensor(pad_color, dtype=image.dtype, device=device).view(1, 3, 1, 1) + out_image = pad_color_tensor.expand(B, 3, target_height, target_width) + out_image[:, :, pad_t:pad_t + new_h, pad_l:pad_l + new_w] = resized_image + + if mask is not None: + out_mask = torch.zeros((B, 1, target_height, target_width), dtype=mask.dtype, device=device) + out_mask[:, :, pad_t:pad_t + new_h, pad_l:pad_l + new_w] = resized_mask + else: + out_mask = None + + out_image = out_image.movedim(1, -1) + if out_mask is not None: + out_mask = out_mask.movedim(1, -1).squeeze(1) + + return (out_image, out_mask, target_width, target_height) + +# Export the nodes NODE_CLASS_MAPPINGS = { - "Wan22EasyPrompt": Wan22EasyPrompt + "Wan22EasyPrompt": Wan22EasyPrompt, + "ImageAspect": ImageAspect } NODE_DISPLAY_NAME_MAPPINGS = { - "Wan22EasyPrompt": "Wan 2.2 Easy Prompt" + "Wan22EasyPrompt": "Wan 2.2 Easy Prompt", + "ImageAspect": "Image Aspect" } \ No newline at end of file