560 lines
21 KiB
Python
560 lines
21 KiB
Python
"""Headless test suite for the H3 Ref2VA Accelerator runtime.
|
|
|
|
Runs without ComfyUI or a GPU: `comfy` is stubbed, tensors are tiny CPU tensors,
|
|
and the block patches are driven the same way ComfyUI's dit patches_replace
|
|
mechanism drives them (args dict with "img", extra dict with "original_block").
|
|
|
|
Usage: python tests/test_runtime.py
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import importlib.util
|
|
import logging
|
|
import math
|
|
import sys
|
|
import types
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Stub comfy before importing nodes.py
|
|
# ---------------------------------------------------------------------------
|
|
|
|
comfy_pkg = types.ModuleType("comfy")
|
|
model_prefetch = types.ModuleType("comfy.model_prefetch")
|
|
|
|
|
|
def _stub_make_prefetch_queue(queue, device, transformer_options):
|
|
return list(queue)
|
|
|
|
|
|
model_prefetch.make_prefetch_queue = _stub_make_prefetch_queue
|
|
|
|
patcher_extension = types.ModuleType("comfy.patcher_extension")
|
|
|
|
|
|
class WrappersMP:
|
|
DIFFUSION_MODEL = "diffusion_model"
|
|
OUTER_SAMPLE = "outer_sample"
|
|
|
|
|
|
patcher_extension.WrappersMP = WrappersMP
|
|
|
|
comfy_pkg.model_prefetch = model_prefetch
|
|
comfy_pkg.patcher_extension = patcher_extension
|
|
sys.modules["comfy"] = comfy_pkg
|
|
sys.modules["comfy.model_prefetch"] = model_prefetch
|
|
sys.modules["comfy.patcher_extension"] = patcher_extension
|
|
|
|
import torch # noqa: E402
|
|
|
|
_NODES_PATH = Path(__file__).resolve().parent.parent / "nodes.py"
|
|
_spec = importlib.util.spec_from_file_location("h3_ref2va_nodes", _NODES_PATH)
|
|
nodes = importlib.util.module_from_spec(_spec)
|
|
sys.modules["h3_ref2va_nodes"] = nodes
|
|
_spec.loader.exec_module(nodes)
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Fixtures
|
|
# ---------------------------------------------------------------------------
|
|
|
|
ROWS, DIM = 16, 6
|
|
SEGMENTS = [(0, 8, "video"), (8, 12, "audio"), (12, 14, "ref_img"), (14, 16, "ref_audio")]
|
|
|
|
|
|
def make_layout(latent_frames=4):
|
|
return SimpleNamespace(segments=list(SEGMENTS), signature=(1, latent_frames))
|
|
|
|
|
|
def payload(layout=None):
|
|
return {"refs": True, "layout": layout if layout is not None else make_layout()}
|
|
|
|
|
|
def const_blocks(seed=0, dtype=torch.float32):
|
|
"""Blocks that add a fixed per-block delta in place (input-independent tail)."""
|
|
g = torch.Generator().manual_seed(seed)
|
|
deltas = [torch.randn(ROWS, DIM, generator=g, dtype=torch.float32).to(dtype) for _ in range(4)]
|
|
|
|
def make(i):
|
|
def block(args):
|
|
args["img"].add_(deltas[i])
|
|
return {"img": args["img"]}
|
|
return block
|
|
|
|
return [make(i) for i in range(4)], deltas
|
|
|
|
|
|
def doubling_block0_blocks(seed=0):
|
|
"""Block 0 doubles the hidden state (residual == original input); blocks 1-3 add constants."""
|
|
g = torch.Generator().manual_seed(seed)
|
|
deltas = [torch.randn(ROWS, DIM, generator=g) for _ in range(4)]
|
|
|
|
def block0(args):
|
|
args["img"].add_(args["img"]) # x -> 2x, residual == x
|
|
return {"img": args["img"]}
|
|
|
|
def make(i):
|
|
def block(args):
|
|
args["img"].add_(deltas[i])
|
|
return {"img": args["img"]}
|
|
return block
|
|
|
|
return [block0, make(1), make(2), make(3)], deltas
|
|
|
|
|
|
def make_runtime(config=None, storage=None, start_sigma=10.0, end_sigma=0.0,
|
|
blocks=None, debug=False, tail_rescale=False,
|
|
cpu_tail_compute=None):
|
|
config = config or nodes.PRESETS[nodes.BALANCED]
|
|
storage = storage or nodes.CPU_STORAGE
|
|
blocks = blocks if blocks is not None else [object() for _ in range(4)]
|
|
return nodes.Ref2VAUltraSafeBlockCacheRuntime(
|
|
config=config, start_sigma=start_sigma, end_sigma=end_sigma,
|
|
block_count=4, storage=storage, debug=debug,
|
|
block_modules=blocks, tail_rescale=tail_rescale,
|
|
cpu_tail_compute=cpu_tail_compute or nodes.CPU_TAIL_SAFE,
|
|
)
|
|
|
|
|
|
def native(blocks, x):
|
|
img = x.clone()
|
|
for b in blocks:
|
|
img = b({"img": img})["img"]
|
|
return img
|
|
|
|
|
|
def run_step(runtime, patches, blocks, x, sigma, pl):
|
|
"""Drive one model call through the diffusion wrapper + block patches,
|
|
mirroring ComfyUI's calling convention."""
|
|
dwrap = nodes.make_diffusion_wrapper(runtime)
|
|
timestep = torch.tensor([sigma * 1000.0])
|
|
topts = {"uuids": ("u1",), "sigmas": torch.tensor([sigma])}
|
|
|
|
def executor(x_in, t_in, control, transformer_options, minimax_payload=None):
|
|
img = x_in.clone()
|
|
for i, p in enumerate(patches):
|
|
img = p({"img": img}, {"original_block": blocks[i]})["img"]
|
|
return img
|
|
|
|
return dwrap(executor, x, timestep, None, topts, pl)
|
|
|
|
|
|
def build(runtime):
|
|
return [nodes.make_block_patch(runtime, i, 3) for i in range(4)]
|
|
|
|
|
|
class LogCapture(logging.Handler):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.records = []
|
|
|
|
def emit(self, record):
|
|
self.records.append(record.getMessage())
|
|
|
|
def __enter__(self):
|
|
logging.getLogger().addHandler(self)
|
|
return self
|
|
|
|
def __exit__(self, *exc):
|
|
logging.getLogger().removeHandler(self)
|
|
return False
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def test_exact_cache_reconstruction_both_storages():
|
|
"""With input-independent blocks 1-3, a cached step must equal native bit-for-bit."""
|
|
for storage in (nodes.CPU_STORAGE, nodes.GPU_STORAGE):
|
|
blocks, _ = const_blocks()
|
|
rt = make_runtime(storage=storage, blocks=blocks)
|
|
patches = build(rt)
|
|
g = torch.Generator().manual_seed(7)
|
|
x1 = torch.randn(ROWS, DIM, generator=g)
|
|
x2 = torch.randn(ROWS, DIM, generator=g)
|
|
|
|
out1 = run_step(rt, patches, blocks, x1, 0.9, payload())
|
|
assert torch.equal(out1, native(blocks, x1)), "full step must match native"
|
|
assert rt.full_steps == 1 and rt.cached_steps == 0
|
|
|
|
out2 = run_step(rt, patches, blocks, x2, 0.8, payload())
|
|
assert rt.cached_steps == 1, f"expected a cache hit ({storage})"
|
|
# Exact in exact arithmetic; float32 add-order differs, so compare to rounding.
|
|
ref = native(blocks, x2)
|
|
assert torch.allclose(out2, ref, atol=1e-5, rtol=1e-5), (
|
|
f"cached step diverged ({storage}): max abs diff "
|
|
f"{(out2 - ref).abs().max().item():.3e}")
|
|
|
|
|
|
def test_max_consecutive_hits_enforced():
|
|
blocks, _ = const_blocks()
|
|
rt = make_runtime(blocks=blocks)
|
|
patches = build(rt)
|
|
x = torch.randn(ROWS, DIM)
|
|
sigmas = [0.9, 0.8, 0.7, 0.6]
|
|
for s in sigmas:
|
|
run_step(rt, patches, blocks, x.clone(), s, payload())
|
|
# zero-diff every step: hit pattern must alternate full/cache/full/cache
|
|
assert rt.full_steps == 2 and rt.cached_steps == 2
|
|
assert rt.cache_step_numbers == [2, 4]
|
|
|
|
|
|
def test_window_gating():
|
|
blocks, _ = const_blocks()
|
|
rt = make_runtime(blocks=blocks, start_sigma=0.5, end_sigma=0.2)
|
|
patches = build(rt)
|
|
x = torch.randn(ROWS, DIM)
|
|
run_step(rt, patches, blocks, x.clone(), 0.9, payload())
|
|
run_step(rt, patches, blocks, x.clone(), 0.8, payload()) # outside window
|
|
assert rt.cached_steps == 0, "no hits outside the sigma window"
|
|
run_step(rt, patches, blocks, x.clone(), 0.4, payload()) # inside window
|
|
assert rt.cached_steps == 1
|
|
|
|
|
|
def test_sigma_increase_resets_context():
|
|
blocks, _ = const_blocks()
|
|
rt = make_runtime(blocks=blocks)
|
|
patches = build(rt)
|
|
x = torch.randn(ROWS, DIM)
|
|
run_step(rt, patches, blocks, x.clone(), 0.9, payload())
|
|
run_step(rt, patches, blocks, x.clone(), 0.8, payload())
|
|
assert rt.cached_steps == 1
|
|
run_step(rt, patches, blocks, x.clone(), 0.95, payload()) # new run: sigma went up
|
|
assert rt.cached_steps == 1, "reset step must run full"
|
|
assert rt.full_steps == 2
|
|
|
|
|
|
def test_input_signature_change_resets_context():
|
|
blocks, _ = const_blocks()
|
|
rt = make_runtime(blocks=blocks)
|
|
patches = build(rt)
|
|
run_step(rt, patches, blocks, torch.randn(ROWS, DIM), 0.9, payload())
|
|
# Same declining sigma, different shape: must not attempt to compare/cache.
|
|
blocks2, _ = const_blocks()
|
|
|
|
def run_shape(rows):
|
|
dwrap = nodes.make_diffusion_wrapper(rt)
|
|
topts = {"uuids": ("u1",), "sigmas": torch.tensor([0.8])}
|
|
|
|
def executor(x_in, t_in, control, transformer_options, minimax_payload=None):
|
|
return x_in.clone()
|
|
|
|
dwrap(executor, torch.randn(rows, DIM), torch.tensor([800.0]), None, topts, payload())
|
|
|
|
run_shape(ROWS + 4)
|
|
ctx = rt.contexts[("u1",)]
|
|
assert ctx.previous_first_residual is None or ctx.previous_first_residual.shape[0] != ROWS
|
|
assert rt.cached_steps == 0
|
|
|
|
|
|
def test_no_refs_idle_and_summary():
|
|
blocks, _ = const_blocks()
|
|
rt = make_runtime(blocks=blocks)
|
|
patches = build(rt)
|
|
x = torch.randn(ROWS, DIM)
|
|
out = run_step(rt, patches, blocks, x, 0.9, {}) # no payload
|
|
assert torch.equal(out, native(blocks, x)), "idle path must be native"
|
|
assert rt.full_steps == 0 and rt.cached_steps == 0
|
|
assert rt.steps_without_refs == 1
|
|
s = rt.summary("test")
|
|
assert "without a Ref2VA payload" in s, s
|
|
|
|
|
|
def test_mixed_idle_summary_line():
|
|
blocks, _ = const_blocks()
|
|
rt = make_runtime(blocks=blocks)
|
|
patches = build(rt)
|
|
x = torch.randn(ROWS, DIM)
|
|
run_step(rt, patches, blocks, x.clone(), 0.9, {})
|
|
run_step(rt, patches, blocks, x.clone(), 0.8, payload())
|
|
run_step(rt, patches, blocks, x.clone(), 0.7, payload())
|
|
s = rt.summary("test")
|
|
assert "steps without Ref2VA payload: 1" in s, s
|
|
assert f"CPU tail compute: {nodes.CPU_TAIL_SAFE}" in s, s
|
|
assert "fast GPU tail path:" not in s, s # Safe CPU never probes/uses the GPU fast path.
|
|
|
|
|
|
|
|
def test_cpu_tail_safe_skips_gpu_headroom_probe():
|
|
blocks, _ = const_blocks()
|
|
rt = make_runtime(blocks=blocks, storage=nodes.CPU_STORAGE,
|
|
cpu_tail_compute=nodes.CPU_TAIL_SAFE)
|
|
patches = build(rt)
|
|
old = nodes._gpu_headroom_ok
|
|
|
|
def should_not_run(*_a, **_k):
|
|
raise AssertionError("Safe CPU must not probe GPU headroom")
|
|
|
|
nodes._gpu_headroom_ok = should_not_run
|
|
try:
|
|
run_step(rt, patches, blocks, torch.randn(ROWS, DIM), 0.9, payload())
|
|
finally:
|
|
nodes._gpu_headroom_ok = old
|
|
assert rt.full_steps == 1
|
|
assert rt.gpu_tail_steps == 0
|
|
|
|
|
|
def test_cpu_tail_auto_consults_gpu_headroom_probe():
|
|
blocks, _ = const_blocks()
|
|
rt = make_runtime(blocks=blocks, storage=nodes.CPU_STORAGE,
|
|
cpu_tail_compute=nodes.CPU_TAIL_AUTO)
|
|
patches = build(rt)
|
|
old = nodes._gpu_headroom_ok
|
|
called = {"value": False}
|
|
|
|
def probe(*_a, **_k):
|
|
called["value"] = True
|
|
return False # headless CPU test intentionally exercises fallback.
|
|
|
|
nodes._gpu_headroom_ok = probe
|
|
try:
|
|
run_step(rt, patches, blocks, torch.randn(ROWS, DIM), 0.9, payload())
|
|
finally:
|
|
nodes._gpu_headroom_ok = old
|
|
assert called["value"], "Auto mode must consult the headroom gate"
|
|
summary = rt.summary("test")
|
|
assert f"CPU tail compute: {nodes.CPU_TAIL_AUTO}" in summary
|
|
assert "fast GPU tail path: 0/1 full steps" in summary
|
|
|
|
|
|
def test_production_widget_defaults():
|
|
spec = nodes.ApplyH3Ref2VAUltraSafeBlockCache.INPUT_TYPES()
|
|
mode_choices, mode_opts = spec["required"]["mode"]
|
|
storage_choices, storage_opts = spec["required"]["cache_storage"]
|
|
tail_choices, tail_opts = spec["optional"]["cpu_tail_compute"]
|
|
_rescale_type, rescale_opts = spec["optional"]["tail_rescale"]
|
|
|
|
assert mode_choices[0] == nodes.BALANCED
|
|
assert mode_opts["default"] == nodes.BALANCED
|
|
assert storage_choices[0] == nodes.CPU_STORAGE
|
|
assert storage_opts["default"] == nodes.CPU_STORAGE
|
|
assert tail_choices == [nodes.CPU_TAIL_SAFE, nodes.CPU_TAIL_AUTO]
|
|
assert tail_opts["default"] == nodes.CPU_TAIL_SAFE
|
|
assert tail_opts.get("advanced") is True
|
|
assert rescale_opts["default"] is False
|
|
assert rescale_opts.get("advanced") is True
|
|
|
|
|
|
def test_missing_metric_warning_once():
|
|
blocks, _ = const_blocks()
|
|
rt = make_runtime(blocks=blocks)
|
|
patches = build(rt)
|
|
bad_layout = make_layout(latent_frames=5) # 8 video rows % 5 != 0 -> temporal None
|
|
x = torch.randn(ROWS, DIM)
|
|
with LogCapture() as cap:
|
|
for s in (0.9, 0.8, 0.7):
|
|
run_step(rt, patches, blocks, x.clone(), s, payload(bad_layout))
|
|
warnings = [m for m in cap.records if "required guard metric(s) unavailable" in m]
|
|
assert len(warnings) == 1, f"expected exactly one warning, got {len(warnings)}"
|
|
assert "temporal" in warnings[0]
|
|
assert rt.cached_steps == 0, "caching must stay disabled when a required metric is missing"
|
|
|
|
|
|
def test_fp32_metric_precision():
|
|
torch.manual_seed(11)
|
|
cur = torch.randn(4096, dtype=torch.float32).to(torch.bfloat16)
|
|
prev = (torch.randn(4096, dtype=torch.float32) * 0.9).to(torch.bfloat16)
|
|
got = nodes._ratio(cur, prev)
|
|
ref = float((cur.double() - prev.double()).abs().mean() / prev.double().abs().mean().clamp(min=1e-8))
|
|
rel = abs(got - ref) / max(abs(ref), 1e-12)
|
|
assert rel < 1e-4, f"global ratio drifted from fp64 reference: rel={rel:.2e}"
|
|
# Fast path and ranged path must now agree (both fp32-accumulated).
|
|
ranged = nodes._ratio(cur, prev, [(0, 4096)])
|
|
assert abs(got - ranged) / max(abs(got), 1e-12) < 1e-4, (got, ranged)
|
|
|
|
|
|
def test_temporal_ratio_math():
|
|
layout = make_layout(latent_frames=4)
|
|
prev = torch.ones(ROWS, DIM)
|
|
cur = prev.clone()
|
|
cur[2:4] += 0.5 # frame 1 of the video segment (rows_per_frame == 2)
|
|
got = nodes._temporal_video_ratio(cur, prev, layout)
|
|
assert got is not None and abs(got - 0.5) < 1e-6, got
|
|
|
|
|
|
def test_tail_rescale_math_and_cache_integrity():
|
|
# Block 0 doubles (residual == input), so per-segment residual energy ratios are controllable.
|
|
for storage in (nodes.CPU_STORAGE, nodes.GPU_STORAGE):
|
|
blocks, deltas = doubling_block0_blocks()
|
|
cfg = nodes.PresetConfig(
|
|
global_threshold=0.5, video_threshold=0.5, audio_threshold=0.5,
|
|
visual_ref_threshold=0.5, audio_ref_threshold=0.5, temporal_threshold=0.5,
|
|
start_percent=0.0, end_percent=1.0, max_consecutive_hits=1,
|
|
)
|
|
rt = make_runtime(config=cfg, storage=storage, blocks=blocks, tail_rescale=True)
|
|
patches = build(rt)
|
|
|
|
x1 = torch.rand(ROWS, DIM) + 0.5 # strictly positive for clean energy ratios
|
|
factors = {"video": 1.05, "audio": 1.30, "ref_img": 1.04, "ref_audio": 1.03}
|
|
x2 = x1.clone()
|
|
for a, b, kind in SEGMENTS:
|
|
x2[a:b] *= factors[kind]
|
|
|
|
run_step(rt, patches, blocks, x1.clone(), 0.9, payload())
|
|
tail_before = rt.contexts[("u1",)].remaining_blocks_residual.clone()
|
|
out2 = run_step(rt, patches, blocks, x2.clone(), 0.8, payload())
|
|
assert rt.cached_steps == 1, f"expected hit ({storage})"
|
|
|
|
expected = x2 * 2 # block-0 output
|
|
tail = sum(deltas[1:])
|
|
clamped = {"video": 1.05, "audio": 1.10, "ref_img": 1.04, "ref_audio": 1.03} # audio clamps
|
|
for a, b, kind in SEGMENTS:
|
|
expected[a:b] += tail[a:b] * clamped[kind]
|
|
assert torch.allclose(out2, expected, atol=1e-4), f"rescaled output mismatch ({storage})"
|
|
|
|
# The persistent cached tail must never be mutated by rescaling.
|
|
assert torch.equal(rt.contexts[("u1",)].remaining_blocks_residual, tail_before), storage
|
|
|
|
|
|
def test_observe_only_never_caches_but_measures():
|
|
blocks, _ = const_blocks()
|
|
rt = make_runtime(config=nodes.PRESETS[nodes.OBSERVE], blocks=blocks)
|
|
patches = build(rt)
|
|
x = torch.randn(ROWS, DIM)
|
|
for s in (0.9, 0.8, 0.7):
|
|
out = run_step(rt, patches, blocks, x.clone(), s, payload())
|
|
assert torch.equal(out, native(blocks, x))
|
|
assert rt.cached_steps == 0
|
|
assert len(rt.metric_history) == 2 # first step has nothing to compare against
|
|
|
|
|
|
def test_prefetch_capture_and_neutralize():
|
|
blocks, _ = const_blocks()
|
|
rt = make_runtime(blocks=blocks)
|
|
q = list(range(4)) # stand-in queue entries
|
|
# Wrong identity: not captured.
|
|
out = rt.capture_prefetch_queue(lambda queue, d, t: list(queue), [object()] * 4, None, {})
|
|
assert rt.prefetch_queue is None
|
|
# Matching identity: captured, then neutralized keeping entry 0.
|
|
out = rt.capture_prefetch_queue(lambda queue, d, t: q, blocks, None, {})
|
|
assert rt.prefetch_queue is q
|
|
rt.neutralize_tail_prefetch()
|
|
assert q[0] == 0 and all(v is None for v in q[1:])
|
|
assert rt.prefetch_suppressed_steps == 1
|
|
|
|
|
|
def test_sample_wrapper_without_prefetch_module():
|
|
blocks, _ = const_blocks()
|
|
rt = make_runtime(blocks=blocks)
|
|
old = nodes._PREFETCH_AVAILABLE
|
|
nodes._PREFETCH_AVAILABLE = False
|
|
try:
|
|
wrap = nodes.make_sample_wrapper(rt, nodes.BALANCED)
|
|
with LogCapture() as cap:
|
|
result = wrap(lambda *a, **k: "sentinel")
|
|
assert result == "sentinel"
|
|
assert any("without tail-prefetch suppression" in m for m in cap.records)
|
|
finally:
|
|
nodes._PREFETCH_AVAILABLE = old
|
|
|
|
|
|
def test_sample_wrapper_restores_prefetch_and_summarizes():
|
|
blocks, _ = const_blocks()
|
|
rt = make_runtime(blocks=blocks)
|
|
original = nodes.comfy.model_prefetch.make_prefetch_queue
|
|
wrap = nodes.make_sample_wrapper(rt, nodes.BALANCED)
|
|
|
|
def executor(*a, **k):
|
|
assert nodes.comfy.model_prefetch.make_prefetch_queue is not original, "patched during sampling"
|
|
return "ok"
|
|
|
|
with LogCapture() as cap:
|
|
assert wrap(executor) == "ok"
|
|
assert nodes.comfy.model_prefetch.make_prefetch_queue is original, "restored after sampling"
|
|
assert any("Accelerator v" in m and "result" in m for m in cap.records)
|
|
|
|
|
|
def _fake_model(block_count=4, existing_dit=None):
|
|
fake_blocks = [object() for _ in range(block_count)]
|
|
H3 = type("MiniMaxH3Model", (), {})
|
|
dm = H3()
|
|
dm.blocks = fake_blocks
|
|
sampling = SimpleNamespace(percent_to_sigma=lambda p: 1.0 - p)
|
|
|
|
class FakePatcher:
|
|
def __init__(self):
|
|
self.model_options = {"transformer_options": {"patches_replace": {"dit": dict(existing_dit or {})}}}
|
|
self.patches = {}
|
|
self.wrappers = {}
|
|
|
|
def get_model_object(self, name):
|
|
return {"diffusion_model": dm, "model_sampling": sampling}[name]
|
|
|
|
def clone(self):
|
|
return self
|
|
|
|
def set_model_patch_replace(self, fn, kind, block_kind, index):
|
|
self.patches[(kind, block_kind, index)] = fn
|
|
|
|
def add_wrapper_with_key(self, wrapper_type, key, fn):
|
|
self.wrappers[(wrapper_type, key)] = fn
|
|
|
|
return FakePatcher()
|
|
|
|
|
|
def test_apply_plumbing_and_conflict_detection():
|
|
node = nodes.ApplyH3Ref2VAUltraSafeBlockCache()
|
|
m = _fake_model()
|
|
(patched,) = node.apply(m, nodes.BALANCED, nodes.CPU_STORAGE, False, tail_rescale=True)
|
|
assert len(patched.patches) == 4
|
|
kinds = {k[0] for k in patched.wrappers}
|
|
assert kinds == {WrappersMP.DIFFUSION_MODEL, WrappersMP.OUTER_SAMPLE}
|
|
|
|
conflicted = _fake_model(existing_dit={("double_block", 1): lambda a, e: a})
|
|
try:
|
|
node.apply(conflicted, nodes.BALANCED, nodes.CPU_STORAGE, False)
|
|
raise AssertionError("expected conflict ValueError")
|
|
except ValueError as e:
|
|
assert "existing DiT block replacement" in str(e)
|
|
|
|
class NotH3:
|
|
blocks = [object(), object()]
|
|
bad = _fake_model()
|
|
bad.get_model_object = lambda name: {"diffusion_model": NotH3(), "model_sampling": None}[name]
|
|
try:
|
|
node.apply(bad, nodes.BALANCED, nodes.CPU_STORAGE, False)
|
|
raise AssertionError("expected class ValueError")
|
|
except ValueError as e:
|
|
assert "MiniMaxH3Model" in str(e)
|
|
|
|
try:
|
|
node.apply(_fake_model(), nodes.CUSTOM, nodes.CPU_STORAGE, False,
|
|
start_percent=0.9, end_percent=0.5)
|
|
raise AssertionError("expected window ValueError")
|
|
except ValueError as e:
|
|
assert "start_percent" in str(e)
|
|
|
|
|
|
def test_sigma_extraction_prefers_transformer_options():
|
|
rt = nodes.Ref2VAUltraSafeBlockCacheRuntime
|
|
t = torch.tensor([800.0])
|
|
assert abs(rt._extract_sigma(t, {"sigmas": torch.tensor([0.42])}) - 0.42) < 1e-6
|
|
assert abs(rt._extract_sigma(t, {}) - 0.8) < 1e-6 # x1000 fallback
|
|
assert abs(rt._extract_sigma(t, {"sigmas": torch.tensor([])}) - 0.8) < 1e-6
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Runner
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def main():
|
|
logging.getLogger().setLevel(logging.DEBUG)
|
|
logging.getLogger().handlers = [logging.NullHandler()]
|
|
tests = [(n, fn) for n, fn in sorted(globals().items()) if n.startswith("test_") and callable(fn)]
|
|
failed = 0
|
|
for name, fn in tests:
|
|
try:
|
|
fn()
|
|
print(f"PASS {name}")
|
|
except Exception as e: # noqa: BLE001
|
|
failed += 1
|
|
print(f"FAIL {name}: {e.__class__.__name__}: {e}")
|
|
print(f"\n{len(tests) - failed}/{len(tests)} passed")
|
|
return 1 if failed else 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|