From 1e2ac5ca67f1c452db9178af28cce939aa689737 Mon Sep 17 00:00:00 2001 From: digital-garbage <133010426+digital-garbage@users.noreply.github.com> Date: Thu, 5 Nov 2026 21:32:34 +0200 Subject: [PATCH] Sparse attention method veda for MiniMax H3 (experimental), with an optional predictor download Learned tile-sparse attention (Veda, arXiv 2605.30325): tiles, tile plans, predictor scoring and selection in plain torch (tested on CPU, including a real tiny H3 with a stand-in kernel), an INT8 Triton kernel derived from SageAttention (untested on a GPU), installed through the same override and per-step wrapper as SLA. Runs at the predictor's trained budget; LoRA, MXFP8, Two stages, dense steps and the attention backend compose. Settings > Models offers the official Hugging Face download of the predictor on a click, with its licence beside it; taggers and this share one downloader (core/fetch.py). --- CHANGELOG.md | 2 + README.md | 1 + app/ui/features/settings/models.js | 29 +- app/ui/features/settings/models.test.js | 40 ++ app/ui/shell/api.js | 2 + core/fetch.py | 37 ++ core/tests/test_shipped_nodes.py | 2 +- modules/_core.py | 6 +- modules/loaders/VEDA-NOTICE.txt | 74 +++ modules/loaders/diffusion_model/__init__.py | 17 + modules/loaders/diffusion_model/nodes.py | 28 +- modules/loaders/sla_attention.py | 13 +- modules/loaders/tests/test_sla_attention.py | 4 +- modules/loaders/tests/test_veda.py | 488 ++++++++++++++++++++ modules/loaders/tests/test_veda_download.py | 154 ++++++ modules/loaders/veda_attention.py | 194 ++++++++ modules/loaders/veda_core.py | 425 +++++++++++++++++ modules/loaders/veda_download.py | 74 +++ modules/loaders/veda_kernel.py | 126 +++++ modules/training/lora/taggers.py | 22 +- 20 files changed, 1700 insertions(+), 38 deletions(-) create mode 100644 core/fetch.py create mode 100644 modules/loaders/VEDA-NOTICE.txt create mode 100644 modules/loaders/tests/test_veda.py create mode 100644 modules/loaders/tests/test_veda_download.py create mode 100644 modules/loaders/veda_attention.py create mode 100644 modules/loaders/veda_core.py create mode 100644 modules/loaders/veda_download.py create mode 100644 modules/loaders/veda_kernel.py diff --git a/CHANGELOG.md b/CHANGELOG.md index f9a10a6c..219e4930 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -19,6 +19,8 @@ ### Added +- **Sparse attention method "veda"** (H3 loader, SLA on, method veda; Experimental, UNTESTED on a GPU). A small learned predictor (Veda, arXiv 2605.30325) picks which 128-token tiles of the attention map matter and an INT8 Triton kernel computes only those, at the predictor's own trained budget (90% for text-to-video). It needs CUDA and Triton and a predictor file in ComfyUI/models/veda, chosen in "veda_predictor" (auto takes the text-to-video file, else the first). Works with LoRA (it sees q/k/v after the LoRA), MXFP8 (weights only: q/k/v stay bf16), Two stages (a size the predictor was not trained for uses the nearest tile plan, and the log says so), dense steps and the attention backend, which keeps the text refiner, masked calls and short sequences. sla_sparsity, the rising share and the references choice are not read in this first version. When Veda cannot run, the status says why and the chosen backend is installed on its own. One log line per run in the SLA line's wording (calls, tile pairs kept, asked vs real sparsity, share of full attention computed, steps at full attention). The Triton kernel is checked against an exact reference on the GPU before first use; the torch half (tiles, predictor, selection, the override on a real tiny H3) is tested on CPU. +- **Veda predictor download** (Settings > Models, shown only for method veda and only while models/veda holds no predictor): a button per official file (text-to-video, reference-to-video) fetches it straight from Hugging Face on your click, with the MiniMax H3 Community License named and linked beside it. FunUI does not host or bundle the file. The download streams to a `.part` file and is kept only when its size and checksum match; taggers and this share one downloader (`core/fetch.py`). - **Run meter** (Settings > Engine > Sampling, off by default). Logs how grainy the picture is, across it and over time, and how much it still changes at each step, then compares with the previous run of the same seed. It adds no model call, so flipping one setting and rendering again shows what that setting did. - **Post settings for the finished video** (Settings > Engine > Post, all off or neutral by default): Smooth jagged edges (Light or Strong; FXAA-style, only strongly contrasting edge pixels are softened), Auto contrast, Auto brightness and Colour strength. Contrast and brightness are measured once on the whole clip and applied as one fixed curve, so nothing flickers from frame to frame. They run on the frames before encoding, so NVENC, its CPU fallback and every codec get the same picture. Save Video reads them from the pipeline's settings object (older saved pipelines are linked to it automatically). Tried on the CPU only. - **LoRA Loader: "Apply to stage"** (First, Second or Both, per slot). Shown only while the Two stages module is present and not switched off for the project; the feature itself may stay off in the sampler settings. A LoRA set to First works only in the smaller first stage, one set to Second only in the full-size steps; Two stages hands each stage the model with its own LoRAs, and ComfyUI re-patches the weights between the stages. Only the model changes for First or Second (the text encoder is encoded before sampling, so its strength is not used). When Two stages is off, or a run cannot be cut in two, those LoRAs are not applied and the Log says so. Both is the old behaviour. Tried on a small model on the CPU with every feature switched on; not on a GPU yet. diff --git a/README.md b/README.md index cfa9331f..8c783ecc 100644 --- a/README.md +++ b/README.md @@ -159,6 +159,7 @@ I express my deepest gratitude to: - [ComfyUI-KJNodes](https://github.com/kijai/ComfyUI-KJNodes) — utility nodes used by the H3 pipeline presets; - [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); +- [Veda](https://github.com/veda-sparse/Veda-on-ComfyUI) by veda-sparse (MIT; paper "Veda: Scalable Video Diffusion via Distilled Sparse Attention", Han, Yang, Hu, Mei, Jiang and Qi, ICML 2026, arXiv 2605.30325) and its training code [Miowtion](https://github.com/veda-sparse/Miowtion) (MIT) — the "veda" method of the loader's sparse attention (tiling, tile plans, predictor scoring and tile selection) is rewritten from their code; its INT8 kernel derives from [SageAttention](https://github.com/thu-ml/SageAttention) 1.0.6 by Jintao Zhang and Haofeng Huang (BSD-3-Clause). Licence notices in `modules/loaders/VEDA-NOTICE.txt`. The Veda predictor files are not part of FunUI: they stay under the MiniMax H3 Community License and are downloaded by you, on your click, from their official Hugging Face repositories; - ComfyUI-MiniMaxH3_LatentUpscaler by Tr1dae — the tiled H3 decode and the X2 Detail VAE output handling are ported from its "VAE Decode (fast)"; - [JoyAI-Echo](https://github.com/jd-opensource/JoyAI-Echo) by jd-opensource — the cross-shot memory idea behind the Guide frames "Memory" mode; - the authors of the papers behind the sampling modules: ALG (arXiv 2506.08456), FreeSliders (2511.00103), Decisiveness (2510.01184), BachVid (2510.21696) and KV-Lock "When to Lock Attention" (2603.09657) for the Identity bank, 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); diff --git a/app/ui/features/settings/models.js b/app/ui/features/settings/models.js index 78241485..732a37f8 100644 --- a/app/ui/features/settings/models.js +++ b/app/ui/features/settings/models.js @@ -123,6 +123,31 @@ export const models = (app) => function mount() { return row; }; + // The Veda predictor (method veda): shown only while models/veda holds none. The file is not ours to host (MiniMax H3 Community + // License), so the one way offered is a fetch from the official Hugging Face address on this click, with the licence beside it. + const vedaRow = () => { + const row = c.region.stack({ gap: "sm" }); + let timer = null, fetching = false; + const show = async () => { + let r; + try { r = await api.vedaState(); } catch (err) { return row.set([c.hint.default({ text: `Could not ask about the Veda predictor: ${err.message}` })]); } + if (gone) return; + if ((r.files || []).length) { if (fetching) { fetching = false; specs = {}; describe().then(draw); } return row.set([]); } // it arrived: the file list is read again + const busy = r.models.some((m) => m.state === "downloading"); + const mb = (n) => `${Math.round(n / 1048576)} MB`; + row.set([c.banner.info({ text: "Veda needs its predictor file in ComfyUI/models/veda, and none is there. You can fetch it straight from Hugging Face; nothing is downloaded until you click." }), + ...r.models.map((m) => c.settingsRow.default({ label: `${m.title} · ${mb(m.size)}`, hint: m.state === "error" ? `Failed: ${m.error}` : m.state === "downloading" ? `${Math.floor(100 * m.done / m.size)}%` : m.hint, + control: c.button.sm({ label: m.state === "downloading" ? "Downloading…" : m.state === "error" ? "Try again" : "Download", tone: "primary", disabled: busy, + onClick: async () => { try { await api.vedaDownload(m.id); fetching = true; } catch (err) { c.toast.warn({ text: err.message }); } show(); } }) })), + c.settingsRow.default({ label: `By downloading you accept the ${r.license.name}.`, hint: "FunUI does not host or bundle the file. It has its own terms of use, including who may use it and where.", + control: c.button.sm({ label: "Read the licence", tone: "ghost", onClick: () => app.dom.openTab(r.license.url) }) })]); + clearTimeout(timer); + timer = busy ? setTimeout(() => { if (!gone && row.node && row.node.isConnected) show(); }, 1500) : null; // a redrawn page leaves its old row behind: that one stops asking + }; + show(); + return row; + }; + const slotBlock = (slot) => { const spec = specs[slot.node]; if (spec === undefined) return [c.hint.default({ text: "Loading…" })]; @@ -143,7 +168,9 @@ export const models = (app) => function mount() { const canTake = !appOwned && (w.type !== "COMBO" || (w.choices || []).length) && sourcesFor(slot, w.type, slotsNow(), specs).length; return canTake ? [row, c.field.default({ label: "↳ or take it from", hint: "Another node's output can drive this value; one Primitive can drive several inputs.", control: wiring(slot, w.name, w.type, null, "A value (typed here)") })] : [row]; }); - return [head, structure(slot), ...sockets, ...rows, ...(!rows.length && !sockets.length ? [c.hint.default({ text: "Nothing to configure on this node." })] : [])]; + const veda = (spec.widgets || []).find((w) => w.name === "veda_predictor"); + const vedaOn = veda && !gone.has(veda.name) && showing(veda, slot, spec); + return [head, structure(slot), ...sockets, ...rows, ...(vedaOn ? [vedaRow()] : []), ...(!rows.length && !sockets.length ? [c.hint.default({ text: "Nothing to configure on this node." })] : [])]; }; // One node as a card: what it is, the one or two values that matter, and whether it needs attention. diff --git a/app/ui/features/settings/models.test.js b/app/ui/features/settings/models.test.js index f4c9019a..a53f3882 100644 --- a/app/ui/features/settings/models.test.js +++ b/app/ui/features/settings/models.test.js @@ -142,3 +142,43 @@ test("a field that belongs to a module is drawn only while that module is here a assert.ok((await draw((id) => id === "two_stage")).includes("apply_to_stage")); assert.ok(!(await draw(() => false)).includes("apply_to_stage")); }); + +test("the Veda predictor row appears only for method veda with no file, names the licence beside the buttons, and downloads only on a click", async () => { + const widgets = [{ name: "sla", type: "BOOLEAN" }, { name: "sla_method", type: "COMBO", choices: ["sla", "veda"], shows: [{ input: "sla", when: [true] }] }, + { name: "veda_predictor", type: "COMBO", choices: ["auto"], shows: [{ input: "sla", when: [true] }, { input: "sla_method", when: ["veda"] }] }]; + const state = { files: [], license: { name: "MiniMax H3 Community License", url: "https://huggingface.co/x/blob/main/LICENSE" }, + models: [{ id: "t2va", title: "T2VA / FL2VA", hint: "text", file: "a.safetensors", size: 275415648, state: "absent", done: 0, error: "" }, + { id: "r2va", title: "R2VA", hint: "refs", file: "b.safetensors", size: 275415752, state: "absent", done: 0, error: "" }] }; + const asked = [], fetched = []; + const open = async (inputs) => { + const slots = [{ id: "model", node: "FunUIDiffusionModelLoader", inputs }]; + const ps = { slots: () => slots, snapshot: () => ({ slots }), ensureLoaded: async () => {}, subscribe: () => () => {}, restore: async () => ({ refused: [] }), + loading: () => false, loadError: () => null, saveNotes: () => [], incomplete: () => [], refused: () => [], queueable: () => true, removedIds: () => [], offered: () => [] }; + const api = { describeNodes: async () => ({ nodes: { FunUIDiffusionModelLoader: { widgets } } }), packProviders: async () => ({}), + vedaState: async () => { asked.push(1); return state; }, vedaDownload: async (id) => { fetched.push(id); state.models[0].state = "downloading"; return { started: true }; } }; + const page = models({ dom: createDom(), pipeline: ps, api })(); + document.body.append(page.node); + await new Promise((r) => setTimeout(r, 0)); + page.node.querySelector(".cx-tile-node").click(); + await new Promise((r) => setTimeout(r, 5)); + return page; + }; + let page = await open({ sla: true, sla_method: "sla" }); + assert.equal(asked.length, 0, "no other method asks, or shows anything"); + assert.doesNotMatch(page.node.textContent, /Veda needs its predictor/); + page.node.remove(); + page = await open({ sla: true, sla_method: "veda" }); + assert.match(page.node.textContent, /none is there/); + assert.match(page.node.textContent, /By downloading you accept the MiniMax H3 Community License/); + assert.match(page.node.textContent, /275 MB|263 MB/); + assert.deepEqual(fetched, [], "looking never downloads"); + [...page.node.querySelectorAll("button")].find((b) => b.textContent === "Download").click(); + await new Promise((r) => setTimeout(r, 5)); + assert.deepEqual(fetched, ["t2va"]); + assert.ok([...page.node.querySelectorAll("button")].every((b) => b.textContent !== "Download" || b.disabled), "one download at a time"); + state.files = ["a.safetensors"]; state.models[0].state = "ready"; + page.node.remove(); + page = await open({ sla: true, sla_method: "veda" }); + assert.doesNotMatch(page.node.textContent, /Veda needs its predictor/, "a predictor in the folder: no row"); + page.node.remove(); +}); diff --git a/app/ui/shell/api.js b/app/ui/shell/api.js index 6dc28267..029e0af1 100644 --- a/app/ui/shell/api.js +++ b/app/ui/shell/api.js @@ -77,6 +77,8 @@ export const api = { forgetShotMemory: (kind, name) => call("POST", "/api/m/conditioning_shot_camera/forget", { kind, name }), tasteKeys: () => call("GET", "/api/m/taste/keys"), torchBuild: () => call("GET", "/api/m/loader_diffusion_model/torch"), + vedaState: () => call("GET", "/api/m/loader_diffusion_model/veda"), + vedaDownload: (id) => call("POST", "/api/m/loader_diffusion_model/veda/download", { id }), blockInfluence: (key) => call("GET", `/api/m/block_influence/status?key=${encodeURIComponent(key || "default")}`), blockInfluenceGroups: (key) => call("GET", `/api/m/block_influence/groups?key=${encodeURIComponent(key || "default")}`), clearBlockInfluence: (key) => call("POST", "/api/m/block_influence/clear", { key: key || "default" }), diff --git a/core/fetch.py b/core/fetch.py new file mode 100644 index 00000000..aced9cbb --- /dev/null +++ b/core/fetch.py @@ -0,0 +1,37 @@ +"""Download one URL to one file, safely. The only downloader FunUI has: taggers and the Veda predictor use it. + +Streams to a `.part` file and renames only when the file is complete, so a cut connection or a wrong file +never leaves anything that looks finished. A size or a SHA-256 that is known in advance is checked before the +rename; any failure removes the `.part` and raises OSError with a sentence for the person. +""" +import hashlib +import urllib.request + + +def to_file(url, dest, progress=lambda done, total: None, size=None, sha256=None, timeout=60): + """Fetch `url` into `dest` (a pathlib.Path). `progress(done, total)` is called per megabyte; total is 0 when + the server does not say. `size` / `sha256`: what the finished file must be.""" + part = dest.with_name(dest.name + ".part") + digest = hashlib.sha256() + try: + with urllib.request.urlopen(urllib.request.Request(url, headers={"User-Agent": "FunUI"}), timeout=timeout) as res, open(part, "wb") as out: + total = int(res.headers.get("Content-Length") or 0) + done = 0 + while True: + chunk = res.read(1 << 20) + if not chunk: + break + out.write(chunk) + digest.update(chunk) + done += len(chunk) + progress(done, total) + if total and done != total: + raise OSError(f"the download stopped at {done} of {total} bytes") + if size is not None and done != size: + raise OSError(f"the file is {done} bytes, expected {size}: it was not kept") + if sha256 is not None and digest.hexdigest() != sha256: + raise OSError("the file does not match its published checksum: it was not kept") + except BaseException: + part.unlink(missing_ok=True) + raise + part.replace(dest) diff --git a/core/tests/test_shipped_nodes.py b/core/tests/test_shipped_nodes.py index 1c961512..ecae3a26 100644 --- a/core/tests/test_shipped_nodes.py +++ b/core/tests/test_shipped_nodes.py @@ -83,7 +83,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_sparsity_start", "mxfp8", + "sla_tail", "sla_stabilize_motion", "sla_method", "sla_tau", "sla_sparsity_start", "mxfp8", "veda_predictor", ], "FunUILoraLoader": [ "model", "lora_name", "strength_model", "clip", "strength_clip", "apply_to_stage", diff --git a/modules/_core.py b/modules/_core.py index f89b22e3..5cf8c485 100644 --- a/modules/_core.py +++ b/modules/_core.py @@ -11,8 +11,8 @@ leading underscore keeps it out of the module scan. """ try: # root on sys.path: tests, dev server - from core import addons, chain, config, contract, control, dit_hooks, graph, input_steer, jobs, log, media, patching, projects, rated_dial, registry, relations, run, schema, shortcuts, streams, sysinfo, temp_files, traits, vae_tiles, video_pipe, nvidia + from core import addons, chain, config, contract, control, dit_hooks, fetch, graph, input_steer, jobs, log, media, patching, projects, rated_dial, registry, relations, run, schema, shortcuts, streams, sysinfo, temp_files, traits, vae_tiles, video_pipe, nvidia except ImportError: # inside ComfyUI: a subpackage of the pack - from ..core import addons, chain, config, contract, control, dit_hooks, graph, input_steer, jobs, log, media, patching, projects, rated_dial, registry, relations, run, schema, shortcuts, streams, sysinfo, temp_files, traits, vae_tiles, video_pipe, nvidia + from ..core import addons, chain, config, contract, control, dit_hooks, fetch, graph, input_steer, jobs, log, media, patching, projects, rated_dial, registry, relations, run, schema, shortcuts, streams, sysinfo, temp_files, traits, vae_tiles, video_pipe, nvidia -__all__ = ["addons", "chain", "config", "contract", "control", "dit_hooks", "graph", "input_steer", "jobs", "log", "media", "patching", "projects", "rated_dial", "registry", "relations", "run", "schema", "shortcuts", "streams", "sysinfo", "temp_files", "traits", "vae_tiles", "video_pipe", "nvidia"] +__all__ = ["addons", "chain", "config", "contract", "control", "dit_hooks", "fetch", "graph", "input_steer", "jobs", "log", "media", "patching", "projects", "rated_dial", "registry", "relations", "run", "schema", "shortcuts", "streams", "sysinfo", "temp_files", "traits", "vae_tiles", "video_pipe", "nvidia"] diff --git a/modules/loaders/VEDA-NOTICE.txt b/modules/loaders/VEDA-NOTICE.txt new file mode 100644 index 00000000..4c0e4211 --- /dev/null +++ b/modules/loaders/VEDA-NOTICE.txt @@ -0,0 +1,74 @@ +Veda sparse attention (veda_core.py, veda_kernel.py, veda_attention.py) +====================================================================== + +These files are a rewrite of parts of Veda-on-ComfyUI +(https://github.com/veda-sparse/Veda-on-ComfyUI, MIT), which in turn rewrites +Miowtion's miowtion/veda (https://github.com/veda-sparse/Miowtion, MIT). +Paper: "Veda: Scalable Video Diffusion via Distilled Sparse Attention", +Han, Yang, Hu, Mei, Jiang, Qi, ICML 2026 (arXiv 2605.30325). + +The INT8 arithmetic in veda_kernel.py is derived from SageAttention 1.0.6 +(quant_per_block.py and attn_qk_int8_per_block.py), through Veda's fork of it. + +The predictor weights are NOT part of FunUI and are never shipped with it. +They are downloaded by the person using FunUI, on their own click, from the +official Hugging Face repositories, and stay under the MiniMax H3 Community +License of their base model. + + +---- Veda-on-ComfyUI: MIT License -------------------------------------------- + +MIT License + +Copyright (c) 2026 veda-sparse + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. + + +---- SageAttention 1.0.6: BSD 3-Clause License -------------------------------- + +BSD 3-Clause License + +Copyright (c) 2024, Jintao Zhang, Haofeng Huang. +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +* Redistributions of source code must retain the above copyright notice, this + list of conditions and the following disclaimer. + +* Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. + +* Neither the name of the copyright holder nor the names of its + contributors may be used to endorse or promote products derived from + this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. \ No newline at end of file diff --git a/modules/loaders/diffusion_model/__init__.py b/modules/loaders/diffusion_model/__init__.py index a0b15329..002181fc 100644 --- a/modules/loaders/diffusion_model/__init__.py +++ b/modules/loaders/diffusion_model/__init__.py @@ -18,5 +18,22 @@ def routes(table, base, web): from ..common import slow_int8_build return web.json_response({"slow_int8": slow_int8_build()}) + @table.get(base + "/veda") + async def _veda(_req): + """The Veda predictor files that are in models/veda, the two official downloads with their state, and the licence.""" + from .. import veda_download + return web.json_response(veda_download.listing()) + + @table.post(base + "/veda/download") + async def _veda_download(req): + """Fetch one official predictor straight from Hugging Face: {"id": "t2va" | "r2va"}. Only ever started by a click.""" + from .. import veda_download + try: + body = await req.json() + except ValueError: + body = None + started = veda_download.download(str(body.get("id") or "") if isinstance(body, dict) else "") + return web.json_response({"started": started}) + PROVIDES = {"routes": routes} diff --git a/modules/loaders/diffusion_model/nodes.py b/modules/loaders/diffusion_model/nodes.py index 6b58d63a..f4bdfadd 100644 --- a/modules/loaders/diffusion_model/nodes.py +++ b/modules/loaders/diffusion_model/nodes.py @@ -12,7 +12,7 @@ import folder_paths from comfy_api.latest import io from ..._core import log -from .. import gguf_support, sla_attention +from .. import gguf_support, sla_attention, veda_attention from .. import int8_convrot as int8_convrot_module from .. import mxfp8 as mxfp8_module from ..common import (COMPUTE_DTYPES, head_bank_metadata, WEIGHT_DTYPES, attention_choices, fp32_over_bf16, unsupported_dtypes, @@ -42,7 +42,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", "mxfp8", "compute_dtype", "fp16_accumulation", "attention", - "sla", "sla_method", "sla_engine", "sla_sparsity", "sla_sparsity_start", "sla_tau", "sla_block_size", + "sla", "sla_method", "veda_predictor", "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") @@ -109,7 +109,7 @@ class FunUIDiffusionModelLoader(io.ComfyNode): "text refiner, which must never be sparsified, and " "low-resolution runs where block selection costs more " "than it saves."), - io.Int.Input("sla_dense_last_steps", extra_dict=shows(("sla", [True]), ("sla_method", ["sla", "sol-attn"])), + io.Int.Input("sla_dense_last_steps", extra_dict=shows(("sla", [True]), ("sla_method", ["sla", "sol-attn", "veda"])), default=sla_attention.SLA_DEFAULTS["dense_last_steps"], min=0, max=8, optional=True, tooltip="Run the last N sampling steps at full attention. 0 " @@ -133,7 +133,7 @@ class FunUIDiffusionModelLoader(io.ComfyNode): "runs instead where it is missing, and the status says so). " "triton: FunUI's own kernel, the only one with block size and " "stabilize motion."), - io.String.Input("sla_dense_steps", extra_dict=shows(("sla", [True]), ("sla_method", ["sla", "sol-attn"])), default=sla_attention.SLA_DEFAULTS["dense_steps"], + io.String.Input("sla_dense_steps", extra_dict=shows(("sla", [True]), ("sla_method", ["sla", "sol-attn", "veda"])), default=sla_attention.SLA_DEFAULTS["dense_steps"], optional=True, tooltip="Steps at full attention, counted from 1: '1' is the first " "step, '1-2' the first two, '1,6' the first and sixth; blank or '0', " @@ -166,7 +166,12 @@ class FunUIDiffusionModelLoader(io.ComfyNode): "comfy_kitchen engine. vsa: video cubes, only for models trained " "for it (the status says when the model lacks VSA's layers); run " "by ComfyUI's own Model Sparse Attention, keeping 1 - sla_sparsity " - "of the cubes."), + "of the cubes. veda (experimental, Triton kernel untested on a GPU): " + "a small learned predictor (Veda, arXiv 2605.30325, a file " + "in models/veda) picks the 128-token tiles at its own trained " + "budget, 90% for text-to-video; sla_sparsity and the rising " + "share are not used. Needs CUDA " + "and Triton."), io.Float.Input("sla_tau", extra_dict=shows(("sla", [True]), ("sla_method", ["sol-attn"])), default=sla_attention.SLA_DEFAULTS["tau"], min=0.0, max=4.0, step=0.05, optional=True, tooltip="sol-attn only: higher skips more. About 1.0 keeps ~16% of key " @@ -188,6 +193,12 @@ class FunUIDiffusionModelLoader(io.ComfyNode): "skipped and the status says so. Slightly different pictures: listen to " "the sound first, it spoils before the picture does. A LoRA applied to " "these layers slows them back down."), + # Added last: saved workflows store widget values by position. + io.Combo.Input("veda_predictor", extra_dict=shows(("sla", [True]), ("sla_method", ["veda"])), options=veda_attention.predictor_choices(), + default="auto", optional=True, + tooltip="veda: the predictor file in ComfyUI/models/veda. auto takes the text-to-video one " + "when it is there, else the first. Use the R2VA file for reference-to-video runs. " + "Not there yet? The row below offers the download."), ], outputs=[ io.Model.Output(display_name="model"), @@ -204,7 +215,8 @@ class FunUIDiffusionModelLoader(io.ComfyNode): 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, - sla_sparsity_start: float = None, mxfp8: bool = False) -> io.NodeOutput: + sla_sparsity_start: float = None, mxfp8: bool = False, + veda_predictor: str = None) -> io.NodeOutput: notes = [f"FunUI Diffusion Model Loader | {model_name}"] accumulation = set_fp16_accumulation(fp16_accumulation) @@ -323,7 +335,7 @@ class FunUIDiffusionModelLoader(io.ComfyNode): protect_audio=sla_protect_audio, dense_fn=override, dense_label=attention, engine=sla_engine, dense_steps=sla_dense_steps, references=sla_references, - tail=sla_tail, stabilize_motion=sla_stabilize_motion, method=sla_method, tau=sla_tau) + tail=sla_tail, stabilize_motion=sla_stabilize_motion, method=sla_method, tau=sla_tau, veda_predictor=veda_predictor) notes.append(sla_note) if not installed: if override is not None: @@ -338,7 +350,7 @@ class FunUIDiffusionModelLoader(io.ComfyNode): # Never reach into the model's shape to describe it: a log line that # assumes structure can fail the load it was only meant to narrate. kind = type(getattr(model, "model", model)).__name__ - applied_attention = "sla_h3" if installed else ( + applied_attention = ("veda" if sla_method == "veda" else "sla_h3") if installed else ( attention if override is not None else "default (as launched)") # The file's own storage, not the widget: an int8 file loaded at "bf16" still runs int8. weights = weight_dtype if not stored else ( diff --git a/modules/loaders/sla_attention.py b/modules/loaders/sla_attention.py index 3535c70e..3ced34e3 100644 --- a/modules/loaders/sla_attention.py +++ b/modules/loaders/sla_attention.py @@ -63,12 +63,14 @@ SLA_DEFAULTS = { "references": "off", "tail": False, "stabilize_motion": False, + "veda_predictor": "auto", } ENGINES = ("comfy_kitchen", "triton") # sla: each query block keeps the top (1 - sparsity) of key blocks, what SLA turbo LoRAs were distilled against. # sol-attn: comfy_kitchen picks per head and block by a threshold (tau), training-free. vsa: FastVideo's cube # tiling, for models trained for it -- ComfyUI's own Model Sparse Attention runs it, handed the settings here. -METHODS = ("sla", "sol-attn", "vsa") +# veda: a learned predictor picks the tiles (veda_attention.py), its own file in models/veda, Triton kernel. +METHODS = ("sla", "sol-attn", "vsa", "veda") REFERENCES = ("off", "light", "heavy") _LIGHT_KEEP = 0.15 # "light": each reference span keeps its best 15% of blocks for every query _CK_BLOCK = 64 # sol_attn's fixed key block @@ -501,7 +503,7 @@ def make_wrapper(state, cfg): raise if counted and state["step"] >= n_steps and not state["summed"]: - _summarise(state, cfg) + cfg.get("summarise", _summarise)(state, cfg) state["summed"] = True state["history"].clear() # only needed step to step, not while the model sits cached return out @@ -512,7 +514,7 @@ def make_wrapper(state, cfg): 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): + references=None, tail=None, stabilize_motion=None, method=None, tau=None, veda_predictor=None): """Give `model` block-sparse H3 attention. Returns (model, status line, installed). Weights are untouched: this installs an attention override and a per-step wrapper on @@ -525,7 +527,7 @@ def install_sla(model, sparsity_ratio=None, sparsity_start=None, block_size=None 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) + stabilize_motion=stabilize_motion, method=method, tau=tau, veda_predictor=veda_predictor) cfg = {key: SLA_DEFAULTS[key] if value is None else value for key, value in given.items()} if not bool(cfg["enabled"]): @@ -534,6 +536,9 @@ def install_sla(model, sparsity_ratio=None, sparsity_start=None, block_size=None return model, "SLA: off (dense baseline)", False if not is_h3_model(model): return model, ("SLA: skipped — not a MiniMax H3 model"), False + if cfg["method"] == "veda": # its own kernel and its own checks: neither engine below is used + from . import veda_attention + return veda_attention.install_veda(model, cfg, dense_fn, dense_label) triton_ok, ck_ok = sla_available(), ck_available() if not (triton_ok or ck_ok): return model, ("SLA: skipped — this machine has neither comfy_kitchen's sol_attn nor CUDA+Triton"), False diff --git a/modules/loaders/tests/test_sla_attention.py b/modules/loaders/tests/test_sla_attention.py index b9d9e869..83f3eadc 100644 --- a/modules/loaders/tests/test_sla_attention.py +++ b/modules/loaders/tests/test_sla_attention.py @@ -267,7 +267,7 @@ def test_the_defaults_are_the_validated_ones(): "protect_audio": True, "enabled": True, "engine": "comfy_kitchen", "dense_steps": "1", "method": "sla", "tau": 1.3, "references": "off", - "tail": False, "stabilize_motion": False} + "tail": False, "stabilize_motion": False, "veda_predictor": "auto"} def test_turning_sla_off_leaves_the_settings_in_place(): @@ -503,6 +503,8 @@ def test_the_loader_draws_each_sla_field_only_in_its_mode(): assert rules["sla_sparsity_start"]["sla_method"] == ["sla"] # the rising share is the fixed-share method's assert rules["sla_block_size"]["sla_engine"] == ["triton"] and rules["sla_tail"]["sla_engine"] == ["comfy_kitchen"] assert all(r.get("sla") == [True] for name, r in rules.items() if name.startswith("sla_")) + assert rules["veda_predictor"]["sla_method"] == ["veda"] + assert "veda" in rules["sla_dense_steps"]["sla_method"] and "veda" not in rules["sla_references"]["sla_method"] # veda: trained budget, no references choice assert rules["sla"] == {} and rules["int8_convrot"] == {} diff --git a/modules/loaders/tests/test_veda.py b/modules/loaders/tests/test_veda.py new file mode 100644 index 00000000..3fdfb338 --- /dev/null +++ b/modules/loaders/tests/test_veda.py @@ -0,0 +1,488 @@ +"""Veda learned sparse attention: the torch half on CPU, the override on a REAL tiny H3, every combination. + +The Triton kernel cannot run on this machine: the end-to-end tests hand the install a stand-in kernel (the fp32 +block-sparse reference), so what they prove is everything around it: tiles, predictor, selection, the override +contract, the log line on the real run path. The kernel itself is checked only where CUDA and Triton exist. +""" +import json +import logging +import sys +import types + +import pytest +import torch + +from modules.loaders import sla_attention as sla, veda_attention as va, veda_core as core, veda_download + +HEADS, DIM, LAYERS = 2, 32, 8 # the tiny H3: two heads of 32; eight blocks so MXFP8 has a middle +BLOCKS, KCALLS = 4, 8 # the fixture's blocks, and the kernel calls per pass (two head groups each) +GRID = (10, 8, 8) # the video's token grid the test predictor was "trained" on: 640 rows +CUDA = torch.cuda.is_available() +try: + import triton # noqa: F401 + HAS_TRITON = True +except Exception: # noqa: BLE001 + HAS_TRITON = False + + +# ── helpers ────────────────────────────────────────────────────────────────────────────────────────────────── + +def _write_predictor(folder, name="p.safetensors", ratio=0.5, tiles=None, layers=LAYERS, heads=HEADS, dim=DIM, grid=GRID, + dtype="float8_e4m3fn", meta=None, drop=None): + from safetensors.torch import save_file + g = torch.Generator().manual_seed(1) + tensors = {} + for layer in range(layers): + for kind in ("proj_q", "proj_k"): + value = torch.randn(heads, 3 * dim, dim, generator=g) * 0.1 + key = f"layers.{layer}.{kind}" + if dtype == "float8_e4m3fn": + scale = value.abs().amax(dim=(1, 2)) / 448.0 + tensors[key], tensors[key + ".__scale"] = (value / scale.view(-1, 1, 1)).to(torch.float8_e4m3fn), scale + else: + tensors[key] = value.to(torch.bfloat16) + for key in drop or (): + tensors.pop(key) + plan = {"g": {"geometry": "g", "grid": list(grid), "shapes": ["2x8x8", "4x4x8"], "head_shape": [[i % 2 for i in range(heads)]] * layers}} + budget = {"target_budget_kind": "tiles", "target_budget_value": str(tiles)} if tiles else {"keep_ratio": str(ratio)} + metadata = {"format": core.FORMAT, "num_layers": str(layers), "num_heads": str(heads), "head_dim": str(dim), "dtype": dtype, + "plans": json.dumps(plan), "step": "600", **budget, **(meta or {})} + save_file(tensors, str(folder / name), metadata=metadata) + return folder / name + + +@pytest.fixture +def folder(tmp_path, monkeypatch): + (tmp_path / "veda").mkdir() + monkeypatch.setattr(veda_download, "folder", lambda: tmp_path / "veda") + return tmp_path / "veda" + + +def _spec(refs=((40, (1, 4, 4)),), target=(62, (3, 5, 6)), seq_len=152): + return seq_len, target, tuple(refs), () + + +def _qkv(seq_len, heads=2, dim=32, dtype=torch.float32, device="cpu", seed=0): + g = torch.Generator().manual_seed(seed) + return tuple((torch.randn(seq_len, heads, dim, generator=g) * 1.5).to(device=device, dtype=dtype) for _ in range(3)) + + +def _dense(q, k, v): + p = torch.softmax(torch.einsum("qhd,khd->hqk", q.float(), k.float()) * q.shape[-1] ** -0.5, dim=-1) + return torch.einsum("hqk,khd->qhd", p, v.float()) + + +def _plan_for(grid, heads=2, layers=1): + return core._plan("t", grid, [(2, 8, 8), (4, 4, 8)], [[i % 2 for i in range(heads)]] * layers) + + +def _bundle(layers=1, heads=2, dim=32, generated=(0.5, None), reference=None): + g = torch.Generator().manual_seed(2) + mk = lambda: [(torch.randn(heads, 3 * dim, dim, generator=g) * 0.1).to(torch.bfloat16) for _ in range(layers)] + return types.SimpleNamespace(layers=layers, heads=heads, dim=dim, generated=generated, reference=reference or generated, + proj_q=mk(), proj_k=mk(), plans=[_plan_for((3, 5, 6), heads, layers)]) + + +# ── tiles ──────────────────────────────────────────────────────────────────────────────────────────────────── + +def test_every_row_is_in_exactly_one_tile_and_real_rows_come_first(): + for shape in ((1, 8, 16), (4, 4, 8), (2, 8, 8)): + tiles = core.span_tiles(100, (3, 5, 6), shape) + real = tiles[tiles >= 0] + assert sorted(real.tolist()) == list(range(100, 190)) + count = (tiles >= 0).sum(1) + assert all((tiles[i, :c] >= 0).all() and (tiles[i, c:] < 0).all() for i, c in enumerate(count.tolist())) + + +def test_the_layout_is_a_permutation_and_gather_then_scatter_gives_every_row_back(): + lay = core.build_layout([(40, (1, 4, 4), (1, 8, 16)), (62, (3, 5, 6), (4, 4, 8))], 152) + assert lay.n_ref == 1 and lay.n_video == 3 and lay.n_global == 1 and lay.ref_tokens == 16 and lay.target_tokens == 90 + q = _qkv(152)[0] + heads = torch.arange(2) + out = torch.zeros(153, 2, 32) + core.scatter_(out, core.gather(q, lay, heads), lay, heads) + assert torch.equal(out[:152], q) + + +def test_overlapping_spans_are_refused(): + with pytest.raises(core.VedaError): + core.build_layout([(0, (3, 5, 6), (4, 4, 8)), (50, (3, 5, 6), (4, 4, 8))], 200) + + +# ── predictor and selection ────────────────────────────────────────────────────────────────────────────────── + +def test_pooling_and_logits_match_a_plain_computation_with_partial_tiles(): + lay = core.build_layout([(0, (3, 5, 6), (4, 4, 8))], 100) # 4 tiles, 90 real rows + 10 global + q = _qkv(100)[0] + t = core.gather(q, lay, torch.arange(2)) + feats = core.pool(t, lay) + for tile in range(lay.n_video): + rows = t[tile * 128:(tile + 1) * 128][:int(lay.valid_count[tile])] + want = torch.cat([rows.mean(0), rows.amax(0), rows.amin(0)], dim=-1) + assert torch.allclose(feats[:, tile], want, atol=1e-5) + pq, pk = (torch.randn(2, 96, 32) for _ in range(2)) + got = core.logits(feats, feats, pq, pk) + qh = torch.bmm(feats, pq) + feats[..., :32] + kh = torch.bmm(feats, pk) + feats[..., :32] + assert torch.allclose(got, torch.bmm(qh, kh.transpose(1, 2)) / 32 ** 0.5, atol=1e-4) + + +def test_selection_keeps_the_own_tile_every_global_tile_and_the_budget_on_average(): + lay = core.build_layout([(0, (10, 8, 8), (4, 4, 8))], 700) + scores = torch.randn(2, lay.n_video, lay.n_video) + budget = (0.5, None) + index, keep = core.select(scores, lay, core.column_blocks(lay, budget, budget)) + mask = core.block_mask(index, keep, lay) + n = lay.n_video + assert mask[:, torch.arange(n), torch.arange(n)].all() # a tile always sees itself + assert mask[:, :, n:].all() and mask[:, n:, :].all() # global rows and columns are dense + want = 0.5 * (-(-640 // 128)) ** 2 / n # equal-cost budget, per query row + assert abs(keep.sum(-1).float().mean().item() - want) < 0.1 + assert not mask[:, :n, :n][:, :, ~lay.kv_ok[:n]].any() # empty key tiles are never chosen + + +def test_keeping_everything_is_exactly_dense_attention_with_and_without_tiled_references(): + spec = _spec() + q, k, v = _qkv(152) + want = _dense(q, k, v) + for refs_budget in ((1.0, None), (None, 1000)): # global refs, and the top-k path with k = all + got = core.attention(q, k, v, 0, spec, _plan_for((3, 5, 6)), _bundle(), ((1.0, None), refs_budget), core.reference_attend, {}, {}) + assert torch.allclose(got, want, atol=2e-5), refs_budget + + +def test_a_sparse_budget_differs_from_dense_stays_finite_and_counts_what_it_computed(): + spec, (q, k, v), stats = _spec(), _qkv(152), {} + got = core.attention(q, k, v, 0, spec, _plan_for((3, 5, 6)), _bundle(), ((0.3, None), (0.3, None)), core.reference_attend, {}, stats) + assert torch.isfinite(got).all() and not torch.allclose(got, _dense(q, k, v), atol=1e-3) + assert 0.0 < float(stats["kept"]) / stats["pairs"] < 1.0 + + +# ── the predictor file ─────────────────────────────────────────────────────────────────────────────────────── + +def test_a_predictor_file_loads_with_its_fp8_scales_and_budget(folder): + b = core.load_bundle(str(_write_predictor(folder))) + assert (b.layers, b.heads, b.dim, b.generated) == (LAYERS, HEADS, DIM, (0.5, None)) and b.proj_q[0].dtype == torch.bfloat16 + assert core.load_bundle(str(_write_predictor(folder, "t.safetensors", tiles=32))).generated == (None, 32.0) + assert core.load_bundle(str(_write_predictor(folder, "b.safetensors", dtype="bfloat16"))).proj_k[3].shape == (HEADS, 3 * DIM, DIM) + assert "8 layers x 2 heads, trained 50% sparse" in core.describe_bundle(b) + + +@pytest.mark.parametrize("kw,words", [ + (dict(meta={"format": "other"}), "not a Veda predictor"), + (dict(drop=["layers.2.proj_q"]), "layers.2.proj_q missing"), + (dict(drop=["layers.0.proj_k.__scale"]), "no per-head fp8 scale"), + (dict(meta={"dtype": "int3"}), "unknown storage dtype"), + (dict(meta={"plans": json.dumps({})}), "no tile plans"), + (dict(heads=3, meta={"num_heads": "2"}), "missing or not"), +]) +def test_a_wrong_predictor_file_is_refused_in_words(folder, kw, words): + path = _write_predictor(folder, **kw) + with pytest.raises(core.VedaError, match=words): + core.load_bundle(str(path)) + + +def test_a_file_that_is_not_safetensors_is_refused(folder): + (folder / "x.safetensors").write_bytes(b"not a model") + with pytest.raises(core.VedaError, match="not a readable safetensors"): + core.load_bundle(str(folder / "x.safetensors")) + + +def test_the_plan_is_the_nearest_aspect_then_duration_and_says_when_it_is_not_exact(): + plans = [core._plan("a", (10, 8, 8), [(2, 8, 8)], [[0]]), core._plan("b", (10, 6, 12), [(2, 8, 8)], [[0]])] + assert core.choose_plan(plans, (10, 8, 8))[1:] == (True, "trained for this size (256x256, 0.9 s)") + plan, exact, how = core.choose_plan(plans, (10, 6, 11)) + assert (plan.name, exact) == ("b", False) and how.startswith("nearest trained size: 384x192") + assert core.choose_plan(plans, (10, 12, 6))[0].name == "b_T" # portrait: b turned on its side + + +def test_the_real_h3_layout_maps_to_a_target_grid_and_row_major_references(comfyui): + from comfy.ldm.minimax.model import PackedLayout + seq, (start, grid), refs, skipped = core.describe_layout( + PackedLayout(5, 10, 16, 16, 2, refs=[{"kind": "image", "latent_h": 8, "latent_w": 8}])) + assert grid == (10, 8, 8) and start == seq - 640 and len(refs) == 1 and refs[0][1] == (1, 4, 4) and skipped == () + with pytest.raises(core.VedaError): + core.describe_layout(types.SimpleNamespace()) + + +# ── device and dtype follow the input (the meta device stands in for a GPU) ───────────────────────────────────── + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +def test_every_tensor_the_call_makes_follows_the_inputs_device_and_dtype(dtype): + q, k, v = _qkv(152, dtype=dtype, device="meta") + cache = {} + out = core.attention(q, k, v, 0, _spec(), _plan_for((3, 5, 6)), _bundle(), ((0.5, None), (0.5, None)), + lambda qt, kt, vt, mask, lay: torch.empty_like(qt), cache, {}) + assert out.device.type == "meta" and out.dtype == dtype and out.shape == q.shape + for lay in cache.values(): + assert all(t.device.type == "meta" for t in vars(lay).values() if isinstance(t, torch.Tensor)) + + +# ── the override, on a real tiny H3 ────────────────────────────────────────────────────────────────────────────── + +class _Kernel: + """Stands in for veda_kernel: the fp32 reference, recording what it was handed.""" + def __init__(self): + self.calls, self.dtypes, self.density, self.gap, self.qsum = 0, set(), [], [], [] + + def attend(self, q, k, v, mask, lay): + self.calls += 1 + self.dtypes.add(q.dtype) + self.density.append(mask.float().mean().item()) + self.qsum.append(q.float().abs().sum().item()) + out = core.reference_attend(q, k, v, mask, lay) + every = torch.ones_like(mask) & lay.kv_ok + self.gap.append((out.float() - core.reference_attend(q, k, v, every, lay).float()).abs().max().item()) + return out + + check = staticmethod(lambda device: None) + + +@pytest.fixture +def veda_on(monkeypatch, folder): + """Triton and CUDA 'present' on this machine, the kernel the stand-in.""" + kernel = _Kernel() + monkeypatch.setattr(sla, "sla_available", lambda: True) + monkeypatch.setitem(sys.modules, "modules.loaders.veda_kernel", kernel) + monkeypatch.setattr("modules.loaders.veda_kernel", kernel, raising=False) + return kernel + + +def _big(tiny, bf16=True, size=(10, 16, 16)): + """The tiny H3 on a video of 640 tokens (10 x 8 x 8), in bf16 so the override accepts q/k/v.""" + dm = tiny.patcher.model.diffusion_model + dtype = torch.bfloat16 if bf16 else torch.float32 + if bf16: + dm.to(dtype) + tiny.video = torch.randn(1, 24, *size).to(dtype) + tiny.audio = tiny.audio.to(dtype) + tiny.context = tiny.context.to(dtype) + return tiny + + +def _install(tiny, **kw): + given = dict(method="veda", min_seq_len=0, dense_steps="", sparsity_ratio=0.9) + patched, note, installed = sla.install_sla(tiny.patcher, **{**given, **kw}) + return patched, note, installed + + +def _last_step(tiny, patched): + return tiny.run(patched, sigma=0.5, sigmas=torch.tensor([1.0, 0.5, 0.0])) # step 2 of 2: the run's summary is written + + +def test_keeping_everything_on_a_real_h3_gives_the_dense_picture(veda_on, folder, comfyui, tiny_h3): + _write_predictor(folder, ratio=1.0) + tiny = _big(tiny_h3) + dense = _last_step(tiny, tiny.patcher)[0] + patched, note, installed = _install(tiny) + assert installed and "predictor p.safetensors" in note and "EXPERIMENTAL" in note + got = _last_step(tiny, patched)[0] + assert veda_on.calls == KCALLS and veda_on.dtypes == {torch.bfloat16} # every block, handed bf16 q/k/v + assert torch.allclose(got.float(), dense.float(), atol=0.1) + + +def test_a_sparse_run_fires_logs_the_line_on_the_real_path_and_differs(veda_on, folder, comfyui, tiny_h3, caplog): + _write_predictor(folder, ratio=0.3) + tiny = _big(tiny_h3) + dense = _last_step(tiny, tiny.patcher)[0] + patched, _, installed = _install(tiny) + with caplog.at_level(logging.INFO): + got = _last_step(tiny, patched)[0] + assert installed and torch.isfinite(got.float()).all() and torch.isfinite(dense.float()).all() + assert max(veda_on.density) < 0.95 and max(veda_on.gap) > 1e-3 # tiles were skipped, and it shows in the output + line = next(r.getMessage() for r in caplog.records if r.getMessage().startswith("[FunUI Veda] triton-int8")) + assert f"{BLOCKS} calls" in line and "256x256" in line and "of full attention" in line and "trained for this size" in line + assert "asked trained (70% sparse)" in line and "sparse, asked" in line and "global tiles pinned" in line and "p.safetensors" in line + assert "1 dense fall-throughs (1 no layout for this call)" in line # the text refiner's call went to the backend + assert "full attention on no step of 2" in caplog.text + + +def test_a_second_run_logs_again_and_a_new_size_uses_the_nearest_plan(veda_on, folder, comfyui, tiny_h3, caplog): + """Two stages: the same patched model samples a small picture, then a larger one.""" + _write_predictor(folder, ratio=0.5) + tiny = _big(tiny_h3) + patched, _, _ = _install(tiny) + with caplog.at_level(logging.INFO): + _last_step(tiny, patched) + tiny.video = torch.randn(1, 24, 10, 24, 24).to(torch.bfloat16) + _last_step(tiny, patched) + lines = [r.getMessage() for r in caplog.records if r.getMessage().startswith("[FunUI Veda] triton-int8")] + assert len(lines) == 2 and "trained for this size" in lines[0] and "nearest trained size" in lines[1] and "384x384" in lines[1] + + +def test_dense_steps_stay_with_the_backend_and_the_log_names_them(veda_on, folder, comfyui, tiny_h3, caplog): + _write_predictor(folder, ratio=0.5) + tiny = _big(tiny_h3) + patched, note, _ = _install(tiny, dense_steps="1") + assert "steps 1" in note + schedule = torch.tensor([1.0, 0.6, 0.3, 0.0]) + tiny.run(patched, sigma=1.0, sigmas=schedule) # step 1: dense on purpose + assert veda_on.calls == 0 + with caplog.at_level(logging.INFO): + for sigma in (0.6, 0.3): + tiny.run(patched, sigma=sigma, sigmas=schedule) + assert veda_on.calls == 2 * KCALLS and "full attention on steps 1 of 3" in caplog.text and "dense step" in caplog.text + + +def test_veda_runs_at_the_trained_budget_whatever_sparsity_or_ramp_are_set(veda_on, folder, comfyui, tiny_h3, caplog): + """First version: Veda's own defaults only, so the comparison with the current SLA workflow is of like with unlike on + purpose. sla_sparsity, the rising share and the references choice are not read, and the status says so.""" + _write_predictor(folder, ratio=0.5) + schedule = torch.tensor([1.0, 0.75, 0.5, 0.25, 0.0]) + tiny = _big(tiny_h3) + steps = [] + for kw in ({}, dict(sparsity_start=0.6, sparsity_ratio=0.85, references="heavy")): + patched, note, _ = _install(tiny, **kw) + veda_on.density.clear() + with caplog.at_level(logging.INFO): + for sigma in schedule[:-1].tolist(): + tiny.run(patched, sigma=sigma, sigmas=schedule) + steps.append(list(veda_on.density)) + assert steps[0] == steps[1] and steps[0][:2] == steps[0][-2:] # the same tiles every step, with or without the settings + assert "sla_sparsity, sla_sparsity_start, sla_references: not used by veda" in note and "sparsity trained (50% sparse)" in note + assert "asked trained (50% sparse)" in caplog.text + + +def test_the_attention_explorer_runs_first_and_veda_sees_what_it_changed(veda_on, folder, comfyui, tiny_h3): + """The explorer's override wraps the one installed at load (core/dit_hooks.add_attention_override): it edits q, then + calls the rest of the chain, which is Veda. Veda must receive the edited q and still take every call.""" + from core import dit_hooks + _write_predictor(folder, ratio=0.5) + tiny = _big(tiny_h3) + patched, _, _ = _install(tiny) + seen = [] + veda_on.attend = lambda q, k, v, mask, lay: (seen.append(q.float().abs().mean().item()), core.reference_attend(q, k, v, mask, lay))[1] + _last_step(tiny, patched) + plain, calls = list(seen), veda_on.calls + seen.clear() + explorer = dit_hooks.add_attention_override + explorer(patched, "explorer", lambda func, q, k, v, heads, **kw: func(q * 1.5, k, v, heads, **kw)) + _last_step(tiny, patched) + assert len(seen) == len(plain) and all(abs(b / a - 1.5) < 0.05 for a, b in zip(plain, seen)) + + +def test_a_lora_on_the_attention_projection_changes_q_k_v_and_veda_still_takes_the_call(veda_on, folder, comfyui, tiny_h3): + _write_predictor(folder, ratio=0.5) + tiny = _big(tiny_h3) + _last_step(tiny, _install(tiny)[0]) + base_q = veda_on.qsum[0] + key = "diffusion_model.blocks.0.attn.qkv_proj.weight" + weight = tiny.patcher.model.diffusion_model.blocks[0].attn.qkv_proj.weight + lora = tiny.patcher.clone() + lora.add_patches({key: ("diff", (torch.randn_like(weight) * 0.3,))}, 1.0) + veda_on.calls, veda_on.qsum = 0, [] + patched, _, installed = _install(types.SimpleNamespace(patcher=lora)) + assert installed + lora.patch_model(load_weights=True) + try: + got = tiny.run(patched, sigma=0.5, sigmas=torch.tensor([1.0, 0.5, 0.0]))[0] + finally: + lora.unpatch_model() + assert veda_on.calls == KCALLS and torch.isfinite(got.float()).all() + assert abs(veda_on.qsum[0] - base_q) > 1e-3 * base_q # block 0's q is the LoRA's, and Veda got it + + +def test_mxfp8_weights_leave_q_k_v_in_bf16_so_veda_and_mxfp8_compose(veda_on, folder, comfyui, tiny_h3): + """MXFP8 stores weights only. The attention override meets q/k/v after the projection, in the compute dtype, + which is what the predictor was trained on: with MXFP8 blocks in the model it is still bf16 and still taken.""" + pytest.importorskip("comfy.quant_ops") + from comfy.quant_ops import QUANT_ALGOS + if "mxfp8" not in QUANT_ALGOS: + pytest.skip("this comfy_kitchen has no MXFP8") + import comfy.ops + from comfy.ldm.minimax.model import MiniMaxH3Model + from modules.loaders import mxfp8 + _write_predictor(folder, ratio=0.5) + tiny = _big(type(tiny_h3)(blocks=8)) + sd = {k: v.clone() for k, v in tiny.patcher.model.diffusion_model.state_dict().items()} + converted, _ = mxfp8.quantize_state_dict(sd) + assert converted > 0 + qm = MiniMaxH3Model(hidden_size=64, num_layers=LAYERS, token_refiner_num_layers=1, num_attention_heads=2, attention_head_dim=32, + ffn_hidden_size=64, text_dim=48, timestep_input_dim=32, time_embed_hidden_size=64, time_embed_dim=32, + rope_inv_freq_len=4, dtype=torch.bfloat16, operations=comfy.ops.mixed_precision_ops({})) + qm.load_state_dict(sd, strict=False) + assert qm.blocks[3].attn.qkv_proj.quant_format == "mxfp8" + qm.rope.inv_freq.copy_(tiny.patcher.model.diffusion_model.rope.inv_freq) + for p in qm.parameters(): + p.requires_grad_(False) + tiny.patcher.model.diffusion_model = qm + patched, _, installed = _install(tiny) + try: + out = _last_step(tiny, patched)[0] + except ValueError: + # This machine's CPU cannot run the MXFP8 matmul (Blackwell GPUs do). Blocks 0 and 1 stay bf16 and run, so Veda's calls + # for them were made before the first MXFP8 layer failed: what the kernel was handed is still checked. + assert installed and veda_on.calls >= 2 and veda_on.dtypes == {torch.bfloat16} + return + assert installed and veda_on.calls == 16 and veda_on.dtypes == {torch.bfloat16} and torch.isfinite(out.float()).all() + + +# ── when Veda cannot run, it says why and the backend stays ────────────────────────────────────────────────────── + +def test_install_says_why_it_skipped_and_hands_the_backend_back(monkeypatch, folder, comfyui, tiny_h3): + tiny = _big(tiny_h3) + monkeypatch.setattr(sla, "sla_available", lambda: False) + out, note, installed = _install(tiny) + assert not installed and out is tiny.patcher and "needs CUDA and Triton" in note + monkeypatch.setattr(sla, "sla_available", lambda: True) + out, note, installed = _install(tiny) + assert not installed and "no predictor in ComfyUI/models/veda" in note + _write_predictor(folder, "x.safetensors", meta={"format": "nope"}) + assert "not a Veda predictor" in _install(tiny)[1] + assert "gone.safetensors is not in" in _install(tiny, veda_predictor="gone.safetensors")[1] + + +def test_a_predictor_made_for_other_heads_runs_dense_and_the_log_says_why(veda_on, folder, comfyui, tiny_h3, caplog): + _write_predictor(folder, heads=4, dim=32) + tiny = _big(tiny_h3) + patched, _, installed = _install(tiny) + with caplog.at_level(logging.INFO): + _last_step(tiny, patched) + assert installed and veda_on.calls == 0 + assert "installed but never invoked" in caplog.text and "2 heads of 32, predictor has 4 of 32" in caplog.text + + +def test_a_failing_kernel_costs_speed_not_the_run_and_is_named(veda_on, folder, comfyui, tiny_h3, caplog): + _write_predictor(folder, ratio=0.5) + veda_on.attend = lambda *a: (_ for _ in ()).throw(RuntimeError("no such kernel")) + tiny = _big(tiny_h3) + patched, _, _ = _install(tiny) + with caplog.at_level(logging.INFO): + out = _last_step(tiny, patched)[0] + assert torch.isfinite(out.float()).all() and "RuntimeError: no such kernel" in caplog.text + + +def test_the_self_check_failing_is_reported_and_runs_dense(veda_on, folder, comfyui, tiny_h3, caplog): + _write_predictor(folder) + veda_on.check = staticmethod(lambda device: "the Triton kernel ignores the block mask on this device") + tiny = _big(tiny_h3) + with caplog.at_level(logging.INFO): + _last_step(tiny, _install(tiny)[0]) + assert "ignores the block mask" in caplog.text and veda_on.calls == 0 + + +def test_float32_and_short_sequences_stay_with_the_backend(veda_on, folder, comfyui, tiny_h3): + _write_predictor(folder) + tiny = _big(tiny_h3, bf16=False) # float32 q/k/v: never the kernel + _last_step(tiny, _install(tiny)[0]) + assert veda_on.calls == 0 + tiny = _big(tiny_h3.__class__(), bf16=True) + _last_step(tiny, _install(tiny, min_seq_len=100000)[0]) + assert veda_on.calls == 0 + + +def test_the_loader_passes_the_predictor_through_and_names_it(veda_on, folder, comfyui, tiny_h3, monkeypatch): + from modules.loaders.diffusion_model.nodes import FunUIDiffusionModelLoader as L + ids = [i.id for i in L.define_schema().inputs] + assert ids[-1] == "veda_predictor" # added last: saved workflows read by position + _write_predictor(folder, "a.safetensors") + _write_predictor(folder, "b.safetensors", ratio=0.2) + _, note, installed = _install(_big(tiny_h3), veda_predictor="b.safetensors") + assert installed and "predictor b.safetensors" in note + + +# ── the kernel, where there is one ─────────────────────────────────────────────────────────────────────────────── + +@pytest.mark.skipif(not (CUDA and HAS_TRITON), reason="needs CUDA + Triton") +def test_the_triton_kernel_passes_its_own_self_check_on_this_gpu(): + from modules.loaders import veda_kernel + assert veda_kernel.check(torch.device("cuda")) is None diff --git a/modules/loaders/tests/test_veda_download.py b/modules/loaders/tests/test_veda_download.py new file mode 100644 index 00000000..e86ca0dc --- /dev/null +++ b/modules/loaders/tests/test_veda_download.py @@ -0,0 +1,154 @@ +"""The optional predictor download: one downloader (core/fetch), only on a click, from the official address, and every way it can fail.""" +import asyncio +import hashlib +import http.server +import threading +import time + +import pytest + +from core import fetch +from modules.loaders import veda_download as vd + +BODY = b"x" * 3_000_000 + + +@pytest.fixture +def dl(tmp_path): + (tmp_path / "dl").mkdir() + return tmp_path / "dl" + + +@pytest.fixture +def server(): + """A local server with a few misbehaving files: /ok, /short (cut off after half), /gone (404).""" + class H(http.server.BaseHTTPRequestHandler): + def log_message(self, *a): + pass + + def do_GET(self): + if self.path == "/gone": + return self.send_error(404) + self.send_response(200) + self.send_header("Content-Length", str(len(BODY))) + self.end_headers() + self.wfile.write(BODY if self.path == "/ok" else BODY[:len(BODY) // 2]) # /short: the connection closes early + + srv = http.server.ThreadingHTTPServer(("127.0.0.1", 0), H) + threading.Thread(target=srv.serve_forever, daemon=True).start() + yield f"http://127.0.0.1:{srv.server_port}" + srv.shutdown() + srv.server_close() + + +def _left(folder): + return sorted(p.name for p in folder.iterdir()) + + +def test_a_good_download_is_renamed_only_when_complete_and_reports_progress(server, dl): + seen = [] + fetch.to_file(f"{server}/ok", dl / "m.safetensors", lambda d, t: seen.append((d, t)), size=len(BODY), sha256=hashlib.sha256(BODY).hexdigest()) + assert _left(dl) == ["m.safetensors"] and (dl / "m.safetensors").read_bytes() == BODY + assert seen[-1] == (len(BODY), len(BODY)) and len(seen) >= 3 + + +@pytest.mark.parametrize("path,kw,words", [ + ("/short", {}, "stopped at"), # a cut connection + ("/ok", {"size": 1234}, "expected 1234"), # the wrong file size + ("/ok", {"sha256": "0" * 64}, "checksum"), # the right size, the wrong bytes + ("/gone", {}, "404"), # the file is not there +]) +def test_a_bad_download_leaves_nothing_behind_and_says_why(server, dl, path, kw, words): + with pytest.raises(OSError, match=words): + fetch.to_file(server + path, dl / "m.safetensors", **kw) + assert _left(dl) == [] # no file, and no .part that could be mistaken for one + + +def test_no_network_is_an_oserror_and_leaves_nothing(dl): + with pytest.raises(OSError): + fetch.to_file("http://127.0.0.1:9/none", dl / "m.safetensors", timeout=2) + assert _left(dl) == [] + + +@pytest.fixture +def veda_dir(tmp_path, monkeypatch): + monkeypatch.setattr(vd, "folder", lambda: tmp_path / "models" / "veda") + vd._state.clear() + return tmp_path / "models" / "veda" + + +def _wait(model_id): + for _ in range(200): + if vd._state[model_id]["state"] != "downloading": + return vd._state[model_id] + time.sleep(0.01) + raise AssertionError("the download never finished") + + +def test_the_listing_names_the_official_files_sizes_and_the_licence(veda_dir): + rows = {r["id"]: r for r in vd.listing()["models"]} + assert rows["t2va"]["file"] == "minimax_h3_t2va_veda_8nfe_600step_preview_fp8.safetensors" and rows["t2va"]["size"] == 275415648 + assert rows["r2va"]["file"] == "minimax_h3_r2va_veda_preview_fp8.safetensors" and rows["r2va"]["size"] == 275415752 + assert rows["t2va"]["url"] == ("https://huggingface.co/Veda-Sparse/Minimax-H3-T2VA-Veda-8NFE-600Step-Preview/resolve/main/" + "minimax_h3_t2va_veda_8nfe_600step_preview_fp8.safetensors") + assert rows["r2va"]["url"].startswith("https://huggingface.co/Veda-Sparse/Minimax-H3-R2VA-Veda-Preview/resolve/main/") + assert vd.listing()["license"]["name"] == "MiniMax H3 Community License" and vd.listing()["license"]["url"].endswith("/blob/main/LICENSE") + assert {r["state"] for r in rows.values()} == {"absent"} + + +def test_download_goes_to_models_veda_with_the_pinned_size_and_checksum(veda_dir): + calls = [] + + def get(url, dest, progress, size=None, sha256=None): + calls.append((url, dest, size, sha256)) + progress(10, size) + dest.write_bytes(b"model") + assert vd.download("t2va", get=get) is True + assert _wait("t2va")["state"] == "ready" + url, dest, size, sha = calls[0] + assert url.startswith("https://huggingface.co/") and dest == veda_dir / "minimax_h3_t2va_veda_8nfe_600step_preview_fp8.safetensors" + assert size == 275415648 and len(sha) == 64 + assert vd.files() == [dest.name] and next(r for r in vd.listing()["models"] if r["id"] == "t2va")["state"] == "ready" + assert vd.download("t2va", get=get) is False # already there: nothing is fetched twice + + +def test_unknown_ids_are_refused_and_a_failure_is_shown_beside_the_button(veda_dir): + assert vd.download("../../etc", get=lambda *a, **k: None) is False + assert vd.download("", get=lambda *a, **k: None) is False + + def boom(url, dest, progress, size=None, sha256=None): + raise OSError("the download stopped at 5 of 10 bytes") + assert vd.download("r2va", get=boom) is True + state = _wait("r2va") + assert state["state"] == "error" and "stopped at 5 of 10" in state["error"] + row = next(r for r in vd.listing()["models"] if r["id"] == "r2va") + assert row["state"] == "error" and "stopped at" in row["error"] and vd.files() == [] + assert vd.download("r2va", get=lambda url, dest, p, size=None, sha256=None: dest.write_bytes(b"x")) is True # retry is allowed + + +def test_a_part_file_is_not_a_predictor(veda_dir): + veda_dir.mkdir(parents=True) + (veda_dir / "a.safetensors.part").write_bytes(b"half") + (veda_dir / "b.safetensors").write_bytes(b"x") + assert vd.files() == ["b.safetensors"] + + +def test_the_routes_list_and_start_but_never_fetch_on_their_own(veda_dir, monkeypatch): + from aiohttp import web + from aiohttp.test_utils import TestClient, TestServer + from modules.loaders import diffusion_model + started = [] + monkeypatch.setattr(vd, "download", lambda model_id, get=None: started.append(model_id) or True) + + async def go(): + app, table = web.Application(), web.RouteTableDef() + diffusion_model.routes(table, "/m", web) + app.add_routes(table) + async with TestClient(TestServer(app)) as c: + listing = await (await c.get("/m/veda")).json() + assert started == [] # looking never downloads + go_ = await (await c.post("/m/veda/download", json={"id": "t2va"})).json() + junk = await c.post("/m/veda/download", data=b"not json") + return listing, go_, junk.status + listing, go_, junk = asyncio.run(go()) + assert [m["id"] for m in listing["models"]] == ["t2va", "r2va"] and go_ == {"started": True} and started == ["t2va", ""] and junk == 200 diff --git a/modules/loaders/veda_attention.py b/modules/loaders/veda_attention.py new file mode 100644 index 00000000..b953270b --- /dev/null +++ b/modules/loaders/veda_attention.py @@ -0,0 +1,194 @@ +"""Veda learned sparse attention for MiniMax H3, as the loader's SLA method "veda". + +Veda (Han, Yang, Hu, Mei, Jiang, Qi; MIT; arXiv 2605.30325) ported from Veda-on-ComfyUI: see veda_core.py for +the method and VEDA-NOTICE.txt for the licences. EXPERIMENTAL and UNTESTED on a GPU: the pure-torch half is +tested on CPU, the Triton kernel (veda_kernel.py) only ever ran on the author's machine of origin. + +It rides the same hook as SLA (`transformer_options["optimized_attention_override"]`, plus SLA's per-step wrapper +for the step counter, the dense steps and the rising share), so it needs no weight change and meets q/k/v only after +the LoRA-patched projections: LoRA, int8_convrot and MXFP8 (weights only) leave what it sees as bf16/fp16 q/k/v, which +is what the predictor was trained on. Calls it does not take (the text refiner, masked calls, short sequences, dense +steps, a head count the predictor was not made for) go to the chosen attention backend, and the log says how many +and why. When Veda cannot run at all, the loader says why and the backend is installed on its own. +""" +import collections +import logging + +from . import sla_attention as sla, veda_core as core, veda_download + +_NAMED = ("auto",) + + +def predictor_choices(): + try: + return list(_NAMED) + veda_download.files() + except Exception: # noqa: BLE001 - no ComfyUI folders here (tests, dev server): a list of nothing is fine + return list(_NAMED) + + +def _pick(name): + """The predictor file's path, or (None, why). "auto": the T2VA file when it is there, else the first.""" + try: + have = veda_download.files() + except Exception: # noqa: BLE001 + have = [] + if not have: + return None, "no predictor in ComfyUI/models/veda (the Veda predictor row in the app offers the download)" + if name in have: + return veda_download.folder() / name, None + if name not in ("", *_NAMED): + return None, f"the predictor {name} is not in ComfyUI/models/veda" + t2va = veda_download.MODELS[0]["file"] + return veda_download.folder() / (t2va if t2va in have else have[0]), None + + +def _fold(veda): + """The step that just ended: its work joins the run's totals (one read of the device per step, not per call), and + its own kept share is what the summary reports, as SLA reports its last step.""" + stats = veda["stats"] + if stats.get("dense"): + veda["run"]["work"] += float(stats["work"]) + veda["run"]["dense"] += stats["dense"] + veda["last"] = (float(stats["kept"]) / stats["pairs"] if stats.get("pairs") else 1.0, stats.get("global", 0)) + stats.clear() + + +def make_override(state, cfg, bundle, attend, dense_fn=None, dense_label=None): + """The `optimized_attention_override`. `attend` is the kernel callable, or None for veda_kernel.attend (checked once + per device, on the first call, against the fp32 reference). A kernel that fails costs speed, not the run.""" + import torch + veda = state["veda"] + ok_dtypes = (torch.bfloat16, torch.float16) + + def override(func, q, k, v, heads, mask=None, attn_precision=None, skip_reshape=False, skip_output_reshape=False, **kwargs): + def dense(why): + state["dense"] += 1 + veda["why"][why] += 1 + run = (lambda *a, **kw: dense_fn(func, *a, **kw)) if dense_fn is not None else func + return run(q, k, v, heads, mask=mask, attn_precision=attn_precision, skip_reshape=skip_reshape, + skip_output_reshape=skip_output_reshape, **kwargs) + + if state["backend"] is None: + state["backend"] = dense_label or getattr(func, "__name__", repr(func)) + to = kwargs.get("transformer_options") or {} + layout, layer = to.get("minimax_h3_layout"), to.get("block_index") + if mask is not None: + return dense("masked call") + if not skip_reshape or q.ndim != 4 or q.shape[0] != 1 or q.shape != k.shape or q.shape != v.shape: + return dense("not H3's packed self-attention") + if q.shape[2] < cfg["min_seq_len"]: + return dense("shorter than min_seq_len") + if layout is None or getattr(layout, "seq_len", None) != q.shape[2]: + return dense("no layout for this call") + if q.dtype not in ok_dtypes: + return dense(f"{q.dtype} q/k/v") + if to.get("_funui_sla_dense", False): + return dense("dense step") + if not isinstance(layer, int): + layer = state["layer"] + state["layer"] += 1 + if veda["step"] != state["step"]: + _fold(veda) + veda["step"] = state["step"] + if layer >= bundle.layers: + return dense("layer outside the predictor") + if q.shape[1] != bundle.heads or q.shape[3] != bundle.dim: + return dense(f"{q.shape[1]} heads of {q.shape[3]}, predictor has {bundle.heads} of {bundle.dim}") + if state["failed"] is not None: + return dense("kernel failed earlier") + try: + fn = attend + if fn is None: + from . import veda_kernel + fn = veda_kernel.attend + if str(q.device) not in veda["checked"]: + veda["checked"][str(q.device)] = veda_kernel.check(q.device) + if veda["checked"][str(q.device)]: + raise RuntimeError(veda["checked"][str(q.device)]) + spec = veda["specs"].get(id(layout)) + if spec is None or spec[0] is not layout: + spec = veda["specs"][id(layout)] = (layout, core.describe_layout(layout)) + spec = spec[1] + plan = veda["plans"].get(spec[1][1]) + if plan is None: + plan = veda["plans"][spec[1][1]] = core.choose_plan(bundle.plans, spec[1][1]) + with sla.pause_malloc_graph(): # the cached tile layouts outlive the call: made outside ComfyUI's compiler graph + out = core.attention(q[0].transpose(0, 1), k[0].transpose(0, 1), v[0].transpose(0, 1), layer, spec, plan[0], bundle, + (bundle.generated, bundle.reference), fn, veda["tiles"], veda["stats"]) + veda["geometry"] = (spec, plan[2]) + state["calls"] += 1 + state["seq"] = q.shape[2] + return out.transpose(0, 1).unsqueeze(0) if skip_output_reshape else out.reshape(1, q.shape[2], -1) + except Exception as exc: # noqa: BLE001 - a bad kernel must not kill the run + state["failed"] = f"{exc.__class__.__name__}: {exc}" + logging.warning("[FunUI Veda] sparse attention failed, dense from here this run: %s", state["failed"], exc_info=True) + return dense("kernel failed") + + return override + + +def summarise(state, cfg): + """One line per sampling run (the wrapper calls it when the run ends), then the per-run caches are let go.""" + veda, stats = state["veda"], state["veda"]["stats"] + why = ", ".join(f"{n} {w}" for w, n in veda["why"].most_common()) + if state["calls"] == 0: + logging.warning("[FunUI Veda] installed but never invoked: attention was NOT sparsified (%d dense calls%s). Check that the " + "model reaching the sampler is the one this loader produced.", state["dense"], f": {why}" if why else "") + else: + # Same fields in the same order as the "[FunUI SLA]" line, so the two read side by side. Tile pairs are the video x + # video ones, counted over the whole run (SLA counts its last step); the global tiles (text, audio) are always exact. + _fold(veda) + spec, how = veda["geometry"] + done = min(1.0, veda["run"]["work"] / veda["run"]["dense"]) if veda["run"]["dense"] else 0.0 + kept, pinned = veda["last"] + logging.info("[FunUI Veda] %s | %d calls | S=%d | video tile pairs %.1f%% kept in the last step (%.1f%% sparse, asked %s) | %d global tiles pinned " + "| %d dense fall-throughs%s | displaced %s | computed %.1f%% of full attention | video %s, tile plan: %s | %s", + veda["kernel"], state["calls"], state["seq"], kept * 100.0, (1.0 - kept) * 100.0, cfg["asked"], pinned, + state["dense"], f" ({why})" if why else "", state["backend"] or "?", done * 100.0, core.describe_grid(spec[1][1]), how, + cfg["predictor_name"]) + logging.info("[FunUI Veda] full attention on %s of %d", f"steps {sla._runs(state['dense_at'])}" if state["dense_at"] else "no step", state["n"]) + if state["failed"] is not None: + logging.warning("[FunUI Veda] the kernel fell back to dense at least once: %s", state["failed"]) + stats.clear() + veda["run"].update(work=0.0, dense=0.0) + veda["step"] = None + veda["tiles"].clear() + veda["why"].clear() + veda["specs"].clear() + + +def install_veda(model, cfg, dense_fn=None, dense_label=None, attend=None): + """Give `model` Veda attention. Returns (model, status, installed), the contract of sla.install_sla: `installed` is + False when Veda did not take, with the reason in the status, so the caller installs the chosen backend alone. + `attend` replaces the Triton kernel (tests only).""" + if attend is None and not sla.sla_available(): + return model, "SLA: veda needs CUDA and Triton (its kernel is Triton's), not here", False + path, why = _pick(cfg.get("veda_predictor")) + if path is None: + return model, f"SLA: veda skipped, {why}", False + try: + bundle = core.load_bundle(str(path)) + except (core.VedaError, OSError) as exc: + return model, f"SLA: veda skipped, {exc}", False + unused = [n for n, on in (("sla_sparsity", float(cfg["sparsity_ratio"]) != sla.SLA_DEFAULTS["sparsity_ratio"]), + ("sla_sparsity_start", float(cfg["sparsity_start"]) > 0.0), ("sla_references", cfg["references"] != "off")) if on] + cfg = {**cfg, "sparsity_ratio": sla.SLA_DEFAULTS["sparsity_ratio"], "sparsity_start": 0.0, "references": "off", + "min_seq_len": int(cfg["min_seq_len"]), "dense_last_steps": int(cfg["dense_last_steps"]), + "predictor_name": path.name, "engine": "veda", "summarise": summarise} + cfg["dense_steps"], bad = sla.parse_steps(cfg["dense_steps"]) + notes = [f"dense_steps: ignored {', '.join(bad)} (steps count from 1: '1' is the first, '1-2' the first two)"] if bad else [] + if unused: + notes.append(f"{', '.join(unused)}: not used by veda, which runs at the predictor's trained budget") + state = sla.new_state() + state["veda"] = {"stats": {}, "tiles": {}, "specs": {}, "plans": {}, "checked": {}, "why": collections.Counter(), "geometry": None, "step": None, "run": {"work": 0.0, "dense": 0.0}, "last": (1.0, 0), "kernel": "triton-int8" if attend is None else "stand-in kernel"} + patched = sla.patching.clone(model) + to = patched.model_options.get("transformer_options", {}).copy() + to["optimized_attention_override"] = make_override(state, cfg, bundle, attend, dense_fn, dense_label) + patched.model_options["transformer_options"] = to + patched.add_wrapper_with_key("diffusion_model", "funui_sla_state", sla.make_wrapper(state, cfg)) + cfg["asked"] = f"trained ({core.describe_budget(bundle.generated)})" + return patched, "\n".join([ + f"SLA on | veda (EXPERIMENTAL, Triton kernel untested on a GPU) | predictor {core.describe_bundle(bundle)} | sparsity {cfg['asked']} | " + f"references {core.describe_budget(bundle.reference)} | min_seq_len={cfg['min_seq_len']} | full attention on steps " + f"{sla._runs(i + 1 for i in cfg['dense_steps']) or 'none'} and the last {cfg['dense_last_steps']} | dense calls -> {dense_label or 'as launched'}", + *(f"SLA: {n}" for n in notes)]), True diff --git a/modules/loaders/veda_core.py b/modules/loaders/veda_core.py new file mode 100644 index 00000000..2044c3d6 --- /dev/null +++ b/modules/loaders/veda_core.py @@ -0,0 +1,425 @@ +"""Veda sparse attention, the torch half: tiles, plans, predictor, selection, one layer's call. + +Veda (Han, Yang, Hu, Mei, Jiang, Qi; MIT; arXiv 2605.30325) lets a small distilled predictor say +which 128-token 3D tiles of the attention map matter, and computes only the top ~10% of them. +Written here from Veda-on-ComfyUI's core (itself a rewrite of Miowtion's miowtion/veda, MIT), in +this codebase's style; see VEDA-NOTICE.txt. The released predictors were trained against these +exact rules (tile order, pooling, budget, own-tile, global rows), so none of them may drift. + +Pure torch, no ComfyUI import: it runs and is tested on any device. The kernel that attends the +kept tiles is a callable handed in (veda_kernel.attend on a GPU, `reference_attend` in tests). +Every tensor made here is built on, or moved to, the device of the tensor it meets. +""" +import functools +import json +import math +import os +from types import SimpleNamespace + +import torch + +TILE = 128 +FORMAT = "miowtion-veda-predictor-v1" +_STORAGE = {"float32": torch.float32, "bfloat16": torch.bfloat16, "float8_e4m3fn": torch.float8_e4m3fn} +FPS = 24 +_CHUNK_BYTES = 256 * 2**20 # one tile-ordered q/k/v/out copy of a head chunk; four are alive at once + + +class VedaError(ValueError): + """The file or layout cannot be used; the message is shown to the person.""" + + +# ── budgets: (ratio, tiles), exactly one set. ratio = share of key tiles KEPT at equal kernel cost ───────────── + +def sparsity_budget(share): + """Skip `share` (0..1) of the key tiles.""" + return (1.0 - float(share), None) + + +def keeps_all(budget): + return budget[1] is None and budget[0] >= 1.0 + + +def describe_budget(budget): + return f"{budget[1]:g} tiles" if budget[1] is not None else f"{100.0 * (1.0 - min(budget[0], 1.0)):g}% sparse" + + +def _per_row(budget, real_tokens, n_cols): + if budget[1] is not None: + return float(budget[1]) + ideal = math.ceil(real_tokens / TILE) # equal cost: a padded tiling never gets a bigger budget + return budget[0] * ideal * ideal / n_cols + + +def _split(budget, n_cols): + """(k_lo, k_hi, frac) with 1 <= k_lo <= k_hi <= n_cols.""" + if budget >= n_cols: + return n_cols, n_cols, 0.0 + lo = min(max(1, math.floor(budget)), n_cols) + return lo, min(lo + 1, n_cols), round(min(max(budget - lo, 0.0), 1.0), 12) + + +@functools.lru_cache(maxsize=256) +def _bresenham(n_rows, frac): + ramp = torch.floor(torch.arange(n_rows + 1, dtype=torch.float64) * frac) + return ramp[1:] > ramp[:-1] + + +# ── tiles ────────────────────────────────────────────────────────────────────────────────────────────────────── + +def shapes(): + """Every power-of-two (t, h, w) with t * h * w = TILE.""" + return sorted((2 ** i, 2 ** j, TILE >> (i + j)) for i in range(8) for j in range(8 - i)) + + +def _padded(grid, shape): + return tuple(-(-g // s) * s for g, s in zip(grid, shape)) + + +def _count(grid, shape): + p = _padded(grid, shape) + return p[0] * p[1] * p[2] // TILE + + +def least_padding_shape(grid): + """Shape that pads `grid` least (ties: the most cubic, then the first); how reference spans are tiled.""" + fit = [s for s in shapes() if all(a <= b for a, b in zip(s, grid))] + return min(fit or shapes(), key=lambda s: (_count(grid, s), max(s) / min(s), s)) + + +def span_tiles(start, grid, shape): + """[n_tiles, 128] packed row ids of one span, -1 on padding. Box order is (h, w, t) outer to inner, rows + inside a box are t, h, w row-major, real rows compacted to the front (so a tile's real rows are a prefix).""" + t, h, w = grid + tp, hp, wp = _padded(grid, shape) + cells = torch.full((tp, hp, wp), -1, dtype=torch.long) + cells[:t, :h, :w] = start + torch.arange(t * h * w).view(t, h, w) + st, sh, sw = shape + tiles = cells.view(tp // st, st, hp // sh, sh, wp // sw, sw).permute(2, 4, 0, 1, 3, 5).reshape(-1, TILE) + return torch.gather(tiles, 1, torch.argsort((tiles < 0).to(torch.int8), dim=1, stable=True)) + + +def build_layout(spans, seq_len, device="cpu"): + """One permutation of the packed sequence. `spans` = [(start, grid, shape)] in packed order, the target last; + every other row is global (text, audio, untiled conditions) and is chunked into tiles after the spans.""" + if not spans: + raise VedaError("the target span must be tiled") + covered, tiles, stop = torch.zeros(seq_len, dtype=torch.bool), [], 0 + for start, grid, shape in spans: + rows = grid[0] * grid[1] * grid[2] + if start < stop or start + rows > seq_len: + raise VedaError(f"span at {start} overlaps another or runs past {seq_len}") + stop = start + rows + covered[start:stop] = True + tiles.append(span_tiles(start, grid, shape)) + video = torch.cat(tiles) + rest = torch.nonzero(~covered).view(-1) + n_global = -(-rest.numel() // TILE) + tail = torch.full((n_global * TILE,), -1, dtype=torch.long) + tail[:rest.numel()] = rest + perm = torch.cat([video.view(-1), tail]) + count = (perm.view(-1, TILE) >= 0).sum(1).to(torch.int32) + partial = torch.nonzero((count > 0) & (count < TILE)).view(-1) + n_video = video.shape[0] + on = lambda x: x.to(device) + return SimpleNamespace( + n_video=n_video, n_ref=sum(t.shape[0] for t in tiles[:-1]), n_global=n_global, n_tiles=n_video + n_global, + ref_tokens=sum(g[0] * g[1] * g[2] for _, g, _ in spans[:-1]), + target_tokens=spans[-1][1][0] * spans[-1][1][1] * spans[-1][1][2], seq_len=seq_len, + valid_count=on(count), gather_index=on(perm.clamp(min=0)), scatter_index=on(torch.where(perm < 0, seq_len, perm)), + pad_slots=on(torch.nonzero(perm < 0).view(-1)), partial_video=on(partial[partial < n_video]), + kv_ok=on(count > 0), slot_valid=on((perm >= 0).to(torch.int32))) + + +def gather(x, lay, heads): + """[S, H, D] -> [N, H', D] in tile order, padding slots zero.""" + out = x[lay.gather_index[:, None], heads[None, :]] + if lay.pad_slots.numel(): + out.index_fill_(0, lay.pad_slots, 0) + return out + + +def scatter_(out, tiled, lay, heads): + """Tile-ordered rows back into out[S + 1, H, D]; padding lands in the spare row S.""" + out[lay.scatter_index[:, None], heads[None, :]] = tiled + + +# ── predictor: pool each video tile to [mean | max | min], score tile pairs ──────────────────────────────────── + +@torch.no_grad() +def pool(x, lay): + """[N, H', D] tile order -> [H', n_video, 3D] fp32 (mean, max, min over each tile's real rows; 0 for empty).""" + n, (heads, dim) = lay.n_video, x.shape[1:] + tiles = x[:n * TILE].view(n, TILE, heads, dim) + mean = tiles.sum(dim=1, dtype=torch.float32) / lay.valid_count[:n].clamp(min=1).to(torch.float32)[:, None, None] + lo, hi = torch.aminmax(tiles, dim=1) + part = lay.partial_video + if part.numel(): # padding rows are 0, a legal value: the partial tiles' extremes are redone masked + sub = tiles.index_select(0, part) + live = (torch.arange(TILE, device=x.device)[None, :] < lay.valid_count.index_select(0, part)[:, None])[:, :, None, None] + hi.index_copy_(0, part, sub.masked_fill(~live, float("-inf")).amax(dim=1)) + lo.index_copy_(0, part, sub.masked_fill(~live, float("inf")).amin(dim=1)) + feats = torch.cat([mean, hi.float(), lo.float()], dim=-1) + return torch.where(lay.kv_ok[:n, None, None], feats, 0.0).permute(1, 0, 2).contiguous() + + +@torch.no_grad() +def logits(fq, fk, pq, pk): + """[H', n_q, n_k] fp32 block logits: x_hat = pool @ P + mean, q_hat . k_hat / sqrt(D).""" + dim = pq.shape[-1] + qh = torch.bmm(fq, pq.to(device=fq.device, dtype=torch.float32)) + fq[..., :dim] + kh = torch.bmm(fk, pk.to(device=fk.device, dtype=torch.float32)) + fk[..., :dim] + return torch.bmm(qh, kh.transpose(1, 2)) / math.sqrt(dim) + + +# ── selection: top-k per query tile over the video x video quadrant ──────────────────────────────────────────── + +def column_blocks(lay, generated, reference): + """[(start, stop, budget, real tokens)]: the reference tiles (when tiled) and the target, each with its own top-k + (one pooled top-k would starve the references, the target's spatial neighbours always win).""" + target = (lay.n_ref, lay.n_video, generated, lay.target_tokens) + return [target] if lay.n_ref == 0 else [(0, lay.n_ref, reference, lay.ref_tokens), target] + + +@torch.no_grad() +def select(scores, lay, blocks): + """(index [H', n_video, K] int64 key tiles, keep [H', n_video, K] bool). Rules: a query tile always keeps its own + tile and that uses budget, empty key tiles are never kept, the fractional budget is spread by a Bresenham pattern.""" + heads, nv, dev = scores.shape[0], lay.n_video, scores.device + rows = torch.arange(nv, device=dev) + idxs, keeps = [], [] + for start, stop, budget, real in blocks: + n = stop - start + if keeps_all(budget): + idxs.append(torch.arange(start, stop, device=dev).expand(heads, nv, n)) + keeps.append(lay.kv_ok[start:stop].expand(heads, nv, n)) + continue + s = scores[:, :, start:stop].float().clone() + s.masked_fill_(~lay.kv_ok[None, None, start:stop], float("-inf")) + own = ((rows >= start) & (rows < stop))[None, :, None] + col = (rows - start).clamp(0, n - 1)[None, :, None].expand(heads, nv, 1) + s.scatter_(-1, col, torch.where(own, float("inf"), s.gather(-1, col))) + lo, hi, frac = _split(_per_row(budget, real, n), n) + vals, idx = torch.topk(s, hi, dim=-1, sorted=True) + allowed = lo + _bresenham(nv, float(frac)).to(device=dev, dtype=torch.long) + idxs.append(idx + start) + keeps.append((torch.arange(hi, device=dev)[None, None, :] < allowed[None, :, None]) & (vals > float("-inf"))) + return torch.cat(idxs, -1), torch.cat(keeps, -1) + + +@torch.no_grad() +def block_mask(index, keep, lay): + """[H', n_tiles, n_tiles] bool, rows are query tiles. Global tiles are dense both ways.""" + n, nv = lay.n_tiles, lay.n_video + mask = torch.zeros(index.shape[0], n, n, dtype=torch.bool, device=index.device) + mask[:, :nv].scatter_(2, index, keep) + mask[:, :, nv:] = lay.kv_ok[nv:] + mask[:, nv:, :] = lay.kv_ok + return mask + + +def tile_index_list(mask): + """A mask as the per-row list of kept tiles: (index [H', n, widest] int32 ascending, count [H', n] int32).""" + count = mask.sum(-1, dtype=torch.int32) + widest = int(count.max().item()) if count.numel() else 0 + order = torch.argsort(mask.to(torch.int8), dim=-1, descending=True, stable=True) + return order[..., :widest].to(torch.int32).contiguous(), count + + +def reference_attend(q, k, v, mask, lay): + """Exact block-sparse attention in fp32 on tile-ordered [N, H', D]: the ground truth and the test stand-in + for the kernel. O(N^2) memory, never for a real sequence.""" + allowed = mask.repeat_interleave(TILE, 1).repeat_interleave(TILE, 2) & lay.slot_valid.bool()[None, None, :] + probs = torch.softmax((torch.einsum("qhd,khd->hqk", q.float(), k.float()) * q.shape[-1] ** -0.5) + .masked_fill(~allowed, float("-inf")), dim=-1).nan_to_num(0.0) + out = torch.einsum("hqk,khd->qhd", probs, v.float()) * lay.slot_valid.to(torch.float32)[:, None, None] + return out.to(q.dtype) + + +# ── what the packed layout holds ──────────────────────────────────────────────────────────────────────────────── + +def _grid_of(position_ids, a, b): + """(T, H, W) when rows a:b are T, H, W row-major (checked against the positions), else None.""" + pos = position_ids[a:b] + axes = [torch.unique(pos[:, i]) for i in range(3)] + t, h, w = (x.numel() for x in axes) + if t * h * w != b - a: + return None + want = torch.stack(torch.meshgrid(*axes, indexing="ij"), dim=-1).reshape(-1, 3) + return (t, h, w) if torch.equal(want, pos.to(want.dtype)) else None + + +def describe_layout(layout): + """ComfyUI's H3 PackedLayout -> a hashable (seq_len, (target start, grid), ((ref start, grid), ...), skipped notes). + References (cond, ref_img) whose rows are not a verified row-major grid stay global, which is always correct.""" + try: + segments = [(int(a), int(b), kind) for a, b, kind in layout.segments] + _, lt, lh, lw, _ = layout.signature + seq_len = int(layout.seq_len) + except (AttributeError, TypeError, ValueError) as exc: + raise VedaError(f"not a MiniMax H3 packed layout ({exc})") from exc + video = [(a, b) for a, b, kind in segments if kind == "video"] + if len(video) != 1: + raise VedaError(f"expected one target video segment, found {len(video)}") + grid = (int(lt), int(lh) // 2, int(lw) // 2) + if grid[0] * grid[1] * grid[2] != video[0][1] - video[0][0]: + raise VedaError(f"the video segment does not match the latent grid {grid}") + pos, refs, skipped = getattr(layout, "position_ids", None), [], [] + for a, b, kind in segments: + if kind in ("cond", "ref_img"): + g = None if pos is None else _grid_of(pos, a, b) + (refs.append((a, g)) if g else skipped.append(f"{kind}[{a}:{b}]")) + return seq_len, (video[0][0], grid), tuple(refs), tuple(skipped) + + +# ── tile plans and the predictor file ────────────────────────────────────────────────────────────────────────── + +def _plan(name, grid, shape_list, head_shape): + return SimpleNamespace(name=name, grid=tuple(grid), shapes=[tuple(s) for s in shape_list], head_shape=head_shape, groups={}) + + +def _aspect(grid): + return grid[2] / grid[1] + + +def describe_grid(grid): + frames = 17 * ((grid[0] - 2) // 5) + 5 if grid[0] >= 2 else 1 # H3: 17n + 5 frames for a latent length + return f"{grid[2] * 32}x{grid[1] * 32}, {frames / FPS:.1f} s" + + +def choose_plan(plans, grid): + """(plan, exact, words): the plan with the nearest aspect, then nearest duration, then least padding. The H<->W + mirrored plans are candidates too. Never fails: a plan of another size still tiles any grid, it is only + outside what the predictor was trained on, and `words` says so.""" + grid = tuple(int(g) for g in grid) + mirrored = [_plan(p.name + "_T", (p.grid[0], p.grid[2], p.grid[1]), [(t, w, h) for t, h, w in p.shapes], p.head_shape) for p in plans] + target = math.log(_aspect(grid)) + best = min(plans + mirrored, key=lambda p: (round(abs(math.log(_aspect(p.grid)) - target), 3), abs(p.grid[0] - grid[0]), + sum(_count(grid, s) for s in p.shapes))) + if best.grid == grid: + return best, True, f"trained for this size ({describe_grid(grid)})" + return best, False, f"nearest trained size: {describe_grid(best.grid)} (this video: {describe_grid(grid)})" + + +def head_groups(plan, layer): + """[(shape, heads int64 CPU)] of one layer: heads that share a tile shape. Cached on the plan.""" + if layer not in plan.groups: + row = torch.tensor(plan.head_shape[layer]) + plan.groups[layer] = [(plan.shapes[i], torch.nonzero(row == i).view(-1)) for i in sorted(set(plan.head_shape[layer]))] + return plan.groups[layer] + + +def _budget_of(meta, prefix, fallback): + kind = meta.get(f"{prefix}_budget_kind") + if kind is None: + return fallback + try: + value = float(meta[f"{prefix}_budget_value"]) + except (KeyError, ValueError) as exc: + raise VedaError(f"{prefix}_budget_value: {exc}") from exc + if kind not in ("ratio", "tiles"): + raise VedaError(f"{prefix}_budget_kind: unknown kind {kind!r}") + return (value, None) if kind == "ratio" else (None, value) + + +def read_metadata(path): + from safetensors import safe_open + name = os.path.basename(path) + try: + with safe_open(path, framework="pt", device="cpu") as f: + meta = dict(f.metadata() or {}) + except Exception as exc: # noqa: BLE001 - safetensors raises several types + raise VedaError(f"{name} is not a readable safetensors file ({exc})") from exc + if meta.get("format") != FORMAT: + raise VedaError(f"{name} is not a Veda predictor (format {meta.get('format')!r}). Put only Veda predictor files in models/veda.") + return meta + + +def load_bundle(path): + """The predictor file, read strictly: every tensor must match its declared shape. Projections are held in bf16 on the + host and follow q's device at call time, so the predictor does not pin VRAM while ComfyUI offloads the model.""" + from safetensors import safe_open + meta, name = read_metadata(path), os.path.basename(path) + try: + layers, heads, dim = int(meta["num_layers"]), int(meta["num_heads"]), int(meta["head_dim"]) + keep = meta.get("keep_ratio") + generated = _budget_of(meta, "target", None if keep is None else (float(keep), None)) + if generated is None: + raise KeyError("target_budget_kind") + reference = _budget_of(meta, "ref", generated) + plan_json = json.loads(meta["plans"]) + except (KeyError, ValueError) as exc: + raise VedaError(f"{name}: incomplete metadata ({exc})") from exc + stored = meta.get("dtype", "float32") + if stored not in _STORAGE: + raise VedaError(f"{name}: unknown storage dtype {stored!r}") + with safe_open(path, framework="pt", device="cpu") as f: + tensors = {key: f.get_tensor(key) for key in f.keys()} + shape, proj = (heads, 3 * dim, dim), {"proj_q": [], "proj_k": []} + for layer in range(layers): + for kind, out in proj.items(): + key = f"layers.{layer}.{kind}" + value = tensors.pop(key, None) + if value is None or tuple(value.shape) != shape: + raise VedaError(f"{name}: {key} missing or not {list(shape)}") + if stored == "float8_e4m3fn": + scale = tensors.pop(key + ".__scale", None) + if scale is None or tuple(scale.shape) != (heads,): + raise VedaError(f"{name}: {key} has no per-head fp8 scale") + value = value.float() * scale.float().view(-1, 1, 1) + out.append(value.to(torch.bfloat16).contiguous()) + if tensors: + raise VedaError(f"{name}: unexpected tensors {sorted(tensors)[:4]}") + plans = [_plan(p["geometry"], p["grid"], [tuple(int(x) for x in s.split("x")) for s in p["shapes"]], [list(r) for r in p["head_shape"]]) + for p in plan_json.values()] + if not plans: + raise VedaError(f"{name}: no tile plans") + for p in plans: + if len(p.head_shape) != layers or len(p.head_shape[0]) != heads: + raise VedaError(f"{name}: plan {p.name} does not fit {layers} layers x {heads} heads") + return SimpleNamespace(path=path, layers=layers, heads=heads, dim=dim, generated=generated, reference=reference, + plans=plans, proj_q=proj["proj_q"], proj_k=proj["proj_k"], meta=meta) + + +def describe_bundle(b): + trained = describe_budget(b.generated) + ("" if b.reference == b.generated else f" / reference {describe_budget(b.reference)}") + return f"{os.path.basename(b.path)}, {b.layers} layers x {b.heads} heads, trained {trained}" + + +# ── one layer ─────────────────────────────────────────────────────────────────────────────────────────────────── + +@torch.no_grad() +def attention(q, k, v, layer, spec, plan, bundle, budgets, attend, cache, stats): + """One layer of Veda attention. q, k, v: [S, H, D] (post QK-norm and RoPE, any strides, any device). + For each group of heads sharing a tile shape, in memory-bounded chunks: gather to tile order, pool, score, select, + build the block mask, `attend(q_t, k_t, v_t, mask, layout)`, scatter back. Returns [S, H, D] in q's dtype. + `cache` holds the tile layouts (keyed by geometry, shape and device); `stats` accumulates what was computed.""" + seq_len, heads, dim = q.shape + dev = q.device + generated, reference = budgets + tile_refs = not keeps_all(reference) + out = q.new_zeros(seq_len + 1, heads, dim) + pq, pk = bundle.proj_q[layer].to(dev), bundle.proj_k[layer].to(dev) + for shape, group in head_groups(plan, layer): + key = (spec, shape, tile_refs, str(dev)) + if key not in cache: + spans = [(s, g, least_padding_shape(g)) for s, g in spec[2]] if tile_refs else [] + cache[key] = build_layout(spans + [(spec[1][0], spec[1][1], shape)], seq_len, dev) + lay = cache[key] + blocks = column_blocks(lay, generated, reference) + step = max(1, _CHUNK_BYTES // (lay.n_tiles * TILE * dim * q.element_size())) + for hs in group.to(dev).split(step): + q_t, k_t = gather(q, lay, hs), gather(k, lay, hs) + scores = logits(pool(q_t, lay), pool(k_t, lay), pq.index_select(0, hs), pk.index_select(0, hs)) + index, keep = select(scores, lay, blocks) + del scores + kept = keep.sum() + stats["kept"] = stats.get("kept", 0) + kept + stats["pairs"] = stats.get("pairs", 0) + len(hs) * lay.n_video ** 2 + stats["work"] = stats.get("work", 0) + kept + len(hs) * lay.n_global * (lay.n_tiles + lay.n_video) + stats["dense"] = stats.get("dense", 0) + len(hs) * (seq_len / TILE) ** 2 + stats["global"] = lay.n_global + mask = block_mask(index, keep, lay) + o_t = attend(q_t, k_t, gather(v, lay, hs), mask, lay) + scatter_(out, o_t, lay, hs) + return out[:seq_len] diff --git a/modules/loaders/veda_download.py b/modules/loaders/veda_download.py new file mode 100644 index 00000000..2043cffc --- /dev/null +++ b/modules/loaders/veda_download.py @@ -0,0 +1,74 @@ +"""The Veda predictor: where it lives, and an optional download on the person's own click. + +FunUI never hosts, bundles or mirrors the file. Its licence (the MiniMax H3 Community License) has conditions +on distribution, so the only route offered is a fetch straight from the official Hugging Face repository, +started by a button that sits next to the licence name and a link to its text. Streaming, the `.part` name and +the size and checksum checks are core/fetch.py's. +""" +import threading +from pathlib import Path + +from .._core import fetch, log + +SOURCE = "FunUI Veda" +HOST = "https://huggingface.co" +LICENSE = {"name": "MiniMax H3 Community License", + "url": f"{HOST}/Veda-Sparse/Minimax-H3-T2VA-Veda-8NFE-600Step-Preview/blob/main/LICENSE"} +# Sizes and checksums are the repositories' own listing (the Git LFS record of each file). +MODELS = ( + {"id": "t2va", "title": "T2VA / FL2VA", "hint": "Text, or first and last frames, to audio and video.", + "repo": "Veda-Sparse/Minimax-H3-T2VA-Veda-8NFE-600Step-Preview", "file": "minimax_h3_t2va_veda_8nfe_600step_preview_fp8.safetensors", + "size": 275415648, "sha256": "2a8d8845c5342756a2781e8e69563940e4bb573c9a40ebb534915ff8fd76573a"}, + {"id": "r2va", "title": "R2VA", "hint": "Reference images or videos to audio and video.", + "repo": "Veda-Sparse/Minimax-H3-R2VA-Veda-Preview", "file": "minimax_h3_r2va_veda_preview_fp8.safetensors", + "size": 275415752, "sha256": "23edace948693772be33d5cc945d302d2b037a29bd44bd57071bfbdd8b02517d"}, +) +_state, _lock = {}, threading.Lock() + + +def folder() -> Path: + import folder_paths + return Path(folder_paths.models_dir) / "veda" + + +def files() -> list: + """The predictor files in models/veda, by name. A `.part` is not one.""" + return sorted(p.name for p in folder().glob("*.safetensors")) if folder().is_dir() else [] + + +def listing() -> dict: + present = files() + rows = [] + for m in MODELS: + s = _state.get(m["id"]) or {} + row = {k: m[k] for k in ("id", "title", "hint", "file", "size")} + row["url"] = f"{HOST}/{m['repo']}/resolve/main/{m['file']}" + row["state"] = "downloading" if s.get("state") == "downloading" else "ready" if m["file"] in present else "error" if s.get("state") == "error" else "absent" + row.update(done=s.get("done", 0), error=s.get("error", "")) + rows.append(row) + return {"files": present, "models": rows, "license": LICENSE} + + +def download(model_id: str, get=None) -> bool: + """Start the download in a thread. False when the id is unknown, the file is there already, or it is already running.""" + spec = next((m for m in MODELS if m["id"] == model_id), None) + get = get or fetch.to_file + with _lock: + if spec is None or (_state.get(model_id) or {}).get("state") == "downloading" or spec["file"] in files(): + return False + _state[model_id] = {"state": "downloading", "done": 0, "error": ""} + + def work(): + try: + target = folder() + target.mkdir(parents=True, exist_ok=True) + get(f"{HOST}/{spec['repo']}/resolve/main/{spec['file']}", target / spec["file"], + lambda done, total: _state[model_id].update(done=done), size=spec["size"], sha256=spec["sha256"]) + _state[model_id] = {"state": "ready", "done": spec["size"], "error": ""} + log.info(SOURCE, f"Veda predictor {spec['title']} downloaded to {target / spec['file']}") + except Exception as exc: # noqa: BLE001 -- shown beside the button, and in the Log + _state[model_id] = {"state": "error", "done": 0, "error": f"{type(exc).__name__}: {exc}"} + log.warning(SOURCE, f"Veda predictor {spec['title']} could not be downloaded: {exc}") + + threading.Thread(target=work, name=f"funui-veda-{model_id}", daemon=True).start() + return True diff --git a/modules/loaders/veda_kernel.py b/modules/loaders/veda_kernel.py new file mode 100644 index 00000000..0bbbfb9b --- /dev/null +++ b/modules/loaders/veda_kernel.py @@ -0,0 +1,126 @@ +"""Block-sparse INT8 attention in Triton for Veda: walks only the kept 128-token tiles. + +UNTESTED on a GPU: written on a machine with no Triton or CUDA. `check` runs the kernel on a small problem +against the fp32 reference before any real use, and the caller falls back to dense attention with the reason +when it fails. + +Derived from SageAttention 1.0.6 (BSD-3-Clause, Copyright (c) 2024, Jintao Zhang, Haofeng Huang), +quant_per_block.py and attn_qk_int8_per_block.py, through Veda-on-ComfyUI's fork of them (MIT); the licence +texts are in VEDA-NOTICE.txt. Upstream's arithmetic is kept on purpose: INT8 Q and K with one scale per +block, the softmax scale folded into Q's quantisation so the loop can use exp2, P and V in fp16. + +What differs from upstream, as in Veda's fork: + * the key loop walks a per (head, query tile) list of kept tiles instead of the whole sequence; + * padding slots are masked by each tile's real-row count (real rows are a prefix of every tile), where + upstream loads out-of-range keys as zeros and lets them into the softmax; + * a tile is two key blocks of 64 (upstream's block), expanded on the host so the loop stays flat; + * tile-ordered [N, H, D] tensors throughout. Veda's optional TMA path (off by default there) is not ported. +""" +import torch +import triton +import triton.language as tl + +from .veda_core import TILE, block_mask, build_layout, column_blocks, gather, reference_attend, select, tile_index_list + +KEY_BLOCK = 64 +WARPS = 8 if torch.version.hip else 4 # measured by Veda on an RTX 5070 (4); 8 on ROCm +STAGES = 3 +LOG2E = 1.4426950408889634 + + +@triton.jit +def _quantize(X, Xq, Scale, stride_n, stride_h, stride_s, pre_scale, D: tl.constexpr, BLK: tl.constexpr): + """One INT8 scale per (block of BLK slots, head), upstream's rounding.""" + blk, head = tl.program_id(0), tl.program_id(1) + offs = head * stride_h + (blk * BLK + tl.arange(0, BLK))[:, None] * stride_n + tl.arange(0, D)[None, :] + x = tl.load(X + offs).to(tl.float32) * pre_scale + scale = tl.max(tl.abs(x)) / 127.0 + scale = tl.where(scale > 0.0, scale, 1.0) # an all-padding block has no scale; 1.0 keeps zeros zero + q = x / scale + q += 0.5 * tl.where(q >= 0, 1, -1) + tl.store(Xq + offs, q.to(tl.int8)) + tl.store(Scale + head * stride_s + blk, scale) + + +@triton.jit +def _attend(Q, K, V, Q_scale, K_scale, Index, Count, Valid, Out, stride_n, stride_h, stride_vn, stride_vh, + stride_qs, stride_ks, stride_ih, stride_iq, n_q_tiles, D: tl.constexpr, BLK: tl.constexpr, BLOCK_M: tl.constexpr): + tile, head = tl.program_id(0), tl.program_id(1) + offs_m = tile * BLOCK_M + tl.arange(0, BLOCK_M) + offs_n, offs_d = tl.arange(0, BLK), tl.arange(0, D) + q = tl.load(Q + head * stride_h + offs_m[:, None] * stride_n + offs_d[None, :]) + q_scale = tl.load(Q_scale + head * stride_qs + tile) + m_i = tl.full([BLOCK_M], -1e30, dtype=tl.float32) # finite: an all-padding block must not make -inf - -inf + l_i = tl.zeros([BLOCK_M], dtype=tl.float32) + acc = tl.zeros([BLOCK_M, D], dtype=tl.float32) + count = tl.load(Count + head * n_q_tiles + tile) + row = Index + head * stride_ih + tile * stride_iq + for i in range(0, count): + block = tl.load(row + i) + start = block * BLK + k = tl.load(K + head * stride_h + (start + offs_n)[None, :] * stride_n + offs_d[:, None]) + qk = tl.dot(q, k).to(tl.float32) * q_scale * tl.load(K_scale + head * stride_ks + block) + qk = tl.where(offs_n[None, :] < tl.load(Valid + block), qk, -float("inf")) + m_ij = tl.maximum(m_i, tl.max(qk, 1)) + p = tl.math.exp2(qk - m_ij[:, None]) + alpha = tl.math.exp2(m_i - m_ij) + l_i = l_i * alpha + tl.sum(p, 1) + v = tl.load(V + head * stride_vh + (start + offs_n)[:, None] * stride_vn + offs_d[None, :]) + acc = acc * alpha[:, None] + acc += tl.dot(p.to(tl.float16), v, out_dtype=tl.float16) + m_i = m_ij + l_i = tl.where(l_i > 0.0, l_i, 1.0) # an empty row would divide by zero; selection always keeps the diagonal + tl.store(Out + head * stride_h + offs_m[:, None] * stride_n + offs_d[None, :], (acc / l_i[:, None]).to(Out.type.element_ty)) + + +def _quant(x, block, pre_scale=1.0): + slots, heads, dim = x.shape + out = torch.empty_like(x, dtype=torch.int8) + scale = torch.empty((heads, slots // block), dtype=torch.float32, device=x.device) + _quantize[(slots // block, heads)](x, out, scale, x.stride(0), x.stride(1), scale.stride(0), pre_scale=pre_scale, D=dim, BLK=block, num_warps=4) + return out, scale + + +def attend(q, k, v, mask, lay): + """The block-sparse attention callable veda_core.attention takes. q, k, v: [N, H', D] contiguous, bf16/fp16, tile order; + mask: [H', n_tiles, n_tiles] bool. Returns [N, H', D] in q's dtype.""" + slots, heads, dim = q.shape + index, count = tile_index_list(mask & lay.kv_ok) + sub = TILE // KEY_BLOCK # kept tiles -> kept key blocks, the kernel's unit + blocks = ((index.to(torch.int32) * sub)[..., None] + torch.arange(sub, device=q.device, dtype=torch.int32)).flatten(-2).contiguous() + valid = (lay.valid_count[:, None] - torch.arange(sub, device=q.device) * KEY_BLOCK).clamp_(0, KEY_BLOCK).flatten().to(torch.int32).contiguous() + q8, q_scale = _quant(q, TILE, dim ** -0.5 * LOG2E) + k8, k_scale = _quant(k, KEY_BLOCK) + v16, out = v.to(torch.float16), torch.empty_like(q) + _attend[(slots // TILE, heads)]( + q8, k8, v16, q_scale, k_scale, blocks, (count * sub).contiguous(), valid, out, q.stride(0), q.stride(1), v16.stride(0), + v16.stride(1), q_scale.stride(0), k_scale.stride(0), blocks.stride(0), blocks.stride(1), n_q_tiles=slots // TILE, D=dim, + BLK=KEY_BLOCK, BLOCK_M=TILE, num_warps=WARPS, num_stages=STAGES) + return out + + +def check(device, dtype=torch.bfloat16): + """Run the kernel on a small problem (history span, target span with partial tiles, global rows, per-row budgets) + and compare it with the fp32 reference. Returns None when it holds, else why not (a kernel that ignores the mask + fails too: an earlier FlashAttention build accepted one and ignored it).""" + spans = [(70, (1, 8, 16), (1, 8, 16)), (238, (5, 8, 8), (4, 4, 8))] + lay = build_layout(spans, 238 + 320, device) + gen = torch.Generator().manual_seed(0) + x = torch.randn(3, lay.seq_len, 2, 128, generator=gen) * 1.5 # peaky softmax: skipped blocks matter + heads = torch.arange(2, device=device) + q, k, v = (gather(t.to(device=device, dtype=dtype), lay, heads) for t in x) + scores = torch.randn(2, lay.n_video, lay.n_video, generator=gen).to(device) + mask = block_mask(*select(scores, lay, column_blocks(lay, (None, 2), (None, 1))), lay) + try: + got = attend(q, k, v, mask, lay) + except Exception as exc: # noqa: BLE001 - kernels fail in many ways; report, do not crash + return f"the Triton kernel failed its self-test: {type(exc).__name__}: {exc}" + real = lay.slot_valid.bool() + want = reference_attend(q, k, v, mask, lay)[real].float() + dense = reference_attend(q, k, v, torch.ones_like(mask) & lay.kv_ok, lay)[real].float() + got = got[real].float() + if (got - dense).abs().max() < 0.5 * (dense - want).abs().max(): + return "the Triton kernel ignores the block mask on this device (computes dense attention)" + limit = 0.05 * max(1.0, want.abs().max().item()) # INT8 with per-block scales: ~1-3% is the format's floor + err = (got - want).abs().max().item() + return None if err <= limit else f"the Triton kernel gives wrong results on this device (max error {err:.3g}, allowed {limit:.3g})" diff --git a/modules/training/lora/taggers.py b/modules/training/lora/taggers.py index 047fd2f0..d6e88e2e 100644 --- a/modules/training/lora/taggers.py +++ b/modules/training/lora/taggers.py @@ -10,9 +10,8 @@ from __future__ import annotations import csv import re import threading -import urllib.request -from ..._core import config, log +from ..._core import config, fetch, log SOURCE = "FunUI LoRA training" FILES = ("model.onnx", "selected_tags.csv") @@ -65,24 +64,7 @@ def listing() -> list[dict]: return rows -def _fetch(url: str, dest, progress) -> None: - """Stream a URL to a file through a .part name, so a cut connection leaves nothing that looks finished.""" - req = urllib.request.Request(url, headers={"User-Agent": "FunUI"}) - part = dest.with_name(dest.name + ".part") - with urllib.request.urlopen(req, timeout=60) as res, open(part, "wb") as out: - total = int(res.headers.get("Content-Length") or 0) - done = 0 - while True: - chunk = res.read(1 << 20) - if not chunk: - break - out.write(chunk) - done += len(chunk) - progress(done, total) - if total and done != total: - part.unlink(missing_ok=True) - raise OSError(f"the download stopped at {done} of {total} bytes") - part.replace(dest) +_fetch = fetch.to_file def download(tagger_id: str, fetch=None) -> bool: