1098 lines
42 KiB
Python
1098 lines
42 KiB
Python
import importlib.util
|
|
import json
|
|
import pathlib
|
|
import sys
|
|
import types
|
|
import unittest
|
|
from unittest import mock
|
|
|
|
|
|
ROOT = pathlib.Path(__file__).resolve().parents[1]
|
|
|
|
|
|
def load_plugin_package():
|
|
spec = importlib.util.spec_from_file_location(
|
|
"lls_node_test_refactor",
|
|
ROOT / "__init__.py",
|
|
submodule_search_locations=[str(ROOT)],
|
|
)
|
|
module = importlib.util.module_from_spec(spec)
|
|
sys.modules["lls_node_test_refactor"] = module
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
class TaggedValue:
|
|
def __init__(self, label):
|
|
self.label = label
|
|
|
|
def __repr__(self):
|
|
return self.label
|
|
|
|
|
|
class FolderPathsStub:
|
|
def __init__(self):
|
|
self._files = {
|
|
"checkpoints": [
|
|
"sd15.safetensors",
|
|
"sdxl_turbo.safetensors",
|
|
"flux-dev-checkpoint.safetensors",
|
|
],
|
|
"diffusion_models": [
|
|
"flux1-schnell.safetensors",
|
|
],
|
|
"vae": [
|
|
"ae.safetensors",
|
|
],
|
|
"text_encoders": [
|
|
"clip_l.safetensors",
|
|
"clip_g.safetensors",
|
|
"t5xxl_fp16.safetensors",
|
|
],
|
|
"upscale_models": [],
|
|
}
|
|
|
|
def get_filename_list(self, category):
|
|
return list(self._files.get(category, []))
|
|
|
|
def get_full_path_or_raise(self, category, name):
|
|
if name in self._files.get(category, []):
|
|
return f"/fake/{category}/{name}"
|
|
raise FileNotFoundError(name)
|
|
|
|
def get_full_path(self, category, name):
|
|
if name in self._files.get(category, []):
|
|
return f"/fake/{category}/{name}"
|
|
return None
|
|
|
|
def get_folder_paths(self, category):
|
|
return [f"/fake/{category}"]
|
|
|
|
|
|
class ComfySDStub:
|
|
def load_checkpoint_guess_config(
|
|
self,
|
|
ckpt_path,
|
|
output_vae=True,
|
|
output_clip=True,
|
|
embedding_directory=None,
|
|
):
|
|
if ckpt_path.endswith("sd15.safetensors"):
|
|
return (TaggedValue("MODEL::SD15"), TaggedValue("CLIP::SD15"), TaggedValue("VAE::SD15"))
|
|
if ckpt_path.endswith("sdxl_turbo.safetensors"):
|
|
return (TaggedValue("MODEL::SDXL_TURBO"), None, None)
|
|
if ckpt_path.endswith("flux-dev-checkpoint.safetensors"):
|
|
return (TaggedValue("MODEL::FLUX_DEV"), None, None)
|
|
raise AssertionError(f"unexpected checkpoint path: {ckpt_path}")
|
|
|
|
def load_diffusion_model(self, unet_path, model_options=None, disable_dynamic=False):
|
|
return TaggedValue(f"MODEL::{pathlib.Path(unet_path).name}")
|
|
|
|
|
|
class CoreNodesStub:
|
|
save_calls = []
|
|
preview_calls = []
|
|
|
|
class VAELoader:
|
|
@staticmethod
|
|
def vae_list(_self):
|
|
return ["ae.safetensors"]
|
|
|
|
def load_vae(self, vae_name):
|
|
return (TaggedValue(f"VAE::{vae_name}"),)
|
|
|
|
class CLIPLoader:
|
|
def load_clip(self, clip_name, type="stable_diffusion", device="default"):
|
|
return (TaggedValue(f"CLIP::{type}::{clip_name}"),)
|
|
|
|
class DualCLIPLoader:
|
|
def load_clip(self, clip_name1, clip_name2, type, device="default"):
|
|
return (TaggedValue(f"CLIP::{type}::{clip_name1}+{clip_name2}"),)
|
|
|
|
class SaveImage:
|
|
def save_images(self, images, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None):
|
|
CoreNodesStub.save_calls.append(
|
|
{
|
|
"images": images,
|
|
"filename_prefix": filename_prefix,
|
|
"prompt": prompt,
|
|
"extra_pnginfo": extra_pnginfo,
|
|
}
|
|
)
|
|
return {
|
|
"ui": {
|
|
"images": [
|
|
{
|
|
"filename": f"{filename_prefix}.png",
|
|
"subfolder": "",
|
|
"type": "output",
|
|
"extra_pnginfo": extra_pnginfo,
|
|
"prompt": prompt,
|
|
}
|
|
]
|
|
}
|
|
}
|
|
|
|
class PreviewImage:
|
|
def save_images(self, images, prompt=None, extra_pnginfo=None):
|
|
CoreNodesStub.preview_calls.append(
|
|
{
|
|
"images": images,
|
|
"prompt": prompt,
|
|
"extra_pnginfo": extra_pnginfo,
|
|
}
|
|
)
|
|
return {
|
|
"ui": {
|
|
"images": [
|
|
{
|
|
"filename": "preview.png",
|
|
"subfolder": "",
|
|
"type": "temp",
|
|
"extra_pnginfo": extra_pnginfo,
|
|
"prompt": prompt,
|
|
}
|
|
]
|
|
}
|
|
}
|
|
|
|
|
|
class RecordingClip:
|
|
def __init__(self, token_keys, calls=None):
|
|
self.token_keys = tuple(token_keys)
|
|
self.calls = [] if calls is None else calls
|
|
self.layer = None
|
|
|
|
def clone(self):
|
|
cloned = RecordingClip(self.token_keys, self.calls)
|
|
cloned.layer = self.layer
|
|
return cloned
|
|
|
|
def clip_layer(self, layer):
|
|
self.layer = layer
|
|
|
|
def tokenize(self, text):
|
|
tokens = {"text": text}
|
|
for key in self.token_keys:
|
|
tokens[key] = [f"{key}::{text}"]
|
|
return tokens
|
|
|
|
def encode_from_tokens_scheduled(self, tokens, add_dict=None):
|
|
self.calls.append(
|
|
{
|
|
"tokens": dict(tokens),
|
|
"add_dict": dict(add_dict or {}),
|
|
"layer": self.layer,
|
|
}
|
|
)
|
|
return [[f"cond::{tokens.get('text', '')}", {"pooled_output": dict(add_dict or {})}]]
|
|
|
|
|
|
class FakeVAE:
|
|
def __init__(self, encode_channels=4, encode_ratio=8, vae_name="vae.safetensors"):
|
|
self.encode_channels = encode_channels
|
|
self.encode_ratio = encode_ratio
|
|
self.encoded_shapes = []
|
|
self._lls_vae_name = vae_name
|
|
|
|
def encode(self, pixels):
|
|
self.encoded_shapes.append(tuple(pixels.shape))
|
|
batch, height, width, _channels = pixels.shape
|
|
return FakeTensor(
|
|
(
|
|
batch,
|
|
self.encode_channels,
|
|
max(1, height // self.encode_ratio),
|
|
max(1, width // self.encode_ratio),
|
|
)
|
|
)
|
|
|
|
def decode(self, latent_tensor):
|
|
batch, _channels, height, width = latent_tensor.shape
|
|
return FakeTensor((batch, height, width, 3))
|
|
|
|
|
|
class FakeTensor:
|
|
def __init__(self, shape):
|
|
self.shape = tuple(shape)
|
|
self.dtype = "float32"
|
|
self.layout = "strided"
|
|
self.device = "cpu"
|
|
self.is_nested = False
|
|
|
|
def size(self):
|
|
return self.shape
|
|
|
|
def movedim(self, source, destination):
|
|
rank = len(self.shape)
|
|
source = source if source >= 0 else rank + source
|
|
destination = destination if destination >= 0 else rank + destination
|
|
order = list(range(rank))
|
|
moved_axis = order.pop(source)
|
|
order.insert(destination, moved_axis)
|
|
return FakeTensor(tuple(self.shape[index] for index in order))
|
|
|
|
def reshape(self, *shape):
|
|
return FakeTensor(shape)
|
|
|
|
|
|
class FakeTorch:
|
|
float32 = "float32"
|
|
|
|
@staticmethod
|
|
def zeros(shape, device=None, dtype=None, layout=None):
|
|
return FakeTensor(shape)
|
|
|
|
|
|
class FakeComfyUtils:
|
|
def __init__(self):
|
|
self.calls = []
|
|
|
|
def common_upscale(self, samples, width, height, upscale_method, crop):
|
|
self.calls.append(
|
|
{
|
|
"input_shape": tuple(samples.shape),
|
|
"width": width,
|
|
"height": height,
|
|
"upscale_method": upscale_method,
|
|
"crop": crop,
|
|
}
|
|
)
|
|
batch, channels, _source_height, _source_width = samples.shape
|
|
return FakeTensor((batch, channels, height, width))
|
|
|
|
|
|
class TestLoaderPromptRefactor(unittest.TestCase):
|
|
def test_task_nodes_are_not_registered_in_contextless_mode(self):
|
|
plugin = load_plugin_package()
|
|
|
|
self.assertNotIn("LLSTaskController", plugin.NODE_CLASS_MAPPINGS)
|
|
self.assertNotIn("LLSTaskInspector", plugin.NODE_CLASS_MAPPINGS)
|
|
|
|
def test_loader_schema_uses_standard_ports_only(self):
|
|
plugin = load_plugin_package()
|
|
node_cls = plugin.NODE_CLASS_MAPPINGS["LLSSimpleCheckpointLoader"]
|
|
required = node_cls.INPUT_TYPES()["required"]
|
|
|
|
self.assertEqual(
|
|
node_cls.RETURN_TYPES,
|
|
("MODEL", "CLIP", "VAE", "CLIP"),
|
|
)
|
|
self.assertEqual(
|
|
node_cls.RETURN_NAMES,
|
|
("model", "clip", "vae", "text_encoder"),
|
|
)
|
|
self.assertEqual(
|
|
tuple(required.keys()),
|
|
(
|
|
"ckpt_name",
|
|
"model_family",
|
|
"load_mode",
|
|
"vae_source",
|
|
"text_encoder_source",
|
|
"external_vae_name",
|
|
"external_text_encoder_1",
|
|
"external_text_encoder_2",
|
|
),
|
|
)
|
|
self.assertNotIn("optional", node_cls.INPUT_TYPES())
|
|
|
|
def test_universal_loader_schema_exposes_single_text_encoder_output(self):
|
|
plugin = load_plugin_package()
|
|
node_cls = plugin.NODE_CLASS_MAPPINGS["LLSUniversalModelLoader"]
|
|
required = node_cls.INPUT_TYPES()["required"]
|
|
|
|
self.assertEqual(
|
|
node_cls.RETURN_TYPES,
|
|
("MODEL", "CLIP", "VAE", "STRING"),
|
|
)
|
|
self.assertEqual(
|
|
node_cls.RETURN_NAMES,
|
|
("model", "text_encoder", "vae", "model_info"),
|
|
)
|
|
self.assertEqual(
|
|
tuple(required.keys()),
|
|
(
|
|
"model_name",
|
|
"model_family",
|
|
"load_mode",
|
|
"vae_source",
|
|
"text_encoder_source",
|
|
"text_encoder_1",
|
|
"text_encoder_2",
|
|
"vae_name",
|
|
),
|
|
)
|
|
self.assertEqual(required["model_family"][1]["default"], "Auto")
|
|
self.assertEqual(required["text_encoder_source"][1]["default"], "auto")
|
|
|
|
def test_loader_tags_sd15_resources_without_context_output(self):
|
|
load_plugin_package()
|
|
from lls_node_test_refactor.model_loader import nodes as loader_nodes
|
|
|
|
with mock.patch.object(loader_nodes, "folder_paths", FolderPathsStub()), mock.patch.object(
|
|
loader_nodes,
|
|
"comfy_sd",
|
|
ComfySDStub(),
|
|
), mock.patch.object(loader_nodes, "comfy_core_nodes", CoreNodesStub()):
|
|
node = loader_nodes.LLSSimpleCheckpointLoader()
|
|
model, clip, vae, text_encoder = node.load_checkpoint(
|
|
"sd15.safetensors",
|
|
"SD1.5",
|
|
"simple",
|
|
"auto",
|
|
"auto",
|
|
"(auto)",
|
|
"(auto)",
|
|
"(auto)",
|
|
)
|
|
|
|
self.assertEqual(model.label, "MODEL::SD15")
|
|
self.assertEqual(clip.label, "CLIP::SD15")
|
|
self.assertIs(text_encoder, clip)
|
|
self.assertEqual(vae.label, "VAE::SD15")
|
|
self.assertEqual(model._lls_family, "SD1.5")
|
|
self.assertEqual(model._lls_model_name, "sd15.safetensors")
|
|
self.assertEqual(clip._lls_family, "SD1.5")
|
|
self.assertEqual(clip._lls_text_encoder_type, "clip")
|
|
self.assertEqual(vae._lls_family, "SD1.5")
|
|
|
|
def test_universal_loader_loads_sd15_into_single_text_encoder_port(self):
|
|
load_plugin_package()
|
|
from lls_node_test_refactor.model_loader import nodes as loader_nodes
|
|
|
|
with mock.patch.object(loader_nodes, "folder_paths", FolderPathsStub()), mock.patch.object(
|
|
loader_nodes,
|
|
"comfy_sd",
|
|
ComfySDStub(),
|
|
), mock.patch.object(loader_nodes, "comfy_core_nodes", CoreNodesStub()):
|
|
node = loader_nodes.LLSUniversalModelLoader()
|
|
model, text_encoder, vae, model_info = node.load(
|
|
"sd15.safetensors",
|
|
"SD1.5",
|
|
"simple",
|
|
"auto",
|
|
"auto",
|
|
"(auto)",
|
|
"(auto)",
|
|
"(auto)",
|
|
)
|
|
|
|
payload = json.loads(model_info)
|
|
self.assertEqual(model.label, "MODEL::SD15")
|
|
self.assertEqual(text_encoder.label, "CLIP::SD15")
|
|
self.assertEqual(vae.label, "VAE::SD15")
|
|
self.assertEqual(payload["model_family"], "SD1.5")
|
|
self.assertEqual(payload["checkpoint_name"], "sd15.safetensors")
|
|
self.assertEqual(payload["text_encoder_source"], "embedded")
|
|
self.assertEqual(payload["vae_source"], "embedded")
|
|
|
|
def test_loader_tags_flux_resources_and_keeps_standard_text_encoder_alias(self):
|
|
load_plugin_package()
|
|
from lls_node_test_refactor.model_loader import nodes as loader_nodes
|
|
|
|
with mock.patch.object(loader_nodes, "folder_paths", FolderPathsStub()), mock.patch.object(
|
|
loader_nodes,
|
|
"comfy_sd",
|
|
ComfySDStub(),
|
|
), mock.patch.object(loader_nodes, "comfy_core_nodes", CoreNodesStub()):
|
|
node = loader_nodes.LLSSimpleCheckpointLoader()
|
|
model, clip, vae, text_encoder = node.load_checkpoint(
|
|
"diffusion_models/flux1-schnell.safetensors",
|
|
"FLUX_SCHNELL",
|
|
"advanced",
|
|
"external",
|
|
"external",
|
|
"ae.safetensors",
|
|
"clip_l.safetensors",
|
|
"t5xxl_fp16.safetensors",
|
|
)
|
|
|
|
self.assertEqual(model.label, "MODEL::flux1-schnell.safetensors")
|
|
self.assertEqual(clip.label, "CLIP::flux::clip_l.safetensors+t5xxl_fp16.safetensors")
|
|
self.assertIs(text_encoder, clip)
|
|
self.assertEqual(vae.label, "VAE::ae.safetensors")
|
|
self.assertEqual(model._lls_family, "FLUX_SCHNELL")
|
|
self.assertEqual(clip._lls_text_encoder_type, "flux_clip_l_t5xxl")
|
|
self.assertEqual(clip._lls_model_name, "diffusion_models/flux1-schnell.safetensors")
|
|
self.assertEqual(vae._lls_vae_name, "ae.safetensors")
|
|
|
|
def test_universal_loader_loads_flux_dual_encoders_into_single_text_encoder_port(self):
|
|
load_plugin_package()
|
|
from lls_node_test_refactor.model_loader import nodes as loader_nodes
|
|
|
|
with mock.patch.object(loader_nodes, "folder_paths", FolderPathsStub()), mock.patch.object(
|
|
loader_nodes,
|
|
"comfy_sd",
|
|
ComfySDStub(),
|
|
), mock.patch.object(loader_nodes, "comfy_core_nodes", CoreNodesStub()):
|
|
node = loader_nodes.LLSUniversalModelLoader()
|
|
model, text_encoder, vae, model_info = node.load(
|
|
"diffusion_models/flux1-schnell.safetensors",
|
|
"FLUX_SCHNELL",
|
|
"advanced",
|
|
"external",
|
|
"external",
|
|
"clip_l.safetensors",
|
|
"t5xxl_fp16.safetensors",
|
|
"ae.safetensors",
|
|
)
|
|
|
|
payload = json.loads(model_info)
|
|
self.assertEqual(model.label, "MODEL::flux1-schnell.safetensors")
|
|
self.assertEqual(text_encoder.label, "CLIP::flux::clip_l.safetensors+t5xxl_fp16.safetensors")
|
|
self.assertEqual(vae.label, "VAE::ae.safetensors")
|
|
self.assertEqual(payload["model_family"], "FLUX_SCHNELL")
|
|
self.assertEqual(payload["text_encoder_name_1"], "clip_l.safetensors")
|
|
self.assertEqual(payload["text_encoder_name_2"], "t5xxl_fp16.safetensors")
|
|
self.assertEqual(payload["text_encoder_name"], "clip_l.safetensors, t5xxl_fp16.safetensors")
|
|
self.assertEqual(payload["vae_name"], "ae.safetensors")
|
|
|
|
def test_loader_raises_clear_error_when_flux_text_encoder_is_missing(self):
|
|
load_plugin_package()
|
|
from lls_node_test_refactor.model_loader import nodes as loader_nodes
|
|
|
|
folder_paths = FolderPathsStub()
|
|
folder_paths._files["text_encoders"] = ["clip_l.safetensors"]
|
|
with mock.patch.object(loader_nodes, "folder_paths", folder_paths), mock.patch.object(
|
|
loader_nodes,
|
|
"comfy_sd",
|
|
ComfySDStub(),
|
|
), mock.patch.object(loader_nodes, "comfy_core_nodes", CoreNodesStub()):
|
|
node = loader_nodes.LLSSimpleCheckpointLoader()
|
|
with self.assertRaisesRegex(RuntimeError, "t5xxl_fp8_e4m3fn.safetensors|t5xxl_fp16.safetensors"):
|
|
node.load_checkpoint(
|
|
"diffusion_models/flux1-schnell.safetensors",
|
|
"FLUX_DEV",
|
|
"simple",
|
|
"external",
|
|
"external",
|
|
"ae.safetensors",
|
|
"clip_l.safetensors",
|
|
"(auto)",
|
|
)
|
|
|
|
def test_prompt_encode_schema_uses_standard_clip_ports_and_prompt_info(self):
|
|
plugin = load_plugin_package()
|
|
node_cls = plugin.NODE_CLASS_MAPPINGS["LLSSimplePromptEncode"]
|
|
required = node_cls.INPUT_TYPES()["required"]
|
|
optional = node_cls.INPUT_TYPES()["optional"]
|
|
|
|
self.assertEqual(node_cls.RETURN_TYPES, ("CONDITIONING", "CONDITIONING", "STRING"))
|
|
self.assertEqual(required["model_family"][1]["default"], "Auto")
|
|
self.assertEqual(optional["text_encoder"][0], "CLIP")
|
|
self.assertEqual(optional["clip"][0], "CLIP")
|
|
self.assertNotIn("task_context", optional)
|
|
|
|
def test_universal_prompt_encode_schema_requires_single_text_encoder_input(self):
|
|
plugin = load_plugin_package()
|
|
node_cls = plugin.NODE_CLASS_MAPPINGS["LLSUniversalPromptEncode"]
|
|
required = node_cls.INPUT_TYPES()["required"]
|
|
optional = node_cls.INPUT_TYPES()["optional"]
|
|
|
|
self.assertEqual(node_cls.RETURN_TYPES, ("CONDITIONING", "CONDITIONING", "STRING"))
|
|
self.assertEqual(required["text_encoder"][0], "CLIP")
|
|
self.assertEqual(required["clip_skip"][1]["default"], -1)
|
|
self.assertNotIn("clip", required)
|
|
self.assertEqual(optional["model_info"][0], "STRING")
|
|
|
|
def test_prompt_encode_dispatches_sdxl_using_inferred_family_defaults(self):
|
|
plugin = load_plugin_package()
|
|
node_cls = plugin.NODE_CLASS_MAPPINGS["LLSSimplePromptEncode"]
|
|
node = node_cls()
|
|
clip = RecordingClip(("g", "l"))
|
|
|
|
positive, negative, prompt_info = node.encode(
|
|
clip=clip,
|
|
positive_prompt="a castle",
|
|
negative_prompt="low quality",
|
|
clip_skip=-1,
|
|
model_family="Auto",
|
|
)
|
|
|
|
self.assertEqual(len(clip.calls), 2)
|
|
self.assertEqual(
|
|
clip.calls[0]["add_dict"],
|
|
{
|
|
"width": 1024,
|
|
"height": 1024,
|
|
"crop_w": 0,
|
|
"crop_h": 0,
|
|
"target_width": 1024,
|
|
"target_height": 1024,
|
|
},
|
|
)
|
|
payload = json.loads(prompt_info)
|
|
self.assertEqual(payload["model_family"], "SDXL")
|
|
self.assertEqual(payload["prompt_mode"], "sdxl")
|
|
self.assertEqual(payload["clip_skip"], -1)
|
|
self.assertEqual(positive[0][0], "cond::a castle")
|
|
self.assertEqual(negative[0][0], "cond::low quality")
|
|
|
|
def test_prompt_encode_dispatches_flux_and_neutralizes_negative_prompt(self):
|
|
plugin = load_plugin_package()
|
|
node_cls = plugin.NODE_CLASS_MAPPINGS["LLSSimplePromptEncode"]
|
|
node = node_cls()
|
|
clip = RecordingClip(("t5xxl",))
|
|
|
|
positive, negative, prompt_info = node.encode(
|
|
clip=clip,
|
|
positive_prompt="a robot in a forest",
|
|
negative_prompt="ugly, blurry",
|
|
clip_skip=-1,
|
|
model_family="Auto",
|
|
)
|
|
|
|
self.assertEqual(len(clip.calls), 2)
|
|
self.assertEqual(clip.calls[0]["add_dict"]["guidance"], 3.5)
|
|
self.assertEqual(clip.calls[1]["tokens"]["text"], "")
|
|
payload = json.loads(prompt_info)
|
|
self.assertEqual(payload["model_family"], "FLUX_DEV")
|
|
self.assertEqual(payload["negative_mode"], "ignored_for_flux")
|
|
self.assertEqual(payload["prompt_mode"], "flux")
|
|
self.assertEqual(positive[0][0], "cond::a robot in a forest")
|
|
self.assertEqual(negative[0][0], "cond::")
|
|
|
|
def test_universal_prompt_encode_dispatches_sdxl_from_model_info(self):
|
|
plugin = load_plugin_package()
|
|
node_cls = plugin.NODE_CLASS_MAPPINGS["LLSUniversalPromptEncode"]
|
|
node = node_cls()
|
|
clip = RecordingClip(("g", "l"))
|
|
|
|
positive, negative, prompt_info = node.encode(
|
|
text_encoder=clip,
|
|
positive_prompt="a castle",
|
|
negative_prompt="low quality",
|
|
clip_skip=-1,
|
|
model_info=json.dumps({"model_family": "SDXL", "checkpoint_name": "sdxl_turbo.safetensors"}),
|
|
)
|
|
|
|
self.assertEqual(len(clip.calls), 2)
|
|
self.assertEqual(clip.calls[0]["add_dict"]["width"], 1024)
|
|
payload = json.loads(prompt_info)
|
|
self.assertEqual(payload["model_family"], "SDXL")
|
|
self.assertEqual(payload["prompt_mode"], "sdxl")
|
|
self.assertEqual(payload["checkpoint_name"], "sdxl_turbo.safetensors")
|
|
self.assertEqual(positive[0][0], "cond::a castle")
|
|
self.assertEqual(negative[0][0], "cond::low quality")
|
|
|
|
def test_universal_prompt_encode_dispatches_flux_from_model_info(self):
|
|
plugin = load_plugin_package()
|
|
node_cls = plugin.NODE_CLASS_MAPPINGS["LLSUniversalPromptEncode"]
|
|
node = node_cls()
|
|
clip = RecordingClip(("t5xxl",))
|
|
|
|
positive, negative, prompt_info = node.encode(
|
|
text_encoder=clip,
|
|
positive_prompt="a robot in a forest",
|
|
negative_prompt="ugly, blurry",
|
|
clip_skip=-1,
|
|
model_info=json.dumps(
|
|
{
|
|
"model_family": "FLUX_DEV",
|
|
"checkpoint_name": "diffusion_models/flux1-schnell.safetensors",
|
|
"text_encoder_name": "clip_l.safetensors, t5xxl_fp16.safetensors",
|
|
}
|
|
),
|
|
)
|
|
|
|
self.assertEqual(len(clip.calls), 2)
|
|
self.assertEqual(clip.calls[0]["add_dict"]["guidance"], 3.5)
|
|
self.assertEqual(clip.calls[1]["tokens"]["text"], "")
|
|
payload = json.loads(prompt_info)
|
|
self.assertEqual(payload["model_family"], "FLUX_DEV")
|
|
self.assertEqual(payload["prompt_mode"], "flux")
|
|
self.assertEqual(payload["checkpoint_name"], "diffusion_models/flux1-schnell.safetensors")
|
|
self.assertEqual(payload["text_encoder_name"], "clip_l.safetensors, t5xxl_fp16.safetensors")
|
|
self.assertEqual(positive[0][0], "cond::a robot in a forest")
|
|
self.assertEqual(negative[0][0], "cond::")
|
|
|
|
def test_empty_latent_uses_model_family_inference_for_flux_shape(self):
|
|
load_plugin_package()
|
|
from lls_node_test_refactor.latent import nodes as latent_nodes
|
|
|
|
node = latent_nodes.LLSSimpleEmptyLatent()
|
|
model = TaggedValue("MODEL::flux1-schnell.safetensors")
|
|
model._lls_family = "FLUX_SCHNELL"
|
|
|
|
with mock.patch.object(latent_nodes, "torch", FakeTorch()):
|
|
latent, width, height, latent_info = node.create_empty_latent(
|
|
"Family Default",
|
|
512,
|
|
512,
|
|
2,
|
|
model_family="Auto",
|
|
model=model,
|
|
)
|
|
|
|
payload = json.loads(latent_info)
|
|
self.assertEqual((width, height), (1024, 1024))
|
|
self.assertEqual(tuple(latent["samples"].shape), (2, 128, 64, 64))
|
|
self.assertEqual(latent["downscale_ratio_spacial"], 16)
|
|
self.assertEqual(payload["model_family"], "FLUX_SCHNELL")
|
|
self.assertEqual(payload["task_mode"], "txt2img")
|
|
self.assertEqual(payload["batch_size"], 2)
|
|
|
|
def test_empty_latent_schema_accepts_optional_img2img_inputs(self):
|
|
plugin = load_plugin_package()
|
|
node_cls = plugin.NODE_CLASS_MAPPINGS["LLSSimpleEmptyLatent"]
|
|
required = node_cls.INPUT_TYPES()["required"]
|
|
optional = node_cls.INPUT_TYPES()["optional"]
|
|
|
|
self.assertEqual(
|
|
tuple(required.keys()),
|
|
(
|
|
"size_preset",
|
|
"width",
|
|
"height",
|
|
"batch_size",
|
|
"model_family",
|
|
"resize_mode",
|
|
),
|
|
)
|
|
self.assertEqual(required["resize_mode"][0], ["keep_aspect", "crop_center", "stretch", "none"])
|
|
self.assertEqual(optional["image"][0], "IMAGE")
|
|
self.assertEqual(optional["vae"][0], "VAE")
|
|
self.assertEqual(optional["model"][0], "MODEL")
|
|
|
|
def test_empty_latent_can_encode_image_using_family_default_size(self):
|
|
load_plugin_package()
|
|
from lls_node_test_refactor.latent import nodes as latent_nodes
|
|
from lls_node_test_refactor.image import nodes as image_nodes
|
|
|
|
comfy_utils = FakeComfyUtils()
|
|
model = TaggedValue("MODEL::SDXL")
|
|
model._lls_family = "SDXL"
|
|
vae = FakeVAE(vae_name="embedded")
|
|
|
|
with mock.patch.object(image_nodes, "comfy_utils", comfy_utils):
|
|
node = latent_nodes.LLSSimpleEmptyLatent()
|
|
latent, width, height, latent_info = node.create_empty_latent(
|
|
"Family Default",
|
|
512,
|
|
512,
|
|
3,
|
|
model_family="Auto",
|
|
resize_mode="keep_aspect",
|
|
model=model,
|
|
image=FakeTensor((1, 768, 1536, 3)),
|
|
vae=vae,
|
|
)
|
|
|
|
payload = json.loads(latent_info)
|
|
self.assertEqual((width, height), (1024, 512))
|
|
self.assertEqual(tuple(latent["samples"].shape), (1, 4, 64, 128))
|
|
self.assertEqual(latent["downscale_ratio_spacial"], 8)
|
|
self.assertEqual(vae.encoded_shapes, [(1, 512, 1024, 3)])
|
|
self.assertEqual(payload["task_mode"], "img2img")
|
|
self.assertEqual(payload["latent_source"], "image_encode")
|
|
self.assertEqual(payload["model_family"], "SDXL")
|
|
self.assertEqual(payload["size_preset"], "Family Default")
|
|
self.assertEqual(payload["batch_size"], 1)
|
|
|
|
def test_empty_latent_can_encode_image_using_custom_size(self):
|
|
load_plugin_package()
|
|
from lls_node_test_refactor.latent import nodes as latent_nodes
|
|
from lls_node_test_refactor.image import nodes as image_nodes
|
|
|
|
comfy_utils = FakeComfyUtils()
|
|
vae = FakeVAE()
|
|
|
|
with mock.patch.object(image_nodes, "comfy_utils", comfy_utils):
|
|
node = latent_nodes.LLSSimpleEmptyLatent()
|
|
latent, width, height, latent_info = node.create_empty_latent(
|
|
"Custom",
|
|
640,
|
|
768,
|
|
4,
|
|
model_family="SD1.5",
|
|
resize_mode="stretch",
|
|
image=FakeTensor((1, 512, 512, 3)),
|
|
vae=vae,
|
|
)
|
|
|
|
payload = json.loads(latent_info)
|
|
self.assertEqual((width, height), (640, 768))
|
|
self.assertEqual(tuple(latent["samples"].shape), (1, 4, 96, 80))
|
|
self.assertEqual(vae.encoded_shapes, [(1, 768, 640, 3)])
|
|
self.assertEqual(payload["task_mode"], "img2img")
|
|
self.assertEqual(payload["size_preset"], "Custom")
|
|
self.assertEqual(payload["resize_mode"], "stretch")
|
|
|
|
def test_empty_latent_requires_vae_when_image_is_connected(self):
|
|
load_plugin_package()
|
|
from lls_node_test_refactor.latent import nodes as latent_nodes
|
|
|
|
node = latent_nodes.LLSSimpleEmptyLatent()
|
|
|
|
with self.assertRaisesRegex(RuntimeError, "Missing VAE"):
|
|
node.create_empty_latent(
|
|
"Custom",
|
|
640,
|
|
640,
|
|
1,
|
|
model_family="SD1.5",
|
|
resize_mode="stretch",
|
|
image=FakeTensor((1, 512, 512, 3)),
|
|
)
|
|
|
|
def test_ksampler_reads_model_and_latent_inference_for_img2img(self):
|
|
load_plugin_package()
|
|
from lls_node_test_refactor.sampling import nodes as sampling_nodes
|
|
|
|
recorded = {}
|
|
model = TaggedValue("MODEL::SDXL_TURBO")
|
|
model._lls_family = "SDXL_TURBO"
|
|
|
|
def fake_common_ksampler(**kwargs):
|
|
recorded.update(kwargs)
|
|
return kwargs["latent"]
|
|
|
|
with mock.patch.object(sampling_nodes, "comfy_sample", object()), mock.patch.object(
|
|
sampling_nodes,
|
|
"comfy_samplers",
|
|
object(),
|
|
), mock.patch.object(
|
|
sampling_nodes,
|
|
"_common_ksampler",
|
|
side_effect=fake_common_ksampler,
|
|
):
|
|
node = sampling_nodes.LLSSimpleKSampler()
|
|
latent, sample_info = node.sample(
|
|
model=model,
|
|
positive=[["pos", {}]],
|
|
negative=[["neg", {}]],
|
|
latent_image={"samples": FakeTensor((1, 4, 64, 64)), "source": "image_encode"},
|
|
model_family="Auto",
|
|
quality_preset="Family Default",
|
|
seed=123,
|
|
steps=30,
|
|
cfg=7.0,
|
|
sampler_name="euler_ancestral",
|
|
scheduler="karras",
|
|
denoise=1.0,
|
|
flux_guidance=9.0,
|
|
)
|
|
|
|
sample_data = json.loads(sample_info)
|
|
self.assertEqual(recorded["steps"], 4)
|
|
self.assertEqual(recorded["cfg"], 1.0)
|
|
self.assertEqual(recorded["denoise"], 1.0)
|
|
self.assertEqual(sample_data["family"], "SDXL_TURBO")
|
|
self.assertEqual(sample_data["task_mode"], "img2img")
|
|
self.assertEqual(sample_data["sampler_name"], "euler")
|
|
self.assertEqual(latent["samples"].shape, (1, 4, 64, 64))
|
|
|
|
def test_ksampler_schema_keeps_legacy_widget_order_for_workflow_compatibility(self):
|
|
plugin = load_plugin_package()
|
|
node_cls = plugin.NODE_CLASS_MAPPINGS["LLSSimpleKSampler"]
|
|
required = node_cls.INPUT_TYPES()["required"]
|
|
|
|
self.assertEqual(
|
|
tuple(required.keys()),
|
|
(
|
|
"model",
|
|
"positive",
|
|
"negative",
|
|
"latent_image",
|
|
"quality_preset",
|
|
"seed",
|
|
"steps",
|
|
"cfg",
|
|
"sampler_name",
|
|
"scheduler",
|
|
"denoise",
|
|
"denoise_mode",
|
|
"adapter_mode",
|
|
"flux_guidance",
|
|
"model_family",
|
|
),
|
|
)
|
|
self.assertEqual(required["model_family"][1]["default"], "Auto")
|
|
|
|
def test_ksampler_accepts_blank_flux_guidance_for_legacy_workflows(self):
|
|
load_plugin_package()
|
|
from lls_node_test_refactor.sampling import nodes as sampling_nodes
|
|
|
|
required = sampling_nodes.LLSSimpleKSampler.INPUT_TYPES()["required"]
|
|
flux_input_type, flux_input_opts = required["flux_guidance"]
|
|
self.assertEqual(flux_input_type, "STRING,FLOAT,INT")
|
|
self.assertEqual(flux_input_opts["widgetType"], "FLOAT")
|
|
self.assertTrue(sampling_nodes.LLSSimpleKSampler.VALIDATE_INPUTS(flux_guidance=""))
|
|
|
|
recorded = {}
|
|
model = TaggedValue("MODEL::FLUX_DEV")
|
|
model._lls_family = "FLUX_DEV"
|
|
|
|
def fake_common_ksampler(**kwargs):
|
|
recorded.update(kwargs)
|
|
return kwargs["latent"]
|
|
|
|
fake_node_helpers = types.SimpleNamespace(
|
|
conditioning_set_values=lambda conditioning, values: [{"conditioning": conditioning, **values}]
|
|
)
|
|
with mock.patch.object(sampling_nodes, "comfy_sample", object()), mock.patch.object(
|
|
sampling_nodes,
|
|
"comfy_samplers",
|
|
object(),
|
|
), mock.patch.object(
|
|
sampling_nodes,
|
|
"_common_ksampler",
|
|
side_effect=fake_common_ksampler,
|
|
), mock.patch.object(
|
|
sampling_nodes,
|
|
"node_helpers",
|
|
fake_node_helpers,
|
|
):
|
|
node = sampling_nodes.LLSSimpleKSampler()
|
|
latent, sample_info = node.sample(
|
|
model=model,
|
|
positive=[["pos", {}]],
|
|
negative=[["neg", {}]],
|
|
latent_image={"samples": FakeTensor((1, 128, 64, 64)), "source": "empty_latent"},
|
|
model_family="Auto",
|
|
quality_preset="Manual",
|
|
seed=123,
|
|
steps=20,
|
|
cfg=1.0,
|
|
sampler_name="euler",
|
|
scheduler="simple",
|
|
denoise=1.0,
|
|
flux_guidance="",
|
|
)
|
|
|
|
sample_data = json.loads(sample_info)
|
|
self.assertEqual(recorded["positive"][0]["guidance"], 3.5)
|
|
self.assertEqual(recorded["negative"][0]["guidance"], 3.5)
|
|
self.assertEqual(sample_data["guidance"], 3.5)
|
|
self.assertEqual(latent["samples"].shape, (1, 128, 64, 64))
|
|
|
|
def test_vae_decode_returns_decode_info_for_flux(self):
|
|
plugin = load_plugin_package()
|
|
node_cls = plugin.NODE_CLASS_MAPPINGS["LLSSimpleVAEDecode"]
|
|
node = node_cls()
|
|
|
|
image, decode_info = node.decode(
|
|
{"samples": FakeTensor((1, 128, 64, 64)), "downscale_ratio_spacial": 16},
|
|
FakeVAE(vae_name="ae.safetensors"),
|
|
)
|
|
|
|
payload = json.loads(decode_info)
|
|
self.assertEqual(tuple(image.shape), (1, 64, 64, 3))
|
|
self.assertEqual(payload["width"], 1024)
|
|
self.assertEqual(payload["height"], 1024)
|
|
self.assertEqual(payload["vae_name"], "ae.safetensors")
|
|
self.assertEqual(payload["decode_stage"], "vae_decode")
|
|
|
|
def test_vae_encode_schema_exposes_img2img_inputs_without_context_port(self):
|
|
plugin = load_plugin_package()
|
|
node_cls = plugin.NODE_CLASS_MAPPINGS["LLSSimpleVAEEncode"]
|
|
required = node_cls.INPUT_TYPES()["required"]
|
|
optional = node_cls.INPUT_TYPES()["optional"]
|
|
|
|
self.assertEqual(node_cls.CATEGORY, "LLS/Image")
|
|
self.assertEqual(node_cls.RETURN_TYPES, ("LATENT", "INT", "INT", "STRING"))
|
|
self.assertEqual(required["image"][0], "IMAGE")
|
|
self.assertEqual(required["vae"][0], "VAE")
|
|
self.assertEqual(required["resize_mode"][0], ["keep_aspect", "crop_center", "stretch", "none"])
|
|
self.assertEqual(required["size_source"][0], ["input_image", "custom", "model_recommended"])
|
|
self.assertEqual(required["model_family"][1]["default"], "Auto")
|
|
self.assertEqual(optional["model"][0], "MODEL")
|
|
self.assertEqual(optional["clip"][0], "CLIP")
|
|
self.assertNotIn("task_context", optional)
|
|
|
|
def test_vae_encode_uses_inferred_model_recommended_size_for_img2img(self):
|
|
load_plugin_package()
|
|
from lls_node_test_refactor.image import nodes as image_nodes
|
|
|
|
comfy_utils = FakeComfyUtils()
|
|
vae = FakeVAE()
|
|
node = image_nodes.LLSSimpleVAEEncode()
|
|
model = TaggedValue("MODEL::SDXL")
|
|
model._lls_family = "SDXL"
|
|
vae._lls_vae_name = "embedded"
|
|
|
|
with mock.patch.object(image_nodes, "comfy_utils", comfy_utils):
|
|
latent, width, height, latent_info = node.encode(
|
|
image=FakeTensor((1, 768, 1536, 3)),
|
|
vae=vae,
|
|
resize_mode="keep_aspect",
|
|
size_source="model_recommended",
|
|
width=512,
|
|
height=512,
|
|
model_family="Auto",
|
|
model=model,
|
|
)
|
|
|
|
payload = json.loads(latent_info)
|
|
self.assertEqual((width, height), (1024, 512))
|
|
self.assertEqual(tuple(latent["samples"].shape), (1, 4, 64, 128))
|
|
self.assertEqual(latent["downscale_ratio_spacial"], 8)
|
|
self.assertEqual(vae.encoded_shapes, [(1, 512, 1024, 3)])
|
|
self.assertEqual(payload["task_mode"], "img2img")
|
|
self.assertEqual(payload["latent_source"], "image_encode")
|
|
self.assertEqual(payload["model_family"], "SDXL")
|
|
self.assertEqual(payload["vae_name"], "embedded")
|
|
self.assertEqual(
|
|
comfy_utils.calls,
|
|
[
|
|
{
|
|
"input_shape": (1, 3, 768, 1536),
|
|
"width": 1024,
|
|
"height": 512,
|
|
"upscale_method": "bilinear",
|
|
"crop": "disabled",
|
|
}
|
|
],
|
|
)
|
|
|
|
def test_vae_encode_input_image_mode_can_passthrough_aligned_size(self):
|
|
load_plugin_package()
|
|
from lls_node_test_refactor.image import nodes as image_nodes
|
|
|
|
vae = FakeVAE()
|
|
node = image_nodes.LLSSimpleVAEEncode()
|
|
latent, width, height, latent_info = node.encode(
|
|
image=FakeTensor((1, 512, 768, 3)),
|
|
vae=vae,
|
|
resize_mode="none",
|
|
size_source="input_image",
|
|
width=1024,
|
|
height=1024,
|
|
model_family="SD1.5",
|
|
)
|
|
|
|
payload = json.loads(latent_info)
|
|
self.assertEqual((width, height), (768, 512))
|
|
self.assertEqual(tuple(latent["samples"].shape), (1, 4, 64, 96))
|
|
self.assertEqual(latent["downscale_ratio_spacial"], 8)
|
|
self.assertEqual(vae.encoded_shapes, [(1, 512, 768, 3)])
|
|
self.assertEqual(payload["latent_source"], "image_encode")
|
|
self.assertEqual(payload["size_source"], "input_image")
|
|
self.assertEqual(payload["resize_mode"], "none")
|
|
|
|
def test_save_image_schema_uses_info_strings_only(self):
|
|
plugin = load_plugin_package()
|
|
node_cls = plugin.NODE_CLASS_MAPPINGS["LLSSaveImage"]
|
|
required = node_cls.INPUT_TYPES()["required"]
|
|
optional = node_cls.INPUT_TYPES()["optional"]
|
|
|
|
self.assertEqual(required["image"][0], "IMAGE")
|
|
self.assertEqual(required["filename_prefix"][0], "STRING")
|
|
self.assertEqual(required["output_mode"][0], ["save", "preview_only"])
|
|
self.assertEqual(required["save_metadata"][0], "BOOLEAN")
|
|
self.assertNotIn("model", optional)
|
|
self.assertNotIn("clip", optional)
|
|
self.assertNotIn("vae", optional)
|
|
self.assertEqual(optional["prompt_info"], ("STRING", {"forceInput": True}))
|
|
self.assertEqual(optional["latent_info"], ("STRING", {"forceInput": True}))
|
|
self.assertEqual(optional["sample_info"], ("STRING", {"forceInput": True}))
|
|
self.assertEqual(optional["decode_info"], ("STRING", {"forceInput": True}))
|
|
self.assertEqual(optional["upscale_info"], ("STRING", {"forceInput": True}))
|
|
|
|
def test_save_image_merges_metadata_from_info_strings(self):
|
|
load_plugin_package()
|
|
from lls_node_test_refactor.image import nodes as image_nodes
|
|
|
|
with mock.patch.object(image_nodes, "comfy_core_nodes", CoreNodesStub()):
|
|
node = image_nodes.LLSSaveImage()
|
|
result = node.save(
|
|
image=FakeTensor((1, 512, 512, 3)),
|
|
filename_prefix="LLS",
|
|
save_metadata=True,
|
|
prompt_info=json.dumps(
|
|
{
|
|
"positive_prompt": "cat",
|
|
"negative_prompt": "bad",
|
|
"model_family": "SD1.5",
|
|
"checkpoint_name": "sd15.safetensors",
|
|
"text_encoder_name": "clip_l.safetensors",
|
|
}
|
|
),
|
|
sample_info=json.dumps(
|
|
{
|
|
"seed": 42,
|
|
"steps": 20,
|
|
"cfg": 7.0,
|
|
"guidance": None,
|
|
"sampler_name": "euler",
|
|
"scheduler": "normal",
|
|
"denoise": 1.0,
|
|
}
|
|
),
|
|
decode_info=json.dumps(
|
|
{
|
|
"vae_name": "VAE::SD15",
|
|
"width": 512,
|
|
"height": 512,
|
|
"batch_size": 1,
|
|
}
|
|
),
|
|
upscale_info=json.dumps({"mode": "none", "scale": 1.0}),
|
|
)
|
|
|
|
metadata = result["ui"]["images"][0]["extra_pnginfo"]["lls_metadata"]
|
|
self.assertEqual(metadata["positive_prompt"], "cat")
|
|
self.assertEqual(metadata["checkpoint_name"], "sd15.safetensors")
|
|
self.assertEqual(metadata["text_encoder_name"], "clip_l.safetensors")
|
|
self.assertEqual(metadata["vae_name"], "VAE::SD15")
|
|
self.assertEqual(metadata["steps"], 20)
|
|
self.assertEqual(metadata["upscale_mode"], "none")
|
|
|
|
def test_save_image_preview_only_uses_native_preview_node_without_saving(self):
|
|
load_plugin_package()
|
|
from lls_node_test_refactor.image import nodes as image_nodes
|
|
|
|
CoreNodesStub.save_calls = []
|
|
CoreNodesStub.preview_calls = []
|
|
|
|
with mock.patch.object(image_nodes, "comfy_core_nodes", CoreNodesStub()):
|
|
node = image_nodes.LLSSaveImage()
|
|
result = node.save(
|
|
image=FakeTensor((1, 512, 512, 3)),
|
|
filename_prefix="LLS",
|
|
save_metadata=True,
|
|
output_mode="preview_only",
|
|
)
|
|
|
|
self.assertEqual(CoreNodesStub.save_calls, [])
|
|
self.assertEqual(len(CoreNodesStub.preview_calls), 1)
|
|
self.assertEqual(result["ui"]["images"][0]["type"], "temp")
|
|
self.assertIsNone(result["ui"]["images"][0]["extra_pnginfo"])
|
|
|
|
def test_generation_config_uses_model_inference_without_context(self):
|
|
plugin = load_plugin_package()
|
|
node_cls = plugin.NODE_CLASS_MAPPINGS["LLSGenerationConfig"]
|
|
node = node_cls()
|
|
required = node_cls.INPUT_TYPES()["required"]
|
|
optional = node_cls.INPUT_TYPES()["optional"]
|
|
model = TaggedValue("MODEL::FLUX_DEV")
|
|
model._lls_family = "FLUX_DEV"
|
|
|
|
self.assertEqual(required["model_family"][1]["default"], "Auto")
|
|
self.assertEqual(optional["model"][0], "MODEL")
|
|
self.assertEqual(optional["clip"][0], "CLIP")
|
|
self.assertEqual(
|
|
tuple(required.keys()),
|
|
("quality_preset", "size_preset", "model_family"),
|
|
)
|
|
|
|
width, height, steps, cfg, guidance, sampler_name, scheduler, denoise, config_info = node.execute(
|
|
quality_preset="Family Default",
|
|
size_preset="Family Default",
|
|
model_family="Auto",
|
|
model=model,
|
|
)
|
|
|
|
payload = json.loads(config_info)
|
|
self.assertEqual((width, height), (1024, 1024))
|
|
self.assertEqual((steps, cfg, guidance), (20, 1.0, 3.5))
|
|
self.assertEqual(sampler_name, "euler")
|
|
self.assertEqual(scheduler, "simple")
|
|
self.assertEqual(denoise, 1.0)
|
|
self.assertEqual(payload["family"], "FLUX_DEV")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|