Vito Sansevero
eaf03874b3
- Add a native resolution COMBO (shipped sizes only) so a size can be picked and saved in the workflow without the panel; custom or unparseable labels fall back to width/height on both sides. - Mirror the dropdown with the chips: options follow the active chip, the value tracks the current size, and picking a label writes width/height. - Draw a decorative ratio shape on each ratio chip and show the first families warning in the status line.
120 lines
3.9 KiB
Python
120 lines
3.9 KiB
Python
import asyncio
|
||
import json
|
||
import logging
|
||
from unittest import mock
|
||
|
||
from kikotools.tools.model_resolution_picker import families
|
||
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_resolution_options_come_from_shipped_families():
|
||
options, config = ModelResolutionPickerNode.INPUT_TYPES()["required"]["resolution"]
|
||
assert options[0] == "custom"
|
||
assert config["default"] == "custom"
|
||
assert "1216×832 - 3:2 (1.0MP)" in options
|
||
assert options[1:] == sorted(set(options[1:]))
|
||
|
||
|
||
def test_resolution_input_sits_between_model_and_width():
|
||
assert list(ModelResolutionPickerNode.INPUT_TYPES()["required"])[:3] == [
|
||
"model",
|
||
"resolution",
|
||
"width",
|
||
]
|
||
|
||
|
||
def test_resolution_label_overrides_width_height():
|
||
_, w, h = ModelResolutionPickerNode().create(
|
||
"SDXL", "1216×832 - 3:2 (1.0MP)", 1024, 1024, 1
|
||
)
|
||
assert (w, h) == (1216, 832)
|
||
|
||
|
||
def test_custom_uses_width_height():
|
||
_, w, h = ModelResolutionPickerNode().create("SDXL", "custom", 1024, 768, 1)
|
||
assert (w, h) == (1024, 768)
|
||
|
||
|
||
def test_unparseable_resolution_behaves_like_custom():
|
||
_, w, h = ModelResolutionPickerNode().create("SDXL", "stale text", 1024, 768, 1)
|
||
assert (w, h) == (1024, 768)
|
||
|
||
|
||
def test_validate_inputs_accepts_unlisted_resolution():
|
||
assert ModelResolutionPickerNode.VALIDATE_INPUTS("SDXL", "9999×1 - 1:1 (0.0MP)")
|
||
|
||
|
||
def test_execution_rounds_to_family_step_and_marks_latent():
|
||
latent, w, h = ModelResolutionPickerNode().create(
|
||
"HunyuanImage 2.1", "custom", 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", "custom", 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", "custom", 1003, 1003, 1
|
||
)
|
||
assert (w, h) == (1000, 1000) # fallback step 8
|
||
assert "Renamed Family" in caplog.text
|
||
|
||
|
||
def test_validate_inputs_accepts_any_label():
|
||
assert ModelResolutionPickerNode.VALIDATE_INPUTS("Renamed Family", "stale") is True
|
||
|
||
|
||
def test_alias_label_resolves_to_its_family():
|
||
sdxl = families.resolve_family("SDXL")
|
||
patched = sdxl._replace(aliases=("SDXL 1.0",))
|
||
others = tuple(f for f in families.FAMILIES if f.label != "SDXL")
|
||
with mock.patch.object(families, "FAMILIES", (patched,) + others):
|
||
assert families.resolve_family("SDXL 1.0").label == "SDXL"
|
||
latent, w, h = ModelResolutionPickerNode().create(
|
||
"SDXL 1.0", "custom", 1216, 832, 1
|
||
)
|
||
assert (w, h) == (1216, 832)
|
||
assert latent["samples"].shape[-2:] == (104, 152)
|
||
|
||
|
||
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"
|
||
)
|