84 lines
2.9 KiB
Python
84 lines
2.9 KiB
Python
import json
|
|
|
|
from nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS, TelegramAutoReport
|
|
|
|
|
|
def test_auto_report_is_registered():
|
|
assert NODE_CLASS_MAPPINGS["TelegramAutoReport"] is TelegramAutoReport
|
|
assert NODE_DISPLAY_NAME_MAPPINGS["TelegramAutoReport"] == "Telegram Auto Workflow Report"
|
|
|
|
|
|
def test_extracts_generation_values_from_comfy_prompt_graph():
|
|
node = TelegramAutoReport()
|
|
payload = {
|
|
"prompt": {
|
|
"1": {"class_type": "CLIPTextEncode", "inputs": {"text": "a neon city"}},
|
|
"2": {"class_type": "CLIPTextEncode", "inputs": {"text": "low quality"}},
|
|
"3": {
|
|
"class_type": "KSampler",
|
|
"inputs": {
|
|
"positive": ["1", 0],
|
|
"negative": ["2", 0],
|
|
"seed": 123,
|
|
"steps": 30,
|
|
"cfg": 7.5,
|
|
"sampler_name": "euler",
|
|
"scheduler": "normal",
|
|
},
|
|
},
|
|
"4": {
|
|
"class_type": "CheckpointLoaderSimple",
|
|
"inputs": {"ckpt_name": "models/checkpoints/base.safetensors"},
|
|
},
|
|
"5": {
|
|
"class_type": "LoraLoader",
|
|
"inputs": {"lora_name": "styles/detail.safetensors", "strength_model": 0.65},
|
|
},
|
|
"6": {
|
|
"class_type": "EmptyLatentVideo",
|
|
"inputs": {"width": 832, "height": 480, "length": 81, "fps": 16},
|
|
},
|
|
}
|
|
}
|
|
|
|
extracted = node._extract_from_payload(payload)
|
|
|
|
assert extracted["positive_prompt"] == "a neon city"
|
|
assert extracted["negative_prompt"] == "low quality"
|
|
assert extracted["seed"] == 123
|
|
assert extracted["steps"] == 30
|
|
assert extracted["cfg"] == 7.5
|
|
assert extracted["sampler"] == "euler"
|
|
assert extracted["scheduler"] == "normal"
|
|
assert extracted["base_model"] == "base.safetensors"
|
|
assert extracted["extra_loras"] == "detail.safetensors: 0.65"
|
|
assert extracted["width"] == 832
|
|
assert extracted["height"] == 480
|
|
assert extracted["frames"] == 81
|
|
assert extracted["fps"] == 16
|
|
|
|
|
|
def test_extracts_metadata_from_companion_json(tmp_path):
|
|
video_path = tmp_path / "clip.mp4"
|
|
metadata_path = tmp_path / "clip.json"
|
|
video_path.write_bytes(b"video")
|
|
metadata_path.write_text(
|
|
json.dumps(
|
|
{
|
|
"prompt": {
|
|
"1": {"class_type": "CLIPTextEncode", "inputs": {"text": "prompt text"}},
|
|
"2": {
|
|
"class_type": "KSampler",
|
|
"inputs": {"positive": ["1", 0], "seed": 456},
|
|
},
|
|
}
|
|
}
|
|
),
|
|
encoding="utf-8",
|
|
)
|
|
|
|
extracted = TelegramAutoReport()._extract_from_metadata_files([str(video_path)])
|
|
|
|
assert extracted["positive_prompt"] == "prompt text"
|
|
assert extracted["seed"] == 456
|