digital-garbageandClaude Sonnet 5
79ca8cfb6b
Loader-level attention option: mean-pool Q and a smoothed K into blocks, score them with one small matmul, attend only the top 1-sparsity fraction of key blocks per query block, via a vendored Triton kernel. Runs alongside the chosen attention backend rather than replacing it (one override slot; the backend becomes SLA's dense fall-through for the text refiner, masked calls, and trailing dense steps), and refuses cleanly -- naming why -- on anything that is not MiniMax H3 or a machine without CUDA+Triton. modules/loaders/sla_block_map.py and sla_kernel.py are near-verbatim ports of v4's vendored files (no FunPack-specific imports, so nothing to adapt); sla_attention.py is the loader-facing glue (H3 gate, override composition, per-step wrapper), same content as v4's with a "not a module" note added to match gguf_support.py's convention. Wired into FunPackDiffusionModelLoader as six new optional widgets, following v4's loaders.py wiring but expressed in v5's io.Schema style. The Triton kernel itself is unexercisable on this Mac (no CUDA) -- matches how real generation is verified in this project (rented GPU, not local). Ported the CPU-testable half of v4's test suite (the override contract, the H3 gate, the per-step wrapper); the three kernel-correctness tests carry the same CUDA+Triton skip guard v4 used. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
104 lines
3.7 KiB
Python
104 lines
3.7 KiB
Python
"""The nodes FunPack actually ships, checked against ComfyUI itself.
|
|
|
|
Needs a ComfyUI source tree, because a node's schema is only meaningful in terms
|
|
of comfy_api's own types. Skipped where there is none.
|
|
"""
|
|
|
|
import pytest
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def shipped(comfyui):
|
|
from core import nodes as nodes_mod
|
|
from core import registry as registry_mod
|
|
|
|
registry = registry_mod.scan()
|
|
nodes, rejected = nodes_mod.collect(registry)
|
|
return registry, nodes, rejected
|
|
|
|
|
|
def _schema(node):
|
|
return node.GET_SCHEMA()
|
|
|
|
|
|
def _input_ids(schema):
|
|
return [getattr(i, "id", None) for i in (schema.inputs or [])]
|
|
|
|
|
|
def test_nothing_shipped_is_rejected(shipped):
|
|
_, _, rejected = shipped
|
|
assert rejected == [], f"a shipped node did not register: {rejected}"
|
|
|
|
|
|
def test_every_shipped_node_is_a_real_comfy_node(shipped, comfyui):
|
|
from comfy_api.latest import io
|
|
_, nodes, _ = shipped
|
|
for node in nodes:
|
|
assert issubclass(node, io.ComfyNode), f"{node.__name__} is not an io.ComfyNode"
|
|
assert isinstance(_schema(node), io.Schema)
|
|
|
|
|
|
# Each entry is a promise to everyone's saved workflows: `node_id` is stored as
|
|
# `class_type`, and input ids are stored as key names. Changing either silently
|
|
# breaks graphs that already exist, so this test is meant to FAIL on a rename and
|
|
# make that a decision rather than an accident.
|
|
FROZEN = {
|
|
"FunPackCheckpointLoader": ["ckpt_name"],
|
|
"FunPackCLIPLoader": [
|
|
"clip_name1", "clip_name2", "clip_name3", "clip_name4", "type", "device",
|
|
],
|
|
"FunPackSampler": [
|
|
"model", "positive", "negative", "latent", "seed", "steps", "cfg",
|
|
"sampler_name", "scheduler", "denoise", "settings",
|
|
],
|
|
"FunPackLoadModifiers": ["model", "settings"],
|
|
"FunPackModifierSettings": ["settings"],
|
|
"FunPackDecode": ["samples", "vae", "model", "audio_vae"],
|
|
"FunPackEmptyLatent": ["model", "width", "height", "length", "batch_size"],
|
|
"FunPackDiffusionModelLoader": [
|
|
"model_name", "weight_dtype", "compute_dtype", "attention", "fp16_accumulation",
|
|
"sla", "sla_sparsity", "sla_block_size", "sla_protect_audio", "sla_min_seq_len",
|
|
"sla_dense_last_steps",
|
|
],
|
|
"FunPackLoraLoader": [
|
|
"model", "lora_name", "strength_model", "clip", "strength_clip",
|
|
],
|
|
"FunPackVAELoader": ["vae_name", "dtype"],
|
|
"FunPackLoadMedia": ["media_id"],
|
|
}
|
|
|
|
|
|
def test_shipped_node_ids_and_inputs_are_frozen(shipped):
|
|
_, nodes, _ = shipped
|
|
actual = {_schema(n).node_id: _input_ids(_schema(n)) for n in nodes}
|
|
assert actual == FROZEN, (
|
|
"a shipped node's id or input names changed. Every saved workflow refers "
|
|
"to these by name; update FROZEN only if you intend to break them."
|
|
)
|
|
|
|
|
|
def test_a_setting_never_becomes_a_node_socket(shipped):
|
|
"""The rule that keeps the sampler from growing 70 widgets, and keeps a new
|
|
setting from rotting saved graphs: settings are a payload, not sockets."""
|
|
registry, _, _ = shipped
|
|
for spec in registry.specs.values():
|
|
keys = set(spec.settings)
|
|
for node in spec.nodes:
|
|
ids = set(_input_ids(_schema(node))) - {None}
|
|
overlap = keys & ids
|
|
assert not overlap, (
|
|
f"{spec.id}'s node {node.__name__} exposes settings as sockets: "
|
|
f"{sorted(overlap)}"
|
|
)
|
|
|
|
|
|
def test_the_extension_hands_comfyui_exactly_what_was_collected(shipped, comfyui):
|
|
import asyncio
|
|
from core import nodes as nodes_mod
|
|
|
|
extension = nodes_mod.extension()
|
|
assert extension is not None
|
|
listed = asyncio.run(extension.get_node_list())
|
|
_, nodes, _ = shipped
|
|
assert {_schema(n).node_id for n in listed} == {_schema(n).node_id for n in nodes}
|