216 lines
8.3 KiB
Python
216 lines
8.3 KiB
Python
"""Tests for ModelsSamplingNode."""
|
||
|
||
import json
|
||
import struct
|
||
|
||
import pytest
|
||
|
||
from mp_nodes.models_sampling import ModelsSamplingNode
|
||
|
||
|
||
def run(mode, **kwargs):
|
||
return ModelsSamplingNode().process(mode=mode, theme="(use pack default)", **kwargs)
|
||
|
||
|
||
def _make_safetensors(tmp_path, name, metadata=None, tensor_keys=None):
|
||
"""Create a minimal valid safetensors file with the given metadata block."""
|
||
header = {}
|
||
if metadata is not None:
|
||
header["__metadata__"] = metadata
|
||
for k in tensor_keys or []:
|
||
header[k] = {"dtype": "F32", "shape": [1], "data_offsets": [0, 4]}
|
||
header_bytes = json.dumps(header).encode("utf-8")
|
||
# Pad to 8-byte alignment
|
||
pad = (8 - len(header_bytes) % 8) % 8
|
||
header_bytes += b" " * pad
|
||
body = b"\x00\x00\x00\x00" * len(tensor_keys or [])
|
||
path = tmp_path / name
|
||
with open(path, "wb") as f:
|
||
f.write(struct.pack("<Q", len(header_bytes)))
|
||
f.write(header_bytes)
|
||
f.write(body)
|
||
return str(path)
|
||
|
||
|
||
class TestModelIntel:
|
||
def test_safetensors_metadata(self, tmp_path):
|
||
path = _make_safetensors(
|
||
tmp_path, "model.safetensors",
|
||
metadata={"author": "test", "trigger_words": "cat dog"},
|
||
tensor_keys=["w0", "w1", "w2"],
|
||
)
|
||
info = run("safetensors_metadata", path=path)[8]
|
||
assert info["metadata"]["author"] == "test"
|
||
assert info["tensor_count"] == 3
|
||
|
||
def test_safetensors_metadata_missing_file(self):
|
||
with pytest.raises(RuntimeError, match="not found"):
|
||
run("safetensors_metadata", path="/no/such/file.safetensors")
|
||
|
||
def test_model_fingerprint_is_deterministic(self, tmp_path):
|
||
path = tmp_path / "x.safetensors"
|
||
path.write_bytes(b"hello world" * 1000)
|
||
a = run("model_fingerprint", path=str(path))[3]
|
||
b = run("model_fingerprint", path=str(path))[3]
|
||
assert a == b
|
||
assert len(a) == 64 # sha256 hex
|
||
|
||
def test_list_installed_models(self, tmp_path):
|
||
for name in ("a.safetensors", "b.safetensors", "c.txt"):
|
||
(tmp_path / name).write_text("x")
|
||
info = run("list_installed_models", folder=str(tmp_path), extensions=".safetensors")[8]
|
||
assert info["count"] == 2
|
||
assert "a.safetensors" in info["files"]
|
||
assert "c.txt" not in info["files"]
|
||
|
||
def test_list_installed_models_missing_folder(self):
|
||
info = run("list_installed_models", folder="/nope", extensions=".safetensors")[8]
|
||
assert info["count"] == 0
|
||
|
||
def test_lora_extract_triggers(self, tmp_path):
|
||
path = _make_safetensors(
|
||
tmp_path, "lora.safetensors",
|
||
metadata={"trigger_words": "cat, dog, bird"},
|
||
)
|
||
triggers = run("lora_extract_triggers", path=path)[3]
|
||
assert "cat" in triggers
|
||
assert "dog" in triggers
|
||
|
||
|
||
class TestResolution:
|
||
def test_aspect_ratio_16_9(self):
|
||
out = run("aspect_ratio_pick", preset="16:9", long_side=1024)
|
||
assert out[5] == 1024 # width
|
||
assert out[6] == 576 # height
|
||
|
||
def test_aspect_ratio_portrait(self):
|
||
out = run("aspect_ratio_pick", preset="9:16 (portrait)", long_side=1024)
|
||
assert out[5] == 576 # width
|
||
assert out[6] == 1024 # height
|
||
|
||
def test_aspect_ratio_square(self):
|
||
out = run("aspect_ratio_pick", preset="1:1 (square)", long_side=512)
|
||
assert out[5] == 512 and out[6] == 512
|
||
|
||
def test_sdxl_bucket_index_zero(self):
|
||
out = run("sdxl_bucket_pick", index=0)
|
||
assert (out[5], out[6]) == (1024, 1024)
|
||
|
||
def test_sdxl_bucket_index_wraps(self):
|
||
out = run("sdxl_bucket_pick", index=11) # 11 % 9 = 2
|
||
assert (out[5], out[6]) == (896, 1152)
|
||
|
||
def test_snap_to_multiple(self):
|
||
assert run("snap_to_multiple", value=1023, multiple=64)[5] == 1024
|
||
assert run("snap_to_multiple", value=999, multiple=8)[5] == 1000
|
||
assert run("snap_to_multiple", value=1024, multiple=64)[5] == 1024
|
||
|
||
def test_megapixel_calculator_1mp_16_9(self):
|
||
out = run("megapixel_calculator", target_megapixels=1.0, ratio_w=16, ratio_h=9, snap_multiple=8)
|
||
# ~ 1MP at 16:9 = 1336 × 752 (snapped)
|
||
actual_mp = out[8]["actual_megapixels"]
|
||
assert 0.9 <= actual_mp <= 1.1
|
||
|
||
|
||
class TestSeed:
|
||
def test_seed_fixed_returns_input(self):
|
||
assert run("seed_cycle", seed=42, cycle_mode="fixed", string_input="")[5] == 42
|
||
|
||
def test_seed_increment(self):
|
||
assert run("seed_cycle", seed=99, cycle_mode="increment", string_input="")[5] == 100
|
||
|
||
def test_seed_decrement(self):
|
||
assert run("seed_cycle", seed=99, cycle_mode="decrement", string_input="")[5] == 98
|
||
|
||
def test_seed_from_string_deterministic(self):
|
||
a = run("seed_cycle", seed=0, cycle_mode="from_string", string_input="hello")[5]
|
||
b = run("seed_cycle", seed=0, cycle_mode="from_string", string_input="hello")[5]
|
||
assert a == b
|
||
assert 0 <= a <= 0xFFFFFFFF
|
||
|
||
def test_seed_from_string_changes_with_input(self):
|
||
a = run("seed_cycle", seed=0, cycle_mode="from_string", string_input="hello")[5]
|
||
b = run("seed_cycle", seed=0, cycle_mode="from_string", string_input="world")[5]
|
||
assert a != b
|
||
|
||
def test_multi_seed_batch(self):
|
||
out = run("multi_seed_batch", base_seed=100, count=4)[8]
|
||
assert out["seeds"] == [100, 101, 102, 103]
|
||
assert out["count"] == 4
|
||
|
||
def test_seed_history_records(self):
|
||
run("seed_cycle", seed=1, cycle_mode="fixed")
|
||
run("seed_cycle", seed=2, cycle_mode="fixed")
|
||
info = run("seed_history")[8]
|
||
assert info["count"] >= 2
|
||
assert info["history"][-1] == 2
|
||
|
||
|
||
class TestSampling:
|
||
def test_sampler_pick(self):
|
||
assert run("sampler_pick", sampler="euler")[4] == "euler"
|
||
|
||
def test_scheduler_pick(self):
|
||
assert run("scheduler_pick", scheduler="karras")[3] == "karras"
|
||
|
||
def test_sampler_params_bundle(self):
|
||
out = run("sampler_params_bundle", sampler="dpmpp_2m", scheduler="karras",
|
||
steps=30, cfg=8.0, denoise=0.8)
|
||
assert out[7] == 8.0 # cfg as FLOAT
|
||
bundle = out[8]
|
||
assert bundle["sampler"] == "dpmpp_2m"
|
||
assert bundle["scheduler"] == "karras"
|
||
assert bundle["steps"] == 30
|
||
assert bundle["denoise"] == 0.8
|
||
|
||
|
||
class TestOpaqueOutputs:
|
||
"""ModelsSamplingNode has MODEL/CLIP/VAE in RETURN_TYPES. None of our ops fill them.
|
||
Verify they pad to None (not raise)."""
|
||
|
||
def test_unfilled_opaque_outputs_are_none(self):
|
||
out = run("sampler_pick", sampler="euler")
|
||
# MODEL/CLIP/VAE are at indices 0, 1, 2
|
||
assert out[0] is None
|
||
assert out[1] is None
|
||
assert out[2] is None
|
||
|
||
|
||
class TestModelResolution:
|
||
def test_qwen_image_16_9(self):
|
||
out = run("model_resolution_pick", model_preset="Qwen-Image — 16:9 1664x928")
|
||
assert (out[5], out[6]) == (1664, 928)
|
||
assert out[8]["preset"] == "Qwen-Image — 16:9 1664x928"
|
||
|
||
def test_sd15_native(self):
|
||
out = run("model_resolution_pick", model_preset="SD 1.5 — 512x512")
|
||
assert (out[5], out[6]) == (512, 512)
|
||
|
||
def test_wan_5b_portrait(self):
|
||
out = run("model_resolution_pick", model_preset="Wan 2.2 TI2V-5B — 720P portrait 704x1280")
|
||
assert (out[5], out[6]) == (704, 1280)
|
||
|
||
def test_widget_appended_last_and_does_not_shadow_aspect_ratio_preset(self):
|
||
optional = ModelsSamplingNode.INPUT_TYPES()["optional"]
|
||
assert list(optional)[-1] == "model_preset"
|
||
assert "16:9" in optional["preset"][0]
|
||
|
||
def test_every_preset_dimension_is_positive(self):
|
||
from mp_nodes.models_sampling.operations.model_resolution import _MODEL_RESOLUTIONS
|
||
assert all(w > 0 and h > 0 for w, h in _MODEL_RESOLUTIONS.values())
|
||
|
||
|
||
class TestSamplerLists:
|
||
def test_samplers_include_current_core_names(self):
|
||
from mp_nodes.models_sampling.operations.seed_sampling import _SAMPLERS
|
||
for name in ("euler_cfg_pp", "res_multistep", "er_sde", "sa_solver", "cfgpp_ud10_ab", "uni_pc_bh2"):
|
||
assert name in _SAMPLERS
|
||
assert len(_SAMPLERS) == len(set(_SAMPLERS))
|
||
|
||
def test_schedulers_include_current_core_names(self):
|
||
from mp_nodes.models_sampling.operations.seed_sampling import _SCHEDULERS
|
||
assert "linear_quadratic" in _SCHEDULERS and "kl_optimal" in _SCHEDULERS
|
||
|
||
def test_new_scheduler_passes_through(self):
|
||
assert run("scheduler_pick", scheduler="kl_optimal")[3] == "kl_optimal"
|