Blocks settings adapt to the detected model: block count read from the file header, hint and clamp in the panel

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Tw9fjYZhHqk196LykrWy3T

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.

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();

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") {

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");
+});

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,

core/dit_hooks.py

@@ -129,6 +129,19 @@ def block_count(model) -> int:
         return 0
 
 
+def count_indexed(keys, prefixes):
+    """How many numbered entries (`<prefix>0.`, `<prefix>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]).
 

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"),

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}

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

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,

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},

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},

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},

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},

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},

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},