diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 5ab513c..eeae02d 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -31,7 +31,35 @@ jobs: run: python -m compileall -q . - name: Run registry security style checks run: python -m ruff check --select E702 . - - name: Verify ComfyUI node-pack registration and schema - run: python -m pytest -q tests/test_nodes.py + - name: Smoke-test ComfyUI node-pack registration and schema + run: | + python - <<'PY' + import importlib.util + import sys + from pathlib import Path + + pack_dir = Path.cwd() + spec = importlib.util.spec_from_file_location( + "comfyui_hicache_smoke", + pack_dir / "__init__.py", + submodule_search_locations=[str(pack_dir)], + ) + if spec is None or spec.loader is None: + raise RuntimeError("could not load the node pack") + pack = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = pack + spec.loader.exec_module(pack) + + assert "HiCacheAccelerate" in pack.NODE_CLASS_MAPPINGS + assert "HiCacheAccelerate" in pack.NODE_DISPLAY_NAME_MAPPINGS + node = pack.NODE_CLASS_MAPPINGS["HiCacheAccelerate"] + schema = node.INPUT_TYPES() + required = schema["required"] + assert required["pipeline"][0] == "HY3DMODEL" + assert set(required["method"][0]) == {"hermite", "dmd", "auto"} + assert node.RETURN_TYPES == ("HY3DMODEL",) + PY + - name: Run budget-bridge integration check + run: python -m pytest -q tests/test_budget_bridge.py - name: Build wheel from clean metadata run: python -m pip wheel . --no-deps --wheel-dir dist diff --git a/tests/test_budget_bridge.py b/tests/test_budget_bridge.py new file mode 100644 index 0000000..7ec9c5f --- /dev/null +++ b/tests/test_budget_bridge.py @@ -0,0 +1,45 @@ +"""Integration checks for the shared hicache-pp budget contract.""" + +import importlib.util +import pathlib + +import torch + + +PACK_DIR = pathlib.Path(__file__).resolve().parents[1] + + +def _load_patch(): + spec = importlib.util.spec_from_file_location("hicache_budget_bridge_patch", PACK_DIR / "hicache_patch.py") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +class _DiT(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("base", torch.ones(1, 2)) + + def forward(self, x, timestep, cond, guidance=None): + return self.base * (1.0 + float(timestep.reshape(-1)[0])) + + +class _Pipe: + def __init__(self): + self.model = _DiT() + + +def test_budget_cap_forces_full_fallback_and_manifest_counts_actual_decisions(): + patch = _load_patch() + pipeline = patch.apply_hicache(_Pipe(), method="hermite", interval=4, max_horizon=1) + for timestep in torch.linspace(0.0, 1.0, 10): + pipeline.model(torch.ones(1, 2), timestep, None) + + report = pipeline.model.budget_manifest + assert report["budget"]["max_horizon"] == 1 + assert report["identity"]["stage"] == "shape" + assert report["counts"]["fallback"] > 0 + assert report["counts"]["full"] + report["counts"]["forecast"] + report["counts"]["fallback"] == 10 + assert all("path" not in event for event in report["events"]) + diff --git a/tests/test_nodes.py b/tests/test_nodes.py deleted file mode 100644 index 7e2d6ed..0000000 --- a/tests/test_nodes.py +++ /dev/null @@ -1,106 +0,0 @@ -"""Node-pack registration tests — run standalone, no ComfyUI install. - -ComfyUI loads a custom node pack by importing the package and reading -``NODE_CLASS_MAPPINGS`` / ``NODE_DISPLAY_NAME_MAPPINGS``; these tests verify -that contract plus the INPUT_TYPES schema shape ComfyUI expects. -""" -import importlib.util -import pathlib -import sys - -import pytest - -PACK_DIR = pathlib.Path(__file__).resolve().parents[1] - - -def _load_pack(): - """Import the node pack the way ComfyUI does (as a package by path).""" - name = "comfyui_hicache_under_test" - if name in sys.modules: - return sys.modules[name] - spec = importlib.util.spec_from_file_location( - name, PACK_DIR / "__init__.py", - submodule_search_locations=[str(PACK_DIR)], - ) - mod = importlib.util.module_from_spec(spec) - sys.modules[name] = mod - spec.loader.exec_module(mod) - return mod - - -def test_pack_imports_without_comfy(): - assert "comfy" not in sys.modules, "test must run without ComfyUI" - pack = _load_pack() - assert hasattr(pack, "NODE_CLASS_MAPPINGS") - assert hasattr(pack, "NODE_DISPLAY_NAME_MAPPINGS") - assert "comfy" not in sys.modules, "node pack must not import comfy internals" - - -def test_node_mappings_consistent(): - pack = _load_pack() - assert "HiCacheAccelerate" in pack.NODE_CLASS_MAPPINGS - for key in pack.NODE_CLASS_MAPPINGS: - assert key in pack.NODE_DISPLAY_NAME_MAPPINGS - - -def test_input_types_schema(): - pack = _load_pack() - cls = pack.NODE_CLASS_MAPPINGS["HiCacheAccelerate"] - schema = cls.INPUT_TYPES() - req = schema["required"] - - # model input + the contract from the task spec: method/interval/warmup - assert req["pipeline"][0] == "HY3DMODEL" - assert set(req["method"][0]) == {"hermite", "dmd", "auto"} - assert req["method"][1]["default"] in req["method"][0] - assert req["interval"][0] == "INT" - assert req["interval"][1]["min"] >= 1 - assert req["warmup_steps"][0] == "INT" - assert req["warmup_steps"][1]["min"] == 0 - assert "enable" in schema["optional"] - - # ComfyUI node contract - assert cls.RETURN_TYPES == ("HY3DMODEL",) - assert callable(getattr(cls, cls.FUNCTION)) - assert isinstance(cls.CATEGORY, str) - - -def test_node_function_patches_and_unpatches(): - import torch - - pack = _load_pack() - node = pack.NODE_CLASS_MAPPINGS["HiCacheAccelerate"]() - - class FakePipe: - model = torch.nn.Linear(2, 2) - - pipe = FakePipe() - original = pipe.model - - (out,) = node.patch(pipe, method="hermite", interval=4, warmup_steps=2) - assert out is not pipe, "node must return a copy, not mutate (ComfyUI caching)" - assert pipe.model is original, "input pipeline must stay unpatched" - assert getattr(out.model, "_hicache_is_patch", False) - assert out.model.inner is original - - # re-patch with new params replaces (not nests) the patch - (out2,) = node.patch(out, method="dmd", interval=3, warmup_steps=1) - assert out2.model.inner is original - assert out2.model.method == "dmd" and out2.model.interval == 3 - assert out.model.method == "hermite", "earlier output keeps its config" - - # enable=False restores the original model - (out3,) = node.patch(out2, enable=False) - assert out3.model is original - - # interval=1 means "no caching" -> also unpatched (and an unpatched - # input passes through unchanged) - (out4,) = node.patch(pipe, interval=1) - assert out4 is pipe and out4.model is original - - -def test_node_rejects_bad_pipeline(): - pack = _load_pack() - node = pack.NODE_CLASS_MAPPINGS["HiCacheAccelerate"]() - with pytest.raises(TypeError, match="no `.model` attribute"): - node.patch(object())