Files
olivv-cs-ComfyUI-FunPack/tests/test_scene_builder.py
T

361 lines
12 KiB
Python

import json
import sys
import types
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
class FakeRoutes:
def get(self, _path):
return lambda func: func
def post(self, _path):
return lambda func: func
sys.modules.setdefault(
"folder_paths",
types.SimpleNamespace(get_user_directory=lambda: ""),
)
sys.modules.setdefault(
"server",
types.SimpleNamespace(PromptServer=types.SimpleNamespace(instance=types.SimpleNamespace(routes=FakeRoutes()))),
)
import templates
from conditioning import FunPackVideoRefinerV2
from templates import FunPackSceneBuilder, load_scene_db, normalize_scene_memory_items, save_scene_db
def use_tmp_scene_store(monkeypatch, tmp_path):
monkeypatch.setattr(templates, "template_store_dir", lambda: str(tmp_path))
monkeypatch.setattr(
templates,
"refinement_state_path",
lambda key, namespace, prefix="refine": str(tmp_path / "refinements" / f"{prefix}_{namespace}_{key}.json"),
)
monkeypatch.setattr(
FunPackVideoRefinerV2,
"_v2_state_path",
lambda self, key: str(tmp_path / "refinements" / f"refine_v2_clip_{key}.json"),
)
def test_scene_builder_collects_inputs_as_memory_but_outputs_selection(monkeypatch, tmp_path):
use_tmp_scene_store(monkeypatch, tmp_path)
builder = FunPackSceneBuilder()
payload = {
"positive_phrases": ["red dress", "walking pose"],
"negative_phrases": ["bad hands"],
}
input_lora_stack = {
"version": 2,
"mode": "ltx2",
"per_block": True,
"loras": [{"name": "vampire.safetensors", "model_weight": 0.8}],
}
positive, negative, lora_stack, status = builder.build_scene(
scene="-None-",
scene_name="Vampire Courtyard",
aliases="gothic vampire",
action="save",
mode="Manual",
intent_prompt="",
positive_prompt="woman in courtyard, red dress, moonlight",
negative_prompt="blurry, bad hands",
refinement_key="scene_key",
scene_payload=json.dumps(payload),
lora_stack=input_lora_stack,
)
assert positive == "red dress, walking pose"
assert negative == "bad hands"
assert "moonlight" not in positive
assert lora_stack is input_lora_stack
assert "Manual saved" in status
data = load_scene_db("scene_key")
assert "woman in courtyard" in data["universal_memory"]
assert "blurry" in data["universal_memory"]
assert data["scenes"]["Vampire Courtyard"]["positive_phrases"] == ["red dress", "walking pose"]
def test_scene_builder_public_interface_is_prompt_stack_status_only():
assert FunPackSceneBuilder.RETURN_NAMES == (
"positive_prompt",
"negative_prompt",
"lora_stack",
"status",
)
assert FunPackSceneBuilder.RETURN_TYPES == (
"STRING",
"STRING",
"FUNPACK_LORA_STACK",
"STRING",
)
inputs = FunPackSceneBuilder.INPUT_TYPES()
required = inputs["required"]
optional = inputs["optional"]
assert "sigmas" not in required
assert "sigmas" not in optional
assert "positive_conditioning" not in required
assert "positive_conditioning" not in optional
assert "refinement_key" not in FunPackSceneBuilder.RETURN_NAMES
assert "scene_name" not in FunPackSceneBuilder.RETURN_NAMES
assert required["scene_positive"][1]["multiline"] is False
assert required["scene_negative"][1]["multiline"] is False
def test_scene_builder_manual_outputs_exact_composed_text(monkeypatch, tmp_path):
use_tmp_scene_store(monkeypatch, tmp_path)
builder = FunPackSceneBuilder()
outputs = builder.build_scene(
scene="-None-",
scene_name="Editable Scene",
aliases="",
action="save",
mode="Manual",
scene_positive="red dress. walking pose, custom order",
scene_negative="bad hands. blur",
positive_prompt="source memory only, should not output",
negative_prompt="source negative only",
refinement_key="scene_key",
scene_payload=json.dumps({
"positive_phrases": ["ignored payload phrase"],
"negative_phrases": ["ignored payload negative"],
}),
)
assert outputs[0] == "red dress. walking pose, custom order"
assert outputs[1] == "bad hands. blur"
assert "source memory only" not in outputs[0]
data = load_scene_db("scene_key")
scene = data["scenes"]["Editable Scene"]
assert scene["positive_text"] == "red dress. walking pose, custom order"
assert scene["negative_text"] == "bad hands. blur"
assert "source memory only" in data["universal_memory"]
def test_scene_builder_auto_exact_and_fuzzy_match_saved_scene(monkeypatch, tmp_path):
use_tmp_scene_store(monkeypatch, tmp_path)
builder = FunPackSceneBuilder()
saved_payload = {
"positive_phrases": ["anthro dragon", "mossy scales"],
"negative_phrases": ["extra limbs"],
}
manual_payload = {
"positive_phrases": ["manual fallback"],
"negative_phrases": [],
}
builder.build_scene(
scene="-None-",
scene_name="Forest Dragon",
aliases="woodland beast",
action="save",
mode="Manual",
intent_prompt="",
positive_prompt="anthro dragon, mossy scales",
negative_prompt="extra limbs",
refinement_key="dragon_key",
scene_payload=json.dumps(saved_payload),
)
exact = builder.build_scene(
scene="-None-",
scene_name="",
aliases="",
action="load",
mode="Auto",
intent_prompt="please apply woodland beast now",
positive_prompt="ignored prompt",
negative_prompt="ignored negative",
refinement_key="dragon_key",
scene_payload=json.dumps(manual_payload),
)
assert exact[0] == "anthro dragon, mossy scales"
assert exact[1] == "extra limbs"
assert exact[2] is None
assert "Auto exact" in exact[3]
fuzzy = builder.build_scene(
scene="-None-",
scene_name="",
aliases="",
action="load",
mode="Auto",
intent_prompt="make it forest and dragon",
positive_prompt="ignored prompt",
negative_prompt="",
refinement_key="dragon_key",
scene_payload=json.dumps(manual_payload),
)
assert fuzzy[0] == "anthro dragon, mossy scales"
assert "Auto fuzzy" in fuzzy[3]
def test_scene_builder_auto_no_match_falls_back_to_manual(monkeypatch, tmp_path):
use_tmp_scene_store(monkeypatch, tmp_path)
builder = FunPackSceneBuilder()
payload = {
"positive_phrases": ["manual only"],
"negative_phrases": ["manual negative"],
}
outputs = builder.build_scene(
scene="-None-",
scene_name="",
aliases="",
action="load",
mode="Auto",
intent_prompt="no known scene here",
positive_prompt="source memory only",
negative_prompt="bad source memory",
refinement_key="",
scene_payload=json.dumps(payload),
)
assert outputs[0] == "manual only"
assert outputs[1] == "manual negative"
assert "Manual" in outputs[3]
def test_scene_builder_learning_mode_collects_memory_and_passes_inputs_through(monkeypatch, tmp_path):
use_tmp_scene_store(monkeypatch, tmp_path)
builder = FunPackSceneBuilder()
lora_stack = {"loras": [{"name": "active.safetensors"}]}
positive_prompt = " first prompt,\nsecond prompt "
negative_prompt = "bad anatomy, blur"
outputs = builder.build_scene(
scene="-None-",
scene_name="",
aliases="",
action="load",
mode="Learning",
intent_prompt="ignored scene trigger",
positive_prompt=positive_prompt,
negative_prompt=negative_prompt,
refinement_key="learn_key",
scene_payload=json.dumps({
"positive_phrases": ["manual phrase"],
"negative_phrases": ["manual negative"],
}),
lora_stack=lora_stack,
)
assert outputs[0] == positive_prompt
assert outputs[1] == negative_prompt
assert outputs[2] is lora_stack
assert "Learning" in outputs[3]
data = load_scene_db("learn_key")
assert "first prompt" in data["universal_memory"]
assert "second prompt" in data["universal_memory"]
assert "bad anatomy" in data["universal_memory"]
assert data["scenes"] == {}
def test_scene_builder_collects_prompt_words_from_sentence_chunks(monkeypatch, tmp_path):
use_tmp_scene_store(monkeypatch, tmp_path)
builder = FunPackSceneBuilder()
builder.build_scene(
scene="-None-",
scene_name="",
aliases="",
action="load",
mode="Learning",
positive_prompt="person smoking in the rain",
negative_prompt="",
refinement_key="word_key",
)
data = load_scene_db("word_key")
assert "person smoking in the rain" in data["universal_memory"]
assert "person" in data["universal_memory"]
assert "smoking" in data["universal_memory"]
assert "the" not in data["universal_memory"]
def test_refiner_reset_preserves_scene_builder_memory_in_refinement_key(monkeypatch, tmp_path):
use_tmp_scene_store(monkeypatch, tmp_path)
builder = FunPackSceneBuilder()
builder.build_scene(
scene="-None-",
scene_name="Memory Scene",
aliases="",
action="save",
mode="Manual",
positive_prompt="silver hair, rain street",
negative_prompt="blur",
refinement_key="shared_key",
scene_payload=json.dumps({
"positive_phrases": ["silver hair"],
"negative_phrases": ["blur"],
}),
)
refiner = FunPackVideoRefinerV2()
reset_state, reset_status = refiner._v2_load_state("shared_key", reset_session=True)
assert reset_status == "fresh"
assert reset_state["prompt_histories"] == {}
assert reset_state["global"]["phrase_memory"] == {}
assert reset_state["scene_builder"]["scenes"]["Memory Scene"]["positive_phrases"] == ["silver hair"]
assert "silver hair" in reset_state["scene_builder"]["universal_memory"]
def test_scene_builder_database_items_keep_categories_and_wildcards():
memory = normalize_scene_memory_items([
{"text": "red dress", "category": "appearance", "wildcard": True},
{"text": "blue dress", "category": "appearance", "wildcard": True},
{"text": "bad hands", "category": "negative"},
])
assert memory["red dress"]["category"] == "appearance"
assert memory["red dress"]["wildcard"] is True
assert memory["blue dress"]["wildcard"] is True
assert memory["bad hands"]["wildcard"] is False
assert memory["bad hands"]["source"] == "negative"
def test_scene_builder_wildcard_outputs_one_adjacent_matching_phrase(monkeypatch, tmp_path):
use_tmp_scene_store(monkeypatch, tmp_path)
builder = FunPackSceneBuilder()
data = load_scene_db("wild_key")
data["universal_memory"] = normalize_scene_memory_items([
{"text": "red dress", "category": "appearance", "wildcard": True},
{"text": "blue dress", "category": "appearance", "wildcard": True},
{"text": "green dress", "category": "appearance", "wildcard": True},
])
save_scene_db(data, "wild_key")
monkeypatch.setattr(templates.random, "choice", lambda choices: "blue dress")
outputs = builder.build_scene(
scene="-None-",
scene_name="Wildcard Scene",
aliases="",
action="load",
mode="Manual",
scene_positive="red dress, blue dress, walking pose, green dress",
scene_negative="",
refinement_key="wild_key",
)
assert outputs[0] == "blue dress, walking pose, green dress"
assert "red dress" not in outputs[0]
def test_scene_builder_legacy_wildcard_group_migrates_to_checkbox():
memory = normalize_scene_memory_items([
{"text": "red dress", "category": "appearance", "wildcard_group": "outfit"},
])
assert memory["red dress"]["wildcard"] is True