feat(kikotools): add Resolution Calculator tool
This commit is contained in:
Vito Sansevero
@@ -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"]
|
||||
@@ -0,0 +1,7 @@
|
||||
"""
|
||||
Base classes and utilities for all ComfyAssets tools
|
||||
"""
|
||||
|
||||
from .base_node import ComfyAssetsBaseNode
|
||||
|
||||
__all__ = ["ComfyAssetsBaseNode"]
|
||||
@@ -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", ()),
|
||||
}
|
||||
@@ -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
|
||||
@@ -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"]
|
||||
@@ -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
|
||||
@@ -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",
|
||||
}
|
||||
Reference in New Issue
Block a user