213 lines
7.2 KiB
Python
213 lines
7.2 KiB
Python
from smart_metadata_reader.graph import GraphIndex
|
|
|
|
|
|
def select_sampler(prompt):
|
|
from smart_metadata_reader.sampler_selector import select_final_sampler
|
|
|
|
return select_final_sampler(GraphIndex(prompt=prompt, workflow=None))
|
|
|
|
|
|
def test_selects_ksampler_upstream_of_save_image_via_vae_decode():
|
|
prompt = {
|
|
"1": {"class_type": "KSampler", "inputs": {"positive": ["9", 0]}},
|
|
"2": {"class_type": "VAEDecode", "inputs": {"samples": ["1", 0]}},
|
|
"3": {"class_type": "SaveImage", "inputs": {"images": ["2", 0]}},
|
|
"9": {"class_type": "CLIPTextEncode", "inputs": {"text": "positive"}},
|
|
}
|
|
|
|
selection = select_sampler(prompt)
|
|
|
|
assert selection.sampler.node_id == "1"
|
|
assert selection.confidence == 1.0
|
|
assert "SaveImage 3.images <- VAEDecode 2.samples <- KSampler 1" in selection.debug_trace
|
|
|
|
|
|
def test_selects_ksampler_upstream_of_preview_image_via_vae_decode():
|
|
prompt = {
|
|
"1": {"class_type": "KSampler", "inputs": {"positive": ["9", 0]}},
|
|
"2": {"class_type": "VAEDecode", "inputs": {"samples": ["1", 0]}},
|
|
"3": {"class_type": "PreviewImage", "inputs": {"images": ["2", 0]}},
|
|
"9": {"class_type": "CLIPTextEncode", "inputs": {"text": "positive"}},
|
|
}
|
|
|
|
selection = select_sampler(prompt)
|
|
|
|
assert selection.sampler.node_id == "1"
|
|
assert "PreviewImage 3.images" in selection.debug_trace
|
|
|
|
|
|
def test_selects_sampler_through_image_passthrough_before_save_image():
|
|
prompt = {
|
|
"1": {"class_type": "KSampler", "inputs": {"positive": ["9", 0]}},
|
|
"2": {"class_type": "VAEDecode", "inputs": {"samples": ["1", 0]}},
|
|
"3": {"class_type": "ImagePassThrough", "inputs": {"image": ["2", 0]}},
|
|
"4": {"class_type": "SaveImage", "inputs": {"images": ["3", 0]}},
|
|
"9": {"class_type": "CLIPTextEncode", "inputs": {"text": "positive"}},
|
|
}
|
|
|
|
selection = select_sampler(prompt)
|
|
|
|
assert selection.sampler.node_id == "1"
|
|
assert "ImagePassThrough 3.image" in selection.debug_trace
|
|
|
|
|
|
def test_multiple_samplers_selects_one_connected_to_final_save_image():
|
|
prompt = {
|
|
"1": {"class_type": "KSampler", "inputs": {"positive": ["9", 0]}},
|
|
"2": {"class_type": "VAEDecode", "inputs": {"samples": ["1", 0]}},
|
|
"3": {"class_type": "SaveImage", "inputs": {"images": ["2", 0]}},
|
|
"8": {"class_type": "KSampler", "inputs": {"positive": ["9", 0]}},
|
|
"9": {"class_type": "CLIPTextEncode", "inputs": {"text": "positive"}},
|
|
}
|
|
|
|
selection = select_sampler(prompt)
|
|
|
|
assert selection.sampler.node_id == "1"
|
|
assert "Selected KSampler 1 because it is upstream of final SaveImage" in selection.debug_trace
|
|
|
|
|
|
def test_fallback_selects_sampler_with_conditioning_inputs_when_no_output_node_exists():
|
|
prompt = {
|
|
"5": {"class_type": "PrimitiveString", "inputs": {"text": "not a sampler"}},
|
|
"6": {
|
|
"class_type": "KSampler",
|
|
"inputs": {"positive": ["9", 0], "negative": ["10", 0]},
|
|
},
|
|
"9": {"class_type": "CLIPTextEncode", "inputs": {"text": "positive"}},
|
|
"10": {"class_type": "CLIPTextEncode", "inputs": {"text": "negative"}},
|
|
}
|
|
|
|
selection = select_sampler(prompt)
|
|
|
|
assert selection.sampler.node_id == "6"
|
|
assert selection.confidence < 1.0
|
|
assert "sampler selected by fallback heuristic" in selection.debug_trace
|
|
|
|
|
|
def test_extracts_ksampler_settings():
|
|
from smart_metadata_reader.adapters.samplers import extract_sampler_settings
|
|
from smart_metadata_reader.models import NodeRecord
|
|
|
|
sampler = NodeRecord(
|
|
node_id="1",
|
|
class_type="KSampler",
|
|
inputs={
|
|
"seed": 123,
|
|
"steps": 30,
|
|
"cfg": 7.5,
|
|
"sampler_name": "euler",
|
|
"scheduler": "normal",
|
|
"denoise": 0.85,
|
|
},
|
|
)
|
|
|
|
assert extract_sampler_settings(sampler) == {
|
|
"seed": 123,
|
|
"steps": 30,
|
|
"cfg": 7.5,
|
|
"sampler_name": "euler",
|
|
"scheduler": "normal",
|
|
"denoise": 0.85,
|
|
}
|
|
|
|
|
|
def test_extracts_ksampler_advanced_and_custom_field_variants():
|
|
from smart_metadata_reader.adapters.samplers import extract_sampler_settings
|
|
from smart_metadata_reader.models import NodeRecord
|
|
|
|
advanced = NodeRecord(
|
|
node_id="2",
|
|
class_type="KSamplerAdvanced",
|
|
inputs={
|
|
"noise_seed": 456,
|
|
"steps": 22,
|
|
"cfg_scale": 6.0,
|
|
"sampler": "dpmpp_2m",
|
|
"schedule": "karras",
|
|
},
|
|
)
|
|
custom = NodeRecord(
|
|
node_id="3",
|
|
class_type="SamplerCustom",
|
|
inputs={
|
|
"noise_seed": 789,
|
|
"steps": 18,
|
|
"cfg": 5.5,
|
|
"sampler_name": "custom_sampler",
|
|
"scheduler": "sgm_uniform",
|
|
},
|
|
)
|
|
|
|
assert extract_sampler_settings(advanced) == {
|
|
"seed": 456,
|
|
"steps": 22,
|
|
"cfg": 6.0,
|
|
"sampler_name": "dpmpp_2m",
|
|
"scheduler": "karras",
|
|
}
|
|
assert extract_sampler_settings(custom) == {
|
|
"seed": 789,
|
|
"steps": 18,
|
|
"cfg": 5.5,
|
|
"sampler_name": "custom_sampler",
|
|
"scheduler": "sgm_uniform",
|
|
}
|
|
|
|
|
|
def test_selects_ultimate_sd_upscale_sampler_like_node_upstream_of_save_image():
|
|
prompt = {
|
|
"1": {"class_type": "CLIPTextEncode", "inputs": {"text": "positive"}},
|
|
"2": {"class_type": "CLIPTextEncode", "inputs": {"text": "negative"}},
|
|
"7": {
|
|
"class_type": "UltimateSDUpscale",
|
|
"inputs": {
|
|
"positive": ["1", 0],
|
|
"negative": ["2", 0],
|
|
"model": ["9", 0],
|
|
"vae": ["10", 0],
|
|
"seed": 123,
|
|
"steps": 12,
|
|
"cfg": 2.5,
|
|
"sampler_name": "euler",
|
|
"scheduler_name": "normal",
|
|
"tile_width": 768,
|
|
"seam_fix_mode": "band pass",
|
|
},
|
|
},
|
|
"8": {"class_type": "SaveImage", "inputs": {"images": ["7", 0]}},
|
|
"9": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": "model.safetensors"}},
|
|
"10": {"class_type": "VAELoader", "inputs": {"vae_name": "vae.safetensors"}},
|
|
}
|
|
|
|
selection = select_sampler(prompt)
|
|
|
|
assert selection.sampler.node_id == "7"
|
|
assert selection.sampler.class_type == "UltimateSDUpscale"
|
|
assert selection.selected_by == "output_chain"
|
|
assert "SaveImage 8.images <- UltimateSDUpscale 7" in selection.debug_trace
|
|
|
|
|
|
def test_unsupported_output_chain_error_includes_diagnostics():
|
|
from smart_metadata_reader.sampler_selector import select_final_sampler
|
|
|
|
prompt = {
|
|
"1": {
|
|
"class_type": "UnsupportedGenerator",
|
|
"inputs": {"image": ["3", 0], "tile_width": 768},
|
|
},
|
|
"2": {"class_type": "SaveImage", "inputs": {"images": ["1", 0]}},
|
|
"3": {"class_type": "LoadImage", "inputs": {"image": "source.png"}},
|
|
}
|
|
|
|
try:
|
|
select_final_sampler(GraphIndex(prompt=prompt, workflow=None))
|
|
except ValueError as exc:
|
|
message = str(exc)
|
|
else:
|
|
raise AssertionError("Expected unsupported output chain to fail")
|
|
|
|
assert "no supported sampler candidate found" in message
|
|
assert "final output nodes found: SaveImage 2" in message
|
|
assert "SaveImage 2.images <- UnsupportedGenerator 1" in message
|
|
assert "input keys: image, tile_width" in message
|