469 lines
19 KiB
Python
469 lines
19 KiB
Python
from mjr_am_backend.features.geninfo import api_node_extractor as x
|
|
from mjr_am_backend.features.geninfo import parser_impl as p
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Detection helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def test_detect_api_node_known_providers():
|
|
assert x._detect_api_node("geminiimagenode") == "google_gemini"
|
|
assert x._detect_api_node("openaigptimage1") == "openai"
|
|
assert x._detect_api_node("openaiimagenode") == "openai"
|
|
assert x._detect_api_node("claudenode") == "anthropic"
|
|
assert x._detect_api_node("bytedanceseedancenode") == "bytedance_seedance"
|
|
assert x._detect_api_node("klingtexttovideo") == "kling_ai"
|
|
assert x._detect_api_node("lumaimagetovideo") == "luma_dream_machine"
|
|
assert x._detect_api_node("minimaxhailuo") == "minimax_hailuo"
|
|
assert x._detect_api_node("runwayimagetovideo") == "runway"
|
|
assert x._detect_api_node("vidutexttovideo") == "vidu"
|
|
assert x._detect_api_node("moonvalleytxt2video") == "moonvalley"
|
|
assert x._detect_api_node("pixversetexttovideo") == "pixverse"
|
|
assert x._detect_api_node("ideogramv2") == "ideogram"
|
|
assert x._detect_api_node("recraftv4") == "recraft"
|
|
assert x._detect_api_node("grokimage") == "xai_grok"
|
|
assert x._detect_api_node("ltxvapitext") == "ltxv_api"
|
|
assert x._detect_api_node("elevenlabstextto") == "eleven_labs"
|
|
assert x._detect_api_node("wan2texttovideoapi") == "alibaba_wan"
|
|
assert x._detect_api_node("happyhorseimagetovideoapi") == "happy_horse"
|
|
assert x._detect_api_node("fluxkontext") == "black_forest_labs"
|
|
assert x._detect_api_node("stabilitystableimage") == "stability_ai"
|
|
assert x._detect_api_node("unknownnode") is None
|
|
|
|
|
|
def test_detect_api_node_specificity_order():
|
|
# More specific fragments must match before the broad catch-all
|
|
assert x._detect_api_node("bytedanceseedream") == "bytedance_seedance"
|
|
assert x._detect_api_node("bytedance2referencenode") == "bytedance_seedance"
|
|
assert x._detect_api_node("geminiimage2node") == "google_gemini"
|
|
assert x._detect_api_node("veo3videogeneration") == "google_veo"
|
|
|
|
|
|
def test_provider_from_model_name():
|
|
assert x._provider_from_model_name("seedance-2.0-pro") == "bytedance_seedance"
|
|
assert x._provider_from_model_name("gemini-2.0-flash") == "google_gemini"
|
|
assert x._provider_from_model_name("veo-3") == "google_veo"
|
|
assert x._provider_from_model_name("gpt-4o") == "openai"
|
|
assert x._provider_from_model_name("dall-e-3") == "openai"
|
|
assert x._provider_from_model_name("claude-3-opus") == "anthropic"
|
|
assert x._provider_from_model_name("flux-kontext-pro") == "black_forest_labs"
|
|
assert x._provider_from_model_name("stable-diffusion-xl") == "stability_ai"
|
|
assert x._provider_from_model_name("happyhorse-1.0-i2v") == "happy_horse"
|
|
assert x._provider_from_model_name("wan-2.1") == "alibaba_wan"
|
|
assert x._provider_from_model_name("kling-v2") == "kling_ai"
|
|
assert x._provider_from_model_name("luma-dream-machine") == "luma_dream_machine"
|
|
assert x._provider_from_model_name("minimax-01") == "minimax_hailuo"
|
|
assert x._provider_from_model_name("runway-gen3") == "runway"
|
|
assert x._provider_from_model_name("vidu-q1") == "vidu"
|
|
assert x._provider_from_model_name("ideogram-v3") == "ideogram"
|
|
assert x._provider_from_model_name("recraft-v3") == "recraft"
|
|
assert x._provider_from_model_name("grok-aurora") == "xai_grok"
|
|
assert x._provider_from_model_name("pixverse-v4") == "pixverse"
|
|
assert x._provider_from_model_name("unknown-model-xyz") is None
|
|
|
|
|
|
def test_provider_from_model_name_wan_no_media_type_required():
|
|
# Wan provider must be detected even without "video" or "image" in name
|
|
assert x._provider_from_model_name("wan-2.1-14b") == "alibaba_wan"
|
|
assert x._provider_from_model_name("wanx-turbo") == "alibaba_wan"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Node classification
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def test_is_api_image_node_known_class_type():
|
|
node = {"class_type": "GeminiImageNode", "inputs": {}}
|
|
assert x._is_api_image_node(node) is True
|
|
|
|
|
|
def test_is_api_image_node_heuristic_fallback():
|
|
node = {
|
|
"class_type": "CustomExternalGen",
|
|
"inputs": {"prompt": "a cat", "model": "custom-v1"},
|
|
}
|
|
assert x._is_api_image_node(node) is True
|
|
|
|
|
|
def test_is_api_image_node_heuristic_rejects_sampler_node():
|
|
node = {
|
|
"class_type": "CustomSampler",
|
|
"inputs": {"prompt": "a cat", "model": "v1", "cfg": 7.5, "steps": 20},
|
|
}
|
|
assert x._is_api_image_node(node) is False
|
|
|
|
|
|
def test_is_api_image_node_rejects_non_api():
|
|
node = {"class_type": "CLIPTextEncode", "inputs": {"text": "hello"}}
|
|
assert x._is_api_image_node(node) is False
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Field extraction helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def test_extract_api_node_prompt_priority():
|
|
assert x._extract_api_node_prompt({"prompt": "a cat"}) == "a cat"
|
|
assert x._extract_api_node_prompt({"model.prompt": "a galloping horse"}) == "a galloping horse"
|
|
assert x._extract_api_node_prompt({"text": "a dog"}) == "a dog"
|
|
assert x._extract_api_node_prompt({"positive_prompt": "sun"}) == "sun"
|
|
assert x._extract_api_node_prompt({}) is None
|
|
assert x._extract_api_node_prompt({"prompt": " "}) is None
|
|
|
|
|
|
def test_extract_api_node_negative():
|
|
assert x._extract_api_node_negative({"negative_prompt": "blur"}) == "blur"
|
|
assert x._extract_api_node_negative({"model.negative_prompt": "washed out"}) == "washed out"
|
|
assert x._extract_api_node_negative({"negative": "ugly"}) == "ugly"
|
|
assert x._extract_api_node_negative({}) is None
|
|
|
|
|
|
def test_extract_api_node_prompt_resolves_linked_text():
|
|
nodes = {
|
|
"1": {"class_type": "PrimitiveStringMultiline", "inputs": {"value": "galloping chrome horse"}},
|
|
}
|
|
assert x._extract_api_node_prompt({"model.prompt": ["1", 0]}, nodes) == "galloping chrome horse"
|
|
|
|
|
|
def test_extract_api_node_model():
|
|
node = {"class_type": "GeminiImageNode"}
|
|
m = x._extract_api_node_model({"model": "gemini-2.0-flash"}, "1", node)
|
|
assert m is not None
|
|
assert m["name"] == "gemini-2.0-flash"
|
|
assert m["confidence"] == "high"
|
|
assert x._extract_api_node_model({}, "1", node) is None
|
|
assert x._extract_api_node_model({"model": " "}, "1", node) is None
|
|
|
|
|
|
def test_extract_api_system_prompt():
|
|
assert x._extract_api_system_prompt({"system_prompt": "be concise"}) == "be concise"
|
|
assert x._extract_api_system_prompt({"system": "you are helpful"}) == "you are helpful"
|
|
assert x._extract_api_system_prompt({}) is None
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Engine type refinement
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def test_refine_engine_type_txt2img():
|
|
out = {"engine": {"type": "txt2img"}}
|
|
x._refine_engine_type(out, "openaitexttoimageapi", {})
|
|
assert out["engine"]["type"] == "txt2img"
|
|
|
|
|
|
def test_refine_engine_type_img2vid():
|
|
out = {"engine": {"type": "txt2vid"}}
|
|
x._refine_engine_type(out, "klingimagetovideo", {})
|
|
assert out["engine"]["type"] == "img2vid"
|
|
|
|
|
|
def test_refine_engine_type_img2vid_from_video_with_image_input():
|
|
out = {"engine": {"type": "txt2vid"}}
|
|
x._refine_engine_type(out, "somevideogen", {"image": ["1", 0]})
|
|
assert out["engine"]["type"] == "img2vid"
|
|
|
|
|
|
def test_refine_engine_type_txt2vid_from_video_no_image_input():
|
|
out = {"engine": {"type": "txt2img"}}
|
|
x._refine_engine_type(out, "somevideogen", {"prompt": "a sunset"})
|
|
assert out["engine"]["type"] == "txt2vid"
|
|
|
|
|
|
def test_refine_engine_type_audio():
|
|
out = {"engine": {"type": "txt2img"}}
|
|
x._refine_engine_type(out, "sonilotexttomusic", {})
|
|
assert out["engine"]["type"] == "audio"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Seedance helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def test_seedance_size_from_resolution():
|
|
assert x._seedance_size_from_resolution("720p") == (1280, 720)
|
|
assert x._seedance_size_from_resolution("1080p") == (1920, 1080)
|
|
assert x._seedance_size_from_resolution("480p") == (854, 480)
|
|
assert x._seedance_size_from_resolution("4k") == (3840, 2160)
|
|
assert x._seedance_size_from_resolution("360p") is None
|
|
|
|
|
|
def test_seedance_media_type_txt2vid():
|
|
assert x._seedance_media_type({"prompt": "a cat"}) == "txt2vid"
|
|
|
|
|
|
def test_seedance_media_type_img2vid():
|
|
ins = {"model.reference_images": ["1", 0]}
|
|
assert x._seedance_media_type(ins) == "img2vid"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Full fallback entry-point
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _make_gemini_nodes(prompt="a golden retriever in sunlight", model="gemini-2.0-flash-preview-image-generation"):
|
|
return {
|
|
"1": {
|
|
"class_type": "GeminiImageNode",
|
|
"inputs": {"prompt": prompt, "model": model, "seed": 42},
|
|
},
|
|
"2": {"class_type": "SaveImage", "inputs": {"images": ["1", 0]}},
|
|
}
|
|
|
|
|
|
def test_extract_api_node_geninfo_fallback_basic():
|
|
nodes = _make_gemini_nodes()
|
|
result = x._extract_api_node_geninfo_fallback(nodes, None)
|
|
assert result is not None
|
|
assert result["positive"]["value"] == "a golden retriever in sunlight"
|
|
assert result["engine"]["api_provider"] == "google_gemini"
|
|
assert result["engine"]["sampler_mode"] == "api"
|
|
assert result["checkpoint"]["name"] == "gemini-2.0-flash-preview-image-generation"
|
|
assert result["seed"]["value"] == 42
|
|
|
|
|
|
def test_extract_api_node_geninfo_fallback_no_api_nodes():
|
|
nodes = {
|
|
"1": {"class_type": "CLIPTextEncode", "inputs": {"text": "hello"}},
|
|
}
|
|
assert x._extract_api_node_geninfo_fallback(nodes, None) is None
|
|
|
|
|
|
def test_extract_api_node_geninfo_fallback_returns_none_if_no_meaningful_fields():
|
|
# Node is detected as API type but has no extractable data
|
|
nodes = {"1": {"class_type": "GeminiImageNode", "inputs": {}}}
|
|
result = x._extract_api_node_geninfo_fallback(nodes, None)
|
|
assert result is None
|
|
|
|
|
|
def test_extract_api_node_geninfo_fallback_with_workflow_meta():
|
|
nodes = _make_gemini_nodes()
|
|
meta = {"title": "Test Workflow"}
|
|
result = x._extract_api_node_geninfo_fallback(nodes, meta)
|
|
assert result is not None
|
|
assert result.get("metadata") == meta
|
|
|
|
|
|
def test_extract_api_node_geninfo_fallback_multiple_nodes():
|
|
nodes = {
|
|
"1": {
|
|
"class_type": "GeminiImageNode",
|
|
"inputs": {"prompt": "first prompt", "model": "gemini-2.0"},
|
|
},
|
|
"2": {
|
|
"class_type": "OpenAIImageNode",
|
|
"inputs": {"prompt": "second prompt", "model": "dall-e-3"},
|
|
},
|
|
"3": {"class_type": "SaveImage", "inputs": {"images": ["1", 0]}},
|
|
}
|
|
result = x._extract_api_node_geninfo_fallback(nodes, None)
|
|
assert result is not None
|
|
extra = result.get("extra_api_passes", [])
|
|
assert len(extra) == 1
|
|
assert extra[0]["positive"]["value"] == "second prompt"
|
|
|
|
|
|
def test_extract_api_node_geninfo_fallback_openai():
|
|
nodes = {
|
|
"1": {
|
|
"class_type": "OpenAIImageNode",
|
|
"inputs": {"prompt": "a serene mountain lake", "model": "gpt-image-1", "seed": 7},
|
|
},
|
|
}
|
|
result = x._extract_api_node_geninfo_fallback(nodes, None)
|
|
assert result is not None
|
|
assert result["engine"]["api_provider"] == "openai"
|
|
assert result["positive"]["value"] == "a serene mountain lake"
|
|
|
|
|
|
def test_extract_api_node_geninfo_fallback_openai_gpt_image_linked_prompt():
|
|
nodes = {
|
|
"1": {"class_type": "PrimitiveStringMultiline", "inputs": {"value": "shoe"}},
|
|
"2": {
|
|
"class_type": "StringReplace",
|
|
"inputs": {
|
|
"string": "Create a product technical analysis infographic for [YOUR PRODUCT].",
|
|
"find": "[YOUR PRODUCT]",
|
|
"replace": ["1", 0],
|
|
},
|
|
},
|
|
"3": {
|
|
"class_type": "OpenAIGPTImage1",
|
|
"inputs": {
|
|
"prompt": ["2", 0],
|
|
"seed": 126220530,
|
|
"quality": "low",
|
|
"background": "auto",
|
|
"size": "1024x1024",
|
|
"n": 1,
|
|
"model": "gpt-image-2",
|
|
"custom_width": 1024,
|
|
"custom_height": 1024,
|
|
"image": ["4", 0],
|
|
},
|
|
},
|
|
"4": {"class_type": "LoadImage", "inputs": {"image": "reference.png"}},
|
|
"5": {"class_type": "SaveImage", "inputs": {"images": ["3", 0]}},
|
|
}
|
|
result = x._extract_api_node_geninfo_fallback(nodes, None)
|
|
assert result is not None
|
|
assert result["engine"]["api_provider"] == "openai"
|
|
assert result["engine"]["type"] == "txt2img"
|
|
assert result["positive"]["value"] == "Create a product technical analysis infographic for shoe."
|
|
assert result["checkpoint"]["name"] == "gpt-image-2"
|
|
assert result["seed"]["value"] == 126220530
|
|
assert result["quality"]["value"] == "low"
|
|
assert result["batch_size"]["value"] == 1
|
|
assert result["size"]["width"] == 1024
|
|
assert result["size"]["height"] == 1024
|
|
assert result["inputs"][0]["filename"] == "reference.png"
|
|
|
|
|
|
def test_extract_api_node_geninfo_fallback_seedance_dotted_keys():
|
|
nodes = {
|
|
"1": {
|
|
"class_type": "ByteDanceSeedanceNode",
|
|
"inputs": {
|
|
"model": "Seedance 2.0",
|
|
"model.prompt": "cinematic ocean sunset",
|
|
"model.negative_prompt": "blur, noisy",
|
|
"model.resolution": "1080p",
|
|
"model.duration": 5,
|
|
"seed": 99,
|
|
},
|
|
},
|
|
}
|
|
result = x._extract_api_node_geninfo_fallback(nodes, None)
|
|
assert result is not None
|
|
assert result["engine"]["api_provider"] == "bytedance_seedance"
|
|
assert result["positive"]["value"] == "cinematic ocean sunset"
|
|
assert result["negative"]["value"] == "blur, noisy"
|
|
assert result["size"]["width"] == 1920
|
|
assert result["seed"]["value"] == 99
|
|
|
|
|
|
def test_extract_api_node_geninfo_fallback_happy_horse_dotted_keys():
|
|
nodes = {
|
|
"13": {
|
|
"class_type": "HappyHorseImageToVideoApi",
|
|
"inputs": {
|
|
"model": "happyhorse-1.0-i2v",
|
|
"model.prompt": "a happy horse running in a meadow",
|
|
"model.resolution": "720p",
|
|
"model.duration": 5,
|
|
"seed": 962403639,
|
|
"first_frame": ["14", 0],
|
|
},
|
|
},
|
|
"14": {
|
|
"class_type": "LoadImage",
|
|
"inputs": {"image": "horse.png"},
|
|
},
|
|
}
|
|
result = x._extract_api_node_geninfo_fallback(nodes, None)
|
|
assert result is not None
|
|
assert result["engine"]["api_provider"] == "happy_horse"
|
|
assert result["engine"]["type"] == "img2vid"
|
|
assert result["positive"]["value"] == "a happy horse running in a meadow"
|
|
assert result["checkpoint"]["name"] == "happyhorse-1.0-i2v"
|
|
assert result["resolution"]["value"] == "720p"
|
|
assert result["size"]["width"] == 1280
|
|
assert result["size"]["height"] == 720
|
|
assert result["duration"]["value"] == 5
|
|
assert result["seed"]["value"] == 962403639
|
|
|
|
|
|
def test_extract_api_node_geninfo_fallback_with_negative_prompt():
|
|
nodes = {
|
|
"1": {
|
|
"class_type": "StabilityStableImage",
|
|
"inputs": {
|
|
"prompt": "portrait of a warrior",
|
|
"negative_prompt": "ugly, deformed",
|
|
"model": "stable-diffusion-xl-1024-v1-0",
|
|
},
|
|
},
|
|
}
|
|
result = x._extract_api_node_geninfo_fallback(nodes, None)
|
|
assert result is not None
|
|
assert result["negative"]["value"] == "ugly, deformed"
|
|
|
|
|
|
def test_extract_api_node_geninfo_fallback_size_from_width_height():
|
|
nodes = {
|
|
"1": {
|
|
"class_type": "RecraftV4",
|
|
"inputs": {"prompt": "abstract art", "model": "recraft-v3", "width": 1024, "height": 768},
|
|
},
|
|
}
|
|
result = x._extract_api_node_geninfo_fallback(nodes, None)
|
|
assert result is not None
|
|
assert result["size"]["width"] == 1024
|
|
assert result["size"]["height"] == 768
|
|
|
|
|
|
def test_extract_api_node_geninfo_fallback_wan_provider():
|
|
nodes = {
|
|
"1": {
|
|
"class_type": "Wan2TextToVideoAPI",
|
|
"inputs": {"prompt": "a butterfly emerging from cocoon", "model": "wan-2.1"},
|
|
},
|
|
}
|
|
result = x._extract_api_node_geninfo_fallback(nodes, None)
|
|
assert result is not None
|
|
assert result["engine"]["api_provider"] == "alibaba_wan"
|
|
|
|
|
|
def test_extract_api_node_geninfo_fallback_aspect_ratio():
|
|
nodes = {
|
|
"1": {
|
|
"class_type": "KlingTextToVideo",
|
|
"inputs": {"prompt": "flying eagle", "model": "kling-v2", "aspect_ratio": "16:9"},
|
|
},
|
|
}
|
|
result = x._extract_api_node_geninfo_fallback(nodes, None)
|
|
assert result is not None
|
|
assert result.get("aspect_ratio", {}).get("value") == "16:9"
|
|
|
|
|
|
def test_extract_api_node_geninfo_fallback_full_integration_via_parser():
|
|
"""End-to-end: parse_geninfo_from_prompt must surface API node data when no sampler present."""
|
|
nodes = {
|
|
"1": {
|
|
"class_type": "GeminiImageNode",
|
|
"inputs": {"prompt": "a red fox in a snowy forest", "model": "gemini-2.0-flash", "seed": 12},
|
|
},
|
|
"2": {"class_type": "SaveImage", "inputs": {"images": ["1", 0]}},
|
|
}
|
|
res = p.parse_geninfo_from_prompt(nodes)
|
|
assert res.ok
|
|
data = res.data
|
|
assert isinstance(data, dict)
|
|
assert data.get("positive", {}).get("value") == "a red fox in a snowy forest"
|
|
assert data.get("engine", {}).get("api_provider") == "google_gemini"
|
|
|
|
|
|
def test_parse_geninfo_from_prompt_surfaces_linked_happy_horse_prompt():
|
|
nodes = {
|
|
"1": {"class_type": "PrimitiveStringMultiline", "inputs": {"value": "cinematic happy horse across a foggy valley"}},
|
|
"13": {
|
|
"class_type": "HappyHorseImageToVideoApi",
|
|
"inputs": {
|
|
"model": "happyhorse-1.0-i2v",
|
|
"model.prompt": ["1", 0],
|
|
"model.resolution": "720P",
|
|
"model.duration": 5,
|
|
"seed": 962403639,
|
|
"first_frame": ["14", 0],
|
|
},
|
|
},
|
|
"14": {"class_type": "LoadImage", "inputs": {"image": "horse.png"}},
|
|
"15": {"class_type": "VHS_VideoCombine", "inputs": {"images": ["13", 0]}},
|
|
}
|
|
|
|
res = p.parse_geninfo_from_prompt(nodes)
|
|
assert res.ok
|
|
data = res.data
|
|
assert isinstance(data, dict)
|
|
assert data.get("engine", {}).get("api_provider") == "happy_horse"
|
|
assert data.get("positive", {}).get("value") == "cinematic happy horse across a foggy valley"
|
|
assert data.get("checkpoint", {}).get("name") == "happyhorse-1.0-i2v"
|