diff --git a/kikotools/__init__.py b/kikotools/__init__.py index f93cc80..0e7d1dc 100644 --- a/kikotools/__init__.py +++ b/kikotools/__init__.py @@ -26,6 +26,7 @@ from .tools.local_image_loader import LocalImageLoaderNode from .tools.project_file_path import ProjectFilePathNode from .tools.model_compare_sheet import ModelCompareSheetNode from .tools.model_downloader import ModelDownloaderNode +from .tools.model_resolution_picker import ModelResolutionPickerNode from .tools.resolution_calculator import ResolutionCalculatorNode from .tools.sampler_combo import SamplerComboCompactNode, SamplerComboNode from .tools.seed_history import SeedHistoryNode @@ -54,6 +55,7 @@ NODE_CLASS_MAPPINGS = { "SamplerCombo": SamplerComboNode, "SamplerComboCompact": SamplerComboCompactNode, "EmptyLatentBatch": EmptyLatentBatchNode, + "ModelResolutionPicker": ModelResolutionPickerNode, "KikoSaveImage": KikoSaveImageNode, "ImageToMultipleOf": ImageToMultipleOfNode, "ImageScaleDownBy": ImageScaleDownByNode, @@ -92,6 +94,7 @@ NODE_DISPLAY_NAME_MAPPINGS = { "SamplerCombo": "Sampler Combo", "SamplerComboCompact": "Sampler Combo (Compact)", "EmptyLatentBatch": "Empty Latent Batch", + "ModelResolutionPicker": "Model Resolution Picker", "KikoSaveImage": "Kiko Save Image", "ImageToMultipleOf": "Image to Multiple of", "ImageScaleDownBy": "Image Scale Down By", diff --git a/kikotools/tools/model_resolution_picker/__init__.py b/kikotools/tools/model_resolution_picker/__init__.py index a03b493..9821f21 100644 --- a/kikotools/tools/model_resolution_picker/__init__.py +++ b/kikotools/tools/model_resolution_picker/__init__.py @@ -1 +1,5 @@ """Model Resolution Picker module.""" + +from .node import ModelResolutionPickerNode + +__all__ = ["ModelResolutionPickerNode"] diff --git a/kikotools/tools/model_resolution_picker/node.py b/kikotools/tools/model_resolution_picker/node.py new file mode 100644 index 0000000..d26a3af --- /dev/null +++ b/kikotools/tools/model_resolution_picker/node.py @@ -0,0 +1,108 @@ +"""Model Resolution Picker node and its families route for ComfyUI.""" + +import logging +from typing import Any, Dict, Tuple + +from aiohttp import web + +from ...base.base_node import ComfyAssetsBaseNode +from ..empty_latent_batch.logic import create_empty_latent_batch +from .families import ( + DEFAULT_FAMILY, + families_payload, + family_labels, + resolve_family, +) +from .logic import round_to_step + +logger = logging.getLogger(__name__) + +FALLBACK_STEP = 8 + + +class ModelResolutionPickerNode(ComfyAssetsBaseNode): + """Pick a model family and size, and get an empty latent plus width/height.""" + + @classmethod + def INPUT_TYPES(cls) -> Dict[str, Any]: + """Define the input types for the ComfyUI node.""" + size_options = {"default": 1024, "min": 64, "max": 8192, "step": 8} + return { + "required": { + "model": ( + family_labels(), + { + "default": DEFAULT_FAMILY, + "tooltip": "Model family; the chips list its preferred sizes", + }, + ), + "width": ( + "INT", + { + **size_options, + "tooltip": "Custom width in pixels; rounded to the model " + "family's step", + }, + ), + "height": ( + "INT", + { + **size_options, + "tooltip": "Custom height in pixels; rounded to the model " + "family's step", + }, + ), + "batch_size": ( + "INT", + { + "default": 1, + "min": 1, + "max": 64, + "step": 1, + "tooltip": "Number of empty latents in the batch", + }, + ), + } + } + + RETURN_TYPES = ("LATENT", "INT", "INT") + RETURN_NAMES = ("latent", "width", "height") + FUNCTION = "create" + CATEGORY = "🫶 ComfyAssets/🖼️ Resolution" + + def create( + self, model: str, width: int, height: int, batch_size: int + ) -> Tuple[Dict[str, Any], int, int]: + """Round the size to the family's step and build the empty latent.""" + family = resolve_family(model) + if family is None: + logger.warning( + "Unknown model family '%s'; using step %d", model, FALLBACK_STEP + ) + step = family.step if family else FALLBACK_STEP + new_width = round_to_step(width, step) + new_height = round_to_step(height, step) + if (new_width, new_height) != (width, height): + self.log_info( + f"Rounded {width}x{height} to {new_width}x{new_height} " + f"(step {step})" + ) + latent = create_empty_latent_batch(new_width, new_height, batch_size) + return (latent, new_width, new_height) + + +async def get_families(request: Any) -> web.Response: + """Serve the model family table as JSON.""" + return web.json_response(families_payload()) + + +# Setup API routes +try: + import server + + server.PromptServer.instance.routes.get("/kikotools/model_resolution/families")( + get_families + ) +except ImportError: + # Server not available during testing + pass diff --git a/tests/unit/tools/model_resolution_picker/test_node.py b/tests/unit/tools/model_resolution_picker/test_node.py new file mode 100644 index 0000000..3437fb5 --- /dev/null +++ b/tests/unit/tools/model_resolution_picker/test_node.py @@ -0,0 +1,57 @@ +import asyncio +import json +import logging + +from kikotools.tools.model_resolution_picker.node import ( + ModelResolutionPickerNode, + get_families, +) + + +def test_inputs_and_outputs(): + inputs = ModelResolutionPickerNode.INPUT_TYPES()["required"] + assert inputs["model"][1]["default"] == "SDXL" + assert "HunyuanImage 2.1" in inputs["model"][0] + width = inputs["width"][1] + assert (width["default"], width["min"], width["max"], width["step"]) == ( + 1024, + 64, + 8192, + 8, + ) + assert ModelResolutionPickerNode.RETURN_TYPES == ("LATENT", "INT", "INT") + assert ModelResolutionPickerNode.CATEGORY == "🫶 ComfyAssets/🖼️ Resolution" + + +def test_execution_rounds_to_family_step_and_marks_latent(): + latent, w, h = ModelResolutionPickerNode().create("HunyuanImage 2.1", 70, 1000, 2) + assert (w, h) == (64, 992) + assert latent["samples"].shape == (2, 4, 124, 8) + assert latent["downscale_ratio_spacial"] == 8 + + +def test_exact_sizes_pass_through(): + _, w, h = ModelResolutionPickerNode().create("Flux.2 / Klein / MOSS", 1248, 832, 1) + assert (w, h) == (1248, 832) + + +def test_unknown_label_keeps_size_and_warns(caplog): + with caplog.at_level(logging.WARNING): + _, w, h = ModelResolutionPickerNode().create("Renamed Family", 1000, 1000, 1) + assert (w, h) == (1000, 1000) + assert "Renamed Family" in caplog.text + + +def test_families_route_returns_payload(): + response = asyncio.run(get_families(None)) + assert response.status == 200 + assert json.loads(response.text)["default"] == "SDXL" + + +def test_registered_in_kikotools(): + from kikotools import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + + assert NODE_CLASS_MAPPINGS["ModelResolutionPicker"] is ModelResolutionPickerNode + assert ( + NODE_DISPLAY_NAME_MAPPINGS["ModelResolutionPicker"] == "Model Resolution Picker" + )