264 lines
11 KiB
Python
264 lines
11 KiB
Python
import json
|
|
import types
|
|
import unittest
|
|
from unittest import mock
|
|
|
|
try:
|
|
from .test_repair_helpers import FakeLatentTensor, load_plugin_package
|
|
except ImportError: # pragma: no cover - discovery mode imports from top level
|
|
from test_repair_helpers import FakeLatentTensor, load_plugin_package
|
|
|
|
|
|
class TestRepairSampler(unittest.TestCase):
|
|
def test_sampler_schema_exposes_repair_inputs(self):
|
|
plugin = load_plugin_package()
|
|
node_cls = plugin.NODE_CLASS_MAPPINGS["LLSSimpleKSampler"]
|
|
schema = node_cls.INPUT_TYPES()
|
|
|
|
self.assertEqual(schema["required"]["denoise_mode"][0], ["manual", "auto_from_repair"])
|
|
self.assertEqual(schema["required"]["adapter_mode"][0], ["auto", "sd_classic", "flux", "sd3", "qwen", "zimage"])
|
|
self.assertEqual(schema["optional"]["repair_info"][0], "LLS_REPAIR_INFO")
|
|
self.assertEqual(schema["optional"]["guidance_stack"][0], "LLS_GUIDANCE_STACK")
|
|
self.assertEqual(schema["optional"]["model_info"][0], "STRING")
|
|
|
|
def test_sampler_keeps_manual_denoise_without_repair_info(self):
|
|
load_plugin_package()
|
|
from lls_node_test_repair.sampling import nodes as sampling_nodes
|
|
|
|
node = sampling_nodes.LLSSimpleKSampler()
|
|
latent = {"samples": FakeLatentTensor((1, 4, 64, 64)), "source": "empty_latent"}
|
|
with mock.patch.object(sampling_nodes, "comfy_sample", object()), mock.patch.object(
|
|
sampling_nodes,
|
|
"comfy_samplers",
|
|
types.SimpleNamespace(KSampler=types.SimpleNamespace(SAMPLERS=["euler"], SCHEDULERS=["karras"])),
|
|
), mock.patch.object(
|
|
sampling_nodes,
|
|
"_common_ksampler",
|
|
return_value={"samples": "done"},
|
|
) as common:
|
|
result_latent, sample_info = node.sample(
|
|
model=types.SimpleNamespace(_lls_family="SD1.5"),
|
|
positive="positive",
|
|
negative="negative",
|
|
latent_image=latent,
|
|
quality_preset="Manual",
|
|
seed=7,
|
|
steps=20,
|
|
cfg=7.0,
|
|
sampler_name="euler",
|
|
scheduler="karras",
|
|
denoise=0.33,
|
|
denoise_mode="manual",
|
|
adapter_mode="auto",
|
|
flux_guidance=3.5,
|
|
model_family="Auto",
|
|
repair_info=None,
|
|
guidance_stack=None,
|
|
model_info=None,
|
|
)
|
|
|
|
self.assertEqual(result_latent["samples"], "done")
|
|
self.assertEqual(common.call_args.kwargs["denoise"], 0.33)
|
|
payload = json.loads(sample_info)
|
|
self.assertFalse(payload["repair_mode"])
|
|
self.assertFalse(payload["guidance_used"])
|
|
|
|
def test_sampler_uses_repair_denoise_when_requested(self):
|
|
load_plugin_package()
|
|
from lls_node_test_repair.sampling import nodes as sampling_nodes
|
|
|
|
node = sampling_nodes.LLSSimpleKSampler()
|
|
latent = {"samples": FakeLatentTensor((1, 4, 64, 64)), "source": "repair_prepare_region"}
|
|
with mock.patch.object(sampling_nodes, "comfy_sample", object()), mock.patch.object(
|
|
sampling_nodes,
|
|
"comfy_samplers",
|
|
types.SimpleNamespace(KSampler=types.SimpleNamespace(SAMPLERS=["euler"], SCHEDULERS=["karras"])),
|
|
), mock.patch.object(
|
|
sampling_nodes,
|
|
"_common_ksampler",
|
|
return_value={"samples": "done"},
|
|
) as common:
|
|
_result_latent, sample_info = node.sample(
|
|
model=types.SimpleNamespace(_lls_family="SDXL"),
|
|
positive="positive",
|
|
negative="negative",
|
|
latent_image=latent,
|
|
quality_preset="Manual",
|
|
seed=7,
|
|
steps=20,
|
|
cfg=7.0,
|
|
sampler_name="euler",
|
|
scheduler="karras",
|
|
denoise=0.33,
|
|
denoise_mode="auto_from_repair",
|
|
adapter_mode="auto",
|
|
flux_guidance=3.5,
|
|
model_family="Auto",
|
|
repair_info={
|
|
"repair_scope": "region",
|
|
"repair_kernel": "latent_mask",
|
|
"recommended_denoise": 0.61,
|
|
"model_family": "SDXL",
|
|
"model_role": "normal",
|
|
},
|
|
guidance_stack={"kind": "placeholder"},
|
|
model_info=None,
|
|
)
|
|
|
|
self.assertEqual(common.call_args.kwargs["denoise"], 0.61)
|
|
payload = json.loads(sample_info)
|
|
self.assertTrue(payload["repair_mode"])
|
|
self.assertTrue(payload["guidance_used"])
|
|
self.assertEqual(payload["repair_scope"], "region")
|
|
self.assertEqual(payload["repair_kernel"], "latent_mask")
|
|
|
|
def test_sampler_patches_native_flux_repair_model_before_sampling(self):
|
|
load_plugin_package()
|
|
from lls_node_test_repair.sampling import nodes as sampling_nodes
|
|
|
|
node = sampling_nodes.LLSSimpleKSampler()
|
|
latent = {"samples": FakeLatentTensor((1, 16, 64, 64)), "source": "repair_prepare_region"}
|
|
model = object()
|
|
patched_model = object()
|
|
with mock.patch.object(sampling_nodes, "comfy_sample", object()), mock.patch.object(
|
|
sampling_nodes,
|
|
"comfy_samplers",
|
|
types.SimpleNamespace(KSampler=types.SimpleNamespace(SAMPLERS=["euler"], SCHEDULERS=["simple"])),
|
|
), mock.patch.object(
|
|
sampling_nodes,
|
|
"_apply_repair_model_patch",
|
|
return_value=patched_model,
|
|
) as patch_model, mock.patch.object(
|
|
sampling_nodes,
|
|
"_common_ksampler",
|
|
return_value={"samples": "done"},
|
|
) as common:
|
|
_result_latent, sample_info = node.sample(
|
|
model=model,
|
|
positive="positive",
|
|
negative="negative",
|
|
latent_image=latent,
|
|
quality_preset="Manual",
|
|
seed=7,
|
|
steps=20,
|
|
cfg=1.0,
|
|
sampler_name="euler",
|
|
scheduler="simple",
|
|
denoise=0.33,
|
|
denoise_mode="auto_from_repair",
|
|
adapter_mode="auto",
|
|
flux_guidance=3.5,
|
|
model_family="Auto",
|
|
repair_info={
|
|
"repair_scope": "region",
|
|
"repair_kernel": "native_fill",
|
|
"recommended_denoise": 0.61,
|
|
"model_family": "FLUX_DEV",
|
|
"model_role": "fill",
|
|
"backend_name": "flux",
|
|
"execution_path": "native_repair",
|
|
"model_patch": "differential_diffusion",
|
|
"model_patch_strength": 1.0,
|
|
},
|
|
guidance_stack=None,
|
|
model_info=None,
|
|
)
|
|
|
|
patch_model.assert_called_once()
|
|
self.assertIs(common.call_args.kwargs["model"], patched_model)
|
|
payload = json.loads(sample_info)
|
|
self.assertEqual(payload["repair_kernel"], "native_fill")
|
|
self.assertEqual(payload["repair_backend"], "flux")
|
|
self.assertEqual(payload["repair_execution_path"], "native_repair")
|
|
|
|
def test_sampler_recovers_shifted_values_from_broken_widget_order(self):
|
|
load_plugin_package()
|
|
from lls_node_test_repair.sampling import nodes as sampling_nodes
|
|
|
|
node = sampling_nodes.LLSSimpleKSampler()
|
|
latent = {"samples": FakeLatentTensor((1, 4, 64, 64)), "source": "repair_prepare_region"}
|
|
with mock.patch.object(sampling_nodes, "comfy_sample", object()), mock.patch.object(
|
|
sampling_nodes,
|
|
"comfy_samplers",
|
|
types.SimpleNamespace(KSampler=types.SimpleNamespace(SAMPLERS=["euler"], SCHEDULERS=["karras"])),
|
|
), mock.patch.object(
|
|
sampling_nodes,
|
|
"_common_ksampler",
|
|
return_value={"samples": "done"},
|
|
) as common:
|
|
_result_latent, sample_info = node.sample(
|
|
model=types.SimpleNamespace(_lls_family="FLUX_DEV"),
|
|
positive="positive",
|
|
negative="negative",
|
|
latent_image=latent,
|
|
quality_preset="Manual",
|
|
seed=7,
|
|
steps=20,
|
|
cfg=7.0,
|
|
sampler_name="euler",
|
|
scheduler="karras",
|
|
denoise=0.33,
|
|
denoise_mode=4.2,
|
|
adapter_mode="Auto",
|
|
flux_guidance="auto_from_repair",
|
|
model_family="auto",
|
|
repair_info={
|
|
"repair_scope": "region",
|
|
"repair_kernel": "latent_mask",
|
|
"recommended_denoise": 0.61,
|
|
"model_family": "FLUX_DEV",
|
|
"model_role": "normal",
|
|
},
|
|
guidance_stack=None,
|
|
model_info=None,
|
|
)
|
|
|
|
self.assertEqual(common.call_args.kwargs["denoise"], 0.61)
|
|
payload = json.loads(sample_info)
|
|
self.assertEqual(payload["family"], "FLUX_DEV")
|
|
self.assertEqual(payload["guidance"], 4.2)
|
|
self.assertTrue(payload["repair_mode"])
|
|
|
|
def test_sampler_rejects_unsupported_qwen_adapter(self):
|
|
load_plugin_package()
|
|
from lls_node_test_repair.sampling import nodes as sampling_nodes
|
|
|
|
node = sampling_nodes.LLSSimpleKSampler()
|
|
latent = {"samples": FakeLatentTensor((1, 4, 64, 64)), "source": "repair_prepare_region"}
|
|
with mock.patch.object(sampling_nodes, "comfy_sample", object()), mock.patch.object(
|
|
sampling_nodes,
|
|
"comfy_samplers",
|
|
types.SimpleNamespace(KSampler=types.SimpleNamespace(SAMPLERS=["euler"], SCHEDULERS=["karras"])),
|
|
):
|
|
with self.assertRaisesRegex(RuntimeError, "QWEN"):
|
|
node.sample(
|
|
model=types.SimpleNamespace(_lls_family="SD1.5"),
|
|
positive="positive",
|
|
negative="negative",
|
|
latent_image=latent,
|
|
quality_preset="Manual",
|
|
seed=7,
|
|
steps=20,
|
|
cfg=7.0,
|
|
sampler_name="euler",
|
|
scheduler="karras",
|
|
denoise=0.33,
|
|
denoise_mode="auto_from_repair",
|
|
adapter_mode="auto",
|
|
flux_guidance=3.5,
|
|
model_family="Auto",
|
|
repair_info={
|
|
"repair_scope": "region",
|
|
"repair_kernel": "vae_inpaint",
|
|
"recommended_denoise": 0.55,
|
|
"model_family": "QWEN",
|
|
"model_role": "unknown",
|
|
},
|
|
guidance_stack=None,
|
|
model_info=None,
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|