From b24e7cb2a2a4300440ee5ffdbece083baa182027 Mon Sep 17 00:00:00 2001 From: Vito Sansevero Date: Sat, 14 Jun 2025 07:59:34 -0700 Subject: [PATCH] feat(kikotools): add Resolution Calculator tool --- kikotools/__init__.py | 17 ++ kikotools/base/__init__.py | 7 + kikotools/base/base_node.py | 76 ++++++++ kikotools/tools/__init__.py | 6 + .../tools/resolution_calculator/__init__.py | 8 + .../tools/resolution_calculator/logic.py | 166 ++++++++++++++++ kikotools/tools/resolution_calculator/node.py | 179 ++++++++++++++++++ 7 files changed, 459 insertions(+) create mode 100644 kikotools/__init__.py create mode 100644 kikotools/base/__init__.py create mode 100644 kikotools/base/base_node.py create mode 100644 kikotools/tools/__init__.py create mode 100644 kikotools/tools/resolution_calculator/__init__.py create mode 100644 kikotools/tools/resolution_calculator/logic.py create mode 100644 kikotools/tools/resolution_calculator/node.py diff --git a/kikotools/__init__.py b/kikotools/__init__.py new file mode 100644 index 0000000..3982143 --- /dev/null +++ b/kikotools/__init__.py @@ -0,0 +1,17 @@ +""" +KikoTools package initialization and node registry +Handles automatic discovery and registration of all ComfyAssets tools +""" + +from .tools.resolution_calculator import ResolutionCalculatorNode + +# ComfyUI node registration mappings +NODE_CLASS_MAPPINGS = { + "ResolutionCalculator": ResolutionCalculatorNode, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "ResolutionCalculator": "Resolution Calculator", +} + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/kikotools/base/__init__.py b/kikotools/base/__init__.py new file mode 100644 index 0000000..ca904ff --- /dev/null +++ b/kikotools/base/__init__.py @@ -0,0 +1,7 @@ +""" +Base classes and utilities for all ComfyAssets tools +""" + +from .base_node import ComfyAssetsBaseNode + +__all__ = ["ComfyAssetsBaseNode"] diff --git a/kikotools/base/base_node.py b/kikotools/base/base_node.py new file mode 100644 index 0000000..cecba5c --- /dev/null +++ b/kikotools/base/base_node.py @@ -0,0 +1,76 @@ +""" +Base node class for all ComfyAssets tools +Provides consistent categorization and shared functionality +""" + +from typing import Dict, Any, Optional +import logging + +logger = logging.getLogger(__name__) + + +class ComfyAssetsBaseNode: + """ + Base class for all ComfyAssets nodes + + Provides: + - Consistent "ComfyAssets" categorization + - Standardized error handling and logging + - Common validation patterns + - Consistent return type handling + """ + + CATEGORY = "ComfyAssets" + + def validate_inputs(self, **kwargs) -> None: + """ + Common input validation logic + Override in subclasses for specific validation needs + + Args: + **kwargs: Input parameters to validate + + Raises: + ValueError: If validation fails + """ + pass + + def handle_error( + self, error_msg: str, exception: Optional[Exception] = None + ) -> None: + """ + Standardized error handling with logging + + Args: + error_msg: Human-readable error message + exception: Optional original exception for logging + """ + logger.error(f"{self.__class__.__name__}: {error_msg}") + if exception: + logger.exception(f"Original exception: {exception}") + raise ValueError(error_msg) + + def log_info(self, message: str) -> None: + """ + Standardized info logging + + Args: + message: Information message to log + """ + logger.info(f"{self.__class__.__name__}: {message}") + + @classmethod + def get_node_info(cls) -> Dict[str, Any]: + """ + Get standardized node information for debugging/introspection + + Returns: + Dict containing node metadata + """ + return { + "class_name": cls.__name__, + "category": getattr(cls, "CATEGORY", "Unknown"), + "function": getattr(cls, "FUNCTION", "Unknown"), + "return_types": getattr(cls, "RETURN_TYPES", ()), + "return_names": getattr(cls, "RETURN_NAMES", ()), + } diff --git a/kikotools/tools/__init__.py b/kikotools/tools/__init__.py new file mode 100644 index 0000000..f420954 --- /dev/null +++ b/kikotools/tools/__init__.py @@ -0,0 +1,6 @@ +""" +Individual tool implementations +Each tool is a self-contained module with node and logic components +""" + +# Import statements will be added as tools are implemented diff --git a/kikotools/tools/resolution_calculator/__init__.py b/kikotools/tools/resolution_calculator/__init__.py new file mode 100644 index 0000000..7085b1b --- /dev/null +++ b/kikotools/tools/resolution_calculator/__init__.py @@ -0,0 +1,8 @@ +""" +Resolution Calculator tool +Calculates upscaled dimensions from image or latent inputs with proper scaling factors +""" + +from .node import ResolutionCalculatorNode + +__all__ = ["ResolutionCalculatorNode"] diff --git a/kikotools/tools/resolution_calculator/logic.py b/kikotools/tools/resolution_calculator/logic.py new file mode 100644 index 0000000..d89e8db --- /dev/null +++ b/kikotools/tools/resolution_calculator/logic.py @@ -0,0 +1,166 @@ +""" +Core calculation logic for Resolution Calculator +Pure functions for dimension extraction and scaling calculations +""" + +import torch +from typing import Tuple, Optional, Union, Dict, Any + + +def extract_dimensions( + image: Optional[torch.Tensor] = None, + latent: Optional[Dict[str, torch.Tensor]] = None, +) -> Tuple[int, int]: + """ + Extract width and height from IMAGE or LATENT tensor + + Args: + image: Optional IMAGE tensor in ComfyUI format [batch, height, width, channels] + latent: Optional LATENT dict with 'samples' tensor [batch, channels, height/8, width/8] + + Returns: + Tuple of (width, height) as integers + + Raises: + ValueError: If neither image nor latent is provided + """ + if image is not None: + # IMAGE tensor format: [batch, height, width, channels] + if len(image.shape) != 4: + raise ValueError( + f"Expected IMAGE tensor with 4 dimensions, got {len(image.shape)}" + ) + + _, height, width, _ = image.shape + return int(width), int(height) + + elif latent is not None: + # LATENT format: {"samples": [batch, channels, height/8, width/8]} + if "samples" not in latent: + raise ValueError("LATENT dict must contain 'samples' key") + + samples = latent["samples"] + if len(samples.shape) != 4: + raise ValueError( + f"Expected LATENT samples tensor with 4 dimensions, got {len(samples.shape)}" + ) + + _, _, latent_height, latent_width = samples.shape + + # Latent dimensions are 1/8 of actual image dimensions + width = int(latent_width * 8) + height = int(latent_height * 8) + return width, height + + else: + raise ValueError("Either image or latent must be provided") + + +def ensure_divisible_by_8(width: int, height: int) -> Tuple[int, int]: + """ + Ensure dimensions are divisible by 8 (ComfyUI requirement) + Rounds to the nearest multiple of 8 + + Args: + width: Input width + height: Input height + + Returns: + Tuple of (width, height) both divisible by 8 + """ + # Round to nearest multiple of 8 + # Formula: ((value + 4) // 8) * 8 + # This rounds 0-3 down, 4-7 up, ensuring nearest multiple + new_width = ((width + 4) // 8) * 8 + new_height = ((height + 4) // 8) * 8 + + return int(new_width), int(new_height) + + +def calculate_scaled_dimensions( + width: int, height: int, scale_factor: float +) -> Tuple[int, int]: + """ + Calculate new dimensions with scale factor and ensure divisible by 8 + + Args: + width: Original width + height: Original height + scale_factor: Scaling factor (e.g., 1.5, 2.0, 3.0) + + Returns: + Tuple of (new_width, new_height) both divisible by 8 + """ + if scale_factor <= 0: + raise ValueError(f"Scale factor must be positive, got {scale_factor}") + + # Calculate new dimensions + new_width = int(width * scale_factor) + new_height = int(height * scale_factor) + + # Ensure divisible by 8 + return ensure_divisible_by_8(new_width, new_height) + + +def validate_scale_factor( + scale_factor: float, min_scale: float = 0.1, max_scale: float = 8.0 +) -> None: + """ + Validate scale factor is within reasonable bounds + + Args: + scale_factor: Scale factor to validate + min_scale: Minimum allowed scale factor + max_scale: Maximum allowed scale factor + + Raises: + ValueError: If scale factor is out of bounds + """ + if not isinstance(scale_factor, (int, float)): + raise ValueError( + f"Scale factor must be a number, got {type(scale_factor).__name__}" + ) + + if scale_factor < min_scale: + raise ValueError( + f"Scale factor {scale_factor} is too small (minimum: {min_scale})" + ) + + if scale_factor > max_scale: + raise ValueError( + f"Scale factor {scale_factor} is too large (maximum: {max_scale})" + ) + + +def calculate_resolution_from_input( + scale_factor: float, + image: Optional[torch.Tensor] = None, + latent: Optional[Dict[str, torch.Tensor]] = None, +) -> Tuple[int, int]: + """ + Main function to calculate resolution from input tensor and scale factor + Combines all the logic steps into a single function + + Args: + scale_factor: Scaling factor + image: Optional IMAGE tensor + latent: Optional LATENT dict + + Returns: + Tuple of (width, height) scaled and divisible by 8 + + Raises: + ValueError: For various validation errors + """ + # Validate scale factor + validate_scale_factor(scale_factor) + + # Extract original dimensions + original_width, original_height = extract_dimensions(image=image, latent=latent) + + # Calculate scaled dimensions + new_width, new_height = calculate_scaled_dimensions( + original_width, original_height, scale_factor + ) + + return new_width, new_height diff --git a/kikotools/tools/resolution_calculator/node.py b/kikotools/tools/resolution_calculator/node.py new file mode 100644 index 0000000..7adda87 --- /dev/null +++ b/kikotools/tools/resolution_calculator/node.py @@ -0,0 +1,179 @@ +""" +Resolution Calculator ComfyUI Node +Provides ComfyUI interface for calculating upscaled dimensions +""" + +import torch +from typing import Dict, Any, Tuple, Optional + +from ...base import ComfyAssetsBaseNode +from .logic import calculate_resolution_from_input + + +class ResolutionCalculatorNode(ComfyAssetsBaseNode): + """ + ComfyUI node for calculating upscaled resolution from image or latent inputs + + Inputs: + - scale_factor (FLOAT): Scaling factor (1.0 to 8.0) + - image (IMAGE, optional): Input image tensor + - latent (LATENT, optional): Input latent tensor + + Outputs: + - width (INT): Calculated width + - height (INT): Calculated height + """ + + @classmethod + def INPUT_TYPES(cls) -> Dict[str, Any]: + """ + Define ComfyUI input interface + + Returns: + Dict with required and optional input specifications + """ + return { + "required": { + "scale_factor": ( + "FLOAT", + { + "default": 2.0, + "min": 1.0, + "max": 8.0, + "step": 0.1, + "display": "slider", + "tooltip": "Factor to scale the resolution by (e.g., 2.0 for 2x upscale)", + }, + ), + }, + "optional": { + "image": ( + "IMAGE", + {"tooltip": "Input image to calculate dimensions from"}, + ), + "latent": ( + "LATENT", + {"tooltip": "Input latent to calculate dimensions from"}, + ), + }, + } + + RETURN_TYPES = ("INT", "INT") + RETURN_NAMES = ("width", "height") + FUNCTION = "calculate_resolution" + + def calculate_resolution( + self, + scale_factor: float, + image: Optional[torch.Tensor] = None, + latent: Optional[Dict[str, torch.Tensor]] = None, + ) -> Tuple[int, int]: + """ + Calculate upscaled resolution from input tensor and scale factor + + Args: + scale_factor: Scaling factor to apply + image: Optional IMAGE tensor [batch, height, width, channels] + latent: Optional LATENT dict with 'samples' tensor + + Returns: + Tuple of (width, height) as integers, both divisible by 8 + + Raises: + ValueError: If validation fails or no input provided + """ + try: + # Validate inputs using base class + self.validate_inputs(scale_factor=scale_factor, image=image, latent=latent) + + # Log the operation + input_type = ( + "IMAGE" + if image is not None + else "LATENT" if latent is not None else "NONE" + ) + self.log_info( + f"Calculating resolution with scale_factor={scale_factor}, input_type={input_type}" + ) + + # Calculate the resolution + width, height = calculate_resolution_from_input( + scale_factor=scale_factor, image=image, latent=latent + ) + + # Log the result + self.log_info(f"Calculated resolution: {width}x{height}") + + return width, height + + except Exception as e: + # Handle and re-raise with context + error_msg = f"Failed to calculate resolution: {str(e)}" + self.handle_error(error_msg, e) + + def validate_inputs( + self, + scale_factor: float, + image: Optional[torch.Tensor] = None, + latent: Optional[Dict[str, torch.Tensor]] = None, + ) -> None: + """ + Validate inputs specific to resolution calculator + + Args: + scale_factor: Scale factor to validate + image: Optional image tensor + latent: Optional latent dict + + Raises: + ValueError: If validation fails + """ + # Check that at least one input is provided + if image is None and latent is None: + raise ValueError("Either 'image' or 'latent' input must be provided") + + # Validate scale factor type + if not isinstance(scale_factor, (int, float)): + raise ValueError( + f"scale_factor must be a number, got {type(scale_factor).__name__}" + ) + + # Additional tensor validation + if image is not None: + if not isinstance(image, torch.Tensor): + raise ValueError( + f"image must be a torch.Tensor, got {type(image).__name__}" + ) + + if len(image.shape) != 4: + raise ValueError( + f"image tensor must have 4 dimensions [batch, height, width, channels], got {len(image.shape)}" + ) + + if latent is not None: + if not isinstance(latent, dict): + raise ValueError(f"latent must be a dict, got {type(latent).__name__}") + + if "samples" not in latent: + raise ValueError("latent dict must contain 'samples' key") + + samples = latent["samples"] + if not isinstance(samples, torch.Tensor): + raise ValueError( + f"latent['samples'] must be a torch.Tensor, got {type(samples).__name__}" + ) + + if len(samples.shape) != 4: + raise ValueError( + f"latent samples tensor must have 4 dimensions [batch, channels, height, width], got {len(samples.shape)}" + ) + + +# Node class mappings for ComfyUI registration +NODE_CLASS_MAPPINGS = { + "ResolutionCalculator": ResolutionCalculatorNode, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "ResolutionCalculator": "Resolution Calculator", +}