diff --git a/CHANGELOG.md b/CHANGELOG.md index 00bfeac3..c2358b16 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Added (experimental) + +- **SLA rising sparsity** (diffusion model loader, `sla_sparsity_start`): skip fewer attention blocks on early steps and more on late ones, in a straight line up to `sla_sparsity`. Off at 0. Untested on H3. Idea from GSRS (arXiv 2610.21129); only the schedule is taken, not its other parts. + ## [5.0.3] "Classy Chopper" - 2026-10-25 ### Changed diff --git a/README.md b/README.md index e038e2d8..94eedfa8 100644 --- a/README.md +++ b/README.md @@ -160,7 +160,7 @@ I express my deepest gratitude to: - [ComfyUI-GGUF](https://github.com/city96/ComfyUI-GGUF) — reads quantized `.gguf` encoders when it is installed; - ComfyUI-H3-SLA-Attention (MIT) — the SLA block-sparse attention for MiniMax H3 is ported from it, and its kernel and block map come from [LightX2V](https://github.com/ModelTC/LightX2V) (Apache-2.0); - ComfyUI-MiniMaxH3_LatentUpscaler by Tr1dae — the tiled H3 decode and the X2 Detail VAE output handling are ported from its "VAE Decode (fast)"; -- the authors of the papers behind the sampling modules: ALG (arXiv 2506.08456), FreeSliders (2511.00103), Decisiveness (2510.01184), momentum (2602.20360), "Pulling The REINS" (2606.17257) and STAS (2603.17825); +- the authors of the papers behind the sampling modules: ALG (arXiv 2506.08456), FreeSliders (2511.00103), Decisiveness (2510.01184), momentum (2602.20360), "Pulling The REINS" (2606.17257) STAS (2603.17825) and GSRS (2610.21129, Dai, Chen, Wen, Zhuang, Chen, Zeng: the rising sparsity schedule in SLA); - [OpenCut](https://github.com/opencut-app/opencut-classic) — the in-browser non-linear video editor whose UI and interaction patterns inspired FunUI's timeline studio; - [Inter](https://github.com/rsms/inter) and IBM Plex Mono — the fonts of the app (SIL Open Font License; copies in `app/assets/fonts`); - [musubi-tuner](https://github.com/kohya-ss/musubi-tuner) by kohya-ss (Apache-2.0) — the LoRA training add-on's loop (which layers get a LoRA, per-block checkpointing, noise sampling and loss weighting, optimizer and learning rate) and its block swap (a number of transformer blocks kept in system memory) follow its MiniMax H3 training recipe; diff --git a/core/tests/test_shipped_nodes.py b/core/tests/test_shipped_nodes.py index e6fad601..e0211483 100644 --- a/core/tests/test_shipped_nodes.py +++ b/core/tests/test_shipped_nodes.py @@ -81,7 +81,7 @@ FROZEN = { "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", "int8_convrot", "sla_engine", "sla_dense_steps", "sla_references", - "sla_tail", "sla_stabilize_motion", "sla_method", "sla_tau", + "sla_tail", "sla_stabilize_motion", "sla_method", "sla_tau", "sla_sparsity_start", ], "FunUILoraLoader": [ "model", "lora_name", "strength_model", "clip", "strength_clip", diff --git a/docs/testing.md b/docs/testing.md index ad0a87b8..197d67f5 100644 --- a/docs/testing.md +++ b/docs/testing.md @@ -12,6 +12,7 @@ COMFYUI_ROOT=/path/to/ComfyUI python -m pytest - `pyproject.toml` sets `--import-mode=importlib` and `testpaths = ["core/tests", "modules"]`. The root `__init__.py` is ComfyUI's entry point and cannot be imported on its own, which is why the mode matters. - Tests sit beside what they test: `core/tests/` and `modules///tests/`. - `COMFYUI_ROOT` points at a ComfyUI checkout (it must contain `comfy/supported_models.py`). Without it, the default is `~/Documents/ComfyUI`. With none found, tests that need ComfyUI skip, and a module that imports `comfy` at load is simply absent. +- On a machine with no GPU (a cloud session), `sh tools/test_env.sh` sets all of this up (CPU torch, pytest, a ComfyUI checkout) and prints the exact command; `tools/cpuforce.py` is the plugin that forces the CPU. Run it once in the environment's setup script and a session starts ready to test. - ComfyUI needs `torch`. A CPU build is enough for the tests. If torch complains about CUDA, force the CPU device for the test run rather than editing a test. - Tests that need `aiohttp` skip themselves with `importorskip` when it is missing. diff --git a/modules/loaders/diffusion_model/nodes.py b/modules/loaders/diffusion_model/nodes.py index ac9cdef0..8a954356 100644 --- a/modules/loaders/diffusion_model/nodes.py +++ b/modules/loaders/diffusion_model/nodes.py @@ -41,7 +41,7 @@ class FunUIDiffusionModelLoader(io.ComfyNode): # How the form lists them (core/widgets): the file, how it is stored and computed, attention, then SLA in the # order its choices narrow it. Not the schema's order, which saved canvas workflows read by position. WIDGET_ORDER = ("model_name", "weight_dtype", "int8_convrot", "compute_dtype", "fp16_accumulation", "attention", - "sla", "sla_method", "sla_engine", "sla_sparsity", "sla_tau", "sla_block_size", + "sla", "sla_method", "sla_engine", "sla_sparsity", "sla_sparsity_start", "sla_tau", "sla_block_size", "sla_dense_steps", "sla_dense_last_steps", "sla_min_seq_len", "sla_protect_audio", "sla_references", "sla_tail", "sla_stabilize_motion") @@ -170,6 +170,14 @@ class FunUIDiffusionModelLoader(io.ComfyNode): step=0.05, optional=True, tooltip="sol-attn only: higher skips more. About 1.0 keeps ~16% of key " "blocks exact, 1.5 ~7%, 2.0 ~3%."), + # Added after everything: saved workflows store widget values by position. + io.Float.Input("sla_sparsity_start", extra_dict=shows(("sla", [True]), ("sla_method", ["sla"])), default=sla_attention.SLA_DEFAULTS["sparsity_start"], + min=0.0, max=0.95, step=0.05, optional=True, + tooltip="Skip fewer blocks early and more late: the share rises in a straight " + "line from this value on the first step to sla_sparsity on the last " + "(attention concentrates as denoising goes, GSRS, arXiv 2610.21129). " + "0 is off: one fixed share all run. Untested on H3; a try is 0.70 " + "with sla_sparsity 0.90. Early steps cost more time than late ones."), ], outputs=[ io.Model.Output(display_name="model"), @@ -185,7 +193,8 @@ class FunUIDiffusionModelLoader(io.ComfyNode): sla_dense_last_steps: int = None, int8_convrot: bool = False, sla_engine: str = None, sla_dense_steps: str = None, sla_references: str = None, sla_tail: bool = None, sla_stabilize_motion: bool = None, - sla_method: str = None, sla_tau: float = None) -> io.NodeOutput: + sla_method: str = None, sla_tau: float = None, + sla_sparsity_start: float = None) -> io.NodeOutput: notes = [f"FunUI Diffusion Model Loader | {model_name}"] accumulation = set_fp16_accumulation(fp16_accumulation) @@ -282,7 +291,7 @@ class FunUIDiffusionModelLoader(io.ComfyNode): if sla: model, sla_note, installed = sla_attention.install_sla( model, - sparsity_ratio=sla_sparsity, block_size=sla_block_size, + sparsity_ratio=sla_sparsity, sparsity_start=sla_sparsity_start, block_size=sla_block_size, min_seq_len=sla_min_seq_len, dense_last_steps=sla_dense_last_steps, protect_audio=sla_protect_audio, dense_fn=override, dense_label=attention, diff --git a/modules/loaders/sla_attention.py b/modules/loaders/sla_attention.py index 33ebe76d..3535c70e 100644 --- a/modules/loaders/sla_attention.py +++ b/modules/loaders/sla_attention.py @@ -50,6 +50,7 @@ _H3_MODEL_CLASS = "comfy.ldm.minimax.model.MiniMaxH3Model" # 3.7x the attention throughput of stock ComfyUI at sparsity 0.90 / block 64. SLA_DEFAULTS = { "sparsity_ratio": 0.90, + "sparsity_start": 0.0, "block_size": 64, "min_seq_len": 8192, "dense_last_steps": 0, @@ -229,6 +230,7 @@ def new_state(): "layer": 0, # this step's call count: the Nth call is the same layer every step "history": {}, # layer -> last step's choices, for stabilize_motion "order": None, # comfy_kitchen: (S, protected spans, key order, sink blocks), same every layer + "sparsity": None, # this step's share of blocks skipped (step_sparsity), set by the wrapper "last": None, # last step index read from the schedule "dense_at": set(), # steps (counted from 1) that ran at full attention this run "n": 0, # this run's step count @@ -236,6 +238,16 @@ def new_state(): } +def step_sparsity(cfg, step, n_steps): + """The share of key blocks skipped on `step` (counted from 1). Falls back to the one fixed value; with a + start below it, rises in a straight line from the start (first step) to `sparsity_ratio` (last), the way + attention concentrates as denoising goes (GSRS, arXiv 2610.21129: dense early, sparser late).""" + end, start = cfg["sparsity_ratio"], cfg["sparsity_start"] + if not 0.0 < start < end or n_steps < 2: + return end + return start + (end - start) * min(1.0, max(0.0, (step - 1) / (n_steps - 1))) + + def _summarise(state, cfg): """One line per sampling run. Never one per block — there are 50 of those.""" if state["calls"] == 0: @@ -245,11 +257,13 @@ def _summarise(state, cfg): "one this loader produced.", state["dense"]) return real = 1.0 - (state["kept"] / state["blocks"]) if state["blocks"] else 0.0 + ramp = step_sparsity(cfg, 1, 2) != cfg["sparsity_ratio"] logging.info( - "[FunUI SLA] %s | %d calls | S=%d | blocks %d/%d kept (%.1f%% sparse, asked %.0f%%) " + "[FunUI SLA] %s | %d calls | S=%d | blocks %d/%d kept in the last step (%.1f%% sparse, asked %s) " "| %d pinned | %d dense fall-throughs | displaced %s", cfg["engine"], state["calls"], state["seq"], state["kept"], state["blocks"], real * 100.0, - cfg["sparsity_ratio"] * 100.0, state["pinned"], state["dense"], state["backend"] or "?") + f"{cfg['sparsity_start'] * 100.0:.0f}% rising to {cfg['sparsity_ratio'] * 100.0:.0f}%" if ramp + else f"{cfg['sparsity_ratio'] * 100.0:.0f}%", state["pinned"], state["dense"], state["backend"] or "?") logging.info("[FunUI SLA] full attention on %s of %d", f"steps {_runs(state['dense_at'])}" if state["dense_at"] else "no step", state["n"]) if state["failed"] is not None: @@ -268,7 +282,7 @@ def _mean_pool_torch(x, blk): return x.reshape(B, H, -1, blk, D).sum(dim=3) / counts # the last, short block averages its real rows only -def _ck_attention(state, cfg, qb, kb, vb, keep, refs, scale): +def _ck_attention(state, cfg, qb, kb, vb, keep, refs, scale, topk_ratio): """comfy_kitchen's sol_attn: one 64-token key-block range (`sink_blocks`) is always exact. Attention does not care about key order, so when the protected blocks are not one run, K and V are put in an order where they are -- a copy per call, needed only then. -> (output, blocks, pinned).""" @@ -311,7 +325,7 @@ def _ck_attention(state, cfg, qb, kb, vb, keep, refs, scale): out = ck.sol_attn(q, k, v, tau=cfg["tau"], scale=scale, sink_blocks=sink, sink_q=[0, 0]) else: out = ck.sol_attn(q, k, v, scale=scale, sink_blocks=sink, sink_q=[0, 0], - topk_ratio=1.0 - cfg["sparsity_ratio"], tail=cfg["tail"]) + topk_ratio=topk_ratio, tail=cfg["tail"]) return (out.to(qb.dtype) if cast else out), nk, sink[1] - sink[0] @@ -326,7 +340,6 @@ def make_override(state, cfg, dense_fn=None, dense_label=None): """ torch = _torch() ok_dtypes = (torch.bfloat16, torch.float16) - topk_ratio = 1.0 - cfg["sparsity_ratio"] blkq = cfg["block_size"] # BLKK=64 is not a typo. On sm_120 the 128x128 tile needs 160 KB of shared memory # against a ~99 KB limit and cannot launch at all; 128x64 both fits and measured @@ -365,6 +378,7 @@ def make_override(state, cfg, dense_fn=None, dense_label=None): try: B, H, S, D = q.shape + topk_ratio = 1.0 - (cfg["sparsity_ratio"] if state["sparsity"] is None else state["sparsity"]) layer = state["layer"] state["layer"] = layer + 1 @@ -379,7 +393,7 @@ def make_override(state, cfg, dense_fn=None, dense_label=None): keep = tuple(keep if cfg["protect_audio"] else ()) + tuple(refs if heavy else ()) if cfg["engine"] == "comfy_kitchen": - out, blocks, pinned = _ck_attention(state, cfg, qb, kb, vb, keep, refs, D ** -0.5) + out, blocks, pinned = _ck_attention(state, cfg, qb, kb, vb, keep, refs, D ** -0.5, topk_ratio) kept = max(1, round(topk_ratio * blocks)) + pinned else: try: @@ -470,6 +484,7 @@ def make_wrapper(state, cfg): to["_funui_sla_dense"] = bool( (last > 0 and state["step"] > n_steps - last) or (state["step"] - 1) in dense_steps) state["n"] = n_steps + state["sparsity"] = step_sparsity(cfg, state["step"], n_steps) if to["_funui_sla_dense"] and counted: state["dense_at"].add(state["step"]) @@ -494,7 +509,7 @@ def make_wrapper(state, cfg): return wrapper -def install_sla(model, sparsity_ratio=None, block_size=None, min_seq_len=None, +def install_sla(model, sparsity_ratio=None, sparsity_start=None, block_size=None, min_seq_len=None, dense_last_steps=None, protect_audio=None, enabled=None, dense_fn=None, dense_label=None, engine=None, dense_steps=None, references=None, tail=None, stabilize_motion=None, method=None, tau=None): @@ -507,7 +522,7 @@ def install_sla(model, sparsity_ratio=None, block_size=None, min_seq_len=None, `installed` is False whenever SLA did not take, so the caller can fall back to installing the chosen backend on its own rather than leaving the model with nothing. """ - given = dict(sparsity_ratio=sparsity_ratio, block_size=block_size, min_seq_len=min_seq_len, + given = dict(sparsity_ratio=sparsity_ratio, sparsity_start=sparsity_start, block_size=block_size, min_seq_len=min_seq_len, dense_last_steps=dense_last_steps, protect_audio=protect_audio, enabled=enabled, engine=engine, dense_steps=dense_steps, references=references, tail=tail, stabilize_motion=stabilize_motion, method=method, tau=tau) @@ -528,11 +543,13 @@ def install_sla(model, sparsity_ratio=None, block_size=None, min_seq_len=None, cfg["method"] = SLA_DEFAULTS["method"] if cfg["method"] == "vsa": return _install_vsa(model, cfg, dense_fn, dense_label) - cfg.update(sparsity_ratio=float(cfg["sparsity_ratio"]), block_size=int(cfg["block_size"]), + cfg.update(sparsity_ratio=float(cfg["sparsity_ratio"]), sparsity_start=float(cfg["sparsity_start"]), block_size=int(cfg["block_size"]), min_seq_len=int(cfg["min_seq_len"]), dense_last_steps=int(cfg["dense_last_steps"]), protect_audio=bool(cfg["protect_audio"]), tail=bool(cfg["tail"]), stabilize_motion=bool(cfg["stabilize_motion"]), references=cfg["references"] if cfg["references"] in REFERENCES else "off") + if cfg["method"] != "sla": + cfg["sparsity_start"] = 0.0 # a rising share is the fixed-share method's; sol-attn routes by threshold if cfg["engine"] not in ENGINES: cfg["engine"] = SLA_DEFAULTS["engine"] if cfg["engine"] == "comfy_kitchen" and not ck_ok: @@ -568,6 +585,8 @@ def install_sla(model, sparsity_ratio=None, block_size=None, min_seq_len=None, blk = cfg["block_size"] if cfg["engine"] == "triton" else 64 how = f"tau={cfg['tau']:.2f}" if cfg["method"] == "sol-attn" else f"sparsity={cfg['sparsity_ratio']:.2f}" + if cfg["method"] == "sla" and step_sparsity(cfg, 1, 2) != cfg["sparsity_ratio"]: + how = f"sparsity={cfg['sparsity_start']:.2f} rising to {cfg['sparsity_ratio']:.2f}" return patched, (f"SLA on | {cfg['method']} on {cfg['engine']} | {how} BLK={blk} " f"min_seq_len={cfg['min_seq_len']} dense_last_steps={cfg['dense_last_steps']} " f"full attention on steps {_runs(i + 1 for i in cfg['dense_steps']) or 'none'} and the last {cfg['dense_last_steps']} protect_audio={cfg['protect_audio']} " diff --git a/modules/loaders/tests/test_sla_attention.py b/modules/loaders/tests/test_sla_attention.py index f8d4cc57..86df2ffb 100644 --- a/modules/loaders/tests/test_sla_attention.py +++ b/modules/loaders/tests/test_sla_attention.py @@ -262,7 +262,7 @@ def test_sla_reports_when_the_machine_cannot_run_it(): def test_the_defaults_are_the_validated_ones(): """0.90 / 64 / protect on: the settings the port was measured at.""" - assert sla.SLA_DEFAULTS == {"sparsity_ratio": 0.90, "block_size": 64, + assert sla.SLA_DEFAULTS == {"sparsity_ratio": 0.90, "sparsity_start": 0.0, "block_size": 64, "min_seq_len": 8192, "dense_last_steps": 0, "protect_audio": True, "enabled": True, "engine": "comfy_kitchen", "dense_steps": "1", "method": "sla", "tau": 1.3, @@ -546,3 +546,22 @@ def test_the_torch_pool_averages_real_rows_only_so_light_runs_without_triton(): assert torch.allclose(pooled[0, h, 0], (x[0, 0, h] + x[0, 1, h]) / 2) assert torch.allclose(pooled[0, h, 1], (x[0, 2, h] + x[0, 3, h]) / 2) assert torch.allclose(pooled[0, h, 2], x[0, 4, h]) # the short block: its one real row + + +def test_sparsity_rises_from_the_start_value_to_the_set_one_and_zero_means_fixed(): + """GSRS: skip fewer blocks early, more late. 0 (the default) must leave the old fixed share untouched.""" + cfg = _cfg(sparsity_ratio=0.9, sparsity_start=0.7) + assert [round(sla.step_sparsity(cfg, i, 4), 3) for i in (1, 2, 3, 4)] == [0.7, 0.767, 0.833, 0.9] + assert [sla.step_sparsity(_cfg(sparsity_ratio=0.9), i, 4) for i in (1, 2, 3, 4)] == [0.9] * 4 + assert sla.step_sparsity(_cfg(sparsity_ratio=0.9, sparsity_start=0.95), 1, 4) == 0.9 # a start above the end is no ramp + assert sla.step_sparsity(cfg, 1, 1) == 0.9 # one step: nothing to ramp over + + +def test_install_drops_the_ramp_for_methods_that_do_not_use_a_fixed_share(monkeypatch): + monkeypatch.setattr(sla, "sla_available", lambda: True) + monkeypatch.setattr(sla, "ck_available", lambda: True) + model = _Patcher("MiniMaxH3Model") + _, note, ok = sla.install_sla(model, sparsity_start=0.7, method="sla", engine="triton") + assert ok and "sparsity=0.70 rising to 0.90" in note + _, note, ok = sla.install_sla(model, sparsity_start=0.7, method="sol-attn") + assert ok and "rising" not in note diff --git a/tools/cpuforce.py b/tools/cpuforce.py new file mode 100644 index 00000000..4b9fb36a --- /dev/null +++ b/tools/cpuforce.py @@ -0,0 +1,17 @@ +"""pytest plugin: run ComfyUI on the CPU, so the suite works on a machine without a GPU. + + COMFYUI_ROOT=/path/to/ComfyUI PYTHONPATH=tools python -m pytest -p cpuforce +""" +import os +import sys + + +def pytest_configure(config): + root = os.environ.get("COMFYUI_ROOT") + if root and root not in sys.path: + sys.path.insert(0, root) + try: + from comfy.cli_args import args + except ImportError: + return + args.cpu = True diff --git a/tools/test_env.sh b/tools/test_env.sh new file mode 100755 index 00000000..d4ce4ec7 --- /dev/null +++ b/tools/test_env.sh @@ -0,0 +1,12 @@ +#!/bin/sh +# Test environment for a machine with no GPU (a cloud session): CPU torch, pytest, a ComfyUI checkout. +# Safe to run again; each step is skipped when already there. Prints the command that runs the suite. +set -e +COMFY=${COMFYUI_ROOT:-$HOME/ComfyUI} +# The three together, from the CPU index: ComfyUI's requirements would otherwise pull a CUDA torchvision that no longer matches. +python3 -c "import torch, torchvision" 2>/dev/null || pip install -q --index-url https://download.pytorch.org/whl/cpu torch torchvision torchaudio +python3 -c "import pytest" 2>/dev/null || pip install -q pytest +[ -f "$COMFY/comfy/supported_models.py" ] || git clone -q --depth 1 https://github.com/comfyanonymous/ComfyUI.git "$COMFY" +pip install -q -r "$COMFY/requirements.txt" +HERE=$(cd "$(dirname "$0")" && pwd) +echo "COMFYUI_ROOT=$COMFY PYTHONPATH=$HERE python3 -m pytest -p cpuforce -q"