danrh
07f3dc254f
Introduce model-specific target profiles, fail-closed grounding and validation, content-addressed Director caching, LTX-2.5 contracts, comparison workflows, benchmarks, schemas, documentation, and regression coverage.
501 lines
19 KiB
Python
501 lines
19 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import copy
|
|
import importlib.util
|
|
import json
|
|
import math
|
|
import os
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
import uuid
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
|
|
import torch
|
|
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
COMFY_ROOT = ROOT.parents[1]
|
|
|
|
|
|
def load_nodes_module():
|
|
if str(COMFY_ROOT) not in sys.path:
|
|
sys.path.insert(0, str(COMFY_ROOT))
|
|
package_name = f"diffusiongemma_director_cache_{uuid.uuid4().hex}"
|
|
spec = importlib.util.spec_from_file_location(
|
|
package_name,
|
|
ROOT / "__init__.py",
|
|
submodule_search_locations=[str(ROOT)],
|
|
)
|
|
if spec is None or spec.loader is None:
|
|
raise RuntimeError("could not load DiffusionGemma Prompt Builder")
|
|
package = importlib.util.module_from_spec(spec)
|
|
sys.modules[package_name] = package
|
|
spec.loader.exec_module(package)
|
|
return sys.modules[f"{package_name}.nodes"]
|
|
|
|
|
|
class DirectorCacheTests(unittest.TestCase):
|
|
@classmethod
|
|
def setUpClass(cls) -> None:
|
|
cls.nodes = load_nodes_module()
|
|
|
|
def runtime(self, *, unload_policy: str = "keep_loaded", model_path: str = "unused"):
|
|
return self.nodes.RuntimeConfig(
|
|
model_path=model_path,
|
|
backend="transformers_inprocess",
|
|
dtype="auto",
|
|
quantization="modelopt_nvfp4",
|
|
local_files_only=True,
|
|
unload_policy=unload_policy,
|
|
max_memory_gb=20.0,
|
|
status={"ready": True, "supports_pixels": True},
|
|
)
|
|
|
|
def context(self, images=None):
|
|
return self.nodes.GemmaContext(
|
|
user_prompt="A short verified scene.",
|
|
images=images,
|
|
source="image" if images is not None else "none",
|
|
media_metadata={
|
|
"source": "image" if images is not None else "none",
|
|
"reference_image_count": 1 if images is not None else 0,
|
|
},
|
|
)
|
|
|
|
def target(self):
|
|
return self.nodes.TargetProfileConfig(target_profile="ltx")
|
|
|
|
@staticmethod
|
|
def successful_result():
|
|
report = {
|
|
"schema": "dg-grounding-report/1",
|
|
"config": {"mode": "strict"},
|
|
"analysis_status": "not_applicable",
|
|
"decision": "not_applicable",
|
|
"would_block": False,
|
|
"grounding_guard_would_block": False,
|
|
"blocked_reasons": [],
|
|
"attempt_count": 0,
|
|
"retry_reasons": [],
|
|
"model_call_count": 1,
|
|
}
|
|
metadata = {
|
|
"ready_for_generation": True,
|
|
"json_parse_valid": True,
|
|
"used_template_fallback": False,
|
|
"salvage_warning": "",
|
|
"grounding_guard": report,
|
|
}
|
|
packet = {
|
|
"ltx_prompt": "A complete verified local video prompt.",
|
|
"ideogram_prompt": "",
|
|
"minimax_h3_prompt": "",
|
|
"negative_prompt": "",
|
|
"scene_segments": [],
|
|
"metadata": metadata,
|
|
}
|
|
return packet, "raw", "reasoning", "", True, "", metadata
|
|
|
|
def generate(self, node, runtime, *, guard=None, cache_mode="reuse"):
|
|
return node.generate(
|
|
runtime,
|
|
self.context(),
|
|
self.target(),
|
|
temperature=0.45,
|
|
creativity_mode="editorial",
|
|
creative_strength=0.6,
|
|
thinking_mode="off",
|
|
max_new_tokens=2048,
|
|
director_cache_mode=cache_mode,
|
|
grounding_guard_config=guard or {"mode": "strict"},
|
|
unique_id="node-does-not-affect-disk-key",
|
|
)
|
|
|
|
def test_second_call_hits_disk_without_running_director_again(self) -> None:
|
|
fixed_key = "a" * 64
|
|
with tempfile.TemporaryDirectory() as temporary, patch.dict(
|
|
os.environ,
|
|
{"DG_DIRECTOR_CACHE_DIR": temporary},
|
|
), patch.object(
|
|
self.nodes,
|
|
"_director_cache_key",
|
|
return_value=fixed_key,
|
|
), patch.object(
|
|
self.nodes,
|
|
"_run_generation_packet",
|
|
side_effect=lambda *_args, **_kwargs: copy.deepcopy(self.successful_result()),
|
|
) as director:
|
|
node = self.nodes.DiffusionGemmaCoTGenerator()
|
|
first = self.generate(node, self.runtime())
|
|
second = self.generate(node, self.runtime())
|
|
|
|
self.assertEqual(director.call_count, 1)
|
|
self.assertEqual(json.loads(first[3])["director_cache"]["status"], "stored")
|
|
second_metadata = json.loads(second[3])
|
|
self.assertTrue(second_metadata["director_cache"]["hit"])
|
|
self.assertTrue(second_metadata["director_runtime"]["cache_hit"])
|
|
self.assertEqual(second_metadata["director_runtime"]["generation_seconds"], 0.0)
|
|
self.assertEqual(second[1], "reasoning")
|
|
self.assertEqual(second[2], "raw")
|
|
|
|
def test_cache_hit_honors_unload_after_run(self) -> None:
|
|
fixed_key = "b" * 64
|
|
with tempfile.TemporaryDirectory() as temporary, patch.dict(
|
|
os.environ,
|
|
{"DG_DIRECTOR_CACHE_DIR": temporary},
|
|
), patch.object(
|
|
self.nodes,
|
|
"_director_cache_key",
|
|
return_value=fixed_key,
|
|
), patch.object(
|
|
self.nodes,
|
|
"_run_generation_packet",
|
|
side_effect=lambda *_args, **_kwargs: copy.deepcopy(self.successful_result()),
|
|
), patch.object(self.nodes, "_release_transformers_runtime") as release:
|
|
node = self.nodes.DiffusionGemmaCoTGenerator()
|
|
runtime = self.runtime(unload_policy="unload_after_run")
|
|
self.generate(node, runtime)
|
|
self.generate(node, runtime)
|
|
|
|
# The mocked first Director call does not own the real wrapper finally;
|
|
# the second, cached call explicitly enforces the lifecycle policy.
|
|
release.assert_called_once_with()
|
|
|
|
def test_corrupt_entry_causes_safe_recomputation(self) -> None:
|
|
fixed_key = "c" * 64
|
|
with tempfile.TemporaryDirectory() as temporary, patch.dict(
|
|
os.environ,
|
|
{"DG_DIRECTOR_CACHE_DIR": temporary},
|
|
), patch.object(
|
|
self.nodes,
|
|
"_director_cache_key",
|
|
return_value=fixed_key,
|
|
), patch.object(
|
|
self.nodes,
|
|
"_run_generation_packet",
|
|
side_effect=lambda *_args, **_kwargs: copy.deepcopy(self.successful_result()),
|
|
) as director:
|
|
node = self.nodes.DiffusionGemmaCoTGenerator()
|
|
self.generate(node, self.runtime())
|
|
self.nodes._director_disk_cache_path(fixed_key).write_text("{", encoding="utf-8")
|
|
result = self.generate(node, self.runtime())
|
|
|
|
self.assertEqual(director.call_count, 2)
|
|
self.assertFalse(json.loads(result[3])["director_cache"]["hit"])
|
|
|
|
def test_detailed_trace_bypasses_cache_side_effects(self) -> None:
|
|
with tempfile.TemporaryDirectory() as temporary, patch.dict(
|
|
os.environ,
|
|
{"DG_DIRECTOR_CACHE_DIR": temporary},
|
|
), patch.object(
|
|
self.nodes,
|
|
"_run_generation_packet",
|
|
side_effect=lambda *_args, **_kwargs: copy.deepcopy(self.successful_result()),
|
|
) as director:
|
|
node = self.nodes.DiffusionGemmaCoTGenerator()
|
|
guard = {"mode": "strict", "save_detailed_trace": True}
|
|
first = self.generate(node, self.runtime(), guard=guard)
|
|
second = self.generate(node, self.runtime(), guard=guard)
|
|
|
|
self.assertEqual(director.call_count, 2)
|
|
self.assertEqual(
|
|
json.loads(first[3])["director_cache"]["status"],
|
|
"bypassed_for_detailed_trace",
|
|
)
|
|
self.assertEqual(
|
|
json.loads(second[3])["director_cache"]["status"],
|
|
"bypassed_for_detailed_trace",
|
|
)
|
|
self.assertEqual(list(Path(temporary).glob("**/*.json")), [])
|
|
|
|
def test_reuse_is_stable_while_explicit_rerun_modes_bypass_native_cache(self) -> None:
|
|
first = self.nodes.DiffusionGemmaCoTGenerator.IS_CHANGED(
|
|
director_cache_mode="reuse"
|
|
)
|
|
second = self.nodes.DiffusionGemmaCoTGenerator.IS_CHANGED(
|
|
director_cache_mode="reuse"
|
|
)
|
|
self.assertEqual(first, second)
|
|
self.assertEqual(
|
|
first,
|
|
f"director-reuse:{self.nodes.DIRECTOR_EXECUTION_CONTRACT_REVISION}",
|
|
)
|
|
for mode in ("refresh", "off"):
|
|
with self.subTest(director_cache_mode=mode):
|
|
value = self.nodes.DiffusionGemmaCoTGenerator.IS_CHANGED(
|
|
director_cache_mode=mode
|
|
)
|
|
self.assertTrue(math.isnan(value))
|
|
|
|
def test_detailed_trace_guard_bypasses_native_cache(self) -> None:
|
|
stable = self.nodes.DiffusionGemmaGroundingGuardSettings.IS_CHANGED(
|
|
save_detailed_trace=False
|
|
)
|
|
self.assertEqual(stable, "grounding-guard-config-v1")
|
|
changed = self.nodes.DiffusionGemmaGroundingGuardSettings.IS_CHANGED(
|
|
save_detailed_trace=True
|
|
)
|
|
self.assertTrue(math.isnan(changed))
|
|
linked = self.nodes.DiffusionGemmaGroundingGuardSettings.IS_CHANGED(
|
|
save_detailed_trace=None
|
|
)
|
|
self.assertTrue(math.isnan(linked))
|
|
|
|
def test_comfy_signature_cache_isolates_an_ltx_only_edit_from_h3(self) -> None:
|
|
# Other test modules install lightweight Comfy stubs during discovery.
|
|
# Temporarily replace that module family with the real local ComfyUI
|
|
# implementation, then restore the discovery-time modules afterward.
|
|
module_prefixes = (
|
|
"comfy",
|
|
"comfy_api",
|
|
"comfy_execution",
|
|
"execution",
|
|
"nodes",
|
|
)
|
|
saved_modules = {
|
|
name: module
|
|
for name, module in tuple(sys.modules.items())
|
|
if any(name == prefix or name.startswith(f"{prefix}.") for prefix in module_prefixes)
|
|
}
|
|
for name in saved_modules:
|
|
sys.modules.pop(name, None)
|
|
|
|
def restore_modules() -> None:
|
|
for name in tuple(sys.modules):
|
|
if any(
|
|
name == prefix or name.startswith(f"{prefix}.")
|
|
for prefix in module_prefixes
|
|
):
|
|
sys.modules.pop(name, None)
|
|
sys.modules.update(saved_modules)
|
|
|
|
self.addCleanup(restore_modules)
|
|
import nodes as comfy_nodes
|
|
from comfy_execution.caching import CacheKeySetInputSignature, LRUCache
|
|
from comfy_execution.graph import DynamicPrompt
|
|
from execution import CacheEntry, IsChangedCache
|
|
|
|
class StableSource:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {"required": {"value": ("STRING",)}}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
FUNCTION = "run"
|
|
|
|
class StablePass:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {"required": {"source": ("STRING",)}}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
FUNCTION = "run"
|
|
|
|
class StableMerge:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"left": ("STRING",),
|
|
"right": ("STRING",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
FUNCTION = "run"
|
|
|
|
class StableOutput(StablePass):
|
|
OUTPUT_NODE = True
|
|
|
|
class_types = {
|
|
"DGCacheTestSource": StableSource,
|
|
"DGCacheTestPass": StablePass,
|
|
"DGCacheTestMerge": StableMerge,
|
|
"DGCacheTestOutput": StableOutput,
|
|
"DiffusionGemmaCoTGenerator": self.nodes.DiffusionGemmaCoTGenerator,
|
|
"DiffusionGemmaGroundingGuardSettings": (
|
|
self.nodes.DiffusionGemmaGroundingGuardSettings
|
|
),
|
|
}
|
|
base_prompt = {
|
|
"shared": {
|
|
"class_type": "DGCacheTestSource",
|
|
"inputs": {"value": "shared creative brief"},
|
|
},
|
|
"510": {
|
|
"class_type": "DGCacheTestSource",
|
|
"inputs": {"value": "LTX profile A"},
|
|
},
|
|
"186": {
|
|
"class_type": "DGCacheTestMerge",
|
|
"inputs": {
|
|
"left": ["shared", 0],
|
|
"right": ["510", 0],
|
|
},
|
|
},
|
|
"ltx_output": {
|
|
"class_type": "DGCacheTestOutput",
|
|
"inputs": {"source": ["186", 0]},
|
|
},
|
|
"182": {
|
|
"class_type": "DGCacheTestSource",
|
|
"inputs": {"value": "model config"},
|
|
},
|
|
"513": {
|
|
"class_type": "DGCacheTestPass",
|
|
"inputs": {"source": ["shared", 0]},
|
|
},
|
|
"515": {
|
|
"class_type": "DiffusionGemmaGroundingGuardSettings",
|
|
"inputs": {"save_detailed_trace": False},
|
|
},
|
|
"518": {
|
|
"class_type": "DGCacheTestSource",
|
|
"inputs": {"value": "H3 target profile"},
|
|
},
|
|
"516": {
|
|
"class_type": "DiffusionGemmaCoTGenerator",
|
|
"inputs": {
|
|
"model_config": ["182", 0],
|
|
"gemma_context": ["513", 0],
|
|
"target_profile_config": ["518", 0],
|
|
"grounding_guard_config": ["515", 0],
|
|
"temperature": 0.45,
|
|
"creativity_mode": "wild",
|
|
"creative_strength": 0.2,
|
|
"thinking_mode": "off",
|
|
"max_new_tokens": 2048,
|
|
"director_cache_mode": "reuse",
|
|
},
|
|
},
|
|
"517": {
|
|
"class_type": "DGCacheTestPass",
|
|
"inputs": {"source": ["516", 0]},
|
|
},
|
|
"h3_output": {
|
|
"class_type": "DGCacheTestOutput",
|
|
"inputs": {"source": ["517", 0]},
|
|
},
|
|
}
|
|
|
|
async def cache_hits_after(changed_prompt):
|
|
cache = LRUCache(CacheKeySetInputSignature, max_size=1)
|
|
first = DynamicPrompt(copy.deepcopy(base_prompt))
|
|
await cache.set_prompt(
|
|
first,
|
|
base_prompt.keys(),
|
|
IsChangedCache("first", first, cache),
|
|
)
|
|
cache.clean_unused()
|
|
for node_id in base_prompt:
|
|
cache.set_local(node_id, CacheEntry(ui=None, outputs=[node_id]))
|
|
|
|
second = DynamicPrompt(changed_prompt)
|
|
await cache.set_prompt(
|
|
second,
|
|
changed_prompt.keys(),
|
|
IsChangedCache("second", second, cache),
|
|
)
|
|
cache.clean_unused()
|
|
return {
|
|
node_id: cache.get_local(node_id) is not None
|
|
for node_id in changed_prompt
|
|
}
|
|
|
|
with patch.dict(comfy_nodes.NODE_CLASS_MAPPINGS, class_types):
|
|
ltx_edit = copy.deepcopy(base_prompt)
|
|
ltx_edit["510"]["inputs"]["value"] = "LTX profile B"
|
|
ltx_hits = asyncio.run(cache_hits_after(ltx_edit))
|
|
self.assertFalse(ltx_hits["510"])
|
|
self.assertFalse(ltx_hits["186"])
|
|
self.assertFalse(ltx_hits["ltx_output"])
|
|
for node_id in ("shared", "513", "515", "516", "517", "h3_output"):
|
|
self.assertTrue(ltx_hits[node_id], node_id)
|
|
|
|
shared_edit = copy.deepcopy(base_prompt)
|
|
shared_edit["shared"]["inputs"]["value"] = "changed shared brief"
|
|
shared_hits = asyncio.run(cache_hits_after(shared_edit))
|
|
for node_id in ("shared", "513", "516", "517", "h3_output"):
|
|
self.assertFalse(shared_hits[node_id], node_id)
|
|
|
|
trace_edit = copy.deepcopy(base_prompt)
|
|
trace_edit["515"]["inputs"]["save_detailed_trace"] = True
|
|
trace_hits = asyncio.run(cache_hits_after(trace_edit))
|
|
for node_id in ("515", "516", "517", "h3_output"):
|
|
self.assertFalse(trace_hits[node_id], node_id)
|
|
|
|
refresh_edit = copy.deepcopy(base_prompt)
|
|
refresh_edit["516"]["inputs"]["director_cache_mode"] = "refresh"
|
|
refresh_hits = asyncio.run(cache_hits_after(refresh_edit))
|
|
for node_id in ("516", "517", "h3_output"):
|
|
self.assertFalse(refresh_hits[node_id], node_id)
|
|
|
|
def test_cache_key_tracks_pixels_and_semantics_but_not_unload_policy(self) -> None:
|
|
with tempfile.TemporaryDirectory() as temporary:
|
|
checkpoint = Path(temporary) / "model.gguf"
|
|
checkpoint.write_bytes(b"small-test-checkpoint")
|
|
images = torch.zeros((1, 2, 2, 3), dtype=torch.float32)
|
|
context = self.context(images)
|
|
target = self.target()
|
|
common = dict(
|
|
temperature=0.45,
|
|
creativity_mode="editorial",
|
|
creative_strength=0.6,
|
|
thinking_mode="off",
|
|
max_new_tokens=2048,
|
|
)
|
|
key = self.nodes._director_cache_key(
|
|
self.runtime(model_path=str(checkpoint)),
|
|
context,
|
|
target,
|
|
{"mode": "strict", "seed": 7},
|
|
**common,
|
|
)
|
|
unload_key = self.nodes._director_cache_key(
|
|
self.runtime(
|
|
unload_policy="unload_after_run",
|
|
model_path=str(checkpoint),
|
|
),
|
|
context,
|
|
target,
|
|
{"mode": "strict", "seed": 7},
|
|
**common,
|
|
)
|
|
changed_images = images.clone()
|
|
changed_images[0, 0, 0, 0] = 1.0
|
|
pixel_key = self.nodes._director_cache_key(
|
|
self.runtime(model_path=str(checkpoint)),
|
|
self.context(changed_images),
|
|
target,
|
|
{"mode": "strict", "seed": 7},
|
|
**common,
|
|
)
|
|
duration_target = self.nodes.TargetProfileConfig(
|
|
target_profile="ltx",
|
|
target_duration_seconds=15.0,
|
|
)
|
|
target_key = self.nodes._director_cache_key(
|
|
self.runtime(model_path=str(checkpoint)),
|
|
context,
|
|
duration_target,
|
|
{"mode": "strict", "seed": 7},
|
|
**common,
|
|
)
|
|
|
|
self.assertEqual(key, unload_key)
|
|
self.assertNotEqual(key, pixel_key)
|
|
self.assertNotEqual(key, target_key)
|
|
|
|
def test_rank_three_tensor_hashes_as_one_image(self) -> None:
|
|
rank_three = torch.zeros((2, 3, 3), dtype=torch.float32)
|
|
rank_four = torch.zeros((2, 2, 3, 3), dtype=torch.float32)
|
|
self.assertEqual(len(self.nodes._ordered_tensor_content_hashes(rank_three)), 1)
|
|
self.assertEqual(len(self.nodes._ordered_tensor_content_hashes(rank_four)), 2)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|