from workflow_extractor import WorkflowExtractor def test_qwen_image_21_uses_distinct_conditioning_outputs(): prompt = { "1": {"class_type": "UNETLoader", "inputs": {"unet_name": "qwen_image_2.1.safetensors"}}, "2": {"class_type": "String Literal", "inputs": {"string": "linked positive"}}, "3": {"class_type": "TextEncodeQwenImage21", "inputs": { "prompt": ["2", 0], "negative_prompt": "direct negative", }}, "4": {"class_type": "KSampler", "inputs": { "model": ["1", 0], "positive": ["3", 0], "negative": ["3", 1], "seed": 42, "steps": 25, "cfg": 1, "sampler_name": "euler", "scheduler": "simple", "denoise": 1, }}, "5": {"class_type": "VAEDecode", "inputs": {"samples": ["4", 0]}}, "6": {"class_type": "MetaHubSaveNode", "inputs": {"images": ["5", 0]}}, } data, missing = WorkflowExtractor(prompt).extract(save_node_id="6") assert data["positive"] == "linked positive" assert data["negative"] == "direct negative" assert data["model_name"] == "qwen_image_2.1.safetensors" assert (data["seed"], data["steps"], data["cfg"]) == (42, 25, 1.0) assert (data["sampler_name"], data["scheduler"]) == ("euler", "simple") assert "positive" not in missing and "negative" not in missing prompt["3"]["inputs"]["negative_prompt"] = "" data, missing = WorkflowExtractor(prompt).extract(save_node_id="6") assert data["negative"] == "" assert "negative" not in missing prompt["3"]["inputs"]["negative_prompt"] = ["8", 0] prompt["7"] = {"class_type": "String Literal", "inputs": {"string": ""}} prompt["8"] = {"class_type": "ComfySwitchNode", "inputs": { "switch": False, "on_false": ["7", 0], "on_true": ["9", 0], }} prompt["9"] = {"class_type": "PrimitiveStringMultiline", "inputs": { "value": "inactive negative", }} data, missing = WorkflowExtractor(prompt).extract(save_node_id="6") assert data["negative"] == "" assert "negative" not in missing prompt["3"]["inputs"]["negative_prompt"] = ["7", 0] data, missing = WorkflowExtractor(prompt).extract(save_node_id="6") assert data["negative"] == "" assert "negative" not in missing def test_qwen_image_21_prompts_stay_separate_through_cfg_guider(): prompt = { "1": {"class_type": "TextEncodeQwenImage21", "inputs": { "prompt": "positive text", "negative_prompt": "negative text", }}, "2": {"class_type": "CFGGuider", "inputs": { "positive": ["1", 0], "negative": ["1", 1], "cfg": 1, }}, "3": {"class_type": "SamplerCustomAdvanced", "inputs": { "guider": ["2", 0], "noise": ["4", 0], "sigmas": ["5", 0], }}, "4": {"class_type": "RandomNoise", "inputs": {"noise_seed": 42}}, "5": {"class_type": "BasicScheduler", "inputs": {"steps": 25}}, } data, missing = WorkflowExtractor(prompt).extract() assert data["positive"] == "positive text" assert data["negative"] == "negative text" assert "positive" not in missing and "negative" not in missing 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_traces_sampler_from_model_3d_input(): prompt = { "1": {"class_type": "KSampler", "inputs": {"seed": 111}}, "2": {"class_type": "KSampler", "inputs": {"seed": 222}}, "3": {"class_type": "SyntheticMeshGenerator", "inputs": {"samples": ["2", 0]}}, "4": {"class_type": "MetaHubSave3DModel", "inputs": {"model_3d": ["3", 0]}}, } extractor = WorkflowExtractor(prompt) data, _missing = extractor.extract(save_node_id="4") assert data["seed"] == 222 def test_workflow_extractor_traces_vae_from_model_3d_input(): prompt = { "1": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": "model.safetensors"}}, "2": {"class_type": "KSampler", "inputs": {"model": ["1", 0]}}, "3": {"class_type": "VAEDecode", "inputs": {"samples": ["2", 0], "vae": ["1", 2]}}, "4": {"class_type": "SyntheticMeshGenerator", "inputs": {"images": ["3", 0]}}, "5": {"class_type": "MetaHubSave3DModel", "inputs": {"model_3d": ["4", 0]}}, } extractor = WorkflowExtractor(prompt) data, _missing = extractor.extract(save_node_id="5") assert data["vae_name"] == "model.safetensors" def test_workflow_extractor_traces_source_image_from_model_3d_input(): prompt = { "1": { "class_type": "LoadImage", "inputs": {"image": "inputs/source.png"}, }, "2": { "class_type": "SyntheticImageToMesh", "inputs": {"image": ["1", 0]}, }, "3": { "class_type": "MetaHubSave3DModel", "inputs": {"model_3d": ["2", 0]}, }, } extractor = WorkflowExtractor(prompt) data, _missing = extractor.extract(save_node_id="3") assert data["generation_type"] == "image2model3d" assert data["source_image"] == { "fileName": "source.png", "relativePath": "inputs/source.png", "nodeId": "1", "nodeType": "LoadImage", } 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_detects_rgthree_power_lora_objects(): prompt = { "67": { "class_type": "Power Lora Loader (rgthree)", "inputs": { "lora_1": { "on": True, "lora": "Z-Detail-Slider.safetensors", "strength": 0.5325, }, "lora_2": { "on": True, "lora": "zy_CinematicShot_zit.safetensors", "strength": 0.6875, }, "lora_3": { "on": False, "lora": "disabled.safetensors", "strength": 1.0, }, }, } } extractor = WorkflowExtractor(prompt) data, _missing = extractor.extract() assert data["lora_list"] == [ {"name": "Z-Detail-Slider.safetensors", "weight": 0.5325}, {"name": "zy_CinematicShot_zit.safetensors", "weight": 0.6875}, ] def test_workflow_extractor_reads_style_prompt_encoder_through_basic_guider(): prompt = { "61": { "class_type": "SamplerCustomAdvanced", "inputs": { "noise": ["63", 0], "guider": ["101", 0], "sampler": ["108", 0], "sigmas": ["62", 0], "latent_image": ["113", 0], }, }, "101": { "class_type": "BasicGuider", "inputs": { "model": ["100", 0], "conditioning": ["114", 0], }, }, "114": { "class_type": "StylePromptEncoder2 //ZImagePowerNodes", "inputs": { "style": "\"Production Photo\"", "text": "A white cat on the roof of a brick house", "clip": ["67", 1], }, }, } extractor = WorkflowExtractor(prompt) data, missing = extractor.extract() assert data["positive"] == "A white cat on the roof of a brick house" assert "positive" not in missing 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": "", "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" def test_workflow_extractor_resolves_custom_advanced_sampler_chain(): prompt = { "6": { "class_type": "CLIPTextEncode", "inputs": {"text": "positive prompt", "clip": ["97", 0]}, }, "39": { "class_type": "VAELoader", "inputs": {"vae_name": "diffusion_pytorch_model.safetensors"}, }, "61": { "class_type": "SamplerCustomAdvanced", "inputs": { "noise": ["63", 0], "guider": ["101", 0], "sampler": ["108", 0], "sigmas": ["62", 0], "latent_image": ["104", 0], }, }, "62": { "class_type": "BasicScheduler", "inputs": {"scheduler": "beta", "steps": 9, "denoise": 1.0, "model": ["100", 0]}, }, "63": { "class_type": "RandomNoise", "inputs": {"noise_seed": 604786586961193}, }, "71": { "class_type": "VAEDecode", "inputs": {"samples": ["61", 0], "vae": ["39", 0]}, }, "91": { "class_type": "UNETLoader", "inputs": {"unet_name": "jibMixZIT_v10.safetensors"}, }, "100": { "class_type": "ModelSamplingAuraFlow", "inputs": {"shift": 5.0, "model": ["91", 0]}, }, "101": { "class_type": "BasicGuider", "inputs": {"model": ["100", 0], "conditioning": ["6", 0]}, }, "104": { "class_type": "EmptyLatentImage", "inputs": {"width": 1024, "height": 1024, "batch_size": 1}, }, "108": { "class_type": "ClownSampler_Beta", "inputs": {"eta": 0.23, "sampler_name": "linear/ralston_2s", "seed": 141088624843721}, }, "112": { "class_type": "MetaHubSaveImage", "inputs": {"images": ["71", 0]}, }, } extractor = WorkflowExtractor(prompt) data, missing = extractor.extract(save_node_id="112") assert data["seed"] == 604786586961193 assert data["steps"] == 9 assert data["sampler_name"] == "linear/ralston_2s" assert data["scheduler"] == "beta" assert data["denoise"] == 1.0 assert data["positive"] == "positive prompt" assert data["model_name"] == "jibMixZIT_v10.safetensors" assert data["vae_name"] == "diffusion_pytorch_model.safetensors" assert {"seed", "steps", "sampler_name", "scheduler", "denoise", "positive", "model_name", "vae_name"}.isdisjoint(missing) def test_workflow_extractor_follows_krea2_false_switch_to_prompt_text(): expected = "A studio portrait with dramatic directional lighting" prompt = { "67": {"class_type": "PrimitiveBoolean", "inputs": {"value": False}}, "70": { "class_type": "ComfySwitchNode", "inputs": {"switch": ["67", 0], "on_false": ["71", 0], "on_true": ["61", 0]}, }, "71": {"class_type": "PrimitiveStringMultiline", "inputs": {"value": expected}}, "72": { "class_type": "KSampler", "inputs": { "seed": 123, "steps": 10, "cfg": 1, "sampler_name": "euler", "scheduler": "simple", "positive": ["73", 0], }, }, "73": {"class_type": "CLIPTextEncode", "inputs": {"text": ["70", 0]}}, } extracted, missing = WorkflowExtractor(prompt).extract() assert extracted["positive"] == expected assert "positive" not in missing def test_workflow_extractor_preserves_empty_selected_switch_branch(): prompt = { "1": {"class_type": "CLIPTextEncode", "inputs": {"text": "positive prompt"}}, "2": {"class_type": "PrimitiveStringMultiline", "inputs": {"value": "inactive negative prompt"}}, "3": {"class_type": "ConditioningZeroOut", "inputs": {}}, "4": { "class_type": "ComfySwitchNode", "inputs": {"switch": False, "on_false": ["3", 0], "on_true": ["2", 0]}, }, "5": {"class_type": "CLIPTextEncode", "inputs": {"text": ["4", 0]}}, "6": { "class_type": "KSampler", "inputs": { "seed": 123, "steps": 10, "cfg": 1, "sampler_name": "euler", "scheduler": "simple", "positive": ["1", 0], "negative": ["5", 0], }, }, } extracted, missing = WorkflowExtractor(prompt).extract() assert extracted["negative"] == "" assert "negative" not in missing