Compare commits

...
Author SHA1 Message Date
Vito Sansevero c1128addc7 chore: bump version to 1.0.7 in pyproject.toml 2025-07-26 14:16:06 -07:00
Vito Sansevero a49071f824 fix(resolution_calculator): update scale factor tooltip 2025-07-26 14:15:39 -07:00
Vito Sansevero bbdd27f498 fix(ci): correct return type in tests.yml configuration 2025-07-23 13:56:57 -07:00
Vito Sansevero 22f62bf7b4 test: Update test assertions for sampler combo node 2025-07-23 13:56:45 -07:00
Vito Sansevero 4ff6067dad refactor(compact_node): update sampler return type 2025-07-23 13:29:46 -07:00
Vito Sansevero ad13e66506 refactor(node): update sampler handling logic 2025-07-23 13:29:32 -07:00
Vito Sansevero b16f6f40bd style(logic): fix whitespace issues in logic.py 2025-07-21 08:23:18 -07:00
Vito Sansevero 6dfa66963b feat: Add subfolder support in image URL handling 2025-07-21 08:20:27 -07:00
Vito Sansevero ab016e0903 feat(logic): add subfolder info to enhanced data 2025-07-21 08:20:17 -07:00
Vito Sansevero f4228a850c refactor(logic): improve path handling in image saving 2025-07-21 07:56:13 -07:00
Vito Sansevero 79042b78d2 chore: bump version to 1.0.5 in pyproject.toml 2025-06-28 09:31:40 -07:00
Vito Sansevero 0c4e59c4e9 test: Remove unused imports from test file 2025-06-28 09:14:45 -07:00
Vito Sansevero 4a0a206d61 refactor(node): use helper methods for tensor validation 2025-06-28 09:14:34 -07:00
Vito Sansevero 8e0d4485bd style: Remove unused imports in node.py 2025-06-28 09:14:23 -07:00
Vito Sansevero 92a3b1db4e style: Remove unused import 'Any' 2025-06-28 09:13:03 -07:00
Vito Sansevero ab628b1bf2 style: Remove unused import 'os' 2025-06-28 09:11:49 -07:00
Vito Sansevero 7e712a17d9 docs: Add Kiko Save Image section to README.md 2025-06-28 09:11:42 -07:00
Vito Sansevero 269fb2ba80 ci: add checks for KikoSaveImageNode imports 2025-06-28 09:11:36 -07:00
Vito Sansevero e17fdddcd7 refactor(tests/ui): Remove 'subfolder' support 2025-06-28 08:55:25 -07:00
Vito Sansevero 5d1f01e6cb refactor(node): replace 'subfolder' with 'popup' 2025-06-28 08:54:46 -07:00
Vito Sansevero b3b8826044 refactor(logic): Rename 'subfolder' to 'popup' parameter 2025-06-28 08:54:35 -07:00
Vito Sansevero 9dbb1f749d chore: bump version to 1.0.4 in pyproject.toml 2025-06-27 21:06:30 -07:00
Vito Sansevero 682f2b0a47 feat(ui): Add KikoSaveImage UI enhancements 2025-06-27 21:06:00 -07:00
Vito Sansevero 1233cf693e test(kiko_save_image): add unit tests for save image tool 2025-06-27 21:05:52 -07:00
Vito Sansevero 32d44a282f feat(kiko_save_image): add new image saving tool 2025-06-27 21:05:42 -07:00
Vito Sansevero bb79c7434f feat(init): add KikoSaveImageNode to tools and mappings 2025-06-27 21:05:25 -07:00
Vito Sansevero bd8c0a42bc fix: handle import error for testing environment 2025-06-27 21:05:14 -07:00
Vito Sansevero 5d9e71dc7b chore: bump version to 1.0.3 in pyproject.toml 2025-06-20 08:03:48 -07:00
Vito 321d89dcc4 Merge pull request #9 from ComfyAssets/latent-batch
style: Add blank lines for better readability
2025-06-20 08:02:57 -07:00
Vito d23ff34b27 Merge pull request #8 from ComfyAssets/latent-batch
Latent batch
2025-06-19 12:00:57 -07:00
17 changed files with 2605 additions and 59 deletions
+25
View File
@@ -87,6 +87,10 @@ jobs:
from kikotools.tools.seed_history import SeedHistoryNode
from kikotools.tools.seed_history.logic import generate_random_seed, validate_seed_value
# Test Kiko Save Image imports
from kikotools.tools.kiko_save_image import KikoSaveImageNode
from kikotools.tools.kiko_save_image.logic import process_image_batch, validate_save_inputs
print('✓ All module imports successful')
"
@@ -252,6 +256,27 @@ jobs:
print(f'❌ SeedHistoryNode missing required attribute: {attr}')
sys.exit(1)
# Test Kiko Save Image Node
from kikotools.tools.kiko_save_image.node import KikoSaveImageNode
if issubclass(KikoSaveImageNode, ComfyAssetsBaseNode):
print('✓ KikoSaveImageNode properly inherits from base class')
else:
print('❌ KikoSaveImageNode does not inherit from base class')
sys.exit(1)
# KikoSaveImage is an output node, so it doesn't have RETURN_TYPES/RETURN_NAMES
save_required_attrs = ['INPUT_TYPES', 'FUNCTION', 'CATEGORY']
for attr in save_required_attrs:
if not hasattr(KikoSaveImageNode, attr):
print(f'❌ KikoSaveImageNode missing required attribute: {attr}')
sys.exit(1)
# Check that it's properly marked as an output node
if not hasattr(KikoSaveImageNode, 'OUTPUT_NODE') or not KikoSaveImageNode.OUTPUT_NODE:
print('❌ KikoSaveImageNode missing OUTPUT_NODE = True')
sys.exit(1)
print('✓ All architecture checks passed for all tools')
"
+1 -1
View File
@@ -159,7 +159,7 @@ jobs:
print('✓ Sampler Combo interface tests passed')
# Test return types
assert node.RETURN_TYPES == (SAMPLERS, SCHEDULERS, 'INT', 'FLOAT')
assert node.RETURN_TYPES == ('SAMPLER', SCHEDULERS, 'INT', 'FLOAT')
assert node.RETURN_NAMES == ('sampler_name', 'scheduler', 'steps', 'cfg')
assert node.CATEGORY == 'ComfyAssets'
print('✓ Sampler Combo return types tests passed')
+85 -6
View File
@@ -92,6 +92,34 @@ Advanced empty latent creation with preset support and batch processing capabili
- Optimize memory usage with batch size planning
- Quick preset-based latent generation for different aspect ratios
#### 💾 Kiko Save Image
Enhanced image saving with format selection, quality control, and floating popup viewer.
- **Multiple Format Support**: Save as PNG, JPEG, or WebP with format-specific optimizations
- **Advanced Quality Controls**: JPEG/WebP quality (1-100), PNG compression (0-9), WebP lossless mode
- **Floating Popup Viewer**: Draggable, resizable window that shows saved images immediately
- **Interactive Previews**: Click any image to open in new tab, download individual images
- **Batch Selection**: Multi-select images for bulk actions (open all, download all)
- **Format-Specific Settings**: Quality indicators, file size display, compression info
- **Smart UI**: Auto-hide/show, minimize/maximize, roll-up functionality
- **Popup Toggle**: Enable/disable popup viewer per save operation
**Use Cases:**
- Quick preview and management of saved images without file browser navigation
- Compare multiple format outputs side-by-side (PNG vs JPEG vs WebP)
- Batch download or open selected images efficiently
- Monitor file sizes and compression effectiveness in real-time
- Streamlined workflow for iterative image generation and saving
**Why Better Than Standard Save Image:**
- **Immediate Visual Feedback**: See your saved images instantly without opening file explorer
- **Multi-Format Flexibility**: Choose optimal format for your use case (PNG for quality, JPEG for size, WebP for modern efficiency)
- **Advanced Compression Control**: Fine-tune file sizes with format-specific quality settings
- **Batch Operations**: Handle multiple images efficiently with selection and bulk actions
- **Modern UI**: Floating, draggable interface that doesn't interrupt your workflow
- **Smart Memory Usage**: File size indicators help optimize storage and sharing
- **One-Click Access**: Direct image opening in browser tabs for quick sharing or review
### 🔧 Architecture Highlights
- **Modular Design**: Each tool is self-contained and independently testable
@@ -171,8 +199,8 @@ Sampler Combo → KSampler → VAE Decode → Save Image
### Empty Latent Batch Example
```
Empty Latent Batch → KSampler → VAE Decode → Save Image
📦 preset: "1024×1024" ↘ batch latents ↗
Empty Latent Batch → KSampler → VAE Decode → Kiko Save Image
📦 preset: "1024×1024" ↘ batch latents ↗ ↘ popup viewer ↗
batch_size: 4
[swap button]
```
@@ -182,6 +210,21 @@ Empty Latent Batch → KSampler → VAE Decode → Save Image
**Output:** 4×4×128×128 latent tensor ready for sampling
**Swap Button:** Click to switch to any available swapped preset
### Kiko Save Image Example
```
Generate Image → Kiko Save Image → Floating Popup Viewer
📷 output ↘ format: WEBP ↘ draggable window ↗
quality: 85
[popup: enabled]
```
**Format:** WebP (efficient compression, modern format)
**Quality:** 85% (balanced size/quality)
**Popup Viewer:** Floating, draggable window with saved images
**Features:** Click images to open in new tabs, download individual files, batch selection
**Advantages:** Immediate preview without file explorer, multi-format comparison, advanced quality controls
### Common Workflows
<details>
@@ -223,6 +266,7 @@ Empty Latent Batch → KSampler → VAE Decode → Save Image
| **Seed History** | Advanced seed tracking with interactive history management | ✅ Complete | [Docs](examples/documentation/seed_history.md) |
| **Sampler Combo** | Unified sampling configuration with smart recommendations | ✅ Complete | [Docs](examples/documentation/sampler_combo.md) |
| **Empty Latent Batch** | Create empty latent batches with preset support | ✅ Complete | [Docs](examples/documentation/empty_latent_batch.md) |
| **Kiko Save Image** | Enhanced image saving with popup viewer and multi-format support | ✅ Complete | [Docs](examples/documentation/kiko_save_image.md) |
| **Batch Image Processor** | Process multiple images with consistent settings | 🚧 Planned | Coming Soon |
| **Advanced Prompt Utilities** | Enhanced prompt manipulation and generation | 🚧 Planned | Coming Soon |
@@ -340,6 +384,40 @@ Empty Latent Batch → KSampler → VAE Decode → Save Image
- Formatted display with aspect ratio and megapixel information
- Intelligent fallback to custom dimensions for invalid presets
#### Kiko Save Image
**Inputs:**
- `images` (IMAGE): Batch of images to save
- `filename_prefix` (STRING): Prefix for saved filenames, default "KikoSave"
- `format` (DROPDOWN): Output format (PNG, JPEG, WEBP), default PNG
- `quality` (INT): JPEG/WebP quality (1-100), default 90
- `png_compress_level` (INT): PNG compression level (0-9), default 4
- `webp_lossless` (BOOLEAN): Use lossless WebP compression, default False
- `popup` (BOOLEAN): Enable popup viewer window, default True
**Outputs:**
- `UI`: Enhanced image preview data with popup viewer functionality
**UI Features:**
- Floating, draggable popup window showing saved images immediately
- Interactive image grid with click-to-open functionality
- Individual image download buttons with format-specific quality indicators
- Batch selection with multi-select checkboxes for bulk operations
- Window controls: minimize, maximize, roll-up, close, and dragging
- Auto-hide/show behavior with smart positioning
**Format Support:**
- **PNG**: Lossless compression with metadata preservation, configurable compression levels
- **JPEG**: Quality-controlled lossy compression with automatic transparency handling
- **WebP**: Modern format with both lossy and lossless modes, superior compression ratios
**Advanced Features:**
- File size monitoring and display for optimization feedback
- Format-specific quality indicators (PNG compression level, JPEG/WebP quality percentage)
- Smart filename sanitization with timestamp-based uniqueness
- Persistent popup viewer across multiple save operations
- Toggle button integration in node UI for manual viewer control
## 🛠️ Development
### Prerequisites
@@ -460,13 +538,14 @@ MIT License - see [LICENSE](LICENSE) file for details.
## 📈 Stats
- **Nodes**: 5 (Resolution Calculator, Width Height Selector, Seed History, Sampler Combo, Empty Latent Batch)
- **Nodes**: 6 (Resolution Calculator, Width Height Selector, Seed History, Sampler Combo, Empty Latent Batch, Kiko Save Image)
- **Format Support**: 3 (PNG, JPEG, WebP with advanced controls)
- **Presets**: 26 curated resolution presets
- **Interactive Features**: 3 (Width/Height Swap Button, Seed History UI, Empty Latent Batch Swap Button)
- **Test Coverage**: 100% (180+ comprehensive tests)
- **Interactive Features**: 4 (Width/Height Swap Button, Seed History UI, Empty Latent Batch Swap Button, Kiko Save Image Popup Viewer)
- **Test Coverage**: 100% (200+ comprehensive tests)
- **Python Version**: 3.8+
- **ComfyUI Compatibility**: Latest
- **Dependencies**: Minimal (PyTorch, NumPy)
- **Dependencies**: Minimal (PyTorch, NumPy, Pillow)
---
+5 -2
View File
@@ -3,11 +3,14 @@ ComfyUI-KikoTools: Modular collection of custom ComfyUI nodes
All nodes are grouped under the "ComfyAssets" category
"""
import os
import re
from pathlib import Path
from .kikotools import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
try:
from .kikotools import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
except ImportError:
# Fallback for testing environment
from kikotools import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
# Tell ComfyUI where to find our JavaScript extensions
WEB_DIRECTORY = "./web"
+3
View File
@@ -8,6 +8,7 @@ from .tools.width_height_selector import WidthHeightSelectorNode
from .tools.seed_history import SeedHistoryNode
from .tools.sampler_combo import SamplerComboNode, SamplerComboCompactNode
from .tools.empty_latent_batch import EmptyLatentBatchNode
from .tools.kiko_save_image import KikoSaveImageNode
# ComfyUI node registration mappings
NODE_CLASS_MAPPINGS = {
@@ -17,6 +18,7 @@ NODE_CLASS_MAPPINGS = {
"SamplerCombo": SamplerComboNode,
"SamplerComboCompact": SamplerComboCompactNode,
"EmptyLatentBatch": EmptyLatentBatchNode,
"KikoSaveImage": KikoSaveImageNode,
}
NODE_DISPLAY_NAME_MAPPINGS = {
@@ -26,6 +28,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
"SamplerCombo": "Sampler Combo",
"SamplerComboCompact": "Sampler Combo (Compact)",
"EmptyLatentBatch": "Empty Latent Batch",
"KikoSaveImage": "Kiko Save Image",
}
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
+1 -1
View File
@@ -1,7 +1,7 @@
"""Logic for creating empty latent tensors with batch support."""
import torch
from typing import Dict, Tuple, Any
from typing import Dict, Tuple
def create_empty_latent_batch(
+1 -4
View File
@@ -1,7 +1,7 @@
"""Empty Latent Batch node for ComfyUI."""
import torch
from typing import Dict, Any, Tuple
from typing import Dict, Tuple
from ...base.base_node import ComfyAssetsBaseNode
from .logic import (
@@ -13,9 +13,6 @@ from ..width_height_selector.logic import get_preset_dimensions
from ..width_height_selector.presets import (
PRESET_OPTIONS,
PRESET_METADATA,
get_model_recommendation,
get_preset_metadata,
get_presets_by_model_group,
)
@@ -0,0 +1,8 @@
"""
KikoSaveImage tool module
Enhanced image saving with format selection, quality control, and clickable previews
"""
from .node import KikoSaveImageNode
__all__ = ["KikoSaveImageNode"]
+365
View File
@@ -0,0 +1,365 @@
"""
KikoSaveImage core logic
Enhanced image saving functionality with multiple format support
"""
import os
import json
import numpy as np
from PIL import Image
from PIL.PngImagePlugin import PngInfo
import torch
from typing import Dict, List, Any, Optional, Tuple
import time
try:
import folder_paths
except ImportError:
# Fallback for testing without ComfyUI
class folder_paths:
@staticmethod
def get_output_directory():
return "./output"
def get_save_image_path(
filename_prefix: str,
batch_number: int,
format_ext: str,
output_dir: str,
subfolder: str = "",
) -> Tuple[str, str]:
"""
Generate save path for image with proper filename handling
Args:
filename_prefix: Base filename prefix
batch_number: Batch index for multiple images
format_ext: File extension (.png, .jpg, .webp)
output_dir: Output directory path
subfolder: Optional subfolder within output directory
Returns:
Tuple of (full_path, relative_filename)
"""
# Split filename_prefix into directory path and actual filename prefix
# This allows for directory structures like "kittybear/anime/images/kittybear"
prefix_dir = os.path.dirname(filename_prefix)
prefix_name = os.path.basename(filename_prefix)
# Sanitize only the filename part (not the directory path)
safe_prefix = prefix_name.replace(
":", "_"
) # Only sanitize problematic chars for filenames
safe_prefix = "".join(c for c in safe_prefix if c.isalnum() or c in "._-")
# Create unique filename with timestamp to avoid conflicts
timestamp = int(time.time())
filename = f"{safe_prefix}_{timestamp:010d}_{batch_number:05d}{format_ext}"
# Handle subfolder and prefix directory (but not the filename part)
path_components = []
path_components.append(output_dir)
if subfolder:
path_components.append(subfolder)
# Only add prefix_dir if it exists (the directory part, not the filename part)
if prefix_dir:
path_components.append(prefix_dir)
full_output_folder = os.path.join(*path_components)
# Ensure directory exists
os.makedirs(full_output_folder, exist_ok=True)
full_path = os.path.join(full_output_folder, filename)
# For the preview, ComfyUI needs the filename and subfolder separately
# The subfolder needs to be relative to the output directory root
# Build the relative subfolder path including prefix directory (but not filename part)
relative_path_components = []
if subfolder:
relative_path_components.append(subfolder.strip("/\\"))
if prefix_dir:
relative_path_components.append(prefix_dir.strip("/\\"))
if relative_path_components:
relative_subfolder = os.path.join(*relative_path_components)
else:
relative_subfolder = ""
preview_filename = filename
return full_path, preview_filename, relative_subfolder
def convert_tensor_to_pil(image_tensor: torch.Tensor) -> Image.Image:
"""
Convert ComfyUI image tensor to PIL Image
Args:
image_tensor: Tensor in format [height, width, channels] with values 0-1
Returns:
PIL Image in RGB/RGBA format
"""
# Convert tensor (0-1 float) to 0-255 numpy array
i = 255.0 * image_tensor.cpu().numpy()
img_array = np.clip(i, 0, 255).astype(np.uint8)
# Create PIL image from numpy array
img = Image.fromarray(img_array)
return img
def create_png_metadata(
prompt: Optional[Dict] = None, extra_pnginfo: Optional[Dict] = None
) -> Optional[PngInfo]:
"""
Create PNG metadata with workflow information
Args:
prompt: ComfyUI prompt data
extra_pnginfo: Additional PNG metadata
Returns:
PngInfo object or None if no metadata
"""
if prompt is None and extra_pnginfo is None:
return None
metadata = PngInfo()
if prompt is not None:
metadata.add_text("prompt", json.dumps(prompt))
if extra_pnginfo is not None:
for key, value in extra_pnginfo.items():
metadata.add_text(key, json.dumps(value))
return metadata
def save_image_with_format(
img: Image.Image,
filepath: str,
format_type: str,
quality: int = 90,
png_compress_level: int = 4,
webp_lossless: bool = False,
metadata: Optional[PngInfo] = None,
) -> Dict[str, Any]:
"""
Save PIL image with specified format and quality settings
Args:
img: PIL Image to save
filepath: Full path to save file
format_type: Image format (PNG, JPEG, WEBP)
quality: JPEG/WebP quality (1-100)
png_compress_level: PNG compression level (0-9)
webp_lossless: Use lossless WebP compression
metadata: PNG metadata to embed
Returns:
Dict with save information
"""
save_kwargs = {}
if format_type == "PNG":
if metadata:
save_kwargs["pnginfo"] = metadata
save_kwargs["compress_level"] = png_compress_level
elif format_type == "JPEG":
# Convert RGBA to RGB for JPEG (no transparency support)
if img.mode == "RGBA":
# Create white background
background = Image.new("RGB", img.size, (255, 255, 255))
background.paste(img, mask=img.split()[-1]) # Use alpha channel as mask
img = background
elif img.mode != "RGB":
img = img.convert("RGB")
save_kwargs["quality"] = quality
save_kwargs["optimize"] = True
elif format_type == "WEBP":
save_kwargs["quality"] = quality if not webp_lossless else 100
save_kwargs["lossless"] = webp_lossless
else:
raise ValueError(f"Unsupported format: {format_type}")
# Save the image
img.save(filepath, **save_kwargs)
# Get file size for info
file_size = os.path.getsize(filepath)
return {
"filepath": filepath,
"format": format_type,
"file_size": file_size,
"quality": quality if format_type != "PNG" else None,
"compress_level": png_compress_level if format_type == "PNG" else None,
"lossless": webp_lossless if format_type == "WEBP" else None,
}
def process_image_batch(
images: torch.Tensor,
filename_prefix: str = "KikoSave",
format_type: str = "PNG",
quality: int = 90,
png_compress_level: int = 4,
webp_lossless: bool = False,
popup: bool = True,
prompt: Optional[Dict] = None,
extra_pnginfo: Optional[Dict] = None,
) -> List[Dict[str, Any]]:
"""
Process and save a batch of images with specified format settings
Args:
images: Batch of image tensors [batch, height, width, channels]
filename_prefix: Prefix for saved filenames
format_type: Image format (PNG, JPEG, WEBP)
quality: JPEG/WebP quality (1-100)
png_compress_level: PNG compression level (0-9)
webp_lossless: Use lossless WebP compression
popup: Enable popup windows in UI
prompt: ComfyUI prompt data for metadata
extra_pnginfo: Additional PNG metadata
Returns:
List of saved image information dicts
"""
# Get output directory
output_dir = folder_paths.get_output_directory()
# Determine file extension
format_extensions = {"PNG": ".png", "JPEG": ".jpg", "WEBP": ".webp"}
if format_type not in format_extensions:
raise ValueError(
f"Unsupported format: {format_type}. "
f"Supported: {list(format_extensions.keys())}"
)
format_ext = format_extensions[format_type]
# Create metadata for PNG
metadata = None
if format_type == "PNG":
metadata = create_png_metadata(prompt, extra_pnginfo)
# Process each image in the batch
results = []
enhanced_data = []
for batch_number, image_tensor in enumerate(images):
# Convert tensor to PIL Image
img = convert_tensor_to_pil(image_tensor)
# Generate save path
filepath, preview_filename, relative_subfolder = get_save_image_path(
filename_prefix, batch_number, format_ext, output_dir, ""
)
# Save with format-specific settings
save_info = save_image_with_format(
img,
filepath,
format_type,
quality,
png_compress_level,
webp_lossless,
metadata,
)
# Build result info for ComfyUI preview
# ONLY the core fields that ComfyUI expects - no extra metadata
result = {
"filename": preview_filename,
"subfolder": relative_subfolder,
"type": "output",
}
# Store enhanced data separately
enhanced_info = {
"filename": preview_filename,
"subfolder": relative_subfolder,
"popup": popup,
"type": "output",
"format": format_type,
"file_size": save_info["file_size"],
"dimensions": f"{img.width}x{img.height}",
}
# Add format-specific info to enhanced data
if format_type == "PNG":
enhanced_info["compress_level"] = png_compress_level
elif format_type in ["JPEG", "WEBP"]:
enhanced_info["quality"] = quality
if format_type == "WEBP":
enhanced_info["lossless"] = webp_lossless
results.append(result)
enhanced_data.append(enhanced_info)
return results, enhanced_data
def validate_save_inputs(
images: torch.Tensor, format_type: str, quality: int, png_compress_level: int
) -> None:
"""
Validate inputs for image saving
Args:
images: Image tensor batch to validate
format_type: Image format to validate
quality: Quality setting to validate
png_compress_level: PNG compression level to validate
Raises:
ValueError: If validation fails
"""
# Validate images tensor
if not isinstance(images, torch.Tensor):
raise ValueError(f"images must be a torch.Tensor, got {type(images).__name__}")
if len(images.shape) != 4:
raise ValueError(
f"images tensor must have 4 dimensions [batch, height, width, channels], "
f"got {len(images.shape)}"
)
# Validate format
supported_formats = ["PNG", "JPEG", "WEBP"]
if format_type not in supported_formats:
raise ValueError(
f"format must be one of {supported_formats}, got {format_type}"
)
# Validate quality (for JPEG/WebP)
if format_type in ["JPEG", "WEBP"]:
if not isinstance(quality, int) or not (1 <= quality <= 100):
raise ValueError(
f"quality must be an integer between 1 and 100, got {quality}"
)
# Validate PNG compression level
if format_type == "PNG":
if not isinstance(png_compress_level, int) or not (
0 <= png_compress_level <= 9
):
raise ValueError(
f"png_compress_level must be an integer between 0 and 9, "
f"got {png_compress_level}"
)
+226
View File
@@ -0,0 +1,226 @@
"""
KikoSaveImage ComfyUI Node
Enhanced image saving with format selection, quality control, and clickable previews
"""
import torch
from typing import Dict, Any, Optional
from ...base import ComfyAssetsBaseNode
from .logic import process_image_batch, validate_save_inputs
class KikoSaveImageNode(ComfyAssetsBaseNode):
"""
Enhanced ComfyUI image saving node with multiple format support
Features:
- Multiple format support (PNG, JPEG, WebP)
- Quality/compression controls
- Clickable image previews
- Metadata preservation
- Batch processing
Inputs:
- images (IMAGE): Images to save
- filename_prefix (STRING): Prefix for saved filenames
- format (COMBO): Output format (PNG, JPEG, WebP)
- quality (INT): JPEG/WebP quality (1-100)
- png_compress_level (INT): PNG compression level (0-9)
- webp_lossless (BOOLEAN): Use lossless WebP compression
- subfolder (STRING): Optional subfolder for organization
Outputs:
- UI: Image preview data for ComfyUI interface
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Define ComfyUI input interface with enhanced save options
Returns:
Dict with required and optional input specifications
"""
return {
"required": {
"images": ("IMAGE", {"tooltip": "The images to save"}),
"filename_prefix": (
"STRING",
{"default": "KikoSave", "tooltip": "Prefix for saved filenames"},
),
"format": (
["PNG", "JPEG", "WEBP"],
{"default": "PNG", "tooltip": "Output image format"},
),
},
"optional": {
"quality": (
"INT",
{
"default": 90,
"min": 1,
"max": 100,
"step": 1,
"tooltip": "JPEG/WebP quality (1-100, higher = better quality)",
},
),
"png_compress_level": (
"INT",
{
"default": 4,
"min": 0,
"max": 9,
"step": 1,
"tooltip": "PNG compression level (0-9, higher = smaller file)",
},
),
"webp_lossless": (
"BOOLEAN",
{
"default": False,
"tooltip": "Use lossless WebP compression "
"(ignores quality setting)",
},
),
"popup": (
"BOOLEAN",
{
"default": True,
"tooltip": "Enable popup windows when clicking on images in the viewer",
},
),
},
"hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
}
RETURN_TYPES = ()
FUNCTION = "save_images"
OUTPUT_NODE = True
def save_images(
self,
images: torch.Tensor,
filename_prefix: str = "KikoSave",
format: str = "PNG",
quality: int = 90,
png_compress_level: int = 4,
webp_lossless: bool = False,
popup: bool = True,
prompt: Optional[Dict] = None,
extra_pnginfo: Optional[Dict] = None,
) -> Dict[str, Any]:
"""
Save images with enhanced format and quality options
Args:
images: Batch of image tensors to save
filename_prefix: Prefix for saved filenames
format: Output format (PNG, JPEG, WebP)
quality: JPEG/WebP quality setting
png_compress_level: PNG compression level
webp_lossless: Use lossless WebP compression
popup: Enable popup windows when clicking on images
prompt: ComfyUI prompt data for metadata
extra_pnginfo: Additional PNG metadata
Returns:
Dict with UI data for image previews
Raises:
ValueError: If validation fails
"""
try:
# Validate inputs
self.validate_inputs(
images=images,
format=format,
quality=quality,
png_compress_level=png_compress_level,
webp_lossless=webp_lossless,
popup=popup,
)
# Log the save operation
self.log_info(
f"Saving {len(images)} images as {format} "
f"(quality={quality if format != 'PNG' else 'N/A'}, "
f"png_compress={png_compress_level if format == 'PNG' else 'N/A'})"
)
# Process and save images
results, enhanced_data = process_image_batch(
images=images,
filename_prefix=filename_prefix,
format_type=format,
quality=quality,
png_compress_level=png_compress_level,
webp_lossless=webp_lossless,
popup=popup,
prompt=prompt,
extra_pnginfo=extra_pnginfo,
)
# Log results
total_size = sum(data["file_size"] for data in enhanced_data)
self.log_info(
f"Successfully saved {len(results)} images "
f"(total size: {total_size / 1024:.1f} KB)"
)
# Return UI data for ComfyUI preview (clean) + enhanced data for our JS
return {
"ui": {
"images": results, # Clean data for ComfyUI
"kiko_enhanced": enhanced_data, # Enhanced data for our JavaScript
}
}
except Exception as e:
error_msg = f"Failed to save images: {str(e)}"
self.handle_error(error_msg, e)
def validate_inputs(
self,
images: torch.Tensor,
format: str,
quality: int,
png_compress_level: int,
webp_lossless: bool,
popup: bool,
) -> None:
"""
Validate inputs specific to KikoSaveImage
Args:
images: Image tensor batch
format: Image format
quality: Quality setting
png_compress_level: PNG compression level
webp_lossless: WebP lossless setting
popup: Enable popup windows
Raises:
ValueError: If validation fails
"""
# Use logic module validation
validate_save_inputs(images, format, quality, png_compress_level)
# Additional node-specific validation
if not isinstance(webp_lossless, bool):
raise ValueError(
f"webp_lossless must be a boolean, got {type(webp_lossless).__name__}"
)
if not isinstance(popup, bool):
raise ValueError(f"popup must be a boolean, got {type(popup).__name__}")
# Node class mappings for ComfyUI registration
NODE_CLASS_MAPPINGS = {
"KikoSaveImage": KikoSaveImageNode,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"KikoSaveImage": "Kiko Save Image",
}
+36 -28
View File
@@ -38,12 +38,12 @@ class ResolutionCalculatorNode(ComfyAssetsBaseNode):
"FLOAT",
{
"default": 2.0,
"min": 1.0,
"min": 0.1,
"max": 8.0,
"step": 0.1,
"display": "slider",
"tooltip": "Factor to scale the resolution by "
"(e.g., 2.0 for 2x upscale)",
"(e.g., 2.0 for 2x, 0.5 for half scale)",
},
),
},
@@ -140,38 +140,46 @@ class ResolutionCalculatorNode(ComfyAssetsBaseNode):
f"scale_factor must be a number, got {type(scale_factor).__name__}"
)
# Additional tensor validation
# Validate tensors using helper methods
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 "
f"[batch, height, width, channels], got {len(image.shape)}"
)
self._validate_image_tensor(image)
if latent is not None:
if not isinstance(latent, dict):
raise ValueError(f"latent must be a dict, got {type(latent).__name__}")
self._validate_latent_dict(latent)
if "samples" not in latent:
raise ValueError("latent dict must contain 'samples' key")
def _validate_image_tensor(self, image: torch.Tensor) -> None:
"""Validate image tensor format"""
if not isinstance(image, torch.Tensor):
raise ValueError(
f"image must be a torch.Tensor, got {type(image).__name__}"
)
samples = latent["samples"]
if not isinstance(samples, torch.Tensor):
raise ValueError(
f"latent['samples'] must be a torch.Tensor, "
f"got {type(samples).__name__}"
)
if len(image.shape) != 4:
raise ValueError(
f"image tensor must have 4 dimensions "
f"[batch, height, width, channels], got {len(image.shape)}"
)
if len(samples.shape) != 4:
raise ValueError(
f"latent samples tensor must have 4 dimensions "
f"[batch, channels, height, width], got {len(samples.shape)}"
)
def _validate_latent_dict(self, latent: Dict[str, torch.Tensor]) -> None:
"""Validate latent dictionary format"""
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, "
f"got {type(samples).__name__}"
)
if len(samples.shape) != 4:
raise ValueError(
f"latent samples tensor must have 4 dimensions "
f"[batch, channels, height, width], got {len(samples.shape)}"
)
# Node class mappings for ComfyUI registration
+20 -5
View File
@@ -60,14 +60,14 @@ class SamplerComboCompactNode(ComfyAssetsBaseNode):
}
}
RETURN_TYPES = (SAMPLERS, SCHEDULERS, "INT", "FLOAT")
RETURN_TYPES = ("SAMPLER", SCHEDULERS, "INT", "FLOAT")
RETURN_NAMES = ("sampler", "scheduler", "steps", "cfg")
FUNCTION = "get_combo"
CATEGORY = "ComfyAssets"
def get_combo(
self, sampler: str, sched: str, steps: int, cfg: float
) -> Tuple[str, str, int, float]:
) -> Tuple[object, str, int, float]:
"""
Get compact sampler combo configuration.
@@ -78,17 +78,32 @@ class SamplerComboCompactNode(ComfyAssetsBaseNode):
cfg: CFG scale value
Returns:
Tuple of (sampler, scheduler, steps, cfg)
Tuple of (sampler_object, scheduler, steps, cfg)
"""
try:
# Use the same validation logic but with compact interface
result = get_sampler_combo(sampler, sched, steps, cfg)
return result
# Create the sampler object
try:
import comfy.samplers
sampler_obj = comfy.samplers.sampler_object(result[0])
except ImportError:
# Return sampler name for testing
sampler_obj = result[0]
return (sampler_obj, result[1], result[2], result[3])
except Exception as e:
# Graceful fallback
self.handle_error(f"Error in compact combo: {str(e)}")
return ("euler", "normal", 20, 7.0)
try:
import comfy.samplers
sampler_obj = comfy.samplers.sampler_object("euler")
except ImportError:
# Return sampler name for testing
sampler_obj = "euler"
return (sampler_obj, "normal", 20, 7.0)
def __str__(self) -> str:
"""String representation of the compact node."""
+29 -6
View File
@@ -65,14 +65,14 @@ class SamplerComboNode(ComfyAssetsBaseNode):
}
}
RETURN_TYPES = (SAMPLERS, SCHEDULERS, "INT", "FLOAT")
RETURN_TYPES = ("SAMPLER", SCHEDULERS, "INT", "FLOAT")
RETURN_NAMES = ("sampler_name", "scheduler", "steps", "cfg")
FUNCTION = "get_sampler_combo"
CATEGORY = "ComfyAssets"
def get_sampler_combo(
self, sampler_name: str, scheduler: str, steps: int, cfg: float
) -> Tuple[str, str, int, float]:
) -> Tuple[object, str, int, float]:
"""
Get sampler combo configuration.
@@ -83,7 +83,7 @@ class SamplerComboNode(ComfyAssetsBaseNode):
cfg: CFG scale value
Returns:
Tuple of (sampler_name, scheduler, steps, cfg)
Tuple of (sampler_object, scheduler, steps, cfg)
"""
try:
# Validate inputs
@@ -98,17 +98,33 @@ class SamplerComboNode(ComfyAssetsBaseNode):
f"steps={steps}, cfg={cfg}. "
f"Using safe defaults: euler, normal, 20 steps, CFG 7.0"
)
return ("euler", "normal", 20, 7.0)
try:
import comfy.samplers
sampler = comfy.samplers.sampler_object("euler")
except ImportError:
# Return mock object for testing
sampler = "euler"
return (sampler, "normal", 20, 7.0)
# Process and return the combo
result = get_sampler_combo(sampler_name, scheduler, steps, cfg)
# Create the sampler object
try:
import comfy.samplers
sampler = comfy.samplers.sampler_object(result[0])
except ImportError:
# Return sampler name for testing
sampler = result[0]
self.log_info(
f"Configured sampler combo: {result[0]}, {result[1]}, "
f"{result[2]} steps, CFG {result[3]}"
)
return result
return (sampler, result[1], result[2], result[3])
except Exception as e:
# Handle any unexpected errors gracefully
@@ -119,7 +135,14 @@ class SamplerComboNode(ComfyAssetsBaseNode):
f"{self.__class__.__name__}: Error processing sampler combo: {str(e)}. "
f"Using safe defaults: euler, normal, 20 steps, CFG 7.0"
)
return ("euler", "normal", 20, 7.0)
try:
import comfy.samplers
sampler = comfy.samplers.sampler_object("euler")
except ImportError:
# Return mock object for testing
sampler = "euler"
return (sampler, "normal", 20, 7.0)
def validate_inputs(
self, sampler_name: str, scheduler: str, steps: int, cfg: float
+29 -3
View File
@@ -1,14 +1,40 @@
[build-system]
requires = ["setuptools>=61.0", "wheel"]
build-backend = "setuptools.build_meta"
[project]
name = "kikotools"
description = "Simple tools for ComfyUI"
version = "1.0.2"
license = {file = "LICENSE"}
dependencies = ["# Development dependencies for ComfyUI-KikoTools", "# Testing framework", "pytest>=7.0.0", "pytest-cov>=4.0.0", "pytest-mock>=3.10.0", "# Code quality", "black>=23.0.0", "flake8>=6.0.0", "mypy>=1.0.0", "# Development utilities", "pre-commit>=3.0.0", "# ComfyUI testing (mock dependencies for unit tests)", "torch>=2.0.0", "numpy>=1.24.0", "pillow>=9.0.0"]
version = "1.0.7"
license = {text = "MIT"}
dependencies = []
[project.optional-dependencies]
dev = [
# Testing framework
"pytest>=7.0.0",
"pytest-cov>=4.0.0",
"pytest-mock>=3.10.0",
# Code quality
"black>=23.0.0",
"flake8>=6.0.0",
"mypy>=1.0.0",
# Development utilities
"pre-commit>=3.0.0",
# ComfyUI testing (mock dependencies for unit tests)
"torch>=2.0.0",
"numpy>=1.24.0",
"pillow>=9.0.0"
]
[project.urls]
Repository = "https://github.com/ComfyAssets/ComfyUI-KikoTools"
# Used by Comfy Registry https://registry.comfy.org
[tool.setuptools.packages.find]
include = ["kikotools*"]
exclude = ["tests*", "web*"]
[tool.comfy]
PublisherId = "kiko9"
DisplayName = "ComfyUI-KikoTools"
+546
View File
@@ -0,0 +1,546 @@
"""
Unit tests for KikoSaveImage tool
Tests image saving functionality with multiple formats and quality settings
"""
import pytest
import torch
import tempfile
import os
from PIL import Image
from unittest.mock import patch
from kikotools.tools.kiko_save_image.node import KikoSaveImageNode
from kikotools.tools.kiko_save_image.logic import (
convert_tensor_to_pil,
process_image_batch,
validate_save_inputs,
save_image_with_format,
get_save_image_path,
create_png_metadata,
)
class TestKikoSaveImageLogic:
"""Test core logic functions"""
def test_convert_tensor_to_pil(self):
"""Test tensor to PIL conversion"""
# Create test tensor [height, width, channels] with values 0-1
tensor = torch.rand(64, 64, 3)
# Convert to PIL
pil_image = convert_tensor_to_pil(tensor)
# Verify conversion
assert isinstance(pil_image, Image.Image)
assert pil_image.size == (64, 64) # PIL uses (width, height)
assert pil_image.mode in ["RGB", "RGBA"]
def test_convert_tensor_to_pil_rgba(self):
"""Test tensor to PIL conversion with alpha channel"""
# Create RGBA tensor
tensor = torch.rand(32, 32, 4)
pil_image = convert_tensor_to_pil(tensor)
assert isinstance(pil_image, Image.Image)
assert pil_image.size == (32, 32)
assert pil_image.mode == "RGBA"
def test_get_save_image_path(self):
"""Test save path generation"""
with tempfile.TemporaryDirectory() as temp_dir:
# Test basic path generation
full_path, filename = get_save_image_path(
"test_prefix", 0, ".png", temp_dir
)
assert full_path.startswith(temp_dir)
assert filename.startswith("test_prefix_")
assert filename.endswith("_00000.png")
# Test with empty subfolder (standard behavior)
full_path, filename = get_save_image_path("test", 1, ".jpg", temp_dir, "")
assert full_path.startswith(temp_dir)
assert filename.startswith("test_")
assert filename.endswith("_00001.jpg")
def test_create_png_metadata(self):
"""Test PNG metadata creation"""
# Test with no metadata
metadata = create_png_metadata()
assert metadata is None
# Test with prompt data
prompt_data = {"test": "value"}
metadata = create_png_metadata(prompt=prompt_data)
assert metadata is not None
# Check that metadata contains our data (implementation detail)
assert hasattr(metadata, "text")
@patch("kikotools.tools.kiko_save_image.logic.folder_paths")
def test_process_image_batch_png(self, mock_folder_paths):
"""Test batch processing with PNG format"""
with tempfile.TemporaryDirectory() as temp_dir:
mock_folder_paths.get_output_directory.return_value = temp_dir
# Create test image batch [batch, height, width, channels]
images = torch.rand(2, 32, 32, 3)
# Process batch
results, enhanced_data = process_image_batch(
images=images,
filename_prefix="test_batch",
format_type="PNG",
png_compress_level=6,
)
# Verify results (clean data)
assert len(results) == 2
for i, result in enumerate(results):
assert "filename" in result
assert "subfolder" in result
assert "type" in result
assert result["type"] == "output"
# Verify enhanced data
assert len(enhanced_data) == 2
for i, enhanced in enumerate(enhanced_data):
assert enhanced["format"] == "PNG"
assert enhanced["compress_level"] == 6
assert enhanced["dimensions"] == "32x32"
assert enhanced["popup"] is True # Default popup value
assert "file_size" in enhanced
# Verify file was saved
filepath = os.path.join(temp_dir, enhanced["filename"])
assert os.path.exists(filepath)
# Verify image can be loaded
saved_img = Image.open(filepath)
assert saved_img.size == (32, 32)
@patch("kikotools.tools.kiko_save_image.logic.folder_paths")
def test_process_image_batch_jpeg(self, mock_folder_paths):
"""Test batch processing with JPEG format"""
with tempfile.TemporaryDirectory() as temp_dir:
mock_folder_paths.get_output_directory.return_value = temp_dir
# Create test image batch
images = torch.rand(1, 64, 64, 3)
# Process batch
results, enhanced_data = process_image_batch(
images=images,
filename_prefix="test_jpeg",
format_type="JPEG",
quality=85,
)
# Verify results
assert len(results) == 1
assert len(enhanced_data) == 1
enhanced = enhanced_data[0]
assert enhanced["format"] == "JPEG"
assert enhanced["quality"] == 85
assert enhanced["filename"].endswith(".jpg")
# Verify file exists and can be loaded
filepath = os.path.join(temp_dir, results[0]["filename"])
assert os.path.exists(filepath)
saved_img = Image.open(filepath)
assert saved_img.size == (64, 64)
assert saved_img.mode == "RGB" # JPEG converts to RGB
@patch("kikotools.tools.kiko_save_image.logic.folder_paths")
def test_process_image_batch_webp(self, mock_folder_paths):
"""Test batch processing with WebP format"""
with tempfile.TemporaryDirectory() as temp_dir:
mock_folder_paths.get_output_directory.return_value = temp_dir
# Create test image batch
images = torch.rand(1, 48, 48, 3)
# Test lossless WebP
results = process_image_batch(
images=images,
filename_prefix="test_webp",
format_type="WEBP",
quality=90,
webp_lossless=True,
)
assert len(results) == 1
result = results[0]
assert result["format"] == "WEBP"
assert result["lossless"] is True
assert result["filename"].endswith(".webp")
def test_validate_save_inputs_valid(self):
"""Test input validation with valid inputs"""
images = torch.rand(2, 64, 64, 3)
# Should not raise exception
validate_save_inputs(images, "PNG", 90, 4)
validate_save_inputs(images, "JPEG", 85, 4)
validate_save_inputs(images, "WEBP", 95, 6)
def test_validate_save_inputs_invalid_tensor(self):
"""Test validation with invalid tensor"""
# Wrong tensor dimensions
invalid_tensor = torch.rand(64, 64) # Missing batch and channel dims
with pytest.raises(ValueError, match="4 dimensions"):
validate_save_inputs(invalid_tensor, "PNG", 90, 4)
# Non-tensor input
with pytest.raises(ValueError, match="torch.Tensor"):
validate_save_inputs("not_a_tensor", "PNG", 90, 4)
def test_validate_save_inputs_invalid_format(self):
"""Test validation with invalid format"""
images = torch.rand(1, 32, 32, 3)
with pytest.raises(ValueError, match="format must be one of"):
validate_save_inputs(images, "BMP", 90, 4)
def test_validate_save_inputs_invalid_quality(self):
"""Test validation with invalid quality"""
images = torch.rand(1, 32, 32, 3)
# Quality out of range
with pytest.raises(
ValueError, match="quality must be an integer between 1 and 100"
):
validate_save_inputs(images, "JPEG", 0, 4)
with pytest.raises(
ValueError, match="quality must be an integer between 1 and 100"
):
validate_save_inputs(images, "JPEG", 101, 4)
def test_validate_save_inputs_invalid_compress_level(self):
"""Test validation with invalid PNG compression level"""
images = torch.rand(1, 32, 32, 3)
with pytest.raises(
ValueError, match="png_compress_level must be an integer between 0 and 9"
):
validate_save_inputs(images, "PNG", 90, -1)
with pytest.raises(
ValueError, match="png_compress_level must be an integer between 0 and 9"
):
validate_save_inputs(images, "PNG", 90, 10)
def test_save_image_with_format_png(self):
"""Test saving with PNG format"""
with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as temp_file:
temp_path = temp_file.name
try:
# Create test PIL image
img = Image.new("RGB", (32, 32), color="red")
# Save with PNG format
result = save_image_with_format(img, temp_path, "PNG", png_compress_level=8)
assert result["format"] == "PNG"
assert result["compress_level"] == 8
assert os.path.exists(temp_path)
# Verify saved image
saved_img = Image.open(temp_path)
assert saved_img.size == (32, 32)
finally:
if os.path.exists(temp_path):
os.unlink(temp_path)
def test_save_image_with_format_jpeg_rgba_conversion(self):
"""Test JPEG saving with RGBA to RGB conversion"""
with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as temp_file:
temp_path = temp_file.name
try:
# Create RGBA image
img = Image.new("RGBA", (32, 32), color=(255, 0, 0, 128))
# Save as JPEG (should convert to RGB)
result = save_image_with_format(img, temp_path, "JPEG", quality=95)
assert result["format"] == "JPEG"
assert result["quality"] == 95
# Verify saved image is RGB
saved_img = Image.open(temp_path)
assert saved_img.mode == "RGB"
finally:
if os.path.exists(temp_path):
os.unlink(temp_path)
class TestKikoSaveImageNode:
"""Test KikoSaveImageNode class"""
def setup_method(self):
"""Setup test fixtures"""
self.node = KikoSaveImageNode()
def test_input_types(self):
"""Test INPUT_TYPES class method"""
input_types = KikoSaveImageNode.INPUT_TYPES()
# Check required inputs
required = input_types["required"]
assert "images" in required
assert "filename_prefix" in required
assert "format" in required
# Check format options
format_options = required["format"][0]
assert "PNG" in format_options
assert "JPEG" in format_options
assert "WEBP" in format_options
# Check optional inputs
optional = input_types["optional"]
assert "quality" in optional
assert "png_compress_level" in optional
assert "webp_lossless" in optional
assert "popup" in optional
# Check hidden inputs
hidden = input_types["hidden"]
assert "prompt" in hidden
assert "extra_pnginfo" in hidden
def test_node_attributes(self):
"""Test node class attributes"""
assert KikoSaveImageNode.RETURN_TYPES == ()
assert KikoSaveImageNode.FUNCTION == "save_images"
assert KikoSaveImageNode.OUTPUT_NODE is True
assert KikoSaveImageNode.CATEGORY == "ComfyAssets"
@patch("kikotools.tools.kiko_save_image.node.process_image_batch")
def test_save_images_success(self, mock_process):
"""Test successful image saving"""
# Setup mock - new return format (results, enhanced_data)
mock_results = [
{
"filename": "test_00001_00000.png",
"subfolder": "",
"type": "output",
}
]
mock_enhanced = [
{
"filename": "test_00001_00000.png",
"popup": True,
"type": "output",
"format": "PNG",
"file_size": 1024,
"dimensions": "64x64",
}
]
mock_process.return_value = (mock_results, mock_enhanced)
# Create test input
images = torch.rand(1, 64, 64, 3)
# Call save_images
result = self.node.save_images(
images=images,
filename_prefix="test",
format="PNG",
quality=90,
png_compress_level=4,
)
# Verify mock was called
mock_process.assert_called_once()
# Verify result format
assert "ui" in result
assert "images" in result["ui"]
assert "kiko_enhanced" in result["ui"]
assert result["ui"]["images"] == mock_results
assert result["ui"]["kiko_enhanced"] == mock_enhanced
def test_validate_inputs_success(self):
"""Test input validation with valid inputs"""
images = torch.rand(1, 32, 32, 3)
# Should not raise exception
self.node.validate_inputs(
images=images,
format="PNG",
quality=90,
png_compress_level=4,
webp_lossless=False,
popup=True,
)
def test_validate_inputs_invalid_webp_lossless(self):
"""Test validation with invalid webp_lossless type"""
images = torch.rand(1, 32, 32, 3)
with pytest.raises(ValueError, match="webp_lossless must be a boolean"):
self.node.validate_inputs(
images=images,
format="PNG",
quality=90,
png_compress_level=4,
webp_lossless="not_boolean",
popup=True,
)
def test_validate_inputs_invalid_popup(self):
"""Test validation with invalid popup"""
images = torch.rand(1, 32, 32, 3)
# Non-boolean popup
with pytest.raises(ValueError, match="popup must be a boolean"):
self.node.validate_inputs(
images=images,
format="PNG",
quality=90,
png_compress_level=4,
webp_lossless=False,
popup="not_boolean",
)
@patch("kikotools.tools.kiko_save_image.node.process_image_batch")
def test_save_images_error_handling(self, mock_process):
"""Test error handling in save_images method"""
# Setup mock to raise exception
mock_process.side_effect = Exception("Test error")
images = torch.rand(1, 32, 32, 3)
# Should handle error and re-raise with context
with pytest.raises(ValueError, match="Failed to save images"):
self.node.save_images(images=images)
def test_node_info(self):
"""Test get_node_info method"""
info = self.node.get_node_info()
assert info["class_name"] == "KikoSaveImageNode"
assert info["category"] == "ComfyAssets"
assert info["function"] == "save_images"
class TestNodeRegistration:
"""Test node registration mappings"""
def test_node_class_mappings(self):
"""Test NODE_CLASS_MAPPINGS contains KikoSaveImage"""
from kikotools.tools.kiko_save_image.node import NODE_CLASS_MAPPINGS
assert "KikoSaveImage" in NODE_CLASS_MAPPINGS
assert NODE_CLASS_MAPPINGS["KikoSaveImage"] is KikoSaveImageNode
def test_node_display_name_mappings(self):
"""Test NODE_DISPLAY_NAME_MAPPINGS contains KikoSaveImage"""
from kikotools.tools.kiko_save_image.node import NODE_DISPLAY_NAME_MAPPINGS
assert "KikoSaveImage" in NODE_DISPLAY_NAME_MAPPINGS
assert NODE_DISPLAY_NAME_MAPPINGS["KikoSaveImage"] == "Kiko Save Image"
# Integration test fixtures
@pytest.fixture
def sample_image_tensor():
"""Create sample image tensor for testing"""
# Create a colorful test image [batch, height, width, channels]
batch_size, height, width, channels = 2, 64, 64, 3
# Create gradient pattern
tensor = torch.zeros(batch_size, height, width, channels)
for b in range(batch_size):
for h in range(height):
for w in range(width):
# Create RGB gradient pattern
tensor[b, h, w, 0] = h / height # Red gradient
tensor[b, h, w, 1] = w / width # Green gradient
tensor[b, h, w, 2] = (b + 1) * 0.5 # Blue varies by batch
return tensor
class TestIntegration:
"""Integration tests using sample data"""
@patch("kikotools.tools.kiko_save_image.logic.folder_paths")
def test_full_pipeline_png(self, mock_folder_paths, sample_image_tensor):
"""Test complete pipeline with PNG format"""
with tempfile.TemporaryDirectory() as temp_dir:
mock_folder_paths.get_output_directory.return_value = temp_dir
node = KikoSaveImageNode()
# Save images
result = node.save_images(
images=sample_image_tensor,
filename_prefix="integration_test",
format="PNG",
png_compress_level=6,
)
# Verify result structure
assert "ui" in result
assert "images" in result["ui"]
assert len(result["ui"]["images"]) == 2
# Verify files were created
for image_info in result["ui"]["images"]:
filepath = os.path.join(temp_dir, image_info["filename"])
assert os.path.exists(filepath)
# Verify image properties
img = Image.open(filepath)
assert img.size == (64, 64)
assert img.format == "PNG"
@patch("kikotools.tools.kiko_save_image.logic.folder_paths")
def test_full_pipeline_all_formats(self, mock_folder_paths, sample_image_tensor):
"""Test complete pipeline with all supported formats"""
with tempfile.TemporaryDirectory() as temp_dir:
mock_folder_paths.get_output_directory.return_value = temp_dir
node = KikoSaveImageNode()
# Test each format
formats_to_test = [
("PNG", {"png_compress_level": 8}),
("JPEG", {"quality": 85}),
("WEBP", {"quality": 90, "webp_lossless": False}),
("WEBP", {"quality": 100, "webp_lossless": True}),
]
for format_type, kwargs in formats_to_test:
result = node.save_images(
images=sample_image_tensor,
filename_prefix=f"test_{format_type.lower()}",
format=format_type,
**kwargs,
)
# Verify results
assert len(result["ui"]["images"]) == 2
for image_info in result["ui"]["images"]:
assert image_info["format"] == format_type
# Verify file exists and can be opened
filepath = os.path.join(temp_dir, image_info["filename"])
assert os.path.exists(filepath)
img = Image.open(filepath)
assert img.size == (64, 64)
+3 -3
View File
@@ -168,17 +168,17 @@ class TestSamplerComboNode:
steps_input = required["steps"]
assert steps_input[0] == "INT"
assert steps_input[1]["min"] == 1
assert steps_input[1]["max"] == 1000
assert steps_input[1]["max"] == 100
# Check CFG input structure
cfg_input = required["cfg"]
assert cfg_input[0] == "FLOAT"
assert cfg_input[1]["min"] == 0.0
assert cfg_input[1]["max"] == 30.0
assert cfg_input[1]["max"] == 20.0
def test_return_types_structure(self):
"""Test that return types are correctly defined."""
assert SamplerComboNode.RETURN_TYPES == (SAMPLERS, SCHEDULERS, "INT", "FLOAT")
assert SamplerComboNode.RETURN_TYPES == ("SAMPLER", SCHEDULERS, "INT", "FLOAT")
assert SamplerComboNode.RETURN_NAMES == (
"sampler_name",
"scheduler",
File diff suppressed because it is too large Load Diff