diff --git a/CHANGELOG.md b/CHANGELOG.md index 455ca917..632d3fa0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,7 @@ ### Added +- **Block settings know the model**: the file's header gives its block count, the Blocks hint says "This model has blocks 0-47", and a list is cut to those blocks when you commit it. - **Temporal style: auto** picks a style for each shot from its own words (motion words punch, calm words hold, loop words loop). - **Randomizer, REINS and Block repeat on LTX** (new traits `block_scope` and `block_rows`); the other block-hook features stay H3-only. Whole-span repeat runs as each-block on LTX. - **Attention explorer on LTX**: the depth groups and the clip regions (four time bands of the picture, and the sound) work; Prompt focus stays H3-only and is left out on LTX. diff --git a/app/ui/features/settings/engine.js b/app/ui/features/settings/engine.js index 74bdc66c..759ad783 100644 --- a/app/ui/features/settings/engine.js +++ b/app/ui/features/settings/engine.js @@ -31,7 +31,7 @@ export const engine = (app) => function mount() { checked: ps.cameraOff(), onChange: (on) => ps.setCameraOff(on).then(draw) }) : null, ...(category === "guidance" && gOff ? [c.hint.default({ text: "Guidance enhancements are disabled for this project. Their settings are hidden, not lost." })] : withSettings.filter((m) => (m.category || "") === category).flatMap((m) => [c.label.section({ text: m.title || m.id }), - ...Object.entries(m.settings).filter(([, spec]) => whenSatisfied(values, m.id, spec.when)).map(([name, spec]) => settingRow(m.id, name, spec, values, change))])), + ...Object.entries(m.settings).filter(([, spec]) => whenSatisfied(values, m.id, spec.when)).map(([name, spec]) => settingRow(m.id, name, spec, values, change, ps.blocks?.()))])), ].filter(Boolean)); } draw(); diff --git a/app/ui/features/settings/field.js b/app/ui/features/settings/field.js index 77881791..69aff834 100644 --- a/app/ui/features/settings/field.js +++ b/app/ui/features/settings/field.js @@ -5,12 +5,28 @@ import { composer as c } from "../../composer/composer.js"; /** A `when` clause names sibling settings on the same module that must hold a value for the row to show. */ export const whenSatisfied = (values, id, when) => !when || Object.entries(when).every(([k, v]) => (Array.isArray(v) ? v : [v]).includes((values[id] || {})[k])); -export function settingRow(id, name, spec, values, onChange) { +/** A block list ("2-10,29-46") held to a model's blocks 0..count-1: entries past the end are cut, a range is shortened, an entry wholly past the end is dropped. */ +export function clampBlocks(text, count) { + if (!Number.isInteger(count) || count < 1) return text; + const last = count - 1, out = []; + for (const part of String(text || "").split(",").map((x) => x.trim()).filter(Boolean)) { + const m = /^(\d+)(?:\s*-\s*(\d+))?$/.exec(part); + if (!m) { out.push(part); continue; } // not a number or range: left for the module to say + const lo = Number(m[1]), hi = m[2] === undefined ? lo : Number(m[2]); + const a = Math.min(lo, hi), b = Math.min(Math.max(lo, hi), last); + if (a > last) continue; + out.push(a === b ? String(a) : `${a}-${b}`); + } + return out.join(","); +} + +export function settingRow(id, name, spec, values, onChange, blocks = null) { const label = spec.label || name; - const hint = spec.hint ? spec.hint + (spec.unit ? ` (${spec.unit})` : "") : spec.unit || ""; + const fits = spec.model_blocks && Number.isInteger(blocks) && blocks > 0; + const hint = (spec.hint ? spec.hint + (spec.unit ? ` (${spec.unit})` : "") : spec.unit || "") + (fits ? ` This model has blocks 0-${blocks - 1}.` : ""); const now = values[id] && values[id][name]; const value = now !== undefined ? now : spec.default; - const set = (v) => onChange(id, name, v); + const set = (v) => onChange(id, name, fits && typeof v === "string" ? clampBlocks(v, blocks) : v); let control; if (spec.type === "bool") return c.toggle.default({ label, hint, checked: Boolean(value), onChange: set }); // a toggle carries its own label and hint if (spec.type === "int" || spec.type === "float") { diff --git a/app/ui/features/settings/field.test.js b/app/ui/features/settings/field.test.js new file mode 100644 index 00000000..4d941a58 --- /dev/null +++ b/app/ui/features/settings/field.test.js @@ -0,0 +1,17 @@ +import test from "node:test"; +import assert from "node:assert/strict"; +import { clampBlocks } from "./field.js"; + +test("a block list is held to the model's blocks", () => { + assert.equal(clampBlocks("2-10,29-46", 48), "2-10,29-46"); + assert.equal(clampBlocks("2-10,29-46", 40), "2-10,29-39"); + assert.equal(clampBlocks("43-44", 40), ""); + assert.equal(clampBlocks("5,50,60-70,38-45", 40), "5,38-39"); + assert.equal(clampBlocks("30-20", 25), "20-24"); +}); + +test("nothing is changed when the count is unknown, and odd entries are left for the module to say", () => { + assert.equal(clampBlocks("2-99", null), "2-99"); + assert.equal(clampBlocks("2-99", 0), "2-99"); + assert.equal(clampBlocks("x, 3", 10), "x,3"); +}); diff --git a/app/ui/shell/pipeline_state.js b/app/ui/shell/pipeline_state.js index 0fff546b..8960da87 100644 --- a/app/ui/shell/pipeline_state.js +++ b/app/ui/shell/pipeline_state.js @@ -43,13 +43,16 @@ export function createPipelineState(rawApi, { answerMs = ANSWER_MS } = {}) { // Which traits the pipeline's chosen model has; null when that is not knowable (no file, or no model // module recognised it, or the one that did names no traits): null means "do not filter", while a list, // even an empty one, hides every module needing a trait it lacks. undefined: the probe failed, try again. + let modelBlocks = null; // how many transformer blocks the chosen file has, when its header says let modelNote = null; // why nothing is hidden for the chosen file, when no model module recognised it async function probeModelTraits(pipelineSlots) { const filename = currentModelFile(pipelineSlots); modelNote = null; + modelBlocks = null; if (!filename) return null; try { const data = await API.probeFamily(filename); + modelBlocks = Number.isInteger(data.blocks) && data.blocks > 0 ? data.blocks : null; if (!data.detected && data.reason) modelNote = `${data.reason}, so every module is offered, including ones this model cannot use.`; return data.detected && Array.isArray(data.traits) ? data.traits : null; } catch (_) { @@ -755,6 +758,8 @@ export function createPipelineState(rawApi, { answerMs = ANSWER_MS } = {}) { skipped: () => skipped, /** The traits the chosen model has (null: not known, so nothing is hidden for it). */ traits: () => lastTraits || null, + /** How many transformer blocks the chosen model file has (null: not known). */ + blocks: () => modelBlocks, incomplete: () => incomplete, refused: () => refused, queueable: () => queueable, diff --git a/core/dit_hooks.py b/core/dit_hooks.py index 40b30df8..3d4703e7 100644 --- a/core/dit_hooks.py +++ b/core/dit_hooks.py @@ -129,6 +129,19 @@ def block_count(model) -> int: return 0 +def count_indexed(keys, prefixes): + """How many numbered entries (`0.`, `1.`...) a header has under any of these prefixes, or None. + A model module answers `probe_blocks` with this over the name its transformer blocks carry.""" + seen = set() + for key in keys: + for prefix in prefixes: + if key.startswith(prefix): + head = key[len(prefix):].split(".", 1)[0] + if head.isdigit(): + seen.add(int(head)) + return len(seen) or None + + def parse_blocks(spec, n_blocks: int): """"40" | "38-42" | "10,40,44" | "" -> (sorted indices, [problems]). diff --git a/core/probe.py b/core/probe.py index d454be27..374d283a 100644 --- a/core/probe.py +++ b/core/probe.py @@ -125,8 +125,16 @@ def detect(path: Path, registry=None) -> dict: except Exception as exc: # noqa: BLE001 from . import log log.once(f"provider:{spec.id}.probe_traits:{type(exc).__name__}", log.WARNING, f"{spec.id}.probe_traits", f"failed while answering -- {type(exc).__name__}: {exc}") + blocks = None + probe_blocks = spec.provides.get("probe_blocks") + if probe_blocks is not None: + try: + blocks = probe_blocks(keyset) + except Exception as exc: # noqa: BLE001 + from . import log + log.once(f"provider:{spec.id}.probe_blocks:{type(exc).__name__}", log.WARNING, f"{spec.id}.probe_blocks", f"failed while answering -- {type(exc).__name__}: {exc}") return {"module": spec.id, "title": spec.title, "detected": True, - "reason": f"{p.name}: {spec.title}", "traits": traits} + "reason": f"{p.name}: {spec.title}", "traits": traits, "blocks": blocks} return {"module": None, "title": None, "detected": False, "reason": (f"{p.name}: could not be checked, {', '.join(broken)} failed while looking at it (see the Log)" if broken else f"{p.name}: no installed model module recognises this file"), diff --git a/modules/models/ltx/__init__.py b/modules/models/ltx/__init__.py index b66bc297..3b73bf7f 100644 --- a/modules/models/ltx/__init__.py +++ b/modules/models/ltx/__init__.py @@ -29,7 +29,7 @@ import comfy.model_management import torch from comfy.nested_tensor import NestedTensor -from ..._core import log, streams as _streams, traits as _traits, vae_tiles +from ..._core import dit_hooks as _dit_hooks, log, streams as _streams, traits as _traits, vae_tiles from .pipeline import optional_nodes as _optional_nodes, presets as _pipeline_presets, upgrade as _upgrade has_block = _traits.has_block @@ -70,6 +70,11 @@ def probe_traits(keys) -> list: return ["audio_stream", "temporal_latent", "temporal_compression", "ltx_av", "block_scope", "block_rows"] +def probe_blocks(keys): + """How many transformer blocks the file has, read off its header.""" + return _dit_hooks.count_indexed(keys, tuple(f"{p}transformer_blocks." for p in _PREFIXES)) + + def traits(model): if not is_ltx(model): return () @@ -150,5 +155,5 @@ def decode(latent, model=None, vae=None, audio_vae=None, tile_size=0): TRAITS = traits PROVIDES = {"empty_latent": empty_latent, "decode": decode, "detect": detect, - "probe_traits": probe_traits, "pipeline_presets": _pipeline_presets, "optional_nodes": _optional_nodes, + "probe_traits": probe_traits, "probe_blocks": probe_blocks, "pipeline_presets": _pipeline_presets, "optional_nodes": _optional_nodes, "video_stream": _streams.av_video_stream, "upgrade_slots": _upgrade} diff --git a/modules/models/ltx/tests/test_ltx.py b/modules/models/ltx/tests/test_ltx.py index 331e8e37..ad26e7ef 100644 --- a/modules/models/ltx/tests/test_ltx.py +++ b/modules/models/ltx/tests/test_ltx.py @@ -276,3 +276,18 @@ def test_an_image_to_video_pipeline_saved_before_the_aspect_source_gets_it(): assert latent["inputs"]["width"] == ["aspect", 0] and not any(r["input"] in ("width", "height") for r in latent["roles"]) assert upgrade(fixed) == fixed # and it is done once assert upgrade(now) == now + + +def test_the_block_count_is_read_off_the_header(comfyui): + m = _ltx() + keys = {f"model.diffusion_model.transformer_blocks.{i}.attn1.to_q.weight" for i in range(48)} | {"model.diffusion_model.audio_adaln_single.linear.weight"} + assert m.probe_blocks(keys) == 48 + assert m.probe_blocks({"audio_adaln_single.linear.weight"}) is None + assert m.probe_blocks({"transformer_blocks.0.x", "transformer_blocks.9.x"}) == 2 # counts distinct numbers, as the file names them + + +def test_the_h3_block_count_ignores_its_token_refiner(comfyui): + from importlib import import_module + h3 = import_module("modules.models.minimax_h3") + keys = {f"blocks.{i}.attn.q.weight" for i in range(50)} | {f"token_refiner.blocks.{i}.x" for i in range(2)} + assert h3.probe_blocks(keys) == 50 diff --git a/modules/models/minimax_h3/__init__.py b/modules/models/minimax_h3/__init__.py index 839f0850..ffa65d30 100644 --- a/modules/models/minimax_h3/__init__.py +++ b/modules/models/minimax_h3/__init__.py @@ -28,7 +28,7 @@ import torch from comfy.nested_tensor import NestedTensor from comfy_extras.nodes_minimax_h3 import temporal_shape -from ..._core import log, traits as _traits +from ..._core import dit_hooks as _dit_hooks, log, traits as _traits from . import vae_decode as _decode from .training import adapt as _training_adapt from .pipeline import optional_nodes as _optional_nodes, presets as _pipeline_presets @@ -102,6 +102,11 @@ def detect(keys) -> bool: H3_ONLY = ("shot_prompts", "sparse_attention") +def probe_blocks(keys): + """How many transformer blocks the file has, read off its header.""" + return _dit_hooks.count_indexed(keys, tuple(f"{p}blocks." for p in _PREFIXES)) + + def probe_traits(keys) -> list: """What `traits()` would say for a loaded H3 model, read before anything loads. @@ -233,7 +238,7 @@ def empty_latent(model, width, height, length, batch_size=1): TRAITS = traits PROVIDES = {"training_adapt": _training_adapt, "empty_latent": empty_latent, "decode": decode, "detect": detect, - "pipeline_presets": _pipeline_presets, "optional_nodes": _optional_nodes, "probe_traits": probe_traits, + "pipeline_presets": _pipeline_presets, "optional_nodes": _optional_nodes, "probe_traits": probe_traits, "probe_blocks": probe_blocks, "target_rows": _target_rows, "anchor_pin": _anchor_pin, "latent_upscaler": _load_upscaler, "rescale_conditioning": _rescale_pins, "text_tokenizer": _prompt_layout.text_tokenizer, diff --git a/modules/sampling/attention_bandit/__init__.py b/modules/sampling/attention_bandit/__init__.py index 6030fc4a..a257c180 100644 --- a/modules/sampling/attention_bandit/__init__.py +++ b/modules/sampling/attention_bandit/__init__.py @@ -81,7 +81,7 @@ SETTINGS = { "when": {"enabled": True}, }, "blocks": { - "type": "text", "default": "2-10,29-46", + "type": "text", "model_blocks": True, "default": "2-10,29-46", "label": "Blocks", "hint": "e.g. 2-10,29-46. Split by depth into groups that learn separately.", "when": {"enabled": True}, diff --git a/modules/sampling/attention_temperature/__init__.py b/modules/sampling/attention_temperature/__init__.py index 632f804b..b172df2a 100644 --- a/modules/sampling/attention_temperature/__init__.py +++ b/modules/sampling/attention_temperature/__init__.py @@ -37,7 +37,7 @@ SETTINGS = { "when": {"enabled": True}, }, "blocks": { - "type": "text", "default": "2-10", + "type": "text", "model_blocks": True, "default": "2-10", "label": "Blocks", "hint": "e.g. 2-10 or 29-46. Keep away from blocks a steering feature uses.", "when": {"enabled": True}, diff --git a/modules/sampling/block_repeat/__init__.py b/modules/sampling/block_repeat/__init__.py index d354eed2..82bf5b8f 100644 --- a/modules/sampling/block_repeat/__init__.py +++ b/modules/sampling/block_repeat/__init__.py @@ -44,7 +44,7 @@ SETTINGS = { "hint": "Re-rolls part of the model's thinking for a different take on the same prompt.", }, "blocks": { - "type": "text", "default": "31-40", + "type": "text", "model_blocks": True, "default": "31-40", "label": "Blocks", "hint": "e.g. 31-40, or 20,35. Somewhere in 20-41 works best.", "when": {"enabled": True}, diff --git a/modules/sampling/identity_bank/__init__.py b/modules/sampling/identity_bank/__init__.py index 86b0a82f..7206bcd8 100644 --- a/modules/sampling/identity_bank/__init__.py +++ b/modules/sampling/identity_bank/__init__.py @@ -45,7 +45,7 @@ SETTINGS = { "when": {"enabled": True}, }, "blocks": { - "type": "text", "default": "20,30", + "type": "text", "model_blocks": True, "default": "20,30", "label": "Blocks", "hint": "Where the look is kept and applied, e.g. 14,20,30. Each block kept costs about 100 MB per remembered clip.", "when": {"enabled": True}, diff --git a/modules/sampling/q_steer/__init__.py b/modules/sampling/q_steer/__init__.py index 1d270e9a..50569416 100644 --- a/modules/sampling/q_steer/__init__.py +++ b/modules/sampling/q_steer/__init__.py @@ -39,7 +39,7 @@ SETTINGS = { "when": {"enabled": True}, }, "blocks": { - "type": "text", "default": "43-44", + "type": "text", "model_blocks": True, "default": "43-44", "label": "Blocks", "hint": "Where it steers, e.g. 43-44 or 0,1,47-49. Every block keeps learning either way.", "when": {"enabled": True}, diff --git a/modules/sampling/reins/__init__.py b/modules/sampling/reins/__init__.py index 4715aeea..927a6791 100644 --- a/modules/sampling/reins/__init__.py +++ b/modules/sampling/reins/__init__.py @@ -39,7 +39,7 @@ SETTINGS = { "when": {"enabled": True}, }, "blocks": { - "type": "text", "default": "43-44", + "type": "text", "model_blocks": True, "default": "43-44", "label": "Blocks", "hint": "Where it steers, e.g. 43-44 or 0. Every block keeps learning either way.", "when": {"enabled": True},