diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 8bbdc4c..f215f7c 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -24,7 +24,7 @@ jobs: pip install -e ".[test]" - name: Run CPU unit tests run: python -m pytest -q - - name: Run registry security style checks - run: python -m ruff check --select E702 . + - name: Run static checks + run: python -m ruff check . - name: Build wheel from clean metadata run: python -m pip wheel . --no-deps --wheel-dir dist diff --git a/pyproject.toml b/pyproject.toml index 2d08202..476dccd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,12 +1,17 @@ # Packaging metadata for the registry.comfy.org Comfy Registry. # Field layout follows https://docs.comfy.org/registry/specifications +[build-system] +requires = ["setuptools>=77"] +build-backend = "setuptools.build_meta" + [project] name = "comfyui-trellis2-hicache" description = "Training-free TRELLIS.2 image-to-3D acceleration: forecast the flow-matching velocity on skipped DiT steps (HiCache Hermite / HiCache++ DMD, via hicache-pp) across the sparse-structure and shape-SLaT stages (texture optional). ~2x faster, near-lossless. Pairs with visualbruno/ComfyUI-Trellis2." version = "0.1.3" keywords = ["comfyui", "diffusion-acceleration", "caching", "training-free", "trellis2", "image-to-3d"] -license = { file = "LICENSE" } +license = "MIT" +license-files = ["LICENSE"] requires-python = ">=3.9" dependencies = ["hicache-pp>=1.1.0"] classifiers = ["Development Status :: 4 - Beta"] diff --git a/tests/test_patch.py b/tests/test_patch.py index efab53a..ab66f0a 100644 --- a/tests/test_patch.py +++ b/tests/test_patch.py @@ -4,7 +4,8 @@ No ComfyUI, no GPU, no TRELLIS: a dummy DiT and a fake SparseTensor stand in for the real models, so the scheduling / CFG-routing / run-boundary / copy-on-patch logic is checked deterministically on CPU. """ -import sys, os +import os +import sys sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) import torch @@ -201,6 +202,78 @@ def test_lazy_pipeline_guard_then_load_still_triggers(): assert replacement.inner is patched.loads[1] +def test_lazy_reapply_updates_config_across_unload_reload(): + """Re-patching a loaded lazy slot must update its future reload config.""" + class LazyPipe: + def __init__(self): + self.models = { + "sparse_structure_flow_model": None, + "shape_slat_flow_model_512": None, + "shape_slat_flow_model_1024": None, + } + + first = apply_hicache( + LazyPipe(), method="hermite", interval=3, stages="sparse_structure", + ) + first.models["sparse_structure_flow_model"] = DummyDiT() + assert first.models["sparse_structure_flow_model"].interval == 3 + + second = apply_hicache( + first, method="dmd", interval=5, stages="sparse_structure", + ) + loaded = second.models["sparse_structure_flow_model"] + assert isinstance(loaded, HiCacheModelPatch) + assert loaded.interval == 5 + assert loaded.method == "dmd" + assert second.models._pending["sparse_structure_flow_model"]["interval"] == 5 + assert second.models._pending["sparse_structure_flow_model"]["method"] == "dmd" + + second.models["sparse_structure_flow_model"] = None + second.models["sparse_structure_flow_model"] = DummyDiT() + reloaded = second.models["sparse_structure_flow_model"] + assert isinstance(reloaded, HiCacheModelPatch) + assert reloaded.interval == 5 + assert reloaded.method == "dmd" + + +def test_lazy_subset_reapply_preserves_unselected_pending_stages(): + """A subset re-patch must not disable deferred wrapping for other stages.""" + class LazyPipe: + def __init__(self): + self.models = { + "sparse_structure_flow_model": None, + "shape_slat_flow_model_512": None, + "shape_slat_flow_model_1024": None, + "tex_slat_flow_model_512": None, + "tex_slat_flow_model_1024": None, + } + + all_stages = apply_hicache( + LazyPipe(), method="hermite", interval=3, stages="all", + ) + all_stages.models["sparse_structure_flow_model"] = DummyDiT() + + sparse_only = apply_hicache( + all_stages, method="dmd", interval=5, stages="sparse_structure", + ) + assert set(sparse_only.models._pending) == set(all_stages.models) + assert sparse_only.models._pending["sparse_structure_flow_model"]["interval"] == 5 + assert sparse_only.models._pending["shape_slat_flow_model_512"]["interval"] == 3 + assert sparse_only.models._pending["tex_slat_flow_model_1024"]["method"] == "hermite" + + sparse_only.models["shape_slat_flow_model_512"] = DummyDiT() + shape = sparse_only.models["shape_slat_flow_model_512"] + assert isinstance(shape, HiCacheModelPatch) + assert shape.interval == 3 + assert shape.method == "hermite" + + sparse_only.models["tex_slat_flow_model_1024"] = DummyDiT() + texture = sparse_only.models["tex_slat_flow_model_1024"] + assert isinstance(texture, HiCacheModelPatch) + assert texture.interval == 3 + assert texture.method == "hermite" + + def test_remove_hicache_clears_pending_before_any_load(): """Disabling a pending patch must not re-enable it on a later assignment.""" class LazyPipe: diff --git a/tests/validate_gpu.py b/tests/validate_gpu.py index 92482e7..7d51c87 100644 --- a/tests/validate_gpu.py +++ b/tests/validate_gpu.py @@ -10,7 +10,10 @@ Run from a hermit-trellis2 checkout (which vendors the trellis2 package): PYTHONPATH=.:../comfyui-trellis2-hicache python ../comfyui-trellis2-hicache/tests/validate_gpu.py \ --ckpt /home/krishi/workspace/data/weights/TRELLIS.2-4B --image """ -import argparse, time, sys, types +import argparse +import sys +import time +import types def _stub_render_deps(): @@ -60,8 +63,10 @@ def main(): v_fast = verts(m_fast) def chamfer(a, b, k=20000): - if a.shape[0] > k: a = a[torch.randperm(a.shape[0], device=a.device)[:k]] - if b.shape[0] > k: b = b[torch.randperm(b.shape[0], device=b.device)[:k]] + if a.shape[0] > k: + a = a[torch.randperm(a.shape[0], device=a.device)[:k]] + if b.shape[0] > k: + b = b[torch.randperm(b.shape[0], device=b.device)[:k]] return float((torch.cdist(a, b).min(1).values.mean() + torch.cdist(b, a).min(1).values.mean()) / 2) cham = chamfer(v_stock, v_fast)