284 lines
8.6 KiB
Python
284 lines
8.6 KiB
Python
from workflow_extractor import WorkflowExtractor
|
|
|
|
|
|
def test_workflow_extractor_basic_prompt():
|
|
prompt = {
|
|
"1": {
|
|
"class_type": "CheckpointLoaderSimple",
|
|
"inputs": {"ckpt_name": "model.safetensors"},
|
|
},
|
|
"2": {
|
|
"class_type": "CLIPTextEncode",
|
|
"inputs": {"text": "positive", "clip": ["1", 1]},
|
|
},
|
|
"3": {
|
|
"class_type": "CLIPTextEncode",
|
|
"inputs": {"text": "negative", "clip": ["1", 1]},
|
|
},
|
|
"4": {
|
|
"class_type": "KSampler",
|
|
"inputs": {
|
|
"seed": 123,
|
|
"steps": 20,
|
|
"cfg": 7.5,
|
|
"sampler_name": "euler",
|
|
"scheduler": "normal",
|
|
"denoise": 1.0,
|
|
"model": ["1", 0],
|
|
"positive": ["2", 0],
|
|
"negative": ["3", 0],
|
|
},
|
|
},
|
|
"5": {
|
|
"class_type": "VAEDecode",
|
|
"inputs": {"samples": ["4", 0], "vae": ["1", 2]},
|
|
},
|
|
"6": {
|
|
"class_type": "MetaHubSaveNode",
|
|
"inputs": {"images": ["5", 0]},
|
|
},
|
|
}
|
|
|
|
extractor = WorkflowExtractor(prompt)
|
|
data, missing = extractor.extract(save_node_id="6")
|
|
|
|
assert not missing
|
|
assert data["seed"] == 123
|
|
assert data["steps"] == 20
|
|
assert data["cfg"] == 7.5
|
|
assert data["sampler_name"] == "euler"
|
|
assert data["scheduler"] == "normal"
|
|
assert data["denoise"] == 1.0
|
|
assert data["positive"] == "positive"
|
|
assert data["negative"] == "negative"
|
|
assert data["model_name"] == "model.safetensors"
|
|
assert data["vae_name"] == "model.safetensors"
|
|
|
|
|
|
def test_workflow_extractor_lora_detection():
|
|
prompt = {
|
|
"1": {
|
|
"class_type": "LoraLoader",
|
|
"inputs": {"lora_name": "detail.safetensors", "strength_model": 0.8},
|
|
}
|
|
}
|
|
|
|
extractor = WorkflowExtractor(prompt)
|
|
data, missing = extractor.extract()
|
|
|
|
assert "loras" not in missing
|
|
assert data["lora_list"] == [{"name": "detail.safetensors", "weight": 0.8}]
|
|
|
|
|
|
def test_workflow_extractor_detects_flux_unet_model_and_vae():
|
|
prompt = {
|
|
"1": {
|
|
"class_type": "UNETLoader",
|
|
"inputs": {"unet_name": "flux-model.safetensors"},
|
|
},
|
|
"2": {
|
|
"class_type": "VAELoader",
|
|
"inputs": {"vae_name": "flux-vae.safetensors"},
|
|
},
|
|
"3": {
|
|
"class_type": "CLIPTextEncode",
|
|
"inputs": {"text": "positive", "clip": ["9", 0]},
|
|
},
|
|
"4": {
|
|
"class_type": "BasicScheduler",
|
|
"inputs": {"model": ["1", 0], "steps": 8, "scheduler": "normal"},
|
|
},
|
|
"5": {
|
|
"class_type": "KSampler",
|
|
"inputs": {
|
|
"seed": 123,
|
|
"steps": 8,
|
|
"cfg": 1,
|
|
"sampler_name": "euler",
|
|
"scheduler": "normal",
|
|
"model": ["1", 0],
|
|
"positive": ["3", 0],
|
|
"negative": ["3", 0],
|
|
"latent_image": ["8", 0],
|
|
},
|
|
},
|
|
"6": {
|
|
"class_type": "VAEDecode",
|
|
"inputs": {"samples": ["5", 0], "vae": ["2", 0]},
|
|
},
|
|
"7": {
|
|
"class_type": "MetaHubSaveNode",
|
|
"inputs": {"images": ["6", 0]},
|
|
},
|
|
"8": {"class_type": "EmptyLatentImage", "inputs": {}},
|
|
"9": {"class_type": "CLIPLoader", "inputs": {}},
|
|
}
|
|
|
|
extractor = WorkflowExtractor(prompt)
|
|
data, _missing = extractor.extract(save_node_id="7")
|
|
|
|
assert data["model_name"] == "flux-model.safetensors"
|
|
assert data["vae_name"] == "flux-vae.safetensors"
|
|
|
|
|
|
def test_workflow_extractor_detects_lora_alternate_keys():
|
|
prompt = {
|
|
"1": {
|
|
"class_type": "PowerLoraLoader",
|
|
"inputs": {"lora_name_1": "style.safetensors", "strength": 0.65},
|
|
}
|
|
}
|
|
|
|
extractor = WorkflowExtractor(prompt)
|
|
data, _missing = extractor.extract()
|
|
|
|
assert data["lora_list"] == [{"name": "style.safetensors", "weight": 0.65}]
|
|
|
|
|
|
def test_workflow_extractor_resolves_connected_seed_and_prompt_strings():
|
|
prompt = {
|
|
"1": {
|
|
"class_type": "UNETLoader",
|
|
"inputs": {"unet_name": "zImageTurboNSFW_82BF16FP8.safetensors"},
|
|
},
|
|
"2": {
|
|
"class_type": "Lora Loader (LoraManager)",
|
|
"inputs": {
|
|
"text": "<lora:Z-Detail-Slider:0.95>",
|
|
"loras": {
|
|
"__value__": [
|
|
{
|
|
"name": "Z-Detail-Slider",
|
|
"strength": 0.75,
|
|
"active": True,
|
|
"clipStrength": 0.75,
|
|
}
|
|
]
|
|
},
|
|
"model": ["1", 0],
|
|
},
|
|
},
|
|
"3": {
|
|
"class_type": "easy positive",
|
|
"inputs": {"positive": "rabbit"},
|
|
},
|
|
"4": {
|
|
"class_type": "JoinStrings",
|
|
"inputs": {"delimiter": " ", "string1": ["3", 0]},
|
|
},
|
|
"5": {
|
|
"class_type": "easy stylesSelector",
|
|
"inputs": {"styles": "fooocus_styles", "select_styles": [], "positive": ["4", 0]},
|
|
},
|
|
"6": {
|
|
"class_type": "CLIPTextEncode",
|
|
"inputs": {"text": ["5", 0], "clip": ["2", 1]},
|
|
},
|
|
"7": {
|
|
"class_type": "ConditioningZeroOut",
|
|
"inputs": {"conditioning": ["6", 0]},
|
|
},
|
|
"8": {
|
|
"class_type": "SeedGenerator",
|
|
"inputs": {"seed": 862877771869053},
|
|
},
|
|
"9": {
|
|
"class_type": "ModelSamplingAuraFlow",
|
|
"inputs": {"shift": 3.0, "model": ["2", 0]},
|
|
},
|
|
"10": {
|
|
"class_type": "KSampler",
|
|
"inputs": {
|
|
"seed": ["8", 0],
|
|
"steps": 9,
|
|
"cfg": 1.0,
|
|
"sampler_name": "euler",
|
|
"scheduler": "simple",
|
|
"denoise": 1.0,
|
|
"model": ["9", 0],
|
|
"positive": ["6", 0],
|
|
"negative": ["7", 0],
|
|
},
|
|
},
|
|
"11": {
|
|
"class_type": "VAEDecode",
|
|
"inputs": {"samples": ["10", 0], "vae": ["12", 0]},
|
|
},
|
|
"12": {
|
|
"class_type": "VAELoader",
|
|
"inputs": {"vae_name": "ae.safetensors"},
|
|
},
|
|
"13": {
|
|
"class_type": "MetaHubSaveNode",
|
|
"inputs": {"images": ["11", 0]},
|
|
},
|
|
}
|
|
|
|
extractor = WorkflowExtractor(prompt)
|
|
data, missing = extractor.extract(save_node_id="13")
|
|
|
|
assert data["seed"] == 862877771869053
|
|
assert data["positive"] == "rabbit"
|
|
assert data["negative"] == ""
|
|
assert data["model_name"] == "zImageTurboNSFW_82BF16FP8.safetensors"
|
|
assert data["vae_name"] == "ae.safetensors"
|
|
assert data["lora_list"] == [{"name": "Z-Detail-Slider", "weight": 0.75}]
|
|
assert "seed" not in missing
|
|
assert "positive" not in missing
|
|
|
|
|
|
|
|
def test_workflow_extractor_detects_img2img_lineage():
|
|
prompt = {
|
|
"1": {
|
|
"class_type": "LoadImage",
|
|
"inputs": {"image": "inputs/base.png"},
|
|
},
|
|
"2": {
|
|
"class_type": "VAEEncode",
|
|
"inputs": {"pixels": ["1", 0], "vae": ["5", 2]},
|
|
},
|
|
"3": {
|
|
"class_type": "KSampler",
|
|
"inputs": {
|
|
"seed": 123,
|
|
"steps": 20,
|
|
"cfg": 7.5,
|
|
"sampler_name": "euler",
|
|
"scheduler": "normal",
|
|
"denoise": 0.45,
|
|
"model": ["5", 0],
|
|
"positive": ["6", 0],
|
|
"negative": ["7", 0],
|
|
"latent_image": ["2", 0],
|
|
},
|
|
},
|
|
"4": {
|
|
"class_type": "VAEDecode",
|
|
"inputs": {"samples": ["3", 0], "vae": ["5", 2]},
|
|
},
|
|
"5": {
|
|
"class_type": "CheckpointLoaderSimple",
|
|
"inputs": {"ckpt_name": "model.safetensors"},
|
|
},
|
|
"6": {
|
|
"class_type": "CLIPTextEncode",
|
|
"inputs": {"text": "positive", "clip": ["5", 1]},
|
|
},
|
|
"7": {
|
|
"class_type": "CLIPTextEncode",
|
|
"inputs": {"text": "negative", "clip": ["5", 1]},
|
|
},
|
|
"8": {
|
|
"class_type": "MetaHubSaveNode",
|
|
"inputs": {"images": ["4", 0]},
|
|
},
|
|
}
|
|
|
|
extractor = WorkflowExtractor(prompt)
|
|
data, _missing = extractor.extract(save_node_id="8")
|
|
|
|
assert data["generation_type"] == "img2img"
|
|
assert data["source_image"]["fileName"] == "base.png"
|
|
assert data["source_image"]["relativePath"] == "inputs/base.png"
|